reference_lpips.py

"""
Compute LPIPS similarity between generated images and each prompt's reference image.

输入是 pia_extract.py 的输出目录：
    output_xxx/
      generated/<prompt_id>/seed_*.png
      results/

脚本会读取 prompts.json 中的 image 字段，在 results 下写出：
    reference_lpips_summary.jsonl
    reference_lpips_scores.jsonl
"""

import argparse
import json
import shutil
from pathlib import Path

import numpy as np
import torch
from PIL import Image
from tqdm import tqdm


def load_lpips_model(net: str, device: torch.device):
    try:
        import lpips
    except ImportError as exc:
        raise SystemExit("缺少 lpips 依赖，请先安装：pip install lpips") from exc

    model = lpips.LPIPS(net=net)
    model.eval().to(device)
    return model


def load_prompts(path: Path, max_prompts: int | None = None) -> list[dict]:
    with open(path, "r", encoding="utf-8") as f:
        prompts = json.load(f)
    if max_prompts is not None:
        prompts = prompts[:max_prompts]
    return prompts


def load_metadata(path: Path) -> dict[str, dict]:
    records = {}
    if not path.exists():
        return records

    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            if not line.strip():
                continue
            record = json.loads(line)
            records[Path(record["path"]).name] = record
    return records


def image_to_tensor(path: Path, image_size: int, device: torch.device) -> torch.Tensor:
    image = Image.open(path).convert("RGB").resize((image_size, image_size), Image.BICUBIC)
    array = np.asarray(image).astype(np.float32) / 127.5 - 1.0
    tensor = torch.from_numpy(array).permute(2, 0, 1).unsqueeze(0)
    return tensor.to(device=device)


@torch.no_grad()
def compute_lpips_scores(
    model,
    reference_path: Path,
    image_paths: list[Path],
    image_size: int,
    batch_size: int,
    device: torch.device,
) -> list[tuple[Path, float]]:
    reference = image_to_tensor(reference_path, image_size, device)
    scores = []

    for start in range(0, len(image_paths), batch_size):
        batch_paths = image_paths[start:start + batch_size]
        batch = torch.cat(
            [image_to_tensor(path, image_size, device) for path in batch_paths],
            dim=0,
        )
        reference_batch = reference.expand(batch.shape[0], -1, -1, -1)
        distances = model(reference_batch, batch).flatten().detach().float().cpu().tolist()
        scores.extend(zip(batch_paths, distances))

    return scores


def best_by_base_seed(score_records: list[dict], top_k: int) -> list[dict]:
    best = {}
    for record in score_records:
        base_seed = record.get("base_seed")
        if base_seed is None:
            continue
        previous = best.get(base_seed)
        if previous is None or record["lpips"] < previous["lpips"]:
            best[base_seed] = record
    return sorted(best.values(), key=lambda item: item["lpips"])[:top_k]


def write_jsonl(path: Path, records: list[dict]) -> None:
    with open(path, "w", encoding="utf-8") as f:
        for record in records:
            f.write(json.dumps(record, ensure_ascii=False) + "\n")


def lpips_tag(value: float) -> str:
    return f"{value:.4f}".replace(".", "p")


def copy_top_k_images(records: list[dict], topk_root: Path, prompt_id: str) -> list[dict]:
    prompt_topk_dir = topk_root / prompt_id
    prompt_topk_dir.mkdir(parents=True, exist_ok=True)

    copied = []
    for rank, record in enumerate(records, 1):
        src = Path(record["image_path"])
        dst_name = f"rank_{rank:02d}_lpips_{lpips_tag(record['lpips'])}_{src.name}"
        dst = prompt_topk_dir / dst_name
        shutil.copy2(src, dst)

        copied_record = dict(record)
        copied_record["rank"] = rank
        copied_record["saved_path"] = str(dst)
        copied.append(copied_record)

    return copied


