aero-recognize / app.py
chiyoi's picture
update
139dd3e
raw history blame
No virus
1.11 kB
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()