Create app.py

app.py

@@ -0,0 +1,90 @@
+import cv2
+import gradio as gr
+from PIL import Image
+from transformers import BridgeTowerForImageAndTextRetrieval, BridgeTowerProcessor
+
+model_id = "BridgeTower/bridgetower-large-itm-mlm-gaudi"
+processor = BridgeTowerProcessor.from_pretrained(model_id)
+model = BridgeTowerForImageAndTextRetrieval.from_pretrained(model_id)
+
+# Process a frame
+def process_frame(image, texts):
+    scores = {}
+    texts = texts.split(",")
+    for t in texts:
+        encoding = processor(image, t, return_tensors="pt")
+        outputs = model(**encoding)
+        scores[t] = "{:.2f}".format(outputs.logits[0, 1].item())
+        # sort scores in descending order
+        scores = dict(sorted(scores.items(), key=lambda item: item[1], reverse=True))
+    return scores
+
+
+# Process a video
+def process(video, text, sample_rate, min_score):
+    video = cv2.VideoCapture(video)
+    fps = round(video.get(cv2.CAP_PROP_FPS))
+    frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
+    length = frames // fps
+    print(f"{fps} fps, {frames} frames, {length} seconds")
+
+    frame_count = 0
+    clips = []
+    clip_images = []
+    clip_started = False
+    while True:
+        ret, frame = video.read()
+        if not ret:
+            break
+
+        if frame_count % (fps * sample_rate) == 0:
+            frame = Image.fromarray(frame)
+            score = process_frame(frame, text)
+            # print(f"{frame_count} {scores}")
+
+            if float(score[text]) > min_score:
+                if clip_started:
+                    end_time = frame_count / fps
+                else:
+                    clip_started = True
+                    start_time = frame_count / fps
+                    end_time = start_time
+                    start_score = score[text]
+                    clip_images.append(frame)
+            elif clip_started:
+                clip_started = False
+                end_time = frame_count / fps
+                clips.append((start_score, start_time, end_time))
+        frame_count += 1
+    return clip_images, clips
+
+
+# Inputs
+video = gr.Video(label="Video")
+text = gr.Text(label="Text query")
+sample_rate = gr.Number(value=5, label="Sample rate (1 frame every 'n' seconds)")
+min_score = gr.Number(value=3, label="Minimum score")
+
+# Output
+gallery = gr.Gallery(label="Images")
+clips = gr.Text(label="Clips (score, start time, end time)")
+
+description = "This Space lets you run semantic search on a video."
+
+iface = gr.Interface(
+    description=description,
+    fn=process,
+    inputs=[video, text, sample_rate, min_score],
+    outputs=[gallery, clips],
+    examples=[
+        [
+            "video.mp4",
+            "wild bears",
+            5,
+            3,
+        ]
+    ],
+    allow_flagging="never",
+)
+
+iface.launch()