About training datasets

[EleFreelance]

Big Guy, that's an amazing piece of work, but I was wondering what data set did you use to train your U^2net model?

[samhaswon]

Thanks! It's a custom dataset I made myself. I'm still working on it, hoping to reduce bias and increase accuracy.

[EleFreelance]

Thank you very much for your reply, do I have a chance to share your custom data set? I want to replicate your effect and do some exploration of model structure optimization.

[samhaswon]

Not currently. I'm still working on it, so I would just be changing it again. If you want something to test with, you can use an already public dataset like this one:
https://zenodo.org/records/5834916#.YkYQCShBw2x

It is technically different but it would give you a good starting point. This dataset doesn't include the eyes and lips if I remember correctly. Regardless, it would give you a good starting point for doing like you are talking about. You could even use the weights I have released to make your own dataset with a bit of cleanup.

[EleFreelance]

Ok, you gave me a very good idea. but I have a question about Google's MediaPipe model that you mentioned in the readme. Which one is better for creating the training datasets, or the model you release?In other words, is your current release model more effective than Google's MediaPipe model for skin detection?

[samhaswon]

I would just use one of the ones I released. You can try the Google one if you want, the code is right there, but I couldn't ever get what I would call "good" results. I wouldn't even call what they have in the documentation for it good. But, it's a tiny model for what it is doing.

[samhaswon]

Also, if you're building an unlabelled dataset, this might help you out: https://github.com/samhaswon/bing_image_downloader
Or, you could use what I forked it from: https://github.com/gurugaurav/bing_image_downloader

My fork is more optimized for Stable Diffusion dataset building whereas the upstream repo will bring in more watermarked images (but less overall) if you are looking for that. Mine also allows you to use a list of queries (demoed using a file) to reduce the duplicates you would have to filter through and make it easier to build a large dataset.

[FaisalAldoraihem]

Stumbled on this model and loving the results!

wondering if you're more comfortable sharing the dataset used now, though? i'd like to try training a more mobile-friendly model

[samhaswon]

Glad it's working well for you! I'm still not comfortable with sharing the dataset, but I am working on a new release. If you want, I could train something smaller while I'm at it. I also plan to release quantized versions this time if memory is the issue.

Alternatively, you could use the PyTorch weights of the U2NetP model (the small, more mobile-friendly one) to change the input size to 320x320 and re-export to ONNX to reduce the computation time for the model. You can see the onnx export code here: https://github.com/samhaswon/rembg-trainer-cuda/blob/cc3d788d82576a281ea254b1e850e914889b6954/u2net_train.py#L229

If you do use Torch, you can grab the model class here: https://github.com/samhaswon/rembg-trainer-cuda/tree/master/model
Then loading is basically just:

# import the model and torch then:
net = U2NET(3, 1)
net.load_state_dict(torch.load(MODEL_PATH, map_location=torch.device(DEVICE)))

For reference, here are inference times from my laptop using the CPU and PyTorch.

Model	Input Size	Time per Inference (s)
U2NET	320x320	0.4973
U2NET	512x512	1.4684
U2NET	1024x1024	6.1395
U2NETP	320x320	0.2565
U2NETP	512x512	0.7775
U2NETP	1024x1024	3.3855

Using ONNX will be faster based on my usage of the model. My main usage is the ONNX U2NetP model at 512x512, which takes 0.5015s for my full pipeline on the same CPU and 0.3645s on an old Intel iGPU (HD Graphics 530). According to potentially unreliable sources, that iGPU has about half of the theoretical performance of my phone's, so definitely try to use the GPU if you can.

[FaisalAldoraihem]

A quantized would definitely be fun to test. Thank you for the offer!

And I'll definitely give the weight switch a shot. For now, I'm using a DeepLabV3 head classifier with mobilenet_v3 as the backbone since I would be running it on four separate cropped images at most from a camera feed.

[samhaswon]

If you want to continue this on Discord, which would work better for me, my username is the same there as it is here (samhaswon).

I've also got a model training that's deeplabv3_mobilenet_v3_large from PyTorch, so we'll see how that does soon.

[samhaswon]

@FaisalAldoraihem This next release is taking longer than I expected it to, so I added some weights for the DeepLabV3 + MobileNetV3 model to a prerelease: https://github.com/samhaswon/skin_segmentation/releases/tag/v0.0.2-pre

[FaisalAldoraihem]

That's perfect, thank you!

[samhaswon]

You're welcome!