app.py

import torch
import clip
import cv2, yt_dlp
from PIL import Image,ImageDraw, ImageFont
import wget
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, ydl_opts={}, format_note='240p', ext='mp4', max_size = 500000000):
    defaults = ['480p', '360p','240p','144p']
    ydl_opts = ydl_opts
    ydl = yt_dlp.YoutubeDL(ydl_opts)
    info_dict = ydl.extract_info(url, download=False)
    formats = info_dict.get('formats', None)
    # filter out formats we can't process
    formats = [f for f in formats if f['ext'] == ext 
               and f['vcodec'].split('.')[0] != 'av01' 
               and f['filesize'] is not None and f['filesize'] <= max_size]
    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]
    
    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)
  
def download_video(url):
    # create "videos" foder for saved videos
    path_videos = Path('videos')
    try:
      path_videos.mkdir(parents=True)
    except FileExistsError:
      pass
    # clear the "videos" folder 
    videos_to_keep = ['LPGUtKKO6F8', 'wJCXBGPo5c8','RicOR0oHATY']
    if len(list(path_videos.glob('*'))) > 10:
        for path_video in path_videos.glob('*'):
            if path_video.stem not in set(videos_to_keep):
                path_video.unlink()
                print(f'removed video {path_video}')
    # select format to download for given video
    # by default select 240p and .mp4 
    try:
      format, format_id, fps = select_video_format(url)
      ydl_opts = {
        'format':format_id,
        'outtmpl': "videos/%(id)s.%(ext)s"}

      with yt_dlp.YoutubeDL(ydl_opts) as ydl:
        try:
          ydl.cache.remove()
          meta = ydl.extract_info(url)
          title = meta['title']
          thumb_url = meta['thumbnail']
          thumb_location = f"videos/{meta['id']}.{thumb_url.split('.')[-1]}"
          wget.download(thumb_url, out=thumb_location)
          save_location = 'videos/' + meta['id'] + '.' + meta['ext']
        except yt_dlp.DownloadError as error:
          print(f'error with download_video function: {error}')
          save_location = None
          thumb_location = None
          title = None
    except IndexError as err:
      print(f"can't find suitable video formats. we are not able to process video larger than 95 Mib at the moment")
      fps, save_location, thumb_location, title = None, None, None, None
    return(fps, save_location, thumb_location, title)

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)
        count += 1
    cap.release()


def vid2frames(url, sampling_interval=1):
  # create folder for extracted frames - if folder exists, delete and create a new one
    path_frames = Path('frames')
    try:
        path_frames.mkdir(parents=True)
    except FileExistsError:
        shutil.rmtree(path_frames)
        path_frames.mkdir(parents=True)
 
    # download the video 
    fps, video, thumb, title = download_video(url)
    if video is not None: 
      if fps is None: fps = 30
      skip_frames = int(fps * 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, path_frames, n_workers), range(n_workers))
    else:
      skip_frames, path_frames = None, None
    return(title, thumb, skip_frames, path_frames)


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, bs=526):
    title, thumb, skip_frames, path_frames= vid2frames(url,sampling_interval)
    print(f"searching {title}")
    if path_frames is not None:
      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),dtype=torch.float32).to(device)
      print(f"encoding images, 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
      print(f'encoding search query')
      thumb = Image.open(thumb).convert("RGB")
      batch_thumb = torch.tensor(np.stack([preprocess(thumb)])).to(device)
      with torch.no_grad():
          thumb_features = model.encode_image(batch_thumb).to(dtype=torch.float32)
          thumb_features /= thumb_features.norm(dim=-1, keepdim=True)
    
      similarity = (100.0 * image_features @ thumb_features.T)
      values, indices = similarity.topk(4, dim=0)
      print(f"indices for best matches{indices}")
      print(f"filenames for best matches {[filenames[i]for i in indices]}")
      best_frames = [Image.open(filenames[ind]).convert("RGB") for ind in indices]
      times = [f'{datetime.timedelta(seconds = round(ind[0].item() * sampling_interval,2))}' for ind in indices]
      image_output = captioned_strip(best_frames,title, times,2)
      print('task complete')
    else:
      title = "not able to download video"
      thumb = 0
      image_output = None

    return(title, thumb, image_output)

inputs = [gr.Textbox(label="Give us the link to your youtube video! (maximum size 50 MB)"),
          gr.Number(1,label='sampling interval (seconds)')]
outputs = [
    gr.HTML(label=""),  # To be used as title
    gr.Image(label="Original thumbnail"),
    gr.Image(label="Top matches in the video"),
]

article = "Adapted from [It Happened One Frame](https://huggingface.co/spaces/YiYiXu/it-happened-one-frame-2)."

gr.Interface(
    run_inference,
    inputs=inputs,
    outputs=outputs,
    title="Locate YouTube Thumbnail In Video",
    description='A CLIP-based app that find the YouTube video frame that matches its thumbnail.',
    article = article,
    examples=[
        ['https://youtu.be/LPGUtKKO6F8', 1],
        ['https://youtu.be/wJCXBGPo5c8', 1],
        ['https://youtu.be/RicOR0oHATY', 1]
    ]
).launch()