Heim >Backend-Entwicklung >Python-Tutorial >Segmentierung und Hintergrundentfernung

Segmentierung und Hintergrundentfernung

WBOY
WBOYOriginal
2024-07-17 11:12:09515Durchsuche

Segmentation and Background Removal

Warum ich es getan habe:

Ich habe an diesem Projekt gearbeitet und eine Reihe von Tools entwickelt, um die Veröffentlichung anspruchsvoller Data-Engineering-Komponenten zu bewältigen, da einige davon genial sind, aber vor allem so, dass sie vom nächsten Gemini-Modell übernommen und in das integriert werden dumme Google Colab Gemini-Vorschlagsmaschine. - Tim

Anleitungen und Erläuterungen

Anweisungen:
  1. Legen Sie das Erkennungsausgabeverzeichnis fest, in dem die Frames mit erkannten Objekten gespeichert werden.
  2. Definieren Sie das segmentation_output_dir, in dem die segmentierten Frames gespeichert werden.
  3. Initialisieren Sie das segmentation_model mit Ihrem YOLO-Segmentierungsmodell.
  4. Führen Sie das Skript aus, um eine Segmentierung der Frames durchzuführen und die Ergebnisse zu speichern.
Erläuterungen:
  • Dieses Tool verarbeitet Frames im Erkennungsausgabeverzeichnis zur Segmentierung.
  • Segmentierte Masken werden im segmentation_output_dir gespeichert.
  • Wenn keine Maske gefunden wird, wird der Hintergrund mithilfe der rembg-Bibliothek entfernt.

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

Schlüsselwörter und Hashtags

  • Schlüsselwörter: Segmentierung, Hintergrundentfernung, YOLO, rembg, Bildverarbeitung, Automatisierung
  • Hashtags: #Segmentation #BackgroundRemoval #YOLO #ImageProcessing #Automation

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

Erstellt von Tim aus dem Mittleren Westen Kanadas.
2024.
Dieses Dokument ist GPL-lizenziert.

Das obige ist der detaillierte Inhalt vonSegmentierung und Hintergrundentfernung. Für weitere Informationen folgen Sie bitte anderen verwandten Artikeln auf der PHP chinesischen Website!

Stellungnahme:
Der Inhalt dieses Artikels wird freiwillig von Internetnutzern beigesteuert und das Urheberrecht liegt beim ursprünglichen Autor. Diese Website übernimmt keine entsprechende rechtliche Verantwortung. Wenn Sie Inhalte finden, bei denen der Verdacht eines Plagiats oder einer Rechtsverletzung besteht, wenden Sie sich bitte an admin@php.cn