Update app

app.py

@@ -148,8 +148,8 @@ def extract_images_and_embeds(video_id, video_path, subtitles, output, expanded=
         if success:
             img_fname = f'{video_id}_{idx:06d}'
             img_fpath = os.path.join(output, 'frames', img_fname + '.jpg')
-            image = maintain_aspect_ratio_resize(image, height=350)     # save frame as JPEG file
-            cv2.imwrite( img_fpath, image)     # save frame as JPEG file
+            # image = maintain_aspect_ratio_resize(image, height=350)     # save frame as JPEG file
+            # cv2.imwrite( img_fpath, image)     # save frame as JPEG file
 		    
             count += 1
             anno.append({
@@ -181,6 +181,7 @@ def extract_images_and_embeds(video_id, video_path, subtitles, output, expanded=
                     'text': batch_list[i]['text'],
                     'image_filepath': batch_list[i]['image_filepath'],
                     'start_time': batch_list[i]['start_time'],
+                    'frame_no': frame_no,
                 })
             batch_list = []
 
@@ -195,6 +196,7 @@ def extract_images_and_embeds(video_id, video_path, subtitles, output, expanded=
                 'text': batch_list[i]['text'],
                 'image_filepath': batch_list[i]['image_filepath'],
                 'start_time': batch_list[i]['start_time'],
+                'frame_no': frame_no,
             })
 
     with open(os.path.join(output, 'annotations.json'), 'w') as fh:
@@ -203,7 +205,9 @@ def extract_images_and_embeds(video_id, video_path, subtitles, output, expanded=
     with open(os.path.join(output, 'embeddings.pkl'), 'wb') as fh:
         pickle.dump(embeddings, fh)
 
-def run_query(video_id, text_query, path='/tmp'):
+def run_query(video_path, text_query, path='/tmp'):
+    
+    vidcap = cv2.VideoCapture(video_path)
     
     embeddings_filepath = os.path.join(path, 'embeddings.pkl')
     faiss_filepath = os.path.join(path, 'faiss_index.pkl')
@@ -228,7 +232,14 @@ def run_query(video_id, text_query, path='/tmp'):
     print('Running FAISS search')
     _, I = faiss_index.search(emb_query, 6) 
 
-    clip_images = [embeddings[idx]['image_filepath'] for idx in I[0]]
+    clip_images = []
+    for idx in I[0]:
+        frame_no = embeddings[idx]['frame_no']
+        vidcap.set(1, frame_no)    # added this line 
+        success, image = vidcap.read()
+        clip_images.append(image)
+
+    # clip_images = [embeddings[idx]['image_filepath'] for idx in I[0]]
     transcripts = [f"({embeddings[idx]['start_time']}) {embeddings[idx]['text']}" for idx in I[0]]
     return clip_images, transcripts
 
@@ -253,12 +264,12 @@ def get_video_id_from_url(video_url):
             return url.path.split('/')[2]
         if url.path[:3] == '/v/':
             return url.path.split('/')[2]
-            
+
     return None
 
 
 def process(video_url, text_query):
-    tmp_dir = os.path.join(os.getcwd(), 'cache')
+    tmp_dir = os.environ.get('TMPDIR', '/tmp')
     video_id = get_video_id_from_url(video_url)
     output_dir = os.path.join(tmp_dir, video_id)
     video_file = download_video(video_url, path=output_dir)
@@ -303,4 +314,4 @@ with gr.Blocks() as demo:
         outputs=[video_player, gallery],
     )
 
-demo.launch()
+demo.launch(share=True)