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)