Skip to content

Commit

Permalink
ultralytics#20 Added extra traditional CV input
Browse files Browse the repository at this point in the history
  • Loading branch information
@manole-alexandru
manole-alexandru committed Apr 25, 2023
1 parent 5e299b3 commit a0a1074
Showing 1 changed file with 32 additions and 10 deletions.
42 changes: 32 additions & 10 deletions utils/dataloaders.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -659,7 +659,7 @@ def __getitem__(self, index):
mosaic = self.mosaic and random.random() < hyp['mosaic']
if mosaic:
# Load mosaic
img, labels, seg = self.load_mosaic(index)
img, labels, seg, pre = self.load_mosaic(index)
shapes = None

# MixUp augmentation
Expand All @@ -671,14 +671,15 @@ def __getitem__(self, index):

else:
# Load image
img, (h0, w0), (h, w), seg = self.load_image(index)
img, (h0, w0), (h, w), seg, pre = self.load_image(index)

# Letterbox
shape = self.batch_shapes[self.batch[index]] if self.rect else self.img_size # final letterboxed shape
img, ratio, pad = letterbox(img, shape, auto=False, scaleup=self.augment)
# mask_shape = [shape[0] // 2, shape[1] // 2] if self.rect else self.img_size // 2
mask_shape = [shape[0], shape[1]] if self.rect else self.img_size
seg, _, _ = letterbox(seg, mask_shape, auto=False, scaleup=self.augment, color=(0, 0, 0))
pre, _, _ = letterbox(pre, mask_shape, auto=False, scaleup=self.augment, color=(0, 0, 0))
shapes = (h0, w0), ((h / h0, w / w0), pad) # for COCO mAP rescaling

labels = self.labels[index].copy()
Expand All @@ -696,12 +697,15 @@ def __getitem__(self, index):
seg, _, _ = random_perspective(seg, labels, degrees=hyp['degrees'], translate=hyp['translate'], scale=hyp['scale'],
shear=hyp['shear'], perspective=hyp['perspective'], random_parameters=random_parameters,
border_value=(0, 0, 0))
pre, _, _ = random_perspective(pre, labels, degrees=self.hyp['degrees'],
translate=self.hyp['translate'],
scale=self.hyp['scale'], shear=self.hyp['shear'],
perspective=self.hyp['perspective'],
border=self.mosaic_border, random_parameters=random_parameters)
if index < 50:
cv2.imwrite('runs/image' + str(index) + '.jpg', img)
cv2.imwrite('runs/mask' + str(index) + '.png', seg)
h, w = seg.shape[0], seg.shape[1]
resized = cv2.resize(seg, (w * 2, h * 2), interpolation=cv2.INTER_AREA)
cv2.imwrite('runs/maskr' + str(index) + '.png', resized)
cv2.imwrite('runs/pre' + str(index) + '.png', pre)

nl = len(labels) # number of labels
if nl:
Expand All @@ -720,13 +724,15 @@ def __getitem__(self, index):
if random.random() < hyp['flipud']:
img = np.flipud(img)
seg = np.flipud(seg)
pre = np.flipud(pre)
if nl:
labels[:, 2] = 1 - labels[:, 2]

# Flip left-right
if random.random() < hyp['fliplr']:
img = np.fliplr(img)
seg = np.fliplr(seg)
pre = np.fliplr(pre)
if nl:
labels[:, 1] = 1 - labels[:, 1]

Expand All @@ -740,7 +746,12 @@ def __getitem__(self, index):

# Convert
img = img.transpose((2, 0, 1))[::-1] # HWC to CHW, BGR to RGB
pre = cv2.cvtColor(pre, cv2.COLOR_BGR2GRAY)
pre = np.expand_dims(pre, axis=0)
img = np.concatenate(img, pre, axis=1)
print('EXTRA INPUT IMAGE SHAPE:', img.shape)
img = np.ascontiguousarray(img)

if not np.all((seg >= 0)):
print('\n POST AUGMENTATION MASK ERROR \n')
# Convert Mask
Expand All @@ -756,6 +767,7 @@ def load_image(self, i):
# Loads 1 image from dataset index 'i', returns (im, original hw, resized hw)
im, f, fn = self.ims[i], self.im_files[i], self.npy_files[i],
seg = None
preprocessed = None
if im is None: # not cached in RAM
if fn.exists(): # load npy
im = np.load(fn)
Expand All @@ -766,8 +778,14 @@ def load_image(self, i):
seg_path = '/'.join(split_path)
seg = cv2.imread(seg_path)

