HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

Concat all segments by default for multi-part masks (#16826)

Browse source 
Co-authored-by: Ultralytics Assistant <135830346+UltralyticsAssistant@users.noreply.github.com>
Co-authored-by: Laughing-q <1185102784@qq.com>
Co-authored-by: Laughing <61612323+Laughing-q@users.noreply.github.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Mohammed Yasin 2024-11-25 23:39:14 +08:00 • committed by GitHub
parent c5fd0bb48a
commit 0727b43ca6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 10 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

14
ultralytics/utils/ops.py
View file

@ -783,23 +783,29 @@ def regularize_rboxes(rboxes):
return torch.stack([x, y, w_, h_, t], dim=-1) # regularized boxes return torch.stack([x, y, w_, h_, t], dim=-1) # regularized boxes
def masks2segments(masks, strategy="largest"): def masks2segments(masks, strategy="all"):
""" """
It takes a list of masks(n,h,w) and returns a list of segments(n,xy). It takes a list of masks(n,h,w) and returns a list of segments(n,xy).
Args: Args:
masks (torch.Tensor): the output of the model, which is a tensor of shape (batch_size, 160, 160) masks (torch.Tensor): the output of the model, which is a tensor of shape (batch_size, 160, 160)
strategy (str): 'concat' or 'largest'. Defaults to largest strategy (str): 'all' or 'largest'. Defaults to all
Returns: Returns:
segments (List): list of segment masks segments (List): list of segment masks
""" """
from ultralytics.data.converter import merge_multi_segment
segments = [] segments = []
for x in masks.int().cpu().numpy().astype("uint8"): for x in masks.int().cpu().numpy().astype("uint8"):
c = cv2.findContours(x, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0] c = cv2.findContours(x, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0]
if c: if c:
if strategy == "concat": # concatenate all segments if strategy == "all": # merge and concatenate all segments
c = np.concatenate([x.reshape(-1, 2) for x in c]) c = (
np.concatenate(merge_multi_segment([x.reshape(-1, 2) for x in c]))
if len(c) > 1
else c[0].reshape(-1, 2)
)
elif strategy == "largest": # select largest segment elif strategy == "largest": # select largest segment
c = np.array(c[np.array([len(x) for x in c]).argmax()]).reshape(-1, 2) c = np.array(c[np.array([len(x) for x in c]).argmax()]).reshape(-1, 2)
else: else: