detection / app.py
Sirivennela's picture
Update app.py
b6dc84d verified
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()