lbhd/visualization.py

import numpy as np
import matplotlib.pyplot as plt

def highlight_score_plot(basename: str, model_scores: dict, num_frames: int, directory=None):
    fig = plt.gcf()
    fig.set_size_inches(20, 10)
    n_clips = len(list(model_scores.values())[0])
    x = np.arange(n_clips)
    x_labels = [f'{frame_id}~{frame_id + num_frames - 1}' for frame_id in range(n_clips)]

    for model, scores in model_scores.items():
        plt.plot(x, scores, 'o-', linewidth=4.0, markersize=10.0, label = model)
    plt.legend()
    plt.title(basename)
    plt.xticks(x, x_labels, rotation=90)
    plt.xlabel('Duration')
    plt.ylabel('Score')
    if directory!= None:
        plt.savefig(f'{directory}/highlight_score.jpg', dpi=100)
    else:
        plt.show()
    plt.close()

def plot_loss(train_loss, val_loss, exp_name: str, directory=None):
    fig = plt.gcf()
    fig.set_size_inches(20, 10)

    x = np.arange(len(train_loss))
    x_labels = x
    
    plt.plot(x, train_loss, 'o-', linewidth=4.0, markersize=10.0, label = f'{exp_name}_train')
    plt.plot(x, val_loss, 'o-', linewidth=4.0, markersize=10.0, label = f'{exp_name}_val')

    plt.legend()
    plt.title(f'{exp_name} training/validation loss')
    plt.xticks(x, x_labels)
    plt.xlabel('Epochs')
    plt.ylabel('Loss')
    if directory!= None:
        plt.savefig(f'{directory}/loss.jpg', dpi=100)
    else:
        plt.show()
    plt.close()

def plot_scores(scores, title=None, directory=None):
    import matplotlib.pyplot as plt
    import numpy as np
    fig = plt.gcf()
    fig.set_size_inches(20, 10)

    x = np.arange(len(scores))
    x_labels = [f'{x[i]}~{x[i]+2}' for i in range(len(x))]
    
    plt.plot(x, scores, 'o-', linewidth=4.0, markersize=10.0)

    plt.grid()
    if not title:
        plt.title(f'scores')
    else:
        plt.title(title)
    plt.xticks(x, x_labels, rotation=90)
    plt.xlabel('time')
    plt.ylabel('scores')
    if directory!= None:
        if not title:
            plt.savefig(f'{directory}/scores.jpg', dpi=100)
        else:
            plt.savefig(f'{directory}/{title}.jpg', dpi=100)
    else:
        plt.show()
    plt.close()

def extract_clips(video_path, intervals, path='../test.mp4'):
    from moviepy.editor import VideoFileClip, concatenate_videoclips
    video = VideoFileClip(video_path)
    clips = [video.subclip(start_sec, end_sec) if end_sec < video.duration else video.subclip(start_sec, video.duration) for start_sec, end_sec in intervals]
    highlight = concatenate_videoclips(clips)
    highlight = highlight.set_audio(highlight.audio)
    highlight.write_videofile(path,audio=True, audio_codec='aac', fps=30, logger=None)
    return highlight