scripts/project_reference_texture.py

"""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)