Ukawka
248 lines
12 KiB
Python
248 lines
12 KiB
Python
"""Prepare bounded completion canvases and merge generated pixels conservatively."""
|
|
import argparse
|
|
import html
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
from normalize_reference import prepare_rgba
|
|
|
|
|
|
def write_json(file, value):
|
|
Path(file).write_text(json.dumps(value, ensure_ascii=False, indent=2), encoding='utf-8')
|
|
|
|
|
|
def read_json(file):
|
|
return json.loads(Path(file).read_text(encoding='utf-8-sig'))
|
|
|
|
|
|
def localize(array, origin, size):
|
|
"""Copy source coordinates to a padded local canvas, without scaling."""
|
|
out = np.zeros((size, size) + array.shape[2:], dtype=array.dtype)
|
|
ox, oy = origin
|
|
h, w = array.shape[:2]
|
|
x0, y0, x1, y1 = max(0, ox), max(0, oy), min(w, ox + size), min(h, oy + size)
|
|
if x1 > x0 and y1 > y0:
|
|
out[y0-oy:y1-oy, x0-ox:x1-ox] = array[y0:y1, x0:x1]
|
|
return out
|
|
|
|
|
|
def allowed_regions(part, plan, masks, visible, radius):
|
|
h, w = visible.shape
|
|
distance = cv2.distanceTransform((~visible).astype(np.uint8), cv2.DIST_L2, 5)
|
|
band = (distance <= radius) & ~visible
|
|
regions, warnings = {}, []
|
|
flags = part['completion']
|
|
if flags['occlusion_fill']:
|
|
occluders = np.zeros_like(visible)
|
|
for name in part['occluded_by']:
|
|
if name in masks:
|
|
occluders |= masks[name]
|
|
regions['occlusion_fill'] = band & occluders
|
|
if not regions['occlusion_fill'].any():
|
|
warnings.append('NO_OCCLUSION_REGION')
|
|
anchors = {}
|
|
if flags['joint_extension_parent']:
|
|
anchors['joint_extension_parent'] = [part['pivot_norm']]
|
|
if flags['joint_extension_child']:
|
|
anchors['joint_extension_child'] = [p['pivot_norm'] for p in plan['parts']
|
|
if p['parent'] == part['id'] and p['rig_role'] == 'bone_part']
|
|
for kind, points in anchors.items():
|
|
disks = np.zeros((h, w), np.uint8)
|
|
for x, y in points:
|
|
cv2.circle(disks, (min(w-1, round(x*w)), min(h-1, round(y*h))), radius, 1, -1)
|
|
regions[kind] = band & (disks > 0)
|
|
if not regions[kind].any():
|
|
warnings.append(f'NO_{kind.upper()}_REGION')
|
|
allowed = np.zeros_like(visible)
|
|
for region in regions.values():
|
|
allowed |= region
|
|
return allowed, regions, warnings
|
|
|
|
|
|
def prepare(directory, selected_ids, size=1024, radius=32):
|
|
out = Path(directory)
|
|
plan, segmentation = read_json(out/'semantic-parts.json'), read_json(out/'segmentation.json')
|
|
source = np.array(Image.open(out/'source.png').convert('RGBA'))
|
|
h, w = source.shape[:2]
|
|
known = {p['id']: p for p in plan['parts']}
|
|
if set(selected_ids) - set(known):
|
|
raise ValueError('Unknown requested part ID')
|
|
records = {p['id']: p for p in segmentation['parts']}
|
|
if len(records) != len(segmentation['parts']) or set(records) - set(known):
|
|
raise ValueError('Segmentation IDs do not match semantic plan')
|
|
masks = {}
|
|
for name in records:
|
|
mask = np.array(Image.open(out/'masks'/f'{name}.png').convert('L'))
|
|
if mask.shape != (h, w) or not np.isin(mask, [0, 255]).all():
|
|
raise ValueError(f'{name}: mask must be binary and match the source canvas')
|
|
mask = mask > 0
|
|
if not mask.any() or (mask & (source[:, :, 3] == 0)).any():
|
|
raise ValueError(f'{name}: empty mask or mask outside original alpha')
|
|
masks[name] = mask
|
|
tasks, skipped = [], []
|
|
for part in plan['parts']:
|
|
name = part['id']
|
|
if selected_ids and name not in selected_ids:
|
|
continue
|
|
if not part['needs_completion']:
|
|
skipped.append({'id': name, 'reason': 'NO_COMPLETION_NEEDED'})
|
|
continue
|
|
if name not in masks or part['visibility'] == 'hidden':
|
|
skipped.append({'id': name, 'reason': 'NO_VISIBLE_SAM_MASK'})
|
|
continue
|
|
visible = masks[name]
|
|
allowed, regions, warnings = allowed_regions(part, plan, masks, visible, radius)
|
|
if not allowed.any():
|
|
skipped.append({'id': name, 'reason': 'NO_ALLOWED_EXTENSION', 'warnings': warnings})
|
|
continue
|
|
ys, xs = np.nonzero(visible | allowed)
|
|
bw, bh = int(xs.max()-xs.min()+1), int(ys.max()-ys.min()+1)
|
|
if max(bw, bh) + 32 > size:
|
|
raise ValueError(f'{name}: completion canvas too small; increase --completion-size (no resizing is performed)')
|
|
origin = [int(xs.min()) - (size-bw)//2, int(ys.min()) - (size-bh)//2]
|
|
folder = out/'tasks'/name
|
|
folder.mkdir(parents=True, exist_ok=True)
|
|
cut = source.copy()
|
|
cut[~visible] = 0
|
|
original = localize(cut, origin, size)
|
|
editable = localize(allowed, origin, size)
|
|
Image.fromarray(original).save(folder/'visible.png')
|
|
white = Image.new('RGBA', (size, size), 'white')
|
|
Image.alpha_composite(white, Image.fromarray(original)).convert('RGB').save(folder/'input.png')
|
|
context = Image.fromarray(localize(source, origin, size))
|
|
Image.alpha_composite(white, context).convert('RGB').save(folder/'context.png')
|
|
edit_mask = np.full((size, size, 4), 255, np.uint8)
|
|
edit_mask[editable, 3] = 0
|
|
Image.fromarray(edit_mask).save(folder/'edit-mask.png')
|
|
Image.fromarray(editable.astype(np.uint8)*255).save(folder/'allowed-mask.png')
|
|
local_regions, region_bounds = {}, {}
|
|
for kind, region in regions.items():
|
|
local = localize(region, origin, size)
|
|
Image.fromarray(local.astype(np.uint8)*255).save(folder/f'{kind}.png')
|
|
local_regions[kind] = int(local.sum())
|
|
ry, rx = np.nonzero(local)
|
|
region_bounds[kind] = [int(rx.min()), int(ry.min()), int(rx.max())+1, int(ry.max())+1] if len(rx) else None
|
|
task = {'id': name, 'semantic': part['semantic'], 'completion': part['completion'],
|
|
'origin_xy': origin, 'canvas_size': [size, size], 'source_size': [w, h],
|
|
'allowed_pixels': int(editable.sum()), 'regions': local_regions,
|
|
'region_bounds_xyxy': region_bounds, 'occluded_by': part['occluded_by'],
|
|
'warnings': records[name].get('warnings', []) + warnings,
|
|
'semantic_verified': False, 'status': 'prepared'}
|
|
write_json(folder/'task.json', task)
|
|
tasks.append(task)
|
|
result = {'schema_version': 1, 'tasks': tasks, 'skipped': skipped,
|
|
'quality_status': 'needs_review', 'canvas_size': size, 'extension_radius': radius}
|
|
write_json(out/'completion-plan.json', result)
|
|
return result
|
|
|
|
|
|
def merge_candidate(folder, candidate_file, output_file):
|
|
folder = Path(folder)
|
|
task = read_json(folder/'task.json')
|
|
original = np.array(Image.open(folder/'visible.png').convert('RGBA'))
|
|
allowed = np.array(Image.open(folder/'allowed-mask.png').convert('L')) > 0
|
|
candidate = Image.open(candidate_file).convert('RGBA')
|
|
qa = {'accepted': False, 'semantic_verified': False, 'reasons': []}
|
|
if list(candidate.size) != task['canvas_size']:
|
|
qa['reasons'].append('CANVAS_SIZE_MISMATCH')
|
|
return qa
|
|
try:
|
|
candidate, preprocessing = prepare_rgba(candidate)
|
|
except ValueError as error:
|
|
qa['reasons'].append('CANNOT_EXTRACT_FOREGROUND')
|
|
qa['detail'] = str(error)
|
|
return qa
|
|
generated = np.array(candidate)
|
|
visible = original[:, :, 3] > 0
|
|
opaque = original[:, :, 3] >= 240
|
|
if not opaque.any():
|
|
qa['reasons'].append('NO_OPAQUE_ALIGNMENT_ANCHOR')
|
|
return qa
|
|
# Detect gross movement/redesign before restoring original pixels, otherwise
|
|
# restoration could hide a shifted generation attached at the wrong place.
|
|
white = Image.new('RGBA', candidate.size, 'white')
|
|
generated_rgb = np.array(Image.alpha_composite(white, candidate))[:, :, :3].astype(float)
|
|
mae = float(np.abs(generated_rgb[opaque] - original[:, :, :3][opaque]).mean())
|
|
retained = float((generated[:, :, 3][opaque] >= 240).mean())
|
|
proposed = (generated[:, :, 3] > 0) & ~visible
|
|
added = proposed & allowed
|
|
_, labels = cv2.connectedComponents((visible | added).astype(np.uint8), connectivity=8)
|
|
touching = np.unique(labels[visible])
|
|
connected = added & np.isin(labels, touching[touching != 0])
|
|
region_additions = {kind: int((connected & (np.array(Image.open(folder/f'{kind}.png')) > 0)).sum())
|
|
for kind in task['regions']}
|
|
missing_regions = [kind for kind, pixels in task['regions'].items() if not pixels or not region_additions[kind]]
|
|
qa.update(alignment_mae=mae, visible_retention=retained,
|
|
new_pixels=int(connected.sum()), discarded_outside_allowed=int((proposed & ~allowed).sum()),
|
|
disconnected_pixels=int((added & ~connected).sum()), region_additions=region_additions,
|
|
unfulfilled_regions=missing_regions, preprocessing=preprocessing)
|
|
if mae > 30 or retained < .95:
|
|
qa['reasons'].append('VISIBLE_ALIGNMENT_MISMATCH')
|
|
if not connected.any():
|
|
qa['reasons'].append('NO_CONNECTED_NEW_PIXELS')
|
|
if qa['reasons']:
|
|
return qa
|
|
merged = np.zeros_like(original)
|
|
merged[connected] = generated[connected]
|
|
merged[visible] = original[visible]
|
|
qa['visible_pixels_exact'] = bool(np.array_equal(merged[visible], original[visible]))
|
|
qa['accepted'] = True
|
|
qa['quality_status'] = 'needs_review'
|
|
Image.fromarray(merged).save(output_file)
|
|
return qa
|
|
|
|
|
|
def finalize(directory):
|
|
out = Path(directory)
|
|
result = read_json(out/'completion-result.json')
|
|
rows = []
|
|
for task in result['parts']:
|
|
name = task['id']
|
|
qa = task.get('qa', {})
|
|
if task.get('selected_file'):
|
|
im = Image.open(out/task['selected_file']).convert('RGBA')
|
|
bbox = im.getbbox()
|
|
crop = im.crop(bbox)
|
|
(out/'parts').mkdir(exist_ok=True)
|
|
crop.save(out/'parts'/f'{name}.png')
|
|
meta = read_json(out/'tasks'/name/'task.json')
|
|
ox, oy = meta['origin_xy']
|
|
task.update(image=f'parts/{name}.png', crop_xyxy=[bbox[0]+ox, bbox[1]+oy, bbox[2]+ox, bbox[3]+oy])
|
|
view = f'<img src="{task["image"]}">' if task.get('image') else ''
|
|
warnings = html.escape(', '.join(task.get('upstream_warnings', [])))
|
|
rows.append(f'<tr><td>{html.escape(name)}<p>{warnings}</p><a href="tasks/{name}/task.json">任务参数</a> · <a href="tasks/{name}/prompt.txt">提示词</a></td><td>{html.escape(task["status"])}</td><td><a href="tasks/{name}/visible.png"><img src="tasks/{name}/visible.png"></a></td><td><a href="tasks/{name}/allowed-mask.png"><img src="tasks/{name}/allowed-mask.png"></a><p><a href="tasks/{name}/context.png">原图局部参考</a></p></td><td>{view}</td><td><pre>{html.escape(json.dumps(qa, ensure_ascii=False, indent=2))}</pre></td></tr>')
|
|
write_json(out/'completion-result.json', result)
|
|
(out/'report.html').write_text('''<!doctype html><html lang="zh-CN"><meta charset="utf-8"><title>部件补全</title>
|
|
<style>body{font:16px system-ui;margin:24px}td,th{border:1px solid #ccc;padding:8px}img{max-width:260px;max-height:300px;background:#eee}pre{white-space:pre-wrap}table{border-collapse:collapse}</style>
|
|
<h1>第三阶段：部件补全</h1><p>新增区域受允许掩码限制，原图可见像素逐像素恢复。几何通过仅表示候选可供复核，不代表遮挡和关节纹理已正确补全。上游 SAM 错分仍需修正。</p>
|
|
<table><tr><th>部件 / 上游告警</th><th>状态</th><th>原可见区域</th><th>允许新增区域（白色）</th><th>补全候选</th><th>几何检查</th></tr>'''+''.join(rows)+'</table><h2>跳过</h2><pre>'+html.escape(json.dumps(result['skipped'], ensure_ascii=False, indent=2))+'</pre></html>', encoding='utf-8')
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('action', choices=['prepare', 'merge', 'finalize'])
|
|
parser.add_argument('--out', required=True)
|
|
parser.add_argument('--part', action='append', default=[])
|
|
parser.add_argument('--size', type=int, default=1024)
|
|
parser.add_argument('--radius', type=int, default=32)
|
|
parser.add_argument('--candidate')
|
|
parser.add_argument('--merged')
|
|
parser.add_argument('--qa')
|
|
args = parser.parse_args()
|
|
if args.action == 'prepare':
|
|
if not 128 <= args.size <= 4096 or not 1 <= args.radius <= 256:
|
|
parser.error('size must be 128..4096 and radius 1..256')
|
|
prepare(args.out, args.part, args.size, args.radius)
|
|
elif args.action == 'merge':
|
|
write_json(args.qa, merge_candidate(args.out, args.candidate, args.merged))
|
|
else:
|
|
finalize(args.out)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|