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
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-----------
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!