import gradio as gr
from ultralytics import YOLO
import cv2
import os

model = YOLO('fire_forest_detection.pt')

def get_prediction(videos):
  cap = cv2.VideoCapture(videos)

  # Create a video writer object
  fourcc = cv2.VideoWriter_fourcc(*'XVID')
  fps = cap.get(cv2.CAP_PROP_FPS)
  size = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
  output_path = "output.mp4"
  out = cv2.VideoWriter(output_path, fourcc, fps, size)

  while True:
    # Capture the next frame from the video
    ret, frame = cap.read()

    # If the frame is empty, break the loop
    if not ret:
        break

    # Perform object detection on the frame
    results = model(frame)[0]

    im = results.plot()

    out.write(im)

  cap.release()
  cv2.destroyAllWindows()
  
  return output_path


app = gr.Interface(
    get_prediction,
    gr.Video(),
    "playable_video",
    title = "Early Forest Fire Detector",
    description = "Early forest fire detector from Unmanned Aerial Vehicle (UAV) videos as input. Please kindly wait 10-15 seconds for Gradio to load the video example",
    examples = [os.path.join(os.path.abspath(''), "sample_video_5s.mp4")]
)

app.launch(debug=True)