from __future__ import annotations import argparse import json import math from datetime import datetime from pathlib import Path from typing import Any from PIL import Image, ImageDraw PROJECT_ROOT = Path(__file__).resolve().parent.parent RUNS_ROOT = PROJECT_ROOT / "runs" / "validation_compare" DEFAULT_REFERENCE_PATH = ( PROJECT_ROOT / "runs" / "validation_eval" / "qwen3_5_35b_a3b_exp_a_official" / "reference" / "reference_subset.jsonl" ) DEFAULT_EXP_A_PATH = ( PROJECT_ROOT / "runs" / "validation_eval" / "qwen3_5_35b_a3b_exp_a_official" / "predictions" / "predictions.jsonl" ) DEFAULT_EXP_B_PATH = ( PROJECT_ROOT / "runs" / "validation_eval" / "qwen3_5_35b_a3b_exp_b_attr" / "predictions" / "predictions.jsonl" ) DEFAULT_MAX_IMAGES = 128 PANEL_SPACING = 24 HEADER_HEIGHT = 84 FOOTER_HEIGHT = 54 BACKGROUND_COLOR = (245, 243, 238) TEXT_COLOR = (28, 26, 24) SUBTEXT_COLOR = (88, 84, 78) GT_COLOR = (28, 138, 68) PRED_COLOR = (208, 52, 52) PALETTE = [ (212, 76, 54), (34, 117, 168), (48, 145, 84), (198, 138, 39), (135, 84, 196), (28, 151, 156), ] def _load_jsonl(path: Path) -> list[dict[str, Any]]: rows: list[dict[str, Any]] = [] with path.open("r", encoding="utf-8") as file: for line in file: line = line.strip() if not line: continue rows.append(json.loads(line)) return rows def _group_reference_rows(reference_rows: list[dict[str, Any]]) -> dict[str, dict[str, Any]]: grouped: dict[str, dict[str, Any]] = {} for row in reference_rows: image_path = str(Path(row["image_path"]).resolve()) source_records = sorted( row.get("source_records", []), key=lambda item: int(item.get("target_index", 0)), ) grouped[image_path] = { "image_path": image_path, "source_records": source_records, "image_width": int(source_records[0]["image_width"]) if source_records else None, "image_height": int(source_records[0]["image_height"]) if source_records else None, } return grouped def _group_prediction_rows(prediction_rows: list[dict[str, Any]]) -> dict[str, dict[str, Any]]: grouped: dict[str, dict[str, Any]] = {} for row in prediction_rows: image_path = str(Path(row["image_path"]).resolve()) grouped[image_path] = row return grouped def _normalize_bbox(values: list[Any]) -> tuple[float, float, float, float] | None: if len(values) != 4: return None try: x1, y1, x2, y2 = [float(value) for value in values] except (TypeError, ValueError): return None if x2 <= x1 or y2 <= y1: return None return x1, y1, x2, y2 def _clamp_bbox( bbox: tuple[float, float, float, float], image_width: int, image_height: int, ) -> tuple[float, float, float, float]: x1, y1, x2, y2 = bbox x1 = max(0.0, min(float(image_width - 1), x1)) y1 = max(0.0, min(float(image_height - 1), y1)) x2 = max(x1 + 1.0, min(float(image_width), x2)) y2 = max(y1 + 1.0, min(float(image_height), y2)) return x1, y1, x2, y2 def _fit_image(image: Image.Image, max_width: int, max_height: int) -> Image.Image: scale = min(max_width / image.width, max_height / image.height) scale = min(scale, 1.0) new_width = max(1, int(round(image.width * scale))) new_height = max(1, int(round(image.height * scale))) if new_width == image.width and new_height == image.height: return image.copy() return image.resize((new_width, new_height), Image.Resampling.LANCZOS) def _draw_label(draw: ImageDraw.ImageDraw, x1: float, y1: float, text: str, color: tuple[int, int, int]) -> None: text = text[:72] bbox = draw.textbbox((0, 0), text) padding_x = 6 padding_y = 4 bg_box = ( x1, max(0.0, y1 - (bbox[3] - bbox[1]) - 2 * padding_y), x1 + (bbox[2] - bbox[0]) + 2 * padding_x, y1, ) draw.rectangle(bg_box, fill=color) draw.text((bg_box[0] + padding_x, bg_box[1] + padding_y), text, fill=(255, 255, 255)) def _draw_boxes( image: Image.Image, annotations: list[dict[str, Any]], *, default_color: tuple[int, int, int], use_palette: bool, ) -> Image.Image: preview = image.copy() draw = ImageDraw.Draw(preview) line_width = max(3, int(round(min(preview.size) / 170))) for index, annotation in enumerate(annotations, start=1): bbox_values = annotation.get("bbox") bbox = _normalize_bbox(bbox_values if isinstance(bbox_values, list) else []) if bbox is None: continue bbox = _clamp_bbox(bbox, preview.width, preview.height) color = PALETTE[(index - 1) % len(PALETTE)] if use_palette else default_color draw.rectangle(bbox, outline=color, width=line_width) target_index = annotation.get("target_index", index - 1) maturity = annotation.get("maturity_level") occlusion = annotation.get("occlusion_degree") label_parts = [f"#{target_index}"] if maturity: label_parts.append(str(maturity)) if occlusion: label_parts.append(f"遮挡:{occlusion}") _draw_label(draw, bbox[0], bbox[1], " | ".join(label_parts), color) return preview def _extract_gt_annotations(reference_row: dict[str, Any]) -> list[dict[str, Any]]: annotations: list[dict[str, Any]] = [] for record in reference_row.get("source_records", []): bbox = record.get("bbox") if not isinstance(bbox, list): continue annotations.append( { "target_index": record.get("target_index"), "bbox": bbox, "maturity_level": record.get("maturity_level"), "occlusion_degree": record.get("occlusion_degree"), } ) return annotations def _build_header( canvas: Image.Image, *, title: str, image_name: str, gt_count: int, exp_a_count: int, exp_b_count: int, exp_a_error: str | None, exp_b_error: str | None, ) -> None: draw = ImageDraw.Draw(canvas) draw.text((24, 18), title, fill=TEXT_COLOR) meta = f"{image_name} | GT={gt_count} | A={exp_a_count} | B={exp_b_count}" draw.text((24, 42), meta, fill=SUBTEXT_COLOR) if exp_a_error: draw.text((24, 62), f"A parse_error: {exp_a_error[:120]}", fill=(150, 58, 42)) if exp_b_error: draw.text((24, 80), f"B parse_error: {exp_b_error[:120]}", fill=(150, 58, 42)) def _build_footer(draw: ImageDraw.ImageDraw, top_y: int, labels: list[str], panel_width: int) -> None: for index, label in enumerate(labels): x = 24 + index * (panel_width + PANEL_SPACING) draw.text((x, top_y), label, fill=TEXT_COLOR) def _build_triptych( *, image_path: Path, gt_annotations: list[dict[str, Any]], exp_a_row: dict[str, Any] | None, exp_b_row: dict[str, Any] | None, output_path: Path, ) -> dict[str, Any]: with Image.open(image_path).convert("RGB") as image: display_height = max(320, min(720, image.height)) display_width = max(320, min(960, image.width)) gt_panel = _fit_image( _draw_boxes(image, gt_annotations, default_color=GT_COLOR, use_palette=True), display_width, display_height, ) exp_a_annotations = exp_a_row.get("annotations", []) if exp_a_row else [] exp_b_annotations = exp_b_row.get("annotations", []) if exp_b_row else [] exp_a_panel = _fit_image( _draw_boxes(image, exp_a_annotations, default_color=PRED_COLOR, use_palette=False), display_width, display_height, ) exp_b_panel = _fit_image( _draw_boxes(image, exp_b_annotations, default_color=PRED_COLOR, use_palette=False), display_width, display_height, ) panel_width = max(gt_panel.width, exp_a_panel.width, exp_b_panel.width) panel_height = max(gt_panel.height, exp_a_panel.height, exp_b_panel.height) canvas_width = 24 * 2 + panel_width * 3 + PANEL_SPACING * 2 canvas_height = HEADER_HEIGHT + panel_height + FOOTER_HEIGHT + 24 canvas = Image.new("RGB", (canvas_width, canvas_height), BACKGROUND_COLOR) _build_header( canvas, title="Validation Comparison: GT vs Exp A vs Exp B", image_name=image_path.name, gt_count=len(gt_annotations), exp_a_count=len(exp_a_annotations), exp_b_count=len(exp_b_annotations), exp_a_error=exp_a_row.get("parse_error") if exp_a_row else "missing prediction row", exp_b_error=exp_b_row.get("parse_error") if exp_b_row else "missing prediction row", ) panel_top = HEADER_HEIGHT panel_lefts = [ 24, 24 + panel_width + PANEL_SPACING, 24 + (panel_width + PANEL_SPACING) * 2, ] for left, panel in zip(panel_lefts, [gt_panel, exp_a_panel, exp_b_panel], strict=True): paste_x = left + (panel_width - panel.width) // 2 paste_y = panel_top + (panel_height - panel.height) // 2 canvas.paste(panel, (paste_x, paste_y)) footer_draw = ImageDraw.Draw(canvas) _build_footer( footer_draw, HEADER_HEIGHT + panel_height + 14, ["GT", "Experiment A", "Experiment B"], panel_width, ) output_path.parent.mkdir(parents=True, exist_ok=True) canvas.save(output_path, quality=95) return { "image_path": str(image_path), "comparison_image": str(output_path), "gt_count": len(gt_annotations), "exp_a_count": len(exp_a_annotations), "exp_b_count": len(exp_b_annotations), "exp_a_parse_error": exp_a_row.get("parse_error") if exp_a_row else "missing prediction row", "exp_b_parse_error": exp_b_row.get("parse_error") if exp_b_row else "missing prediction row", } def build_validation_comparison( *, reference_path: Path, exp_a_predictions_path: Path, exp_b_predictions_path: Path, output_dir: Path | None = None, max_images: int = DEFAULT_MAX_IMAGES, ) -> dict[str, Any]: reference_rows = _load_jsonl(reference_path) exp_a_rows = _group_prediction_rows(_load_jsonl(exp_a_predictions_path)) exp_b_rows = _group_prediction_rows(_load_jsonl(exp_b_predictions_path)) reference_grouped = _group_reference_rows(reference_rows) timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") run_dir = output_dir or RUNS_ROOT / f"exp_a_vs_b_{timestamp}" run_dir.mkdir(parents=True, exist_ok=True) image_dir = run_dir / "images" image_dir.mkdir(parents=True, exist_ok=True) manifest_rows: list[dict[str, Any]] = [] sorted_image_paths = sorted(reference_grouped.keys())[:max_images] digits = max(3, int(math.log10(max(1, len(sorted_image_paths)))) + 1) for index, image_path_str in enumerate(sorted_image_paths, start=1): image_path = Path(image_path_str) reference_row = reference_grouped[image_path_str] gt_annotations = _extract_gt_annotations(reference_row) comparison_path = image_dir / f"{index:0{digits}d}__{image_path.stem}.jpg" manifest_rows.append( _build_triptych( image_path=image_path, gt_annotations=gt_annotations, exp_a_row=exp_a_rows.get(image_path_str), exp_b_row=exp_b_rows.get(image_path_str), output_path=comparison_path, ) ) manifest_path = run_dir / "manifest.jsonl" with manifest_path.open("w", encoding="utf-8") as file: for row in manifest_rows: file.write(json.dumps(row, ensure_ascii=False)) file.write("\n") summary = { "created_at": datetime.now().isoformat(timespec="seconds"), "reference_path": str(reference_path.resolve()), "exp_a_predictions_path": str(exp_a_predictions_path.resolve()), "exp_b_predictions_path": str(exp_b_predictions_path.resolve()), "image_count": len(manifest_rows), "run_dir": str(run_dir), "image_dir": str(image_dir), "manifest_path": str(manifest_path), } (run_dir / "summary.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8") return summary def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="生成实验 A / 实验 B / GT 的验证集可视化对比图。") parser.add_argument("--reference-path", default=str(DEFAULT_REFERENCE_PATH), help="reference_subset.jsonl 路径") parser.add_argument("--exp-a-predictions-path", default=str(DEFAULT_EXP_A_PATH), help="实验 A predictions.jsonl") parser.add_argument("--exp-b-predictions-path", default=str(DEFAULT_EXP_B_PATH), help="实验 B predictions.jsonl") parser.add_argument("--output-dir", default=None, help="输出目录；默认写入 runs/validation_compare") parser.add_argument("--max-images", type=int, default=DEFAULT_MAX_IMAGES, help="最多导出多少张图") return parser.parse_args() if __name__ == "__main__": args = parse_args() summary = build_validation_comparison( reference_path=Path(args.reference_path).expanduser().resolve(), exp_a_predictions_path=Path(args.exp_a_predictions_path).expanduser().resolve(), exp_b_predictions_path=Path(args.exp_b_predictions_path).expanduser().resolve(), output_dir=Path(args.output_dir).expanduser().resolve() if args.output_dir else None, max_images=args.max_images, ) print(json.dumps(summary, ensure_ascii=False, indent=2))