Refactor code structure and import configurations

app.py

@@ -1,24 +1,18 @@
 import tempfile
+
+import numpy as np
 import gradio as gr
 import tensorflow as tf
 from moviepy.editor import VideoFileClip
+from moviepy.video.io.ImageSequenceClip import ImageSequenceClip
 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_path = 'weights/classifier-8-epoch10.keras'
 classifier = load_classifier(classifier_path)
 
 print('Load detector.')
@@ -30,23 +24,32 @@ def fn(video: gr.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)
+    process_frame = FrameProcessor(detector, classifier)
+    processed_frames = []
+    for frame in clip.iter_frames():
+      processed_frames.append(process_frame(frame))
+      yield processed_frames[-1], None
+    processed_clip = ImageSequenceClip(processed_frames, clip.fps)
+    processed_clip.audio = clip.audio
     clip.write_videofile(output, fps=clip.fps, audio_codec='aac', logger=None)
-  return output
+  yield processed_frames[-1], output
 
 inputs = gr.Video(sources=['upload'], label='Input Video')
-outputs = gr.Video(interactive=False, label='Aeroplane Position and Action Marked')
+outputs = [
+  gr.Image(interactive=False, label='Last Frame Processed'),
+  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'],
-]
+  ['examples/Zv7GyH-fpEY_2023.0_2033.0.mp4']]
 
 iface = gr.Interface(
+  title='Aeroplane Position and Action Detection',
+  description='Detect aeroplane position and action in a video.',
+  theme='soft',
   fn=fn,
   inputs=inputs,
   outputs=outputs,
   examples=examples,
-)
+  cache_examples=False)
 iface.launch()

configurations.py

@@ -0,0 +1,47 @@
+# Data
+data_dir = 'storage/dataset'
+training_ratio = 0.7
+validation_ratio = 0.02
+num_frames = 8
+frame_step = 1
+resolution = 224
+frame_size = (resolution, resolution)
+
+id_to_name = {
+  0: 'Flying',
+  1: 'Landing',
+  2: 'Other',
+  3: 'Straight Taxiing',
+  4: 'Takeoff',
+  5: 'Turning Maneuver',
+}
+
+name_to_id = {
+  'Flying': 0,
+  'Landing': 1,
+  'Other': 2,
+  'Straight Taxiing': 3,
+  'Takeoff': 4,
+  'Turning Maneuver': 5,
+}
+
+# Model
+model_id = 'a0'
+checkpoint_dir = f'storage/pretrained_weights/movinet_{model_id}_base'
+num_classes = 6
+
+# Inference
+detect_object_frame_steps = 5
+classify_action_frame_steps = 15
+classify_action_num_frames = 8
+
+# Train
+train_id = 8
+batch_size = 16
+learning_rate = 0.001
+epochs = 15
+model_save_path = f'storage/output/classifier-{train_id}.keras'
+log_dir = f'storage/logs/classifier-{train_id}.log'
+
+# Train more
+initial_epoch = 0

core/data.py

@@ -5,11 +5,7 @@ import cv2
 import numpy as np
 import tensorflow as tf
 
-training_ratio = 0.7
-validation_ratio = 0.02
-num_frames = 8
-frame_step = 15
-frame_size = (224, 224)
+from configurations import *
 
 def format_frame(frame):
   frame = tf.image.convert_image_dtype(frame, tf.float32)
@@ -35,37 +31,35 @@ def pick_frames(video: str):
   frames = frames[..., [2, 1, 0]]
   return frames
 
-def Data(data_dir: str):
-  data_dir = Path(data_dir)
+def Data():
+  data_dir_path = Path(data_dir)
   return {
-  'training':{
-    a.name: (lambda ps: ps[:int(len(ps) * training_ratio)])([x for x in a.iterdir()])
-    for a in data_dir.iterdir()
-  },
-  'validation': {
-    a.name: (lambda ps: ps[
-      int(len(ps) * training_ratio) :
-      int(len(ps) * (training_ratio + validation_ratio))
-    ])([x for x in a.iterdir()])
-    for a in data_dir.iterdir()
-    },
+    'training': {
+      a.name: (
+        lambda ps: ps[
+          :int(len(ps) * training_ratio)])(
+        [x for x in a.iterdir()])
+      for a in data_dir_path.iterdir()},
+    'validation': {
+      a.name: (
+        lambda ps: ps[
+          int(len(ps) * training_ratio):
+          int(len(ps) * (training_ratio + validation_ratio))])(
+        [x for x in a.iterdir()])
+      for a in data_dir_path.iterdir()},
+    'testing': {
+      a.name: (
+        lambda ps: ps[
+          int(len(ps) * (training_ratio + validation_ratio)):])(
+        [x for x in a.iterdir()])
+      for a in data_dir_path.iterdir()},
   }
 
