aero-recognize / app.py
chiyoi's picture
fix
4f30d97
raw
history blame
No virus
1.33 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)
id_to_name = {
0: 'Flying',
1: 'Landing',
2: 'Other',
3: 'Straight Taxiing',
4: 'Takeoff',
5: 'Turning Maneuver',
}
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')
examples = [
['examples/ZFLFDfovqls_001310_001320.mp4'], # cspell: disable-line
['examples/Zv7GyH-fpEY_2023.0_2033.0.mp4'],
]
iface = gr.Interface(
fn=fn,
inputs=inputs,
outputs=outputs,
examples=examples,
)
iface.launch()