infer.py

import argparse
from pathlib import Path
import cv2
import numpy as np
from ultralytics import YOLO
import time

def run_inference_on_image(model_path, image_path, conf_threshold=0.5, save_path=None):
    """Run inference on a single image"""
    # Load model
    model = YOLO(model_path)
    
    # Run inference
    start_time = time.time()
    results = model(image_path, conf=conf_threshold)
    inference_time = time.time() - start_time
    
    # Process results
    img = cv2.imread(image_path)
    
    # Draw results on image
    for result in results:
        boxes = result.boxes
        print(f"Detected {len(boxes)} guns in {inference_time:.4f} seconds")
        
        for box in boxes:
            x1, y1, x2, y2 = map(int, box.xyxy[0])
            conf = float(box.conf[0])
            
            # Draw bounding box
            cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.putText(img, f"Gun: {conf:.2f}", (x1, y1 - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
    
    # Save or display the result
    if save_path:
        cv2.imwrite(save_path, img)
        print(f"Result saved to {save_path}")
    else:
        cv2.imshow("Gun Detection Result", img)
        cv2.waitKey(0)
        cv2.destroyAllWindows()

def run_inference_on_video(model_path, video_path, conf_threshold=0.55, save_path=None):
    """Run inference on a video file"""
    # Load model
    model = YOLO(model_path)
    
    # Open video
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"Error: Could not open video {video_path}")
        return
    
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    
    # Create video writer if save_path is provided
    if save_path:
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        writer = cv2.VideoWriter(save_path, fourcc, fps, (width, height))
    
    # Process frames
    frame_count = 0
    total_time = 0
    
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        
        # Start timing
        start_time = time.time()
        
        # Convert BGR to RGB and normalize
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        
        # Run inference
        results = model(frame_rgb, conf=conf_threshold)
        
        # Calculate inference time
        inference_time = time.time() - start_time
        total_time += inference_time
        frame_count += 1
        
        # Draw results on frame
        annotated_frame = frame.copy()
        for result in results:
            for box in result.boxes:
                x1, y1, x2, y2 = map(int, box.xyxy[0])
                conf = float(box.conf[0])
                
                # Filter out low-confidence detections
                if conf < conf_threshold:
                    continue
                
                # Draw bounding box and label
                cv2.rectangle(annotated_frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                cv2.putText(annotated_frame, f"Weapon: {conf:.2f}", (x1, y1 - 10),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
        
        # Add FPS info
        fps_text = f"FPS: {1/inference_time:.1f}"
        cv2.putText(annotated_frame, fps_text, (20, 40),
                    cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
        
        # Save or display the frame
        if save_path:
            writer.write(annotated_frame)
        else:
            cv2.imshow("Gun Detection", annotated_frame)
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
    
    # Release resources
    cap.release()
    if save_path:
        writer.release()
    cv2.destroyAllWindows()
    
    # Print statistics
    avg_fps = frame_count / total_time if total_time > 0 else 0
    print(f"Processed {frame_count} frames in {total_time:.2f} seconds ({avg_fps:.2f} FPS)")

def run_inference_on_webcam(model_path, camera_id=0, conf_threshold=0.55):
    """Run inference on webcam"""
    # Load model
    model = YOLO(model_path)
    
    # Open webcam
    cap = cv2.VideoCapture(camera_id)
    if not cap.isOpened():
        print(f"Error: Could not open webcam {camera_id}")
        return
    
    # Process frames
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        
        # Start timing
        start_time = time.time()
        
        # Convert BGR to RGB and normalize
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        
        # Run inference
        results = model(frame_rgb, conf=conf_threshold)
        
        # Calculate inference time
        inference_time = time.time() - start_time
        
        # Draw results on frame
        annotated_frame = frame.copy()
        for result in results:
            for box in result.boxes:
                x1, y1, x2, y2 = map(int, box.xyxy[0])
                conf = float(box.conf[0])
                
                # Filter out low-confidence detections
                if conf < conf_threshold:
                    continue
                
                # Draw bounding box and label
                cv2.rectangle(annotated_frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                cv2.putText(annotated_frame, f"Weapon: {conf:.2f}", (x1, y1 - 10),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
        
        # Add FPS info
        fps_text = f"FPS: {1/inference_time:.1f}"
        cv2.putText(annotated_frame, fps_text, (20, 40),
                    cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
        
        # Display the frame
        cv2.imshow("Gun Detection (Press 'q' to quit)", annotated_frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    
    # Release resources
    cap.release()
    cv2.destroyAllWindows()

def main():
    # Parse command-line arguments
    parser = argparse.ArgumentParser(description="Run inference with YOLOv8 gun detection model")
    parser.add_argument("--model", type=str, required=True, help="Path to the trained model")
    parser.add_argument("--source", type=str, required=True, 
                        help="Path to image, video file or 'webcam' for live detection")
    parser.add_argument("--conf", type=float, default=0.5, help="Confidence threshold")
    parser.add_argument("--output", type=str, default=None, help="Path to save results")
    
    args = parser.parse_args()
    
    # Run inference based on source type
    if args.source.lower() == "webcam":
        run_inference_on_webcam(args.model, camera_id=0, conf_threshold=args.conf)
    elif args.source.lower().endswith(('.mp4', '.avi', '.mov', '.mkv')):
        run_inference_on_video(args.model, args.source, conf_threshold=args.conf, save_path=args.output)
    else:
        run_inference_on_image(args.model, args.source, conf_threshold=args.conf, save_path=args.output)

if __name__ == "__main__":
    main()