compare_ffprobe_to_annolid.py

import tqdm
import subprocess
from pathlib import Path
import csv

CSV_NAME = "_tracked_annotated"


def get_csv_mp4(root: Path) -> list[Path]:
    csvs = set()
    paths = []

    for f in root.glob("*.csv"):
        csvs.add(f.with_suffix("").name)
    for f in root.glob("*.mp4"):
        if f.with_suffix("").name + CSV_NAME in csvs:
            paths.append(f.with_suffix(""))

    return paths


def parse_csv(filename: Path, instance_name: str = "mouse") -> tuple[int, int, float | None]:
    with open(filename, "r") as fh:
        reader = csv.reader(fh, delimiter=",")
        lines = list(reader)[1:]

    nums = []
    times = []
    for line in lines:
        n, inst, _, _, _, _, *rt = line
        if inst == instance_name:
            nums.append(int(n))
            if rt:
                times.append(float(rt[0]))
    return len(nums), max(nums) + 1, max(times) if times else None


def parse_video(filename: Path, ffprobe: Path) -> tuple[int, float, float]:
    cmd = [
        str(ffprobe), "-v", "error", "-select_streams", "v:0", "-count_frames", "-show_entries",
        "stream=r_frame_rate,nb_read_frames", "-show_entries", "format=duration", "-of",
        "default=noprint_wrappers=1:nokey=1", str(filename),
    ]
    p = subprocess.run(cmd, check=True, capture_output=True).stdout
    rate_f, count, duration = p.decode().splitlines()
    num, denum = map(int, rate_f.split("/"))
    rate = num / denum

    return int(count), rate, float(duration)


def export_frames(root: Path, output_csv: Path, ffprobe: Path) -> None:
    header = [
        "filename", "csv_n_lines", "csv_max_frames", "csv_duration", "csv_calculated_fps", "ff_n_frames", "ff_duration",
        "ff_reported_fps", "ff_calculated_fps"
    ]
    lines = []
    for filename in tqdm.tqdm(get_csv_mp4(root)):
        csv_n_lines, csv_max_frames, csv_dur = parse_csv(filename.parent / (filename.name + f"{CSV_NAME}.csv"))
        ff_n, ff_rate, ff_dur = parse_video(filename.with_suffix(".mp4"), ffprobe)

        line = [
            filename.name, csv_n_lines, csv_max_frames, csv_dur, csv_max_frames / csv_dur if csv_dur else 0,
            ff_n, ff_dur, ff_rate, ff_n / ff_dur,
        ]
        lines.append(list(map(str, line)))

    with open(output_csv, "w") as fh:
        fh.write(",".join(header))
        fh.write("\n")
        for line in lines:
            fh.write(",".join(line))
            fh.write("\n")