aero-recognize / app.py
chiyoi's picture
update: Add examples.
92ddd3a
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()