206 lines
11 KiB
Python
206 lines
11 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 localize, merge_candidate, finalize
|
|
from completion_planning import prepare_adaptive, read, write, review_inputs
|
|
from spine_assembly import load_assets
|
|
from test_completion import make_fixture
|
|
|
|
|
|
def candidate_for(folder, native):
|
|
"""Independent fixture: insert known native ground truth into documented work rectangle."""
|
|
folder = Path(folder)
|
|
task = read(folder / 'task.json')
|
|
rect = task['transform']['working_xyxy']
|
|
canvas = Image.open(folder / 'input.png').convert('RGBA')
|
|
content = Image.fromarray(native).resize((rect[2]-rect[0], rect[3]-rect[1]), Image.Resampling.NEAREST)
|
|
canvas.paste(content, (rect[0], rect[1]))
|
|
target = folder / 'test-candidate.png'
|
|
canvas.save(target)
|
|
return target
|
|
|
|
|
|
def mock_candidate(folder):
|
|
folder = Path(folder)
|
|
native = np.array(Image.open(folder / 'native-visible.png'))
|
|
edit = np.array(Image.open(folder / 'native-allowed.png')) > 0
|
|
native[edit] = [180, 50, 20, 255]
|
|
return candidate_for(folder, native)
|
|
|
|
|
|
class AdaptiveTests(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)
|
|
|
|
def task(self, mode='board', size=256, review=None):
|
|
plan = prepare_adaptive(self.root, ['torso'], size, 16, mode, review)
|
|
return plan['tasks'][0], self.root / 'tasks' / 'torso'
|
|
|
|
def test_small_part_scales_up_roundtrip_and_reference_panels_are_excluded(self):
|
|
task, folder = self.task()
|
|
self.assertGreater(task['transform']['scale'], 1)
|
|
self.assertEqual(len(task['reference_panel_boxes']), 2)
|
|
forward = task['transform']['source_to_canvas']
|
|
inverse = task['transform']['canvas_to_source']
|
|
for x, y in [(20, 20), (49, 39), (0, 79)]:
|
|
mx, my = forward[0]*x+forward[2], forward[4]*y+forward[5]
|
|
self.assertAlmostEqual(inverse[0]*mx+inverse[2], x)
|
|
self.assertAlmostEqual(inverse[4]*my+inverse[5], y)
|
|
native = np.array(Image.open(folder / 'native-visible.png'))
|
|
edit = np.array(Image.open(folder / 'native-allowed.png')) > 0
|
|
native[edit] = [180, 50, 20, 255]
|
|
output = folder / 'merged.png'
|
|
qa = merge_candidate(folder, candidate_for(folder, native), output)
|
|
self.assertTrue(qa['accepted'], qa)
|
|
keep = np.array(Image.open(folder / 'native-keep.png')) > 0
|
|
np.testing.assert_array_equal(np.array(Image.open(output))[keep], native[keep])
|
|
write(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(self.root / 'completion-result.json')['parts'][0]
|
|
self.assertEqual(result['crop_xyxy'], [20, 20, 50, 40])
|
|
self.assertIn('preservation', result)
|
|
self.assertTrue((folder / 'reference.png').exists())
|
|
mask = np.array(Image.open(folder / 'edit-mask.png'))
|
|
for x0, y0, x1, y1 in task['reference_panel_boxes']:
|
|
self.assertTrue(np.all(mask[y0:y1, x0:x1, 3] == 255))
|
|
|
|
def test_actual_hole_creates_task_despite_false_semantic_flag(self):
|
|
plan = read(self.root / 'semantic-parts.json')
|
|
plan['parts'][0].update(needs_completion=False, visibility='visible', occluded_by=[],
|
|
completion={k: False for k in plan['parts'][0]['completion']})
|
|
write(self.root / 'semantic-parts.json', plan)
|
|
mask = np.array(Image.open(self.root / 'masks/torso.png'))
|
|
mask[26:30, 26:30] = 0
|
|
Image.fromarray(mask).save(self.root / 'masks/torso.png')
|
|
task, folder = self.task()
|
|
self.assertEqual(task['regions']['segmentation_gap'], 16)
|
|
native = np.array(Image.open(folder / 'native-visible.png'))
|
|
edit = np.array(Image.open(folder / 'native-allowed.png')) > 0
|
|
native[edit] = [200, 40, 10, 255]
|
|
output = folder / 'merged.png'
|
|
qa = merge_candidate(folder, candidate_for(folder, native), output)
|
|
self.assertTrue(qa['accepted'], qa)
|
|
# Repair restores source texture, even though the model proposed a different colour.
|
|
gap = np.array(Image.open(folder / 'native-segmentation_gap.png')) > 0
|
|
source = np.array(Image.open(folder / 'native-source.png'))
|
|
np.testing.assert_array_equal(np.array(Image.open(output))[gap], source[gap])
|
|
|
|
def test_attachment_cleanup_releases_only_declared_overlay_and_spine_accepts_contract(self):
|
|
# Move collar inside torso to represent an overlaid detail needing a clean base.
|
|
mask = np.zeros((80, 80), np.uint8)
|
|
mask[26:30, 26:30] = 255
|
|
Image.fromarray(mask).save(self.root / 'masks/collar.png')
|
|
task, folder = self.task('both')
|
|
self.assertEqual(task['regions']['attachment_cleanup'], 16)
|
|
native = np.array(Image.open(folder / 'native-visible.png'))
|
|
edit = np.array(Image.open(folder / 'native-allowed.png')) > 0
|
|
native[edit] = [20, 100, 200, 255]
|
|
output = folder / 'merged.png'
|
|
qa = merge_candidate(folder, candidate_for(folder, native), output)
|
|
self.assertTrue(qa['accepted'], qa)
|
|
write(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)
|
|
data = load_assets(self.root, self.root)
|
|
self.assertEqual(data[2]['parts'][0]['source'], 'completion')
|
|
# Contract tampering cannot turn a protected pixel into an untracked edit.
|
|
result = read(self.root / 'completion-result.json')['parts'][0]
|
|
image = np.array(Image.open(self.root / result['image']))
|
|
image[0, 0] = [255, 0, 0, 255]
|
|
Image.fromarray(image).save(self.root / result['image'])
|
|
with self.assertRaisesRegex(ValueError, 'preserve original visible'):
|
|
load_assets(self.root, self.root)
|
|
|
|
def test_review_replacement_removal_and_invalid_partition(self):
|
|
replace = np.zeros((80, 80), np.uint8)
|
|
replace[26:29, 26:29] = 255
|
|
remove = np.zeros_like(replace)
|
|
remove[35:40, 35:40] = 255
|
|
Image.fromarray(replace).save(self.root / 'replace.png')
|
|
Image.fromarray(remove).save(self.root / 'remove.png')
|
|
review = self.root / 'review-input.json'
|
|
write(review, {'schema_version': 1, 'parts': {'torso': {'confirmed': True,
|
|
'replace_mask': 'replace.png', 'remove_mask': 'remove.png', 'extension_radius': 8}}})
|
|
task, folder = self.task(review=review)
|
|
self.assertEqual(task['extension_radius'], 8)
|
|
native = np.array(Image.open(folder / 'native-visible.png'))
|
|
edit = np.array(Image.open(folder / 'native-allowed.png')) > 0
|
|
native[edit] = [180, 50, 20, 255]
|
|
qa = merge_candidate(folder, candidate_for(folder, native), folder / 'merged.png')
|
|
self.assertTrue(qa['accepted'], qa)
|
|
deleted = np.array(Image.open(folder / 'native-remove.png')) > 0
|
|
self.assertTrue(np.all(np.array(Image.open(folder / 'merged.png'))[deleted] == 0))
|
|
Image.fromarray(remove).save(self.root / 'replace.png')
|
|
with self.assertRaisesRegex(ValueError, 'disjoint'):
|
|
self.task(review=review)
|
|
|
|
def test_large_part_downscales_and_invalid_review_cannot_escape(self):
|
|
# Upsize the entire source and masks to exceed request dimensions.
|
|
for file in [self.root / 'source.png', *list((self.root / 'masks').glob('*.png'))]:
|
|
im = Image.open(file)
|
|
im.resize((800, 800), Image.Resampling.NEAREST).save(file)
|
|
task, folder = self.task(size=128)
|
|
self.assertLess(task['transform']['scale'], 1)
|
|
native = np.array(Image.open(folder / 'native-visible.png'))
|
|
native[np.array(Image.open(folder / 'native-allowed.png')) > 0] = [180, 50, 20, 255]
|
|
output = folder / 'merged.png'
|
|
qa = merge_candidate(folder, candidate_for(folder, native), output)
|
|
self.assertTrue(qa['accepted'], qa)
|
|
keep = np.array(Image.open(folder / 'native-keep.png')) > 0
|
|
np.testing.assert_array_equal(np.array(Image.open(output))[keep], native[keep])
|
|
review = self.root / 'bad-review.json'
|
|
write(review, {'schema_version': 1, 'parts': {'torso': {'confirmed': True, 'replace_mask': '../escape.png'}}})
|
|
with self.assertRaisesRegex(ValueError, 'escapes'):
|
|
review_inputs(review, ['torso'], (800, 800))
|
|
|
|
def test_adaptive_rejects_movement_and_missing_cleanup(self):
|
|
task, folder = self.task()
|
|
candidate = mock_candidate(folder)
|
|
image = np.array(Image.open(candidate))
|
|
Image.fromarray(np.roll(image, 30, axis=1)).save(candidate)
|
|
qa = merge_candidate(folder, candidate, folder / 'bad.png')
|
|
self.assertFalse(qa['accepted'])
|
|
self.assertIn('VISIBLE_ALIGNMENT_MISMATCH', qa['reasons'])
|
|
mask = np.zeros((80, 80), np.uint8)
|
|
mask[26:30, 26:30] = 255
|
|
Image.fromarray(mask).save(self.root / 'masks/collar.png')
|
|
_, folder = self.task('both')
|
|
no_fill = np.array(Image.open(folder / 'native-visible.png'))
|
|
qa = merge_candidate(folder, candidate_for(folder, no_fill), folder / 'bad-cleanup.png')
|
|
self.assertFalse(qa['accepted'])
|
|
self.assertIn('INCOMPLETE_ATTACHMENT_CLEANUP', qa['reasons'])
|
|
|
|
def test_duplicate_masks_blocked_until_reviewed(self):
|
|
plan = read(self.root / 'semantic-parts.json')
|
|
plan['parts'][0]['occluded_by'] = []
|
|
write(self.root / 'semantic-parts.json', plan)
|
|
Image.open(self.root / 'masks/torso.png').save(self.root / 'masks/collar.png')
|
|
result = prepare_adaptive(self.root, ['torso'])
|
|
self.assertEqual(result['skipped'][0]['reason'], 'AMBIGUOUS_MASK_REQUIRES_REVIEW')
|
|
|
|
def test_spine_strict_policy_respects_rebuilt_need_even_if_original_flag_was_false(self):
|
|
plan = read(self.root / 'semantic-parts.json')
|
|
plan['parts'][0].update(needs_completion=False, visibility='visible', occluded_by=[],
|
|
completion={k: False for k in plan['parts'][0]['completion']})
|
|
write(self.root / 'semantic-parts.json', plan)
|
|
write(self.root / 'completion-result.json', {'schema_version': 2, 'profile': 'adaptive', 'parts': [
|
|
{'id': 'torso', 'status': 'rejected', 'qa': {'accepted': False}}], 'skipped': []})
|
|
with self.assertRaisesRegex(ValueError, 'Completion required for torso'):
|
|
load_assets(self.root, self.root, fallback_policy='error')
|
|
data = load_assets(self.root, self.root)
|
|
self.assertEqual(data[2]['parts'][0]['reason'], 'COMPLETION_REJECTED')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|