Upload folder using huggingface_hub

shared/README.md

@@ -0,0 +1 @@
+This folder shall have code utilities shared across different tasks.

shared/scripts/upload_data.py

@@ -0,0 +1,109 @@
+"""Uploads dataset to huggingface datasets."""
+import os
+import sys
+
+import pandas as pd
+import numpy as np
+
+from huggingface_hub import HfApi
+import shared.utils as su
+from sound_of_water.data.csv_loader import (
+    load_csv_sound_of_water,
+    configure_paths_sound_of_water,
+)
+
+
+if __name__ == "__main__":
+    api = HfApi()
+
+    data_root = "/work/piyush/from_nfs2/datasets/SoundOfWater"
+    repo_id = "bpiyush/sound-of-water"
+
+    save_splits = False
+    if save_splits:
+        # Load CSV
+        paths = configure_paths_sound_of_water(data_root)
+        df = load_csv_sound_of_water(paths)
+        del df["video_clip_path"]
+        del df["audio_clip_path"]
+        del df["box_path"]
+        del df["mask_path"]
+
+        # Splits
+        train_ids = su.io.load_txt(os.path.join(data_root, "splits/train.txt"))
+        df_train = df[df.item_id.isin(train_ids)]
+        df_train["file_name"] = df_train["item_id"].apply(lambda x: f"videos/{x}.mp4")
+        df_train.to_csv(os.path.join(data_root, "splits/train.csv"), index=False)
+        print(" [:::] Train split saved.")
+
+        test_I_ids = su.io.load_txt(os.path.join(data_root, "splits/test_I.txt"))
+        df_test_I = df[df.item_id.isin(test_I_ids)]
+        df_test_I["file_name"] = df_test_I["item_id"].apply(lambda x: f"videos/{x}.mp4")
+        df_test_I.to_csv(os.path.join(data_root, "splits/test_I.csv"), index=False)
+        print(" [:::] Test I split saved.")
+
+        test_II_ids = su.io.load_txt(os.path.join(data_root, "splits/test_II.txt"))
+        df_test_II = df[df.item_id.isin(test_II_ids)]
+        df_test_II["file_name"] = df_test_II["item_id"].apply(lambda x: f"videos/{x}.mp4")
+        df_test_II.to_csv(os.path.join(data_root, "splits/test_II.csv"), index=False)
+        print(" [:::] Test II split saved.")
+
+        test_III_ids = su.io.load_txt(os.path.join(data_root, "splits/test_III.txt"))
+        df_test_III = df[df.item_id.isin(test_III_ids)]
+        df_test_III["file_name"] = df_test_III["item_id"].apply(lambda x: f"videos/{x}.mp4")
+        df_test_III.to_csv(os.path.join(data_root, "splits/test_III.csv"), index=False)
+        print(" [:::] Test III split saved.")
+
+
+    create_splits = False
+    if create_splits:
+        train_ids = su.io.load_txt(os.path.join(data_root, "splits/train.txt"))
+        train_ids = np.unique(train_ids)
+
+        test_I_ids = su.io.load_txt(os.path.join(data_root, "splits/test_I.txt"))
+        test_I_ids = np.unique(test_I_ids)
+
+        other_ids = np.array(
+            list(set(df.item_id.unique()) - set(train_ids) - set(test_I_ids))
+        )
+        sub_df = df[~df.item_id.isin(set(train_ids) | set(test_I_ids))]
+        X = sub_df[
+            (sub_df.visibility != "transparent") & (sub_df["shape"].isin(["cylindrical", "semiconical"]))
+        ]
+        test_II_ids = list(X.item_id.unique())
+        assert set(test_II_ids).intersection(set(train_ids)) == set()
+        assert set(test_II_ids).intersection(set(test_I_ids)) == set()
+        su.io.save_txt(test_II_ids, os.path.join(data_root, "splits/test_II.txt"))
+
+        X = sub_df[ 
+            (sub_df.visibility.isin(["transparent", "opaque"])) & \
+            (sub_df["shape"].isin(["cylindrical", "semiconical", "bottleneck"]))
+        ]
+        test_III_ids = list(X.item_id.unique())
+        assert set(test_III_ids).intersection(set(train_ids)) == set()
+        assert set(test_III_ids).intersection(set(test_I_ids)) == set()
+        assert set(test_III_ids).intersection(set(test_II_ids)) != set()
+        su.io.save_txt(test_III_ids, os.path.join(data_root, "splits/test_III.txt"))
+
+    upload_file = True
+    if upload_file:
+        file = "README.md"
+        print(f" [:::] Uploading file: {file}")
+        api.upload_file(
+            path_or_fileobj=os.path.join(data_root, file),
+            path_in_repo=file,
+            repo_id=repo_id,
+            repo_type="dataset",
+        )
+
+    upload_folder = False
+    if upload_folder:
+        # Upload splits folder
+        foldername = "annotations"
+        print(f" [:::] Uploading folder: {foldername}")
+        api.upload_folder(
+            folder_path=os.path.join(data_root, foldername),
+            path_in_repo=foldername, # Upload to a specific folder
+            repo_id=repo_id,
+            repo_type="dataset",
+        )

shared/utils/__init__.py

@@ -0,0 +1,12 @@
+import shared.utils.paths as paths
+import shared.utils.log as log
+import shared.utils.io as io
+import shared.utils.audio as audio
+import shared.utils.image as image
+import shared.utils.av as av
+import shared.utils.pandas_utils as pd_utils
+import shared.utils.visualize as visualize
+import shared.utils.metrics as metrics
+import shared.utils.misc as misc
+import shared.utils.keypoint_matching as keypoint_matching
+import shared.utils.physics as physics

shared/utils/audio.py

@@ -0,0 +1,224 @@
+"""Audio utils"""
+import librosa
+import numpy as np
+import matplotlib.pyplot as plt
+
+
+def load_audio(audio_path: str, sr: int = None, max_duration: int = 10., start: int = 0, stop: int = None):
+    """Loads audio and pads/trims it to max_duration"""
+    data, sr = librosa.load(audio_path, sr=sr)
+    
+    if stop is not None:
+        start = int(start * sr)
+        stop = int(stop * sr)
+        data = data[start:stop]
+    
+    # Convert to mono
+    if len(data.shape) > 1:
+        data = np.mean(data, axis=1)
+    
+    n_frames = int(max_duration * sr)
+    if len(data) > n_frames:
+        data = data[:n_frames]
+    elif len(data) < n_frames:
+        data = np.pad(data, (0, n_frames - len(data)), "constant")
+    return data, sr
+    
+
+# def compute_spectrogram(data: np.ndarray, sr: int):
+#     D = librosa.stft(data)  # STFT of y
+#     S_db = librosa.amplitude_to_db(np.abs(D), ref=np.max)
+#     return S_db
+
+
+def compute_spec_freq_mean(S_db: np.ndarray, eps=1e-5):
+    # Compute mean of spectrogram over frequency axis
+    S_db_normalized = (S_db - S_db.mean(axis=1)[:, None]) / (S_db.std(axis=1)[:, None] + eps)
+    S_db_over_time = S_db_normalized.sum(axis=0)
+    return S_db_over_time
+
+
+def process_audiofile(audio_path, functions=["load_audio", "compute_spectrogram", "compute_spec_freq_mean"]):
+    """Processes audio file with a list of functions"""
+    data, sr = load_audio(audio_path)
+    for function in functions:
+        if function == "load_audio":
+            pass
+        elif function == "compute_spectrogram":
+            data = compute_spectrogram(data, sr)
+        elif function == "compute_spec_freq_mean":
+            data = compute_spec_freq_mean(data)
+        else:
+            raise ValueError(f"Unknown function {function}")
+    return data
+
+
+
+"""PyDub's silence detection is based on the energy of the audio signal."""
+import numpy as np
+
+
+def sigmoid(x):
+    return 1 / (1 + np.exp(-x))
+
+
+class SilenceDetector:
+    
+
+    def __init__(self, silence_thresh=-36) -> None:
+        self.silence_thresh = silence_thresh
+    
+    def __call__(self, audio_path: str, start=None, end=None):
+        
+        import pydub
+        from pydub.utils import db_to_float
+        
+        try:
+            waveform = pydub.AudioSegment.from_file(audio_path)
+        except:
+            print("Error loading audio file: ", audio_path)
+            return 100.0
+        
+        start_ms = int(start * 1000) if start else 0
+        end_ms = int(end * 1000) if end else len(waveform)
+        waveform = waveform[start_ms:end_ms]
+        
+        # convert silence threshold to a float value (so we can compare it to rms)
+        silence_thresh = db_to_float(self.silence_thresh) * waveform.max_possible_amplitude
+        
+        if waveform.rms == 0:
+            return 100.0
+
+        silence_prob = sigmoid((silence_thresh - waveform.rms) / waveform.rms)
+
+        # return waveform.rms <= silence_thresh
+        return np.round(100 * silence_prob, 2)
+
+
+def frequency_bin_to_value(bin_index, sr, n_fft):
+    return int(bin_index * sr / n_fft)
+
+
+def time_bin_to_value(bin_index, hop_length, sr):
+    return (bin_index) * (hop_length / sr)
+
+
+def add_time_annotations(ax, nt_bins, hop_length, sr, skip=50):
+    # Show time (s) values on the x-axis
+    t_bins = np.arange(nt_bins)
+    t_vals = np.round(np.array([time_bin_to_value(tb, hop_length, sr) for tb in t_bins]), 1)
+    try:
+        ax.set_xticks(t_bins[::skip], t_vals[::skip])
+    except:
+        pass
+    ax.set_xlabel("Time (s)")
+
+
+def add_freq_annotations(ax, nf_bins, sr, n_fft, skip=50):
+    f_bins = np.arange(nf_bins)
+    f_vals = np.array([frequency_bin_to_value(fb, sr, n_fft) for fb in f_bins])
+    try:
+        ax.set_yticks(f_bins[::skip], f_vals[::skip])
+    except:
+        pass
+    # ax.set_yticks(f_bins[::skip])
+    # ax.set_yticklabels(f_vals[::skip])
+    ax.set_ylabel("Frequency (Hz)")
+
+
+def show_single_spectrogram(
+    spec,
+    sr,
+    n_fft,
+    hop_length,
+    ax=None,
+    fig=None,
+    figsize=(10, 2),
+    cmap="viridis",
+    colorbar=True,
+    show=True,
+    format='%+2.0f dB',
+    xlabel='Time (s)',
+    ylabel="Frequency (Hz)",
+    title=None,
+    show_dom_freq=False,
+):
+
+    if ax is None:
+        fig, ax = plt.subplots(1, 1, figsize=figsize)
+    axim = ax.imshow(spec, origin="lower", cmap=cmap)
+
+    # Show frequency (Hz) values on y-axis
+    nf_bins, nt_bins = spec.shape
+
+    if "frequency" in ylabel.lower():
+        # Add frequency annotation
+        add_freq_annotations(ax, nf_bins, sr, n_fft)
+
+    # Add time annotation
+    add_time_annotations(ax, nt_bins, hop_length, sr)
+
+    ax.set_title(title)
+    ax.set_xlabel(xlabel)
+    ax.set_ylabel(ylabel)
+
+    if colorbar:
+        fig.colorbar(axim, ax=ax, orientation='vertical', fraction=0.01, format=format)
+
+    if show_dom_freq:
+        fmax = spec.argmax(axis=0)
+        ax.scatter(np.arange(spec.shape[1]), fmax, color="white", s=0.2)
+
+    if show:
+        plt.show()
+
+
+def compute_spectrogram(y, n_fft, hop_length, margin, n_mels=None):
+    # STFT
+    D = librosa.stft(y, n_fft=n_fft, hop_length=hop_length)
+
+    # Run HPSS
+    S, _ = librosa.decompose.hpss(D, margin=margin)
+
+    # DB
+    S = librosa.amplitude_to_db(np.abs(S), ref=np.max)
+
+    if n_mels is not None:
+        S = librosa.feature.melspectrogram(S=S, n_mels=n_mels)
+
+    return S
+
+
+def show_spectrogram(S, sr, n_fft=512, hop_length=256, figsize=(10, 3), n_mels=None, ax=None, show=True):
+    if ax is None:
+        fig, ax = plt.subplots(1, 1, figsize=figsize)
+    y_axis = "mel" if n_mels is not None else "linear"
+    librosa.display.specshow(
+        S,
+        sr=sr,
+        hop_length=hop_length,
+        n_fft=n_fft,
+        y_axis=y_axis,
+        x_axis='time',
+        ax=ax,
+    )
+    ax.set_title("LogSpectrogram" if n_mels is None else "LogMelSpectrogram")
+    if show:
+        plt.show()
+
+
+def show_frame_and_spectrogram(frame, S, sr, figsize=(12, 4), show=True, axes=None, **spec_args):
+    if axes is None:
+        fig, axes = plt.subplots(1, 2, figsize=figsize, gridspec_kw={"width_ratios": [0.2, 0.8]})
+    ax = axes[0]
+    ax.imshow(frame)
+    ax.set_xticks([])
+    ax.set_yticks([])
+
+    ax = axes[1]
+    show_spectrogram(S=S, sr=sr, ax=ax, show=False, **spec_args)
+
+    plt.tight_layout()
+
+    if show:
+        plt.show()

shared/utils/av.py

@@ -0,0 +1,93 @@
+"""Audio-visual helper functions."""
+import cv2
+import numpy as np
+import torch
+
+
+def save_video_with_audio(video, audio, output_path):
+    """
+    Saves a video file with audio.
+    
+    Args:
+        video (np.ndarray): Video frames.
+        audio (np.ndarray): Audio samples.
+        output_path (str): Output path.
+    """
+
+    # check the correct shape and format for audio
+    assert isinstance(audio, np.ndarray)
+    assert len(audio.shape) == 2
+    assert audio.shape[1] in [1, 2]
+
+    # create video writer
+    video_writer = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), 30, (video.shape[2], video.shape[1]))
+    # write the image frames to the video
+    for frame in video:
+        video_writer.write(frame)
+    # add the audio data to the video
+    video_writer.write(audio)
+    # release the VideoWriter object
+    video_writer.release()
+
+
+def save_video_from_image_sequence_and_audio(sequence, audio, save_path, video_fps=15, audio_fps=22100):
+    from moviepy.editor import VideoClip, AudioClip, ImageSequenceClip
+    from moviepy.audio.AudioClip import AudioArrayClip
+    
+    assert isinstance(sequence, list) and isinstance(audio, (np.ndarray, torch.Tensor))
+    assert len(audio.shape) == 2 and audio.shape[1] in [1, 2]
+    
+    video_duration = len(sequence) / video_fps
+    audio_duration = len(audio) / audio_fps
+    # # print(f"Video duration: {video_duration:.2f}s, audio duration: {audio_duration:.2f}s")
+    # assert video_duration == audio_duration, \
+    #     f"Video duration ({video_duration}) and audio duration ({audio_duration}) do not match."
+
+    video_clip = ImageSequenceClip(sequence, fps=video_fps)
+    audio_clip = AudioArrayClip(audio, fps=audio_fps)
+    video_clip = video_clip.set_audio(audio_clip)
+    # video_clip.write_videofile(save_path, verbose=True, logger=None, fps=video_fps, audio_fps=audio_fps)
+    video_clip.write_videofile(save_path, verbose=False, logger=None)
+
+
+import cv2, os
+import argparse
+import numpy as np
+from glob import glob
+import librosa
+import subprocess
+
+
+def generate_video(args):
+
+	frames = glob('{}/*.png'.format(args.input_dir))
+	print("Total frames = ", len(frames))
+
+	frames.sort(key = lambda x: int(x.split("/")[-1].split(".")[0]))
+
+	img = cv2.imread(frames[0])
+	print(img.shape)
+	fname = 'inference.avi'
+	video = cv2.VideoWriter(
+        fname, cv2.VideoWriter_fourcc(*'DIVX'), args.fps, (img.shape[1], img.shape[0]),
+    )
+ 
+	for i in range(len(frames)):
+		img = cv2.imread(frames[i])
+		video.write(img)
+	
+	video.release()
+
+	output_file_name = args.output_video
+
+	no_sound_video = output_file_name + '_nosound.mp4'
+	subprocess.call('ffmpeg -hide_banner -loglevel panic -i %s -c copy -an -strict -2 %s' % (fname, no_sound_video), shell=True)
+
+	if args.audio_file is not None:
+		video_output = output_file_name + '.mp4'
+		subprocess.call('ffmpeg -hide_banner -loglevel panic -y -i %s -i %s -strict -2 -q:v 1 %s' % 
+						(args.audio_file, no_sound_video, video_output), shell=True)
+
+		os.remove(no_sound_video)
+	
+	os.remove(fname)

shared/utils/classification.py

@@ -0,0 +1,47 @@
+"""Helper functions for classification tasks."""
+import matplotlib.pyplot as plt
+import pandas as pd
+import numpy as np
+
+
+def plot_metric_curve(
+        xvalues, yvalues, thresholds, title=None,
+        figsize=(8, 7), show_thresholds=True, show_legend=True,
+        ylabel='X', xlabel='Y', ax=None, text_delta=0.01,
+        label="Metric Curve", color="royalblue", show=False,
+        fill=None,
+    ):
+    """Plot a metric curve, e.g., PR curve or ROC curve."""
+
+    if ax is None:
+        fig, ax = plt.subplots(1, 1, figsize=figsize)
+
+    ax.grid(alpha=0.3)
+    ax.set_title(title)
+    ax.set_ylabel(ylabel)
+    ax.set_xlabel(xlabel)
+    
+    ax.plot(xvalues, yvalues, marker='o', label=label, color=color)
+    ax.set_xlim(-0.08, 1.08)
+    ax.set_ylim(-0.08, 1.08)
+    
+    if fill is not None:
+        yticks = ax.get_yticks()
+        ax.fill_between(xvalues, yvalues, "", alpha=0.08, color=color)
+        # Add `fill` inside the curve
+        # Find a single (x, y) s.t. it is inside the curve
+        ax.text(0.4, 0.5, fill, color=color)
+        ax.set_yticks(yticks)
+        ax.set_yticklabels([f"{y:.1f}" for y in yticks])
+        ax.set_ylim(-0.08, 1.08)
+    
+    # Show thresholds
+    if show_thresholds:
+        for x, y, t in zip(xvalues, yvalues, thresholds):
+            ax.text(x + text_delta, y + text_delta, np.round(t, 2), color=color, alpha=0.5)
+    
+    if show_legend:
+        ax.legend()
+    
+    if show:
+        plt.show()

shared/utils/epic.py

@@ -0,0 +1,15 @@
+"""Utils specific for EPIC data."""
+import datetime
+
+
+def timestamp_to_seconds(timestamp: str):
+    # Parse the timestamp string into a datetime object
+    time_obj = datetime.datetime.strptime(timestamp, '%H:%M:%S.%f')
+
+    # Calculate the total number of seconds using the timedelta object
+    total_seconds = time_obj.time().second \
+        + time_obj.time().minute * 60 \
+        + time_obj.time().hour * 3600 \
+        + time_obj.time().microsecond / 1000000
+    
+    return total_seconds

shared/utils/image.py

