Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import cv2 | |
| import os | |
| from ultralytics import YOLO | |
| import tempfile | |
| # Load model | |
| model = YOLO("yolov8n.pt") | |
| def detect_faults(video_file): | |
| # Save uploaded video to a temp file path | |
| temp_input_path = video_file | |
| # Create a temp directory to store the output video | |
| temp_dir = tempfile.mkdtemp() | |
| output_path = os.path.join(temp_dir, "output.mp4") | |
| # Open input video | |
| cap = cv2.VideoCapture(temp_input_path) | |
| width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) | |
| height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) | |
| fps = cap.get(cv2.CAP_PROP_FPS) | |
| fourcc = cv2.VideoWriter_fourcc(*'mp4v') | |
| out = cv2.VideoWriter(output_path, fourcc, fps, (width, height)) | |
| while cap.isOpened(): | |
| success, frame = cap.read() | |
| if not success: | |
| break | |
| results = model.predict(source=frame, conf=0.25, verbose=False) | |
| annotated = results[0].plot() | |
| out.write(annotated) | |
| cap.release() | |
| out.release() | |
| return output_path | |
| # Gradio UI | |
| interface = gr.Interface( | |
| fn=detect_faults, | |
| inputs=gr.Video(label="Upload Thermal Video (MP4)"), | |
| outputs=gr.Video(label="Detected Output"), | |
| title="Thermal Fault Detection", | |
| description="Upload a thermal drone video to detect overheating or burning areas using YOLOv8." | |
| ) | |
| if __name__ == "__main__": | |
| interface.launch() | |