A modified Attention Unet model with post-process for retina optic disc segmention.
The performance of our model on Messidor-dataset：
I use a modified Attention Unet which input of model is 128x128pix image patches. To know more about attn-unet, please see the paper. When sampling the patches, I focus the algorithm get samples around optic disc. The patches is like that:
so the groundtruth is:
The model is trained on DRION dataset. 90 images to train. 19 images to test.
To get the groundtruth of DRION, I write a convert tool, you can find in DRION-DB-PythonTools.
Pretrained can be downloaded here. Extract them to dir Dataset.
When directly use unet model, we often get some error predictions. So I use a post-process algorithm:
- predicted area can't be to small.
- minimum bounding rectangle's height/width or width/height should be in 0.45~2.5
lefted area is the final output. The problem of this algorithm is that the parameters not self-adjusting, so you have to change them if input image is larger or smaller than before.
The structure is based on my own DL_Segmention_Template. Difference between this project and the template is that we have metric module in dir: perception/metric/. To get more Information about the structure please see readme in DL_Segmention_Template.
You can find model parameter in configs/segmention_config.json.
please run main_trainer.py first time, then you will get data_route in experiment dir. Put your data in there, now you can run main_trainer.py again to train a model.
The model is trained with DRION dataset on my own desktop (intel i7-7700hq, 24g, gtx1050 2g) within 30 minutes. Dataset
If u want to test your own image, put your image to (OpticDisc)/test/origin,and change the img_type of predict settings in configs/segmention_config.json, run main_test.py to get your result. The result is in (OpticDisc)/test/result