feat: add pipeline to predict

app.py

@@ -1,20 +1,30 @@
 import gradio as gr 
+from run_on_video.run import MomentDETRPredictor
+from moviepy.video.io.ffmpeg_tools import ffmpeg_extract_subclip
 
-TITLE = """<h2 align="center"> ✍️ Highlight Detection with MomentDETR </h2>"""
+ckpt_path = "run_on_video/moment_detr_ckpt/model_best.ckpt"
+clip_model_name_or_path = "ViT-B/32"
+device = 'cpu'
 
-def submit_video(input_video, retrieval_text):
-    print(input_video)
-    print(retrieval_text)
-    return input_video
+moment_detr_predictor = MomentDETRPredictor(
+    ckpt_path=ckpt_path,
+    clip_model_name_or_path=clip_model_name_or_path,
+    device=device
+)
 
+def trim_video(video_path, start, end, output_file='result.mp4'):
+    ffmpeg_extract_subclip(video_path, start, end, targetname=output_file)
+    return output_file
 
 with gr.Blocks() as demo:
-    gr.HTML(TITLE)
+    output_videos = gr.State([])
+    moment_scores = gr.State([])
+    gr.HTML("""<h2 align="center"> ✍️ Highlight Detection with MomentDETR </h2>""")
     with gr.Row():
         with gr.Blocks():
             with gr.Column():
-                gr.Markdown("### Input Video")
-                input_video = gr.PlayableVideo().style(height=500)
+                gr.HTML("""<h3 align="center"> Input Video </h3>""")
+                input_video = gr.PlayableVideo()
                 retrieval_text = gr.Textbox(
                     placeholder="What should be highlighted?",
                     visible=True
@@ -22,17 +32,59 @@ with gr.Blocks() as demo:
                 submit =gr.Button("Submit")
         with gr.Blocks():
             with gr.Column(): 
-                gr.Markdown("### Results")
-                with gr.Row():
-                    output_video = gr.PlayableVideo().style(height=500)
+                gr.HTML("""<h3 align="center"> Highlight Videos </h3>""")
+                display_score = gr.Markdown("### Moment Score: ")
+                radio_button = gr.Radio(
+                    choices=[i for i in range(10)], 
+                    label="Moments", 
+                    value=0
+                )
+                playable_video = gr.PlayableVideo(visible=True)
 
+        def update_video_player(radio_value, output_videos, moment_scores):
+            return {
+                playable_video: output_videos[radio_value], 
+                display_score: f'### Moment Score: {moment_scores[radio_value]}'
+            }
+                
+    def submit_video(input_video, retrieval_text):
+        print(f'== video path: {input_video}')
+        print(f'== retrieval_text: {retrieval_text}')
+        if retrieval_text is None:
+            retrieval_text = ''
+        predictions, video_frames = moment_detr_predictor.localize_moment(
+            video_path=input_video, 
+            query_list=[retrieval_text]
+        )
+        pred_windows = [[pred[0], pred[1]]for pred in predictions[0]['pred_relevant_windows']]
+        scores = [pred[-1] for pred in predictions[0]['pred_relevant_windows']]
 
+        print(f'== predict start end time: {pred_windows}')
+        print(f'== prediction scores: {scores}')
+        output_files = [ trim_video(
+            video_path=input_video, 
+            start=pred_windows[i][0], 
+            end=pred_windows[i][1],
+            output_file=f'{i}.mp4'
+        ) for i in range(10)]
+        print(f'== output_files: {output_files}')
+        return { 
+            output_videos: output_files, 
+            moment_scores: scores,
+            playable_video:  output_files[0],
+            display_score: f'### Moment Score: {scores[0]}'
+        }
 
+    radio_button.change(
+        fn=update_video_player, 
+        inputs=[radio_button, output_videos, moment_scores],
+        outputs=[playable_video, display_score]
+    )
 
     submit.click(
-            fn=submit_video, 
-            inputs=[input_video, retrieval_text], 
-            outputs=[output_video]
-        )
+        fn=submit_video, 
+        inputs=[input_video, retrieval_text], 
+        outputs=[output_videos, moment_scores, playable_video, display_score]
+    )
 
 demo.launch()

requirements.txt

@@ -11,4 +11,5 @@ pandas
 ffmpeg-python
 ftfy
 regex
-Pillow 
+Pillow 
+moviepy

run_on_video/data_utils.py

@@ -3,7 +3,7 @@ import os
 import numpy as np
 import ffmpeg
 import math
-import clip
+import run_on_video.clip as clip
 
 
 class ClipFeatureExtractor:

run_on_video/run.py

@@ -1,7 +1,7 @@
 import torch
 
-from data_utils import ClipFeatureExtractor
-from model_utils import build_inference_model
+from run_on_video.data_utils import ClipFeatureExtractor
+from run_on_video.model_utils import build_inference_model
 from utils.tensor_utils import pad_sequences_1d
 from moment_detr.span_utils import span_cxw_to_xx
 from utils.basic_utils import l2_normalize_np_array