Pembahagian dan Pembuangan Latar Belakang

Mengapa saya melakukannya:

Saya sedang mengusahakan projek ini dan membangunkan sekumpulan alatan untuk melalui penerbitan komponen kejuruteraan data tugas berat kerana sesetengah daripadanya adalah bijak, tetapi kebanyakannya, supaya mereka tertarik dengan model Gemini seterusnya dan dimasukkan ke dalam enjin cadangan Google Colab Gemini yang bodoh. - Tim

Arahan dan Penerangan

Arahan:

1. Tetapkan detection_output_dir tempat bingkai dengan objek yang dikesan disimpan.
2. Tentukan segmentation_output_dir tempat bingkai yang disegmen akan disimpan.
3. Mulakan model segmentasi dengan model segmentasi YOLO anda.
4. Jalankan skrip untuk melakukan pembahagian pada bingkai dan simpan hasilnya.

Penjelasan:

• Alat ini memproses bingkai dalam detection_output_dir untuk pembahagian.
• Topeng bersegmen disimpan dalam segmentation_output_dir.
• Jika tiada topeng ditemui, latar belakang dialih keluar menggunakan perpustakaan rembg.

Kod:

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.")

Kata kunci dan Hashtag

• Kata kunci: pembahagian, penyingkiran latar belakang, YOLO, rembg, pemprosesan imej, automasi
• Hashtags: #Segmentasi #BackgroundRemoval #YOLO #ImageProcessing #Automation

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

Dicipta oleh Tim dari Midwest Kanada.
2024.
Dokumen ini Berlesen GPL.

Atas ialah kandungan terperinci Pembahagian dan Pembuangan Latar Belakang. Untuk maklumat lanjut, sila ikut artikel berkaitan lain di laman web China PHP!