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()