pipelines/visualize_validation_comparison.py

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))