app.py

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)