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