xclele/yolov10
1
0
Fork 0
mirror of https://github.com/THU-MIG/yolov10.git synced 2025-05-23 21:44:22 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix overlap_mask (#651)

Browse Source 
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Laughing 2023-01-27 03:42:40 +08:00 committed by GitHub
parent 15b3b0365a
commit dc9502c700
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 9 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
ultralytics/yolo/data/dataset.py
View File

@ -145,7 +145,9 @@ class YOLODataset(BaseDataset):
normalize=True, normalize=True,
return_mask=self.use_segments, return_mask=self.use_segments,
return_keypoint=self.use_keypoints, return_keypoint=self.use_keypoints,
batch_idx=True)) batch_idx=True,
mask_ratio=hyp.mask_ratio,
mask_overlap=hyp.overlap_mask))
return transforms return transforms
def close_mosaic(self, hyp): def close_mosaic(self, hyp):
@ -155,7 +157,9 @@ class YOLODataset(BaseDataset):
normalize=True, normalize=True,
return_mask=self.use_segments, return_mask=self.use_segments,
return_keypoint=self.use_keypoints, 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): def update_labels_info(self, label):
"""custom your label format here""" """custom your label format here"""

6
ultralytics/yolo/v8/segment/train.py
View File

@ -110,11 +110,11 @@ class SegLoss(Loss):
target_scores, target_scores_sum, fg_mask) target_scores, target_scores_sum, fg_mask)
for i in range(batch_size): for i in range(batch_size):
if fg_mask[i].sum(): 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: 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: 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]] xyxyn = target_bboxes[i][fg_mask[i]] / imgsz[[1, 0, 1, 0]]
marea = xyxy2xywh(xyxyn)[:, 2:].prod(1) marea = xyxy2xywh(xyxyn)[:, 2:].prod(1)
mxyxy = xyxyn * torch.tensor([mask_w, mask_h, mask_w, mask_h], device=self.device) mxyxy = xyxyn * torch.tensor([mask_w, mask_h, mask_w, mask_h], device=self.device)