app.py

import csv
import os
import tempfile

import gradio as gr
import requests
import torch
import torchvision
import torchvision.transforms as T
from PIL import Image
from featup.util import norm
from torchaudio.functional import resample

from denseav.train import LitAVAligner
from denseav.plotting import plot_attention_video, plot_2head_attention_video, plot_feature_video
from denseav.shared import norm, crop_to_divisor, blur_dim
from os.path import join


mode = "hf"

if mode == "local":
    sample_videos_dir = "samples"
else:
    os.environ['TORCH_HOME'] = '/tmp/.cache'
    os.environ['HF_HOME'] = '/tmp/.cache'
    os.environ['HF_DATASETS_CACHE'] = '/tmp/.cache'
    os.environ['TRANSFORMERS_CACHE'] = '/tmp/.cache'
    os.environ['GRADIO_EXAMPLES_CACHE'] = '/tmp/gradio_cache'
    sample_videos_dir = "/tmp/samples"


def download_video(url, save_path):
    response = requests.get(url)
    with open(save_path, 'wb') as file:
        file.write(response.content)


base_url = "https://marhamilresearch4.blob.core.windows.net/denseav-public/samples/"
sample_videos_urls = {
    "puppies.mp4": base_url + "puppies.mp4",
    "peppers.mp4": base_url + "peppers.mp4",
    "boat.mp4": base_url + "boat.mp4",
    "elephant2.mp4": base_url + "elephant2.mp4",

}

# Ensure the directory for sample videos exists
os.makedirs(sample_videos_dir, exist_ok=True)

# Download each sample video
for filename, url in sample_videos_urls.items():
    save_path = os.path.join(sample_videos_dir, filename)
    # Download the video if it doesn't already exist
    if not os.path.exists(save_path):
        print(f"Downloading {filename}...")
        download_video(url, save_path)
    else:
        print(f"{filename} already exists. Skipping download.")

csv.field_size_limit(100000000)
options = ['language', "sound-language", "sound"]
load_size = 224
plot_size = 224

video_input = gr.Video(label="Choose a video to featurize", height=480)
model_option = gr.Radio(options, value="language", label='Choose a model')

video_output1 = gr.Video(label="Audio Video Attention", height=480)
video_output2 = gr.Video(label="Multi-Head Audio Video Attention (Only Availible for sound_and_language)",
                         height=480)
video_output3 = gr.Video(label="Visual Features", height=480)

models = {o: LitAVAligner.from_pretrained(f"mhamilton723/DenseAV-{o}") for o in options}


def process_video(video, model_option):
    # model = models[model_option].cuda()
    model = models[model_option]

    original_frames, audio, info = torchvision.io.read_video(video, end_pts=10, pts_unit='sec')
    sample_rate = 16000

    if info["audio_fps"] != sample_rate:
        audio = resample(audio, info["audio_fps"], sample_rate)
    audio = audio[0].unsqueeze(0)

    img_transform = T.Compose([
        T.Resize(load_size, Image.BILINEAR),
        lambda x: crop_to_divisor(x, 8),
        lambda x: x.to(torch.float32) / 255,
        norm])

    frames = torch.cat([img_transform(f.permute(2, 0, 1)).unsqueeze(0) for f in original_frames], axis=0)

    plotting_img_transform = T.Compose([
        T.Resize(plot_size, Image.BILINEAR),
        lambda x: crop_to_divisor(x, 8),
        lambda x: x.to(torch.float32) / 255])

    frames_to_plot = plotting_img_transform(original_frames.permute(0, 3, 1, 2))

    with torch.no_grad():
        # audio_feats = model.forward_audio({"audio": audio.cuda()})
        audio_feats = model.forward_audio({"audio": audio})
        audio_feats = {k: v.cpu() for k, v in audio_feats.items()}
        # image_feats = model.forward_image({"frames": frames.unsqueeze(0).cuda()}, max_batch_size=2)
        image_feats = model.forward_image({"frames": frames.unsqueeze(0)}, max_batch_size=2)
        image_feats = {k: v.cpu() for k, v in image_feats.items()}

        sim_by_head = model.sim_agg.get_pairwise_sims(
            {**image_feats, **audio_feats},
            raw=False,
            agg_sim=False,
            agg_heads=False
        ).mean(dim=-2).cpu()

        sim_by_head = blur_dim(sim_by_head, window=3, dim=-1)
        print(sim_by_head.shape)

    temp_video_path_1 = tempfile.mktemp(suffix='.mp4')

    plot_attention_video(
        sim_by_head,
        frames_to_plot,
        audio,
        info["video_fps"],
        sample_rate,
        temp_video_path_1)

    if model_option == "sound_and_language":
        temp_video_path_2 = tempfile.mktemp(suffix='.mp4')

        plot_2head_attention_video(
            sim_by_head,
            frames_to_plot,
            audio,
            info["video_fps"],
            sample_rate,
            temp_video_path_2)

    else:
        temp_video_path_2 = None

    temp_video_path_3 = tempfile.mktemp(suffix='.mp4')
    temp_video_path_4 = tempfile.mktemp(suffix='.mp4')

    plot_feature_video(
        image_feats["image_feats"].cpu(),
        audio_feats['audio_feats'].cpu(),
        frames_to_plot,
        audio,
        info["video_fps"],
        sample_rate,
        temp_video_path_3,
        temp_video_path_4,
    )
    # return temp_video_path_1, temp_video_path_2, temp_video_path_3, temp_video_path_4

    return temp_video_path_1, temp_video_path_2, temp_video_path_3


with gr.Blocks() as demo:
    with gr.Column():
        gr.Markdown("## Visualizing Sound and Language with DenseAV")
        gr.Markdown(
            "This demo allows you to explore the inner attention maps of DenseAV's dense multi-head contrastive operator.")
        with gr.Row():
            with gr.Column(scale=1):
                model_option.render()
            with gr.Column(scale=3):
                video_input.render()
        with gr.Row():
            submit_button = gr.Button("Submit")
        with gr.Row():
            gr.Examples(
                examples=[
                    [join(sample_videos_dir, "puppies.mp4"), "sound_and_language"],
                    [join(sample_videos_dir, "peppers.mp4"), "language"],
                    [join(sample_videos_dir, "elephant2.mp4"), "language"],
                    [join(sample_videos_dir, "boat.mp4"), "language"]

                ],
                inputs=[video_input, model_option]
            )
        with gr.Row():
            video_output1.render()
            video_output2.render()
            video_output3.render()

    submit_button.click(fn=process_video, inputs=[video_input, model_option],
                        outputs=[video_output1, video_output2, video_output3])


if mode == "local":
    demo.launch(server_name="0.0.0.0", server_port=6006, debug=True)
else:
    demo.launch(server_name="0.0.0.0", server_port=7860, debug=True)