Title: Use rembg with different models to remove the background from images in Python

[Pictures with their backgrounds removed by different rembg models in Python]

My previous post Use rembg to remove the background from images in Python did a good job removing the background form one image but had some problems with a second image. It seemed to think parts of the image were shadows so it preserved them as semi-transparent areas.

The rembg library supports multiple artificial intelligence models that have been trained on different kinds of data and they may do a better or worse job at removing the background from a particular image. This example tests eight of the pre-built models to see which most effectively removes the background from the problematic image.

Sessions

Normally when you call rembg's remove function, it implicitly creates a session that loads a default model and then uses that model to process the image. To use a different model, you need to create a session yourself, passing it the name of the model that you would like to use.

When you use a model, the program downloads the model and saves it in your .u2net directory. Some of the larger models take a significant amount of time to download, up to 40 to 50 seconds with my computer and internet connection. Those larger models also take a while to process an image.

If you're processing multiple images, you can make things faster by using the same session for all of them. Then the program doesn't need to load the model each time.

This example, however, processes one image with several different models so I'm not going to do that. Here's the remove_background function revised to take a model as a parameter with the changes highlighted in blue.

def remove_background(infile, outfile, model='u2netp'): curdir = os.getcwd() infile = os.path.join(curdir, infile) outfile = os.path.join(curdir, outfile) # Start the model session. session = new_session(model) input_image = Image.open(infile) output_image = remove(input_image, session=session) output_image.save(outfile)

This version of the function uses the model name to create a session and then passes the session into the remove function.

Calling remove_background

The following code shows how the program tests eight different models for a single image.

class RembgModelsApp: def __init__(self): self.window = tk.Tk() self.window.title('use_rembg_models') self.window.protocol('WM_DELETE_WINDOW', self.kill_callback) # Make a list of models. models = [ 'isnet-general-use', 'sam', 'birefnet-general', 'birefnet-general-lite', 'birefnet-dis', 'birefnet-hrsod', 'birefnet-cod', 'birefnet-massive', ] # Process the image with each of the models. self.images = [] num_cols = 4 row = 0 col = 0 infile = 'cake2.png' for model_name in models: outfile = infile.replace('.png', f'_{model_name}.png') print(f'{infile} --> {outfile}') # Process the file. remove_background(infile, outfile, model_name) # Display the result. image = Image.open(outfile) image = draw_over_checkerboard(image) image = ImageTk.PhotoImage(image) self.images.append(image) label = tk.Label(self.window, image=image) label.grid(row=row, column=col, padx=1, pady=1) # Show the model name in a label below the image. label = tk.Label(self.window, text=model_name) label.grid(row=row+1, column=col, padx=1, pady=1) # Move to the next row if necssary. col += 1 if col >= num_cols: row += 2 col = 0 # Display the window. self.window.focus_force() self.window.mainloop()

The code initializes tkinter as before and then loops through eight model names. For each model, the code calls the remove_background function. It displays the result on top of a checkerboard and shows the model's name in a label below the image. The rest of the code just makes tkinter arrange the results nicely.

Conclusion

Using a model isn't too hard, although it may slow things down for the larger models. If you download the example and try it, you'll see how long the larger models take. The test image is only 300×225 pixels and the program would take much longer for larger images.

If you look at the picture at the top of the post, you'll see that different models provide better or worse results. The birefnet-general-lite model provides a good outcome but the birefnet-massive model gives the best result. As you may guess from its name, it is a big model so it takes a long time to download and a long time to process the image.

For a list of available models, see the Models section on rembg's github page.

Download the example to experiment with it and to see additional details.