tests/test_swift_sft_export.py

from __future__ import annotations

import json
import tempfile
import unittest
from pathlib import Path

from pipelines.export_swift_sft_dataset import run_swift_sft_export


def build_gold_row(image_path: str) -> dict:
    return {
        "image_path": image_path,
        "source_image_path": r"F:\old\img_a.jpg",
        "messages": [
            {"role": "system", "content": "sys"},
            {
                "role": "user",
                "content": [
                    {"type": "image", "image": image_path},
                    {"type": "text", "text": "请输出 annotations。"},
                ],
            },
            {
                "role": "assistant",
                "content": json.dumps(
                    {
                        "annotations": [
                            {
                                "bbox": [100, 200, 300, 400],
                                "maturity_level": "未成熟",
                                "maturity_ratio": 0.1,
                                "occlusion_degree": "轻度",
                                "reasoning": "位于中央果实左侧后方，被茎叶部分遮挡，局部可见绿色轮廓。",
                            },
                            {
                                "bbox": [500, 600, 700, 800],
                                "maturity_level": "完熟",
                                "maturity_ratio": 0.9,
                                "occlusion_degree": "无",
                                "reasoning": "果实颜色鲜艳偏红，无遮挡。",
                            },
                        ]
                    },
                    ensure_ascii=False,
                ),
            },
        ],
        "source_records": [
            {
                "image_path": r"F:\old\img_a.jpg",
                "image_width": 1000,
                "image_height": 1000,
                "target_index": 0,
                "bbox": [10, 20, 30, 40],
                "bbox_1000": [100, 200, 300, 400],
                "maturity_level": "未成熟",
                "maturity_ratio": 0.1,
                "occlusion_degree": "轻度",
                "reasoning": "位于中央果实左侧后方，被茎叶部分遮挡，局部可见绿色轮廓。",
                "is_valid": True,
            },
            {
                "image_path": r"F:\old\img_a.jpg",
                "image_width": 1000,
                "image_height": 1000,
                "target_index": 1,
                "bbox": [50, 60, 70, 80],
                "bbox_1000": [500, 600, 700, 800],
                "maturity_level": "完熟",
                "maturity_ratio": 0.9,
                "occlusion_degree": "无",
                "reasoning": "果实颜色鲜艳偏红，无遮挡。",
                "is_valid": True,
            },
        ],
    }


class SwiftSFTExportTest(unittest.TestCase):
    def test_export_official_mode_builds_grounding_format(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            temp_path = Path(temp_dir)
            gold_path = temp_path / "gold.jsonl"
            output_dir = temp_path / "swift_sft"
            image_path = str((temp_path / "img_a.jpg").resolve())
            gold_path.write_text(json.dumps(build_gold_row(image_path), ensure_ascii=False) + "\n", encoding="utf-8")

            summary = run_swift_sft_export(
                gold_dataset_path=gold_path,
                output_dir=output_dir,
                val_ratio=0.0,
                export_mode="official",
            )

            self.assertEqual(summary["export_mode"], "official")
            row = json.loads((output_dir / "train.jsonl").read_text(encoding="utf-8").splitlines()[0])
            self.assertTrue(row["messages"][0]["content"].startswith("<image>找到图像中的未成熟番茄、半成熟番茄和完熟番茄"))
            assistant = json.loads(row["messages"][1]["content"])
            self.assertEqual(assistant[0], {"bbox_2d": "<bbox>", "label": "<ref-object>"})
            self.assertEqual(row["objects"]["ref"], ["未成熟番茄", "完熟番茄"])
            self.assertEqual(row["objects"]["bbox"], [[10, 20, 30, 40], [50, 60, 70, 80]])

    def test_export_attr_enhanced_mode_adds_attribute_text(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            temp_path = Path(temp_dir)
            gold_path = temp_path / "gold.jsonl"
            output_dir = temp_path / "swift_sft"
            image_path = str((temp_path / "img_a.jpg").resolve())
            gold_path.write_text(json.dumps(build_gold_row(image_path), ensure_ascii=False) + "\n", encoding="utf-8")

            summary = run_swift_sft_export(
                gold_dataset_path=gold_path,
                output_dir=output_dir,
                val_ratio=0.0,
                export_mode="attr_enhanced",
            )

            self.assertEqual(summary["export_mode"], "attr_enhanced")
            row = json.loads((output_dir / "train.jsonl").read_text(encoding="utf-8").splitlines()[0])
            assistant = json.loads(row["messages"][1]["content"])
            self.assertEqual(assistant[0]["bbox_2d"], "<bbox>")
            self.assertEqual(assistant[0]["label"], "<ref-object>")
            self.assertIn("attribute_text", assistant[0])
            self.assertEqual(assistant[0]["attribute_text"], "轻度遮挡,局部可见绿色")
            self.assertEqual(assistant[1]["attribute_text"], "无遮挡,局部可见红色")
            self.assertEqual(row["metadata"]["export_mode"], "attr_enhanced")

    def test_export_attr_occlusion_only_mode_keeps_only_occlusion_text(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            temp_path = Path(temp_dir)
            gold_path = temp_path / "gold.jsonl"
            output_dir = temp_path / "swift_sft"
            image_path = str((temp_path / "img_a.jpg").resolve())
            gold_path.write_text(json.dumps(build_gold_row(image_path), ensure_ascii=False) + "\n", encoding="utf-8")

            summary = run_swift_sft_export(
                gold_dataset_path=gold_path,
                output_dir=output_dir,
                val_ratio=0.0,
                export_mode="attr_occlusion_only",
            )

            self.assertEqual(summary["export_mode"], "attr_occlusion_only")
            row = json.loads((output_dir / "train.jsonl").read_text(encoding="utf-8").splitlines()[0])
            assistant = json.loads(row["messages"][1]["content"])
            self.assertEqual(assistant[0]["attribute_text"], "轻度遮挡")
            self.assertEqual(assistant[1]["attribute_text"], "无遮挡")
            self.assertEqual(row["metadata"]["semantic_hint_mode"], "short_prompt_hints_plus_occlusion_attribute_text")

    def test_export_supports_fixed_val_count(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            temp_path = Path(temp_dir)
            gold_path = temp_path / "gold.jsonl"
            output_dir = temp_path / "swift_sft"
            rows = []
            for index in range(5):
                image_path = str((temp_path / f"img_{index}.jpg").resolve())
                rows.append(build_gold_row(image_path))
            gold_path.write_text("\n".join(json.dumps(row, ensure_ascii=False) for row in rows) + "\n", encoding="utf-8")

            summary = run_swift_sft_export(
                gold_dataset_path=gold_path,
                output_dir=output_dir,
                val_ratio=0.01,
                val_count=2,
                export_mode="official",
            )

            self.assertEqual(summary["val_count"], 2)
            self.assertEqual(summary["sample_counts"]["val"], 2)
            self.assertEqual(summary["sample_counts"]["train"], 3)
            self.assertEqual(len((output_dir / "val.jsonl").read_text(encoding="utf-8").splitlines()), 2)


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