Training on custom data and the predictions are like gray-scale instead of full black and white mask.

[Sukanth07]

So i prepared my model where it has images of solid background colors and i want to remove only that background while preserving the other design elements including the abstract shapes and splashes. I used the BiRefNet-general-epoch_244.pth and also the model.safetensor file from Huggingface as the pretrained base model and the swin_large as the backbone. And set batch size to 1 for 2500 image-mask pairs dataset with gradient accumulation step to 4. All images are 1024x1024 resolution. And i set finetune to 200 epochs which will take around 6 days to complete the finetuning. But the problem is, i am getting the prediction results as gray-scale image instead of full black and white mask. I didnt altered any of the model architecture and pre&post processing steps. But i dont know why i get this issue. Can anyone help me with this.

[ZhengPeng7]

What😂? You mean the outputs of the model are gray-scale images instead of the mask?
I didn't use grad-acc, but it should not cause that consequence.

If you didn't change any codes of the model, you may need to check your GT data -- if they are gray-scale images.
Besides, you can use my colab demo or the notebooks in the tutorials folder, set your weights as the one to check if the problem in in the training or some pre/post-processing.

[Sukanth07]

Thanks for the reply @ZhengPeng7 . By the way due to my GPU constraints i set that training settings let me migrate to cloud and check once. And my GTs are proper black and white masks as in the uploaded screenshot. Let me debug and try once again and if any changes or result i get i will post here. Thanks once again!

[ZhengPeng7]

You're welcome :)

And 6 days is too long. If you have ~500 images for training, and since it's fine-tuning with BiRefNet-general (if you download the BiRefNet-general-epoch_244.pth locally, you don't need the weights on HF -- duplicated), you can train on custom data for only 30 epochs, which should already show noticeable effects.

[ZhengPeng7]

Besides, your GTs seem to be matting-level annotations? 0~1 instead of 0 / 1. If so, using the matting version of BiRefNet weights may help. And remember to set the task in config.py to Matting.

[Sukanth07]

I successfully trained the matting model on my dataset with a batch size of 2 for 50 epochs with 2500 custom dataset. The results are fantastic—I'm seeing significant improvements and exceptional quality on most samples!

However, I’ve noticed some limitations in challenging cases, particularly:

• Fine structures like thin hairs or fur often have pixelated edges or incomplete masks.
• Some background color artifacts remain visible along the fine edges, especially in high-contrast areas.

I understand that under resource constraints and with a single model for whole thing, there are limits to what can be achieved. But is there any recommended approach to further improve edge quality and reduce background color bleeding in these edge cases?

I’ve attached a couple of sample results for reference.

[ZhengPeng7]

Glad to see the improvement from you!

For your two questions:

1. Higher resolution can help solve this kind of thing. But it's higher cost, and you can try the next point.
2. Using some foreground estimation algorithms can help remove some of the background colors in regions with transparency and some white boundaries. You can refer to my codes.