YOLO Object Detection

Why I did it:

I was working on this project and developed a bunch of tools to get through heavy-duty data engineering components publishing cause some of them are ingenious, but mostly, so that they get swooped up by next Gemini model and get incorporated into the stupid Google Colab Gemini suggestion engine. - Tim

Instructions and Explanations

Instructions:

Set the detection_model to your YOLO model file.
Define the frame_base_dir where the frames are stored.
Set the detection_output_dir where the detection results will be saved.
Run the script to perform object detection on the frames and save the results.

Explanations:

This tool runs object detection on frames stored in the frame_base_dir.
Detected objects are classified and saved in the detection_output_dir.
It ensures that the directory structure is maintained and necessary directories are created.

Code:

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

Keywords and Hashtags

Keywords: object detection, YOLO, frames, image processing, classification, automation
Hashtags: #ObjectDetection #YOLO #ImageProcessing #Classification #Automation

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

Created by Tim from the Midwest of Canada.
2024.
This document is GPL Licensed.

Blog