jiayuanz3
commented
9 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!
@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?