import copy import json from pathlib import Path import tempfile import unittest import numpy as np from PIL import Image from completion_geometry import prepare, merge_candidate, localize, read_json, write_json, finalize def make_fixture(directory): root = Path(directory) root.mkdir(parents=True, exist_ok=True) (root/'masks').mkdir(exist_ok=True) source = np.zeros((80, 80, 4), np.uint8) source[20:40, 20:40] = [180, 50, 20, 255] source[20, 20, 3] = 80 source[20:40, 40:50] = [20, 80, 180, 255] Image.fromarray(source).save(root/'source.png') part = {'id': 'torso', 'semantic': 'torso', 'parent': None, 'rig_role': 'bone_part', 'bind_to': None, 'visibility': 'partial', 'bbox_norm_xyxy': [.25, .25, .5, .5], 'pivot_norm': [.4, .4], 'positive_points_norm': [], 'negative_points_norm': [], 'occluded_by': ['collar'], 'needs_completion': True, 'confidence': .8, 'completion': {'occlusion_fill': True, 'joint_extension_parent': False, 'joint_extension_child': False}} collar = copy.deepcopy(part) collar.update(id='collar', semantic='collar', parent='torso', bind_to='torso', rig_role='attachment', visibility='visible', occluded_by=[], needs_completion=False, completion={key: False for key in part['completion']}) write_json(root/'semantic-parts.json', {'schema_version': 2, 'coordinate_space': 'normalized_xy_top_left', 'parts': [part, collar], 'notes': []}) write_json(root/'segmentation.json', {'schema_version': 1, 'parts': [{'id': 'torso', 'warnings': ['SEMANTIC_REVIEW_REQUIRED']}, {'id': 'collar'}]}) for name, x0, x1 in [('torso', 20, 40), ('collar', 40, 50)]: mask = np.zeros((80, 80), np.uint8) mask[20:40, x0:x1] = 255 Image.fromarray(mask).save(root/'masks'/f'{name}.png') # Expected 128 canvas: union bbox (20,20,50,40) -> origin (-29,-34). # Construct candidate independently from production prepare/merge logic. candidate = np.zeros((128, 128, 4), np.uint8) candidate[54:74, 49:79] = [180, 50, 20, 255] candidate[54, 49, 3] = 80 Image.fromarray(candidate).save(root/'candidate.png') return source class CompletionTests(unittest.TestCase): def setUp(self): self.temp = tempfile.TemporaryDirectory() self.addCleanup(self.temp.cleanup) self.root = Path(self.temp.name) self.source = make_fixture(self.root) self.plan = prepare(self.root, ['torso'], size=128, radius=16) self.folder = self.root/'tasks'/'torso' def test_canvas_mapping_and_only_occluder_region_editable(self): task = self.plan['tasks'][0] self.assertEqual(task['origin_xy'], [-29, -34]) self.assertEqual(task['allowed_pixels'], 200) visible = np.array(Image.open(self.folder/'visible.png')) self.assertEqual(visible[54, 49, 3], 80) mask = np.array(Image.open(self.folder/'edit-mask.png')) self.assertTrue(np.all(mask[54:74, 69:79, 3] == 0)) self.assertTrue(np.all(mask[54:74, 49:69, 3] == 255)) self.assertEqual(int((mask[:, :, 3] == 0).sum()), 200) def test_merge_restores_rgba_exactly_and_clips_extra_pixels(self): candidate = np.array(Image.open(self.root/'candidate.png')) candidate[54:74, 49:69, :3] += 5 candidate[10:14, 10:14] = [255, 0, 0, 255] Image.fromarray(candidate).save(self.root/'candidate.png') output = self.folder/'merged.png' qa = merge_candidate(self.folder, self.root/'candidate.png', output) self.assertTrue(qa['accepted']) self.assertEqual(qa['new_pixels'], 200) self.assertEqual(qa['discarded_outside_allowed'], 16) self.assertTrue(qa['visible_pixels_exact']) merged = np.array(Image.open(output)) np.testing.assert_array_equal(merged[54:74, 49:69], self.source[20:40, 20:40]) self.assertEqual(merged[11, 11, 3], 0) write_json(self.root/'completion-result.json', {'parts': [{'id': 'torso', 'status': 'geometry_passed_needs_review', 'selected_file': 'tasks/torso/merged.png', 'qa': qa}], 'skipped': []}) finalize(self.root) result = read_json(self.root/'completion-result.json') self.assertEqual(result['parts'][0]['crop_xyxy'], [20, 20, 50, 40]) def test_reject_resize_shift_no_additions_and_background_only(self): candidate = Image.open(self.root/'candidate.png').convert('RGBA') variants = [(candidate.resize((64, 64)), 'CANVAS_SIZE_MISMATCH'), (Image.fromarray(np.roll(np.array(candidate), 12, axis=1)), 'VISIBLE_ALIGNMENT_MISMATCH'), (Image.open(self.folder/'visible.png'), 'NO_CONNECTED_NEW_PIXELS'), (Image.new('RGB', (128, 128), 'white'), 'CANNOT_EXTRACT_FOREGROUND')] for number, (image, reason) in enumerate(variants): file = self.root/f'bad-{number}.png' image.save(file) output = self.root/f'rejected-{number}.png' qa = merge_candidate(self.folder, file, output) self.assertFalse(qa['accepted']) self.assertIn(reason, qa['reasons']) self.assertFalse(output.exists()) def test_white_background_extraction_preserves_visible_and_adds_only_allowed(self): candidate = Image.open(self.root/'candidate.png').convert('RGBA') rgb = Image.alpha_composite(Image.new('RGBA', candidate.size, 'white'), candidate).convert('RGB') rgb.save(self.root/'white.png') qa = merge_candidate(self.folder, self.root/'white.png', self.folder/'merged.png') self.assertTrue(qa['accepted']) self.assertEqual(qa['new_pixels'], 200) def test_uniform_canvas_recovery_preserves_source_and_still_rejects_changes(self): original_candidate = Image.open(self.root/'candidate.png').convert('RGBA') larger = original_candidate.resize((256, 256), Image.Resampling.NEAREST) file = self.root/'large.png' larger.save(file) output = self.folder/'recovered.png' qa = merge_candidate(self.folder, file, output, 'uniform') self.assertTrue(qa['accepted']) self.assertTrue(qa['canvas_resampled']) self.assertEqual(qa['returned_size'], [256, 256]) self.assertEqual(qa['canvas_scale'], .5) merged = np.array(Image.open(output)) np.testing.assert_array_equal(merged[54:74, 49:69], self.source[20:40, 20:40]) self.assertEqual(Image.open(file).size, (256, 256)) for name, im, reason in [ ('wide', larger.resize((260, 256)), 'CANVAS_ASPECT_RATIO_MISMATCH'), ('shifted', Image.fromarray(np.roll(np.array(larger), 32, axis=1)), 'VISIBLE_ALIGNMENT_MISMATCH'), ]: file = self.root/f'{name}.png' im.save(file) qa = merge_candidate(self.folder, file, self.folder/f'{name}-merged.png', 'uniform') self.assertFalse(qa['accepted']) self.assertIn(reason, qa['reasons']) def test_disconnected_additions_are_not_merged(self): candidate = np.array(Image.open(self.folder/'visible.png')) candidate[54:57, 76:79] = [180, 50, 20, 255] Image.fromarray(candidate).save(self.root/'detached.png') qa = merge_candidate(self.folder, self.root/'detached.png', self.folder/'detached-merged.png') self.assertFalse(qa['accepted']) self.assertEqual(qa['disconnected_pixels'], 9) self.assertIn('NO_CONNECTED_NEW_PIXELS', qa['reasons']) def test_block_missing_occluder_and_reject_canvas_or_masks(self): plan = read_json(self.root/'semantic-parts.json') plan['parts'][0]['occluded_by'] = [] write_json(self.root/'semantic-parts.json', plan) self.assertEqual(prepare(self.root, [], 128, 16)['skipped'][0]['reason'], 'NO_ALLOWED_EXTENSION') with self.assertRaises(ValueError): prepare(self.root, ['missing'], 128, 16) Image.new('L', (12, 12), 255).save(self.root/'masks'/'torso.png') with self.assertRaisesRegex(ValueError, 'mask must'): prepare(self.root, [], 128, 16) def test_joint_regions_use_parent_and_child_pivots(self): plan = read_json(self.root/'semantic-parts.json') plan['parts'][0]['completion'] = {'occlusion_fill': False, 'joint_extension_parent': True, 'joint_extension_child': True} plan['parts'][0]['pivot_norm'] = [.25, .4] child = copy.deepcopy(plan['parts'][0]) child.update(id='arm', parent='torso', pivot_norm=[.5, .4]) plan['parts'].append(child) write_json(self.root/'semantic-parts.json', plan) task = prepare(self.root, ['torso'], 128, 8)['tasks'][0] self.assertGreater(task['regions']['joint_extension_parent'], 0) self.assertGreater(task['regions']['joint_extension_child'], 0) if __name__ == '__main__': unittest.main()