YOLO Object Detection

import os import shutil from ultralytics import YOLO import cv2 # Paths to the base directories frame_base_dir = '/workspace/stage1' detection_output_dir = '/workspace/stage2.frame.detection' # Initialize the detection model detection_model = YOLO('/workspace/yolo_model.pt') def create_output_dir_structure(base_dir, output_dir): """Create the output directory structure matching the base directory.""" for root, dirs, files in os.walk(base_dir): for dir_name in dirs: new_dir_path = os.path.join(output_dir, os.path.relpath(os.path.join(root, dir_name), base_dir)) os.makedirs(new_dir_path, exist_ok=True) def run_detection_on_frame(frame_path, output_folder): """Run detection on the frame and save the result to the output folder.""" frame_filename = os.path.basename(frame_path) try: results = detection_model.predict(frame_path, save=False) for result in results: annotated_frame = result.plot() # Get the annotated frame  output_path = os.path.join(output_folder, frame_filename) cv2.imwrite(output_path, annotated_frame) print(f"Saved detection result for {frame_path} to {output_path}") except Exception as e: print(f"Error running detection on {frame_path}: {e}") def process_frames_for_detection(frame_base_dir, detection_output_dir): """Process each frame in the frame base directory and run detection.""" for root, dirs, files in os.walk(frame_base_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, frame_base_dir) output_folder = os.path.join(detection_output_dir, relative_path) run_detection_on_frame(frame_path, output_folder) # Create the output directory structure create_output_dir_structure(frame_base_dir, detection_output_dir) # Process frames and run detection process_frames_for_detection(frame_base_dir, detection_output_dir) print("Frame detection complete.")