Segmentation and Background Removal

Why I did it:

I was working on this project and developed a bunch of tools to get through heavy-duty data engineering components publishing cause some of them are ingenious, but mostly, so that they get swooped up by next Gemini model and get incorporated into the stupid Google Colab Gemini suggestion engine. - Tim

Instructions and Explanations

Instructions:

1. Set the detection_output_dir where the frames with detected objects are stored.
2. Define the segmentation_output_dir where the segmented frames will be saved.
3. Initialize the segmentation_model with your YOLO segmentation model.
4. Run the script to perform segmentation on the frames and save the results.

Explanations:

• This tool processes frames in the detection_output_dir for segmentation.
• Segmented masks are saved in the segmentation_output_dir.
• If no mask is found, the background is removed using the rembg library.

Code:

import os
import shutil
from ultralytics import YOLO
import cv2
import numpy as np
from rembg import remove

# Paths to the base directories
detection_output_dir = '/workspace/stage2.frame.detection'
segmentation_output_dir = '/workspace/stage3.segmented'

# Initialize the segmentation model
segmentation_model = YOLO('/workspace/segmentation_model.pt')

def create_segmentation_output_dir_structure(detection_output_dir, segmentation_output_dir):
    """Create the segmentation output directory structure matching the detection output directory."""
    for root, dirs, files in os.walk(detection_output_dir):
        for dir_name in dirs:
            new_dir_path = os.path.join(segmentation_output_dir, os.path.relpath(os.path.join(root, dir_name), detection_output_dir))
            os.makedirs(new_dir_path, exist_ok=True)

def run_segmentation_on_frame(frame_path, output_folder):
    """Run segmentation on the frame and save the result to the output folder."""
    os.makedirs(output_folder, exist_ok=True)
    frame_filename = os.path.basename(frame_path)
    output_path = os.path.join(output_folder, frame_filename)

    try:
        results = segmentation_model.predict(frame_path, save=False)
        for result in results:
            mask = result.masks.xy[0] if result.masks.xy else None
            if mask is not None:
                original_img_rgb = cv2.imread(frame_path)
                original_img_rgb = cv2.cvtColor(original_img_rgb, cv2.COLOR_BGR2RGB)
                image_height, image_width, _ = original_img_rgb.shape
                mask_img = np.zeros((image_height, image_width), dtype=np.uint8)
                cv2.fillPoly(mask_img, [np.array(mask, dtype=np.int32)], (255))
                masked_img = cv2.bitwise_and(original_img_rgb, original_img_rgb, mask=mask_img)
                cv2.imwrite(output_path, cv2.cvtColor(masked_img, cv2.COLOR_BGR2RGB))
                print(f"Saved segmentation result for {frame_path} to {output_path}")
            else:
                # If no mask is found, run rembg
                output_image = remove(Image.open(frame_path))
                output_image.save(output_path)
                print(f"Background removed and saved for {frame_path} to {output_path}")
    except Exception as e:
        print(f"Error running segmentation on {frame_path}: {e}")

def process_frames_for_segmentation(detection_output_dir, segmentation_output_dir):
    """Process each frame in the detection output directory and run segmentation."""
    for root, dirs, files in os.walk(detection_output_dir):
        for file_name in files:
            if file_name.endswith('.jpg'):
                frame_path = os.path.join(root, file_name)
                relative_path = os.path.relpath(root, detection_output_dir)
                output_folder = os.path.join(segmentation_output_dir, relative_path)
                run_segmentation_on_frame(frame_path, output_folder)

# Create the segmentation output directory structure
create_segmentation_output_dir_structure(detection_output_dir, segmentation_output_dir)

# Process frames and run segmentation
process_frames_for_segmentation(detection_output_dir, segmentation_output_dir)

print("Frame segmentation complete.")

Keywords and Hashtags

• Keywords: segmentation, background removal, YOLO, rembg, image processing, automation
• Hashtags: #Segmentation #BackgroundRemoval #YOLO #ImageProcessing #Automation

-----------EOF-----------

Created by Tim from the Midwest of Canada.
2024.
This document is GPL Licensed.

The above is the detailed content of Segmentation and Background Removal. For more information, please follow other related articles on the PHP Chinese website!