aero-recognize / app.py
chiyoi's picture
fix
b44c3e2
raw
history blame
No virus
1.36 kB
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')
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()