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("找到图像中的未成熟番茄、半成熟番茄和完熟番茄")) assistant = json.loads(row["messages"][1]["content"]) self.assertEqual(assistant[0], {"bbox_2d": "", "label": ""}) 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"], "") self.assertEqual(assistant[0]["label"], "") 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()