def run(args: argparse.Namespace) -> Path:
    output_root = Path(args.output)
    generated_root = output_root / args.generated_subdir
    results_dir = output_root / args.results_subdir
    results_dir.mkdir(parents=True, exist_ok=True)
    topk_root = results_dir / args.top_k_image_dir

    device = torch.device(args.device)
    model = load_lpips_model(args.net, device)
    prompts = load_prompts(Path(args.prompts), args.max_prompts)
    reference_dir = Path(args.reference_dir)

    summary_records = []
    all_score_records = []

    for item in tqdm(prompts, desc="LPIPS reference matching"):
        prompt_id = str(item["id"])
        reference_name = item.get("image")
        if not reference_name:
            continue

        prompt_dir = generated_root / prompt_id
        reference_path = reference_dir / reference_name
        if not prompt_dir.exists() or not reference_path.exists():
            continue

        image_paths = sorted(prompt_dir.glob(args.image_glob))
        if not image_paths:
            continue

        metadata_by_name = load_metadata(prompt_dir / "metadata.jsonl")
        scores = compute_lpips_scores(
            model=model,
            reference_path=reference_path,
            image_paths=image_paths,
            image_size=args.image_size,
            batch_size=args.batch_size,
            device=device,
        )

        score_records = []
        for image_path, lpips_value in scores:
            metadata = metadata_by_name.get(image_path.name, {})
            record = {
                "id": prompt_id,
                "prompt": item.get("text"),
                "reference_image": reference_name,
                "reference_path": str(reference_path),
                "image_path": str(image_path),
                "image_name": image_path.name,
                "lpips": float(lpips_value),
                "base_seed": metadata.get("base_seed"),
                "branch_index": metadata.get("branch_index"),
                "output_index": metadata.get("output_index"),
            }
            score_records.append(record)

        score_records.sort(key=lambda record: record["lpips"])
        lpips_values = [record["lpips"] for record in score_records]
        top_k = score_records[:args.top_k]
        top_k_distinct_base_seed = best_by_base_seed(score_records, args.top_k)
        saved_top_k = []
        if args.save_top_k_images and top_k:
            saved_top_k = copy_top_k_images(top_k, topk_root, prompt_id)

        summary_records.append({
            "id": prompt_id,
            "prompt": item.get("text"),
            "reference_image": reference_name,
            "reference_path": str(reference_path),
            "num_images": len(score_records),
            "best_lpips": float(min(lpips_values)),
            "mean_lpips": float(np.mean(lpips_values)),
            "median_lpips": float(np.median(lpips_values)),
            "top_k": top_k,
            "top_k_distinct_base_seed": top_k_distinct_base_seed,
            "saved_top_k": saved_top_k,
        })
        all_score_records.extend(score_records)

    summary_path = results_dir / "reference_lpips_summary.jsonl"
    scores_path = results_dir / "reference_lpips_scores.jsonl"
    write_jsonl(summary_path, summary_records)
    write_jsonl(scores_path, all_score_records)

    print(f"LPIPS summary saved: {summary_path}")
    print(f"LPIPS scores saved: {scores_path}")
    if args.save_top_k_images:
        print(f"Top-k images saved under: {topk_root}")
    print(f"Processed prompts: {len(summary_records)}")
    return summary_path


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="计算生成图与 prompt reference image 的 LPIPS 距离")
    parser.add_argument("--output", type=str, required=True, help="pia_extract.py 输出目录")
    parser.add_argument("--prompts", type=str, default="prompts/prompts.json")
    parser.add_argument("--reference_dir", type=str, default="prompts/images")
    parser.add_argument("--generated_subdir", type=str, default="generated")
    parser.add_argument("--results_subdir", type=str, default="results")
    parser.add_argument("--image_glob", type=str, default="seed_*.png")
    parser.add_argument("--max_prompts", type=int, default=None)
    parser.add_argument("--top_k", type=int, default=10)
    parser.add_argument("--save_top_k_images", action=argparse.BooleanOptionalAction, default=True)
    parser.add_argument("--top_k_image_dir", type=str, default="reference_lpips_topk")
    parser.add_argument("--batch_size", type=int, default=16)
    parser.add_argument("--image_size", type=int, default=512)
    parser.add_argument("--device", type=str, default="cuda")
    parser.add_argument("--net", type=str, default="alex", choices=["alex", "vgg", "squeeze"])
    return parser.parse_args()


if __name__ == "__main__":
    run(parse_args())