#!/usr/bin/env python3 import argparse import json import os import time from pathlib import Path import lpips import numpy as np import torch import torch.distributed as dist from decord import VideoReader, cpu from skimage.metrics import structural_similarity def parse_args(): parser = argparse.ArgumentParser( description="Compare paired videos frame by frame with PSNR, SSIM, and LPIPS." ) parser.add_argument("--reference_dir", type=Path, required=True) parser.add_argument("--comparison_dir", type=Path, required=True) parser.add_argument("--output_json", type=Path, required=True) parser.add_argument("--batch_size", type=int, default=8) return parser.parse_args() def init_distributed(): if "LOCAL_RANK" not in os.environ: return 0, 0, 1, torch.device("cuda") local_rank = int(os.environ["LOCAL_RANK"]) torch.cuda.set_device(local_rank) dist.init_process_group(backend="nccl") return dist.get_rank(), local_rank, dist.get_world_size(), torch.device( f"cuda:{local_rank}" ) def list_video_pairs(reference_dir, comparison_dir): reference = {path.name: path for path in reference_dir.glob("*.mp4")} comparison = {path.name: path for path in comparison_dir.glob("*.mp4")} if reference.keys() != comparison.keys(): missing = sorted(reference.keys() - comparison.keys()) extra = sorted(comparison.keys() - reference.keys()) raise ValueError(f"Video sets do not match: missing={missing}, extra={extra}") if not reference: raise ValueError("No paired MP4 videos found") return [(reference[name], comparison[name]) for name in sorted(reference)] @torch.inference_mode() def compare_video(reference_path, comparison_path, lpips_model, device, batch_size): reference_video = VideoReader(str(reference_path), ctx=cpu(0)) comparison_video = VideoReader(str(comparison_path), ctx=cpu(0)) if len(reference_video) != len(comparison_video): raise ValueError( f"Frame-count mismatch for {reference_path.name}: " f"{len(reference_video)} vs. {len(comparison_video)}" ) psnr_values = [] ssim_values = [] lpips_values = [] for start in range(0, len(reference_video), batch_size): end = min(start + batch_size, len(reference_video)) indices = list(range(start, end)) reference_frames = reference_video.get_batch(indices).asnumpy() comparison_frames = comparison_video.get_batch(indices).asnumpy() if reference_frames.shape != comparison_frames.shape: raise ValueError( f"Frame-shape mismatch for {reference_path.name}: " f"{reference_frames.shape} vs. {comparison_frames.shape}" ) reference_tensor = ( torch.from_numpy(np.ascontiguousarray(reference_frames)) .permute(0, 3, 1, 2) .to(device=device, dtype=torch.float32) .div_(255.0) ) comparison_tensor = ( torch.from_numpy(np.ascontiguousarray(comparison_frames)) .permute(0, 3, 1, 2) .to(device=device, dtype=torch.float32) .div_(255.0) ) mse = torch.mean( (reference_tensor - comparison_tensor) ** 2, dim=(1, 2, 3) ) batch_psnr = -10.0 * torch.log10(mse) if not torch.isfinite(batch_psnr).all(): raise ValueError(f"Non-finite PSNR encountered in {reference_path.name}") psnr_values.extend(batch_psnr.cpu().tolist()) batch_lpips = lpips_model( reference_tensor.mul(2.0).sub(1.0), comparison_tensor.mul(2.0).sub(1.0), ) lpips_values.extend(batch_lpips.flatten().cpu().tolist()) for reference_frame, comparison_frame in zip( reference_frames, comparison_frames ): ssim_values.append( float( structural_similarity( reference_frame, comparison_frame, channel_axis=2, data_range=255, ) ) ) return { "video": reference_path.name, "frames": len(reference_video), "psnr": float(np.mean(psnr_values)), "ssim": float(np.mean(ssim_values)), "lpips_alex": float(np.mean(lpips_values)), } def main(): args = parse_args() if args.batch_size <= 0: raise ValueError("--batch_size must be positive") rank, local_rank, world_size, device = init_distributed() started_at = time.perf_counter() pairs = list_video_pairs(args.reference_dir, args.comparison_dir) local_pairs = pairs[rank::world_size] lpips_model = lpips.LPIPS(net="alex").eval().to(device) local_results = [] for index, (reference_path, comparison_path) in enumerate(local_pairs, start=1): local_results.append( compare_video( reference_path, comparison_path, lpips_model, device, args.batch_size, ) ) print( f"[rank {rank}] {index}/{len(local_pairs)} {reference_path.name}", flush=True, ) if dist.is_initialized(): gathered = [None] * world_size if rank == 0 else None dist.gather_object(local_results, gathered, dst=0) if rank == 0: all_results = [ result for rank_results in gathered for result in rank_results ] else: all_results = None else: all_results = local_results if rank == 0: all_results.sort(key=lambda item: item["video"]) frame_count = sum(item["frames"] for item in all_results) weighted_sums = { metric: sum(item[metric] * item["frames"] for item in all_results) for metric in ("psnr", "ssim", "lpips_alex") } aggregate = { metric: weighted_sums[metric] / frame_count for metric in weighted_sums } payload = { "reference_dir": str(args.reference_dir.resolve()), "comparison_dir": str(args.comparison_dir.resolve()), "video_count": len(all_results), "frame_count": frame_count, "aggregation": "arithmetic mean over paired frames", "method": { "psnr": "RGB, data_range=1, per-frame", "ssim": ( "skimage.metrics.structural_similarity, RGB channel_axis=2, " "data_range=255, default window" ), "lpips": "lpips 0.1.4, AlexNet, RGB in [-1, 1], full resolution", }, "aggregate": aggregate, "elapsed_seconds": time.perf_counter() - started_at, "per_video": all_results, } args.output_json.parent.mkdir(parents=True, exist_ok=True) args.output_json.write_text(json.dumps(payload, indent=2) + "\n") print(json.dumps(payload["aggregate"], indent=2), flush=True) print(f"Saved metrics to {args.output_json}", flush=True) if dist.is_initialized(): dist.barrier(device_ids=[local_rank]) dist.destroy_process_group() if __name__ == "__main__": main()