tests/test_gold_builder.py

from __future__ import annotations

import json
import tempfile
import unittest
from pathlib import Path

from pipelines.gold_builder import build_gold_dataset


class GoldBuilderPipelineTest(unittest.TestCase):
    def test_build_gold_dataset_can_rewrite_windows_image_path_to_linux_root(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            temp_path = Path(temp_dir)
            silver_path = temp_path / "silver.jsonl"
            output_dir = temp_path / "gold_release"
            image_root = temp_path / "images"
            image_root.mkdir(parents=True, exist_ok=True)
            image_file = image_root / "example.jpg"
            image_file.write_bytes(b"fake")

            row = {
                "image_path": r"F:\lsy\code\agent-base\clearn_base_data\example.jpg",
                "image_width": 100,
                "image_height": 100,
                "target_index": 0,
                "bbox": [10, 10, 20, 20],
                "bbox_1000": [100, 100, 200, 200],
                "maturity_level": "未成熟",
                "maturity_ratio": 0.1,
                "occlusion_degree": "无",
                "reasoning": "目标一。",
                "is_valid": True,
            }
            silver_path.write_text(json.dumps(row, ensure_ascii=False) + "\n", encoding="utf-8")

            summary = build_gold_dataset(
                silver_dataset_path=silver_path,
                output_dir=output_dir,
                image_root_dir=str(image_root),
                fail_on_missing_image=True,
            )

            self.assertTrue(summary["path_rewrite"]["enabled"])
            self.assertEqual(summary["path_rewrite"]["rewritten_images"], 1)
            self.assertTrue(summary["path_rewrite"]["all_rewritten_paths_exist"])

            gold_rows = [json.loads(line) for line in (output_dir / "gold_dataset.jsonl").read_text(encoding="utf-8").splitlines() if line.strip()]
            self.assertEqual(gold_rows[0]["image_path"], str(image_file.resolve()))
            self.assertEqual(gold_rows[0]["source_image_path"], row["image_path"])
            self.assertEqual(gold_rows[0]["messages"][1]["content"][0]["image"], str(image_file.resolve()))

    def test_build_gold_dataset_groups_records_by_image(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            temp_path = Path(temp_dir)
            silver_path = temp_path / "silver.jsonl"
            output_dir = temp_path / "gold_release"

            rows = [
                {
                    "image_path": str((temp_path / "img_a.jpg").resolve()),
                    "image_width": 100,
                    "image_height": 100,
                    "target_index": 0,
                    "bbox": [10, 10, 20, 20],
                    "bbox_1000": [100, 100, 200, 200],
                    "maturity_level": "未成熟",
                    "maturity_ratio": 0.1,
                    "occlusion_degree": "无",
                    "reasoning": "目标一。",
                    "is_valid": True,
                },
                {
                    "image_path": str((temp_path / "img_a.jpg").resolve()),
                    "image_width": 100,
                    "image_height": 100,
                    "target_index": 1,
                    "bbox": [30, 30, 40, 40],
                    "bbox_1000": [300, 300, 400, 400],
                    "maturity_level": "完熟",
                    "maturity_ratio": 0.9,
                    "occlusion_degree": "轻度",
                    "reasoning": "目标二。",
                    "is_valid": True,
                },
                {
                    "image_path": str((temp_path / "img_b.jpg").resolve()),
                    "image_width": 100,
                    "image_height": 100,
                    "target_index": 0,
                    "bbox": [5, 5, 6, 6],
                    "bbox_1000": [50, 50, 60, 60],
                    "maturity_level": "未成熟",
                    "maturity_ratio": 0.1,
                    "occlusion_degree": "重度",
                    "reasoning": "这个样本无效。",
                    "is_valid": False,
                },
            ]
            with silver_path.open("w", encoding="utf-8") as file:
                for row in rows:
                    file.write(json.dumps(row, ensure_ascii=False))
                    file.write("\n")

            summary = build_gold_dataset(silver_dataset_path=silver_path, output_dir=output_dir, sample_size=1)
            self.assertEqual(summary["gold_images_total"], 1)
            self.assertEqual(summary["gold_annotations_total"], 2)
            self.assertEqual(summary["silver_total_records"], 3)
            self.assertEqual(summary["silver_valid_records"], 2)
            self.assertEqual(summary["silver_invalid_records"], 1)
            self.assertEqual(summary["sample_size"], 1)

            gold_path = output_dir / "gold_dataset.jsonl"
            summary_path = output_dir / "summary.json"
            report_path = output_dir / "report.md"
            sample_path = output_dir / "sample_check.jsonl"

            self.assertTrue(gold_path.exists())
            self.assertTrue(summary_path.exists())
            self.assertTrue(report_path.exists())
            self.assertTrue(sample_path.exists())

            gold_rows = [json.loads(line) for line in gold_path.read_text(encoding="utf-8").splitlines() if line.strip()]
            self.assertEqual(len(gold_rows), 1)
            annotations = json.loads(gold_rows[0]["messages"][2]["content"])["annotations"]
            self.assertEqual(len(annotations), 2)
            self.assertEqual(annotations[0]["bbox"], [100, 100, 200, 200])
            self.assertEqual(annotations[1]["bbox"], [300, 300, 400, 400])

            report_text = report_path.read_text(encoding="utf-8")
            self.assertIn("Gold Release Report", report_text)

            sample_rows = [json.loads(line) for line in sample_path.read_text(encoding="utf-8").splitlines() if line.strip()]
            self.assertEqual(len(sample_rows), 1)
            self.assertEqual(sample_rows[0]["target_indices"], [0, 1])


if __name__ == "__main__":
    unittest.main()