Ukawka
85 lines
4.0 KiB
Python
85 lines
4.0 KiB
Python
import copy
|
|
import json
|
|
from pathlib import Path
|
|
import tempfile
|
|
import unittest
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
from sam_segment import choose_candidate, validate_prompts
|
|
from segmentation_outputs import assess_candidate, render_outputs, write_report
|
|
|
|
|
|
class SegmentationTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.rgba = np.zeros((24, 32, 4), dtype=np.uint8)
|
|
self.rgba[4:20, 5:27] = [120, 50, 200, 255]
|
|
self.rgba[5, 6, 3] = 80
|
|
self.image = Image.fromarray(self.rgba)
|
|
self.mask = self.rgba[:, :, 3] > 0
|
|
|
|
def test_background_leak_measured_before_clipping(self):
|
|
raw = np.ones((24, 32), dtype=bool)
|
|
qa = assess_candidate(raw, self.rgba[:, :, 3], [5, 4, 27, 20], [[6, 5]])
|
|
self.assertGreater(qa['outside_alpha_ratio'], .5)
|
|
self.assertEqual(qa['seed_hit_fraction'], 1)
|
|
self.assertEqual(qa['area_px'], 16 * 22)
|
|
|
|
def test_rank_prefers_local_mask_and_never_invents_empty_mask(self):
|
|
wrong = self.mask.copy()
|
|
local = np.zeros_like(wrong)
|
|
local[8:13, 8:13] = True
|
|
best, _ = choose_candidate([wrong, local], [.95, .7], self.rgba[:, :, 3], [8, 8, 13, 13], [[10, 10]])
|
|
self.assertEqual(best['index'], 1)
|
|
best, _ = choose_candidate([np.zeros_like(local)], [.99], self.rgba[:, :, 3], [8, 8, 13, 13], [[10, 10]])
|
|
self.assertIsNone(best)
|
|
|
|
def test_prompts_reject_transparent_points_mismatched_canvas_and_duplicate_ids(self):
|
|
plan = {'parts': [{'id': 'hand', 'visibility': 'visible'}]}
|
|
prompts = {'width': 32, 'height': 24, 'parts': [{'id': 'hand', 'status': 'candidate',
|
|
'seed_source': 'alpha_roi_distance_transform', 'bbox_px_xyxy': [0, 0, 32, 24],
|
|
'positive_points_px': [[10, 10]]}]}
|
|
validate_prompts(plan, prompts, self.image)
|
|
for variant in range(3):
|
|
bad = copy.deepcopy(prompts)
|
|
if variant == 0:
|
|
bad['parts'][0]['positive_points_px'] = [[0, 0]]
|
|
elif variant == 1:
|
|
bad['width'] = 99
|
|
else:
|
|
bad['parts'].append(copy.deepcopy(bad['parts'][0]))
|
|
with self.assertRaises(ValueError):
|
|
validate_prompts(plan, bad, self.image)
|
|
|
|
def test_render_preserves_cutout_pixels_alpha_and_measures_overlap(self):
|
|
parts = [{'id': 'hand', 'semantic': 'hand', 'mask': self.mask, 'model_score': .9},
|
|
{'id': 'forearm', 'semantic': 'forearm', 'mask': self.mask, 'model_score': .8}]
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
result = render_outputs(self.image, parts, directory)
|
|
self.assertEqual(result['coverage']['coverage_ratio'], 1)
|
|
self.assertEqual(result['coverage']['overlap_pixels'], int(self.mask.sum()))
|
|
self.assertEqual(result['near_duplicate_masks'][0]['iou'], 1)
|
|
self.assertIn('NEAR_DUPLICATE_MASK', result['parts'][0]['warnings'])
|
|
cut = np.array(Image.open(Path(directory) / 'parts/hand.png'))
|
|
np.testing.assert_array_equal(cut, self.rgba[4:20, 5:27])
|
|
self.assertEqual(Image.open(Path(directory) / 'annotated.png').size, self.image.size)
|
|
sheet = np.array(Image.open(Path(directory) / 'parts-sheet.png'))
|
|
self.assertTrue(np.all(sheet[0, 0] == 255))
|
|
a, b = [p['sheet_xyxy'] for p in result['parts']]
|
|
self.assertTrue(a[2] <= b[0] or b[2] <= a[0] or a[3] <= b[1] or b[3] <= a[1])
|
|
result['skipped'] = []
|
|
write_report(directory, result)
|
|
self.assertTrue((Path(directory) / 'report.html').exists())
|
|
json.dumps(result)
|
|
|
|
def test_renderer_rejects_empty_masks_and_unsafe_ids(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
for name, mask in [('hand', np.zeros_like(self.mask)), ('../bad', self.mask)]:
|
|
with self.assertRaises(ValueError):
|
|
render_outputs(self.image, [{'id': name, 'mask': mask}], directory)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|