search

SuperMedIntel / Medical-SAM-Adapter

Adapting Segment Anything Model for Medical Image Segmentation
GNU General Public License v3.0
1.02k stars 90 forks source link

training code only trains REFUGE dataset for cup segmentation #90

Closed LJQCN101 closed 9 months ago

LJQCN101 commented 10 months ago

@jiayuanz3

By looking at the recently supported REFUGE dataset, I noticed that only the masks of cup are passed through the training loop. The masks of disc are not referenced at all.

At line 213 of dataset.py:

return { 'image':img, 'multi_rater': multi_rater_cup, 'multi_rater_disc': multi_rater_disc, 'mask_cup': mask_cup, 'mask_disc': mask_disc, 'label': mask_cup, 'p_label':point_label, 'pt_cup':pt_cup, 'pt_disc':pt_disc, 'pt':pt_cup, 'image_meta_dict':image_meta_dict, }

It is specifically using mask_cup for label. But mask_disc is not referenced outside of dataset.py. So is it not possible to do multi class segmentation?

jiayuanz3 commented 10 months ago

Heya, thanks for reaching out! Yes, the existing code architecture doesn't support multi-class segmentation. The data loader we have is particularly well-suited for tasks involving either disc or cup segmentation. To address this, you might consider adding a multi-head component in the final layer. While we primarily focus on single-class datasets and may not update the original code to include multi-class functionality, we definitely encourage and welcome your contributions. If you can integrate this feature without conflicting with the existing code, we would be more than happy to consider your pull request as a valuable addition from a contributor!