import ffmpeg import torch import youtube_dl import numpy as np import streamlit as st from sentence_transformers import SentenceTransformer, util, models from clip import CLIPModel from PIL import Image @st.cache(allow_output_mutation=True, max_entries=1) def get_model(): txt_model = SentenceTransformer('clip-ViT-B-32-multilingual-v1').to(dtype=torch.float32, device=torch.device('cpu')) clip = CLIPModel() vis_model = SentenceTransformer(modules=[clip]).to(dtype=torch.float32, device=torch.device('cpu')) return txt_model, vis_model def get_embedding(txt_model, vis_model, query, video): text_emb = txt_model.encode(query, device='cpu') # Encode an image: images = [] for img in video: images.append(Image.fromarray(img)) img_embs = vis_model.encode(images, device='cpu') return text_emb, img_embs def find_frames(url, txt_model, vis_model, desc, seconds, top_k): text = st.text("Downloading video (Descargando video)...") probe = ffmpeg.probe(url) video_stream = next((stream for stream in probe['streams'] if stream['codec_type'] == 'video'), None) width = int(video_stream['width']) height = int(video_stream['height']) out, _ = ( ffmpeg .input(url, t=seconds) .output('pipe:', format='rawvideo', pix_fmt='rgb24') .run(capture_stdout=True) ) text.text("Processing video (Procesando video)...") video = ( np .frombuffer(out, np.uint8) .reshape([-1, height, width, 3]) )[::10] txt_embd, img_embds = get_embedding(txt_model, vis_model, desc, video) cos_scores = np.array(util.cos_sim(txt_embd, img_embds)) ids = np.argsort(cos_scores)[0][-top_k:] imgs = [Image.fromarray(video[i]) for i in ids] text.empty() st.image(imgs) with open("HOME.md", "r") as f: HOME_PAGE = f.read() with open("INICIO.md", "r") as f: INICIO_PAGINA = f.read() def main_page(txt_model, vis_model): st.title("Introducing Youtube CLIFS") st.markdown(HOME_PAGE) def inicio_pagina(txt_model, vis_model): st.title("Presentando Youtube CLIFS") st.markdown(INICIO_PAGINA) def clifs_page(txt_model, vis_model): st.title("CLIFS") st.sidebar.markdown("### Controls (Controles):") seconds = st.sidebar.slider( "How many seconds of video to consider? (¿Cuántos segundos de video considerar?)", min_value=10, max_value=120, value=60, step=1, ) top_k = st.sidebar.slider( "Top K", min_value=1, max_value=5, step=1, ) desc = st.sidebar.text_input( "Search Query (Consulta de Búsqueda)", value="Pancake in the shape of an otter", help="Text description of what you want to find in the video (Descripción de texto de que desea encontrar en el video)", ) url = st.sidebar.text_input( "Youtube Video URL (URL del Video de Youtube)", value='https://youtu.be/xUv6XgPwGaQ', help="Youtube video you want to search (Video de Youtube que desea búscar)", ) submit_button = st.sidebar.button("Search (Búscar)") if submit_button: ydl_opts = {"format": "mp4[height=360]"} with youtube_dl.YoutubeDL(ydl_opts) as ydl: info_dict = ydl.extract_info(url, download=False) video_url = info_dict.get("url", None) find_frames(video_url, txt_model, vis_model, desc, seconds, top_k) PAGES = { "Home": main_page, "Inicio": inicio_pagina, "CLIFS": clifs_page } def run(): st.set_page_config(page_title="Youtube CLIFS") # main body txt_model, vis_model = get_model() st.sidebar.title("Navigation (Navegación)") selection = st.sidebar.radio("Go to (Ir a)", list(PAGES.keys())) page = PAGES[selection](txt_model, vis_model) if __name__ == "__main__": run()