Introduction
This guide demonstrates how to replace the background of an image using Python code only, without relying on image editing software like Photoshop. The goal is to keep the subject intact while swapping in an AI-generated background.
While this approach may not be revolutionary, it addresses a common need, so I hope it will be helpful for those with similar requirements.
Input and Output Images
Let’s start with the results.
The following output image was generated from the input image shown below.
Input Image
Output Image
Libraries
Install requests to handle the API calls.
$ pip install requests I verified the version as follows:
$ pip list | grep -e requests requests 2.31.0 API Key
For background generation, we’ll use Stability AI’s Web API.
To access this API, you’ll need to obtain an API Key from their Developer Platform. For pricing, refer to the Pricing page.
To keep your key secure, save it as an environment variable rather than hardcoding it in your code.
In my environment, I use the zshrc settings file.
$ open ~/.zshrc I saved the key under the name STABILITY_API_KEY.
export STABILITY_API_KEY=your_api_key_here Code
Here, we use the Remove Background API to isolate the subject. We then pass the extracted image to the Inpaint API to create the new background.
The prompt used is "Large glass windows with a view of the metropolis behind"
import os import requests # File paths input_path = './input.png' # Original image mask_path = './mask.png' # Mask image (temporarily generated) output_path = './output.png' # Output image # Check for API Key api_key = os.getenv("STABILITY_API_KEY") if api_key is None: raise Exception("Missing Stability API key.") headers = { "Accept": "image/*", "Authorization": f"Bearer {api_key}" } # Call Remove Background API response = requests.post( f"https://api.stability.ai/v2beta/stable-image/edit/remove-background", headers=headers, files={ "image": open(input_path, "rb") }, data={ "output_format": "png" }, ) # Save mask image if response.status_code == 200: with open(mask_path, 'wb') as file: file.write(response.content) else: raise Exception(str(response.json())) # Call Inpaint API response = requests.post( "https://api.stability.ai/v2beta/stable-image/edit/inpaint", headers=headers, files={ "image": open(mask_path, "rb"), }, data={ "prompt": "Large glass windows with a view of the metropolis behind", "output_format": "png", "grow_mask": 0, # Disable blurring around the mask }, ) # Delete mask image os.remove(mask_path) # Save output image if response.status_code == 200: with open(output_path, "wb") as file: file.write(response.content) else: raise Exception(str(response.json())) Using rembg
Another approach for background removal is to use rembg. This method requires only one API call, making it more cost-effective, though it may result in differences in extraction accuracy.
First, install rembg.
$ pip install rembg I verified the version as follows:
$ pip list | grep -e rembg rembg 2.0.59 Here’s the code for this approach:
from rembg import remove import os import requests # File paths input_path = './input.png' # Input image path mask_path = './mask.png' # Mask image path (temporarily generated) output_path = './output.png' # Output image path # Generate mask image with background removed with open(input_path, 'rb') as i: with open(mask_path, 'wb') as o: input_image = i.read() mask_image = remove(input_image) o.write(mask_image) # Check for API Key api_key = os.getenv("STABILITY_API_KEY") if api_key is None: raise Exception("Missing Stability API key.") # Call Inpaint API response = requests.post( "https://api.stability.ai/v2beta/stable-image/edit/inpaint", headers={ "Accept": "image/*", "Authorization": f"Bearer {api_key}" }, files={ "image": open(mask_path, "rb"), }, data={ "prompt": "Large glass windows with a view of the metropolis behind", "output_format": "png", "grow_mask": 0, }, ) # Delete mask image os.remove(mask_path) # Save output image if response.status_code == 200: with open(output_path, "wb") as file: file.write(response.content) else: raise Exception(str(response.json())) Here’s the output image. In this case, the accuracy of the extraction seems satisfactory.
If you set up a local Stable Diffusion environment, you can eliminate API call costs, so feel free to explore that option if it suits your needs.
Conclusion
Being able to achieve this through code alone is highly convenient.
It’s exciting to witness the ongoing improvements in workflow efficiency.
Top comments (0)