fix overlap_mask (ultralytics#651)

Co-authored-by: Glenn Jocher <[email protected]>
bigcash · Jan 26, 2023 · dc9502c · dc9502c
1 parent 15b3b03
commit dc9502c
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 5 deletions.
diff --git a/ultralytics/yolo/data/dataset.py b/ultralytics/yolo/data/dataset.py
@@ -145,7 +145,9 @@ def build_transforms(self, hyp=None):
                    normalize=True,
                    return_mask=self.use_segments,
                    return_keypoint=self.use_keypoints,
-                   batch_idx=True))
+                   batch_idx=True,
+                   mask_ratio=hyp.mask_ratio,
+                   mask_overlap=hyp.overlap_mask))
         return transforms
 
     def close_mosaic(self, hyp):
@@ -155,7 +157,9 @@ def close_mosaic(self, hyp):
                    normalize=True,
                    return_mask=self.use_segments,
                    return_keypoint=self.use_keypoints,
-                   batch_idx=True))
+                   batch_idx=True,
+                   mask_ratio=hyp.mask_ratio,
+                   mask_overlap=hyp.overlap_mask))
 
     def update_labels_info(self, label):
         """custom your label format here"""

diff --git a/ultralytics/yolo/v8/segment/train.py b/ultralytics/yolo/v8/segment/train.py
@@ -110,11 +110,11 @@ def __call__(self, preds, batch):
                                               target_scores, target_scores_sum, fg_mask)
             for i in range(batch_size):
                 if fg_mask[i].sum():
-                    mask_idx = target_gt_idx[i][fg_mask[i]] + 1
+                    mask_idx = target_gt_idx[i][fg_mask[i]]
                     if self.overlap:
-                        gt_mask = torch.where(masks[[i]] == mask_idx.view(-1, 1, 1), 1.0, 0.0)
+                        gt_mask = torch.where(masks[[i]] == (mask_idx + 1).view(-1, 1, 1), 1.0, 0.0)
                     else:
-                        gt_mask = masks[batch_idx == i][mask_idx]
+                        gt_mask = masks[batch_idx.view(-1) == i][mask_idx]
                     xyxyn = target_bboxes[i][fg_mask[i]] / imgsz[[1, 0, 1, 0]]
                     marea = xyxy2xywh(xyxyn)[:, 2:].prod(1)
                     mxyxy = xyxyn * torch.tensor([mask_w, mask_h, mask_w, mask_h], device=self.device)