Spaces:
Runtime error
Runtime error
| import torch | |
| import clip | |
| import cv2, youtube_dl | |
| from PIL import Image,ImageDraw, ImageFont | |
| import os | |
| from functools import partial | |
| from multiprocessing.pool import Pool | |
| import shutil | |
| from pathlib import Path | |
| import numpy as np | |
| import datetime | |
| import gradio as gr | |
| # load model and preprocess | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| model, preprocess = clip.load("ViT-B/32") | |
| def select_video_format(url, format_note='480p', ext='mp4'): | |
| defaults = ['480p', '360p','240p','144p'] | |
| ydl_opts = {} | |
| ydl = youtube_dl.YoutubeDL(ydl_opts) | |
| info_dict = ydl.extract_info(url, download=False) | |
| formats = info_dict.get('formats', None) | |
| available_format_notes = set([f['format_note'] for f in formats]) | |
| if format_note not in available_format_notes: | |
| format_note = [d for d in defaults if d in available_format_notes][0] | |
| formats = [f for f in formats if f['format_note'] == format_note and f['ext'] == ext and f['vcodec'].split('.')[0] != 'av01'] | |
| format = formats[0] | |
| format_id = format.get('format_id', None) | |
| fps = format.get('fps', None) | |
| print(f'format selected: {format}') | |
| return(format, format_id, fps) | |
| # to-do: delete saved videos | |
| # testing aria2c | |
| def download_video(url,format_id, n_keep=10): | |
| ydl_opts = { | |
| 'format':format_id, | |
| 'external_downloader' : 'aria2c', | |
| 'external_downloader_args' :['--max-connection-per-server=16','--dir=videos'], | |
| 'outtmpl': "videos/%(id)s.%(ext)s"} | |
| # create a directory for saved videos | |
| video_path = Path('videos') | |
| try: | |
| video_path.mkdir(parents=True) | |
| except FileExistsError: | |
| pass | |
| with youtube_dl.YoutubeDL(ydl_opts) as ydl: | |
| try: | |
| ydl.cache.remove() | |
| meta = ydl.extract_info(url) | |
| save_location = 'videos/' + meta['id'] + '.' + meta['ext'] | |
| except youtube_dl.DownloadError as error: | |
| print(f'error with download_video function: {error}') | |
| return(save_location) | |
| def process_video_parallel(video, skip_frames, dest_path, num_processes, process_number): | |
| cap = cv2.VideoCapture(video) | |
| frames_per_process = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) // (num_processes) | |
| count = frames_per_process * process_number | |
| cap.set(cv2.CAP_PROP_POS_FRAMES, count) | |
| print(f"worker: {process_number}, process frames {count} ~ {frames_per_process * (process_number + 1)} \n total number of frames: {cap.get(cv2.CAP_PROP_FRAME_COUNT)} \n video: {video}; isOpen? : {cap.isOpened()}") | |
| while count < frames_per_process * (process_number + 1) : | |
| ret, frame = cap.read() | |
| if not ret: | |
| break | |
| if count % skip_frames ==0: | |
| filename =f"{dest_path}/{count}.jpg" | |
| cv2.imwrite(filename, frame) | |
| #print(f"saved {filename}") | |
| count += 1 | |
| cap.release() | |
| def vid2frames(url, sampling_interval=1, ext='mp4'): | |
| # create folder for extracted frames - if folder exists, delete and create a new one | |
| dest_path = Path('frames') | |
| try: | |
| dest_path.mkdir(parents=True) | |
| except FileExistsError: | |
| shutil.rmtree(dest_path) | |
| dest_path.mkdir(parents=True) | |
| # figure out the format for download, | |
| # by default select 480p and .mp4 | |
| format, format_id, fps = select_video_format(url, format_note='480p', ext='mp4') | |
| # download the video | |
| video = download_video(url,format_id) | |
| # calculate skip_frames | |
| try: | |
| skip_frames = int(fps * sampling_interval) | |
| except: | |
| skip_frames = int(30 * sampling_interval) | |
| print(f'video saved at: {video}, fps:{fps}, skip_frames: {skip_frames}') | |
| # extract video frames at given sampling interval with multiprocessing - | |
| n_workers = min(os.cpu_count(), 12) | |
| print(f'now extracting frames with {n_workers} process...') | |
| with Pool(n_workers) as pool: | |
| pool.map(partial(process_video_parallel, video, skip_frames, dest_path, n_workers), range(n_workers)) | |
| return(skip_frames, dest_path) | |
| def captioned_strip(images, caption=None, times=None, rows=1): | |
| increased_h = 0 if caption is None else 30 | |
| w, h = images[0].size[0], images[0].size[1] | |
| img = Image.new("RGB", (len(images) * w // rows, h * rows + increased_h)) | |
| for i, img_ in enumerate(images): | |
| img.paste(img_, (i // rows * w, increased_h + (i % rows) * h)) | |
| if caption is not None: | |
| draw = ImageDraw.Draw(img) | |
| font = ImageFont.truetype( | |
| "/usr/share/fonts/truetype/liberation2/LiberationMono-Bold.ttf", 16 | |
| ) | |
| font_small = ImageFont.truetype("/usr/share/fonts/truetype/liberation2/LiberationMono-Bold.ttf", 12) | |
| draw.text((60, 3), caption, (255, 255, 255), font=font) | |
| for i,ts in enumerate(times): | |
| draw.text(( | |
| (i % rows) * w + 40 , #column poistion | |
| i // rows * h + 33) # row position | |
| , ts, | |
| (255, 255, 255), font=font_small) | |
| return img | |
| def run_inference(url, sampling_interval, search_query, bs=526): | |
| skip_frames, path_frames= vid2frames(url,sampling_interval) | |
| filenames = sorted(path_frames.glob('*.jpg'),key=lambda p: int(p.stem)) | |
| n_frames = len(filenames) | |
| bs = min(n_frames,bs) | |
| print(f"extracted {n_frames} frames, now encoding images") | |
| # encoding images one batch at a time, combine all batch outputs -> image_features, size n_frames x 512 | |
| image_features = torch.empty(size=(n_frames, 512)).to(device) | |
| print(f"batch size :{bs} ; number of batches: {len(range(0, n_frames,bs))}") | |
| for b in range(0, n_frames,bs): | |
| images = [] | |
| # loop through all frames in the batch -> create batch_image_input, size bs x 3 x 224 x 224 | |
| for filename in filenames[b:b+bs]: | |
| image = Image.open(filename).convert("RGB") | |
| images.append(preprocess(image)) | |
| batch_image_input = torch.tensor(np.stack(images)).to(device) | |
| # encoding batch_image_input -> batch_image_features | |
| with torch.no_grad(): | |
| batch_image_features = model.encode_image(batch_image_input) | |
| batch_image_features /= batch_image_features.norm(dim=-1, keepdim=True) | |
| # add encoded image embedding to image_features | |
| image_features[b:b+bs] = batch_image_features | |
| # encoding search query | |
| with torch.no_grad(): | |
| text_features = model.encode_text(clip.tokenize(search_query).to(device)) | |
| text_features /= text_features.norm(dim=-1, keepdim=True) | |
| print(image_features.dtype, text_features.dtype) | |
| similarity = (100.0 * image_features @ text_features.T) | |
| values, indices = similarity.topk(4, dim=0) | |
| best_frames = [Image.open(filenames[ind]).convert("RGB") for ind in indices] | |
| times = [f'{datetime.timedelta(seconds = ind[0].item() * sampling_interval)}' for ind in indices] | |
| image_output = captioned_strip(best_frames,search_query, times,2) | |
| title = search_query | |
| return(title, image_output) | |
| inputs = [gr.inputs.Textbox(label="Give us the link to your youtube video!"), | |
| gr.Number(5,label='sampling interval (seconds)'), | |
| gr.inputs.Textbox(label="What do you want to search?")] | |
| outputs = [ | |
| gr.outputs.HTML(label=""), # To be used as title | |
| gr.outputs.Image(label=""), | |
| ] | |
| gr.Interface( | |
| run_inference, | |
| inputs=inputs, | |
| outputs=outputs, | |
| title="It Happened One Frame", | |
| description='A CLIP-based app that search video frame based on text', | |
| examples=[ | |
| ['https://youtu.be/v1rkzUIL8oc', 1, "James Cagney dancing down the stairs"], | |
| ['https://youtu.be/k4R5wZs8cxI', 1, "James Cagney smashes a grapefruit into Mae Clarke's face"], | |
| ['https://youtu.be/0diCvgWv_ng', 1, "little Deborah practicing her ballet while wearing a tutu in empty restaurant"] | |
| ] | |
| ).launch(debug=True,enable_queue=True) | |