glenn-jocher commited on
Commit
9a9333d
1 Parent(s): b203c9b

add replicate() to datasets.py

Browse files
Files changed (1) hide show
  1. utils/datasets.py +21 -1
utils/datasets.py CHANGED
@@ -310,7 +310,6 @@ class LoadImagesAndLabels(Dataset): # for training/testing
310
  self.mosaic_border = [-img_size // 2, -img_size // 2]
311
  self.stride = stride
312
 
313
-
314
  # Define labels
315
  self.label_files = [x.replace('images', 'labels').replace(os.path.splitext(x)[-1], '.txt')
316
  for x in self.img_files]
@@ -629,6 +628,9 @@ def load_mosaic(self, index):
629
  # np.clip(labels4[:, 1:] - s / 2, 0, s, out=labels4[:, 1:]) # use with center crop
630
  np.clip(labels4[:, 1:], 0, 2 * s, out=labels4[:, 1:]) # use with random_affine
631
 
 
 
 
632
  # Augment
633
  # img4 = img4[s // 2: int(s * 1.5), s // 2:int(s * 1.5)] # center crop (WARNING, requires box pruning)
634
  img4, labels4 = random_affine(img4, labels4,
@@ -641,6 +643,23 @@ def load_mosaic(self, index):
641
  return img4, labels4
642
 
643
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
644
  def letterbox(img, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True):
645
  # Resize image to a 32-pixel-multiple rectangle https://github.com/ultralytics/yolov3/issues/232
646
  shape = img.shape[:2] # current shape [height, width]
@@ -765,6 +784,7 @@ def cutout(image, labels):
765
  box2_area = (b2_x2 - b2_x1) * (b2_y2 - b2_y1) + 1e-16
766
 
767
  # Intersection over box2 area
 
768
  return inter_area / box2_area
769
 
770
  # create random masks
 
310
  self.mosaic_border = [-img_size // 2, -img_size // 2]
311
  self.stride = stride
312
 
 
313
  # Define labels
314
  self.label_files = [x.replace('images', 'labels').replace(os.path.splitext(x)[-1], '.txt')
315
  for x in self.img_files]
 
628
  # np.clip(labels4[:, 1:] - s / 2, 0, s, out=labels4[:, 1:]) # use with center crop
629
  np.clip(labels4[:, 1:], 0, 2 * s, out=labels4[:, 1:]) # use with random_affine
630
 
631
+ # Replicate
632
+ # img4, labels4 = replicate(img4, labels4)
633
+
634
  # Augment
635
  # img4 = img4[s // 2: int(s * 1.5), s // 2:int(s * 1.5)] # center crop (WARNING, requires box pruning)
636
  img4, labels4 = random_affine(img4, labels4,
 
643
  return img4, labels4
644
 
645
 
646
+ def replicate(img, labels):
647
+ # Replicate labels
648
+ h, w = img.shape[:2]
649
+ boxes = labels[:, 1:].astype(int)
650
+ x1, y1, x2, y2 = boxes.T
651
+ s = ((x2 - x1) + (y2 - y1)) / 2 # side length (pixels)
652
+ for i in s.argsort()[:round(s.size * 0.5)]: # smallest indices
653
+ x1b, y1b, x2b, y2b = boxes[i]
654
+ bh, bw = y2b - y1b, x2b - x1b
655
+ yc, xc = int(random.uniform(0, h - bh)), int(random.uniform(0, w - bw)) # offset x, y
656
+ x1a, y1a, x2a, y2a = [xc, yc, xc + bw, yc + bh]
657
+ img[y1a:y2a, x1a:x2a] = img[y1b:y2b, x1b:x2b] # img4[ymin:ymax, xmin:xmax]
658
+ labels = np.append(labels, [[labels[i, 0], x1a, y1a, x2a, y2a]], axis=0)
659
+
660
+ return img, labels
661
+
662
+
663
  def letterbox(img, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True):
664
  # Resize image to a 32-pixel-multiple rectangle https://github.com/ultralytics/yolov3/issues/232
665
  shape = img.shape[:2] # current shape [height, width]
 
784
  box2_area = (b2_x2 - b2_x1) * (b2_y2 - b2_y1) + 1e-16
785
 
786
  # Intersection over box2 area
787
+
788
  return inter_area / box2_area
789
 
790
  # create random masks