A Python script to plot video frame sizes and frame types as a bar chart, calculate and plot 100ms, 1s and 3s moving average bitrate values.

plot-frame-sizes.py

# SPDX-License-Identifier: MPL-2.0
# Author: Maxim Sharabayko
# See https://maxsharabayko.github.io/blog/posts/plot-vc-frame-size/
import click
import pandas as pd
from bokeh.plotting import figure, show, output_file
from bokeh.io import output_notebook
from bokeh.transform import factor_cmap
from bokeh.models import Legend, LegendItem
from bokeh.layouts import column

def load_csv(file_path):
    """Load CSV file into a DataFrame and drop empty columns."""
    """ffprobe -v error -select_streams v:0 -show_entries frame=pkt_pos,pkt_size,pict_type -of csv=p=0 input_file.ts > frame_info.csv"""
    df = pd.read_csv(file_path, header=None, names=['frame_pos', 'frame_size', 'pict_type', 'empty1', 'empty2'])
    df = df.drop(columns=['empty1', 'empty2'])
    # Convert display order to stream order. Drop the old index to use the new order.
    df = df.sort_values(by='frame_pos').reset_index(drop=True)
    return df

def calculate_bitrates(df, fps):
    """Calculate 3s, 1s and 100 ms rolling average bitrate values."""
    fps_div_10 = fps // 10

    # Calculate the rolling sum of 'frame_size' values converted from bytes to megabits.
    df['mbps_mavg_1s'] = df['frame_size'].rolling(window=fps).sum() * 8 / 1e6

    # Similar for 3s moving average.
    df['mbps_mavg_3s'] = df['frame_size'].rolling(window=3*fps).sum() * 8 / 1e6 / 3

    # The same but using 100 millisecond intervals.
    df['mbps_mavg_100ms'] = df['frame_size'].rolling(window=fps_div_10).sum() * 8 / 1e6 * 10 / (fps_div_10 * 10) * fps
    return df

    # Calculate the first missing values of the 'mbps_mavg_1s',
    # place them into df['mbps_mavg_1s']
    #df.loc[:fps, 'mbps_mavg_1s'] = df['frame_size'].expanding(min_periods=1).sum().iloc[:fps] * 8 / 1e6


def calculate_kbyte_frame_sizes(df):
    # Convert to kilobytes
    df['frame_size_kBytes'] = df['frame_size'] / 1000
    return df

def add_color_column(df):
    """Add a color column based on pict_type."""
    colors = {'I': 'red', 'P': 'blue', 'B': 'green'}
    df['color'] = df['pict_type'].map(colors)
    return df

def create_bar_chart(df):
    """Create a bar chart for frame sizes with color legend."""
    # Create the bar chart figure.
    p1 = figure(x_axis_label='Frame Index', y_axis_label='Frame Size, kbytes', title='Frame Sizes by Type', width=1000, height=400)

    # Add bars to the bar chart figure.
    bars = p1.vbar(x=df.index, top=df['frame_size_kBytes'], width=0.9, color=df['color'])
    
    # Find the first index of each pict_type
    i_index = df[df['pict_type'] == 'I'].index[0]
    p_index = df[df['pict_type'] == 'P'].index[0]
    b_index = df[df['pict_type'] == 'B'].index[0]

    # Add a legend to describe color meaning.
    legend_items = [
        LegendItem(label="I Frame", renderers=[bars], index=i_index),
        LegendItem(label="P Frame", renderers=[bars], index=p_index),
        LegendItem(label="B Frame", renderers=[bars], index=b_index)
    ]
    legend1 = Legend(items=legend_items, location="top_right")
    p1.add_layout(legend1)
    
    return p1

def create_bitrates_plot(df, x_range):
    """Create a plot for bitrate values with toggleable lines."""
    # Create the moving average plot figure
    p2 = figure(x_axis_label='Frame Index', y_axis_label='Mbps', title='Bitrate', width=1000, height=400, x_range=x_range)

    # Add the Mbps moving average line to the moving average plot figure
    line1 = p2.line(df.index, df['mbps_mavg_3s'], line_width=2, color='black', legend_label='MAvg Bitrate (3s), Mbps')
    line2 = p2.line(df.index, df['mbps_mavg_1s'], line_width=2, color='green', legend_label='MAvg Bitrate (1s), Mbps')

    # Add the rolling average bitrate line to the moving average plot figure
    line3 = p2.line(df.index, df['mbps_mavg_100ms'], line_width=2, color='blue', legend_label='MAvg Bitrate (100ms), Mbps')

    p2.legend.click_policy = "hide"
    
    return p2

def plot_frame_sizes(df):
    """A function to plot frame sizes using a bar chart and bitrate values right below."""
    p1 = create_bar_chart(df)
    p2 = create_bitrates_plot(df, x_range=p1.x_range)
    
    # Uncomment if you want to use it in a notebook.
    #output_notebook()
    show(column(p1, p2))

@click.command()
@click.argument('file_path')
@click.argument('fps', type=click.IntRange(min=1))
def plot_frame_sizes_from_csv(file_path, fps):
    """A function to plot frame sizes and bitrate values from a CSV file."""
    """pip install bokeh pandas click"""
    df = load_csv(file_path)
    df = calculate_bitrates(df, fps)
    df = calculate_kbyte_frame_sizes(df)
    df = add_color_column(df)
    
    plot_frame_sizes(df)

if __name__ == '__main__':
    plot_frame_sizes_from_csv()