search

King-HAW / GMS

Official repository of Generative Medical Segmentation
https://arxiv.org/abs/2403.18198
39 stars 3 forks source link

How to train multiclass segmentation? #3

Open Jeonghyun66 opened 1 month ago

Jeonghyun66 commented 1 month ago

Hi, I want to train GMS model with msd pancreas dataset. And it has 3 classes(background, pancreas, tumor). How can I train those 3 classes with your code? I don't know which variable indicates the num_classes. Thank you for sharing this code. :)

King-HAW commented 1 month ago

Hi, thanks for your attention. Currently, GMS only supports binary segmentation results. If you want to apply GMS on the multi-label segmentation dataset, 1) you can treat each foreground and background as a binary segmentation task, that being said you have to run GMS two times on the MSD dataset, 2) follow the GSS paper and change the ground truth to the RGB image (e.g., green for the pancreas and red for the tumor), use the RGB ground truth to train the model.