Ukawka
344 lines
19 KiB
Python
344 lines
19 KiB
Python
"""Adaptive completion planning, reference canvases and reversible native-pixel merging."""
|
|
import hashlib
|
|
import json
|
|
import math
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from PIL import Image, ImageDraw
|
|
|
|
from normalize_reference import prepare_rgba
|
|
|
|
|
|
def read(file):
|
|
return json.loads(Path(file).read_text(encoding='utf-8-sig'))
|
|
|
|
|
|
def write(file, value):
|
|
Path(file).write_text(json.dumps(value, ensure_ascii=False, indent=2, allow_nan=False), encoding='utf-8')
|
|
|
|
|
|
def binary(file, shape):
|
|
a = np.array(Image.open(file))
|
|
if a.shape != shape or not np.isin(a, [0, 255]).all():
|
|
raise ValueError(f'Expected full-canvas binary PNG: {file}')
|
|
return a > 0
|
|
|
|
|
|
def review_inputs(file, ids, shape):
|
|
"""Review masks are explicitly authored data; never guess ownership from scores."""
|
|
if not file:
|
|
return {}, {}
|
|
file = Path(file).resolve()
|
|
data = read(file)
|
|
if data.get('schema_version') != 1 or not isinstance(data.get('parts'), dict):
|
|
raise ValueError('Completion review requires schema_version=1 and parts object')
|
|
hashes = {'review.json': hashlib.sha256(file.read_bytes()).hexdigest()}
|
|
result = {}
|
|
for name, entry in data['parts'].items():
|
|
if name not in ids or not isinstance(entry, dict) or entry.get('confirmed') is not True:
|
|
raise ValueError(f'Review entry requires known ID and confirmed=true: {name}')
|
|
if set(entry) - {'confirmed', 'extension_radius', 'repair_mask', 'replace_mask', 'remove_mask'}:
|
|
raise ValueError(f'Unknown review fields: {name}')
|
|
item = {'confirmed': True}
|
|
if 'extension_radius' in entry:
|
|
radius = entry['extension_radius']
|
|
if type(radius) is not int or not 1 <= radius <= 256:
|
|
raise ValueError(f'Invalid reviewed radius: {name}')
|
|
item['extension_radius'] = radius
|
|
for key in ['repair_mask', 'replace_mask', 'remove_mask']:
|
|
if key not in entry:
|
|
continue
|
|
ref = entry[key]
|
|
if not isinstance(ref, str) or Path(ref).is_absolute():
|
|
raise ValueError('Review mask must be relative to review file')
|
|
mask_file = (file.parent / ref).resolve()
|
|
if not mask_file.is_relative_to(file.parent):
|
|
raise ValueError('Review mask path escapes review directory')
|
|
item[key] = binary(mask_file, shape)
|
|
hashes[f'{name}/{key}'] = hashlib.sha256(mask_file.read_bytes()).hexdigest()
|
|
result[name] = item
|
|
return result, hashes
|
|
|
|
|
|
def bounding(mask):
|
|
ys, xs = np.nonzero(mask)
|
|
return [int(xs.min()), int(ys.min()), int(xs.max()) + 1, int(ys.max()) + 1] if len(xs) else None
|
|
|
|
|
|
def enclosed(mask):
|
|
padded = np.pad(mask.astype(np.uint8), 1)
|
|
fill = padded.copy()
|
|
cv2.floodFill(fill, None, (0, 0), 1)
|
|
return ~fill[1:-1, 1:-1].astype(bool)
|
|
|
|
|
|
def white_image(array):
|
|
im = Image.fromarray(array).convert('RGBA')
|
|
return Image.alpha_composite(Image.new('RGBA', im.size, 'white'), im).convert('RGB')
|
|
|
|
|
|
def prepare_adaptive(directory, selected_ids, size=1024, radius=None, reference_mode='board', review_file=None):
|
|
# Import shared native coordinate helpers lazily to avoid a module cycle.
|
|
from completion_geometry import localize, allowed_regions
|
|
out = Path(directory)
|
|
plan, seg = read(out / 'semantic-parts.json'), read(out / 'segmentation.json')
|
|
rgba = np.array(Image.open(out / 'source.png').convert('RGBA'))
|
|
h, w = rgba.shape[:2]
|
|
ids = {p['id']: p for p in plan['parts']}
|
|
if set(selected_ids) - set(ids):
|
|
raise ValueError('Unknown requested part ID')
|
|
records = {p['id']: p for p in seg['parts']}
|
|
if len(records) != len(seg['parts']) or set(records) - set(ids):
|
|
raise ValueError('Invalid segmentation IDs')
|
|
masks = {name: binary(out / 'masks' / f'{name}.png', (h, w)) for name in records}
|
|
for name, mask in masks.items():
|
|
if not mask.any() or np.any(mask & (rgba[:, :, 3] == 0)):
|
|
raise ValueError(f'Empty mask or mask outside original alpha: {name}')
|
|
reviews, review_hashes = review_inputs(review_file, ids, (h, w))
|
|
coverage = np.zeros((h, w), np.uint16)
|
|
for mask in masks.values():
|
|
coverage += mask
|
|
tasks, skipped, assessments = [], [], []
|
|
(out / 'preservation').mkdir(exist_ok=True)
|
|
for part in plan['parts']:
|
|
name = part['id']
|
|
if selected_ids and name not in selected_ids:
|
|
continue
|
|
if name not in masks or part['visibility'] == 'hidden':
|
|
skipped.append({'id': name, 'reason': 'NO_VISIBLE_SAM_MASK'})
|
|
continue
|
|
visible = masks[name]
|
|
box = bounding(visible)
|
|
review = reviews.get(name, {})
|
|
resolved_radius = review.get('extension_radius', radius if radius is not None else
|
|
int(np.clip(math.ceil(min(box[2]-box[0], box[3]-box[1]) * .15), 4, 96)))
|
|
inferred = dict(part)
|
|
inferred['completion'] = {
|
|
'occlusion_fill': bool(part['occluded_by']),
|
|
'joint_extension_parent': part['rig_role'] == 'bone_part' and part['parent'] is not None,
|
|
'joint_extension_child': part['rig_role'] == 'bone_part' and any(
|
|
p['parent'] == name and p['rig_role'] == 'bone_part' for p in plan['parts'])}
|
|
_, regions, warnings = allowed_regions(inferred, plan, masks, visible, resolved_radius)
|
|
# Only enclosed source foreground without another mask owner is an automatic
|
|
# repair candidate. Open gaps and ambiguous ownership stay diagnostic.
|
|
gap = enclosed(visible) & (rgba[:, :, 3] > 0) & (coverage == 0)
|
|
cleanup = np.zeros_like(visible)
|
|
for front in part['occluded_by']:
|
|
p = ids[front]
|
|
if front in masks and p['rig_role'] != 'bone_part' and p['bind_to'] == name:
|
|
cleanup |= visible & masks[front]
|
|
replace = review.get('replace_mask', np.zeros_like(visible))
|
|
removed = review.get('remove_mask', np.zeros_like(visible))
|
|
repair = review.get('repair_mask', np.zeros_like(visible))
|
|
if np.any((replace | removed) & ~visible) or np.any(repair & visible) or np.any(replace & removed):
|
|
raise ValueError(f'{name}: replacement/removal must be disjoint subsets of SAM; repair must be outside SAM')
|
|
regions['segmentation_gap'] = gap | repair
|
|
regions['attachment_cleanup'] = cleanup & ~removed
|
|
regions['reviewed_replacement'] = replace
|
|
allowed = np.zeros_like(visible)
|
|
for region in regions.values():
|
|
allowed |= region
|
|
allowed &= ~removed
|
|
protected = visible & ~allowed & ~removed
|
|
unexpected = []
|
|
duplicates = []
|
|
for other, mask in masks.items():
|
|
if other == name:
|
|
continue
|
|
overlap = int((mask & visible).sum())
|
|
related = other in part['occluded_by'] or name in ids[other]['occluded_by']
|
|
if overlap and not related:
|
|
unexpected.append({'id': other, 'pixels': overlap})
|
|
if overlap / max(1, int((mask | visible).sum())) >= .8:
|
|
duplicates.append(other)
|
|
assessment = {'id': name, 'original_needs_completion': part['needs_completion'],
|
|
'inferred_completion': inferred['completion'], 'radius': resolved_radius,
|
|
'region_pixels': {k: int(v.sum()) for k, v in regions.items()},
|
|
'protected_pixels': int(protected.sum()), 'removed_pixels': int(removed.sum()),
|
|
'unexpected_overlaps': unexpected, 'review_confirmed': bool(review)}
|
|
assessments.append(assessment)
|
|
if duplicates and not review:
|
|
skipped.append({'id': name, 'reason': 'AMBIGUOUS_MASK_REQUIRES_REVIEW', 'overlaps': duplicates})
|
|
continue
|
|
if not protected.any() or not (protected & (rgba[:, :, 3] >= 240)).any():
|
|
skipped.append({'id': name, 'reason': 'NO_PRESERVED_ALIGNMENT_ANCHOR'})
|
|
continue
|
|
if not allowed.any() and not removed.any():
|
|
skipped.append({'id': name, 'reason': 'NO_OBSERVED_EDIT_REGION'})
|
|
continue
|
|
region_box = bounding(visible | allowed)
|
|
bw, bh = region_box[2] - region_box[0], region_box[3] - region_box[1]
|
|
native_size = max(8, math.ceil(max(bw, bh) * 1.15))
|
|
origin = [region_box[0] - (native_size-bw)//2, region_box[1] - (native_size-bh)//2]
|
|
work_size = int(size * .68) if reference_mode == 'board' else size
|
|
offset = [int(size * .02), (size-work_size)//2] if reference_mode == 'board' else [0, 0]
|
|
scale = work_size / native_size
|
|
folder = out / 'tasks' / name
|
|
folder.mkdir(parents=True, exist_ok=True)
|
|
native_source = localize(rgba, origin, native_size)
|
|
native_visible = localize(visible, origin, native_size)
|
|
native_keep = localize(protected, origin, native_size)
|
|
native_edit = localize(allowed, origin, native_size)
|
|
native_remove = localize(removed, origin, native_size)
|
|
cut = native_source.copy()
|
|
cut[~native_keep] = 0
|
|
Image.fromarray(cut).save(folder / 'native-visible.png')
|
|
Image.fromarray(native_source).save(folder / 'native-source.png')
|
|
for label, mask in [('keep', native_keep), ('allowed', native_edit), ('remove', native_remove), ('sam', native_visible)]:
|
|
Image.fromarray(mask.astype(np.uint8)*255).save(folder / f'native-{label}.png')
|
|
|
|
def project(array, is_mask=False):
|
|
small = np.array(Image.fromarray(array.astype(np.uint8)*255 if is_mask else array).resize(
|
|
(work_size, work_size), Image.Resampling.NEAREST if is_mask else Image.Resampling.LANCZOS))
|
|
canvas = np.zeros((size, size) + small.shape[2:], np.uint8)
|
|
canvas[offset[1]:offset[1]+work_size, offset[0]:offset[0]+work_size] = small
|
|
return canvas
|
|
|
|
model_visible = project(cut)
|
|
Image.fromarray(model_visible).save(folder / 'visible.png')
|
|
model_allowed = project(native_edit, True) > 0
|
|
model_remove = project(native_remove, True) > 0
|
|
Image.fromarray(model_allowed.astype(np.uint8)*255).save(folder / 'allowed-mask.png')
|
|
mask_rgba = np.full((size, size, 4), 255, np.uint8)
|
|
mask_rgba[model_allowed | model_remove, 3] = 0
|
|
Image.fromarray(mask_rgba).save(folder / 'edit-mask.png')
|
|
context = white_image(project(native_source))
|
|
context.save(folder / 'context.png')
|
|
full_reference = white_image(rgba)
|
|
full_reference.thumbnail((size, size), Image.Resampling.LANCZOS)
|
|
full_reference.save(folder / 'reference.png')
|
|
model_input = white_image(model_visible)
|
|
panel_boxes = []
|
|
if reference_mode == 'board':
|
|
draw = ImageDraw.Draw(model_input)
|
|
x = int(size * .74)
|
|
draw.line([(x-4, 0), (x-4, size)], fill='#777777', width=1)
|
|
for label, image, top in [('LOCAL REFERENCE', white_image(native_source), 0),
|
|
('FULL REFERENCE', full_reference.copy(), size//2)]:
|
|
draw.text((x, top+2), label, fill='black')
|
|
image.thumbnail((size-x-4, size//2-24), Image.Resampling.LANCZOS)
|
|
model_input.paste(image, (x, top+22))
|
|
panel_boxes.append([x, top+22, x+image.width, top+22+image.height])
|
|
model_input.save(folder / 'input.png')
|
|
bounds, native_counts = {}, {}
|
|
for kind, region in regions.items():
|
|
native = localize(region & allowed, origin, native_size)
|
|
model = project(native, True)
|
|
Image.fromarray(native.astype(np.uint8)*255).save(folder / f'native-{kind}.png')
|
|
Image.fromarray(model).save(folder / f'{kind}.png')
|
|
bounds[kind] = bounding(model > 0)
|
|
native_counts[kind] = int(native.sum())
|
|
preservation = {}
|
|
for key, mask in [('preserve_mask', protected), ('editable_mask', allowed), ('removed_mask', removed)]:
|
|
relative = f'preservation/{name}-{key}.png'
|
|
Image.fromarray(mask.astype(np.uint8)*255).save(out / relative)
|
|
preservation[key] = relative
|
|
if unexpected:
|
|
warnings.append('UNEXPECTED_OVERLAP_REVIEW_REQUIRED')
|
|
if cleanup.any():
|
|
warnings.append('ATTACHMENT_BASE_CLEANUP_NEEDS_REVIEW')
|
|
task = {'id': name, 'semantic': part['semantic'], 'profile': 'adaptive',
|
|
'completion': inferred['completion'], 'original_completion': part['completion'],
|
|
'origin_xy': origin, 'native_size': native_size, 'canvas_size': [size, size], 'source_size': [w, h],
|
|
'extension_radius': resolved_radius, 'allowed_pixels': int(native_edit.sum()),
|
|
'removed_pixels': int(native_remove.sum()), 'regions': native_counts,
|
|
'region_bounds_xyxy': bounds, 'occluded_by': part['occluded_by'],
|
|
'warnings': records[name].get('warnings', []) + warnings, 'preservation': preservation,
|
|
'reference_mode': reference_mode, 'reference_panel_boxes': panel_boxes,
|
|
'transform': {'scale': scale, 'working_xyxy': [*offset, offset[0]+work_size, offset[1]+work_size],
|
|
'source_to_canvas': [scale, 0, offset[0]-scale*origin[0]+(scale-1)/2,
|
|
0, scale, offset[1]-scale*origin[1]+(scale-1)/2],
|
|
'canvas_to_source': [1/scale, 0, origin[0]-offset[0]/scale+(1/scale-1)/2,
|
|
0, 1/scale, origin[1]-offset[1]/scale+(1/scale-1)/2]},
|
|
'semantic_verified': False, 'status': 'prepared'}
|
|
write(folder / 'task.json', task)
|
|
tasks.append(task)
|
|
result = {'schema_version': 2, 'profile': 'adaptive', 'tasks': tasks, 'skipped': skipped,
|
|
'assessments': assessments, 'review_hashes': review_hashes, 'reference_mode': reference_mode,
|
|
'canvas_size': size, 'quality_status': 'needs_review'}
|
|
write(out / 'completion-plan.json', result)
|
|
return result
|
|
|
|
|
|
def merge_adaptive(folder, candidate_file, output_file, size_policy):
|
|
folder = Path(folder)
|
|
task = read(folder / 'task.json')
|
|
candidate = Image.open(candidate_file).convert('RGBA')
|
|
qa = {'accepted': False, 'semantic_verified': False, 'reasons': [], 'profile': 'adaptive',
|
|
'expected_size': task['canvas_size'], 'returned_size': list(candidate.size),
|
|
'size_policy': size_policy, 'canvas_resampled': False, 'transform': task['transform']}
|
|
tw, th = task['canvas_size']
|
|
if candidate.size != (tw, th):
|
|
if size_policy == 'strict' or candidate.width * th != candidate.height * tw:
|
|
qa['reasons'].append('CANVAS_SIZE_MISMATCH' if size_policy == 'strict' else 'CANVAS_ASPECT_RATIO_MISMATCH')
|
|
return qa
|
|
qa.update(canvas_resampled=True, canvas_scale=tw/candidate.width, resampling='LANCZOS')
|
|
candidate = candidate.resize((tw, th), Image.Resampling.LANCZOS)
|
|
candidate.save(Path(output_file).with_name(Path(output_file).stem + '-normalized.png'))
|
|
# Record raw behavior before masking/restoration; an accepted HTTP response is
|
|
# not evidence that the relay/model respected the supplied edit mask.
|
|
expected_input = np.array(Image.open(folder / 'input.png').convert('RGB')).astype(float)
|
|
raw = np.array(white_image(np.array(candidate))).astype(float)
|
|
editable = np.array(Image.open(folder / 'edit-mask.png'))[:, :, 3] == 0
|
|
outside_delta = np.abs(raw - expected_input).mean(axis=2)
|
|
qa['raw_outside_edit_mae'] = float(outside_delta[~editable].mean()) if (~editable).any() else 0
|
|
qa['raw_outside_edit_changed_pixels'] = int(((outside_delta > 30) & ~editable).sum())
|
|
rect = task['transform']['working_xyxy']
|
|
try:
|
|
cropped, preprocessing = prepare_rgba(candidate.crop(rect))
|
|
except ValueError as error:
|
|
qa.update(reasons=['CANNOT_EXTRACT_FOREGROUND'], detail=str(error))
|
|
return qa
|
|
side = task['native_size']
|
|
generated = np.array(cropped.resize((side, side), Image.Resampling.LANCZOS))
|
|
original = np.array(Image.open(folder / 'native-visible.png').convert('RGBA'))
|
|
source = np.array(Image.open(folder / 'native-source.png').convert('RGBA'))
|
|
keep = binary(folder / 'native-keep.png', (side, side))
|
|
allowed = binary(folder / 'native-allowed.png', (side, side))
|
|
removed = binary(folder / 'native-remove.png', (side, side))
|
|
# Judge alignment before the inverse resampling. Alpha at native boundary
|
|
# pixels changes during interpolation even for an exactly aligned candidate.
|
|
model_original = np.array(Image.open(folder / 'visible.png').convert('RGBA').crop(rect))
|
|
model_generated = np.array(cropped)
|
|
opaque = model_original[:, :, 3] >= 240
|
|
if not opaque.any():
|
|
qa['reasons'].append('NO_OPAQUE_ALIGNMENT_ANCHOR')
|
|
return qa
|
|
rgb = np.array(white_image(model_generated)).astype(float)
|
|
mae = float(np.abs(rgb[opaque] - model_original[:, :, :3][opaque]).mean())
|
|
retained = float((model_generated[:, :, 3][opaque] >= 240).mean())
|
|
proposed = (generated[:, :, 3] > 0) & ~keep
|
|
additions = proposed & allowed
|
|
_, labels = cv2.connectedComponents((keep | additions).astype(np.uint8), connectivity=8)
|
|
touching = np.unique(labels[keep])
|
|
connected = additions & np.isin(labels, touching[touching != 0])
|
|
counts = {kind: int((connected & binary(folder / f'native-{kind}.png', (side, side))).sum())
|
|
for kind in task['regions']}
|
|
qa.update(alignment_mae=mae, visible_retention=retained, alignment_space='request_work_panel', new_pixels=int(connected.sum()),
|
|
removed_pixels=int(removed.sum()), region_additions=counts,
|
|
unfulfilled_regions=[k for k, n in task['regions'].items() if n and not counts[k]],
|
|
discarded_outside_allowed=int((proposed & ~allowed).sum()), preprocessing=preprocessing)
|
|
if mae > 30 or retained < .95:
|
|
qa['reasons'].append('VISIBLE_ALIGNMENT_MISMATCH')
|
|
if not connected.any() and not removed.any():
|
|
qa['reasons'].append('NO_CONNECTED_NEW_PIXELS')
|
|
for kind in ['attachment_cleanup', 'reviewed_replacement']:
|
|
if task['regions'][kind] and counts[kind] / task['regions'][kind] < .95:
|
|
qa['reasons'].append(f'INCOMPLETE_{kind.upper()}')
|
|
if qa['reasons']:
|
|
return qa
|
|
merged = np.zeros_like(original)
|
|
merged[connected] = generated[connected]
|
|
# Visible segmentation gaps already have ground-truth texture in source.
|
|
# Restore it where the proposed fill connects, rather than trusting a repaint.
|
|
repair = binary(folder / 'native-segmentation_gap.png', (side, side)) & connected & (source[:, :, 3] > 0)
|
|
merged[repair] = source[repair]
|
|
merged[keep] = original[keep]
|
|
merged[removed] = 0
|
|
qa.update(accepted=True, quality_status='needs_review', visible_pixels_exact=True,
|
|
preservation_scope='preserve_mask', repaired_source_pixels=int(repair.sum()))
|
|
Image.fromarray(merged).save(output_file)
|
|
return qa
|