Remove any Background image - TensorFlowJs

TensorFlow General Discussion
,
1

Friends, I have a big problem.
Today with tensorflowjs/body-pix I can remove the background from an image and leave just one person.

Turns out I have a motorcycle parts system and the background photos have ‘noises’ and I wanted to remove that background and leave just the part so my model can tell me which part it is.

How would I get TensorFlowJs to remove the background and leave just the part once I have the pre-trained model of the parts?

Is there anything like Body-pix that does this?

3

Welcome to the forum Cristian! Thanks for posting your question and for trying out TensorFlow.js!

So our bodypix model is currently trained on human bodies as you have guessed. If you want to make your own segmentation model for a custom thing you will need to start collecting annotated training data (1000s of images) with examples of the thing you are trying to recognize and ensure you have a pixel level segmentation available for each image from which the model can learn from.

Do you have the training data? That would be your first task.

You can then choose a segmentation model you want to retrain for it to learn from your data instead.

To do that you have a number of options:

  1. Find an unfrozen model that has already been defined and retrain it with your data (easier but may have to start in Python first and then convert to TensorFlow.js with our converter).

A good tutorial for (1) is by our community expert Hugo Zanini:

madewithtfjs Interview:

Write up:

Code here:

  1. Write the code yourself to define the model using the TensorFlow.js APIs and train that (harder)

Maybe @Shanqing_Cai has some thoughts here though as he actually worked on BodyPix!

3 Likes
4

Thank you, I’m going to watch the video and delve into the topic to try to carry out this endeavor.

5

Good advice, @Jason ! I think these are the best paths to follow.

From what you said @Cristian_Heleno Cristian, I guess that RefineNet will be able to solve your problem. The model was trained in the KITTI dataset and can recognize the following classes:

  • Airplane
  • Bicycle
  • Bird
  • Boat
  • Bottle
  • Bus
  • Car
  • Cat
  • Chair
  • CowDining Table
  • Dog
  • Horse
  • Motorbike
  • Person
  • Potted Plant
  • Sheep
  • Sofa
  • Train
  • TV

If you need to train a custom model, I recommend using the Computer Vision Annotation Tool (CVAT) on the Semantic Segmentation mode and follow this repo instructions to retrain a RefineNet model using TensorFlow.

Then, you can easily convert the trained model to Tensorflow.js using the tf.js converter.

1 Like
6

Sounds good - do share with us your final result if you get it working for your usecase! Would love to see!

1 Like