"""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'' if task.get('image') else '' warnings = html.escape(', '.join(task.get('upstream_warnings', []))) rows.append(f'{html.escape(name)}

{warnings}

任务参数 · 提示词{html.escape(task["status"])}

原图局部参考

{view}
{html.escape(json.dumps(qa, ensure_ascii=False, indent=2))}
') write_json(out/'completion-result.json', result) (out/'report.html').write_text('''部件补全

第三阶段：部件补全

新增区域受允许掩码限制，原图可见像素逐像素恢复。几何通过仅表示候选可供复核，不代表遮挡和关节纹理已正确补全。上游 SAM 错分仍需修正。

'''+''.join(rows)+'
部件 / 上游告警状态原可见区域允许新增区域（白色）补全候选几何检查

跳过

'+html.escape(json.dumps(result['skipped'], ensure_ascii=False, indent=2))+'
', 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()