Buckets:

fishxinyu
/

flexcombine-bench

Files

xet

fishxinyu/flexcombine-bench / evaluation /VideoAlign /evaluate_dataset.py

fishxinyu

2 months ago

download

raw

6.31 kB

"""Evaluate generated videos against a prompt file using the VideoReward model.
Given a folder of generated videos and a prompt file mapping sample IDs to text
prompts (e.g. datasets/v0/pose_depth/captions.json), scores each video on
Visual Quality (VQ), Motion Quality (MQ), Text Alignment (TA), and their sum
(Overall).
Videos are matched to prompts by filename stem, e.g. <generated-dir>/001.mp4
is matched against captions["001"].
Usage:
    python evaluate_dataset.py \\
        --generated-dir /home/xinyuy/dataset_processing/flexcombine-bench/flexbench/results/v0/pose_depth \\
        --prompt-file /home/xinyuy/dataset_processing/flexcombine-bench/datasets/v0/pose_depth/captions.json \\
        --load-from-pretrained ./checkpoints \\
        --output-dir ./results/pose_depth
"""

import argparse
import csv
import json
from pathlib import Path

import torch
from tqdm import tqdm

from inference import VideoVLMRewardInference

REWARD_KEYS = ["VQ", "MQ", "TA", "Overall"]


def load_prompts(prompt_file: Path) -> dict[str, str]:
    with open(prompt_file, "r") as f:
        prompts = json.load(f)
    return {str(k): v for k, v in prompts.items()}


def evaluate_dataset(
    generated_dir: Path,
    prompt_file: Path,
    load_from_pretrained: str,
    load_from_pretrained_step: int = -1,
    device: str = "cuda:0",
    dtype: torch.dtype = torch.bfloat16,
    batch_size: int = 8,
    use_norm: bool = True,
    output_dir: Path | None = None,
) -> None:
    prompts_by_id = load_prompts(prompt_file)

    generated_files = {f.stem: f for f in sorted(generated_dir.glob("*.mp4"))}
    print(len(generated_files), "generated video(s) found in", generated_dir)
    sample_ids = sorted(set(generated_files) & set(prompts_by_id))

    missing = sorted(set(generated_files) - set(prompts_by_id))
    if missing:
        print(f"[warn] {len(missing)} generated video(s) have no matching prompt, skipping: {missing[:5]}{'...' if len(missing) > 5 else ''}")

    if not sample_ids:
        print(f"No matching samples between {generated_dir} and {prompt_file}")
        return

    print(f"Loading VideoReward model from {load_from_pretrained}...")
    inferencer = VideoVLMRewardInference(
        load_from_pretrained,
        load_from_pretrained_step=load_from_pretrained_step,
        device=device,
        dtype=dtype,
    )

    print(f"Evaluating {len(sample_ids)} sample(s)...")
    per_sample: dict[str, dict] = {}

    for i in tqdm(range(0, len(sample_ids), batch_size)):
        batch_ids = sample_ids[i : i + batch_size]
        batch_video_paths = [str(generated_files[sid]) for sid in batch_ids]
        batch_prompts = [prompts_by_id[sid] for sid in batch_ids]

        with torch.no_grad():
            rewards = inferencer.reward(batch_video_paths, batch_prompts, use_norm=use_norm)

        for sid, reward in zip(batch_ids, rewards):
            per_sample[sid] = {k: reward[k] for k in REWARD_KEYS}

    mean_rewards = {
        k: sum(per_sample[sid][k] for sid in sample_ids) / len(sample_ids) for k in REWARD_KEYS
    }

    print(f"\n{'=' * 50}")
    print(f"Mean metrics over {len(sample_ids)} sample(s):")
    for k, v in mean_rewards.items():
        print(f"  {k}: {v:.4f}")

    if output_dir is not None:
        output_dir.mkdir(parents=True, exist_ok=True)

        json_path = output_dir / "results.json"
        result = {
            "mean": {k: round(v, 4) for k, v in mean_rewards.items()},
            "per_sample": {
                sid: {k: round(v, 4) for k, v in vals.items()} for sid, vals in per_sample.items()
            },
        }
        with open(json_path, "w") as f:
            json.dump(result, f, indent=2)

        csv_path = output_dir / "results.csv"
        with open(csv_path, "w", newline="") as f:
            writer = csv.writer(f)
            writer.writerow(["sample_id"] + REWARD_KEYS)
            writer.writerow(["average"] + [round(mean_rewards[k], 4) for k in REWARD_KEYS])
            for sid in sample_ids:
                writer.writerow([sid] + [round(per_sample[sid][k], 4) for k in REWARD_KEYS])

        print(f"Results saved to {output_dir}")


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Evaluate generated videos with VideoReward against a prompt file.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "--generated-dir",
        required=True,
        help="Directory containing generated .mp4 files (one per sample ID).",
    )
    parser.add_argument(
        "--prompt-file",
        required=True,
        help=(
            "JSON file mapping sample ID -> text prompt, e.g. "
            "datasets/v0/pose_depth/captions.json."
        ),
    )
    parser.add_argument(
        "--load-from-pretrained",
        default="./checkpoints",
        help="Path to the VideoReward checkpoint directory.",
    )
    parser.add_argument(
        "--load-from-pretrained-step",
        type=int,
        default=-1,
        help="Checkpoint step to load, -1 for the latest.",
    )
    parser.add_argument("--device", default="cuda:0", help="Device to run inference on.")
    parser.add_argument(
        "--dtype",
        default="bf16",
        choices=["bf16", "fp16", "fp32"],
        help="Model dtype.",
    )
    parser.add_argument("--batch-size", type=int, default=8, help="Inference batch size.")
    parser.add_argument(
        "--no-norm",
        action="store_true",
        help="Disable rescaling rewards using the checkpoint's inference_config stats.",
    )
    parser.add_argument(
        "--output-dir",
        default=None,
        help="Optional folder to write results.json and results.csv.",
    )

    args = parser.parse_args()

    dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[args.dtype]

    evaluate_dataset(
        generated_dir=Path(args.generated_dir),
        prompt_file=Path(args.prompt_file),
        load_from_pretrained=args.load_from_pretrained,
        load_from_pretrained_step=args.load_from_pretrained_step,
        device=args.device,
        dtype=dtype,
        batch_size=args.batch_size,
        use_norm=not args.no_norm,
        output_dir=Path(args.output_dir) if args.output_dir else None,
    )


