adapt mosaic to img channel count

This commit is contained in:
Glenn Jocher 2020-04-10 16:28:59 -07:00
parent 6736d7d125
commit aa8b1098dd
1 changed files with 1 additions and 1 deletions

View File

@ -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