@@ -0,0 +1,81 @@
+"""Image operations."""
+from copy import deepcopy
+from PIL import Image
+import numpy as np
+
+
+def center_crop(im: Image):
+    width, height = im.size
+    new_width = width if width < height else height
+    new_height = height if height < width else width 
+
+    left = (width - new_width)/2
+    top = (height - new_height)/2
+    right = (width + new_width)/2
+    bottom = (height + new_height)/2
+
+    # Crop the center of the image
+    im = im.crop((left, top, right, bottom))
+    
+    return im
+
+
+def pad_to_square(im: Image, color=(0, 0, 0)):
+    im = deepcopy(im)
+    width, height = im.size
+
+    vert_pad = (max(width, height) - height) // 2
+    hor_pad = (max(width, height) - width) // 2
+    
+    if len(im.mode) == 3:
+        color = (0, 0, 0)
+    elif len(im.mode) == 1:
+        color = 0
+    else:
+        raise ValueError(f"Image mode not supported. Image has {im.mode} channels.")
+    
+    return add_margin(im, vert_pad, hor_pad, vert_pad, hor_pad, color=color)
+
+
+def add_margin(pil_img, top, right, bottom, left, color=(0, 0, 0)):
+    """Ref: https://note.nkmk.me/en/python-pillow-add-margin-expand-canvas/"""
+    width, height = pil_img.size
+    new_width = width + right + left
+    new_height = height + top + bottom
+    result = Image.new(pil_img.mode, (new_width, new_height), color)
+    result.paste(pil_img, (left, top))
+    return result
+
+
+def resize_image(image, new_height, new_width):
+    # Convert the numpy array image to PIL Image
+    pil_image = Image.fromarray(image)
+
+    # Resize the PIL Image
+    resized_image = pil_image.resize((new_width, new_height))
+
+    # Convert the resized PIL Image back to numpy array
+    resized_image_np = np.array(resized_image)
+
+    return resized_image_np
+
+
+def pad_to_width(pil_image, new_width, color=(0, 0, 0)):
+    """Pad the image to the specified width."""
+    # Convert the numpy array image to PIL Image
+    # pil_image = Image.fromarray(image)
+
+    # Get the current width and height of the image
+    width, height = pil_image.size
+    assert new_width > width, f"New width {new_width} is less than the current width {width}."
+
+    # Calculate the padding required
+    hor_pad = new_width - width
+
+    # Add padding to the image
+    padded_image = add_margin(pil_image, 0, hor_pad, 0, 0, color=color)
+
+    # Convert the padded PIL Image back to numpy array
+    # padded_image_np = np.array(padded_image)
+
+    return padded_image

shared/utils/io.py

@@ -0,0 +1,151 @@
+"""
+Utilities for input-output loading/saving.
+"""
+
+from typing import Any, List
+import yaml
+import pickle
+import json
+import pandas as pd
+
+
+class PrettySafeLoader(yaml.SafeLoader):
+    """Custom loader for reading YAML files"""
+    def construct_python_tuple(self, node):
+        return tuple(self.construct_sequence(node))
+
+
+PrettySafeLoader.add_constructor(
+    u'tag:yaml.org,2002:python/tuple',
+    PrettySafeLoader.construct_python_tuple
+)
+
+
+def load_yml(path: str, loader_type: str = 'default'):
+    """Read params from a yml file.
+
+    Args:
+        path (str): path to the .yml file
+        loader_type (str, optional): type of loader used to load yml files. Defaults to 'default'.
+
+    Returns:
+        Any: object (typically dict) loaded from .yml file
+    """
+    assert loader_type in ['default', 'safe']
+
+    loader = yaml.Loader if (loader_type == "default") else PrettySafeLoader
+
+    with open(path, 'r') as f:
+        data = yaml.load(f, Loader=loader)
+
+    return data
+
+
+def save_yml(data: dict, path: str):
+    """Save params in the given yml file path.
+
+    Args:
+        data (dict): data object to save
+        path (str): path to .yml file to be saved
+    """
+    with open(path, 'w') as f:
+        yaml.dump(data, f, default_flow_style=False)
+
+
+def load_pkl(path: str, encoding: str = "ascii"):
+    """Loads a .pkl file.
+
+    Args:
+        path (str): path to the .pkl file
+        encoding (str, optional): encoding to use for loading. Defaults to "ascii".
+
+    Returns:
+        Any: unpickled object
+    """
+    return pickle.load(open(path, "rb"), encoding=encoding)
+
+
+def save_pkl(data: Any, path: str) -> None:
+    """Saves given object into .pkl file
+
+    Args:
+        data (Any): object to be saved
+        path (str): path to the location to be saved at
+    """
+    with open(path, 'wb') as f:
+        pickle.dump(data, f)
+
+
+def load_json(path: str) -> dict:
+    """Helper to load json file"""
+    with open(path, 'rb') as f:
+        data = json.load(f)
+    return data
+
+
+def save_json(data: dict, path: str):
+    """Helper to save `dict` as .json file."""
+    with open(path, 'w') as f:
+        json.dump(data, f)
+
+
+def load_txt(path: str):
+    """Loads lines of a .txt file.
+
+    Args:
+        path (str): path to the .txt file
+
+    Returns:
+        List: lines of .txt file
+    """
+    with open(path) as f:
+        lines = f.read().splitlines()
+    return lines
+
+
+def save_txt(data: dict, path: str):
+    """Writes data (lines) to a txt file.
+
+    Args:
+        data (dict): List of strings
+        path (str): path to .txt file
+    """
+    assert isinstance(data, list)
+
+    lines = "\n".join(data)
+    with open(path, "w") as f:
+        f.write(str(lines))
+
+
+def read_spreadsheet(sheet_id, gid, url=None, drop_na=True, **kwargs):
+    if url is None:
+        BASE_URL = 'https://docs.google.com/spreadsheets/d/'
+        url = BASE_URL + sheet_id + f'/export?gid={gid}&format=csv'
+    df = pd.read_csv(url, **kwargs)
+    
+    if drop_na:
+        # drop all rows which have atleast 1 NaN value
+        df = df.dropna(axis=0)
+
+    return df
+
+
+def load_midi(file, rate=16000):
+    import pretty_midi
+    assert file.endswith('.mid')
+    pm = pretty_midi.PrettyMIDI(file)
+    y = pm.synthesize(fs=rate)
+    return y, rate
+
+
+def load_ptz(path):
+    import gzip
+    import torch
+    with gzip.open(path, 'rb') as f:
+        data = torch.load(f)
+    return data
+
+
+def save_video(frames, path, fps=30):
+    import imageio
+    imageio.mimwrite(path, frames, fps=fps)

shared/utils/keypoint_matching.py

@@ -0,0 +1,330 @@
+"""Implements keypoint matching for a pair of images."""
+import os
+import numpy as np
+import PIL
+import cv2
+import matplotlib.pyplot as plt
+
+
+def show_single_image(img, figsize=(7, 5), title="Single image"):
+    """Displays a single image."""
+    fig = plt.figure(figsize=figsize)
+    plt.axis("off")
+    plt.imshow(img)
+    plt.title(title)
+    plt.show()
+
+
+def show_two_images(img1, img2, title="Two images"):
+    """Displays a pair of images."""
+    fig, ax = plt.subplots(1, 2, figsize=(10, 5), constrained_layout=True)
+
+    ax[0].axis("off")
+    ax[0].imshow(img1)
+
+    ax[1].axis("off")
+    ax[1].imshow(img2)
+
+    plt.suptitle(title)
+    plt.show()
+
+
+def show_three_images(img1, img2, img3, ax1_title="", ax2_title="", ax3_title="", title="Three images"):
+    """Displays a triplet of images."""
+    fig, ax = plt.subplots(1, 3, figsize=(15, 5), constrained_layout=True)
+
+    ax[0].axis("off")
+    ax[0].imshow(img1)
+    ax[0].set_title(ax1_title)
+
+    ax[1].axis("off")
+    ax[1].imshow(img2)
+    ax[1].set_title(ax2_title)
+
+    ax[2].axis("off")
+    ax[2].imshow(img3)
+    ax[2].set_title(ax3_title)
+
+    plt.suptitle(title)
+    plt.show()
+
+
+class KeypointMatcher:
+    """Class for Keypoint matching for a pair of images."""
+
+    def __init__(self, **sift_args) -> None:
+        self.SIFT = cv2.SIFT_create(**sift_args)
+        self.BFMatcher = cv2.BFMatcher()
+    
+    @staticmethod
+    def _check_images(img1: np.ndarray, img2: np.ndarray):
+        assert isinstance(img1, np.ndarray)
+        assert len(img1.shape) == 2
+
+        assert isinstance(img2, np.ndarray)
+        assert len(img2.shape) == 2
+
+        # assert img1.shape == img2.shape
+    
+    @staticmethod
+    def _show_matches(img1, kp1, img2, kp2, matches, K=10, figsize=(10, 5), drawMatches_args=dict(matchesThickness=3, singlePointColor=(0, 0, 0))):
+        """Displays matches found in the image"""
+        selected_matches = np.random.choice(matches, K)
+        img3 = cv2.drawMatches(img1, kp1, img2, kp2, selected_matches, outImg=None, **drawMatches_args)
+        show_single_image(img3, figsize=figsize, title=f"Randomly selected K = {K} matches between the pair of images.")
+        return img3
+
+    def match(self, img1: PIL.Image, img2: PIL.Image, show_matches: bool = True):
+        """Finds, describes and matches keypoints in given pair of images."""
+        
+        img1 = np.array(img1)
+        img1 = cv2.cvtColor(img1, cv2.COLOR_RGB2GRAY)
+        
+        img2 = np.array(img2)
+        img2 = cv2.cvtColor(img2, cv2.COLOR_RGB2GRAY)
+        
+        # check input images
+        self._check_images(img1, img2)
+
+        # find kps and descriptors in each image
+        kp1, des1 = self.SIFT.detectAndCompute(img1, None)
+        kp2, des2 = self.SIFT.detectAndCompute(img2, None)
+
+        # compute matches via Brute-force matching
+        matches = self.BFMatcher.match(des1, des2)
+
+        # sort them in the order of their distance
+        matches = sorted(matches, key = lambda x:x.distance)
+
+        if show_matches:
+            self._show_matches(img1, kp1, img2, kp2, matches)
+
+        return matches, kp1, des1, kp2, des2
+
+
+def warp(im, M, output_shape):
+    out = np.zeros((output_shape[0], output_shape[1]))
+    for i in range(output_shape[0]):
+        for j in range(output_shape[1]):
+            u, v = np.array([[i, j, 0, 0, 1, 0], [0, 0, i, j, 0, 1]]) @ M
+            u = int(round(u))
+            v = int(round(v))
+            if im.shape[0] > u >= 0 and im.shape[1] > v >= 0:
+                out[i, j] = im[u, v]
+
+    return out
+
+
+def project_2d_to_6d(X: np.ndarray):
+    """Projects X (N x 2) to Z (2N x 6) space."""
+    N = len(X)
+    assert X.shape == (N, 2)
+
+    Z = np.zeros((2 * N, 6))
+    # in columns 0 to 2, fill even indexed rows of Z with X, and fill 5th column with 1
+    Z[::2, 0:2] = X
+    Z[::2, 4] = 1.0
+    # in columns 2 to 4, fill odd indexed rows of Z with X
+    Z[1::2, 2:4] = X
+    Z[1::2, 5] = 1.0
+
+    return Z
+
+
+def project_6d_to_2d(Z: np.ndarray):
+    """Projects Z (2N x 6) to X (N x 2) space."""
+    N = len(Z) // 2
+    assert Z.shape == (2 * N, 6)
+
+    X_from_even_rows = Z[::2, 0:2]
+    X_from_odd_rows = Z[1::2, 2:4]
+    assert (X_from_even_rows == X_from_odd_rows).all()
+
+    return X_from_even_rows
+
+
+
+def project_2d_to_1d(X: np.ndarray):
+    """Returns X (N x 2) from Z (2N, 1)"""
+    N = len(X)
+    X_stretched = np.zeros(2 * N)
+    X_stretched[::2] = X[:, 0]
+    X_stretched[1::2] = X[:, 1]
+    return X_stretched
+
+
+def project_1d_to_2d(Z: np.ndarray):
+    """Returns X (N x 2) from Z (2N, 1)"""
+    N = len(Z) // 2
+    assert Z.shape == (2 * N,)
+
+    X = np.zeros((N, 2))
+    X[:, 0] = Z[::2]
+    X[:, 1] = Z[1::2]
+
+    return X
+
+
+def rigid_body_transform(X: np.ndarray, params: np.ndarray):
+    """Performs rigid body transformation of points X (N x 2) using params (6 x 1 flattened)"""
+    N = len(X)
+    assert X.shape == (N, 2)
+
+    X = project_2d_to_6d(X)
+
+    X_transformed = np.matmul(X, params)
+    X_transformed = project_1d_to_2d(X_transformed)
+    assert X_transformed.shape == (N, 2)
+
+    return X_transformed
+
+
+def rigid_body_transform_params(X1: np.ndarray, X2: np.ndarray):
+    """Returns rigid-body transform parameters RT (6 x 1) assuming transformation between X1 and X2"""
+    N = len(X1)
+    assert X1.shape == X2.shape
+    assert X1.shape == (N, 2)
+
+    # X2 = X1 * params => params = psuedoinverse(X1) * X2
+    X1_expanded = project_2d_to_6d(X1)
+    assert X1_expanded.shape == (2 * N, 6)
+
+    X2_stretched = project_2d_to_1d(X2)
+    assert X2_stretched.shape == (2 * N,)
+
+    params = np.dot(np.linalg.pinv(X1_expanded), X2_stretched)
+    return params
+
+
+class ImageAlignment:
+    """Class to perform alignment of a pair of images given keypoints."""
+
+    def __init__(self) -> None:
+        pass
+    
+    @staticmethod
+    def show_transformed_points(img1, img2, X1, kp1, kp2, matches, params, num_inliers, num_to_show=20):
+        import matplotlib.cm as cm
+
+        H1, W1 = img1.shape
+        H2, W2 = img2.shape
+        img = np.hstack([img1, img2])
+
+        random_matches = np.random.choice(matches, num_to_show)
+
+        fig, ax = plt.subplots(1, 1, figsize=(15, 6))
+        colors = cm.rainbow(np.linspace(0, 1, num_to_show))
+
+        for i, match in enumerate(random_matches):
+
+            # select a single match to visualize
+            x1, y1 = kp1[match.queryIdx].pt
+            x2, y2 = kp2[match.trainIdx].pt
+
+            # get (x1, y1) transformed to (x1_transformed, y1_transformed)
+            A = project_2d_to_6d(np.array([[x1, y1]]))
+            (x1_transformed, y1_transformed) = np.dot(A, params)
+
+            ax.imshow(img, cmap="gray")
+            ax.axis("off")
+            ax.scatter(x1_transformed + W1, y1_transformed, s=200, marker="x", color=colors[i])
+            ax.plot(
+                (x1, x1_transformed + W1), (y1, y1_transformed),
+                linestyle="--", color=colors[i], marker="o",
+            )
+
+        ax.set_title(
+            f"Points in image 1 mapped to transformed points estimated by {num_inliers} points.",
+            fontsize=18,
+        )
+
+        os.makedirs("./results/", exist_ok=True)
+        plt.savefig(f"./results/match_transformed_inliers_{num_inliers}.png", bbox_inches="tight")
+        plt.show()
+
+    def ransac(
+            self, img1, kp1, img2, kp2, matches, num_matches=6, max_iter=500,
+            radius_in_px=10, show_transformed=True, inlier_th_for_show=1000
+        ):
+        """Performs RANSAC to find best matches."""
+
+        best_inlier_count = 0
+        best_params = None
+
+        # get coordinates of all points in image 1
+        X1 = np.array([kp1[matches[i].queryIdx].pt for i in range(len(matches))])
+
+        # get coordinates of all points in image 2
+        X2 = np.array([kp2[matches[i].trainIdx].pt for i in range(len(matches))])
+
+        for i in range(max_iter):
+            # choose matches randomly
+            selected_matches = np.random.choice(matches, num_matches)
+
+            # get matched keypoints in img1
+            X1_selected = np.array([kp1[selected_matches[i].queryIdx].pt for i in range(len(selected_matches))])
+
+            # get matched keypoints in img2
+            X2_selected = np.array([kp2[selected_matches[i].trainIdx].pt for i in range(len(selected_matches))])
+
+            # get transformation parameters
+            params = rigid_body_transform_params(X1_selected, X2_selected)
+            
+            # transform X1 to get X2_transformed
+            X2_transformed = rigid_body_transform(X1, params)
+
+            # find inliers
+            diff = np.linalg.norm(X2_transformed - X2, axis=1)
+            indices = diff < radius_in_px
+            num_inliers = sum(indices)
+            if num_inliers > best_inlier_count:
+                print(f"Found {num_inliers} inliers!")
+                best_params = params
+                best_inlier_count = num_inliers
+
+                if show_transformed and num_inliers > inlier_th_for_show:
+                    self.show_transformed_points(img1, img2, X1, kp1, kp2, matches, best_params, num_inliers)
+
+        return best_params
+    
+    def align(
+            self, img1, kp1, img2, kp2, matches, num_matches=6,
+            max_iter=500, show_warped_image=True,
+            save_warped=False, path="results/sample.png",
+            method="custom"
+        ):
+        best_params = self.ransac(img1, kp1, img2, kp2, matches, max_iter=max_iter, num_matches=num_matches)
+
+        # apply the affine transformation using cv2.warpAffine()
+        rows, cols = img1.shape[:2]
+
+        if method == 'custom':
+            img1_warped = warp(img1, best_params, (rows, cols))
+        else:
+            M = np.zeros((2, 3))
+            M[0, :2] = best_params[:2]
+            M[1, :2] = best_params[2:4]
+            M[0, 2] = best_params[4]
+            M[1, 2] = best_params[5]
+            img1_warped = cv2.warpAffine(img1, M, (cols, rows))
+
+        if show_warped_image:
+            show_three_images(
+                img1, img2, img1_warped, title="",
+                ax1_title="Image 1", ax2_title="Image 2", ax3_title="Transformation: Image 1 to Image 2",
+            )
+
+        if save_warped:
+            plt.imsave(path, img1_warped)
+
+        return best_params
+
+
+if __name__ == "__main__":
+    # read & show images
+    boat1 = cv2.imread('boat1.pgm', cv2.IMREAD_GRAYSCALE)
+    boat2 = cv2.imread('boat2.pgm', cv2.IMREAD_GRAYSCALE)
+    show_two_images(boat1, boat2, title="Given pair of images.")
+
+    kp_matcher = KeypointMatcher(contrastThreshold=0.1, edgeThreshold=5)
+    matches, kp1, des1, kp2, des2 = kp_matcher.match(boat1, boat2, show_matches=True)

shared/utils/log.py

