Spaces:
Running
Running
| import json | |
| from pathlib import Path | |
| import tempfile | |
| import unittest | |
| from PIL import Image | |
| from image_cleaner.aspect import assess_aspect | |
| from image_cleaner.config import ASPECT_BUCKETS, Config | |
| from image_cleaner.cli import main | |
| from image_cleaner.images import filter_reasons | |
| from image_cleaner.io import prepare, manifest_items | |
| from image_cleaner.jobs import analyze, merge | |
| class AspectTests(unittest.TestCase): | |
| def match(self, width, height, **kwargs): | |
| return assess_aspect({"width": width, "height": height}, Config(**kwargs)) | |
| def test_strict_presets_orientation_and_all(self): | |
| for bucket in ASPECT_BUCKETS: | |
| w, h = map(int, bucket.split(":")) | |
| self.assertEqual(self.match(w*20, h*20, aspect_buckets="all")["aspect_bucket"], bucket) | |
| self.assertFalse(self.match(60, 40, aspect_buckets="2:3")["aspect_bucket_match"]) | |
| self.assertFalse(self.match(9, 21, aspect_buckets="all")["aspect_bucket_match"]) | |
| self.assertFalse(self.match(1001, 1000, aspect_buckets="1:1")["aspect_bucket_match"]) | |
| def test_inclusive_symmetric_boundary_and_custom_tolerance(self): | |
| for w, h in [(1050, 1000), (1000, 1050)]: | |
| row = self.match(w, h, aspect_buckets="1:1", aspect_mode="loose") | |
| self.assertTrue(row["aspect_bucket_match"]) | |
| self.assertEqual(row["aspect_bucket_deviation"], .05) | |
| for w, h in [(1051, 1000), (1000, 1051)]: | |
| self.assertFalse(self.match(w, h, aspect_buckets="1:1", aspect_mode="loose")["aspect_bucket_match"]) | |
| self.assertTrue(self.match(1100, 1000, aspect_buckets="1:1", aspect_mode="loose", aspect_tolerance=.1)["aspect_bucket_match"]) | |
| self.assertFalse(self.match(1001, 1000, aspect_buckets="1:1", aspect_mode="loose", aspect_tolerance=0)["aspect_bucket_match"]) | |
| def test_overlap_assignment_and_tie(self): | |
| self.assertEqual(self.match(1400, 1000, aspect_buckets="3:2,4:3", aspect_mode="loose", aspect_tolerance=.2)["aspect_bucket"], "4:3") | |
| self.assertEqual(self.match(100, 100, aspect_buckets="3:2,2:3", aspect_mode="loose", aspect_tolerance=.5)["aspect_bucket"], "2:3") | |
| def test_validation_and_normalization(self): | |
| for value in ["", "1:1,", [], "4:5", ["all"], 5]: | |
| with self.assertRaises(ValueError): | |
| Config(aspect_buckets=value) | |
| for value in [None, True, -1, float("nan"), float("inf")]: | |
| with self.assertRaises(ValueError): | |
| Config(aspect_tolerance=value) | |
| self.assertEqual(Config(aspect_buckets="3:2,1:1,1:1").aspect_buckets, ("1:1", "3:2")) | |
| self.assertIsNone(Config(aspect_buckets="none").aspect_buckets) | |
| with self.assertRaises(ValueError): | |
| Config(aspect_mode="invalid") | |
| def test_disabled_clears_stale_assignment(self): | |
| row = {"width": 100, "height": 100, "short_edge": 100, "aspect_bucket": "1:1"} | |
| self.assertEqual(filter_reasons(row, Config()), []) | |
| self.assertIsNone(row["aspect_bucket"]) | |
| self.assertIsNone(row["aspect_bucket_match"]) | |
| def test_sharded_cached_filter_cli_and_reports(self): | |
| temp_root = Path(__file__).resolve().parents[1]/"runs"/"test_tmp" | |
| temp_root.mkdir(parents=True, exist_ok=True) | |
| with tempfile.TemporaryDirectory(dir=temp_root) as temp: | |
| root = Path(temp) | |
| images = root/"images" | |
| images.mkdir() | |
| for i, size in enumerate([(256, 256), (315, 300), (316, 300), (300, 450), (128, 128)]): | |
| Image.new("RGB", size, "blue").save(images/f"{i}.png") | |
| (images/"broken.jpg").write_bytes(b"broken") | |
| manifest = root/"input.jsonl" | |
| prepare(images, manifest) | |
| config = Config(deduplicate=False, near_deduplicate=False) | |
| analyze(manifest, root/"single", config, workers=2) | |
| for rank in range(2): | |
| analyze(manifest, root/"shards", config, workers=2, rank=rank, world_size=2) | |
| # Remove source bytes: repeated filtering must use cached measurements only. | |
| for path in images.iterdir(): | |
| path.unlink() | |
| args = ["filter", "--shards", str(root/"shards"), "--output", str(root/"loose"), | |
| "--aspect-buckets", "1:1,2:3", "--aspect-mode", "loose", "--resolution", "256p"] | |
| self.assertEqual(main(args), 0) | |
| report = json.loads((root/"loose"/"report.json").read_text(encoding="utf-8")) | |
| self.assertEqual(report["kept_count"], 3) | |
| self.assertEqual(report["stage_metrics"]["aspect_bucket_matched_count"], 4) | |
| self.assertEqual(report["stage_metrics"]["aspect_bucket_pass_rate"], .8) | |
| self.assertEqual(report["aspect_bucket_filter"]["counts_before"], {"1:1": 3, "2:3": 1}) | |
| self.assertEqual(report["aspect_bucket_filter"]["counts_after"], {"1:1": 2, "2:3": 1}) | |
| merge(root/"single", root/"equivalent", resolution="256p", aspect_buckets="1:1,2:3", aspect_mode="loose") | |
| for status in ["kept", "rejected", "error", "duplicate"]: | |
| self.assertEqual((root/"loose"/f"{status}.jsonl").read_bytes(), (root/"equivalent"/f"{status}.jsonl").read_bytes()) | |
| merge(root/"shards", root/"strict", aspect_buckets="1:1,2:3") | |
| strict = list(manifest_items(root/"strict"/"rejected.jsonl")) | |
| self.assertEqual(len(strict), 2) | |
| self.assertIn("aspect_bucket_mismatch", strict[0]["reasons"]) | |
| self.assertIsNone(strict[0]["aspect_bucket"]) | |
| self.assertIsNotNone(strict[0]["aspect_bucket_nearest"]) | |
| args[args.index(str(root/"loose"))] = str(root/"tighter") | |
| self.assertEqual(main(args + ["--aspect-tolerance", "0.03"]), 0) | |
| tighter = json.loads((root/"tighter"/"report.json").read_text(encoding="utf-8")) | |
| self.assertEqual(tighter["kept_count"], 2) | |
| if __name__ == "__main__": | |
| unittest.main() | |