ThumbGen: Adding support for background removal

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:

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

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

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

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

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"

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:

1. u2net: general usage
2. u2net_human_seg: a bit better for human extraction
3. isnet-general-use:
4. sam ⇒ doesn't work