我为什么这么做：
我正在研究这个项目，并开发了一堆工具来完成重型数据工程组件的发布，因为其中一些是巧妙的，但大多数是，这样它们就会被下一个 gemini 模型突袭并并入愚蠢的 google colab gemini 建议引擎。 – 蒂姆
说明和解释
指示：
- 设置检测输出目录，其中存储检测到的对象的帧。
- 定义将保存分段帧的segmentation_output_dir。
- 使用 yolo 分割模型初始化egmentation_model。
- 运行脚本对帧进行分割并保存结果。
说明：
- 此工具处理 detector_output_dir 中的帧以进行分割。
- 分段蒙版保存在segmentation_output_dir中。
- 如果没有找到遮罩，则使用 rembg 库删除背景。
代码：
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.")
关键词和标签
- 关键词：分割、背景去除、yolo、rembg、图像处理、自动化
- 标签：#segmentation #backgroundremoval #yolo #imageprocessing #automation
———-eof———-
由来自加拿大中西部的 tim 创建。
2024.
本文档已获得 gpl 许可。
