Updated segmentation dataloader to work with >255 instances (see ultr…

…alytics/yolov5#9461 and PR #9493 of the yolov5 repo) (#866)
WongKinYiu · Oct 4, 2022 · 1dc3461 · 1dc3461
1 parent 89428f8
commit 1dc3461
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/seg/utils/segment/dataloaders.py b/seg/utils/segment/dataloaders.py
@@ -308,7 +308,8 @@ def polygons2masks(img_size, polygons, color, downsample_ratio=1):
 
 def polygons2masks_overlap(img_size, segments, downsample_ratio=1):
     """Return a (640, 640) overlap mask."""
-    masks = np.zeros((img_size[0] // downsample_ratio, img_size[1] // downsample_ratio), dtype=np.uint8)
+    masks = np.zeros((img_size[0] // downsample_ratio, img_size[1] // downsample_ratio),
+            dtype=np.int32 if len(segments) > 255 else np.uint8)
     areas = []
     ms = []
     for si in range(len(segments)):