| """Project a reference onto visible UV texels; retain occluded and shared-backside texels.""" |
| import argparse |
| import json |
| from pathlib import Path |
| import sys |
|
|
| import bpy |
| import numpy as np |
| from mathutils import Vector |
| from mathutils.bvhtree import BVHTree |
|
|
| p = argparse.ArgumentParser() |
| p.add_argument('--folder', required=True) |
| p.add_argument('--front-axis', choices=['-Y', '+Y', '-X', '+X'], default='-Y') |
| p.add_argument('--strength', type=float, default=.9) |
| a = p.parse_args(sys.argv[sys.argv.index('--') + 1:]) |
| folder = Path(a.folder) |
| bpy.ops.object.select_all(action='SELECT') |
| bpy.ops.object.delete(use_global=False) |
| bpy.ops.import_scene.gltf(filepath=str(folder / 'source.glb')) |
| objects = [o for o in bpy.context.scene.objects if o.type == 'MESH'] |
| front = np.array({'-Y': [0, -1, 0], '+Y': [0, 1, 0], '-X': [-1, 0, 0], '+X': [1, 0, 0]}[a.front_axis], dtype=float) |
| right = np.cross(front, [0, 0, 1]) * -1 |
| vertices, triangles, records = [], [], [] |
| offset = 0 |
| for obj in objects: |
| mesh = obj.data |
| mesh.calc_loop_triangles() |
| world = np.array([obj.matrix_world @ v.co for v in mesh.vertices]) |
| normals = np.array([obj.matrix_world.to_3x3().inverted().transposed() @ v.normal for v in mesh.vertices]) |
| normals /= np.maximum(np.linalg.norm(normals, axis=1, keepdims=True), 1e-12) |
| vertices.extend(world.tolist()) |
| triangles.extend(tuple(offset + i for i in t.vertices) for t in mesh.loop_triangles) |
| records.append((obj, world, normals)) |
| offset += len(world) |
| vertices = np.array(vertices) |
| bvh = BVHTree.FromPolygons(vertices.tolist(), triangles, all_triangles=True) |
| height = np.ptp(vertices[:, 2]) |
| horizontal = vertices @ right |
| span = np.ptp(horizontal) |
| horizontal_min = horizontal.min() |
| vertical_min = vertices[:, 2].min() |
| if height <= 0 or span <= 0: |
| raise ValueError('Degenerate projection bounds') |
| reference = np.load(folder / 'reference.npz') |
| rgba, mask, bounds = reference['rgba'], reference['mask'], reference['bounds'] |
| ref_h, ref_w = mask.shape |
| x0, y0, x1, y1 = bounds |
| metrics = {'front_axis': a.front_axis, 'strength': a.strength, |
| 'method': 'per-texel UV projection with BVH occlusion, normal falloff and shared-UV protection', |
| 'material_color_boundaries_protected': True, 'geometry_unchanged': True, 'uv_unchanged': True, 'pbr_channels_preserved': True, 'parts': []} |
| inputs = json.loads((folder / 'inputs.json').read_text()) |
| for item in inputs: |
| if not item.get('selected', False): |
| np.save(folder / ('projected-' + item['file']), np.load(folder / item['file'])) |
| metrics['parts'].append({'part_id': item['part_id'], 'changed_texels': 0, 'skipped': 'outside explicit projection scope'}) |
| continue |
| matches = [r for r in records if r[0].name == item['part_id'] or r[0].name.startswith(item['part_id'] + ' ')] |
| if len(matches) != 1: |
| raise ValueError(f"Cannot identify unique component {item['part_id']}") |
| obj, world, normals = matches[0] |
| region = item.get('reference_box') |
| if region is not None: |
| region = np.asarray(region, dtype=float) |
| if (region.shape != (4,) or not np.isfinite(region).all() or np.any(region < 0) |
| or np.any(region > 1) or region[2] <= region[0] or region[3] <= region[1]): |
| raise ValueError('Invalid normalized reference rectangle') |
| px0, py0, px1, py1 = region * [ref_w - 1, ref_h - 1, ref_w - 1, ref_h - 1] |
| local_horizontal = world @ right |
| local_min, local_span = local_horizontal.min(), np.ptp(local_horizontal) |
| local_z, local_height = world[:, 2].min(), np.ptp(world[:, 2]) |
| if local_span <= 0 or local_height <= 0: |
| raise ValueError('Cannot project a reference rectangle onto degenerate component bounds') |
| else: |
| px0, py0, px1, py1 = x0, y0, x1, y1 |
| local_min, local_span = horizontal_min, span |
| local_z, local_height = vertical_min, height |
| source = np.load(folder / item['file']) |
| h, w = source.shape[:2] |
| weights = np.zeros((h, w), dtype=np.float32) |
| colors = np.zeros((h, w, 3), dtype=np.uint8) |
| blocked = np.zeros((h, w), dtype=bool) |
| coverage = np.zeros((h, w), dtype=bool) |
| replacement_weights = np.zeros((h, w), dtype=np.float32) |
| reference_coordinates = np.zeros((h, w, 2), dtype=np.float32) |
| uv_data = obj.data.uv_layers.active.data |
| for triangle in obj.data.loop_triangles: |
| uv = np.array([uv_data[i].uv for i in triangle.loops]) |
| if np.any(uv < -1e-6) or np.any(uv > 1 + 1e-6): |
| raise ValueError('Tiled UVs are unsupported for reference projection; source retained.') |
| pixel = uv * [w, -h] + [0, h] |
| low = np.maximum(np.floor(pixel.min(0)).astype(int), 0) |
| high = np.minimum(np.ceil(pixel.max(0)).astype(int), [w - 1, h - 1]) |
| if np.any(high < low): |
| continue |
| yy, xx = np.mgrid[low[1]:high[1] + 1, low[0]:high[0] + 1] |
| points = np.stack([xx.ravel() + .5, yy.ravel() + .5], axis=1) |
| v0, v1 = pixel[1] - pixel[0], pixel[2] - pixel[0] |
| determinant = v0[0] * v1[1] - v0[1] * v1[0] |
| if abs(determinant) < 1e-10: |
| continue |
| delta = points - pixel[0] |
| b1 = (delta[:, 0] * v1[1] - delta[:, 1] * v1[0]) / determinant |
| b2 = (v0[0] * delta[:, 1] - v0[1] * delta[:, 0]) / determinant |
| bary = np.stack([1 - b1 - b2, b1, b2], axis=1) |
| inside = np.all(bary >= -1e-7, axis=1) |
| if not inside.any(): |
| continue |
| bary = bary[inside] |
| xx, yy = xx.ravel()[inside], yy.ravel()[inside] |
| coverage[yy, xx] = True |
| ids = list(triangle.vertices) |
| xyz = bary @ world[ids] |
| n = bary @ normals[ids] |
| facing = np.clip((n @ front - .35) / .55, 0, 1) |
| facing = facing * facing * (3 - 2 * facing) |
| rx = np.clip(np.rint(px0 + (xyz @ right - local_min) / local_span * (px1 - px0)).astype(int), 0, ref_w - 1) |
| ry = np.clip(np.rint(py1 - (xyz[:, 2] - local_z) / local_height * (py1 - py0)).astype(int), 0, ref_h - 1) |
| allowance = np.zeros(len(xyz)) |
| replacement = item.get('replace_colors_box') |
| if replacement is not None: |
| left, top, right_edge, bottom = np.asarray(replacement) * [ref_w-1, ref_h-1, ref_w-1, ref_h-1] |
| edge_distance = np.minimum.reduce([rx-left, right_edge-rx, ry-top, bottom-ry]) |
| allowance = np.clip(edge_distance / 2, 0, 1) |
| allowance = allowance * allowance * (3 - 2 * allowance) |
| facing = np.maximum(facing, allowance) |
| roi_weight = np.ones(len(xyz)) |
| if item.get('projection_box') is not None: |
| left, top, right_edge, bottom = np.asarray(item['projection_box']) * [ref_w-1, ref_h-1, ref_w-1, ref_h-1] |
| edge_distance = np.minimum.reduce([rx-left, right_edge-rx, ry-top, bottom-ry]) |
| roi_weight = np.clip(edge_distance / item.get('projection_feather_pixels', 2), 0, 1) |
| roi_weight = roi_weight * roi_weight * (3 - 2 * roi_weight) |
| eligible = (facing > 0) & (roi_weight > 0) & mask[ry, rx] |
| visible = np.zeros(len(xyz), dtype=bool) |
| for index in np.flatnonzero(eligible): |
| point = xyz[index] |
| hit, _, _, _ = bvh.ray_cast(Vector(point + front * height * 2), Vector(-front)) |
| visible[index] = hit is not None and np.linalg.norm(np.array(hit) - point) < height * 1e-4 |
| original_rgb = source[yy, xx, :3].astype(float) / 255 |
| target_rgb = rgba[ry, rx, :3].astype(float) / 255 |
| original_sat = np.ptp(original_rgb, axis=1) / np.maximum(original_rgb.max(1), 1e-6) |
| target_sat = np.ptp(target_rgb, axis=1) / np.maximum(target_rgb.max(1), 1e-6) |
| colored = np.clip((original_sat - .35) / .15, 0, 1) |
| washed_out = np.clip((.45 - target_sat) / .15, 0, 1) |
| color_agreement = 1 - colored * washed_out |
| source_chroma = original_rgb - original_rgb.mean(1, keepdims=True) |
| target_chroma = target_rgb - target_rgb.mean(1, keepdims=True) |
| cosine = np.sum(source_chroma * target_chroma, axis=1) / np.maximum( |
| np.linalg.norm(source_chroma, axis=1) * np.linalg.norm(target_chroma, axis=1), 1e-6) |
| changed_hue = colored * np.clip((target_sat - .35) / .15, 0, 1) |
| color_agreement *= 1 - changed_hue * (1 - np.clip((cosine - .3) / .5, 0, 1)) |
| color_agreement = np.maximum(color_agreement, allowance) |
| weight = facing * visible * color_agreement * a.strength * roi_weight |
| blocked[yy[weight == 0], xx[weight == 0]] = True |
| keep = weight > weights[yy, xx] |
| weights[yy[keep], xx[keep]] = weight[keep] |
| replacement_weights[yy[keep], xx[keep]] = allowance[keep] |
| colors[yy[keep], xx[keep]] = rgba[ry[keep], rx[keep], :3] |
| reference_coordinates[yy[keep], xx[keep]] = np.stack([rx[keep] / ref_w, ry[keep] / ref_h], axis=1) |
| weights[blocked] = 0 |
| harmonization = np.zeros(3) |
| selection = (replacement_weights > .9) & (weights > .8) |
| original_rgb = source[selection, :3].astype(float) |
| target_rgb = colors[selection].astype(float) |
| original_chroma = original_rgb - original_rgb.mean(1, keepdims=True) |
| target_chroma = target_rgb - target_rgb.mean(1, keepdims=True) |
| cosine = np.sum(original_chroma * target_chroma, axis=1) / np.maximum( |
| np.linalg.norm(original_chroma, axis=1) * np.linalg.norm(target_chroma, axis=1), 1e-6) |
| original_sat = np.ptp(original_rgb, axis=1) / np.maximum(original_rgb.max(1), 1) |
| target_sat = np.ptp(target_rgb, axis=1) / np.maximum(target_rgb.max(1), 1) |
| samples = ((cosine > .9) & (original_sat > .15) & (original_sat < .65) |
| & (target_sat > .15) & (target_sat < .65) |
| & (original_rgb.mean(1) > 90) & (target_rgb.mean(1) > 90)) |
| spatial_fit = None |
| if samples.sum() >= 64: |
| harmonization = np.clip(np.median(original_rgb[samples] - target_rgb[samples], axis=0), -24, 24) |
| if item.get('color_matching') == 'spatial': |
| coords = reference_coordinates[selection][samples] |
| center = coords.mean(axis=0) |
| scale = np.maximum(np.ptp(coords, axis=0), 1e-4) |
| design = np.column_stack([np.ones(len(coords)), (coords-center)/scale]) |
| residual = original_rgb[samples] - target_rgb[samples] |
| fit = np.linalg.lstsq(design, residual, rcond=None)[0] |
| errors = np.linalg.norm(design @ fit - residual, axis=1) |
| inliers = errors <= max(8, np.percentile(errors, 75)) |
| if inliers.sum() >= 64 and np.linalg.cond(design[inliers]) < 100: |
| fit = np.linalg.lstsq(design[inliers], residual[inliers], rcond=None)[0] |
| spatial_fit = {'center': center.tolist(), 'scale': scale.tolist(), 'coefficients': fit.tolist(), 'samples': int(inliers.sum())} |
| localized = (replacement_weights > 0) & (weights > 0) |
| target_rgb = colors[localized].astype(float) |
| detail_preservation = np.clip((target_rgb.mean(1) - 40) / 80, 0, 1) |
| correction = harmonization |
| if spatial_fit is not None: |
| coords = reference_coordinates[localized] |
| design = np.column_stack([np.ones(len(coords)), (coords-center)/scale]) |
| correction = np.clip(design @ fit, -48, 48) |
| colors[localized] = np.clip(target_rgb + replacement_weights[localized, None] |
| * detail_preservation[:, None] * correction, 0, 255).astype(np.uint8) |
| padding = 2 if item.get('replace_colors_box') is not None and item.get('projection_box') is None else 0 |
| padded_weights, padded_colors = weights.copy(), colors.copy() |
| for dy in range(-padding, padding + 1): |
| for dx in range(-padding, padding + 1): |
| sy, sx = slice(max(0, -dy), min(h, h-dy)), slice(max(0, -dx), min(w, w-dx)) |
| ty, tx = slice(max(0, dy), min(h, h+dy)), slice(max(0, dx), min(w, w+dx)) |
| keep = (~coverage[ty, tx]) & (weights[sy, sx] > padded_weights[ty, tx]) |
| padded_weights[ty, tx][keep] = weights[sy, sx][keep] |
| padded_colors[ty, tx][keep] = colors[sy, sx][keep] |
| weights, colors = padded_weights, padded_colors |
| changed = weights > 0 |
| result = source.copy() |
| result[:, :, :3] = np.rint(source[:, :, :3] * (1 - weights[:, :, None]) + colors * weights[:, :, None]).astype(np.uint8) |
| np.save(folder / ('projected-' + item['file']), result) |
| assert np.array_equal(source[~changed], result[~changed]) |
| assert np.array_equal(source[:, :, 3], result[:, :, 3]) |
| metrics['parts'].append({'part_id': item['part_id'], 'changed_texels': int(changed.sum()), |
| 'total_texels': int(h * w), 'protected_texels_unchanged': True, |
| 'reference_box': item.get('reference_box'), |
| 'projection_box': item.get('projection_box'), |
| 'projection_feather_pixels': item.get('projection_feather_pixels', 2), |
| 'spatial_color_fit': spatial_fit, |
| 'replace_colors_box': item.get('replace_colors_box'), |
| 'unused_texel_padding': padding, |
| 'local_color_offset': harmonization.tolist(), |
| 'color_match_samples': int(samples.sum()), |
| 'alpha_unchanged': True}) |
| (folder / 'projection.json').write_text(json.dumps(metrics, indent=2)) |
| print(json.dumps(metrics), flush=True) |
|
|
