Update datasets.py

ultralytics · Mar 5, 2021 · 04cb07f · 04cb07f
1 parent e2e785d
commit 04cb07f
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/utils/datasets.py b/utils/datasets.py
@@ -666,7 +666,7 @@ def load_mosaic(self, index):
     labels4, segments4 = [], []
     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.randint(0, self.n - 1) for _ in range(3)]  # 3 additional image indices
+    indices = [index] + [self.indices[random.randint(0, self.n - 1)] for _ in range(3)]  # 3 additional image indices
     for i, index in enumerate(indices):
         # Load image
         img, _, (h, w) = load_image(self, index)
@@ -721,7 +721,7 @@ def load_mosaic9(self, index):
 
     labels9, segments9 = [], []
     s = self.img_size
-    indices = [index] + [random.randint(0, self.n - 1) for _ in range(8)]  # 8 additional image indices
+    indices = [index] + [self.indices[random.randint(0, self.n - 1)] for _ in range(8)]  # 8 additional image indices
     for i, index in enumerate(indices):
         # Load image
         img, _, (h, w) = load_image(self, index)