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

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

print("Tensorflow version " + tf.__version__)

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.')
  output = f'Marked-{str(video)}'
  clip = VideoFileClip(video)
  data_dir = 'storage/dataset'
  id_to_name, _ = ClassMapping(data_dir)
  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 video

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