app.py

"""Streamlit demo to visualize AVTime Dataset."""
import os
from os.path import join, exists, dirname, abspath, basename

import numpy as np
import pandas as pd
import streamlit as st

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

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


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


def read_spreadsheet(sheet_id, gid, **kwargs):
    BASE_URL = 'https://docs.google.com/spreadsheets/d/'
    df = pd.read_csv(BASE_URL + sheet_id + f'/export?gid={gid}&format=csv', **kwargs)
    # drop all rows which have atleast 1 NaN value
    df = df.dropna(axis=0)
    return df


SHEET_ID = "1ejT6SWbihrYzfB_npb0WS2UHd8TIxPZkMpqbJ-M_hfw"
GID = "0"


if __name__ == "__main__":
    # Add a title
    st.set_page_config(layout="wide")
    st.title("Clips from AVTime Dataset 🎬")
    
    if "df" not in st.session_state:
        # Read spreadsheet
        df = read_spreadsheet(sheet_id=SHEET_ID, gid=GID, on_bad_lines='skip')
        st.session_state.df = df
    else:
        df = st.session_state.df

    st.markdown(f"**Total number of relevant clips**: {len(df)}", unsafe_allow_html=True)

    reload_button = st.button("Reload")
    NUM = 9
    indices = np.random.randint(0, len(st.session_state.df), NUM)
    if reload_button:
        indices = np.random.randint(0, len(st.session_state.df), NUM)

    # Create a grid of videos
    grid = make_grid(3, 3)
    per_video_width = 360
    per_video_height = 240

    for idx, index in enumerate(indices):
        sample = st.session_state.df.iloc[index].to_dict()
        video_id = sample["video_id"]
        stime, etime = sample["start_time"], sample["end_time"]
        duration = etime - stime
        desc = sample["description"]
        
        i, j = idx // 3, idx % 3
        grid[i][j].caption(f"Segment duration: {duration}")
        url = f"https://www.youtube.com/embed/{video_id}?start={int(stime)}&end={int(etime)}"
        html_code = f"""
        <iframe height="{per_video_height}" width="{per_video_width}" src="{url}" frameborder="0" allowfullscreen></iframe>
        """
        grid[i][j].markdown(html_code, unsafe_allow_html=True)
        grid[i][j].caption(f"{desc}")