Update app.py

app.py

@@ -8,7 +8,7 @@ import torch
 from transformers import AutoTokenizer, CLIPTextModelWithProjection
 
 
-TITLE="""<h1 style="font-size: 42px;" align="center">Video Retrieval</h1>"""
+TITLE="""<h1 style="font-size: 64px;" align="center">Video Retrieval</h1>"""
 
 DESCRIPTION="""This is a video retrieval demo using [Diangle/clip4clip-webvid](https://huggingface.co/Diangle/clip4clip-webvid)."""
 IMAGE='<div style="text-align: right;"><img src="https://huggingface.co/spaces/Diangle/Clip4Clip-webvid/resolve/main/Searchium.png" width="333" height="216"/>'
@@ -23,7 +23,6 @@ ft_visual_features_database = np.load(ft_visual_features_file)
 database_csv_path = os.path.join(DATA_PATH, 'dataset_v1.csv')
 database_df = pd.read_csv(database_csv_path)
 
-
 class NearestNeighbors:
     """
     Class for NearestNeighbors.   
@@ -56,28 +55,33 @@ class NearestNeighbors:
             sim, idx = self.index.search(q_data, self.n_neighbors)        
         else:            
             if self.metric == 'binary':
-                print('binary search: ')
-                bq_data = np.packbits((q_data > 0.0).astype(bool), axis=1)            
-            print(bq_data.shape, self.index.d)
+                print('This is binary search.')
+                bq_data = np.packbits((q_data > 0.0).astype(bool), axis=1)
             sim, idx = self.index.search(bq_data, max(self.rerank_from, self.n_neighbors))
             
             if self.rerank_from > self.n_neighbors:
-                rerank_data = self.o_data[idx[0]]
-                rerank_search = NearestNeighbors(n_neighbors=self.n_neighbors, metric='cosine') 
-                rerank_search.fit(rerank_data)
-                sim, re_idxs = rerank_search.kneighbors(q_data)
-                idx = [idx[0][re_idxs[0]]]
-        
+                re_sims = np.zeros([len(q_data), self.n_neighbors], dtype=float)
+                re_idxs = np.zeros([len(q_data), self.n_neighbors], dtype=float)
+                for i, q in enumerate(q_data):
+                    rerank_data = self.o_data[idx[i]]
+                    rerank_search = NearestNeighbors(n_neighbors=self.n_neighbors, metric='cosine') 
+                    rerank_search.fit(rerank_data)
+                    re_sim, re_idx = rerank_search.kneighbors(np.asarray([q]))
+                    re_sims[i, :] = re_sim
+                    re_idxs[i, :] = idx[i][re_idx]
+                idx = re_idxs
+                sim = re_sims
+
         return sim, idx
     
+                    
 model = CLIPTextModelWithProjection.from_pretrained("Diangle/clip4clip-webvid")
-tokenizer = AutoTokenizer.from_pretrained("Diangle/clip4clip-webvid")
+tokenizer = CLIPTokenizer.from_pretrained("Diangle/clip4clip-webvid")
     
 def search(search_sentence):
     inputs = tokenizer(text=search_sentence , return_tensors="pt", padding=True)
 
-    outputs = model(input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"], return_dict=False)
-    # Customized projection layer    
+    outputs = model(input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"], return_dict=False)  
     text_projection = model.state_dict()['text_projection.weight']
     text_embeds = outputs[1] @ text_projection
     final_output = text_embeds[torch.arange(text_embeds.shape[0]), inputs["input_ids"].argmax(dim=-1)]
@@ -89,8 +93,15 @@ def search(search_sentence):
     
     nn_search = NearestNeighbors(n_neighbors=5, metric='binary', rerank_from=100)
     nn_search.fit(np.packbits((ft_visual_features_database > 0.0).astype(bool), axis=1), o_data=ft_visual_features_database)
-    sims, idxs = nn_search.kneighbors(sequence_output)  
-    return database_df.iloc[idxs[0]]['contentUrl'].to_list()
+    sims, idxs = nn_search.kneighbors(sequence_output)
+    # print(database_df.iloc[idxs[0]]['contentUrl'])
+    urls = database_df.iloc[idxs[0]]['contentUrl'].to_list()
+    AUTOPLAY_VIDEOS = []
+    for url in urls:
+        AUTOPLAY_VIDEOS.append("""<video controls muted autoplay>
+                    <source src={} type="video/mp4">
+                    </video>""".format(url))
+    return AUTOPLAY_VIDEOS
 
 
 with gr.Blocks() as demo:
@@ -102,12 +113,13 @@ with gr.Blocks() as demo:
         with gr.Column():
             inp = gr.Textbox(placeholder="Write a sentence.")
             btn = gr.Button(value="Retrieve")
-            ex = [["a woman waving to the camera"],["a basketball player performing a slam dunk"], ["how to bake a chocolate cake"], ["birds fly in the sky"]]
+            ex = [["mind-blowing magic tricks"],["baking chocolate cake"], 
+                ["birds fly in the sky"], ["natural wonders of the world"]]
             gr.Examples(examples=ex,
-                    inputs=[inp],
-                    )
+                        inputs=[inp]
+                        )
         with gr.Column():
-            out = [gr.Video(format='mp4') for _ in range(5)]
+            out = [gr.HTML() for _ in range(5)]
         btn.click(search, inputs=inp, outputs=out)
         
 demo.launch()