@@ -0,0 +1,72 @@
+"""Loggers."""
+import os
+from os.path import dirname, realpath, abspath
+from tqdm.auto import tqdm
+import numpy as np
+
+
+curr_filepath = abspath(__file__)
+repo_path = dirname(dirname(dirname(curr_filepath)))
+# repo_path = dirname(dirname(dirname(realpath(__file__))))
+
+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,
+    )
+    tqdm._instances.clear()
+    
+    return iterator
+
+
+def print_retrieval_metrics_for_csv(metrics, scale=100):
+    print_string = [
+        np.round(scale * metrics["R1"], 3),
+        np.round(scale * metrics["R5"], 3),
+        np.round(scale * metrics["R10"], 3),
+    ]
+    if "MR" in metrics:
+        print_string += [metrics["MR"]]
+    print()
+    print("Final metrics: ", ",".join([str(x) for x in print_string]))
+    print()
+
+
+def print_update(update, fillchar=":", color="yellow", pos="center"):
+    from termcolor import colored
+    # add ::: to the beginning and end of the update s.t. the total length of the
+    # update spans the whole terminal
+    try:
+        terminal_width = os.get_terminal_size().columns - 2
+    except:
+        terminal_width = 98
+    if pos == "center":
+        update = update.center(len(update) + 2, " ")
+        update = update.center(terminal_width, fillchar)
+    elif pos == "left":
+        update = update.ljust(terminal_width, fillchar)
+        update = update.ljust(len(update) + 2, " ")
+    elif pos == "right":
+        update = update.rjust(terminal_width, fillchar)
+        update = update.rjust(len(update) + 2, " ")
+    else:
+        raise ValueError("pos must be one of 'center', 'left', 'right'")
+    print(colored(update, color))
+
+
+def json_print(data, indent=4):
+    import json
+    print(json.dumps(data, indent=indent))
+
+
+def get_terminal_width():
+    import shutil
+    return shutil.get_terminal_size().columns
+
+
+if __name__ == "__main__":
+    print("Repo path:", repo_path)
+    

shared/utils/metrics.py

@@ -0,0 +1,458 @@
+"""Helpers for metric functions"""
+import numpy as np
+import pandas as pd
+
+
+def calculate_iou(box1, box2):
+    """
+    Calculate Intersection over Union (IoU) between two bounding boxes.
+
+    Args:
+        box1 (tuple): Coordinates of the first bounding box in the format (x1, y1, x2, y2).
+        box2 (tuple): Coordinates of the second bounding box in the format (x1, y1, x2, y2).
+
+    Returns:
+        float: Intersection over Union (IoU) score.
+    """
+    # Extract coordinates
+    x1, y1, x2, y2 = box1
+    x1_, y1_, x2_, y2_ = box2
+
+    # Calculate the intersection area
+    intersection_area = max(0, min(x2, x2_) - max(x1, x1_)) * max(0, min(y2, y2_) - max(y1, y1_))
+
+    # Calculate the areas of each bounding box
+    box1_area = (x2 - x1) * (y2 - y1)
+    box2_area = (x2_ - x1_) * (y2_ - y1_)
+
+    # Calculate IoU
+    iou = intersection_area / float(box1_area + box2_area - intersection_area)
+
+    return iou
+
+
+def compute_intersection_1d(x, y):
+    # sort the boxes
+    x1, x2 = sorted(x)
+    y1, y2 = sorted(y)
+    
+    # compute the intersection
+    intersection = max(0, min(x2, y2) - max(x1, y1))
+    
+    return intersection
+
+def compute_union_1d(x, y):
+    # sort the boxes
+    x1, x2 = sorted(x)
+    y1, y2 = sorted(y)
+    
+    # compute the union
+    union = max(x2, y2) - min(x1, y1)
+    
+    return union
+
+
+def compute_iou_1d(pred_box, true_box):
+    """
+    Compute IoU for 1D boxes.
+    
+    Args:
+        pred_box (float): Predicted box, [x1, x2]
+        true_box (float): Ground truth box, [x1, x2]
+    
+    Returns:
+        float: IoU
+    """
+    intersection = compute_intersection_1d(pred_box, true_box)
+    union = compute_union_1d(pred_box, true_box)
+    iou = intersection / union
+    return iou
+
+
+def compute_iou_1d_single_candidate_multiple_targets(pred_box, true_boxes):
+    """
+    Compute IoU for 1D boxes.
+    
+    Args:
+        pred_box (float): Predicted box, [x1, x2]
+        true_boxes (np.ndarray): Ground truth boxes, shape: (N, 2)
+    
+    Returns:
+        float: IoU
+    """
+    ious = []
+    for i, true_box in enumerate(true_boxes):
+        ious.append(compute_iou_1d(pred_box, true_box))
+    return np.array(ious)
+
+
+def compute_iou_1d_multiple_candidates_multiple_targets(pred_boxes, true_boxes):
+    """
+    Compute IoU for 1D boxes.
+    
+    Args:
+        pred_boxes (np.ndarray): Predicted boxes, shape: (N, 2)
+        true_boxes (np.ndarray): Ground truth boxes, shape: (N, 2)
+    
+    Returns:
+        float: IoU
+    """
+    iou_matrix = np.zeros((len(pred_boxes), len(true_boxes)))
+    for i, pred_box in enumerate(pred_boxes):
+        for j, true_box in enumerate(true_boxes):
+            iou_matrix[i, j] = compute_iou_1d(pred_box, true_box)
+    return iou_matrix
+
+
+def compute_mean_iou_1d(pred_boxes, gt_boxes, threshold=0.5):
+    """
+    Computes mean IOU for 1D bounding boxes.
+    
+    Args:
+        pred_boxes (np.ndarray): Predicted boxes, shape: (N, 2)
+        gt_boxes (np.ndarray): Ground truth boxes, shape: (N, 2)
+        threshold (float): Threshold to consider a prediction correct
+    
+    Returns:
+        float: Mean IOU
+    """
+    # Compute IoU for each pair of boxes
+    iou_matrix = np.zeros((len(pred_boxes), len(gt_boxes)))
+    for i, pred_box in enumerate(pred_boxes):
+        for j, gt_box in enumerate(gt_boxes):
+            iou_matrix[i, j] = compute_iou_1d(pred_box, gt_box)
+
+    # Compute the max IoU for each predicted box
+    max_iou_indices = np.argmax(iou_matrix, axis=1)
+    max_iou = iou_matrix[np.arange(len(pred_boxes)), max_iou_indices]
+    
+    # For each predicted box, compute TP and FP ground truth boxes
+    tp = np.zeros(len(pred_boxes))
+    fp = np.zeros(len(pred_boxes))
+    iou = np.zeros(len(pred_boxes))
+    
+    tp = np.where(iou_matrix >= threshold, 1, 0)
+    tp = max_iou >= threshold
+    fp = max_iou < threshold
+    iou = max_iou
+    mean_iou = np.mean(iou)
+    import ipdb; ipdb.set_trace()
+
+
+
+
+
+
+
+def calculate_mAP_1d(pred_boxes, pred_scores, true_boxes, iou_thresh=0.5):
+    """Calculate mean average precision for 1D boxes.
+
+    Args:
+        pred_boxes (numpy array): Predicted boxes, shape (num_boxes,)
+        pred_scores (numpy array): Predicted scores, shape (num_boxes,)
+        true_boxes (numpy array): Ground truth boxes, shape (num_boxes,)
+        iou_thresh (float): IoU threshold to consider a prediction correct
+
+    Returns:
+        float: Mean average precision (mAP)
+    """
+    # Sort predicted boxes by score (in descending order)
+    sort_inds = np.argsort(pred_scores)[::-1]
+    pred_boxes = pred_boxes[sort_inds]
+    pred_scores = pred_scores[sort_inds]
+
+    # Compute true positives and false positives at each threshold
+    tp = np.zeros(len(pred_boxes))
+    fp = np.zeros(len(pred_boxes))
+    for i, box in enumerate(pred_boxes):
+        ious = np.abs(box - true_boxes) / np.maximum(1e-9, np.abs(box) + np.abs(true_boxes))
+        if len(ious) > 0:
+            max_iou_idx = np.argmax(ious)
+            if ious[max_iou_idx] >= iou_thresh:
+                if tp[max_iou_idx] == 0:
+                    tp[i] = 1
+                    fp[i] = 0
+                else:
+                    fp[i] = 1
+            else:
+                fp[i] = 1
+
+    # Compute precision and recall at each threshold
+    tp_cumsum = np.cumsum(tp)
+    fp_cumsum = np.cumsum(fp)
+    recall = tp_cumsum / len(true_boxes)
+    precision = tp_cumsum / (tp_cumsum + fp_cumsum)
+
+    # Compute AP as area under precision-recall curve
+    ap = 0
+    for t in np.arange(0, 1.1, 0.1):
+        if np.sum(recall >= t) == 0:
+            p = 0
+        else:
+            p = np.max(precision[recall >= t])
+        ap += p / 11
+
+    return ap
+
+
+def segment_iou(target_segment, candidate_segments):
+    """Compute the temporal intersection over union between a
+    target segment and all the test segments.
+    Parameters
+    ----------
+    target_segment : 1d array
+        Temporal target segment containing [starting, ending] times.
+    candidate_segments : 2d array
+        Temporal candidate segments containing N x [starting, ending] times.
+    Outputs
+    -------
+    tiou : 1d array
+        Temporal intersection over union score of the N's candidate segments.
+    """
+    tt1 = np.maximum(target_segment[0], candidate_segments[:, 0])
+    tt2 = np.minimum(target_segment[1], candidate_segments[:, 1])
+    # Intersection including Non-negative overlap score.
+    segments_intersection = (tt2 - tt1).clip(0)
+    # Segment union.
+    segments_union = (candidate_segments[:, 1] - candidate_segments[:, 0]) \
+      + (target_segment[1] - target_segment[0]) - segments_intersection
+    # Compute overlap as the ratio of the intersection
+    # over union of two segments.
+    tIoU = segments_intersection.astype(float) / segments_union
+    return tIoU
+
+
+def interpolated_prec_rec(prec, rec):
+    """Interpolated AP - VOCdevkit from VOC 2011.
+    """
+    mprec = np.hstack([[0], prec, [0]])
+    mrec = np.hstack([[0], rec, [1]])
+    for i in range(len(mprec) - 1)[::-1]:
+        mprec[i] = max(mprec[i], mprec[i + 1])
+    idx = np.where(mrec[1::] != mrec[0:-1])[0] + 1
+    ap = np.sum((mrec[idx] - mrec[idx - 1]) * mprec[idx])
+    return ap
+
+
+from tqdm import tqdm
+def compute_average_precision_detection(
+        ground_truth,
+        prediction,
+        tiou_thresholds=np.linspace(0.5, 0.95, 10),
+    ):
+    """Compute average precision (detection task) between ground truth and
+    predictions data frames. If multiple predictions occurs for the same
+    predicted segment, only the one with highest score is matches as
+    true positive. This code is greatly inspired by Pascal VOC devkit.
+
+    Ref: https://github.com/zhang-can/CoLA/blob/\
+        d21f1b5a4c6c13f9715cfd4ac1ebcd065d179157/eval/eval_detection.py#L200
+
+    Parameters
+    ----------
+    ground_truth : df
+        Data frame containing the ground truth instances.
+        Required fields: ['video-id', 't-start', 't-end']
+    prediction : df
+        Data frame containing the prediction instances.
+        Required fields: ['video-id, 't-start', 't-end', 'score']
+    tiou_thresholds : 1darray, optional
+        Temporal intersection over union threshold.
+    Outputs
+    -------
+    ap : float
+        Average precision score.
+    """
+    ap = np.zeros(len(tiou_thresholds))
+    if prediction.empty:
+        return ap
+
+    npos = float(len(ground_truth))
+    lock_gt = np.ones((len(tiou_thresholds),len(ground_truth))) * -1
+    # Sort predictions by decreasing score order.
+    sort_idx = prediction['score'].values.argsort()[::-1]
+    prediction = prediction.loc[sort_idx].reset_index(drop=True)
+
+    # Initialize true positive and false positive vectors.
+    tp = np.zeros((len(tiou_thresholds), len(prediction)))
+    fp = np.zeros((len(tiou_thresholds), len(prediction)))
+
+    # Adaptation to query faster
+    ground_truth_gbvn = ground_truth.groupby('video-id')
+
+    # Assigning true positive to truly grount truth instances.
+    for idx, this_pred in prediction.iterrows():
+
+        try:
+            # Check if there is at least one ground truth in the video associated.
+            ground_truth_videoid = ground_truth_gbvn.get_group(this_pred['video-id'])
+        except Exception as e:
+            fp[:, idx] = 1
+            continue
+
+        this_gt = ground_truth_videoid.reset_index()
+        tiou_arr = segment_iou(this_pred[['t-start', 't-end']].values,
+                               this_gt[['t-start', 't-end']].values)
+        # We would like to retrieve the predictions with highest tiou score.
+        tiou_sorted_idx = tiou_arr.argsort()[::-1]
+        for tidx, tiou_thr in enumerate(tiou_thresholds):
+            for jdx in tiou_sorted_idx:
+                if tiou_arr[jdx] < tiou_thr:
+                    fp[tidx, idx] = 1
+                    break
+                if lock_gt[tidx, this_gt.loc[jdx]['index']] >= 0:
+                    continue
+                # Assign as true positive after the filters above.
+                tp[tidx, idx] = 1
+                lock_gt[tidx, this_gt.loc[jdx]['index']] = idx
+                break
+
+            if fp[tidx, idx] == 0 and tp[tidx, idx] == 0:
+                fp[tidx, idx] = 1
+
+    tp_cumsum = np.cumsum(tp, axis=1).astype(float)
+    fp_cumsum = np.cumsum(fp, axis=1).astype(float)
+    recall_cumsum = tp_cumsum / npos
+
+    precision_cumsum = tp_cumsum / (tp_cumsum + fp_cumsum)
+
+    for tidx in range(len(tiou_thresholds)):
+        ap[tidx] = interpolated_prec_rec(precision_cumsum[tidx,:], recall_cumsum[tidx,:])
+
+
+    return ap
+
+
+def ap_wrapper(
+        true_clips,
+        pred_clips,
+        pred_scores,
+        tiou_thresholds=np.linspace(0.5, 0.95, 10),
+    ):
+    assert isinstance(true_clips, np.ndarray)
+    assert len(true_clips.shape) == 2 and true_clips.shape[1] == 2
+    assert isinstance(pred_clips, np.ndarray)
+    assert len(pred_clips.shape) == 2 and pred_clips.shape[1] == 2
+    assert isinstance(pred_scores, np.ndarray)
+    assert len(pred_scores.shape) == 1 and len(pred_scores) == pred_clips.shape[0]
+    
+    true_df = pd.DataFrame(
+        {
+            "video-id": ["video1"] * len(true_clips),
+            "t-start": true_clips[:, 0],
+            "t-end": true_clips[:, 1],
+        }
+    )
+    pred_df = pd.DataFrame(
+        {
+            "video-id": ["video1"] * len(pred_clips),
+            "t-start": pred_clips[:, 0],
+            "t-end": pred_clips[:, 1],
+            "score": pred_scores,
+        }
+    )
+    return compute_average_precision_detection(
+        true_df,
+        pred_df,
+        tiou_thresholds=tiou_thresholds,
+    )
+
+
+def nms_1d(df: pd.DataFrame, score_col="score", iou_thresh=0.5):
+    """Applies NMS on 1D (start, end) box predictions."""
+    columns = set(df.columns)
+    # assert columns == set(["video_id", "start", "end", "score"])
+    assert set(["start", "end", "video_id", score_col]).issubset(columns)
+    video_ids = df["video_id"].unique()
+    
+    # Group by video_id
+    groups = df.groupby("video_id")
+    
+    # Loop over videos
+    keep_indices = []
+    net_success_fraction = []
+    tqdm._instances.clear()
+    iterator = tqdm(
+        video_ids,
+        desc="Applying NMS to each video",
+        bar_format='{l_bar}{bar:10}{r_bar}{bar:-10b}',
+    )
+    for video_id in iterator:
+
+        # Get rows for this video
+        rows = groups.get_group(video_id)
+
+        # Sort by score
+        rows = rows.sort_values(score_col, ascending=False)
+        
+        # Loop over rows until empty
+        n_clips = len(rows)
+        n_clips_selected_in_video = 0
+        while len(rows):
+            
+            # Add top row to keep_indices
+            top_row = rows.iloc[0]
+            keep_indices.append(rows.index[0])
+            n_clips_selected_in_video += 1
+            top_row = top_row.to_dict()
+            
+            top_segment = np.array([top_row["start"], top_row["end"]])
+            rows = rows.iloc[1:]
+            other_segments = rows[["start", "end"]].values
+            iou_values = segment_iou(top_segment, other_segments)
+            
+            # Remove rows IoU > iou_thresh
+            rows = rows[iou_values < iou_thresh]
+
+        net_success_fraction.append(n_clips_selected_in_video / n_clips)
+    net_success_fraction = np.array(net_success_fraction).mean()
+    print("> Net success fraction: {:.2f}".format(net_success_fraction))
+    
+    return keep_indices
+
+
+if __name__ == "__main__":
+    true_clips = np.array(
+        [
+            [0.1, 0.7],
+            [3.4, 7.8],
+            [3.9, 5.4],
+        ]
+    )
+    pred_clips = np.array(
+        [
+            [0.2, 0.8],
+            [3.5, 7.9],
+            [3.9, 5.4],
+            [5.6, 6.7],
+            [6.0, 6.5],
+        ],
+    )
+    pred_scores = np.array([0.9, 0.8, 0.7, 0.6, 0.5])
+    
+    # 1. Check IoU for a single pair of boxes
+    iou = compute_iou_1d(pred_clips[0], true_clips[0])
+    # Manually check that the result is correct
+    # Clips are [0.1, 0.7] and [0.2, 0.8]
+    # Intersection: [0.2, 0.7] - length = 0.5
+    # Union: [0.1, 0.8] - length = 0.7
+    # Ratio: 0.5 / 0.7 = 0.714
+    assert np.isclose(iou, 0.714, 3), "Incorrect IoU"
+    
+    # 2. Check IoU for a single predicted box and multiple ground truth boxes
+    ious = compute_iou_1d_single_candidate_multiple_targets(pred_clips[0], true_clips)
+    assert np.allclose(ious, [0.714, 0.0, 0.0], 3), "Incorrect IoU"
+    
+    # 3. Check mean IoU for multiple predicted boxes and multiple ground truth boxes
+    ious = compute_iou_1d_multiple_candidates_multiple_targets(pred_clips, true_clips)
+    assert ious.shape == (5, 3), "Incorrect shape"
+    
+    ap = ap_wrapper(
+        true_clips,
+        pred_clips,
+        pred_scores,
+        tiou_thresholds=np.linspace(0.5, 0.95, 3),
+    )
+    # Take the mean of the APs across IoU thresholds
+    final_ap = np.mean(ap)
+    import ipdb; ipdb.set_trace()

shared/utils/misc.py

