Spaces:
Sleeping
Sleeping
File size: 4,137 Bytes
6cb50ff a03d512 04f4d0b 8c84287 4527f8f 8c84287 2e47361 8c84287 a668c53 fa66a0f 04f4d0b 8c84287 2e47361 8c84287 2e47361 8c84287 2e47361 8c84287 a03d512 8c84287 4527f8f 8c84287 2e47361 8c84287 2e47361 8c84287 2e47361 8c84287 2e47361 6cb50ff 4527f8f 2e47361 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 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 |
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
import uuid
# 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:
yield "Error: No video uploaded", []
return
# 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():
yield "Error: Could not open video file", []
return
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
detected_frame_paths = [] # Store paths of frames with detections
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)
# Save and yield frame if objects are detected
if results[0].boxes is not None and len(results[0].boxes) > 0:
annotated_frame = results[0].plot()
frame_filename = os.path.join(output_folder, f"frame_{frame_count:04d}.jpg")
cv2.imwrite(frame_filename, annotated_frame)
detected_frame_paths.append(frame_filename)
# Collect confidence scores for plotting
confs = results[0].boxes.conf.cpu().numpy()
confidence_scores.extend(confs)
# Yield current status and gallery
yield f"Processed frame {frame_count} with detections", detected_frame_paths[:]
cap.release()
# Generate confidence score plot if requested
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()
detected_frame_paths.append(graph_path)
# Final yield with all results
status = f"Saved {len(detected_frame_paths)} frames with detections in {output_folder}. {f'Graph saved as {graph_path}' if plot_graphs and confidence_scores else ''}"
yield status, detected_frame_paths
# Gradio interface
with gr.Blocks() as iface:
gr.Markdown("# YOLOv8 Object Detection - Real-time Frame Output")
gr.Markdown("Upload a short video to view frames with detections immediately in a gallery. Optionally generate a confidence score graph.")
with gr.Row():
video_input = gr.Video(label="Upload Video")
output_folder = gr.Textbox(label="Output Folder", value="detected_frames")
plot_graphs = gr.Checkbox(label="Generate Confidence Score Graph", value=False)
submit_button = gr.Button("Process Video")
status_output = gr.Text(label="Status")
gallery_output = gr.Gallery(label="Detected Frames and Graph", preview=True, columns=3)
submit_button.click(
fn=process_video,
inputs=[video_input, output_folder, plot_graphs],
outputs=[status_output, gallery_output],
concurrency_limit=1
)
if __name__ == "__main__":
iface.launch() |