Reorganize interface and code to be more modular and add necessary debian packages to install

app.py

@@ -1,23 +1,22 @@
-from torch._C import device
 import ffmpeg
+import torch
 import youtube_dl
 
 import numpy as np
+import streamlit as st
 
-from PIL import Image
-import requests
-
-import torch
 from sentence_transformers import SentenceTransformer, util, models
 from clip import CLIPModel
-# from sentence_transformers.models import CLIPModel
 from PIL import Image
 
-clip = CLIPModel()
-model = SentenceTransformer(modules=[clip]).to(dtype=torch.float32, device=torch.device('cpu'))
+@st.cache(allow_output_mutation=True, max_entries=1)
+def get_model():
+    clip = CLIPModel()
+    model = SentenceTransformer(modules=[clip]).to(dtype=torch.float32, device=torch.device('cpu'))
+    return model
 
 
-def get_embedding(query, video):
+def get_embedding(model, query, video):
     text_emb = model.encode(query, device='cpu')
 
     # Encode an image:
@@ -28,25 +27,9 @@ def get_embedding(query, video):
 
     return text_emb, img_embs
 
-
-# # Encode an image:
-# url = "http://images.cocodataset.org/val2017/000000039769.jpg"
-# img = Image.fromarray(np.array(Image.open(requests.get(url, stream=True).raw))).convert('RGB')
-# img_emb = model.encode([img, img], device='cpu')
-
-# # Encode text descriptions
-# text_emb = model.encode(['Two dogs in the snow', 'Two cats laying on a sofa',
-#                          'A picture of London at night'], device='cpu')
-
-# # Compute cosine similarities
-# cos_scores = util.cos_sim(img_emb, text_emb)
-# print(cos_scores)
-
-
-def my_hook(d):
+def my_hook(d, model, desc, top_k, text):
     if d['status'] == 'finished':
-        print(d)
-        print('Done downloading, now extracting frames ...')
+        text.text("Processing video...")
         probe = ffmpeg.probe(d["filename"])
         video_stream = next((stream for stream in probe['streams'] if stream['codec_type'] == 'video'), None)
         width = int(video_stream['width'])
@@ -61,27 +44,47 @@ def my_hook(d):
             np
             .frombuffer(out, np.uint8)
             .reshape([-1, height, width, 3])
-        )[::10]
-
-        print(video.shape)
-        txt_embd, img_embds = get_embedding("two white puppies", video)
-        cos_scores = util.cos_sim(txt_embd, img_embds)
-        print(cos_scores)
-
-
-ydl_opts = {"format": "mp4", "progress_hooks": [my_hook], }
-with youtube_dl.YoutubeDL(ydl_opts) as ydl:
-    ydl.download(['https://youtu.be/I3AaW9ZevIU'])
-
-
-# # out, _ = (
-# #     ffmpeg
-# #     .input('in.mp4')
-# #     .output('pipe:', format='rawvideo', pix_fmt='rgb24')
-# #     .run(capture_stdout=True)
-# # )
-# # video = (
-# #     np
-# #     .frombuffer(out, np.uint8)
-# #     .reshape([-1, height, width, 3])
-# )
+        )[::10][:200]
+
+        txt_embd, img_embds = get_embedding(model, desc, video)
+        cos_scores = np.array(util.cos_sim(txt_embd, img_embds))
+        ids = np.argsort(cos_scores)[0][-top_k:]
+
+        imgs = [Image.fromarray(video[i]) for i in ids]
+        text.empty()
+        st.image(imgs)
+
+def run():
+    st.set_page_config(page_title="Youtube CLIFS")
+    # main body
+    model = get_model()
+    
+    st.sidebar.markdown("### Controls:")
+    top_k = st.sidebar.slider(
+        "Top K",
+        min_value=1,
+        max_value=5,
+        step=1,
+    )
+    desc = st.sidebar.text_input(
+        "Search Description",
+        value="Two white puppies",
+        help="Text description of what you want to find in the video",
+    )
+    url = st.sidebar.text_input(
+        "Youtube Video URL",
+        value='https://youtu.be/I3AaW9ZevIU',
+        help="Youtube video you'd like to search through",
+    )
+
+    submit_button = st.sidebar.button("Search")
+    if submit_button:
+        text = st.text("Downloading video...")
+        hook = lambda d: my_hook(d, model, desc, top_k, text)
+        ydl_opts = {"format": "mp4[height=360]", "progress_hooks": [hook], }
+        with youtube_dl.YoutubeDL(ydl_opts) as ydl:
+            ydl.download([url])
+
+
+if __name__ == "__main__":
+    run()

packages.txt

@@ -0,0 +1 @@
+ffmpeg