@@ -0,0 +1,116 @@
+"""Misc utils."""
+import os
+from shared.utils.log import tqdm_iterator
+import numpy as np
+
+
+class AttrDict(dict):
+    def __init__(self, *args, **kwargs):
+        super(AttrDict, self).__init__(*args, **kwargs)
+        self.__dict__ = self
+        
+
+def ignore_warnings(type="ignore"):
+    import warnings
+    warnings.filterwarnings(type)
+    os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+
+def download_youtube_video(youtube_id, ext='mp4', resolution="360p", **kwargs):
+    import pytube
+    video_url = f"https://www.youtube.com/watch?v={youtube_id}"
+    yt = pytube.YouTube(video_url)
+    try:
+        streams = yt.streams.filter(
+            file_extension=ext, res=resolution, progressive=True, **kwargs,
+        )
+        # streams[0].download(output_path=save_dir, filename=f"{video_id}.{ext}")
+        streams[0].download(output_path='/tmp', filename='sample.mp4')
+    except:
+        print("Failed to download video: ", video_url)
+        return None
+    return "/tmp/sample.mp4"
+
+
+def check_audio(video_path):
+    from moviepy.video.io.VideoFileClip import VideoFileClip
+    try:
+        return VideoFileClip(video_path).audio is not None
+    except:
+        return False
+
+
+def check_audio_multiple(video_paths, n_jobs=8):
+    """Parallelly check if videos have audio"""
+    iterator = tqdm_iterator(video_paths, desc="Checking audio")
+    from joblib import Parallel, delayed
+    return Parallel(n_jobs=n_jobs)(
+            delayed(check_audio)(video_path) for video_path in iterator
+        )
+
+
+def num_trainable_params(model, round=3, verbose=True, return_count=False):
+    n_params = sum([p.numel() for p in model.parameters() if p.requires_grad])
+    model_name = model.__class__.__name__
+    if round is not None:
+        value = np.round(n_params / 1e6, round)
+        unit = "M"
+    else:
+        value = n_params
+        unit = ""
+    if verbose:
+        print(f"::: Number of trainable parameters in {model_name}: {value} {unit}")
+    if return_count:
+        return n_params
+
+
+def num_params(model, round=3):
+    n_params = sum([p.numel() for p in model.parameters()])
+    model_name = model.__class__.__name__
+    if round is not None:
+        value = np.round(n_params / 1e6, round)
+        unit = "M"
+    else:
+        value = n_params
+        unit = ""
+    print(f"::: Number of total parameters in {model_name}: {value}{unit}")
+
+
+def fix_seed(seed=42):
+    """Fix all numpy/pytorch/random seeds."""
+    import random
+    import torch
+    import numpy as np
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+
+
+def check_tensor(x):
+    print(x.shape, x.min(), x.max())
+
+
+def find_nearest_indices(a, b):
+    """
+    Finds the indices of the elements in `a` that are closest to each element in `b`.
+
+    Args:
+        a (np.ndarray): The array to search for the closest values.
+        b (np.ndarray): The array of values to search for.
+    
+    Returns:
+        np.ndarray: The indices of the closest values in `a` for each element in `b`.
+    """
+    # Reshape `a` and `b` to make use of broadcasting
+    a = np.array(a)
+    b = np.array(b)
+
+    # Calculate the absolute difference between each element in `b` and all elements in `a`
+    diff = np.abs(a - b[:, np.newaxis])
+
+    # Find the index of the minimum value along the second axis (which corresponds to `a`)
+    indices = np.argmin(diff, axis=1)
+
+    return indices

shared/utils/pandas_utils.py

@@ -0,0 +1,117 @@
+"""Utility functions for pandas operations"""
+
+from typing import List
+import numpy as np
+import pandas as pd
+
+
+def apply_filters(df: pd.DataFrame, filters: dict, reset_index=False):
+    """
+    Filters df based on given filters (key-values pairs).
+    """
+    import omegaconf
+    X = df.copy()
+
+    all_indices = []
+    for col, values in filters.items():
+        if isinstance(values, (list, tuple, np.ndarray, omegaconf.listconfig.ListConfig)):
+            indices = X[col].isin(list(values))
+        else:
+            indices = X[col] == values
+        all_indices.append(indices)
+        # print(col, values, len(indices), sum(indices))
+        # X = X[indices]
+    if len(all_indices):
+        all_indices = np.array(all_indices)
+        indices = np.all(all_indices, axis=0)
+        X = X[indices]
+
+    if reset_index:
+        X = X.reset_index(drop=True)
+
+    return X
+
+
+def apply_antifilters(df: pd.DataFrame, filters: dict, reset_index=False):
+    """
+    Filters df removing rows for given filters (key-values pairs).
+    """
+    X = df.copy()
+
+    for col, values in filters.items():
+        if isinstance(values, (list, tuple, np.ndarray)):
+            indices = X[col].isin(list(values))
+        else:
+            indices = X[col] == values
+        X = X[~indices]
+
+    if reset_index:
+        X = X.reset_index(drop=True)
+
+    return X
+
+
+def custom_eval(x):
+    """Splits string '["a", "b", "c"]' into ["a", "b", "c"]."""
+    if isinstance(x, str):
+        x = x.replace('[', '')
+        x = x.replace(']', '')
+
+        x = x.split(',')
+        x = [y.rstrip().lstrip() for y in x]
+        return x
+    else:
+        return ['NA']
+
+
+def split_column_into_columns(df, column):
+    """
+    For given df, splits `column` containing values like '["a", "b"]'
+    into one-hot subcolumns like a. b with `Yes`/`No` values.
+    """
+    df[column] = df[column].apply(custom_eval)
+
+    unique_values = []
+    for i in range(len(df)):
+        index = df.index[i]
+
+        list_of_values = df.loc[index, column]
+
+        for x in list_of_values:
+            if (x != 'NA') and (x != ''):
+                df.at[index, x] = 'Yes'
+                if x not in unique_values:
+                    unique_values.append(x)
+
+    df[unique_values] = df[unique_values].fillna('No')
+    df[f'any_{column}'] = df[unique_values].apply(
+        lambda x: 'Yes' if 'Yes' in list(x) else 'No', axis=1
+    )
+    return df
+
+
+def custom_read_csv(path: str, columns_to_onehot: List) -> pd.DataFrame:
+    """Custom CSV reader
+
+    Args:
+        path (str): path to .csv file
+        columns_to_onehot (List): list of columns to one-hotify
+
+    Returns:
+        pd.DataFrame: loaded df
+    """
+    df = pd.read_csv(path)
+    for column in columns_to_onehot:
+        df = split_column_into_columns(df, column)
+    return df
+
+
+def split_df(df, test_size=0.2):
+    from sklearn.model_selection import train_test_split
+    # split the dataframe into train and test sets
+    train_df, test_df = train_test_split(df, test_size=test_size, random_state=42)
+
+    # split the train set into train and validation sets
+    train_df, val_df = train_test_split(train_df, test_size=test_size, random_state=42)
+    
+    return train_df, val_df, test_df

shared/utils/paths.py

@@ -0,0 +1,35 @@
+"""Path utils."""
+from os.path import dirname, abspath
+
+
+curr_filepath = abspath(__file__)
+repo_path = dirname(dirname(dirname(curr_filepath)))
+
+
+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/",
+        "vggdev21": "/scratch/shared/beegfs/piyush/datasets/",
+        "node407": "/var/scratch/pbagad/datasets/",
+        "gnodee5": "/scratch/shared/beegfs/piyush/datasets/",
+        "gnodeg2": "/scratch/shared/beegfs/piyush/datasets/",
+        "gnodec2": "/scratch/shared/beegfs/piyush/datasets/",
+        "Piyushs-MacBook-Pro": "/Users/piyush/projects/",
+        "gnodec1": "/scratch/shared/beegfs/piyush/datasets/",
+        "gnodec5": "/scratch/shared/beegfs/piyush/datasets/",
+        "gnodec4": "/scratch/shared/beegfs/piyush/datasets/",
+        "gnoded2": "/scratch/shared/beegfs/piyush/datasets/",
+    }
+    hostname = socket.gethostname()
+    hostname = hostname.split(".")[0]
+    
+    assert hostname in data_root_lib.keys(), \
+        "Hostname {} not in data_root_lib".format(hostname)
+
+    data_root = data_root_lib[hostname]
+    return data_root
+    

shared/utils/physics.py

@@ -0,0 +1,341 @@
+import numpy as np
+import torch
+
+
+# Universal constants
+C = 340. * 100.  # Speed of sound in air (cm/s)
+
+
+def compute_length_of_air_column_cylindrical(
+        timestamps, duration, height, b, **kwargs,
+    ):
+    """
+    Randomly chooses a l(t) curve satisfying the two point equations.
+    """
+    L = height * ( (1 - np.exp(b * (duration - timestamps))) / (1 - np.exp(b * duration)) )
+    return L
+
+
+def compute_axial_frequency_cylindrical(
+        lengths, radius, beta=0.62, mode=1, **kwargs,
+    ):
+    """
+    Computes axial resonance frequency for cylindrical container at given timestamps.
+    """
+    if mode == 1:
+        harmonic_weight = 1.
+    elif mode == 2:
+        harmonic_weight = 3.
+    elif mode == 3:
+        harmonic_weight = 5.
+    else:
+        raise ValueError
+
+    # Compute fundamental frequency curve
+    F0 = harmonic_weight * (0.25 * C) * (1. / (lengths + (beta * radius)))
+
+    return F0
+
+
+def compute_axial_frequency_bottleneck(
+        lengths, radius, height, Rn, Hn, beta_bottle=(0.6 + 8/np.pi), **kwargs,
+    ):
+    # Here, R and H are base radius and height of the bottleneck
+    eps = 1e-6
+    kappa = (0.5 * C / np.pi) * (Rn/radius) * np.sqrt(1 / (Hn + beta_bottle * Rn))
+    frequencies = kappa * np.sqrt(1 / (lengths + eps))
+    return frequencies
+
+
+def compute_f0_cylindrical(Y, rho_g, a, R, H, mode=1, **kwargs,):
+
+    if mode == 1:
+        m = 1.875
+        n = 2
+    elif mode == 2:
+        m = 4.694
+        n = 3
+    elif mode == 3:
+        m = 7.855
+        n = 4
+    else:
+        raise ValueError
+
+    term = ( ((n**2 - 1)**2) + ((m * R/H)**4) ) / (1 + (1./n**2))
+    f0 = (1. / (12 * np.pi)) * np.sqrt(3 * Y / rho_g) * (a / (R**2)) * np.sqrt(term)
+    return f0
+
+
+def compute_xi_cylindrical(rho_l, rho_g, R, a, **kwargs,):
+    """
+    Different papers use different multipliers.
+    For us, using 12. * (4./9.) works best empirically.
+    """
+    xi = 12. * (4. / 9.) * (rho_l/rho_g) * (R/a)
+    return xi
+
+
+def compute_radial_frequency_cylindrical(
+        heights, R, H, Y, rho_g, a, rho_l, power=3, mode=1, **kwargs,
+    ):
+    """
+    Computes radial resonance frequency for cylindrical.
+
+    Args:
+        heights (np.ndarray): height of liquid at pre-defined time stamps
+    """
+    # Only f0 changes for higher modes
+    f0 = compute_f0_cylindrical(Y, rho_g, a, R, H, mode=mode)
+    xi = compute_xi_cylindrical(rho_l, rho_g, R, a)
+    frequencies = f0 / np.sqrt(1 + xi * ((heights/H) ** power) )
+    return frequencies
+
+
+def compute_slant_lengths_semiconical(
+        timestamps, duration, r_top, r_bot, height, **kwargs,
+    ):
+
+    # Top radius / base radius
+    rf =  r_bot / r_top
+
+    # Time fraction
+    tf = timestamps/duration
+    
+    # Height fractions: h(t) / H
+    height_fractions = (1. / (rf - 1)) * (np.cbrt(((rf**3 - 1) * (tf)) + 1) - 1)
+    
+    # Slant air column lengths
+    heights = height_fractions * height
+    slant_lengths = np.sqrt(1 - ((r_top - r_bot) / height)**2) * (height - heights)
+
+    return slant_lengths
+
+
+def compute_axial_frequency_semiconical(slant_lengths, r_top, r_bot, beta=1.28, **kwargs):
+    """
+    Computes axial resonance frequency for cylinder.
+
+    Args:
+        slant_lengths (np.ndarray): slant length of air column
+        r_top (float): top radius
+        r_bot (float): base radius
+        beta (float): end correction coefficient
+    """
+    frequencies_axial = (C / 2) * (1 / (slant_lengths + (beta * (r_bot + r_top))))
+    return frequencies_axial
+
+
+def get_frequencies(
+        t,
+        params,
+        container_shape="cylindrical",
+        harmonic=None,
+        vibration_type="axial",
+        semiconical_as_cylinder=False,
+    ):
+    """
+    Computes requires frequency f(t) for given t.
+    """
+
+    if container_shape == "semiconical":
+        # Makes an assumption that semiconical shape is similar to cylindrical
+        if semiconical_as_cylinder:
+            container_shape = "cylindrical"
+    
+    if (container_shape == "cylindrical") or (container_shape == "bottleneck_as_cylindrical"):
+
+        # Compute length of air column first
+        lengths = compute_length_of_air_column_cylindrical(t, **params)
+
+        if vibration_type == "axial":
+            frequencies = compute_axial_frequency_cylindrical(lengths, **params)
+
+            if harmonic is not None:
+                assert harmonic > 0 and isinstance(harmonic, int)
+                frequencies = frequencies * harmonic
+
+        elif vibration_type == "radial":
+            if harmonic is None:
+                mode = 1
+            else:
+                assert isinstance(harmonic, int)
+                assert harmonic in [1, 2]
+                mode = harmonic + 1
+            frequencies = compute_radial_frequency_cylindrical(
+                lengths, mode=mode, **params,
+            )
+
+        else:
+            raise NotImplementedError
+
+    elif container_shape == "semiconical":
+
+            # Compute length of air column first
+            slant_lengths = compute_slant_lengths_semiconical(t, **params)
+
+            if vibration_type == "axial":
+                frequencies = compute_axial_frequency_semiconical(
+                    slant_lengths, **params,
+                )
+    
+                if harmonic is not None:
+                    assert harmonic > 0 and isinstance(harmonic, int)
+                    frequencies = frequencies * harmonic
+
+            else:
+                raise NotImplementedError
+
+    elif container_shape == "bottleneck":
+
+        # Compute length of air column first assuming 
+        # base of the bottle is a cylindrical
+        lengths = compute_length_of_air_column_cylindrical(t, **params)
+
+        if vibration_type == "axial":
+            frequencies = compute_axial_frequency_bottleneck(
+                lengths, **params,
+            )
+
+            if harmonic is not None:
+                assert harmonic > 0 and isinstance(harmonic, int)
+                frequencies = frequencies * harmonic
+        else:
+            raise NotImplementedError
+
+    else:
+        raise ValueError
+
+    return frequencies
+
+
+def get_params(row, semiconical_as_cylinder=False):
+    m = row["measurements"]
+    duration = row["end_time"] - row["start_time"]
+    params = dict(duration=duration)
+    if row["shape"] == "cylindrical":
+        radius = 0.25 * (m["diameter_top"] + m["diameter_bottom"])
+        height = m["net_height"]
+        params.update(
+            height=height,
+            radius=radius,
+            beta=row.get("beta", 0.62),
+            # Constant flow
+            b=0.01,
+        )
+    elif row["shape"] == "semiconical":
+
+        if semiconical_as_cylinder:
+            # Assume semiconical shape as cylindrical
+            radius = 0.25 * (m["diameter_top"] + m["diameter_bottom"])
+            height = m["net_height"]
+            params.update(
+                height=height,
+                radius=radius,
+                beta=0.62,
+                # Constant flow
+                b=0.01,
+            )
+        else:
+            r_top = 0.5 * m["diameter_top"]
+            r_bot = 0.5 * m["diameter_bottom"]
+            height = m["net_height"]
+            beta = 1.28
+            params.update(
+                r_top=r_top,
+                r_bot=r_bot,
+                height=height,
+                beta=beta,
+            )
+    elif row["shape"] == "bottleneck":
+        radius = 0.5 * m["diameter_bottom"]
+        Rn = 0.5 * m["diameter_top"]
+        Hn = m["neck_height"] 
+        height = m["net_height"] - Hn
+        params.update(
+            height=height,
+            radius=radius,
+            Rn=Rn,
+            Hn=Hn,
+            # Constant flow
+            b=0.01,
+        )
+    elif row["shape"] == "bottleneck_as_cylindrical":
+        # Approximates bottleneck as cylindrical
+        radius = 0.5 * m["diameter_bottom"]
+        height = m["net_height"] + m["neck_height"]
+        params.update(
+            height=height,
+            radius=radius,
+            beta=row.get("beta", 0.62),
+            # Constant flow
+            b=0.01,
+        )
+    else:
+        raise ValueError
+    return params
+
+def frequency_to_wavelength(f):
+    """
+    Converts frequency to wavelength.
+
+    Args:
+        f (float): frequency
+    """
+    return C / f
+
+
+def wavelength_to_frequency(l):
+    """
+    Converts wavelength to frequency.
+
+    Args:
+        l (float): wavelength
+    """
+    return C / l
+
+
+def get_cylinder_radius(m):
+    return 0.25 * (m['diameter_top'] + m['diameter_bottom'])
+
+
+def get_cylinder_height(m):
+    return m['net_height']
+
+
+def get_flow_rate(m, duration):
+    r = get_cylinder_radius(m)
+    h = get_cylinder_height(m)
+    volume = np.pi * (r**2) * h
+    q = volume / duration
+    return q
+
+
+def get_length_of_air_column(m, duration, timestamps):
+    h = get_cylinder_height(m)
+    l = (-h/duration) * timestamps + h
+    l = torch.from_numpy(l)
+    return l
+
+
+def estimate_cylinder_radius(wavelengths, timestamps=None, beta=0.62):
+    radius_pred = ((1. / beta) * (wavelengths[-1] / 4.)).item()
+    return radius_pred
+
+
+def estimate_cylinder_height(wavelengths, timestamps=None, beta=0.62):
+    height_pred = wavelengths[0] / 4. - wavelengths[-1] / 4.
+    return height_pred.item()
+
+
+def estimate_flow_rate(wavelengths, timestamps=None, output_fps=49.):
+    radius = estimate_cylinder_radius(wavelengths)
+    l_pred = (wavelengths - wavelengths[-1]) / 4.
+    slope = np.gradient(l_pred).mean() * output_fps
+    Q_pred = -np.pi * (radius**2) * slope
+    return Q_pred
+
+
+def estimate_length_of_air_column(wavelengths, timestamps=None):
+    l_pred = (wavelengths - wavelengths[-1]) / 4.
+    return l_pred
+

shared/utils/text_basic.py

