add initial app

app.py

@@ -0,0 +1,84 @@
+from datetime import timedelta
+import gradio as gr
+from sentence_transformers import SentenceTransformer
+import torchvision
+from sklearn.metrics.pairwise import cosine_similarity
+import numpy as np
+
+from inference import Inference
+import utils
+
+model_checkpoint = 'saved_model'
+encoder_model_name = 'google/vit-large-patch32-224-in21k'
+decoder_model_name = 'gpt2'
+frame_step = 300
+
+inference = Inference(
+    decoder_model_name=decoder_model_name,
+    model_checkpoint=model_checkpoint,
+)
+
+model = SentenceTransformer('all-mpnet-base-v2')
+
+def search_in_video(video, query):
+    result = torchvision.io.read_video(video)
+    video = result[0]
+    video_fps = result[2]['video_fps']
+
+    video_segments = [
+        video[idx:idx + frame_step, :, :, :] for idx in range(0, video.shape[0], frame_step)
+    ]
+
+    generated_texts = []
+
+    for video_seg in video_segments:
+        pixel_values = utils.video2image(video_seg, encoder_model_name)
+
+        generated_text = inference.generate_text(pixel_values, encoder_model_name)
+        generated_texts.append(generated_text)
+
+    sentences = [query] + generated_texts
+
+    sentence_embeddings = model.encode(sentences)
+
+    similarities = cosine_similarity(
+        [sentence_embeddings[0]],
+        sentence_embeddings[1:]
+    )
+    arg_sorted_similarities = np.argsort(similarities)
+
+    ordered_similarity_scores = similarities[0][arg_sorted_similarities]
+
+    best_video = video_segments[arg_sorted_similarities[0, -1]]
+    torchvision.io.write_video('best.mp4', best_video, video_fps)
+
+    total_frames = video.shape[0]
+
+    video_frame_segs = [
+        [idx, min(idx + frame_step, total_frames)] for idx in range(0, total_frames, frame_step)
+    ]
+    ordered_start_ends = []
+
+    for [start, end] in video_frame_segs:
+        td = timedelta(seconds=(start / video_fps))
+        s = round(td.total_seconds(), 2)
+        
+        td = timedelta(seconds=(end / video_fps))
+        e = round(td.total_seconds(), 2)
+        
+        ordered_start_ends.append(f'{s}:{e}')
+    
+    ordered_start_ends = np.array(ordered_start_ends)[arg_sorted_similarities]
+
+    labels_to_scores = dict(
+        zip(ordered_start_ends[0].tolist(), ordered_similarity_scores[0].tolist())
+    )
+
+    return 'best.mp4', labels_to_scores
+
+app = gr.Interface(
+    fn=search_in_video,
+    inputs=['video', 'text'],
+    outputs=['video', gr.outputs.Label(num_top_classes=3, type='auto')],
+)
+app.launch(share=True)

inference.py

@@ -0,0 +1,29 @@
+import torch
+from transformers import AutoTokenizer, VisionEncoderDecoderModel
+
+import utils
+
+class Inference:
+  def __init__(self, decoder_model_name, model_checkpoint, max_length=32):
+    self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+    self.tokenizer = AutoTokenizer.from_pretrained(decoder_model_name)
+    self.encoder_decoder_model = VisionEncoderDecoderModel.from_pretrained(model_checkpoint)
+    self.encoder_decoder_model.to(self.device)
+
+    self.max_length = max_length
+
+  def generate_text(self, video, encoder_model_name):
+    if isinstance(video, str):
+      pixel_values = utils.video2image_from_path(video, encoder_model_name)
+    else:
+      pixel_values = video
+
+    if not self.tokenizer.pad_token:
+      self.tokenizer.add_special_tokens({'pad_token': '[PAD]'})
+      self.encoder_decoder_model.decoder.resize_token_embeddings(len(self.tokenizer))
+
+    generated_ids = self.encoder_decoder_model.generate(pixel_values.unsqueeze(0).to(self.device), max_length=self.max_length)
+    generated_text = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+
+    return generated_text

requirements.txt

@@ -0,0 +1,4 @@
+nltk==3.7
+tqdm==4.64.0
+scikit-learn==1.1.1
+sentence-transformers==2.2.0

utils.py

@@ -0,0 +1,44 @@
+from transformers import ViTFeatureExtractor
+import torchvision
+import torchvision.transforms.functional as fn
+import torch as th
+import os
+import pickle
+
+
+def video2image_from_path(video_path, feature_extractor_name):
+    video = torchvision.io.read_video(video_path)
+
+    return video2image(video[0], feature_extractor_name)
+
+
+def video2image(video, feature_extractor_name):
+    feature_extractor = ViTFeatureExtractor.from_pretrained(
+        feature_extractor_name
+    )
+
+    vid = th.permute(video, (3, 0, 1, 2))
+    samp = th.linspace(0, vid.shape[1]-1, 49, dtype=th.long)
+    vid = vid[:, samp, :, :]
+
+    im_l = list()
+    for i in range(vid.shape[1]):
+        im_l.append(vid[:, i, :, :])
+
+    inputs = feature_extractor(im_l, return_tensors="pt")
+
+    inputs = inputs['pixel_values']
+
+    im_h = list()
+    for i in range(7):
+        im_v = th.cat((inputs[0+i*7, :, :, :],
+                       inputs[1+i*7, :, :, :],
+                       inputs[2+i*7, :, :, :],
+                       inputs[3+i*7, :, :, :],
+                       inputs[4+i*7, :, :, :],
+                       inputs[5+i*7, :, :, :],
+                       inputs[6+i*7, :, :, :]), 2)
+        im_h.append(im_v)
+    resize = fn.resize(th.cat(im_h, 1), size=[224])
+
+    return resize