pipelines/apply_missing_label_patch_to_reference.py

from __future__ import annotations

import argparse
import json
from collections import Counter, defaultdict
from datetime import datetime
from pathlib import Path
from typing import Any


PROJECT_ROOT = Path(__file__).resolve().parents[1]
DEFAULT_REFERENCE_PATH = (
    PROJECT_ROOT
    / "runs/validation_eval/qwen3_5_35b_a3b_exp_a_official/reference/reference_subset.jsonl"
)
DEFAULT_PATCH_PATH = (
    PROJECT_ROOT
    / "runs/analysis/missing_label_candidates_20260524/review_export/accepted_missing_label_patch.jsonl"
)
DEFAULT_OUTPUT_DIR = PROJECT_ROOT / "runs/gold/gold_release_v1_1_reviewed_val"


def main() -> None:
    """把人工接受的漏标 patch 应用到验证 reference 副本。"""
    parser = argparse.ArgumentParser(description="生成补漏标后的验证 reference。")
    parser.add_argument("--reference-path", default=str(DEFAULT_REFERENCE_PATH))
    parser.add_argument("--patch-path", default=str(DEFAULT_PATCH_PATH))
    parser.add_argument("--output-dir", default=str(DEFAULT_OUTPUT_DIR))
    args = parser.parse_args()

    reference_path = Path(args.reference_path).expanduser().resolve()
    patch_path = Path(args.patch_path).expanduser().resolve()
    output_dir = Path(args.output_dir).expanduser().resolve()
    output_dir.mkdir(parents=True, exist_ok=True)

    reference_rows = _read_jsonl(reference_path)
    patch_rows = _read_jsonl(patch_path)
    patched_rows, applied_rows = _apply_patch(reference_rows=reference_rows, patch_rows=patch_rows)

    reference_out = output_dir / "reference_subset_reviewed.jsonl"
    patch_applied_out = output_dir / "applied_missing_label_patch.jsonl"
    summary_out = output_dir / "summary.json"
    report_out = output_dir / "report.md"

    _write_jsonl(reference_out, patched_rows)
    _write_jsonl(patch_applied_out, applied_rows)

    summary = _build_summary(
        reference_path=reference_path,
        patch_path=patch_path,
        output_dir=output_dir,
        reference_out=reference_out,
        patch_applied_out=patch_applied_out,
        reference_rows=reference_rows,
        patched_rows=patched_rows,
        patch_rows=patch_rows,
        applied_rows=applied_rows,
    )
    summary_out.write_text(json.dumps(summary, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    report_out.write_text(_build_report(summary), encoding="utf-8")
    print(json.dumps(summary, ensure_ascii=False, indent=2))


def _apply_patch(
    *,
    reference_rows: list[dict[str, Any]],
    patch_rows: list[dict[str, Any]],
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
    rows_by_image_key = {Path(str(row["image_path"])).name: row for row in reference_rows}
    patches_by_image_key: defaultdict[str, list[dict[str, Any]]] = defaultdict(list)
    for patch_row in patch_rows:
        patch_type = str(patch_row.get("patch_type", "add_missing_label"))
        if patch_type != "add_missing_label":
            raise ValueError(f"当前 patch 应用器只允许安全追加 add_missing_label，收到: {patch_type}")
        patches_by_image_key[str(patch_row["image_key"])].append(patch_row)

    applied_rows = []
    patched_rows = []
    for row in reference_rows:
        image_key = Path(str(row["image_path"])).name
        row_copy = json.loads(json.dumps(row, ensure_ascii=False))
        source_records = _extract_source_records(row_copy)
        existing_count = len(source_records)
        image_patches = sorted(
            patches_by_image_key.get(image_key, []),
            key=lambda item: int(item.get("target_index", 0)),
        )
        for patch_index, patch_row in enumerate(image_patches):
            source_record = _patch_to_source_record(patch_row, fallback_target_index=existing_count + patch_index)
            source_records.append(source_record)
            applied_rows.append(
                {
                    "candidate_id": patch_row["candidate_id"],
                    "image_key": image_key,
                    "assigned_target_index": source_record["target_index"],
                    "bbox_1000": source_record["bbox_1000"],
                    "maturity_level": source_record["maturity_level"],
                    "occlusion_degree": source_record["occlusion_degree"],
                    "support_count": patch_row.get("support_count"),
                    "support_runs": patch_row.get("support_runs", []),
                }
            )
        source_records = sorted(source_records, key=lambda item: int(item["target_index"]))
        row_copy["source_records"] = source_records
        row_copy["messages"] = _replace_assistant_annotations(row_copy.get("messages", []), source_records)
        row_copy["review_patch_metadata"] = {
            "patch_version": "missing_label_review_20260524",
            "applied_missing_labels": len(image_patches),
            "original_annotation_count": existing_count,
            "reviewed_annotation_count": len(source_records),
        }
        patched_rows.append(row_copy)

    missing_image_keys = sorted(set(patches_by_image_key) - set(rows_by_image_key))
    if missing_image_keys:
        raise ValueError(f"patch 中存在 reference 未覆盖图像: {missing_image_keys[:5]}")
    return patched_rows, applied_rows


def _patch_to_source_record(patch_row: dict[str, Any], fallback_target_index: int) -> dict[str, Any]:
    target_index = int(patch_row.get("target_index", fallback_target_index))
    return {
        "image_path": patch_row["image_path"],
        "image_width": int(patch_row["image_width"]),
        "image_height": int(patch_row["image_height"]),
        "target_index": target_index,
        "bbox": patch_row["bbox"],
        "maturity_level": patch_row["maturity_level"],
        "maturity_ratio": patch_row["maturity_ratio"],
        "occlusion_degree": patch_row["occlusion_degree"],
        "reasoning": patch_row["reasoning"],
        "bbox_1000": patch_row["bbox_1000"],
        "is_valid": True,
        "source": "human_accepted_missing_label_patch",
        "candidate_id": patch_row["candidate_id"],
        "support_count": patch_row.get("support_count"),
        "support_runs": patch_row.get("support_runs", []),
    }


def _replace_assistant_annotations(
    messages: list[dict[str, Any]],
    source_records: list[dict[str, Any]],
) -> list[dict[str, Any]]:
    messages_copy = json.loads(json.dumps(messages, ensure_ascii=False))
    payload = {
        "annotations": [
            {
                "bbox": record["bbox_1000"],
                "maturity_level": record["maturity_level"],
                "maturity_ratio": record["maturity_ratio"],
                "occlusion_degree": record["occlusion_degree"],
                "reasoning": record["reasoning"],
            }
            for record in source_records
        ]
    }
    if messages_copy:
        messages_copy[-1]["content"] = json.dumps(payload, ensure_ascii=False)
    return messages_copy


def _extract_source_records(row: dict[str, Any]) -> list[dict[str, Any]]:
    records = row.get("source_records")
    if isinstance(records, list):
        return records
    messages = row.get("messages", [])
    if not messages:
        return []
    payload = json.loads(messages[-1].get("content", "{}"))
    annotations = payload.get("annotations", [])
    return annotations if isinstance(annotations, list) else []


def _build_summary(
    *,
    reference_path: Path,
    patch_path: Path,
    output_dir: Path,
    reference_out: Path,
    patch_applied_out: Path,
    reference_rows: list[dict[str, Any]],
    patched_rows: list[dict[str, Any]],
    patch_rows: list[dict[str, Any]],
    applied_rows: list[dict[str, Any]],
) -> dict[str, Any]:
    original_count = sum(len(_extract_source_records(row)) for row in reference_rows)
    reviewed_count = sum(len(_extract_source_records(row)) for row in patched_rows)
    applied_by_maturity = Counter(row.get("maturity_level") or "未知" for row in patch_rows)
    applied_by_support = Counter(int(row.get("support_count", 0)) for row in patch_rows)
    return {
        "created_at": datetime.now().isoformat(timespec="seconds"),
        "reference_path": str(reference_path),
        "patch_path": str(patch_path),
        "output_dir": str(output_dir),
        "reference_subset_reviewed_path": str(reference_out),
        "applied_missing_label_patch_path": str(patch_applied_out),
        "image_count": len(patched_rows),
        "original_annotations_total": original_count,
        "applied_missing_labels_total": len(applied_rows),
        "reviewed_annotations_total": reviewed_count,
        "patched_image_total": len({row["image_key"] for row in applied_rows}),
        "applied_by_maturity": dict(applied_by_maturity),
        "applied_by_support_count": dict(sorted(applied_by_support.items())),
        "quality_checks": {
            "annotation_count_matches": reviewed_count == original_count + len(applied_rows),
            "all_patch_rows_applied": len(applied_rows) == len(patch_rows),
            "messages_count_matches_source_records": all(
                _assistant_annotation_count(row) == len(_extract_source_records(row)) for row in patched_rows
            ),
        },
    }


def _assistant_annotation_count(row: dict[str, Any]) -> int:
    messages = row.get("messages", [])
    if not messages:
        return -1
    payload = json.loads(messages[-1].get("content", "{}"))
    annotations = payload.get("annotations", [])
    return len(annotations) if isinstance(annotations, list) else -1


def _build_report(summary: dict[str, Any]) -> str:
    lines = [
        "# Gold v1.1 Reviewed Validation Reference",
        "",
        f"- created_at: `{summary['created_at']}`",
        f"- image_count: `{summary['image_count']}`",
        f"- original_annotations_total: `{summary['original_annotations_total']}`",
        f"- applied_missing_labels_total: `{summary['applied_missing_labels_total']}`",
        f"- reviewed_annotations_total: `{summary['reviewed_annotations_total']}`",
        f"- patched_image_total: `{summary['patched_image_total']}`",
        "",
        "## Artifacts",
        "",
        f"- reference_subset_reviewed_path: `{summary['reference_subset_reviewed_path']}`",
        f"- applied_missing_label_patch_path: `{summary['applied_missing_label_patch_path']}`",
        "",
        "## Quality Checks",
        "",
    ]
    for key, value in summary["quality_checks"].items():
        lines.append(f"- {key}: `{value}`")
    lines.append("")
    return "\n".join(lines)


def _read_jsonl(path: Path) -> list[dict[str, Any]]:
    rows = []
    with path.open("r", encoding="utf-8") as file:
        for line in file:
            line = line.strip()
            if line:
                rows.append(json.loads(line))
    return rows


def _write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None:
    with path.open("w", encoding="utf-8") as file:
        for row in rows:
            file.write(json.dumps(row, ensure_ascii=False))
            file.write("\n")


if __name__ == "__main__":
    main()