DEV Community

Timothy Fosteman
Timothy Fosteman

Posted on

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:
  1. Set the detection_model to your YOLO model file.
  2. Define the frame_base_dir where the frames are stored.
  3. Set the detection_output_dir where the detection results will be saved.
  4. 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.")
Enter fullscreen mode Exit fullscreen mode

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.

Heroku

Simplify your DevOps and maximize your time.

Since 2007, Heroku has been the go-to platform for developers as it monitors uptime, performance, and infrastructure concerns, allowing you to focus on writing code.

Learn More

Top comments (0)

Image of Docusign

🛠️ Bring your solution into Docusign. Reach over 1.6M customers.

Docusign is now extensible. Overcome challenges with disconnected products and inaccessible data by bringing your solutions into Docusign and publishing to 1.6M customers in the App Center.

Learn more