@@ -0,0 +1,44 @@
+"""Utils for processing and encoding text."""
+
+import torch
+
+
+
+def lemmatize_verbs(verbs: list):
+    from nltk.stem import WordNetLemmatizer
+    wnl = WordNetLemmatizer()
+    return [wnl.lemmatize(verb, 'v') for verb in verbs]
+
+
+def lemmatize_adverbs(adverbs: list):
+    from nltk.stem import WordNetLemmatizer
+    wnl = WordNetLemmatizer()
+    return [wnl.lemmatize(adverb, 'r') for adverb in adverbs]
+
+
+class SentenceEncoder:
+
+    def __init__(self, model_name="roberta-base"):
+        from transformers import RobertaTokenizer, RobertaModel
+        if model_name == 'roberta-base':
+            self.tokenizer = RobertaTokenizer.from_pretrained(model_name)
+            self.model = RobertaModel.from_pretrained(model_name)
+    
+    def encode_sentence(self, sentence):
+        inputs = self.tokenizer.encode_plus(
+            sentence, add_special_tokens=True, return_tensors='pt',
+        )
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+        # sentence_embedding = torch.mean(outputs.last_hidden_state, dim=1).squeeze(0)
+        sentence_embedding = outputs.last_hidden_state[:, 0, :]
+        return sentence_embedding
+    
+    def encode_sentences(self, sentences):
+        """Encodes a list of sentences using model."""
+        tokenized_input = self.tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')
+        with torch.no_grad():
+            outputs = self.model(**tokenized_input)
+        embeddings = outputs.last_hidden_state[:, 0, :]
+        return embeddings
+    

shared/utils/visualize.py