if __name__ == "__main__":
    main()

Xet Storage Details

Size:: 6.31 kB
Xet hash:: a812b60f6c4b025f5b2cdc54fc8561c8a739f2b2448e33bd892e4afd80e88158

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	"""Evaluate generated videos against a prompt file using the VideoReward model.

	Given a folder of generated videos and a prompt file mapping sample IDs to text
	prompts (e.g. datasets/v0/pose_depth/captions.json), scores each video on
	Visual Quality (VQ), Motion Quality (MQ), Text Alignment (TA), and their sum
	(Overall).

	Videos are matched to prompts by filename stem, e.g. <generated-dir>/001.mp4
	is matched against captions["001"].

	Usage:
	python evaluate_dataset.py \\
	--generated-dir /home/xinyuy/dataset_processing/flexcombine-bench/flexbench/results/v0/pose_depth \\
	--prompt-file /home/xinyuy/dataset_processing/flexcombine-bench/datasets/v0/pose_depth/captions.json \\
	--load-from-pretrained ./checkpoints \\
	--output-dir ./results/pose_depth
	"""

	import argparse
	import csv
	import json
	from pathlib import Path

	import torch
	from tqdm import tqdm

	from inference import VideoVLMRewardInference

	REWARD_KEYS = ["VQ", "MQ", "TA", "Overall"]


	def load_prompts(prompt_file: Path) -> dict[str, str]:
	with open(prompt_file, "r") as f:
	prompts = json.load(f)
	return {str(k): v for k, v in prompts.items()}


	def evaluate_dataset(
	generated_dir: Path,
	prompt_file: Path,
	load_from_pretrained: str,
	load_from_pretrained_step: int = -1,
	device: str = "cuda:0",
	dtype: torch.dtype = torch.bfloat16,
	batch_size: int = 8,
	use_norm: bool = True,
	output_dir: Path \| None = None,
	) -> None:
	prompts_by_id = load_prompts(prompt_file)

	generated_files = {f.stem: f for f in sorted(generated_dir.glob("*.mp4"))}
	print(len(generated_files), "generated video(s) found in", generated_dir)
	sample_ids = sorted(set(generated_files) & set(prompts_by_id))

	missing = sorted(set(generated_files) - set(prompts_by_id))
	if missing:
	print(f"[warn] {len(missing)} generated video(s) have no matching prompt, skipping: {missing[:5]}{'...' if len(missing) > 5 else ''}")

	if not sample_ids:
	print(f"No matching samples between {generated_dir} and {prompt_file}")
	return

	print(f"Loading VideoReward model from {load_from_pretrained}...")
	inferencer = VideoVLMRewardInference(
	load_from_pretrained,
	load_from_pretrained_step=load_from_pretrained_step,
	device=device,
	dtype=dtype,
	)

	print(f"Evaluating {len(sample_ids)} sample(s)...")
	per_sample: dict[str, dict] = {}

	for i in tqdm(range(0, len(sample_ids), batch_size)):
	batch_ids = sample_ids[i : i + batch_size]
	batch_video_paths = [str(generated_files[sid]) for sid in batch_ids]
	batch_prompts = [prompts_by_id[sid] for sid in batch_ids]

	with torch.no_grad():
	rewards = inferencer.reward(batch_video_paths, batch_prompts, use_norm=use_norm)

	for sid, reward in zip(batch_ids, rewards):
	per_sample[sid] = {k: reward[k] for k in REWARD_KEYS}

	mean_rewards = {
	k: sum(per_sample[sid][k] for sid in sample_ids) / len(sample_ids) for k in REWARD_KEYS
	}

	print(f"\n{'=' * 50}")
	print(f"Mean metrics over {len(sample_ids)} sample(s):")
	for k, v in mean_rewards.items():
	print(f" {k}: {v:.4f}")

	if output_dir is not None:
	output_dir.mkdir(parents=True, exist_ok=True)

	json_path = output_dir / "results.json"
	result = {
	"mean": {k: round(v, 4) for k, v in mean_rewards.items()},
	"per_sample": {
	sid: {k: round(v, 4) for k, v in vals.items()} for sid, vals in per_sample.items()
	},
	}
	with open(json_path, "w") as f:
	json.dump(result, f, indent=2)

	csv_path = output_dir / "results.csv"
	with open(csv_path, "w", newline="") as f:
	writer = csv.writer(f)
	writer.writerow(["sample_id"] + REWARD_KEYS)
	writer.writerow(["average"] + [round(mean_rewards[k], 4) for k in REWARD_KEYS])
	for sid in sample_ids:
	writer.writerow([sid] + [round(per_sample[sid][k], 4) for k in REWARD_KEYS])

	print(f"Results saved to {output_dir}")


	def main() -> None:
	parser = argparse.ArgumentParser(
	description="Evaluate generated videos with VideoReward against a prompt file.",
	formatter_class=argparse.ArgumentDefaultsHelpFormatter,
	)
	parser.add_argument(
	"--generated-dir",
	required=True,
	help="Directory containing generated .mp4 files (one per sample ID).",
	)
	parser.add_argument(
	"--prompt-file",
	required=True,
	help=(
	"JSON file mapping sample ID -> text prompt, e.g. "
	"datasets/v0/pose_depth/captions.json."
	),
	)
	parser.add_argument(
	"--load-from-pretrained",
	default="./checkpoints",
	help="Path to the VideoReward checkpoint directory.",
	)
	parser.add_argument(
	"--load-from-pretrained-step",
	type=int,
	default=-1,
	help="Checkpoint step to load, -1 for the latest.",
	)
	parser.add_argument("--device", default="cuda:0", help="Device to run inference on.")
	parser.add_argument(
	"--dtype",
	default="bf16",
	choices=["bf16", "fp16", "fp32"],
	help="Model dtype.",
	)
	parser.add_argument("--batch-size", type=int, default=8, help="Inference batch size.")
	parser.add_argument(
	"--no-norm",
	action="store_true",
	help="Disable rescaling rewards using the checkpoint's inference_config stats.",
	)
	parser.add_argument(
	"--output-dir",
	default=None,
	help="Optional folder to write results.json and results.csv.",
	)

	args = parser.parse_args()

	dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[args.dtype]

	evaluate_dataset(
	generated_dir=Path(args.generated_dir),
	prompt_file=Path(args.prompt_file),
	load_from_pretrained=args.load_from_pretrained,
	load_from_pretrained_step=args.load_from_pretrained_step,
	device=args.device,
	dtype=dtype,
	batch_size=args.batch_size,
	use_norm=not args.no_norm,
	output_dir=Path(args.output_dir) if args.output_dir else None,
	)


	if __name__ == "__main__":
	main()