search

WAMAWAMA / TNSCUI2020-Seg-Rank1st

This is the source code of the 1st place solution for segmentation task in MICCAI 2020 TN-SCUI challenge.
481 stars 78 forks source link

how to train a multi class segment? #3

Closed hitxujian closed 4 years ago

hitxujian commented 4 years ago

thanks for your sharing! if i want to train a multi class segment,how i custom your code

WAMAWAMA commented 4 years ago

tks for your issue👍 there are 2 key parts you need to modify (without the preprocessing part) :

  1. you can find two dataloader functions in /step2to4_train_validate_inference/loader/data_loader.py, which named get_loader and get_loader_difficult, and you need to modify the code of the 'processing GT' part in this script (as well as the img_mask_aug.py) to ensure that GT is loaded successfully
  2. the loss functions only support binary class segmentation, so you also need to replace the current loss function with a multi-class loss function in the solver.py to ensure that loss is calculated correctly

These changes may not be easy to implement if you are not familiar with my code, so I strongly recommend that you familiarize yourself with the code with the demo dataset (DDTI) before making any changes, and I will consider modifying the code to accommodate multi-class segmentation in the future