Ukawka
181 lines
9.2 KiB
Python
181 lines
9.2 KiB
Python
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()
|