chuanenlin/which-frame
6
1import streamlit as st2import cv23from PIL import Image4import clip as openai_clip5import torch6import math7from humanfriendly import format_timespan8import numpy as np9import time10import os11import yt_dlp12import io13 14EXAMPLE_URL = "https://www.youtube.com/watch?v=zTvJJnoWIPk"15CACHED_DATA_PATH = "cached_data/"16 17device = "cuda" if torch.cuda.is_available() else "cpu"18model, preprocess = openai_clip.load("ViT-B/32", device=device)19 20def fetch_video(url):21 if url != EXAMPLE_URL:22 st.error("Only the example video is supported due to compute constraints.")23 st.stop()24 25 try:26 ydl_opts = {27 'format': 'bestvideo[height<=360][ext=mp4]/best[height<=360]',28 'quiet': True,29 'no_warnings': True,30 'extract_flat': False,31 'no_check_certificates': True32 }33 with yt_dlp.YoutubeDL(ydl_opts) as ydl:34 info = ydl.extract_info(url, download=False)35 video_url = info['url']36 return None, video_url37 38 except Exception as e:39 st.error(f"Error fetching video: {str(e)}")40 st.stop()41 42def extract_frames(video, status_text, progress_bar):43 cap = cv2.VideoCapture(video)44 frames = []45 fps = cap.get(cv2.CAP_PROP_FPS)46 frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))47 step = max(1, round(fps/2))48 total_frames = frame_count // step49 frame_indices = []50 for i in range(0, frame_count, step):51 cap.set(cv2.CAP_PROP_POS_FRAMES, i)52 ret, frame = cap.read()53 if ret:54 frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)55 frames.append(Image.fromarray(frame_rgb))56 frame_indices.append(i)57 58 current_frame = len(frames)59 status_text.text(f'Extracting frames... ({min(current_frame, total_frames)}/{total_frames})')60 progress = min(current_frame / total_frames, 1.0)61 progress_bar.progress(progress)62 63 cap.release()64 return frames, fps, frame_indices65 66def encode_frames(video_frames, status_text):67 batch_size = 25668 batches = math.ceil(len(video_frames) / batch_size)69 video_features = torch.empty([0, 512], dtype=torch.float32).to(device)70 71 for i in range(batches):72 batch_frames = video_frames[i*batch_size : (i+1)*batch_size]73 batch_preprocessed = torch.stack([preprocess(frame) for frame in batch_frames]).to(device)74 with torch.no_grad():75 batch_features = model.encode_image(batch_preprocessed)76 batch_features = batch_features.float()77 batch_features /= batch_features.norm(dim=-1, keepdim=True)78 video_features = torch.cat((video_features, batch_features))79 status_text.text(f'Encoding frames... ({(i+1)*batch_size}/{len(video_frames)})')80 81 return video_features82 83def img_to_bytes(img):84 img_byte_arr = io.BytesIO()85 img.save(img_byte_arr, format='JPEG')86 img_byte_arr = img_byte_arr.getvalue()87 return img_byte_arr88 89def get_youtube_timestamp_url(url, frame_idx, frame_indices):90 frame_count = frame_indices[frame_idx]91 fps = st.session_state.fps92 seconds = frame_count / fps93 seconds_rounded = int(seconds)94 95 if url == EXAMPLE_URL:96 video_id = "zTvJJnoWIPk"97 else:98 try:99 from urllib.parse import urlparse, parse_qs100 parsed_url = urlparse(url)101 video_id = parse_qs(parsed_url.query)['v'][0]102 except:103 return None, None104 105 return f"https://youtu.be/{video_id}?t={seconds_rounded}", seconds106 107def display_results(best_photo_idx, video_frames):108 st.subheader("Top 10 Results")109 for frame_id in best_photo_idx:110 result = video_frames[frame_id]111 st.image(result, width=400)112 113 timestamp_url, seconds = get_youtube_timestamp_url(st.session_state.url, frame_id, st.session_state.frame_indices)114 if timestamp_url:115 st.markdown(f"[▶️ Play video at {format_timespan(int(seconds))}]({timestamp_url})")116 117def text_search(search_query, video_features, video_frames, display_results_count=10):118 display_results_count = min(display_results_count, len(video_frames))119 120 with torch.no_grad():121 text_tokens = openai_clip.tokenize(search_query).to(device)122 text_features = model.encode_text(text_tokens)123 text_features = text_features.float()124 text_features /= text_features.norm(dim=-1, keepdim=True)125 126 video_features = video_features.float()127 128 similarities = (100.0 * video_features @ text_features.T)129 values, best_photo_idx = similarities.topk(display_results_count, dim=0)130 display_results(best_photo_idx, video_frames)131 132def image_search(query_image, video_features, video_frames, display_results_count=10):133 query_image = preprocess(query_image).unsqueeze(0).to(device)134 135 with torch.no_grad():136 image_features = model.encode_image(query_image)137 image_features = image_features.float()138 image_features /= image_features.norm(dim=-1, keepdim=True)139 140 video_features = video_features.float()141 142 similarities = (100.0 * video_features @ image_features.T)143 values, best_photo_idx = similarities.topk(display_results_count, dim=0)144 display_results(best_photo_idx, video_frames)145 146def text_and_image_search(search_query, query_image, video_features, video_frames, display_results_count=10):147 with torch.no_grad():148 text_tokens = openai_clip.tokenize(search_query).to(device)149 text_features = model.encode_text(text_tokens)150 text_features = text_features.float()151 text_features /= text_features.norm(dim=-1, keepdim=True)152 153 query_image = preprocess(query_image).unsqueeze(0).to(device)154 with torch.no_grad():155 image_features = model.encode_image(query_image)156 image_features = image_features.float()157 image_features /= image_features.norm(dim=-1, keepdim=True)158 159 combined_features = (text_features + image_features) / 2160 161 video_features = video_features.float()162 similarities = (100.0 * video_features @ combined_features.T)163 values, best_photo_idx = similarities.topk(display_results_count, dim=0)164 display_results(best_photo_idx, video_frames)165 166def load_cached_data(url):167 if url == EXAMPLE_URL:168 try:169 video_frames = np.load(f"{CACHED_DATA_PATH}example_frames.npy", allow_pickle=True)170 video_features = torch.load(f"{CACHED_DATA_PATH}example_features.pt")171 fps = np.load(f"{CACHED_DATA_PATH}example_fps.npy")172 frame_indices = np.load(f"{CACHED_DATA_PATH}example_frame_indices.npy")173 return video_frames, video_features, fps, frame_indices174 except:175 return None, None, None, None176 return None, None, None, None177 178def save_cached_data(url, video_frames, video_features, fps, frame_indices):179 if url == EXAMPLE_URL:180 os.makedirs(CACHED_DATA_PATH, exist_ok=True)181 np.save(f"{CACHED_DATA_PATH}example_frames.npy", video_frames)182 torch.save(video_features, f"{CACHED_DATA_PATH}example_features.pt")183 np.save(f"{CACHED_DATA_PATH}example_fps.npy", fps)184 np.save(f"{CACHED_DATA_PATH}example_frame_indices.npy", frame_indices)185 186def clear_cached_data():187 if os.path.exists(CACHED_DATA_PATH):188 try:189 for file in os.listdir(CACHED_DATA_PATH):190 file_path = os.path.join(CACHED_DATA_PATH, file)191 if os.path.isfile(file_path):192 os.unlink(file_path)193 os.rmdir(CACHED_DATA_PATH)194 except Exception as e:195 print(f"Error clearing cache: {e}")196 197st.set_page_config(page_title="Which Frame? 🎞️🔍", page_icon = "🔍", layout = "centered", initial_sidebar_state = "collapsed")198 199hide_streamlit_style = """200<style>201/* Hide Streamlit elements */202#MainMenu {visibility: hidden;}203footer {visibility: hidden;}204* {205 font-family: Avenir;206}207.block-container {208 max-width: 800px;209 padding: 2rem 1rem;210}211.stTextInput input {212 border-radius: 8px;213 border: 1px solid #E0E0E0;214 padding: 0.75rem;215 font-size: 1rem;216}217.stRadio [role="radiogroup"] {218 background: #F8F8F8;219 padding: 1rem;220 border-radius: 12px;221}222h1 {text-align: center;}223.css-gma2qf {display: flex; justify-content: center; font-size: 36px; font-weight: bold;}224a:link {text-decoration: none;}225a:hover {text-decoration: none;}226.st-ba {font-family: Avenir;}227.st-button {text-align: center;}228</style>229"""230st.markdown(hide_streamlit_style, unsafe_allow_html=True)231 232# Initialize session state233if 'initialized' not in st.session_state:234 st.session_state.initialized = False235 st.session_state.video_frames = None236 st.session_state.video_features = None237 st.session_state.fps = None238 st.session_state.frame_indices = None239 st.session_state.url = EXAMPLE_URL240 241# Load data on first run242if not st.session_state.initialized:243 cached_frames, cached_features, cached_fps, cached_frame_indices = load_cached_data(EXAMPLE_URL)244 245 if cached_frames is not None:246 st.session_state.video_frames = cached_frames247 st.session_state.video_features = cached_features248 st.session_state.fps = cached_fps249 st.session_state.frame_indices = cached_frame_indices250 st.session_state.initialized = True251 else:252 st.error("Could not load video data. Please contact the administrator.")253 st.stop()254 255st.title("Which Frame? 🎞️🔍")256st.markdown("""257Search a video semantically. For example, which frame has "a person with sunglasses"?258Search using text, images, or a mix of text + image. WhichFrame uses [CLIP](https://github.com/openai/CLIP) for zero-shot frame classification.259""")260 261st.video(EXAMPLE_URL)262st.caption("Note: Try out the code linked at the bottom of the page to run WhichFrame on your own videos.")263 264if st.session_state.initialized:265 search_type = st.radio("Search Method", ["Text Search", "Image Search", "Text + Image Search"], index=0)266 267 if search_type == "Text Search": # Text Search268 text_query = st.text_input("Type a search query (e.g., 'red car' or 'person with sunglasses')")269 if st.button("Search"):270 if not text_query:271 st.error("Please enter a search query first")272 else:273 text_search(text_query, st.session_state.video_features, st.session_state.video_frames)274 elif search_type == "Image Search": # Image Search275 uploaded_file = st.file_uploader("Upload a query image", type=['png', 'jpg', 'jpeg'])276 if uploaded_file is not None:277 query_image = Image.open(uploaded_file).convert('RGB')278 st.image(query_image, caption="Query Image", width=200)279 if st.button("Search"):280 if uploaded_file is None:281 st.error("Please upload an image first")282 else:283 image_search(query_image, st.session_state.video_features, st.session_state.video_frames)284 else: # Text + Image Search285 text_query = st.text_input("Type a search query")286 uploaded_file = st.file_uploader("Upload a query image", type=['png', 'jpg', 'jpeg'])287 if uploaded_file is not None:288 query_image = Image.open(uploaded_file).convert('RGB')289 st.image(query_image, caption="Query Image", width=200)290 291 if st.button("Search"):292 if not text_query or uploaded_file is None:293 st.error("Please provide both text query and image")294 else:295 text_and_image_search(text_query, query_image, st.session_state.video_features, st.session_state.video_frames)296 297st.markdown("---")298st.markdown(299 "By [David Chuan-En Lin](https://chuanenlin.com/). "300 "Play with the code at [https://github.com/chuanenlin/whichframe](https://github.com/chuanenlin/whichframe)."301 "v2 code with better interface and model at [https://github.com/chuanenlin/whichframe-v2](https://github.com/chuanenlin/whichframe-v2)."302)