Remove Image Background with Python & AI

Remove Image Background with Python

Table of Contents

Spread the love

How to Remove Image Background with Python?

I was surfing the internet searching for AI tools, and I stumbled upon an Image Background Removal service called Remove.Bg

It works great. But I was shocked by the pricing. I mean, it charges, 0.2$ per Image!????

Remove Image Background with Python
Remove Image Background with Python

So, I decided to take the challenge and build my own Image Background Remover without relying on any APIs or Third Parties.

I kept searching for a couple of days, I found a great python project on GitHub, Called U2Net.

A Big Shout-out and Thanks to the authors of this project!

I tried my best to make things as simple as possible so that even newbies can use and run the python script I built.

I also created a Short Video showing The Script & How to use it, and I pointed out two special secrets that may help you build an online income stream with the help of this script. Here is the video:????

For example, You can turn this into an online tool, that people pay to use, based on a monthly subscription.

What you can also do is that. You can use the script and turn it into an API and publish it on rapid API and sell it for a monthly recurring income.

rapid api image background remover

In addition to that, you can build a free tools website and monetize it with Ads. Like I do with H-supertools. I build a lot of free tools, and I monetize them with ads.


Let’s get started with the Project.

Remove Image Background with Python

So, let's open our source code with visual studio code

If you don't know How to set up Python with Visual Studio Code.

Here is a tutorial to help you.????

How To Download And Install Python With Visual Studio Code (2022)

The main essence of our script is this magic method here called removebg. We'll pass the Image Path as an argument.

So, two methods are very important in this Python Script. One that saves the output and the other one will create the mask of the image and the image without background and return back.

Remove Image Background with Python

So, in the first part, we are loading the model. I will give you the trained model as well so you can use it directly.

After Loading, we call the removebg function to remove the background image.  

Let's see how it works I will open again and get a sample image from their website

Remove Image Background with Python

So we can compare the results.

I will save this cat image on my desktop.

I will copy the name, and open the script again. And paste the image path here

And simply run the script. 

To see the Results. Go to your Python Project Folder.

Find a Folder named ‘static.'

Now in the static folder, go to inputs.

So, this is the input (this is the main image.)

In the Masks folder, there will be a mask of the same image.  

And finally, this is the result. It's Great! ????

Remove Image Background with Python (Result)

Now let's go to, the result here is almost the same. 

Thanks to everyone who helped in developing this awesome machine learning model.???? 

Remove Image Background with Python (Source Code)

You can Download Full Source Code or Check it out On GitHub

By the way here's the script!


import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import cv2
import uuid
import os

from model import U2NET
from torch.autograd import Variable
from skimage import io, transform
from PIL import Image

# Get The Current Directory
currentDir = os.path.dirname(__file__)

# Functions:
# Save Results

def save_output(image_name, output_name, pred, d_dir, type):
    predict = pred
    predict = predict.squeeze()
    predict_np = predict.cpu().data.numpy()
    im = Image.fromarray(predict_np*255).convert('RGB')
    image = io.imread(image_name)
    imo = im.resize((image.shape[1], image.shape[0]))
    pb_np = np.array(imo)
    if type == 'image':
        # Make and apply mask
        mask = pb_np[:, :, 0]
        mask = np.expand_dims(mask, axis=2)
        imo = np.concatenate((image, mask), axis=2)
        imo = Image.fromarray(imo, 'RGBA')
# Remove Background From Image (Generate Mask, and Final Results)

def removeBg(imagePath):
    inputs_dir = os.path.join(currentDir, 'static/inputs/')
    results_dir = os.path.join(currentDir, 'static/results/')
    masks_dir = os.path.join(currentDir, 'static/masks/')

    # convert string of image data to uint8
    with open(imagePath, "rb") as image:
        f =
        img = bytearray(f)

    nparr = np.frombuffer(img, np.uint8)

    if len(nparr) == 0:
        return '---Empty image---'

    # decode image
        img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        # build a response dict to send back to client
        return "---Empty image---"

    # save image to inputs
    unique_filename = str(uuid.uuid4())
    cv2.imwrite(inputs_dir+unique_filename+'.jpg', img)

    # processing
    image = transform.resize(img, (320, 320), mode='constant')

    tmpImg = np.zeros((image.shape[0], image.shape[1], 3))

    tmpImg[:, :, 0] = (image[:, :, 0]-0.485)/0.229
    tmpImg[:, :, 1] = (image[:, :, 1]-0.456)/0.224
    tmpImg[:, :, 2] = (image[:, :, 2]-0.406)/0.225

    tmpImg = tmpImg.transpose((2, 0, 1))
    tmpImg = np.expand_dims(tmpImg, 0)
    image = torch.from_numpy(tmpImg)

    image = image.type(torch.FloatTensor)
    image = Variable(image)

    d1, d2, d3, d4, d5, d6, d7 = net(image)
    pred = d1[:, 0, :, :]
    ma = torch.max(pred)
    mi = torch.min(pred)
    dn = (pred-mi)/(ma-mi)
    pred = dn

    save_output(inputs_dir+unique_filename+'.jpg', unique_filename +
                '.png', pred, results_dir, 'image')
    save_output(inputs_dir+unique_filename+'.jpg', unique_filename +
                '.png', pred, masks_dir, 'mask')
    return "---Success---"

# ------- Load Trained Model --------
print("---Loading Model---")
model_name = 'u2net'
model_dir = os.path.join(currentDir, 'saved_models',
                         model_name, model_name + '.pth')
net = U2NET(3, 1)
if torch.cuda.is_available():
    net.load_state_dict(torch.load(model_dir, map_location='cpu'))
# ------- Load Trained Model --------

print("---Removing Background...")
# ------- Call The removeBg Function --------
imgPath = "Image_File_Path"  # Change this to your image path


Also, here is some help, if you download the Full Project Folder.

File Structure

when you open the Downloads folder. You'll find this inside:

Saved models



And here is the model which is trained so you can use it directly. 

And we have the script file here.

And the static folder we have the inputs, masks, and results. You can find them here also it's super easy to use.


This is the requirements text file you can use directly to import all the requirements to install them with the pip command.

And after you finish I really advise you to turn this into an API with Flask.

It's super easy.

If you don't know flask, and how you can turn the script into an API and sell it.

Here is a Tutorial.????

How to Deploy Python flask API on a Windows VPS Server (5-min Definitive Guide)

If you have any questions, leave them in the comments.

I'll try my best to get back to you.

Take Care and Happy Coding!

Spread the love