File size: 4,736 Bytes
3f3265f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | import json
import tempfile
import unittest
from pathlib import Path
from PIL import Image
from approach.pipeline_utils import (
enrich_ape_results,
output_path_for_selection,
parse_orienter_image_name,
resolve_image_path,
run_optional_reflection,
select_jsonl_lines,
write_json_atomic,
)
class PipelineUtilsTest(unittest.TestCase):
def test_parse_orienter_image_name_is_strict(self):
self.assertEqual(parse_orienter_image_name("123_4.jpg"), ("123", "4", 123004))
with self.assertRaises(ValueError):
parse_orienter_image_name("123_4_extra.jpg")
with self.assertRaises(ValueError):
parse_orienter_image_name("nested/123_4.jpg")
with self.assertRaises(ValueError):
parse_orienter_image_name("123_1000.jpg")
def test_resolve_image_path_rejects_escape_from_images_directory(self):
with tempfile.TemporaryDirectory() as tmp_dir:
images_dir = Path(tmp_dir) / "images"
images_dir.mkdir()
self.assertEqual(
resolve_image_path(images_dir, "nested/123_4.jpg"),
(images_dir / "nested" / "123_4.jpg").resolve(),
)
with self.assertRaises(ValueError):
resolve_image_path(images_dir, "../123_4.jpg")
with self.assertRaises(ValueError):
resolve_image_path(images_dir, str(Path(tmp_dir) / "123_4.jpg"))
def test_select_jsonl_lines_supports_ranges_or_shards(self):
lines = [f"{idx}\n" for idx in range(10)]
self.assertEqual(select_jsonl_lines(lines, start_index=2, end_index=8), lines[2:8])
self.assertEqual(
select_jsonl_lines(lines, shard_index=1, num_shards=2),
["1\n", "3\n", "5\n", "7\n", "9\n"],
)
with self.assertRaises(ValueError):
select_jsonl_lines(
lines,
start_index=2,
end_index=8,
shard_index=1,
num_shards=2,
)
def test_output_path_for_selection_prevents_shard_overwrite(self):
self.assertEqual(
output_path_for_selection("predictions.json", shard_index=1, num_shards=4),
"predictions.shard01-of-04.json",
)
self.assertEqual(
output_path_for_selection("predictions.json", start_index=100, end_index=200),
"predictions.rows100-200.json",
)
self.assertEqual(output_path_for_selection("predictions.json"), "predictions.json")
def test_enrich_ape_results_keeps_added_metadata(self):
enriched = enrich_ape_results(
[{"bbox": [1, 2, 3, 4], "category_name": "button"}],
"123_4.jpg",
lambda image_name: 123004,
)
self.assertEqual(enriched[0]["image_id"], 123004)
self.assertEqual(enriched[0]["category_id"], "button")
self.assertEqual(enriched[0]["bbox"], [1, 2, 3, 4])
def test_write_json_atomic_writes_one_valid_json_document(self):
with tempfile.TemporaryDirectory() as tmp_dir:
out_path = Path(tmp_dir) / "ape.json"
write_json_atomic(str(out_path), [{"image_id": 1}, {"image_id": 2}])
self.assertEqual(json.loads(out_path.read_text()), [{"image_id": 1}, {"image_id": 2}])
def test_optional_reflection_is_disabled_by_default(self):
self.assertIsNone(
run_optional_reflection(
"unused.png",
[],
detector=lambda candidates, previous: previous,
enabled=False,
)
)
def test_optional_reflection_feeds_back_to_detector(self):
with tempfile.TemporaryDirectory() as tmp_dir:
image_path = Path(tmp_dir) / "image.png"
Image.new("RGB", (10, 10), "white").save(image_path)
detector_calls = []
def advisor(**kwargs):
if detector_calls:
return {"verified": [0], "needs_refinement": [], "feedback": []}
return {"verified": [], "needs_refinement": [0], "feedback": ["button label"]}
def detector(candidates, previous):
detector_calls.append(candidates)
return [{"bbox": [2, 2, 3, 3], "category_name": "button"}]
result = run_optional_reflection(
str(image_path),
[{"bbox": [1, 1, 3, 3], "category_name": "button"}],
detector=detector,
enabled=True,
advisor=advisor,
)
self.assertEqual(detector_calls, [["button label"]])
self.assertFalse(result["max_iterations_reached"])
if __name__ == "__main__":
unittest.main()
|