split_path2 = f.split('/')
split_path2[4] = 'preprocessed'
preprocessed_path = '/'.join(split_path2)
preprocessed = cv2.imread(preprocessed_path)

im = cv2.imread(f) # BGR

assert preprocessed is not None, f'Preprocessed Image Not Found {preprocessed_path}'
assert seg is not None, f'Segmentation Mask Not Found {seg_path}'
assert im is not None, f'Image Not Found {f}'
h0, w0 = im.shape[:2] # orig hw
Expand All @@ -777,7 +795,7 @@ def load_image(self, i):
im = cv2.resize(im, (math.ceil(w0 * r), math.ceil(h0 * r)), interpolation=interp)
seg = cv2.resize(seg, (math.ceil(w0 * r), math.ceil(h0 * r)), interpolation=interp)
return im, (h0, w0), im.shape[:2], seg # im, hw_original, hw_resized
return self.ims[i], self.im_hw0[i], self.im_hw[i], seg # im, hw_original, hw_resized
return self.ims[i], self.im_hw0[i], self.im_hw[i], seg, preprocessed # im, hw_original, hw_resized

def cache_images_to_disk(self, i):
# Saves an image as an *.npy file for faster loading
Expand All @@ -788,19 +806,20 @@ def cache_images_to_disk(self, i):
def load_mosaic(self, index):
# YOLOv5 4-mosaic loader. Loads 1 image + 3 random images into a 4-image mosaic
labels4, segments4 = [], []
seg4 = []
seg4, pre4 = [], []
s = self.img_size
yc, xc = (int(random.uniform(-x, 2 * s + x)) for x in self.mosaic_border) # mosaic center x, y
indices = [index] + random.choices(self.indices, k=3) # 3 additional image indices
random.shuffle(indices)
for i, index in enumerate(indices):
# Load image
img, _, (h, w), seg = self.load_image(index)
img, _, (h, w), seg, preprocessed = self.load_image(index)

# place img in img4
if i == 0: # top left
img4 = np.full((s * 2, s * 2, img.shape[2]), 114, dtype=np.uint8) # base image with 4 tiles
seg4 = np.full((s * 2, s * 2, img.shape[2]), 114, dtype=np.uint8)
pre4 = np.full((s * 2, s * 2, img.shape[2]), 114, dtype=np.uint8)
x1a, y1a, x2a, y2a = max(xc - w, 0), max(yc - h, 0), xc, yc # xmin, ymin, xmax, ymax (large image)
x1b, y1b, x2b, y2b = w - (x2a - x1a), h - (y2a - y1a), w, h # xmin, ymin, xmax, ymax (small image)
elif i == 1: # top right
Expand All @@ -815,6 +834,7 @@ def load_mosaic(self, index):

img4[y1a:y2a, x1a:x2a] = img[y1b:y2b, x1b:x2b] # img4[ymin:ymax, xmin:xmax]
seg4[y1a:y2a, x1a:x2a] = seg[y1b:y2b, x1b:x2b]
pre4[y1a:y2a, x1a:x2a] = preprocessed[y1b:y2b, x1b:x2b]
padw = x1a - x1b
padh = y1a - y1b

Expand Down Expand Up @@ -846,8 +866,10 @@ def load_mosaic(self, index):
seg4, _, _ = random_perspective(seg4, labels4, segments4, degrees=self.hyp['degrees'], translate=self.hyp['translate'],
scale=self.hyp['scale'], shear=self.hyp['shear'], perspective=self.hyp['perspective'],
border=self.mosaic_border, random_parameters=random_parameters)

return img4, labels4, seg4
pre4, _, _ = random_perspective(pre4, labels4, segments4, degrees=self.hyp['degrees'], translate=self.hyp['translate'],
scale=self.hyp['scale'], shear=self.hyp['shear'], perspective=self.hyp['perspective'],
border=self.mosaic_border, random_parameters=random_parameters)
return img4, labels4, seg4, pre4

def load_mosaic9(self, index):
# YOLOv5 9-mosaic loader. Loads 1 image + 8 random images into a 9-image mosaic
Expand Down

0 comments on commit a0a1074

Please sign in to comment.