Spaces:
No application file
No application file
File size: 4,454 Bytes
152e3c7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | import argparse
import os
import sys
import cv2
from ultralytics import YOLO
import numpy as np
# --- Configuration ---
# Set the path for the YOLO model weights
YOLO_MODEL_PATH = r"S:\ano_dec_pro\AnomalyDetectionCVPR2018-Pytorch\yolo_my_model.pt"
# --- Main Detection Function ---
def analyze_video_with_yolo(video_path: str, model_path: str = YOLO_MODEL_PATH, conf_threshold: float = 0.5,return_class=False):
"""
Analyzes a video using a pre-trained YOLO model for object detection
and prints the predicted anomaly class for frames with detections.
:param video_path: Path to the input video file.
:param model_path: Path to the trained YOLO model weights.
:param conf_threshold: Minimum confidence score required for a detection.
"""
if not os.path.exists(video_path):
print(f"[ERROR] Video file not found at: {video_path}")
return
if not os.path.exists(model_path):
print(f"[ERROR] YOLO model not found at: {model_path}")
return
try:
# Load the YOLO model (Assumes it's trained for your 14 anomaly classes)
model = YOLO(model_path, task='detect')
labels = model.names
print(f"[INFO] YOLO Model loaded successfully with {len(labels)} classes.")
except Exception as e:
print(f"[FATAL] Failed to load YOLO model: {e}")
return
cap = cv2.VideoCapture(video_path)
if not cap.isOpened():
print(f"[ERROR] Could not open video stream for: {video_path}")
return
frame_num = 0
detections_found = 0
print("-" * 50)
print(f"Starting analysis of: {os.path.basename(video_path)}")
print("-" * 50)
detected_class = "None"
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
if frame_num % 10 == 0:
results = model(frame, conf=conf_threshold, verbose=False)
if results and len(results[0].boxes) > 0:
detections = results[0].boxes
best_detection = detections[detections.conf.argmax()]
class_idx = int(best_detection.cls.item())
class_name = labels.get(class_idx, "Unknown Class")
detected_class = class_name
frame_num += 1
cap.release()
cv2.destroyAllWindows()
if return_class:
return detected_class
# while cap.isOpened():
# ret, frame = cap.read()
# if not ret:
# break
# # Analyze every 10th frame to speed up testing and reduce redundancy
# if frame_num % 10 == 0:
# # Run inference (set verbose=False to keep the output clean)
# results = model(frame, conf=conf_threshold, verbose=False)
# # Process results
# if results and len(results[0].boxes) > 0:
# detections = results[0].boxes
# # We'll take the highest confidence detection in the frame
# best_detection = detections[detections.conf.argmax()]
# class_idx = int(best_detection.cls.item())
# class_name = labels.get(class_idx, "Unknown Class")
# confidence = best_detection.conf.item()
# print(f"Frame {frame_num:05d}: DETECTED -> {class_name} ({confidence:.2f})")
# detections_found += 1
# frame_num += 1
# cap.release()
# cv2.destroyAllWindows()
print("-" * 50)
print(f"Analysis complete. Total detections reported: {detections_found}")
print("-" * 50)
# if __name__ == '__main__':
# # --- Example Usage (Run this command in your environment) ---
# # Change the video path to a file you want to test!
# example_video = r"S:\ano_dec_pro\dataset\Abuse\Abuse002_x264.mp4"
# # NOTE: In a shell/CLI environment (like VS Code or PowerShell), you typically
# # run this by executing the script and passing arguments:
# # python yolo_detection.py --video_path "S:\ano_dec_pro\dataset\Abuse\Abuse002_x264.mp4"
# # We call the function directly for simple execution in Python interpreter
# print("NOTE: Running the detection script. Please change 'example_video' to your actual file path.")
# analyze_video_with_yolo(example_video)
|