chore: update dependency version, UI

README.md

@@ -1,6 +1,6 @@
 ---
 title: Highlight Detection with MomentDETR 
-emoji: ✍️ 
+emoji: 🎞️
 colorFrom: purple
 colorTo: yellow
 sdk: gradio

app.py

@@ -1,10 +1,15 @@
 import gradio as gr 
 from run_on_video.run import MomentDETRPredictor
 from moviepy.video.io.ffmpeg_tools import ffmpeg_extract_subclip
+import torch 
+
+DESCRIPTION = """
+_This Space demonstrates model [QVHighlights: Detecting Moments and Highlights in Videos via Natural Language Queries](https://arxiv.org/abs/2107.09609), NeurIPS 2021, by [Jie Lei](http://www.cs.unc.edu/~jielei/), [Tamara L. Berg](http://tamaraberg.com/), [Mohit Bansal](http://www.cs.unc.edu/~mbansal/)_
+"""
 
 ckpt_path = "run_on_video/moment_detr_ckpt/model_best.ckpt"
 clip_model_name_or_path = "ViT-B/32"
-device = 'cpu'
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 moment_detr_predictor = MomentDETRPredictor(
     ckpt_path=ckpt_path,
@@ -16,20 +21,24 @@ 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:
+def display_prediction(result):
+    return f'Moment({result[0]} ~ {result[1]}), Score: {result[2]}'
+
+with gr.Blocks(theme=gr.themes.Default()) as demo:
     output_videos = gr.State([])
-    moment_scores = gr.State([])
-    gr.HTML("""<h2 align="center"> ✍️ Highlight Detection with MomentDETR </h2>""")
+    moment_prediction = gr.State([])
+    gr.HTML("""<h2 align="center"> 🎞️  Highlight Detection with MomentDETR </h2>""")
+    gr.Markdown(DESCRIPTION)
     with gr.Column():
         with gr.Row():
             with gr.Blocks():
                 with gr.Column():
                     gr.HTML("""<h3 align="center"> Input Video </h3>""")
-                    input_video = gr.PlayableVideo()
+                    input_video = gr.Video(label="Please input mp4", height=400)
             with gr.Blocks():
                 with gr.Column(): 
                     gr.HTML("""<h3 align="center"> Highlight Videos </h3>""")
-                    playable_video = gr.PlayableVideo()
+                    playable_video = gr.Video(height=400)
         with gr.Row():
             with gr.Column():
                 retrieval_text = gr.Textbox(
@@ -41,15 +50,15 @@ with gr.Blocks() as demo:
             with gr.Column():
                 display_score = gr.Markdown("### Moment Score: ")
                 radio_button = gr.Radio(
-                    choices=[i for i in range(10)], 
+                    choices=[i+1 for i in range(10)], 
                     label="Moments", 
-                    value=0
+                    value=1
                 )
 
-        def update_video_player(radio_value, output_videos, moment_scores):
+        def update_video_player(radio_value, output_videos, moment_prediction):
             return {
-                playable_video: output_videos[radio_value], 
-                display_score: f'### Moment Score: {moment_scores[radio_value]}'
+                playable_video: output_videos[radio_value-1], 
+                display_score: display_prediction(moment_prediction[radio_value-1])
             }
                 
     def submit_video(input_video, retrieval_text):
@@ -61,35 +70,32 @@ with gr.Blocks() as demo:
             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}')
+        predictions = predictions[0]['pred_relevant_windows']
+        pred_windows = [[pred[0], pred[1]]for pred in predictions]
         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,
+            moment_prediction: predictions,
             playable_video:  output_files[0],
-            display_score: f'### Moment Score: {scores[0]}'
+            display_score: display_prediction(predictions[0])
         }
 
     radio_button.change(
         fn=update_video_player, 
-        inputs=[radio_button, output_videos, moment_scores],
+        inputs=[radio_button, output_videos, moment_prediction],
         outputs=[playable_video, display_score]
     )
 
     submit.click(
         fn=submit_video, 
         inputs=[input_video, retrieval_text], 
-        outputs=[output_videos, moment_scores, playable_video, display_score]
+        outputs=[output_videos, moment_prediction, playable_video, display_score]
     )
 
 demo.launch()

requirements.txt

@@ -12,4 +12,5 @@ ffmpeg-python
 ftfy
 regex
 Pillow 
-moviepy
+moviepy==1.0.3
+gradio==3.41.0