Ukawka
130 lines
6.8 KiB
Python
130 lines
6.8 KiB
Python
#!/usr/bin/env python3
|
|
"""Alpha constraints and deterministic seed candidates, never semantic masks."""
|
|
import argparse
|
|
import json
|
|
import math
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
from normalize_reference import prepare_rgba
|
|
|
|
|
|
def pixel(point, width, height):
|
|
# Normalized points map to their containing pixel; 1 maps to the last pixel.
|
|
return min(width - 1, int(point[0] * width)), min(height - 1, int(point[1] * height))
|
|
|
|
|
|
def ground_parts(image, plan):
|
|
alpha = np.array(image.convert('RGBA'))[:, :, 3]
|
|
foreground = alpha > 0
|
|
height, width = alpha.shape
|
|
usable_alpha = bool(foreground.any() and not foreground.all())
|
|
silhouette_distance = cv2.distanceTransform(np.pad(foreground.astype(np.uint8), 1), cv2.DIST_L2, cv2.DIST_MASK_PRECISE)[1:-1, 1:-1]
|
|
diagnostics, prompts = [], []
|
|
for part in plan['parts']:
|
|
box = part['bbox_norm_xyxy']
|
|
entry = {
|
|
'id': part['id'], 'status': 'hidden' if part['visibility'] == 'hidden' else 'blocked',
|
|
'bbox_px_xyxy': None, 'bbox_foreground_ratio': None,
|
|
'foreground_bbox_px_xyxy': None, 'foreground_components': 0,
|
|
'model_positive_points': [], 'candidate_points': [], 'warnings': [],
|
|
'semantic_verified': False,
|
|
}
|
|
for point in part.get('positive_points_norm', []):
|
|
x, y = pixel(point, width, height)
|
|
in_box = bool(box and box[0] <= point[0] <= box[2] and box[1] <= point[1] <= box[3])
|
|
on_foreground = bool(foreground[y, x])
|
|
entry['model_positive_points'].append({
|
|
'point_norm': point, 'point_px': [x, y], 'alpha': int(alpha[y, x]),
|
|
'inside_bbox': in_box, 'on_foreground': on_foreground,
|
|
'status': 'UNTRUSTED_SEED' if usable_alpha and on_foreground and in_box else 'INVALID_SEED',
|
|
})
|
|
if part['visibility'] != 'hidden' and box:
|
|
x0, y0 = math.floor(box[0] * width), math.floor(box[1] * height)
|
|
x1, y1 = math.ceil(box[2] * width), math.ceil(box[3] * height)
|
|
entry['bbox_px_xyxy'] = [x0, y0, x1, y1]
|
|
roi = foreground[y0:y1, x0:x1]
|
|
entry['bbox_foreground_ratio'] = float(roi.mean()) if roi.size else 0
|
|
if not usable_alpha:
|
|
entry['warnings'].append('NO_USABLE_ALPHA')
|
|
elif not roi.any():
|
|
entry['warnings'].append('EMPTY_FOREGROUND_ROI')
|
|
else:
|
|
ys, xs = np.nonzero(roi)
|
|
entry['foreground_bbox_px_xyxy'] = [int(xs.min()) + x0, int(ys.min()) + y0, int(xs.max()) + x0 + 1, int(ys.max()) + y0 + 1]
|
|
count, _ = cv2.connectedComponents(roi.astype(np.uint8), connectivity=8)
|
|
entry['foreground_components'] = count - 1
|
|
# Pad the ROI with background so an opaque ROI has a finite
|
|
# distance, and candidates stay away from both silhouette/box edges.
|
|
distance = cv2.distanceTransform(np.pad(roi.astype(np.uint8), 1), cv2.DIST_L2, cv2.DIST_MASK_PRECISE)[1:-1, 1:-1]
|
|
# Only centers strictly inside the original continuous ROI qualify.
|
|
eligible_x = ((np.arange(x0, x1) + .5) / width >= box[0]) & ((np.arange(x0, x1) + .5) / width <= box[2])
|
|
eligible_y = ((np.arange(y0, y1) + .5) / height >= box[1]) & ((np.arange(y0, y1) + .5) / height <= box[3])
|
|
eligible = roi & eligible_y[:, None] & eligible_x[None, :]
|
|
if eligible.any():
|
|
scores = np.where(eligible, distance, -1)
|
|
cy, cx = np.unravel_index(np.argmax(scores), scores.shape)
|
|
x, y = int(cx) + x0, int(cy) + y0
|
|
entry['candidate_points'] = [{
|
|
'point_px': [x, y], 'point_norm': [(x + .5) / width, (y + .5) / height],
|
|
'distance_to_roi_boundary_px': float(distance[cy, cx]),
|
|
'distance_to_silhouette_px': float(silhouette_distance[y, x]),
|
|
'source': 'alpha_roi_distance_transform',
|
|
}]
|
|
entry['status'] = 'candidate'
|
|
else:
|
|
entry['warnings'].append('NO_PIXEL_CENTER_IN_ROI')
|
|
if entry['bbox_foreground_ratio'] < .5:
|
|
entry['warnings'].append('LOW_FOREGROUND_RATIO')
|
|
if count > 2:
|
|
entry['warnings'].append('MULTIPLE_FOREGROUND_COMPONENTS')
|
|
entry['warnings'].append('SEMANTIC_REVIEW_REQUIRED')
|
|
diagnostics.append(entry)
|
|
prompts.append({
|
|
'id': part['id'], 'text_prompt': part['semantic'], 'status': entry['status'],
|
|
'bbox_px_xyxy': entry['bbox_px_xyxy'],
|
|
'positive_points_px': [p['point_px'] for p in entry['candidate_points']],
|
|
'positive_points_norm': [p['point_norm'] for p in entry['candidate_points']],
|
|
'negative_points_px': [], 'seed_source': 'alpha_roi_distance_transform',
|
|
'semantic_verified': False, 'requires_mask_qa': True,
|
|
})
|
|
points = [p for entry in diagnostics for p in entry['model_positive_points']]
|
|
return {
|
|
'schema_version': 1, 'width': width, 'height': height,
|
|
'alpha_threshold': 0, 'coordinate_mapping': 'floor(norm * size), clamped at size - 1; boxes are half-open',
|
|
'summary': {'model_points': len(points), 'invalid_model_points': sum(p['status'] == 'INVALID_SEED' for p in points),
|
|
'background_model_points': sum(not p['on_foreground'] for p in points),
|
|
'candidate_parts': sum(p['status'] == 'candidate' for p in diagnostics),
|
|
'blocked_parts': sum(p['status'] == 'blocked' for p in diagnostics)},
|
|
'parts': diagnostics,
|
|
}, {'schema_version': 1, 'width': width, 'height': height, 'quality_status': 'needs_review', 'parts': prompts}
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('--input', required=True)
|
|
parser.add_argument('--output', required=True)
|
|
parser.add_argument('--prepare', action='store_true')
|
|
parser.add_argument('--parts')
|
|
parser.add_argument('--prompts')
|
|
args = parser.parse_args()
|
|
with Image.open(args.input) as image:
|
|
if args.prepare:
|
|
prepared, metadata = prepare_rgba(image)
|
|
prepared.save(args.output, format='PNG')
|
|
print(json.dumps(metadata))
|
|
else:
|
|
if not args.parts or not args.prompts:
|
|
parser.error('--parts and --prompts are required for grounding')
|
|
plan = json.loads(Path(args.parts).read_text(encoding='utf-8'))
|
|
diagnostics, prompts = ground_parts(image, plan)
|
|
Path(args.output).write_text(json.dumps(diagnostics, indent=2), encoding='utf-8')
|
|
Path(args.prompts).write_text(json.dumps(prompts, indent=2), encoding='utf-8')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|