====== ThumbGen: Adding support for background removal ======
{{tag>dev python ai background rembg youtube thumbnail thumbgen}}
In this post we will see how to remove the background from a given image. This will be a pretty simple article in fact, but at least it could serve as reference later on how to use this new script command.
====== ======
Youtube video for this article available at:
;#;
{{ youtube>Re2jc4QOQXs?large }}
;#;
==== Initial setup ====
Reference utilities:
* https://www.remove.bg/
* https://removal.ai/
Reference implementation:
* https://github.com/danielgatis/rembg
Our implementation is available as part of the NervProj project: https://github.com/roche-emmanuel/nervproj
Adding the rembg package in **media_env**: media_env:
inherit: default_env
packages:
- moviepy
- Pillow
- ffmpeg-python
- opencv-python
- rembg[gpu]
- scipy
Adding initial script function to remove background:
def remove_background(self, in_file, out_file, model_name):
"""Remove the background from an image file"""
# Usage infos: https://github.com/roche-emmanuel/rembg/blob/main/USAGE.md
if out_file is None:
out_file = self.set_path_extension(in_file, "_nobg.png")
logger.info("Removing background from %s...", in_file)
input_img = Image.open(in_file)
# output_img = remove(input_img)
# model_name = "u2net"
# model_name = "isnet-general-use"
session = new_session(model_name)
output_img = remove(input_img, session=session, only_mask=True)
output_img.save(out_file)
logger.info("Done removing background.")
return True
==== Expanding subject mask ====
I tried to experiment a bit to add support to extend the extracted area with an arbitrary number of pixels:
if max_dist != min_dist:
# Get the mask only:
mask = remove(input_img, session=session, only_mask=True).convert("L")
logger.info("Retrieved foreground mask.")
# Compute the distance to foreground:
dist = self.compute_distance_to_foreground(mask)
# Convert the input image to RGBA:
img = input_img.convert("RGBA")
img_arr = np.array(img)
dist = np.clip(dist, min_dist, max_dist)
alpha = 1.0 - (dist - min_dist) / (max_dist - min_dist)
img_arr[:, :, 3] = (255 * alpha).astype(np.uint8)
# Convert the numpy array back to image:
output_img = Image.fromarray(img_arr)
To activate this we need to specify different values for the ''--mind'' and the ''--maxd'' command line arguments: nvp rembg --mind 5 --maxd 10
==== Background color replacement ====
Added support for background color specification:
def update_bg_color(self, img, bg_color):
"""Replace the background color"""
col = [np.float32(el) / 255.0 for el in bg_color.split(",")]
arr = np.array(img).astype(np.float32) / 255.0
colarr = np.zeros_like(arr)
alpha = arr[:, :, 3]
for i in range(4):
colarr[:, :, i] = col[i]
arr[:, :, i] = arr[:, :, i] * alpha + colarr[:, :, i] * (1.0 - alpha)
arr = (arr * 255.0).astype(np.uint8)
return Image.fromarray(arr)
Which can then be used as this:
nvp rembg -i 2thumbs.png --mind 1 --maxd 1 --bgcolor 255,255,255,255
==== Batch processing and output folder specification ====
Also added support for output folder specification, and using that for batch processing:
if in_file == "all":
# iterate on all image files:
cur_dir = self.get_cwd()
out_dir = self.get_param("out_dir")
self.make_folder(out_dir)
all_files = self.get_all_files(cur_dir, recursive=False)
exts = [".png", ".jpeg", ".jpg"]
for fname in all_files:
ext = self.get_path_extension(fname).lower()
if ext not in exts:
continue
# Check if we already have the text file:
src_file = self.get_path(cur_dir, fname)
# out_file = self.set_path_extension(fname, "_nobg.png")
out_file = self.set_path_extension(fname, ".png")
out_file = self.get_path(out_dir, out_file)
if self.file_exists(out_file):
continue
# Otherwise we process this file:
self.remove_background(
src_file, out_file, model, min_dist, max_dist, bg_color, contour_size, contour_color
)
return True
==== Adding support for contour display ====
Added support to fill only the contour area around the subject with a specific color:
def apply_contour(self, arr, mask, dist, contour_size, contour_color):
"""Apply a contour around the target object"""
col = [np.float32(el) / 255.0 for el in contour_color.split(",")]
logger.info("Applying contour of size %f, with color %s", contour_size, contour_color)
img_arr = arr.astype(np.float32) / 255.0
# Prepare the result image:
res = np.copy(img_arr)
# Fill with the contour color:
idx = dist <= contour_size
alpha = np.array(mask).astype(np.float32) / 255.0
for i in range(4):
# Fill with the contour color:
res[idx, i] = col[i]
# Re-add the subject on top of contour:
res[:, :, i] = img_arr[:, :, i] * alpha + res[:, :, i] * (1.0 - alpha)
return (res * 255.0).astype(np.uint8)
nvp rembg --ctsize 10 -i sideview.png --mind 7 --maxd 10 --model "u2net_human_seg"
==== Notes on the different models ====
They will be stored in **${HOME}/.u2net** (and will only be downloaded the first time they are used)
We can specify the model to use on the command line: nvp rembg --bgcolor 255,255,255,255 --out-dir ../nobg --model "u2net_human_seg"
And tested with different models:
- u2net: general usage
- u2net_human_seg: a bit better for human extraction
- isnet-general-use:
- sam => **doesn't work**