import os
import tempfile
import gradio as gr
import tensorflow as tf
from moviepy.editor import VideoFileClip
from ultralytics import YOLO

from core.model import load_classifier
from core.inference import FrameProcessor

print("Tensorflow version " + tf.__version__)

id_to_name = {
  0: 'Flying',
  1: 'Landing',
  2: 'Other',
  3: 'Straight Taxiing',
  4: 'Takeoff',
  5: 'Turning Maneuver',
}

print('Load classifier.')
classifier_path = 'weights/classifier-7.keras'
classifier = load_classifier(classifier_path)

print('Load detector.')
detector_path = 'weights/yolov8n.pt'
detector = YOLO(detector_path)

def fn(video: gr.Video):
  print('Process video.')
  with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as f:
    output = f.name
    clip = VideoFileClip(video)
    process_frame = FrameProcessor(detector, classifier, id_to_name)
    clip = clip.fl_image(process_frame)
    clip.write_videofile(output, fps=clip.fps, audio_codec='aac', logger=None)
  return output

inputs = gr.Video(sources=['upload'], label='Input Video')
outputs = gr.Video(interactive=False, label='Aeroplane Position and Action Marked')

iface = gr.Interface(
  fn=fn,
  inputs=inputs,
  outputs=outputs,
)
iface.launch()