app.py

"""Streamlit demo to visualize auto-annotated Foley segments from movie clips."""
import os
from os.path import join, exists, dirname, abspath
import json

from tqdm import tqdm
import numpy as np
import pandas as pd
import streamlit as st
from moviepy.video.io.VideoFileClip import VideoFileClip

import warnings
warnings.simplefilter(action='ignore')


curr_filepath = abspath(__file__)
repo_path = dirname(dirname(curr_filepath))


def load_json(path: str) -> dict:
    """Helper to load json file"""
    with open(path, 'rb') as f:
        data = json.load(f)
    return data


def tqdm_iterator(items, desc=None, bar_format=None, **kwargs):
    tqdm._instances.clear()
    iterator = tqdm(
        items,
        desc=desc,
        bar_format='{l_bar}{bar:10}{r_bar}{bar:-10b}',
        **kwargs,
    )
    
    return iterator


def get_data_root_from_hostname():
    import socket

    data_root_lib = {
        "diva": "/ssd/pbagad/datasets/",
        "node": "/var/scratch/pbagad/datasets/",
        "fs4": "/var/scratch/pbagad/datasets/",
    }
    hostname = socket.gethostname()
    hostname = hostname[:4]
    
    data_root = data_root_lib.get(hostname, "NA")
    return data_root


def load_clips_df(df_path, data_dir, verbose=True, use_local=False):
    assert exists(df_path), f"File {df_path} does not exist"
    df = pd.read_csv(df_path)
    print(df.columns)
    if verbose:
        print("Number of clips:", len(df))
    # filter out clips that are not downloaded
    df["video_path"] = df["videoid"].apply(lambda x: join(data_dir, "pytube_videos", f"{x}.mp4"))
    if use_local:
        df = df[df["video_path"].apply(exists)]
    if verbose:
        print("Number of clips (with videos available):", len(df))
    df["audio_path"] = df["videoid"].apply(lambda x: join(data_dir, "pytube_audio", f"{x}.wav"))
    if use_local:
        df = df[df["audio_path"].apply(exists)]
    if verbose:
        print("Number of clips (with audio available):", len(df))
    df["annot_path"] = df["videoid"].apply(lambda x: join(data_dir, "annotations", f"{x}.json"))
    if use_local:
        df = df[df["annot_path"].apply(exists)]
    if verbose:
        print("Number of clips (with annotations available):", len(df))
    return df


def summarize_classification_probs(silence, probs):
    summary = [f"Silence: {silence}"]
    summary += [f"{l.capitalize()}: {p}" for (l, p) in probs]
    return " | ".join(summary)


def cut_video_in_segments(video_path, segments):
    video = VideoFileClip(video_path)
    tmp_dir = os.path.join(os.path.expanduser("~"), "tmp")
    clip_paths = [f"{tmp_dir}/clip_{i}.mp4" for i in range(len(segments))]
    iterator = tqdm_iterator(
        zip(segments, clip_paths), total=len(segments), desc="Preparing clips",
    )
    clips = [
        video.subclip(x, y).write_videofile(f, logger=None, verbose=False) \
            for (x, y), f in iterator
    ]
    return clip_paths


def process_sample(row):

    video_path = row["video_path"]
    audio_path = row["audio_path"]

    annot = load_json(row["annot_filtered"])
    seg_indices = [i for i, flag in enumerate(annot["keep_status"]) if flag]

    keys = ["non_speech_segments", "silence_prob", "audiomae_on_audioset", "duration"]
    for k in keys:
        annot[k] = [x for i, x in enumerate(annot[k]) if i in seg_indices]
    del annot["keep_status"]

    labels = [
        summarize_classification_probs(
            annot["silence_prob"][i], annot["audiomae_on_audioset"][i]
        ) for i in range(len(annot["non_speech_segments"]))
    ]
    clip_paths = cut_video_in_segments(video_path, annot["non_speech_segments"])
    return clip_paths, labels, annot["non_speech_segments"], annot["duration"]


def make_grid(cols,rows):
    grid = [0]*cols
    for i in range(cols):
        with st.container():
            grid[i] = st.columns(rows)
    return grid