@@ -0,0 +1,2208 @@
+"""Helpers for visualization"""
+import os
+import numpy as np
+import matplotlib
+import matplotlib.pyplot as plt
+import cv2
+import PIL
+from PIL import Image, ImageOps, ImageDraw
+from os.path import exists
+import librosa.display
+import pandas as pd
+import itertools
+import librosa
+from tqdm import tqdm
+from IPython.display import Audio, Markdown, display
+from ipywidgets import Button, HBox, VBox, Text, Label, HTML, widgets
+from shared.utils.log import tqdm_iterator
+
+import warnings
+warnings.filterwarnings("ignore")
+
+try:
+    import torchvideotransforms
+except:
+    print("Failed to import torchvideotransforms. Proceeding without.")
+    print("Please install using:")
+    print("pip install git+https://github.com/hassony2/torch_videovision")
+
+
+# define predominanat colors
+COLORS = {
+    "pink": (242, 116, 223),
+    "cyan": (46, 242, 203),
+    "red": (255, 0, 0),
+    "green": (0, 255, 0),
+    "blue": (0, 0, 255),
+    "yellow": (255, 255, 0),
+}
+
+
+def get_predominant_color(color_key, mode="RGB", alpha=0):
+    assert color_key in COLORS.keys(), f"Unknown color key: {color_key}"
+    if mode == "RGB":
+        return COLORS[color_key]
+    elif mode == "RGBA":
+        return COLORS[color_key] + (alpha,)
+
+
+def show_single_image(image: np.ndarray, figsize: tuple = (8, 8), title: str = None, cmap: str = None, ticks=False):
+    """Show a single image."""
+    fig, ax = plt.subplots(1, 1, figsize=figsize)
+
+    if isinstance(image, Image.Image):
+        image = np.asarray(image)
+
+    ax.set_title(title)
+    ax.imshow(image, cmap=cmap)
+    
+    if not ticks:
+        ax.set_xticks([])
+        ax.set_yticks([])
+
+    plt.show()
+
+
+def show_grid_of_images(
+        images: np.ndarray, n_cols: int = 4, figsize: tuple = (8, 8), subtitlesize=14,
+        cmap=None, subtitles=None, title=None, save=False, savepath="sample.png", titlesize=20,
+        ysuptitle=0.8, xlabels=None, sizealpha=0.7, show=True, row_labels=None, aspect=None,
+    ):
+    """Show a grid of images."""
+    n_cols = min(n_cols, len(images))
+
+    copy_of_images = images.copy()
+    for i, image in enumerate(copy_of_images):
+        if isinstance(image, Image.Image):
+            image = np.asarray(image)
+            copy_of_images[i] = image
+
+    if subtitles is None:
+        subtitles = [None] * len(images)
+
+    if xlabels is None:
+        xlabels = [None] * len(images)
+    
+    if row_labels is None:
+        num_rows = int(np.ceil(len(images) / n_cols))
+        row_labels = [None] * num_rows
+
+    n_rows = int(np.ceil(len(images) / n_cols))
+    fig, axes = plt.subplots(n_rows, n_cols, figsize=figsize)
+    if len(images) == 1:
+        axes = np.array([[axes]])
+    for i, ax in enumerate(axes.flat):
+        if i < len(copy_of_images):
+            if len(copy_of_images[i].shape) == 2 and cmap is None:
+                cmap="gray"
+            ax.imshow(copy_of_images[i], cmap=cmap, aspect=aspect)
+            ax.set_title(subtitles[i], fontsize=subtitlesize)
+        ax.set_xlabel(xlabels[i], fontsize=sizealpha * subtitlesize)
+        ax.set_xticks([])
+        ax.set_yticks([])
+
+        col_idx = i % n_cols
+        if col_idx == 0:
+            ax.set_ylabel(row_labels[i // n_cols], fontsize=sizealpha * subtitlesize)
+
+    fig.tight_layout()
+    plt.suptitle(title, y=ysuptitle, fontsize=titlesize)
+    if save:
+        plt.savefig(savepath, bbox_inches='tight')
+    if show:
+        plt.show()
+
+
+
+def add_text_to_image(image, text):
+    from PIL import ImageFont
+    from PIL import ImageDraw
+    
+    # # resize image
+    # image = image.resize((image.size[0] * 2, image.size[1] * 2))
+
+    draw = ImageDraw.Draw(image)
+    font = ImageFont.load_default()
+    # font = ImageFont.load("arial.pil")
+    # font = ImageFont.FreeTypeFont(size=20)
+    # font = ImageFont.truetype("arial.ttf", 28, encoding="unic")
+
+    # change fontsize
+    
+    # select color = black if image is mostly white
+    if np.mean(image) > 200:
+        draw.text((0, 0), text, (0,0,0), font=font)
+    else:
+        draw.text((0, 0), text, (255,255,255), font=font)
+    
+    # draw.text((0, 0), text, (255,255,255), font=font)
+    return image
+
+
+def show_keypoint_matches(
+        img1, kp1, img2, kp2, matches,
+        K=10, figsize=(10, 5), drawMatches_args=dict(matchesThickness=3, singlePointColor=(0, 0, 0)),
+        choose_matches="random",
+    ):
+    """Displays matches found in the pair of images"""
+    if choose_matches == "random":
+        selected_matches = np.random.choice(matches, K)
+    elif choose_matches == "all":
+        K = len(matches)
+        selected_matches = matches
+    elif choose_matches == "topk":
+        selected_matches = matches[:K]
+    else:
+        raise ValueError(f"Unknown value for choose_matches: {choose_matches}")
+
+    # color each match with a different color
+    cmap = matplotlib.cm.get_cmap('gist_rainbow', K)
+    colors = [[int(x*255) for x in cmap(i)[:3]] for i in np.arange(0,K)]
+    drawMatches_args.update({"matchColor": -1, "singlePointColor": (100, 100, 100)})
+    
+    img3 = cv2.drawMatches(img1, kp1, img2, kp2, selected_matches, outImg=None, **drawMatches_args)
+    show_single_image(
+        img3,
+        figsize=figsize,
+        title=f"[{choose_matches.upper()}] Selected K = {K} matches between the pair of images.",
+    )
+    return img3
+
+
+def draw_kps_on_image(image: np.ndarray, kps: np.ndarray, color=COLORS["red"], radius=3, thickness=-1, return_as="PIL"):
+    """
+    Draw keypoints on image.
+
+    Args:
+        image: Image to draw keypoints on.
+        kps: Keypoints to draw. Note these should be in (x, y) format.
+    """
+    if isinstance(image, Image.Image):
+        image = np.asarray(image)
+    if isinstance(color, str):
+        color = PIL.ImageColor.getrgb(color)
+        colors = [color] * len(kps)
+    elif isinstance(color, tuple):
+        colors = [color] * len(kps)
+    elif isinstance(color, list):
+        colors = [PIL.ImageColor.getrgb(c) for c in color]
+    assert len(colors) == len(kps), f"Number of colors ({len(colors)}) must be equal to number of keypoints ({len(kps)})"
+
+    for kp, c in zip(kps, colors):
+        image = cv2.circle(
+            image.copy(), (int(kp[0]), int(kp[1])), radius=radius, color=c, thickness=thickness)
+    
+    if return_as == "PIL":
+        return Image.fromarray(image)
+
+    return image
+
+
+def get_concat_h(im1, im2):
+    """Concatenate two images horizontally"""
+    dst = Image.new('RGB', (im1.width + im2.width, im1.height))
+    dst.paste(im1, (0, 0))
+    dst.paste(im2, (im1.width, 0))
+    return dst
+
+
+def get_concat_v(im1, im2):
+    """Concatenate two images vertically"""
+    dst = Image.new('RGB', (im1.width, im1.height + im2.height))
+    dst.paste(im1, (0, 0))
+    dst.paste(im2, (0, im1.height))
+    return dst
+
+
+def show_images_with_keypoints(images: list, kps: list, radius=15, color=(0, 220, 220), figsize=(10, 8)):
+    assert len(images) == len(kps)
+
+    # generate
+    images_with_kps = []
+    for i in range(len(images)):
+        img_with_kps = draw_kps_on_image(images[i], kps[i], radius=radius, color=color, return_as="PIL")
+        images_with_kps.append(img_with_kps)
+    
+    # show
+    show_grid_of_images(images_with_kps, n_cols=len(images), figsize=figsize)
+
+
+def set_latex_fonts(usetex=True, fontsize=14, show_sample=False, **kwargs):
+    try:
+        plt.rcParams.update({
+            "text.usetex": usetex,
+            "font.family": "serif",
+            # "font.serif": ["Computer Modern Romans"],
+            "font.size": fontsize,
+            **kwargs,
+        })
+        if show_sample:
+            plt.figure()
+            plt.title("Sample $y = x^2$")
+            plt.plot(np.arange(0, 10), np.arange(0, 10)**2, "--o")
+            plt.grid()
+            plt.show()
+    except:
+        print("Failed to setup LaTeX fonts. Proceeding without.")
+        pass
+
+
+
+def plot_2d_points(
+        list_of_points_2d,
+        colors=None,
+        sizes=None,
+        markers=None,
+        alpha=0.75,
+        h=256,
+        w=256,
+        ax=None,
+        save=True,
+        savepath="test.png",
+    ):
+
+    if ax is None:
+        fig, ax = plt.subplots(1, 1)
+    ax.set_xlim([0, w])
+    ax.set_ylim([0, h])
+    
+    if sizes is None:
+        sizes = [0.1 for _ in range(len(list_of_points_2d))]
+    if colors is None:
+        colors = ["gray" for _ in range(len(list_of_points_2d))]
+    if markers is None:
+        markers = ["o" for _ in range(len(list_of_points_2d))]
+
+    for points_2d, color, s, m in zip(list_of_points_2d, colors, sizes, markers):
+        ax.scatter(points_2d[:, 0], points_2d[:, 1], s=s, alpha=alpha, color=color, marker=m)
+    
+    if save:
+        plt.savefig(savepath, bbox_inches='tight')
+
+
+def plot_2d_points_on_image(
+        image,
+        img_alpha=1.0,
+        ax=None,
+        list_of_points_2d=[],
+        scatter_args=dict(),
+    ):
+    if ax is None:
+        fig, ax = plt.subplots(1, 1)
+    ax.imshow(image, alpha=img_alpha)
+    scatter_args["save"] = False
+    plot_2d_points(list_of_points_2d, ax=ax, **scatter_args)
+    
+    # invert the axis
+    ax.set_ylim(ax.get_ylim()[::-1])
+
+
+def compare_landmarks(
+        image, ground_truth_landmarks, v2d, predicted_landmarks,
+        save=False, savepath="compare_landmarks.png", num_kps_to_show=-1,
+        show_matches=True,
+    ):
+
+    # show GT landmarks on image
+    fig, axes = plt.subplots(1, 3, figsize=(11, 4))
+    ax = axes[0]
+    plot_2d_points_on_image(
+        image,
+        list_of_points_2d=[ground_truth_landmarks],
+        scatter_args=dict(sizes=[15], colors=["limegreen"]),
+        ax=ax,
+    )
+    ax.set_title("GT landmarks", fontsize=12)
+    
+    # since the projected points are inverted, using 180 degree rotation about z-axis
+    ax = axes[1]
+    plot_2d_points_on_image(
+        image,
+        list_of_points_2d=[v2d, predicted_landmarks],
+        scatter_args=dict(sizes=[0.08, 15], markers=["o", "x"], colors=["royalblue", "red"]),
+        ax=ax,
+    )
+    ax.set_title("Projection of predicted mesh", fontsize=12)
+    
+    # plot the ground truth and predicted landmarks on the same image
+    ax = axes[2]
+    plot_2d_points_on_image(
+        image,
+        list_of_points_2d=[
+            ground_truth_landmarks[:num_kps_to_show],
+            predicted_landmarks[:num_kps_to_show],
+        ],
+        scatter_args=dict(sizes=[15, 15], markers=["o", "x"], colors=["limegreen", "red"]),
+        ax=ax,
+        img_alpha=0.5,
+    )
+    ax.set_title("GT and predicted landmarks", fontsize=12)
+
+    if show_matches:
+        for i in range(num_kps_to_show):
+            x_values = [ground_truth_landmarks[i, 0], predicted_landmarks[i, 0]]
+            y_values = [ground_truth_landmarks[i, 1], predicted_landmarks[i, 1]]
+            ax.plot(x_values, y_values, color="yellow", markersize=1, linewidth=2.)
+
+    fig.tight_layout()
+    if save:
+        plt.savefig(savepath, bbox_inches="tight")
+        
+
+
+def plot_historgam_values(
+        X, display_vals=False,
+        bins=50, figsize=(8, 5),
+        show_mean=True,
+        xlabel=None, ylabel=None,
+        ax=None, title=None, show=False,
+        **kwargs,
+    ):
+    if ax is None:
+        fig, ax = plt.subplots(1, 1, figsize=figsize)
+
+    ax.hist(X, bins=bins, **kwargs)
+    if title is None:
+        title = "Histogram of values"
+    
+    ax.set_xlabel(xlabel)
+    ax.set_ylabel(ylabel)
+    
+    if display_vals:
+        x, counts = np.unique(X, return_counts=True)
+        # sort_indices = np.argsort(x)
+        # x = x[sort_indices]
+        # counts = counts[sort_indices]
+        # for i in range(len(x)):
+        #     ax.text(x[i], counts[i], counts[i], ha='center', va='bottom')
+    
+    ax.grid(alpha=0.3)
+    
+    if show_mean:
+        mean = np.mean(X)
+        mean_string = f"$\mu$: {mean:.2f}"
+        ax.set_title(title + f" ({mean_string}) ")
+    else:
+        ax.set_title(title)
+    
+    if not show:
+        return ax
+    else:
+        plt.show()
+
+
+"""Helper functions for all kinds of 2D/3D visualization"""
+def bokeh_2d_scatter(x, y, desc, figsize=(700, 700), colors=None, use_nb=False, title="Bokeh scatter plot"):
+    import matplotlib.colors as mcolors
+    from bokeh.plotting import figure, output_file, show, ColumnDataSource
+    from bokeh.models import HoverTool
+    from bokeh.io import output_notebook
+
+    if use_nb:
+        output_notebook()
+
+    # define colors to be assigned
+    if colors is None:
+        # applies the same color
+        # create a color iterator: pick a random color and apply it to all points
+        # colors = [np.random.choice(itertools.cycle(palette))] * len(x)
+        colors = [np.random.choice(["red", "green", "blue", "yellow", "pink", "black", "gray"])] * len(x)
+
+        # # applies different colors
+        # colors = np.array([ [r, g, 150] for r, g in zip(50 + 2*x, 30 + 2*y) ], dtype="uint8")
+
+
+    # define the df of data to plot
+    source = ColumnDataSource(
+            data=dict(
+                x=x,
+                y=y,
+                desc=desc,
+                color=colors,
+            )
+        )
+
+    # define the attributes to show on hover
+    hover = HoverTool(
+            tooltips=[
+                ("index", "$index"),
+                ("(x, y)", "($x, $y)"),
+                ("Desc", "@desc"),
+            ]
+        )
+
+    p = figure(
+        plot_width=figsize[0], plot_height=figsize[1], tools=[hover], title=title,
+    )
+    p.circle('x', 'y', size=10, source=source, fill_color="color")
+    show(p)
+
+
+
+
+def bokeh_2d_scatter_new(
+        df, x, y, hue, label, color_column=None, size_col=None,
+        figsize=(700, 700), use_nb=False, title="Bokeh scatter plot",
+        legend_loc="bottom_left", edge_color="black", audio_col=None,
+    ):
+    from bokeh.plotting import figure, output_file, show, ColumnDataSource
+    from bokeh.models import HoverTool
+    from bokeh.io import output_notebook
+
+    if use_nb:
+        output_notebook()
+
+    assert {x, y, hue, label}.issubset(set(df.keys()))
+
+    if isinstance(color_column, str) and color_column in df.keys():
+        color_column_name = color_column
+    else:
+        import matplotlib.colors as mcolors
+        colors = list(mcolors.BASE_COLORS.keys()) + list(mcolors.TABLEAU_COLORS.values())
+        # colors = list(mcolors.BASE_COLORS.keys())
+        colors = itertools.cycle(np.unique(colors))
+
+        hue_to_color = dict()
+        unique_hues = np.unique(df[hue].values)
+        for _hue in unique_hues:
+            hue_to_color[_hue] = next(colors)
+        df["color"] = df[hue].apply(lambda k: hue_to_color[k])
+        color_column_name = "color"
+    
+    if size_col is not None:
+        assert isinstance(size_col, str) and size_col in df.keys()
+    else:
+        sizes = [10.] * len(df)
+        df["size"] = sizes
+        size_col = "size"
+
+    source = ColumnDataSource(
+        dict(
+            x = df[x].values,
+            y = df[y].values,
+            hue = df[hue].values,
+            label = df[label].values,
+            color = df[color_column_name].values,
+            edge_color = [edge_color] * len(df),
+            sizes = df[size_col].values,
+        )
+    )
+
+    # define the attributes to show on hover
+    hover = HoverTool(
+            tooltips=[
+                ("index", "$index"),
+                ("(x, y)", "($x, $y)"),
+                ("Desc", "@label"),
+                ("Cluster", "@hue"),
+            ]
+        )
+
+    p = figure(
+        plot_width=figsize[0],
+        plot_height=figsize[1],
+        tools=["pan","wheel_zoom","box_zoom","save","reset","help"] + [hover],
+        title=title,
+    )
+    p.circle(
+        'x', 'y', size="sizes",
+        source=source, fill_color="color",
+        legend_group="hue", line_color="edge_color",
+    )
+    p.legend.location = legend_loc
+    p.legend.click_policy="hide"
+
+
+    show(p)
+
+    
+import torch
+def get_sentence_embedding(model, tokenizer, sentence):
+    encoded = tokenizer.encode_plus(sentence, return_tensors="pt")
+
+    with torch.no_grad():
+        output = model(**encoded)
+    
+    last_hidden_state = output.last_hidden_state
+    assert last_hidden_state.shape[0] == 1
+    assert last_hidden_state.shape[-1] == 768
+    
+    # only pick the [CLS] token embedding (sentence embedding)
+    sentence_embedding = last_hidden_state[0, 0]
+    
+    return sentence_embedding
+
+
+def lighten_color(color, amount=0.5):
+    """
+    Lightens the given color by multiplying (1-luminosity) by the given amount.
+    Input can be matplotlib color string, hex string, or RGB tuple.
+
+    Examples:
+    >> lighten_color('g', 0.3)
+    >> lighten_color('#F034A3', 0.6)
+    >> lighten_color((.3,.55,.1), 0.5)
+    """
+    import matplotlib.colors as mc
+    import colorsys
+    try:
+        c = mc.cnames[color]
+    except:
+        c = color
+    c = colorsys.rgb_to_hls(*mc.to_rgb(c))
+    return colorsys.hls_to_rgb(c[0], 1 - amount * (1 - c[1]), c[2])
+
+
+def plot_histogram(df, col, ax=None, color="blue", title=None, xlabel=None, **kwargs):
+    if ax is None:
+        fig, ax = plt.subplots(1, 1, figsize=(5, 4))
+    ax.grid(alpha=0.3)
+    xlabel = col if xlabel is None else xlabel
+    ax.set_xlabel(xlabel)
+    ax.set_ylabel("Frequency")
+    title = f"Historgam of {col}" if title is None else title
+    ax.set_title(title)
+    label = f"Mean: {np.round(df[col].mean(), 1)}"
+    ax.hist(df[col].values, density=False, color=color, edgecolor=lighten_color(color, 0.1), label=label, **kwargs)
+    if "bins" in kwargs:
+        xticks = list(np.arange(kwargs["bins"])[::5])
+        xticks += list(np.linspace(xticks[-1], int(df[col].max()), 5, dtype=int))
+        # print(xticks)
+        ax.set_xticks(xticks)
+    ax.legend()
+    plt.show()
+
+
+def beautify_ax(ax, title=None, titlesize=20, sizealpha=0.7, xlabel=None, ylabel=None):
+    labelsize = sizealpha * titlesize
+    ax.grid(alpha=0.3)
+    ax.set_xlabel(xlabel, fontsize=labelsize)
+    ax.set_ylabel(ylabel, fontsize=labelsize)
+    ax.set_title(title, fontsize=titlesize)
+
+
+
+
+def get_text_features(text: list, model, device, batch_size=16):
+    import clip
+    text_batches = [text[i:i+batch_size] for i in range(0, len(text), batch_size)]
+    text_features = []
+    model = model.to(device)
+    model = model.eval()
+    for batch in tqdm(text_batches, desc="Getting text features", bar_format="{l_bar}{bar:20}{r_bar}"):
+        batch = clip.tokenize(batch).to(device)
+        with torch.no_grad():
+            batch_features = model.encode_text(batch)
+        text_features.append(batch_features.cpu().numpy())
+    text_features = np.concatenate(text_features, axis=0)
+    return text_features
+
+
+from sklearn.manifold import TSNE
+def reduce_dim(X, perplexity=30, n_iter=1000):
+    tsne = TSNE(
+        n_components=2,
+        perplexity=perplexity,
+        n_iter=n_iter,
+        init='pca',
+        # learning_rate="auto",
+    )
+    Z = tsne.fit_transform(X)
+    return Z
+
+
+from IPython.display import Video
+def show_video(video_path):
+    """Show a video in a Jupyter notebook"""
+    assert exists(video_path), f"Video path {video_path} does not exist"
+    
+    # display the video in a Jupyter notebook
+    return Video(video_path, embed=True, width=480)
+    # Video(video_path, embed=True, width=600, height=400)
+    # html_attributes="controls autoplay loop muted"
+
+
+
+
+def show_single_audio(filepath=None, data=None, rate=None, start=None, end=None, label="Sample audio"):        
+    
+    if filepath is None:
+        assert data is not None and rate is not None, "Either filepath or data and rate must be provided"
+        args = dict(data=data, rate=rate)
+    else:
+        assert data is None and rate is None, "Either filepath or data and rate must be provided"
+        data, rate = librosa.load(filepath)
+        # args = dict(filename=filepath)
+        args = dict(data=data, rate=rate)
+    
+    if start is not None and end is not None:
+        start = max(int(start * rate), 0)
+        end = min(int(end * rate), len(data))
+    else:
+        start = 0
+        end = len(data)
+    data = data[start:end]
+    args["data"] = data
+
+    if label is None:
+        label = "Sample audio"
+
+    label = Label(f"{label}")
+    out = widgets.Output()
+    with out:
+        display(Audio(**args))
+    vbox = VBox([label, out])
+    return vbox
+
+
+def show_single_audio_with_spectrogram(filepath=None, data=None, rate=None, label="Sample audio", figsize=(6, 2)):
+    
+    if filepath is None:
+        assert data is not None and rate is not None, "Either filepath or data and rate must be provided"
+    else:
+        data, rate = librosa.load(filepath)
+    
+    # Show audio
+    vbox = show_single_audio(data=data, rate=rate, label=label)
+    # get width of audio widget
+    width = vbox.children[1].layout.width
+
+    # Show spectrogram
+    spec_out = widgets.Output()
+    D = librosa.stft(data)  # STFT of y
+    S_db = librosa.amplitude_to_db(np.abs(D), ref=np.max)
+    with spec_out:
+        fig, ax = plt.subplots(figsize=figsize)
+        img = librosa.display.specshow(
+            S_db,
+            ax=ax,
+            x_axis='time',
+            # y_axis='linear',
+        )
+    # img = widgets.Image.from_file(fig)
+    # import ipdb; ipdb.set_trace()
+    # img = widgets.Image(img)
+    # add image to vbox
+    vbox.children += (spec_out,)
+    return vbox
+
+def show_spectrogram(audio_path=None, data=None, rate=None, figsize=(6, 2), ax=None, show=True):
+    if data is None and rate is None:
+        # Show spectrogram
+        data, rate = librosa.load(audio_path)
+    else:
+        assert audio_path is None, "Either audio_path or data and rate must be provided"
+
+    hop_length = 512
+    D = librosa.stft(data, n_fft=2048, hop_length=hop_length, win_length=2048)  # STFT of y
+    S_db = librosa.amplitude_to_db(np.abs(D), ref=np.max)
+
+    # Create spectrogram plot widget
+    if ax is None:
+        fig, ax = plt.subplots(1, 1, figsize=figsize)
+    im = ax.imshow(S_db, origin='lower', aspect='auto', cmap='inferno')
+
+    # Replace xtixks with time
+    xticks = ax.get_xticks()
+    time_in_seconds = librosa.frames_to_time(xticks, sr=rate, hop_length=hop_length)
+    ax.set_xticklabels(np.round(time_in_seconds, 1))
+    ax.set_xlabel('Time')
+    ax.set_yticks([])
+    if ax is None:
+        plt.close(fig)
+
+    # Create widget output
+    spec_out = widgets.Output()
+    with spec_out:
+        display(fig)
+    return spec_out
+
+
+def show_single_video_and_spectrogram(
+        video_path, audio_path,
+        label="Sample video", figsize=(6, 2),
+        width=480,
+        show_spec_stats=False,
+    ):
+    # Show video
+    vbox = show_single_video(video_path, label=label, width=width)
+    # get width of video widget
+    width = vbox.children[1].layout.width
+
+    # Show spectrogram
+    data, rate = librosa.load(audio_path)
+    hop_length = 512
+    D = librosa.stft(data, n_fft=2048, hop_length=hop_length, win_length=2048)  # STFT of y
+    S_db = librosa.amplitude_to_db(np.abs(D), ref=np.max)
+
+    # Create spectrogram plot widget
+    fig, ax = plt.subplots(1, 1, figsize=figsize)
+    im = ax.imshow(S_db, origin='lower', aspect='auto', cmap='inferno')
+
+    # Replace xtixks with time
+    xticks = ax.get_xticks()
+    time_in_seconds = librosa.frames_to_time(xticks, sr=rate, hop_length=hop_length)
+    ax.set_xticklabels(np.round(time_in_seconds, 1))
+    ax.set_xlabel('Time')
+    ax.set_yticks([])
+    plt.close(fig)
+
+    # Create widget output
+    spec_out = widgets.Output()
+    with spec_out:
+        display(fig)
+    vbox.children += (spec_out,)
+
+    if show_spec_stats:
+        # Compute mean of spectrogram over frequency axis
+        eps = 1e-5
+        S_db_normalized = (S_db - S_db.mean(axis=1)[:, None]) / (S_db.std(axis=1)[:, None] + eps)
+        S_db_over_time = S_db_normalized.sum(axis=0)
+        # Plot S_db_over_time
+        fig, ax = plt.subplots(1, 1, figsize=(6, 2))
+        # ax.set_title("Spectrogram over time")
+        ax.grid(alpha=0.5)
+        x = np.arange(len(S_db_over_time))
+        x = librosa.frames_to_time(x, sr=rate, hop_length=hop_length)
+        x = np.round(x, 1)
+        ax.plot(x, S_db_over_time)
+        ax.set_xlabel('Time')
+        ax.set_yticks([])
+        plt.close(fig)
+        plot_out = widgets.Output()
+        with plot_out:
+            display(fig)
+        vbox.children += (plot_out,)
+
+    return vbox
+
+
+def show_single_spectrogram(
+        filepath=None,
+        data=None,
+        rate=None,
+        start=None,
+        end=None,
+        ax=None,
+        label="Sample spectrogram",
+        figsize=(6, 2),
+        xlabel="Time",
+    ):
+    
+    if filepath is None:
+        assert data is not None and rate is not None, "Either filepath or data and rate must be provided"
+    else:
+        rate = 22050
+        offset = start or 0
+        clip_duration = end - start if end is not None else None
+        data, rate = librosa.load(filepath, sr=rate, offset=offset, duration=clip_duration)
+    
+    # start = 0 if start is None else int(rate * start)
+    # end = len(data) if end is None else int(rate * end)
+    # data = data[start:end]
+    
+    # Show spectrogram
+    spec_out = widgets.Output()
+    D = librosa.stft(data)  # STFT of y
+    S_db = librosa.amplitude_to_db(np.abs(D), ref=np.max)
+
+    if ax is None:
+        fig, ax = plt.subplots(figsize=figsize)
+
+    with spec_out:
+        img = librosa.display.specshow(
+            S_db,
+            ax=ax,
+            x_axis='time',
+            sr=rate,
+            # y_axis='linear',
+        )
+    ax.set_xlabel(xlabel)
+    ax.margins(x=0)
+    plt.subplots_adjust(wspace=0, hspace=0)
+
+    # img = widgets.Image.from_file(fig)
+    # import ipdb; ipdb.set_trace()
+    # img = widgets.Image(img)
+    # add image to vbox
+    vbox = VBox([spec_out])
+    return vbox
+    # return spec_out
+
+
+# from decord import VideoReader
+def show_single_video(filepath, label="Sample video", width=480, fix_resolution=True):
+    
+    if label is None:
+        label = "Sample video"
+    
+    height = None
+    if fix_resolution:
+        aspect_ratio = 16. / 9.
+        height = int(width * (1/ aspect_ratio))
+
+    label = Label(f"{label}")
+    out = widgets.Output()
+    with out:
+        display(Video(filepath, embed=True, width=width, height=height))
+    vbox = VBox([label, out])
+    return vbox
+
+
+def show_grid_of_audio(files, starts=None, ends=None, labels=None, ncols=None, show_spec=False):
+    
+    for f in files:
+        assert os.path.exists(f), f"File {f} does not exist."
+
+    if labels is None:
+        labels = [None] * len(files)
+    
+    if starts is None:
+        starts = [None] * len(files)
+    
+    if ends is None:
+        ends = [None] * len(files)
+
+    assert len(files) == len(labels)
+    
+    if ncols is None:
+        ncols = 3
+    nfiles = len(files)
+    nrows = nfiles // ncols + (nfiles % ncols != 0)
+    # print(nrows, ncols)
+    
+    for i in range(nrows):
+        row_hbox = []
+        for j in range(ncols):
+            idx = i * ncols + j
+            # print(i, j, idx)
+            
+            if idx < len(files):
+                file, label = files[idx], labels[idx]
+                start, end = starts[idx], ends[idx]
+                vbox = show_single_audio(
+                    filepath=file, label=label, start=start, end=end
+                )
+                if show_spec:
+                    spec_box = show_spectrogram(file, figsize=(3.6, 1))
+                    # Add spectrogram to vbox
+                    vbox.children += (spec_box,)
+
+                # if not show_spec:
+                #     vbox = show_single_audio(
+                #         filepath=file, label=label, start=start, end=end
+                #     )
+                # else:
+                #     vbox = show_single_audio_with_spectrogram(
+                #         filepath=file, label=label
+                #     )
+                row_hbox.append(vbox)
+        row_hbox = HBox(row_hbox)
+        display(row_hbox)
+
+
+def show_grid_of_videos(
+        files,
+        cut=False,
+        starts=None,
+        ends=None,
+        labels=None,
+        ncols=None,
+        width_overflow=False,
+        show_spec=False,
+        width_of_screen=1000,
+    ):
+    from moviepy.editor import VideoFileClip
+    
+    for f in files:
+        assert os.path.exists(f), f"File {f} does not exist."
+
+    if labels is None:
+        labels = [None] * len(files)
+    if starts is not None and ends is not None:
+        cut = True
+    if starts is None:
+        starts = [None] * len(files)
+    if ends is None:
+        ends = [None] * len(files)
+
+    assert len(files) == len(labels) == len(starts) == len(ends)
+    
+    # cut the videos to the specified duration
+    if cut:
+        cut_files = []
+        for i, f in enumerate(files):
+            start, end = starts[i], ends[i]
+            
+            tmp_f = os.path.join(os.path.expanduser("~"), f"tmp/clip_{i}.mp4")
+            cut_files.append(tmp_f)
+        
+            video = VideoFileClip(f)
+            start = 0 if start is None else start
+            end = video.duration-1 if end is None else end
+            # print(start, end)
+            video.subclip(start, end).write_videofile(tmp_f, logger=None, verbose=False)
+        files = cut_files
+
+    if ncols is None:
+        ncols = 3
+        width_of_screen = 1000
+
+    # get width of the whole display screen
+    if not width_overflow:
+        width_of_single_video = width_of_screen // ncols
+    else:
+        width_of_single_video = 280
+
+    nfiles = len(files)
+    nrows = nfiles // ncols + (nfiles % ncols != 0)
+    # print(nrows, ncols)
+    
+    for i in range(nrows):
+        row_hbox = []
+        for j in range(ncols):
+            idx = i * ncols + j
+            # print(i, j, idx)
+            
+            if idx < len(files):
+                file, label = files[idx], labels[idx]
+                if not show_spec:
+                    vbox = show_single_video(file, label, width_of_single_video)
+                else:
+                    vbox = show_single_video_and_spectrogram(file, file, width=width_of_single_video, label=label)
+                row_hbox.append(vbox)
+        row_hbox = HBox(row_hbox)
+        display(row_hbox)
+        
+
+
+def preview_video(fp, label="Sample video frames", mode="uniform", frames_to_show=6):
+    from decord import VideoReader
+    
+    assert exists(fp), f"Video does not exist at {fp}"
+    vr = VideoReader(fp)
+
+    nfs = len(vr)
+    fps = vr.get_avg_fps()
+    dur = nfs / fps
+    
+    if mode == "all":
+        frame_indices = np.arange(nfs)
+    elif mode == "uniform":
+        frame_indices = np.linspace(0, nfs - 1, frames_to_show, dtype=int)
+    elif mode == "random":
+        frame_indices = np.random.randint(0, nfs - 1, replace=False)
+        frame_indices = sorted(frame_indices)
+    else:
+        raise ValueError(f"Unknown frame viewing mode {mode}.")
+    
+    # Show grid of image
+    images = vr.get_batch(frame_indices).asnumpy()
+    show_grid_of_images(images, n_cols=len(frame_indices), title=label, figsize=(12, 2.3), titlesize=10)
+
+
+def preview_multiple_videos(fps, labels, mode="uniform", frames_to_show=6):
+    for fp in fps:
+        assert exists(fp), f"Video does not exist at {fp}"
+    
+    for fp, label in zip(fps, labels):
+        preview_video(fp, label, mode=mode, frames_to_show=frames_to_show)
+
+
+
+def show_small_clips_in_a_video(
+        video_path,
+        clip_segments: list,
+        width=360,
+        labels=None,
+        show_spec=False,
+        resize=False,
+    ):
+    from moviepy.editor import VideoFileClip
+    from ipywidgets import Layout
+
+    video = VideoFileClip(video_path)
+    
+    if resize:
+        # Resize the video
+        print("Resizing the video to width", width)
+        video = video.resize(width=width)
+    
+    if labels is None:
+        labels = [
+            f"Clip {i+1} [{clip_segments[i][0]} : {clip_segments[i][1]}]" for i in range(len(clip_segments))
+        ]
+    else:
+        assert len(labels) == len(clip_segments)
+    
+    tmp_dir = os.path.join(os.path.expanduser("~"), "tmp")
+    tmp_clippaths = [f"{tmp_dir}/clip_{i}.mp4" for i in range(len(clip_segments))]
+    
+    iterator = tqdm_iterator(zip(clip_segments, tmp_clippaths), total=len(clip_segments), desc="Preparing clips")
+    clips = [
+        video.subclip(x, y).write_videofile(f, logger=None, verbose=False) \
+            for (x, y), f in iterator
+    ]
+    # show_grid_of_videos(tmp_clippaths, labels, ncols=len(clips), width_overflow=True)
+    hbox = []
+    for i in range(len(clips)):
+        # vbox = show_single_video(tmp_clippaths[i], labels[i], width=280)
+        
+        vbox = widgets.Output()
+        with vbox:
+            if show_spec:
+                display(
+                    show_single_video_and_spectrogram(
+                        tmp_clippaths[i], tmp_clippaths[i],
+                        width=width, figsize=(4.4, 1.5), 
+                    )
+                )
+            else:
+                display(Video(tmp_clippaths[i], embed=True, width=width))
+            # reduce vspace between video and label
+            display(Label(labels[i], layout=Layout(margin="-8px 0px 0px 0px")))
+            # if show_spec:
+            #     display(show_single_spectrogram(tmp_clippaths[i], figsize=(4.5, 1.5)))
+        hbox.append(vbox)
+    hbox = HBox(hbox)
+    display(hbox)
+
+
+def show_single_video_and_audio(
+        video_path, audio_path, label="Sample video and audio",
+        start=None, end=None, width=360, sr=44100, show=True,
+    ):
+    from moviepy.editor import VideoFileClip
+
+    # Load video
+    video = VideoFileClip(video_path)
+    video_args = {"embed": True, "width": width}
+    filepath = video_path
+
+    # Load audio
+    audio_waveform, sr = librosa.load(audio_path, sr=sr)
+    audio_args = {"data": audio_waveform, "rate": sr}
+
+    if start is not None and end is not None:
+        
+        # Cut video from start to end
+        tmp_dir = os.path.join(os.path.expanduser("~"), "tmp")
+        clip_path = os.path.join(tmp_dir, "clip_sample.mp4")
+        video.subclip(start, end).write_videofile(clip_path, logger=None, verbose=False)
+        filepath = clip_path
+        
+        # Cut audio from start to end
+        audio_waveform = audio_waveform[int(start * sr): int(end * sr)]
+        audio_args["data"] = audio_waveform
+
+    out = widgets.Output()
+    with out:
+        label = f"{label} [{start} : {end}]"
+        display(Label(label))
+        display(Video(filepath, **video_args))
+        display(Audio(**audio_args))
+    
+    if show:
+        display(out)
+    else:
+        return out
+
+
+def plot_waveform(waveform, sample_rate, figsize=(10, 2), ax=None, skip=100, show=True, title=None):
+    if isinstance(waveform, torch.Tensor):
+        waveform = waveform.numpy()
+    
+    time_axis = torch.arange(0, len(waveform)) / sample_rate
+    waveform = waveform[::skip]
+    time_axis = time_axis[::skip]
+
+    if len(waveform.shape) == 1:
+        num_channels = 1
+        num_frames = waveform.shape[0]
+        waveform = waveform.reshape(1, num_frames)
+    elif len(waveform.shape) == 2:
+        num_channels, num_frames = waveform.shape
+    else:
+        raise ValueError(f"Waveform has invalid shape {waveform.shape}")
+    
+    if ax is None:
+        figure, axes = plt.subplots(num_channels, 1, figsize=figsize)
+        if num_channels == 1:
+            axes = [axes]
+        for c in range(num_channels):
+            axes[c].plot(time_axis, waveform[c], linewidth=1)
+            axes[c].grid(True)
+            if num_channels > 1:
+                axes[c].set_ylabel(f"Channel {c+1}")
+        figure.suptitle(title)
+    else:
+        assert num_channels == 1
+        ax.plot(time_axis, waveform[0], linewidth=1)
+        ax.grid(True)
+        # ax.set_xticks([])
+        # ax.set_yticks([])
+        # ax.set_xlim(-0.1, 0.1)
+        ax.set_ylim(-0.05, 0.05)
+    
+    if show:
+        plt.show(block=False)
+
+
+def show_waveform_as_image(waveform, sr=16000):
+    """Plots a waveform as plt fig and converts into PIL.Image"""
+    fig, ax = plt.subplots(figsize=(10, 2))
+    plot_waveform(waveform, sr, ax=ax, show=False)
+    fig.canvas.draw()
+    img = Image.frombytes('RGB', fig.canvas.get_width_height(), fig.canvas.tostring_rgb())
+    plt.close(fig)
+    return img
+
+
+def plot_raw_audio_signal_with_markings(signal: np.ndarray, markings: list,
+        title: str = 'Raw audio signal with markings',
+        figsize: tuple = (23, 4),
+    ):
+
+    plt.figure(figsize=figsize)
+    plt.grid()
+
+    plt.plot(signal)
+    for value in markings:
+        plt.axvline(x=value, c='red')
+    plt.xlabel('Time')
+    plt.title(title)
+
+    plt.show()
+    plt.close()
+
+
+def get_concat_h(im1, im2):
+    """Concatenate two images horizontally"""
+    dst = Image.new('RGB', (im1.width + im2.width, im1.height))
+    dst.paste(im1, (0, 0))
+    dst.paste(im2, (im1.width, 0))
+    return dst
+
+
+def concat_images(images):
+    im1 = images[0]
+    canvas_height = max([im.height for im in images])
+    dst = Image.new('RGB', (sum([im.width for im in images]), im1.height))
+    start_width = 0
+    for i, im in enumerate(images):
+        if im.height < canvas_height:
+            start_height = (canvas_height - im.height) // 2
+        else:
+            start_height = 0
+        print(i, start_height)
+        dst.paste(im, (start_width, start_height))
+        start_width += im.width
+    return dst
+
+
+def concat_images_with_border(images, border_width=5, border_color="white"):
+    im1 = images[0]
+    total_width = sum([im.width for im in images]) + (len(images) - 1) * border_width
+    max_height = max([im.height for im in images])
+    dst = Image.new(
+        'RGB',
+        (total_width, max_height),
+        border_color,
+    )
+    start_width = 0
+    uniform_height = im1.height
+    canvas_height = max([im.height for im in images])
+    for i, im in enumerate(images):
+        # if im.height != uniform_height:
+        #     im = resize_height(im.copy(), uniform_height)
+        if im.height < canvas_height:
+            start_height = (canvas_height - im.height) // 2
+            
+            # Pad with zeros at top and bottom
+            im = ImageOps.expand(
+                im, border=(0, start_height, 0, canvas_height - im.height - start_height),
+            )
+            start_height = 0
+        else:
+            start_height = 0
+        dst.paste(im, (start_width, start_height))
+        start_width += im.width + border_width
+    return dst
+
+
+def concat_images_vertically(images):
+    im1 = images[0]
+    dst = Image.new('RGB', (im1.width, sum([im.height for im in images])))
+    start_height = 0
+    for i, im in enumerate(images):
+        dst.paste(im, (0, start_height))
+        start_height += im.height
+    return dst
+
+
+def concat_images_vertically_with_border(images, border_width=5, border_color="white"):
+    im1 = images[0]
+    dst = Image.new('RGB', (im1.width, sum([im.height for im in images]) + (len(images) - 1) * border_width), border_color)
+    start_height = 0
+    for i, im in enumerate(images):
+        dst.paste(im, (0, start_height))
+        start_height += im.height + border_width
+    return dst
+
+
+def get_concat_v(im1, im2):
+    """Concatenate two images vertically"""
+    dst = Image.new('RGB', (im1.width, im1.height + im2.height))
+    dst.paste(im1, (0, 0))
+    dst.paste(im2, (0, im1.height))
+    return dst
+
+
+def set_latex_fonts(usetex=True, fontsize=14, show_sample=False, **kwargs):
+    try:
+        plt.rcParams.update({
+            "text.usetex": usetex,
+            "font.family": "serif",
+            "font.serif": ["Computer Modern Roman"],
+            "font.size": fontsize,
+            **kwargs,
+        })
+        if show_sample:
+            plt.figure()
+            plt.title("Sample $y = x^2$")
+            plt.plot(np.arange(0, 10), np.arange(0, 10)**2, "--o")
+            plt.grid()
+            plt.show()
+    except:
+        print("Failed to setup LaTeX fonts. Proceeding without.")
+        pass
+
+
+def get_colors(num_colors, palette="jet"):
+    cmap = plt.get_cmap(palette)
+    colors = [cmap(i) for i in np.linspace(0, 1, num_colors)]
+    return colors
+
+
+def add_box_on_image(image, bbox, color="red", thickness=3, resized=False, fillcolor=None, fillalpha=0.2):
+    """
+    Adds bounding box on image.
+    
+    Args:
+        image (PIL.Image): image
+        bbox (list): [xmin, ymin, xmax, ymax]
+        color: -
+        thickness: -
+    """
+    image = image.copy().convert("RGB")
+    # color = get_predominant_color(color)
+    color = PIL.ImageColor.getrgb(color)
+    
+    # Apply alpha to fillcolor
+    if fillcolor is not None:
+        if isinstance(fillcolor, str):
+            fillcolor = PIL.ImageColor.getrgb(fillcolor)
+            fillcolor= fillcolor + (int(fillalpha * 255),)
+        elif isinstance(fillcolor, tuple):
+            if len(fillcolor) == 3:
+                fillcolor= fillcolor + (int(fillalpha * 255),)
+            else:
+                pass
+
+    # Create an instance of the ImageDraw class
+    draw = ImageDraw.Draw(image, "RGBA")
+
+    # Draw the bounding box on the image
+    draw.rectangle(bbox, outline=color, width=thickness, fill=fillcolor)
+
+    # Resize
+    new_width, new_height = (320, 240)
+    if resized:
+        image = image.resize((new_width, new_height))
+
+    return image
+
+
+def add_multiple_boxes_on_image(image, bboxes, colors=None, thickness=3, resized=False, fillcolor=None, fillalpha=0.2):
+    image = image.copy().convert("RGB")
+    if colors is None:
+        colors = ["red"] * len(bboxes)
+    for bbox, color in zip(bboxes, colors):
+        image = add_box_on_image(image, bbox, color, thickness, resized, fillcolor, fillalpha)
+    return image
+
+
+def colorize_mask(mask, color="red"):
+    # mask = mask.convert("RGBA")
+    color = PIL.ImageColor.getrgb(color)
+    mask = ImageOps.colorize(mask, (0, 0, 0, 0), color)
+    return mask
+
+
+def add_mask_on_image(image: Image, mask: Image, color="green", alpha=0.5):
+    image = image.copy()
+    mask = mask.copy()
+
+    # get color if it is a string
+    if isinstance(color, str):
+        color = PIL.ImageColor.getrgb(color)
+    # color = get_predominant_color(color)
+    mask = ImageOps.colorize(mask, (0, 0, 0, 0), color)
+
+    mask = mask.convert("RGB")
+    assert (mask.size == image.size)
+    assert (mask.mode == image.mode)
+
+    # Blend the original image and the segmentation mask with a 50% weight
+    blended_image = Image.blend(image, mask, alpha)
+    return blended_image
+
+
+def blend_images(img1, img2, alpha=0.5):
+    # Convert images to RGBA
+    img1 = img1.convert("RGBA")
+    img2 = img2.convert("RGBA")
+    alpha_blended = Image.blend(img1, img2, alpha=alpha)
+    # Convert back to RGB
+    alpha_blended = alpha_blended.convert("RGB")
+    return alpha_blended
+
+
+def visualize_youtube_clip(
+        youtube_id, st, et, label="",
+        show_spec=False,
+        video_width=360, video_height=240,
+    ):
+    
+    url = f"https://www.youtube.com/embed/{youtube_id}?start={int(st)}&end={int(et)}"
+    video_html_code = f"""
+    <iframe height="{video_height}" width="{video_width}" src="{url}" frameborder="0" allowfullscreen></iframe>
+    """
+    label_html_code = f"""<b>Caption</b>: {label} <br> <b>Time</b>: {st} to {et}"""
+    
+    # Show label and video below it
+    label = widgets.HTML(label_html_code)
+    video = widgets.HTML(video_html_code)
+    
+    if show_spec:
+        import pytube
+        import base64
+        from io import BytesIO
+        from moviepy.video.io.VideoFileClip import VideoFileClip
+        from moviepy.audio.io.AudioFileClip import AudioFileClip
+
+        # Load audio directly from youtube
+        video_url = f"https://www.youtube.com/watch?v={youtube_id}"
+        yt = pytube.YouTube(video_url)
+        # Get the audio stream
+        audio_stream = yt.streams.filter(only_audio=True).first()
+
+        # Download audio stream
+        # audio_file = os.path.join("/tmp", "sample_audio.mp3")
+        audio_stream.download(output_path='/tmp', filename='sample.mp4')
+        
+        audio_clip = AudioFileClip("/tmp/sample.mp4")
+        audio_subclip = audio_clip.subclip(st, et)
+        sr = audio_subclip.fps
+        y = audio_subclip.to_soundarray().mean(axis=1)
+        audio_subclip.close()
+        audio_clip.close()
+        
+        # Compute spectrogram in librosa
+        S_db = librosa.power_to_db(librosa.feature.melspectrogram(y, sr=sr), ref=np.max)
+        # Compute width in cms from video_width
+        width = video_width / plt.rcParams["figure.dpi"] + 0.63
+        height = video_height / plt.rcParams["figure.dpi"]
+        out = widgets.Output()
+        with out:
+            fig, ax = plt.subplots(figsize=(width, height))
+            librosa.display.specshow(S_db, sr=sr, x_axis='time', ax=ax)
+            ax.set_ylabel("Frequency (Hz)")
+    else:
+        out = widgets.Output()
+    
+    vbox = widgets.VBox([label, video, out])
+
+    return vbox
+ 
+
+def visualize_pair_of_youtube_clips(clip_a, clip_b):
+    yt_id_a = clip_a["youtube_id"]
+    label_a = clip_a["sentence"]
+    st_a, et_a = clip_a["time"]
+    
+    yt_id_b = clip_b["youtube_id"]
+    label_b = clip_b["sentence"]
+    st_b, et_b = clip_b["time"]
+    
+    # Show the clips side by side
+    clip_a = visualize_youtube_clip(yt_id_a, st_a, et_a, label_a, show_spec=True)
+    # clip_a = widgets.Output()
+    # with clip_a:
+    #     visualize_youtube_clip(yt_id_a, st_a, et_a, label_a, show_spec=True)
+    
+    clip_b = visualize_youtube_clip(yt_id_b, st_b, et_b, label_b, show_spec=True)
+    # clip_b = widgets.Output()
+    # with clip_b:
+    #     visualize_youtube_clip(yt_id_b, st_b, et_b, label_b, show_spec=True)
+
+    hbox = HBox([
+        clip_a, clip_b
+    ])
+    display(hbox)
+    
+
+def plot_1d(x: np.ndarray, figsize=(6, 2), title=None, xlabel=None, ylabel=None, show=True, **kwargs):
+    assert (x.ndim == 1)
+    fig, ax = plt.subplots(figsize=figsize)
+    ax.grid(alpha=0.3)
+    ax.set_title(title)
+    ax.set_xlabel(xlabel)
+    ax.set_ylabel(ylabel)
+    ax.plot(np.arange(len(x)), x, **kwargs)
+    if show:
+        plt.show()
+    else:
+        plt.close()
+    return fig
+
+
+
+def make_grid(cols,rows):
+    import streamlit as st
+    grid = [0]*cols
+    for i in range(cols):
+        with st.container():
+            grid[i] = st.columns(rows)
+    return grid
+
+
+def display_clip(video_path, stime, etime, label=None):
+    """Displays clip at index i."""
+    assert exists(video_path), f"Video does not exist at {video_path}"
+    display(
+        show_small_clips_in_a_video(
+            video_path, [(stime, etime)], labels=[label],
+        ),
+    )
+
+
+def countplot(df, column, title=None, rotation=90, ylabel="Count", figsize=(8, 5), ax=None, show=True, show_counts=False):
+    
+    if ax is None:
+        fig, ax = plt.subplots(figsize=figsize)
+
+    ax.grid(alpha=0.4)
+    ax.set_xlabel(column)
+    ax.set_ylabel(ylabel)
+    ax.set_title(title)
+    
+    data = dict(df[column].value_counts())
+    # Extract keys and values from the dictionary
+    categories = list(data.keys())
+    counts = list(data.values())
+
+    # Create a countplot
+    ax.bar(categories, counts)
+    ax.set_xticklabels(categories, rotation=rotation)
+    
+    # Show count values on top of bars
+    if show_counts:
+        max_v = max(counts)
+        for i, v in enumerate(counts):
+            delta = 0.01 * max_v
+            ax.text(i, v + delta, str(v), ha="center")
+    
+    if show:
+        plt.show()
+
+
+def get_linspace_colors(cmap_name='viridis', num_colors = 10):
+    import matplotlib.colors as mcolors
+
+    # Get the colormap object
+    cmap = plt.cm.get_cmap(cmap_name)
+
+    # Get the evenly spaced indices
+    indices = np.arange(0, 1, 1./num_colors)
+
+    # Get the corresponding colors from the colormap
+    colors = [mcolors.to_hex(cmap(idx)) for idx in indices]
+    
+    return colors
+
+
+def hex_to_rgb(colors):
+    from PIL import ImageColor
+    return [ImageColor.getcolor(c, "RGB") for c in colors]
+
+
+def plot_audio_feature(times, feature, feature_label="Feature", xlabel="Time", figsize=(20, 2)):
+    fig, ax = plt.subplots(1, 1, figsize=figsize)
+    ax.grid(alpha=0.4)
+    ax.set_xlabel(xlabel)
+    ax.set_ylabel(feature_label)
+    ax.set_yticks([])
+    
+    ax.plot(times, feature, '--', linewidth=0.5)
+    plt.show()
+
+
+
+def compute_rms(y, frame_length=512):
+    rms = librosa.feature.rms(y=y, frame_length=frame_length)[0]
+    times = librosa.samples_to_time(frame_length * np.arange(len(rms)))
+    return times, rms
+
+
+def plot_audio_features(path, label, show=True, show_video=True, features=["rms"], frame_length=512, figsize=(5, 2), return_features=False):
+    # Load audio
+    y, sr = librosa.load(path)
+    
+    # Show video
+    if show_video:
+        if show:
+            display(
+                show_single_video_and_spectrogram(
+                    path, path, label=label, figsize=figsize,
+                    width=410,
+                )
+            )
+    else:
+        if show:
+            # Show audio and spectrogram
+            display(
+                show_single_audio_with_spectrogram(path, label=label, figsize=figsize)
+            )
+
+    feature_data = dict() 
+    for f in features:
+        fn = eval(f"compute_{f}")
+        args = dict(y=y, frame_length=frame_length)
+        xvals, yvals = fn(**args)
+        feature_data[f] = (xvals, yvals)
+        
+        if show:
+            display(
+                plot_audio_feature(
+                    xvals, yvals, feature_label=f.upper(), figsize=(figsize[0] - 0.25, figsize[1]),
+                )
+            )
+    
+    if return_features:
+        return feature_data
+
+
+def rescale_frame(frame, scale=1.):
+    """Rescales a frame by a factor of scale."""
+    return frame.resize((int(frame.width * scale), int(frame.height * scale)))
+
+
+def save_gif(images, path, duration=None, fps=30):
+    import imageio
+    images = [np.asarray(image) for image in images]
+    if fps is not None:
+        imageio.mimsave(path, images, fps=fps)
+    else:
+        assert duration is not None
+        imageio.mimsave(path, images, duration=duration)
+
+
+def show_subsampled_frames(frames, n_show, figsize=(15, 3), as_canvas=True):
+    indices = np.arange(len(frames))
+    indices = np.linspace(0, len(frames) - 1, n_show, dtype=int)
+    show_frames = [frames[i] for i in indices]
+    if as_canvas:
+        return concat_images(show_frames)
+    else:
+        show_grid_of_images(show_frames, n_cols=n_show, figsize=figsize, subtitles=indices)
+
+
+def tensor_to_heatmap(x, scale=True, cmap="viridis", flip_vertically=False):
+    import PIL
+    
+    if isinstance(x, torch.Tensor):
+        x = x.numpy()
+    
+    if scale:
+        x = (x - x.min()) / (x.max() - x.min())
+    
+    cm = plt.get_cmap(cmap)
+    if flip_vertically:
+        x = np.flip(x, axis=0) # put low frequencies at the bottom in image
+    x = cm(x)
+    x = (x * 255).astype(np.uint8)
+    if x.shape[-1] == 3:
+        x = PIL.Image.fromarray(x, mode="RGB")
+    elif x.shape[-1] == 4:
+        x = PIL.Image.fromarray(x, mode="RGBA").convert("RGB")
+    else:
+        raise ValueError(f"Invalid shape {x.shape}")
+    return x
+
+
+def batch_tensor_to_heatmap(x, scale=True, cmap="viridis", flip_vertically=False, resize=None):
+    y = []
+    for i in range(len(x)):
+        h = tensor_to_heatmap(x[i], scale, cmap, flip_vertically)
+        if resize is not None:
+            h = h.resize(resize)
+        y.append(h)
+    return y
+
+
+def change_contrast(img, level):
+    factor = (259 * (level + 255)) / (255 * (259 - level))
+    def contrast(c):
+        return 128 + factor * (c - 128)
+    return img.point(contrast)
+
+
+def change_brightness(img, alpha):
+    import PIL
+    enhancer = PIL.ImageEnhance.Brightness(img)
+    # to reduce brightness by 50%, use factor 0.5
+    img = enhancer.enhance(alpha)
+    return img
+
+
+def draw_horizontal_lines(image, y_values, color=(255, 0, 0), colors=None, line_thickness=2):
+    """
+    Draw horizontal lines on a PIL image at specified Y positions.
+
+    Args:
+        image (PIL.Image.Image): The input PIL image.
+        y_values (list or int): List of Y positions where lines will be drawn.
+                               If a single integer is provided, a line will be drawn at that Y position.
+        color (tuple): RGB color tuple (e.g., (255, 0, 0) for red).
+        line_thickness (int): Thickness of the lines.
+
+    Returns:
+        PIL.Image.Image: The PIL image with the drawn lines.
+    """
+    image = image.copy()
+    
+    if isinstance(color, str):
+        color = PIL.ImageColor.getcolor(color, "RGB")
+        
+    if colors is None:
+        colors = [color] * len(y_values)
+    else:
+        if isinstance(colors[0], str):
+            colors = [PIL.ImageColor.getcolor(c, "RGB") for c in colors]
+    
+    if isinstance(y_values, int):
+        y_values = [y_values]
+    
+    # Create a drawing context on the image
+    draw = PIL.ImageDraw.Draw(image)
+
+    if isinstance(y_values, int):
+        y_values = [y_values]
+
+    for y, c in zip(y_values, colors):
+        draw.line([(0, y), (image.width, y)], fill=c, width=line_thickness)
+
+    return image
+
+
+def draw_vertical_lines(image, x_values, color=(255, 0, 0), colors=None, line_thickness=2):
+    """
+    Draw vertical lines on a PIL image at specified X positions.
+
+    Args:
+        image (PIL.Image.Image): The input PIL image.
+        x_values (list or int): List of X positions where lines will be drawn.
+                               If a single integer is provided, a line will be drawn at that X position.
+        color (tuple): RGB color tuple (e.g., (255, 0, 0) for red).
+        line_thickness (int): Thickness of the lines.
+
+    Returns:
+        PIL.Image.Image: The PIL image with the drawn lines.
+    """
+    image = image.copy()
+    
+    if isinstance(color, str):
+        color = PIL.ImageColor.getcolor(color, "RGB")
+        
+    if colors is None:
+        colors = [color] * len(x_values)
+    else:
+        if isinstance(colors[0], str):
+            colors = [PIL.ImageColor.getcolor(c, "RGB") for c in colors]
+    
+    if isinstance(x_values, int):
+        x_values = [x_values]
+    
+    # Create a drawing context on the image
+    draw = PIL.ImageDraw.Draw(image)
+
+    if isinstance(x_values, int):
+        x_values = [x_values]
+
+    for x, c in zip(x_values, colors):
+        draw.line([(x, 0), (x, image.height)], fill=c, width=line_thickness)
+
+    return image
+
+
+def show_arrow_on_image(image, start_loc, end_loc, color="red", thickness=3):
+    """Draw a line on PIL image from start_loc to end_loc."""
+    image = image.copy()
+    color = get_predominant_color(color)
+
+    # Create an instance of the ImageDraw class
+    draw = ImageDraw.Draw(image)
+
+    # Draw the bounding box on the image
+    draw.line([start_loc, end_loc], fill=color, width=thickness)
+
+    return image
+
+
+def draw_arrow_on_image_cv2(image, start_loc, end_loc, color="red", thickness=2, both_ends=False):
+    image = image.copy()
+    image = np.asarray(image)
+    if isinstance(color, str):
+        color = PIL.ImageColor.getcolor(color, "RGB")
+    image = cv2.arrowedLine(image, start_loc, end_loc, color, thickness)
+    if both_ends:
+        image = cv2.arrowedLine(image, end_loc, start_loc, color, thickness)
+    return PIL.Image.fromarray(image)
+
+
+def draw_arrow_with_text(image, start_loc, end_loc, text="", color="red", thickness=2, font_size=20, both_ends=False, delta=5):
+    image = np.asarray(image)
+    if isinstance(color, str):
+        color = PIL.ImageColor.getcolor(color, "RGB")
+
+    # Calculate the center point between start_loc and end_loc
+    center_x = (start_loc[0] + end_loc[0]) // 2
+    center_y = (start_loc[1] + end_loc[1]) // 2
+    center_point = (center_x, center_y)
+
+    # Draw the arrowed line
+    image = cv2.arrowedLine(image, start_loc, end_loc, color, thickness)
+    if both_ends:
+        image = cv2.arrowedLine(image, end_loc, start_loc, color, thickness)
+
+    # Create a PIL image from the NumPy array for drawing text
+    image_with_text = Image.fromarray(image)
+    draw = PIL.ImageDraw.Draw(image_with_text)
+    
+    # Calculate the text size
+    # font = PIL.ImageFont.truetype("arial.ttf", font_size)
+    # This gives an error: "OSError: cannot open resource", as a hack, use the following
+    text_width, text_height = draw.textsize(text)
+    
+    # Calculate the position to center the text
+    text_x = center_x - (text_width // 2) - delta
+    text_y = center_y - (text_height // 2)
+
+    # Draw the text
+    draw.text((text_x, text_y), text, color)
+
+    return image_with_text
+
+
+def draw_arrowed_line(image, start_loc, end_loc, color="red", thickness=2):
+    """
+    Draw an arrowed line on a PIL image from a starting point to an ending point.
+
+    Args:
+        image (PIL.Image.Image): The input PIL image.
+        start_loc (tuple): Starting point (x, y) for the arrowed line.
+        end_loc (tuple): Ending point (x, y) for the arrowed line.
+        color (str): Color of the line (e.g., 'red', 'green', 'blue').
+        thickness (int): Thickness of the line and arrowhead.
+
+    Returns:
+        PIL.Image.Image: The PIL image with the drawn arrowed line.
+    """
+    image = image.copy()
+    if isinstance(color, str):
+        color = PIL.ImageColor.getcolor(color, "RGB")
+    
+    
+    # Create a drawing context on the image
+    draw = ImageDraw.Draw(image)
+
+    # Draw a line from start to end
+    draw.line([start_loc, end_loc], fill=color, width=thickness)
+
+    # Calculate arrowhead points
+    arrow_size = 10  # Size of the arrowhead
+    dx = end_loc[0] - start_loc[0]
+    dy = end_loc[1] - start_loc[1]
+    length = (dx ** 2 + dy ** 2) ** 0.5
+    cos_theta = dx / length
+    sin_theta = dy / length
+    x1 = end_loc[0] - arrow_size * cos_theta
+    y1 = end_loc[1] - arrow_size * sin_theta
+    x2 = end_loc[0] - arrow_size * sin_theta
+    y2 = end_loc[1] + arrow_size * cos_theta
+    x3 = end_loc[0] + arrow_size * sin_theta
+    y3 = end_loc[1] - arrow_size * cos_theta
+
+    # Draw the arrowhead triangle
+    draw.polygon([end_loc, (x1, y1), (x2, y2), (x3, y3)], fill=color)
+
+    return image
+
+
+def center_crop_to_fraction(image, frac=0.5):
+    """Center crop an image to a fraction of its original size."""
+    width, height = image.size
+    new_width = int(width * frac)
+    new_height = int(height * frac)
+    left = (width - new_width) // 2
+    top = (height - new_height) // 2
+    right = (width + new_width) // 2
+    bottom = (height + new_height) // 2
+    return image.crop((left, top, right, bottom))
+
+
+def decord_load_frames(vr, frame_indices):
+    if isinstance(frame_indices, int):
+        frame_indices = [frame_indices]
+    frames = vr.get_batch(frame_indices).asnumpy()
+    frames = [Image.fromarray(frame) for frame in frames]
+    return frames
+
+
+def paste_mask_on_image(original_image, bounding_box, mask):
+    """
+    Paste a 2D mask onto the original image at the location specified by the bounding box.
+
+    Parameters:
+    - original_image (PIL.Image): The original image.
+    - bounding_box (tuple): Bounding box coordinates (left, top, right, bottom).
+    - mask (PIL.Image): The 2D mask.
+
+    Returns:
+    - PIL.Image: Image with the mask pasted on it.
+
+    Example:
+    ```
+    original_image = Image.open('original.jpg')
+    bounding_box = (100, 100, 200, 200)
+    mask = Image.open('mask.png')
+    result_image = paste_mask_on_image(original_image, bounding_box, mask)
+    result_image.show()
+    ```
+    """
+    # Create a copy of the original image to avoid modifying the input image
+    result_image = original_image.copy()
+
+    # Crop the mask to the size of the bounding box
+    mask_cropped = mask.crop((0, 0, bounding_box[2] - bounding_box[0], bounding_box[3] - bounding_box[1]))
+
+    # Paste the cropped mask onto the original image at the specified location
+    result_image.paste(mask_cropped, (bounding_box[0], bounding_box[1]))
+
+    return result_image
+
+
+def display_images_as_video_moviepy(image_list, fps=5, show=True):
+    """
+    Display a list of PIL images as a video in Jupyter Notebook using MoviePy.
+
+    Parameters:
+    - image_list (list): List of PIL images.
+    - fps (int): Frames per second for the video.
+    - show (bool): Whether to display the video in the notebook.
+
+    Example:
+    ```
+    image_list = [Image.open('frame1.jpg'), Image.open('frame2.jpg'), ...]
+    display_images_as_video_moviepy(image_list, fps=10)
+    ```
+    """
+    from IPython.display import display
+    from moviepy.editor import ImageSequenceClip
+
+    image_list = list(map(np.asarray, image_list))
+    clip = ImageSequenceClip(image_list, fps=fps)
+    if show:
+        display(clip.ipython_display(width=200))
+    os.remove("__temp__.mp4")
+
+
+def resize_height(img, H):
+    w, h = img.size
+    asp_ratio = w / h
+    W = np.ceil(asp_ratio * H).astype(int)
+    return img.resize((W, H))
+
+
+def resize_width(img, W):
+    w, h = img.size
+    asp_ratio = w / h
+    H = int(W / asp_ratio)
+    return img.resize((W, H))
+
+
+def resized_minor_side(img, size=256):
+    H, W = img.size
+    if H < W:
+        H_new = size
+        W_new = int(size * W / H)
+        return img.resize((W_new, H_new))
+    else:
+        W_new = size
+        H_new = int(size * H / W)
+        return img.resize((W_new, H_new))
+
+
+def brighten_image(img, alpha=1.2):
+    enhancer = PIL.ImageEnhance.Brightness(img)
+    img = enhancer.enhance(alpha)
+    return img
+
+
+def darken_image(img, alpha=0.8):
+    enhancer = PIL.ImageEnhance.Brightness(img)
+    img = enhancer.enhance(alpha)
+    return img
+
+
+def fig2img(fig):
+    """Convert a Matplotlib figure to a PIL Image and return it"""
+    import io
+    buf = io.BytesIO()
+    fig.savefig(buf)
+    buf.seek(0)
+    img = Image.open(buf)
+    return img
+
+
+def show_temporal_tsne(
+        tsne,
+        timestamps=None,
+        title="tSNE: feature vectors over time",
+        cmap='viridis',
+        ax=None,
+        fig=None,
+        show=True,
+        num_ticks=10,
+        return_as_pil=False,
+        dpi=100,
+        label='Time (s)',
+        figsize=(6, 4),
+        s=None,
+    ):
+
+    if timestamps is None:
+        timestamps = np.arange(len(tsne))
+
+    if ax is None or fig is None:
+        fig, ax = plt.subplots(1, 1, figsize=figsize, dpi=dpi)
+
+    cmap = plt.get_cmap(cmap)
+    scatter = ax.scatter(
+        tsne[:, 0], tsne[:, 1], c=np.arange(len(tsne)), cmap=cmap, s=s,
+        edgecolor='k', linewidth=0.5, 
+    )
+
+    ax.grid(alpha=0.4)
+    ax.set_title(f"{title}", fontsize=11)
+    ax.set_xlabel("$z_{1}$")
+    ax.set_ylabel("$z_{2}$")
+
+    # Create a colorbar
+    cbar = fig.colorbar(scatter, ax=ax, label=label)
+    
+    # Set custom ticks and labels on the colorbar
+    ticks = np.linspace(0, len(tsne) - 1, num_ticks, dtype=int)
+    tick_labels = np.round(timestamps[ticks], 1)
+    cbar.set_ticks(ticks)
+    cbar.set_ticklabels(tick_labels)
+
+    if show:
+        plt.show()
+    else:
+        if return_as_pil:
+            plt.tight_layout(pad=0.2)
+            # fig.canvas.draw()
+            # image = PIL.Image.frombytes(
+            #     'RGB',
+            #     fig.canvas.get_width_height(),
+            #     fig.canvas.tostring_rgb(),
+            # )
+            # return image
+
+            # Return as PIL Image without displaying the plt figure
+            image = fig2img(fig)
+            plt.close(fig)
+            return image
+
+
+def mark_keypoints(image, keypoints, color=(255, 255, 0), radius=1):
+    """
+    Marks keypoints on an image with a given color and radius.
+    
+    :param image: The input PIL image.
+    :param keypoints: A list of (x, y) tuples representing the keypoints.
+    :param color: The color to use for the keypoints (default: red).
+    :param radius: The radius of the circle to draw for each keypoint (default: 5).
+    :return: A new PIL image with the keypoints marked.
+    """
+    # Make a copy of the image to avoid modifying the original
+    image_copy = image.copy()
+
+    # Create a draw object to add graphical elements
+    draw = ImageDraw.Draw(image_copy)
+
+    # Loop through each keypoint and draw a circle
+    for x, y in keypoints:
+        # Draw a circle with the specified radius and color
+        draw.ellipse(
+            (x - radius, y - radius, x + radius, y + radius),
+            fill=color,
+            width=2
+        )
+
+    return image_copy
+
+
+def draw_line_on_image(image, x_coords, y_coords, color=(255, 255, 0), width=3):
+    """
+    Draws a line on an image given lists of x and y coordinates.
+    
+    :param image: The input PIL image.
+    :param x_coords: List of x-coordinates for the line.
+    :param y_coords: List of y-coordinates for the line.
+    :param color: Color of the line in RGB (default is red).
+    :param width: Width of the line (default is 3).
+    :return: The PIL image with the line drawn.
+    """
+    image = image.copy()
+
+    # Ensure the number of x and y coordinates are the same
+    if len(x_coords) != len(y_coords):
+        raise ValueError("x_coords and y_coords must have the same length")
+
+    # Create a draw object to draw on the image
+    draw = ImageDraw.Draw(image)
+
+    # Create a list of (x, y) coordinate tuples
+    coordinates = list(zip(x_coords, y_coords))
+
+    # Draw the line connecting the coordinates
+    draw.line(coordinates, fill=color, width=width)
+
+    return image
+
+
+def add_binary_strip_vertically(
+        image,
+        binary_vector,
+        strip_width=15,
+        one_color="yellow",
+        zero_color="gray",
+):
+    """
+    Add a binary strip to the right side of an image.
+
+    :param image: PIL Image to which the strip will be added.
+    :param binary_vector: Binary vector of length 512 representing the strip.
+    :param strip_width: Width of the strip to be added.
+    :param one_color: Color for "1" pixels (default: red).
+    :param zero_color: Color for "0" pixels (default: white).
+    :return: New image with the binary strip added on the right side.
+    """
+    one_color = PIL.ImageColor.getrgb(one_color)
+    zero_color = PIL.ImageColor.getrgb(zero_color)
+
+    height = image.height
+    if len(binary_vector) != height:
+        raise ValueError("Binary vector must be of length 512")
+
+    # Create a new strip with the specified width and 512 height
+    strip = PIL.Image.new("RGB", (strip_width, height))
+
+    # Fill the strip based on the binary vector
+    pixels = strip.load()
+    for i in range(height):
+        color = one_color if binary_vector[i] == 1 else zero_color
+        for w in range(strip_width):
+            pixels[w, i] = color
+
+    # Combine the original image with the new strip
+    # new_image = PIL.Image.new("RGB", (image.width + strip_width, height))
+    # new_image.paste(image, (0, 0))
+    # new_image.paste(strip, (image.width, 0))
+    new_image = image.copy()
+    new_image.paste(strip, (image.width - strip_width, 0))
+
+    return new_image
+
+
+def add_binary_strip_horizontally(
+    image,
+    binary_vector,
+    strip_height=15,
+    one_color="limegreen",
+    zero_color="gray",
+):
+    """
+    Add a binary strip to the top of an image.
+
+    :param image: PIL Image to which the strip will be added.
+    :param binary_vector: Binary vector of length 512 representing the strip.
+    :param strip_height: Height of the strip to be added.
+    :param one_color: Color for "1" pixels, accepts color names or hex (default: red).
+    :param zero_color: Color for "0" pixels, accepts color names or hex (default: white).
+    :return: New image with the binary strip added at the top.
+    """
+    width = image.width
+    if len(binary_vector) != width:
+        raise ValueError("Binary vector must be of length 512")
+
+    # Convert colors to RGB tuples
+    one_color_rgb = PIL.ImageColor.getrgb(one_color)
+    zero_color_rgb = PIL.ImageColor.getrgb(zero_color)
+
+    # Create a new strip with the specified height and 512 width
+    strip = PIL.Image.new("RGB", (width, strip_height))
+
+    # Fill the strip based on the binary vector
+    pixels = strip.load()
+    for i in range(width):
+        color = one_color_rgb if binary_vector[i] == 1 else zero_color_rgb
+        for h in range(strip_height):
+            pixels[i, h] = color
+
+    # Combine the original image with the new strip
+    # new_image = PIL.Image.new("RGB", (width, image.height + strip_height))
+    # new_image.paste(strip, (0, 0))
+    # new_image.paste(image, (0, strip_height))
+    new_image = image.copy()
+    new_image.paste(strip, (0, 0))
+
+    return new_image
+
+
+# Define a function to increase font sizes for a specific plot
+def increase_font_sizes(ax, font_scale=1.6):
+    for item in ([ax.title, ax.xaxis.label, ax.yaxis.label] +
+                 ax.get_xticklabels() + ax.get_yticklabels()):
+        item.set_fontsize(item.get_fontsize() * font_scale)
+
+
+
+def cut_fraction_of_bbox(image, box, frac=0.7):
+    """
+    Cuts the image such that the box occupies a fraction of the image.
+    """
+    W, H = image.size
+    x1, y1, x2, y2 = box
+    w = x2 - x1
+    h = y2 - y1
+    new_w = int(w / frac)
+    new_h = int(h / frac)
+    x1_new = max(0, x1 - (new_w - w) // 2)
+    x2_new = min(W, x2 + (new_w - w) // 2)
+    y1_new = max(0, y1 - (new_h - h) // 2)
+    y2_new = min(H, y2 + (new_h - h) // 2)
+    return image.crop((x1_new, y1_new, x2_new, y2_new))