Ukawka
144 lines
7.4 KiB
Python
144 lines
7.4 KiB
Python
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_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()
|