if __name__ == "__main__":
    
    # Streamlit app code
    st.set_page_config(layout="wide")
    st.title("Foley Segments from Condensed Movies Dataset 🎬")
    
    st.markdown(
        "> **Note**: This demo shows cut out segments from clips in the [Condensed Movies](https://www.robots.ox.ac.uk/~vgg/data/condensed-movies/) dataset. "\
        "The segments are adjudged to have Foley sounds based on AudioMAE predictions on AudioSet classes. "\
        "The segments with duration not in [2s, 30s] or those with silence or high probability of speech/music sounds are filtered out. "\
        "However, segments can still be noisy. Furthermore, even if a clip has "\
        "Foley, there can still be background music/score which we have not removed yet."
    )
    st.markdown(
        """> <span style="color:red">Warning</span>: Currently, each clip can be played only once. Replaying starts the clip from beginning of the video.""",
        unsafe_allow_html=True
    )
    st.markdown(
        "**Instructions**: Click the **Reload** button to see segments from a new clip. "\
        "Reloading the page is not necessary."
    )
    
    use_local = False
    data_root = get_data_root_from_hostname()
    data_dir = join(data_root, "CondensedMovies")
    video_dir = join(data_dir, "pytube_videos")

    annot_dir = join(repo_path, "external/CondensedMovies/data/metadata/")

    if "subdf" not in st.session_state:
        df = load_clips_df(join(".", "clips.csv"), data_dir, verbose=True)
        df["annot_filtered"] = df["annot_path"].apply(lambda x: x.replace(".json", "_filtered.json"))
        df = df[df["annot_filtered"].apply(exists)]
        df["num_foley_segments"] = df["annot_filtered"].apply(lambda f: sum(load_json(f)["keep_status"]))
        subdf = df[df["num_foley_segments"].apply(lambda x: x > 0)]
        st.session_state.subdf = subdf
        num_foley = subdf["num_foley_segments"].sum()
        st.session_state.num_foley = num_foley
        print("Loaded subdf with {} rows".format(len(subdf)))

    
    reload_button = st.button("Reload")
    # index = 0
    index = np.random.randint(0, len(st.session_state.subdf))
    if reload_button:
        index = np.random.randint(0, len(st.session_state.subdf))

    row = st.session_state.subdf.iloc[index].to_dict()
    if use_local:
        clip_paths, labels, segments, durations = process_sample(row)
    else:
        annot = load_json(row["annot_filtered"])
        seg_indices = [i for i, flag in enumerate(annot["keep_status"]) if flag]
        keys = ["non_speech_segments", "silence_prob", "audiomae_on_audioset", "duration"]
        for k in keys:
            annot[k] = [x for i, x in enumerate(annot[k]) if i in seg_indices]
        del annot["keep_status"]
        labels = [
            summarize_classification_probs(
                annot["silence_prob"][i], annot["audiomae_on_audioset"][i]
            ) for i in range(len(annot["non_speech_segments"]))
        ]
        segments, durations = annot["non_speech_segments"], annot["duration"]
        clip_paths = [f"https://www.youtube.com/watch?v={row['videoid']}"] * len(segments)

    # Make a grid of videos and captions in streamlit
    videos = clip_paths
    video_id = row["videoid"]
    movie = row["title"]
    st.markdown(f"Showing Foley segments from a clip in movie: **{movie}**")

    # Create a grid of videos
    grid = make_grid(3, 3)

    # Add videos to the grid
    for idx in range(0, min(len(videos), 9)):
        i, j = idx // 3, idx % 3
    
        start, end = segments[idx]
        duration = durations[idx]

        grid[i][j].caption(f"Segment duration: {duration}")
        if not use_local:
            url = f"https://www.youtube.com/embed/{video_id}?start={int(start)}&end={int(end)}"
            html_code = f"""
            <iframe height="320" width="420" src="{url}" frameborder="0" allowfullscreen></iframe>
            """
            grid[i][j].markdown(html_code, unsafe_allow_html=True)
        else:
            grid[i][j].video(videos[idx])
        grid[i][j].caption(f"{labels[idx]}")
    
    
    st.markdown("##### Some stats")
    st.write(f"Total number of unique clips: {len(st.session_state.subdf)}")
    st.write(f"Total number of foley segments: {st.session_state.num_foley}")