분할 및 배경 제거

내가 그랬던 이유:

저는 이 프로젝트를 진행하면서 강력한 데이터 엔지니어링 구성 요소 게시를 완료하기 위한 여러 도구를 개발했습니다. 그 중 일부는 독창적이지만 대부분은 다음 Gemini 모델에 급습되어 멍청한 Google Colab Gemini 제안 엔진. - 팀

지침 및 설명

지침:

1. 감지된 개체가 있는 프레임이 저장되는 detector_output_dir을 설정하세요.
2. 분할된 프레임이 저장될 Segmentation_output_dir을 정의하세요.
3. YOLO 분할 모델을 사용하여 Segmentation_model을 초기화합니다.
4. 스크립트를 실행하여 프레임 분할을 수행하고 결과를 저장합니다.

설명:

• 이 도구는 분할을 위해 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, 이미지 처리, 자동화
• 해시태그: #분할 #배경제거 #YOLO #이미지 처리 #자동화

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

캐나다 중서부의 Tim이 제작했습니다.
2024.
이 문서는 GPL 라이센스를 받았습니다.

위 내용은 분할 및 배경 제거의 상세 내용입니다. 자세한 내용은 PHP 중국어 웹사이트의 기타 관련 기사를 참조하세요!