Spaces:
Sleeping
Sleeping
| 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() |