import copy import json from pathlib import Path import shutil import tempfile import unittest import numpy as np from PIL import Image from spine_assembly import load_assets, build_skeleton, export_assets, write_json def make_fixture(directory): """An asymmetric mechanical fixture; no humanoid names or assumed part count.""" root = Path(directory) sam, completion = root / 'sam', root / 'complete' (sam / 'masks').mkdir(parents=True) source = np.zeros((48, 64, 4), np.uint8) specs = [('shell', [16, 12, 30, 24], [40, 70, 130, 255]), ('link_z', [30, 22, 40, 32], [160, 70, 10, 255]), ('tip_q', [40, 28, 45, 36], [80, 170, 30, 255]), ('decal', [19, 15, 22, 18], [200, 50, 130, 255])] for _, (x0, y0, x1, y1), color in specs: source[y0:y1, x0:x1] = color source[12, 16, 3] = 90 source[25, 33] = [230, 140, 40, 255] Image.fromarray(source).save(sam / 'source.png') def part(name, parent, role, pivot, box, hidden=False, needs=False, occluded=None): return {'id': name, 'semantic': name, 'parent': parent, 'rig_role': role, 'bind_to': parent if role != 'bone_part' else None, 'pivot_norm': [pivot[0] / 64, pivot[1] / 48], 'visibility': 'hidden' if hidden else 'visible', 'bbox_norm_xyxy': None if hidden else [box[0]/64, box[1]/48, box[2]/64, box[3]/48], 'positive_points_norm': [], 'negative_points_norm': [], 'occluded_by': occluded or [], 'confidence': .9, 'needs_completion': hidden or needs, 'completion': {'occlusion_fill': hidden, 'joint_extension_parent': needs, 'joint_extension_child': False}} parts = [part('tip_q', 'link_z', 'bone_part', [40, 29], specs[2][1], needs=True), part('decal', 'hub', 'attachment', [19, 15], specs[3][1]), part('shell', 'hub', 'attachment', [20, 16], specs[0][1], occluded=['decal']), part('link_z', 'hub', 'bone_part', [31, 23], specs[1][1], needs=True), part('hub', None, 'bone_part', [21, 17], None, hidden=True, occluded=['shell'])] write_json(sam / 'semantic-parts.json', {'schema_version': 2, 'coordinate_space': 'normalized_xy_top_left', 'parts': parts, 'notes': []}) records = [] for name, (x0, y0, x1, y1), _ in specs: mask = np.zeros(source.shape[:2], np.uint8) mask[y0:y1, x0:x1] = 255 Image.fromarray(mask).save(sam / 'masks' / f'{name}.png') records.append({'id': name, 'mask': f'masks/{name}.png', 'crop_xyxy': [x0, y0, x1, y1]}) write_json(sam / 'segmentation.json', {'schema_version': 1, 'parts': records}) shutil.copytree(sam, completion) (completion / 'parts').mkdir() # Add pixels to the left of link_z; the original visible pixels are exact. candidate = np.zeros((10, 14, 4), np.uint8) candidate[:, :4] = [160, 70, 10, 255] candidate[:, 4:] = source[22:32, 30:40] Image.fromarray(candidate).save(completion / 'parts' / 'link_z.png') write_json(completion / 'completion-result.json', {'parts': [ {'id': 'link_z', 'status': 'geometry_passed_needs_review', 'selected_file': 'tasks/link_z/merged-1.png', 'qa': {'accepted': True}, 'image': 'parts/link_z.png', 'crop_xyxy': [26, 22, 40, 32]}, {'id': 'tip_q', 'status': 'rejected', 'qa': {'accepted': False, 'reasons': ['VISIBLE_ALIGNMENT_MISMATCH']}} ], 'skipped': [{'id': 'hub', 'reason': 'NO_VISIBLE_SAM_MASK'}, {'id': 'shell', 'reason': 'NO_COMPLETION_NEEDED'}, {'id': 'decal', 'reason': 'NO_COMPLETION_NEEDED'}]}) class AssemblyTests(unittest.TestCase): def setUp(self): self.temp = tempfile.TemporaryDirectory() self.addCleanup(self.temp.cleanup) self.root = Path(self.temp.name) make_fixture(self.root) self.sam, self.complete = self.root / 'sam', self.root / 'complete' def load(self, **kwargs): return load_assets(self.sam, self.complete, **kwargs) def result(self): return json.loads((self.complete / 'completion-result.json').read_text()) def test_pose_coordinates_hidden_hierarchy_and_attachment_binding(self): data = self.load() result = export_assets(data, self.root / 'out') self.assertEqual((result['parts'], result['bones'], result['completion_parts'], result['fallback_parts']), (4, 4, 1, 1)) skeleton = build_skeleton(data[1], data[2]) self.assertEqual([b['name'] for b in skeleton['bones']], ['root', 'hub', 'link_z', 'tip_q']) bones = {b['name']: b for b in skeleton['bones']} self.assertAlmostEqual(bones['hub']['x'], -11) self.assertAlmostEqual(bones['hub']['y'], 31) self.assertAlmostEqual(bones['link_z']['x'], 10) self.assertAlmostEqual(bones['link_z']['y'], -6) self.assertAlmostEqual(bones['tip_q']['x'], 9) self.assertAlmostEqual(bones['tip_q']['y'], -6) slots = {s['name']: s for s in skeleton['slots']} self.assertEqual(slots['shell']['bone'], 'hub') self.assertNotIn('hub', slots) self.assertLess(data[2]['draw_order'].index('shell'), data[2]['draw_order'].index('decal')) # Independent forward kinematics reconstructs each crop centre exactly. world = {'root': (0, 0)} for bone in skeleton['bones'][1:]: px, py = world[bone['parent']] world[bone['name']] = px + bone['x'], py + bone['y'] for name, (_, box) in data[3].items(): att = skeleton['skins'][0]['attachments'][name][name] bx, by = world[slots[name]['bone']] self.assertAlmostEqual(att['x'] + bx + 32, (box[0] + box[2]) / 2) self.assertAlmostEqual(48 - att['y'] - by, (box[1] + box[3]) / 2) image = np.array(Image.open(self.root / 'out' / 'assembled.png')) np.testing.assert_array_equal(image[12, 16], [40, 70, 130, 90]) self.assertEqual(skeleton['animations']['check__link_z']['bones']['link_z']['rotate'][1]['value'], 10) def test_rejection_fallback_and_strict_policy(self): data = self.load() assets = {p['id']: p for p in data[2]['parts']} self.assertEqual(assets['tip_q']['source'], 'sam') self.assertEqual(assets['tip_q']['reason'], 'COMPLETION_REJECTED') with self.assertRaisesRegex(ValueError, 'Completion required for tip_q'): self.load(fallback_policy='error') only_sam = load_assets(self.sam) self.assertTrue(all(p['source'] in ['sam', 'hidden'] for p in only_sam[2]['parts'])) def test_mixed_runs_and_changed_masks_rejected(self): mask = self.complete / 'masks' / 'tip_q.png' Image.new('L', (64, 48), 255).save(mask) with self.assertRaisesRegex(ValueError, 'Completion mask mismatch'): self.load() shutil.copy2(self.sam / 'masks' / 'tip_q.png', mask) (self.complete / 'semantic-parts.json').write_text('{}') with self.assertRaisesRegex(ValueError, 'Completion input mismatch'): self.load() def test_bad_crop_alignment_path_and_missing_visible_parts_rejected(self): original = self.result() for box in [[26, 22, 39, 32], [-1, 22, 13, 32], [27, 22, 41, 32]]: result = copy.deepcopy(original) result['parts'][0]['crop_xyxy'] = box write_json(self.complete / 'completion-result.json', result) with self.assertRaises(ValueError): self.load() result = copy.deepcopy(original) result['parts'][0]['image'] = '../sam/source.png' write_json(self.complete / 'completion-result.json', result) with self.assertRaisesRegex(ValueError, 'outside input'): self.load() seg = json.loads((self.sam / 'segmentation.json').read_text()) seg['parts'] = seg['parts'][1:] write_json(self.sam / 'segmentation.json', seg) with self.assertRaisesRegex(ValueError, 'Visible part has no SAM mask'): load_assets(self.sam) def test_draw_order_override_and_unconstrained_overlap(self): data = self.load() self.assertTrue(data[2]['unconstrained_overlaps']) override = self.root / 'order.json' order = ['decal', 'shell', 'link_z', 'tip_q'] write_json(override, {'draw_order': order}) modified = self.load(draw_order_file=override)[2] self.assertEqual(modified['draw_order'], order) self.assertIn(['shell', 'decal'], modified['overridden_occlusion_constraints']) write_json(override, {'draw_order': ['shell'] * 4}) with self.assertRaisesRegex(ValueError, 'exactly once'): self.load(draw_order_file=override) def test_attachment_only_and_root_binding(self): plan = {'parts': [{'id': 'panel', 'rig_role': 'attachment', 'bind_to': None}]} assembly = {'width': 64, 'height': 48, 'draw_order': ['panel'], 'parts': [{'id': 'panel', 'crop_xyxy': [1, 2, 10, 12]}]} skeleton = build_skeleton(plan, assembly) self.assertEqual(skeleton['bones'], [{'name': 'root'}]) self.assertEqual(skeleton['slots'][0]['bone'], 'root') self.assertEqual(skeleton['animations'], {'setup': {}}) if __name__ == '__main__': unittest.main()