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:
- Set the
detection_output_dir
where the frames with detected objects are stored. - Define the
segmentation_output_dir
where the segmented frames will be saved. - Initialize the
segmentation_model
with your YOLO segmentation model. - 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.
Top comments (0)