ncoop57 commited on
Commit
81abbd1
1 Parent(s): 2924167

Add generating GIF from returned frames

Browse files
Files changed (1) hide show
  1. app.py +13 -3
app.py CHANGED
@@ -1,5 +1,7 @@
1
  import ffmpeg
 
2
  import torch
 
3
  import youtube_dl
4
 
5
  import numpy as np
@@ -53,10 +55,18 @@ def find_frames(url, txt_model, vis_model, desc, seconds, top_k):
53
  txt_embd, img_embds = get_embedding(txt_model, vis_model, desc, video)
54
  cos_scores = np.array(util.cos_sim(txt_embd, img_embds))
55
  ids = np.argsort(cos_scores)[0][-top_k:]
56
- imgs = [Image.fromarray(video[i]) for i in ids]
 
 
 
 
 
57
 
58
  gif_runner.empty()
59
  text.empty()
 
 
 
60
  st.image(imgs)
61
 
62
  with open("HOME.md", "r") as f:
@@ -89,8 +99,8 @@ def clifs_page(txt_model, vis_model):
89
  top_k = st.sidebar.slider(
90
  "Top K",
91
  min_value=1,
92
- max_value=5,
93
- value=3,
94
  step=1,
95
  )
96
  desc = st.sidebar.text_input(
 
1
  import ffmpeg
2
+ import os
3
  import torch
4
+ import uuid
5
  import youtube_dl
6
 
7
  import numpy as np
 
55
  txt_embd, img_embds = get_embedding(txt_model, vis_model, desc, video)
56
  cos_scores = np.array(util.cos_sim(txt_embd, img_embds))
57
  ids = np.argsort(cos_scores)[0][-top_k:]
58
+ imgs = [Image.fromarray(video[i]) for i in sorted(ids)]
59
+
60
+ # from: https://stackoverflow.com/a/57751793/5768407
61
+ fname = uuid.uuid4().hex
62
+ imgs[0].save(fp=f"./{fname}.gif", format='GIF', append_images=imgs[1:],
63
+ save_all=True, duration=200, loop=0)
64
 
65
  gif_runner.empty()
66
  text.empty()
67
+ st.image(f"./{fname}.gif")
68
+ # remove the gif from file so we don't build up a bunch of files
69
+ os.remove(f"./{fname}.gif")
70
  st.image(imgs)
71
 
72
  with open("HOME.md", "r") as f:
 
99
  top_k = st.sidebar.slider(
100
  "Top K",
101
  min_value=1,
102
+ max_value=20,
103
+ value=10,
104
  step=1,
105
  )
106
  desc = st.sidebar.text_input(