Spaces:
Runtime error
Runtime error
File size: 2,392 Bytes
8e56e9a 9e4e41a 8e56e9a 08d45c6 a32fa13 2b412f8 518bedc c53c45c a32fa13 c53c45c 956dd18 a07398d 956dd18 c53c45c a07398d 8e56e9a a07398d 956dd18 c53c45c 956dd18 a07398d 956dd18 a07398d 8e56e9a c53c45c a07398d 8e56e9a 956dd18 a07398d 956dd18 a07398d 9e4e41a 956dd18 a07398d 8e56e9a 956dd18 a32fa13 956dd18 8e56e9a 956dd18 |
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 |
import cv2
import gradio as gr
import uuid
import os
import torch
import tempfile
import shutil
from ultralytics import YOLO
# Automatically download your best.pt model from your dataset repo
model_path = "best.pt"
if not os.path.exists(model_path):
os.system("wget https://huggingface.co/datasets/Prasanna1622/solar-fault-dataset/resolve/main/best.pt")
# Initialize the YOLO model
model = YOLO(model_path)
# Inference function
def detect_faults(video_path):
"""
- video_path: the path to the uploaded video file on disk.
- Returns: path to the annotated output.mp4.
"""
try:
# Create a unique RUN directory so YOLO does not overwrite previous results
unique_id = str(uuid.uuid4())[:8]
project_dir = os.path.join("runs", "detect", unique_id)
os.makedirs(project_dir, exist_ok=True)
print(f"🛠️ Running inference, saving to: {project_dir}")
# Run YOLO predict; this saves the annotated video in project_dir/
results = model.predict(
source=video_path, # path to the uploaded video
save=True,
save_txt=False,
conf=0.5,
project=os.path.join("runs", "detect"),
name=unique_id
)
print("✅ YOLO predict() finished.")
# Check if output video exists
original_name = os.path.basename(video_path)
output_video_path = os.path.join("runs", "detect", unique_id, original_name)
print(f"🛠️ Looking for output video at: {output_video_path}")
if os.path.exists(output_video_path):
print("✅ Output video found, returning it.")
return output_video_path
else:
print(f"❌ Output video NOT found at: {output_video_path}")
return "Error: Annotated video not found."
except Exception as e:
# Print the full exception in logs, return a simple string in UI
print(f"❌ Exception during detect_faults: {e}")
return "Error during processing."
# Create Gradio UI
demo = gr.Interface(
fn=detect_faults,
inputs=gr.Video(label="Upload Input Video"),
outputs=gr.Video(label="Detected Output Video"),
title="Solar Panel Fault Detection from Drone Video",
description="Upload a drone video to detect solar panel faults using a YOLOv8 model."
)
if __name__ == "__main__":
demo.launch()
|