search

ZhengPeng7 / BiRefNet

[CAAI AIR'24] Bilateral Reference for High-Resolution Dichotomous Image Segmentation
https://www.birefnet.top
MIT License
991 stars 77 forks source link

Custom dataset fine-tuning #49

Closed htaoruan closed 1 month ago

htaoruan commented 1 month ago

Great work, I'm trying to fine tune a DIS for other categories such as hair, changes were made in dataset.py to make sure the data loads without problems, and would like to ask which model is better to base the fine tuning on, and is the fine tuning setup tested below? --resume ... /weights/BiRefNet-general-epoch_244.pth Anything to note?

ZhengPeng7 commented 1 month ago

What do you mean by other categories? No specified are required in the existing training.

Currently, I have not conducted experiments on which weights are better for fine-tuning... But that 'general one' is the most likely to work best.

And about the hair, which is a relatively special thing here -- do you want to do the matting task (labels are in 0~1) or dichotomous segmentation (labels are in 0 or 1)? If it's in the former condition, I would recommend also trying the weights for the portrait segmentation.

htaoruan commented 1 month ago

Thanks for your reply, would like to confirm that the fine-tuning based on BiRefNet-portrait-epoch_145.pth only requires modifying the resume parameter in mutil-gpu train right?

ZhengPeng7 commented 1 month ago

You may also need to turn off the IoU fine-tuning in config.py, if your data is on matting.

ZhengPeng7 commented 1 month ago

Feel free to re-open it if you still have questions.