SHI-Labs / OneFormer

OneFormer: One Transformer to Rule Universal Image Segmentation, arxiv 2022 / CVPR 2023
https://praeclarumjj3.github.io/oneformer
MIT License
1.41k stars 128 forks source link

Segmentation label map #42

Closed rose-jinyang closed 1 year ago

rose-jinyang commented 1 year ago

Hello How are you? Thanks for contributing to this project. I want to get one segmentation map containing all the class labels rather than binary mask for each instance. How can I get it?

praeclarumjj3 commented 1 year ago

Hi, you can run a loop to aggregate all the predicted binary masks:

# create an all-zeros mask
single_channel_mask = torch.zeros_like(image)

# loop through all instance masks
for cls_id, mask in zip(result.pred_classes, result.pred_masks):
    mask *= cls_id
    single_channel_mask = torch.max(single_channel_mask, mask)
rose-jinyang commented 1 year ago

Thanks for your help.