Heim > Artikel > Backend-Entwicklung > Segmentierung und Hintergrundentfernung
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
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.")
-----------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!