Spaces:
Runtime error
Runtime error
File size: 3,098 Bytes
fc9052c 65be584 fc9052c 65be584 fc9052c 3c00238 4d4428a fc9052c 65be584 3c00238 4d4428a 3c00238 4d4428a 3c00238 fc9052c 3c00238 fc9052c 3c00238 fc9052c 4d4428a 95982ed b862d21 4d4428a fc9052c |
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 |
import gradio as gr
import cv2
from PIL import Image
import numpy as np
from ultralytics import YOLO
from huggingface_hub import hf_hub_download
import os
# Verify paths and Hugging Face repository details
REPO_ID = "StephanST/WALDO30" # Replace with the correct repo ID if different
MODEL_FILENAME = "WALDO30_yolov8m_640x640.pt" # Replace if the filename is different
# Download the model from Hugging Face
try:
model_path = hf_hub_download(repo_id=REPO_ID, filename=MODEL_FILENAME)
print(f"Model downloaded successfully to: {model_path}")
except Exception as e:
raise RuntimeError(f"Failed to download model from Hugging Face. Verify `repo_id` and `filename`. Error: {e}")
# Load the YOLOv8 model
try:
model = YOLO(model_path) # Load the YOLOv8 model
print("Model loaded successfully!")
except Exception as e:
raise RuntimeError(f"Failed to load the YOLO model. Verify the model file at `{model_path}`. Error: {e}")
# Detection function for images
def detect_on_image(image):
try:
results = model(image) # Perform detection
annotated_frame = results[0].plot() # Get annotated image
return Image.fromarray(annotated_frame)
except Exception as e:
raise RuntimeError(f"Error during image processing: {e}")
# Detection function for videos
def detect_on_video(video):
try:
temp_video_path = "processed_video.mp4"
cap = cv2.VideoCapture(video)
fourcc = cv2.VideoWriter_fourcc(*"mp4v")
out = cv2.VideoWriter(temp_video_path, fourcc, cap.get(cv2.CAP_PROP_FPS),
(int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))))
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
results = model(frame) # Perform detection
annotated_frame = results[0].plot() # Get annotated frame
out.write(annotated_frame)
cap.release()
out.release()
return temp_video_path
except Exception as e:
raise RuntimeError(f"Error during video processing: {e}")
# Gradio Interface using Blocks
with gr.Blocks() as app:
gr.Markdown("# Sat ESPR View")
gr.Markdown("Upload an image or video to see object detection results.")
# Image processing block
with gr.Row():
with gr.Column():
image_input = gr.Image(type="pil", label="Upload Image")
image_button = gr.Button("Detect on Image")
with gr.Column():
image_output = gr.Image(type="pil", label="Detected Image")
# Video processing block
with gr.Row():
with gr.Column():
video_input = gr.Video(label="Upload Video")
video_button = gr.Button("Detect on Video")
with gr.Column():
video_output = gr.Video(label="Detected Video")
# Set up events
image_button.click(detect_on_image, inputs=image_input, outputs=image_output)
video_button.click(detect_on_video, inputs=video_input, outputs=video_output)
if __name__ == "__main__":
app.launch()
|