-def ClassMapping(data_dir: str):
-  data_dir = Path(data_dir)
-  id_to_name = sorted([x.name for x in data_dir.iterdir()])
-  name_to_id = {
-    name: i
-    for i, name in enumerate(id_to_name)
-  }
-  return (id_to_name, name_to_id)
-
-def FrameGenerator(data_dir: str, split: Literal['training', 'validation']):
-  _, name_to_id = ClassMapping(data_dir)
-  data = Data(data_dir)
+def FrameGenerator(split: Literal['training', 'validation']):
+  data = Data()
   def generator():
     pairs = [
-      (video, class_name)
+      (str(video), class_name)
       for class_name, videos in data[split].items()
       for video in videos
     ]

core/inference.py

@@ -2,15 +2,12 @@ from imgviz import instances2rgb
 import tensorflow as tf
 import numpy as np
 
+from configurations import *
 from core.data import format_frame
 
 # detections: (classes: list of class_name, boxes: list of [x1, y1, x2, y2])
 # actions: list of f'{action_name}: {confidence}'
 
-detect_object_frame_steps = 5
-classify_action_frame_steps = 15
-classify_action_num_frames = 8
-
 def detect_object(detector, frame):
   result = detector(frame, classes=4, verbose=False)[0]
   classes = result.boxes.cls.numpy()
@@ -21,12 +18,12 @@ def detect_object(detector, frame):
   )
   return detections
 
-def classify_action(classifier, frames, id_to_name):
+def classify_action(classifier, frames):
   actions = []
   frames = np.array(frames)
-  frames = frames[..., [2, 1, 0]]
+  # frames = frames[..., [2, 1, 0]]
   frames = tf.expand_dims(frames, 0)
-  output = classifier(frames, training=False)
+  output = classifier(frames)
   confidences = tf.nn.softmax(output).numpy()[0]
   for (class_id, confidence) in enumerate(confidences):
     other_class_id = 2
@@ -65,7 +62,7 @@ def draw_boxes(frame, detections, actions):
   )
   return frame
 
-def FrameProcessor(detector, classifier, id_to_name):
+def FrameProcessor(detector, classifier):
   current_frame = 0
   frames = []
   actions = []
@@ -80,7 +77,7 @@ def FrameProcessor(detector, classifier, id_to_name):
       detections = detect_object(detector, frame)
     if len(frames) == classify_action_num_frames:
       print(f'Classify action: Until frame {current_frame}')
-      actions = classify_action(classifier, frames, id_to_name)
+      actions = classify_action(classifier, frames)
       frames = []
     frame = draw_boxes(frame, detections, actions)
     return frame

core/model.py

@@ -3,42 +3,37 @@ from tensorflow import keras
 from official.projects.movinet.modeling import movinet
 from official.projects.movinet.modeling import movinet_model
 
-model_id = 'a1'
-num_classes = 6
-num_frames = 8
-resolution = 224
+from configurations import *
 
-batch_size = 32
-learning_rate = 0.001
-backbone_trainable = True
+def load_backbone():
+  return movinet.Movinet()
 
-def build_classifier_with_pretrained_weights(checkpoint_dir: str):
-  backbone = movinet.Movinet(model_id=model_id)
-  backbone.trainable = backbone_trainable
-  model = movinet_model.MovinetClassifier(backbone=backbone, num_classes=600)
+def build_classifier():
+  backbone = load_backbone()
+  model = movinet_model.MovinetClassifier(
+    backbone=backbone,
+    num_classes=600)
   checkpoint_path = tf.train.latest_checkpoint(checkpoint_dir)
   checkpoint = tf.train.Checkpoint(model=model)
   status = checkpoint.restore(checkpoint_path)
   status.assert_existing_objects_matched()
-  model = movinet_model.MovinetClassifier(
-    backbone=backbone,
-    num_classes=num_classes,
-  )
   model.build([batch_size, num_frames, resolution, resolution, 3])
-  return model
+  output = keras.layers.Dense(num_classes)
+  return keras.Sequential(layers=[model, output])
 
-def load_classifier(weights_path: str):
-  backbone = movinet.Movinet(model_id=model_id)
+def load_classifier():
+  backbone = load_backbone()
   model = movinet_model.MovinetClassifier(
     backbone=backbone,
     num_classes=num_classes,
-  )
-  model.build([1, num_frames, resolution, resolution, 3])
-  model.load_weights(weights_path)
+    output_states=True)
+  model.build([batch_size, num_frames, resolution, resolution, 3])
+  output = keras.layers.Dense(num_classes)
+  model = keras.Sequential(layers=[model, output])
+  model.load_weights(model_save_path)
   return model
 
 def compile_classifier(model):
   loss = keras.losses.SparseCategoricalCrossentropy(from_logits=True)
   optimizer = keras.optimizers.Adam(learning_rate=learning_rate)
   model.compile(optimizer=optimizer, loss=loss, metrics=['accuracy'])
-  return model

tmpyt25_04fTEMP_MPY_wvf_snd.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fe6df5c78cebb5a214ed7c83d8826c441d088ba75198a65429dbcb3619959f53
+size 162883

weights/classifier-8-epoch10.keras

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c11a51f095a4902755f375740980537257764536bd089d2ae7cbe1cbb8343184
+size 38477915