Spaces:
Sleeping
Sleeping
File size: 3,150 Bytes
6cb50ff a03d512 04f4d0b 8c84287 4527f8f 8c84287 a668c53 fa66a0f 04f4d0b 8c84287 a03d512 8c84287 4527f8f 8c84287 6cb50ff 4527f8f 8c84287 |
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 |
import cv2
import torch
import gradio as gr
import numpy as np
import os
import matplotlib.pyplot as plt
from ultralytics import YOLO, __version__ as ultralytics_version
# Debug: Check environment
print(f"Torch version: {torch.__version__}")
print(f"Gradio version: {gr.__version__}")
print(f"Ultralytics version: {ultralytics_version}")
print(f"CUDA available: {torch.cuda.is_available()}")
# Load YOLOv8 model
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")
model = YOLO('./data/best.pt').to(device)
def process_video(video, output_folder="detected_frames", plot_graphs=False):
if video is None:
return "Error: No video uploaded"
# Create output folder if it doesn't exist
if not os.path.exists(output_folder):
os.makedirs(output_folder)
cap = cv2.VideoCapture(video)
if not cap.isOpened():
return "Error: Could not open video file"
frame_width, frame_height = 320, 240 # Smaller resolution
frame_count = 0
frame_skip = 5 # Process every 5th frame
max_frames = 100 # Limit for testing
confidence_scores = [] # Store confidence scores for plotting
while True:
ret, frame = cap.read()
if not ret or frame_count > max_frames:
break
frame_count += 1
if frame_count % frame_skip != 0:
continue
frame = cv2.resize(frame, (frame_width, frame_height))
print(f"Processing frame {frame_count}")
# Run YOLOv8 inference
results = model(frame)
annotated_frame = results[0].plot()
# Save annotated frame
frame_filename = os.path.join(output_folder, f"frame_{frame_count:04d}.jpg")
cv2.imwrite(frame_filename, annotated_frame)
# Collect confidence scores for plotting
if results[0].boxes is not None:
confs = results[0].boxes.conf.cpu().numpy()
confidence_scores.extend(confs)
cap.release()
# Generate confidence score plot if requested
graph_path = None
if plot_graphs and confidence_scores:
plt.figure(figsize=(10, 5))
plt.hist(confidence_scores, bins=20, color='blue', alpha=0.7)
plt.title('Distribution of Confidence Scores')
plt.xlabel('Confidence Score')
plt.ylabel('Frequency')
graph_path = os.path.join(output_folder, "confidence_histogram.png")
plt.savefig(graph_path)
plt.close()
return f"Frames saved in {output_folder}. {f'Graph saved as {graph_path}' if graph_path else ''}"
# Gradio interface
iface = gr.Interface(
fn=process_video,
inputs=[
gr.Video(label="Upload Video"),
gr.Textbox(label="Output Folder", value="detected_frames"),
gr.Checkbox(label="Generate Confidence Score Graph", value=False)
],
outputs=gr.Text(label="Status"),
title="YOLOv8 Object Detection - Frames Output",
description="Upload a short video to save detected frames as images and optionally generate a confidence score graph."
)
if __name__ == "__main__":
iface.launch() |