From aa8b1098dd4a3776e2f3ce588d7047932f2f7d74 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 10 Apr 2020 16:28:59 -0700 Subject: [PATCH] adapt mosaic to img channel count --- utils/datasets.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/utils/datasets.py b/utils/datasets.py index 3d831e71..8f1341e3 100755 --- a/utils/datasets.py +++ b/utils/datasets.py @@ -527,7 +527,6 @@ def load_mosaic(self, index): labels4 = [] s = self.img_size xc, yc = [int(random.uniform(s * 0.5, s * 1.5)) for _ in range(2)] # mosaic center x, y - img4 = np.full((s * 2, s * 2, 3), 114, dtype=np.uint8) # base image with 4 tiles indices = [index] + [random.randint(0, len(self.labels) - 1) for _ in range(3)] # 3 additional image indices for i, index in enumerate(indices): # Load image @@ -535,6 +534,7 @@ def load_mosaic(self, 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 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