Update: Add choices for actions.

app.py

@@ -7,7 +7,7 @@ from moviepy.video.io.ImageSequenceClip import ImageSequenceClip
 
 from configuration import Config
 from model import load_classifier, load_detector
-from inference import format_frame, detect_object, classify_action, draw_boxes
+from inference import format_frame, detect_object, classify_action, draw_boxes, draw_classes
 
 config = Config()
 print(f'TensorFlow {tf.__version__}')
@@ -20,9 +20,13 @@ classifier.summary()
 print('Load detector.')
 detector = load_detector(config)
 
-def fn(video: gr.Video):
+def fn(video: gr.Video, actions: list[int]):
   print('Process video.')
-  with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as f:
+  do_detect = 0 in actions
+  do_classify = 1 in actions
+  if not do_detect and not do_classify:
+    return video
+  with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as f:
     output = f.name
     clip = VideoFileClip(video)
     processed_frames = []
@@ -31,15 +35,21 @@ def fn(video: gr.Video):
     detections = ([], [])
     for i, frame in enumerate(clip.iter_frames()):
       if i % config.classify_action_frame_step == 0:
-        frames.append(format_frame(frame, config))
+        if do_classify:
+          frames.append(format_frame(frame, config))
       if i % config.detect_object_frame_step == 0:
         print(f'Detect object: Frame {i}')
-        detections = detect_object(detector, frame)
+        if do_detect:
+          detections = detect_object(detector, frame)
       if len(frames) == config.classify_action_num_frames:
         print(f'Classify action: Until frame {i}')
-        actions = classify_action(classifier, frames, config.id_to_name)
-        frames = []
-      frame = draw_boxes(frame, detections, actions)
+        if do_classify:
+          actions = classify_action(classifier, frames, config.id_to_name)
+          frames = []
+      if do_detect:
+        frame = draw_boxes(frame, detections, actions, do_classify)
+      else:
+        frame = draw_classes(frame, actions)
       processed_frames.append(frame)
       if i % config.yield_frame_steps == 0:
         quality = 9
@@ -53,10 +63,16 @@ def fn(video: gr.Video):
     processed_clip.write_videofile(output, fps=clip.fps, audio_codec='aac', logger=None)
   yield frame, output
 
-inputs = gr.Video(sources=['upload'], label='Input Video')
+inputs = [
+  gr.Video(sources=['upload'], label='输入视频片段'),
+  gr.CheckboxGroup(
+    ['飞机检测', '飞机行为识别'],
+    label='执行任务',
+    info='可以选择仅执行飞机检测任务或仅执行飞机行为识别任务作为演示。',
+    type='index')]
 outputs = [
-  gr.Image(interactive=False, label='Last Frame Processed'),
-  gr.Video(interactive=False, label='Aeroplane Position and Action Marked')]
+  gr.Image(interactive=False, label='最新处理的视频帧'),
+  gr.Video(interactive=False, label='标记飞机位置及行为的视频片段')]
 
 examples = [
   ['examples/ZFLFDfovqls_001310_001320.mp4'], # cspell: disable-line

configuration.py

@@ -29,4 +29,4 @@ class Config:
   detect_object_frame_step = 5
   classify_action_frame_step = frame_step
   classify_action_num_frames = num_frames
-  yield_frame_steps = 5
+  yield_frame_steps = 3

inference.py

@@ -1,6 +1,6 @@
 import tensorflow as tf
 import numpy as np
-from imgviz import instances2rgb
+from imgviz import instances2rgb, label2rgb
 
 from configuration import Config
 
@@ -33,7 +33,7 @@ def classify_action(classifier, frames, id_to_name):
       actions.append(f'{id_to_name[class_id]}: {confidence:.2f}')
   return actions
 
-def draw_boxes(frame, detections, actions):
+def draw_boxes(frame, detections, actions, do_classify):
   (classes, boxes) = detections
   max_area = 0
   max_area_id = 0
@@ -45,9 +45,12 @@ def draw_boxes(frame, detections, actions):
   labels = [0 for _ in classes]
   colormap = [(0x39, 0xc5, 0xbb)]
   line_width = 2
-  captions = [
-    f'{class_name}\n' + '\n'.join(actions if i == max_area_id else [])
-    for (i, class_name) in enumerate(classes)]
+  if not do_classify:
+    captions = classes
+  else:
+    captions = [
+      f'{class_name}\n' + '\n'.join(actions if i == max_area_id else [])
+      for (i, class_name) in enumerate(classes)]
   bboxes = [
     [box[1], box[0], box[3], box[2]]
     for box in boxes]
@@ -61,6 +64,15 @@ def draw_boxes(frame, detections, actions):
     line_width=line_width)
   return frame
 
+def draw_class(frame, actions):
+  label=['Airplane']
+  label_names = ['Airplane\n' + '\n'.join(actions)]
+  frame = label2rgb(
+    image=frame,
+    label=label,
+    label_names=label_names)
+  return frame
+
 def FrameProcessor(detector, classifier, config: Config):
   current_frame = 0
   frames = []