""" Compute LPIPS similarity between generated images and each prompt's reference image. 输入是 pia_extract.py 的输出目录： output_xxx/ generated//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())