892 lines
31 KiB
Python
892 lines
31 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
fill_all_empty_faces_v3_square_harmonic_atlas.py
|
||
|
||
目标：替换原来的“一个缺色孤岛 -> 一个均值色块”的逻辑，
|
||
并保证输出贴图严格保持 1:1 方形。
|
||
|
||
核心逻辑：
|
||
1. 正常 face 保持原图像素位置不变；如果输出画布变大，则同时 remap u/v。
|
||
不能只 remap v，否则输出方形后旧 UV 会整体错位。
|
||
2. missing face 不再按 region 写一个颜色点，也不再整组 face 指向同一个 vt。
|
||
3. 在 3D mesh 拓扑上，对 missing region 的“顶点颜色”做 harmonic/Laplace 扩散：
|
||
- 边界顶点颜色来自相邻正常 face 的原始 texture 采样；
|
||
- 内部顶点颜色由图拉普拉斯扩散得到；
|
||
- 所以颜色从边缘向内部连续变化。
|
||
4. 每个 missing face 在新增 atlas 中分配一个小三角 tile。
|
||
- 输出画布是 square_size x square_size；原贴图放在左上角；
|
||
- atlas 放在原贴图下方的安全区域，不塞进原贴图右侧，避免采样串色；
|
||
- tile 内不是纯色，而是按三个顶点颜色做 barycentric 插值；
|
||
- 相邻 missing face 即使在 atlas 里不相邻，只要共享 3D 顶点，边上的插值颜色也连续；
|
||
- tile 周围做 padding，避免纹理过滤时采到其他 tile。
|
||
|
||
输入输出参数保持原脚本一致：
|
||
--input_obj
|
||
--input_texture
|
||
--missing_faces
|
||
--output_obj
|
||
--output_texture
|
||
|
||
额外参数都是可选。
|
||
"""
|
||
|
||
import argparse
|
||
import os
|
||
import time
|
||
from collections import defaultdict, deque
|
||
from typing import Dict, List, Tuple, Set
|
||
|
||
import cv2
|
||
import numpy as np
|
||
import tqdm
|
||
|
||
|
||
# ============================================================
|
||
# OBJ / texture IO
|
||
# ============================================================
|
||
|
||
def _parse_obj_index(s: str, current_len: int) -> int:
|
||
"""OBJ index: positive is 1-based, negative is relative to current list end."""
|
||
idx = int(s)
|
||
if idx > 0:
|
||
return idx - 1
|
||
if idx < 0:
|
||
return current_len + idx
|
||
raise ValueError("OBJ index 不能为 0")
|
||
|
||
|
||
def read_obj_basic(obj_path: str):
|
||
vertices: List[List[float]] = []
|
||
uvs: List[List[float]] = []
|
||
vertex_indices: List[List[int]] = []
|
||
uv_indices: List[List[int]] = []
|
||
|
||
mtllib_lines: List[str] = []
|
||
usemtl_name = "material_0"
|
||
|
||
with open(obj_path, "r", encoding="utf-8", errors="ignore") as f:
|
||
for line in f:
|
||
if line.startswith("mtllib "):
|
||
mtllib_lines.append(line.strip())
|
||
elif line.startswith("usemtl "):
|
||
parts = line.strip().split(maxsplit=1)
|
||
if len(parts) == 2:
|
||
usemtl_name = parts[1]
|
||
elif line.startswith("v "):
|
||
parts = line.strip().split()
|
||
if len(parts) < 4:
|
||
continue
|
||
vertices.append([float(parts[1]), float(parts[2]), float(parts[3])])
|
||
elif line.startswith("vt "):
|
||
parts = line.strip().split()
|
||
if len(parts) < 3:
|
||
continue
|
||
uvs.append([float(parts[1]), float(parts[2])])
|
||
elif line.startswith("f "):
|
||
parts = line.strip().split()[1:]
|
||
if len(parts) != 3:
|
||
raise ValueError(f"当前脚本只支持三角面，发现非三角 face: {line.strip()}")
|
||
|
||
v_row = []
|
||
vt_row = []
|
||
for token in parts:
|
||
sub = token.split("/")
|
||
if len(sub) < 2 or sub[1] == "":
|
||
raise ValueError(f"face 缺少 vt，无法处理: {line.strip()}")
|
||
v_row.append(_parse_obj_index(sub[0], len(vertices)))
|
||
vt_row.append(_parse_obj_index(sub[1], len(uvs)))
|
||
|
||
vertex_indices.append(v_row)
|
||
uv_indices.append(vt_row)
|
||
|
||
vertices_np = np.asarray(vertices, dtype=np.float32)
|
||
uvs_np = np.asarray(uvs, dtype=np.float32)
|
||
vertex_indices_np = np.asarray(vertex_indices, dtype=np.int64)
|
||
uv_indices_np = np.asarray(uv_indices, dtype=np.int64)
|
||
|
||
if len(vertices_np) == 0 or len(uvs_np) == 0 or len(vertex_indices_np) == 0:
|
||
raise ValueError("OBJ 内没有读到完整的 v / vt / f 数据。")
|
||
|
||
return vertices_np, uvs_np, vertex_indices_np, uv_indices_np, mtllib_lines, usemtl_name
|
||
|
||
|
||
def read_missing_faces(path: str, index_base: int = 0) -> np.ndarray:
|
||
arr = []
|
||
with open(path, "r", encoding="utf-8", errors="ignore") as f:
|
||
for line in f:
|
||
s = line.strip()
|
||
if not s:
|
||
continue
|
||
arr.append(int(s) - index_base)
|
||
return np.asarray(arr, dtype=np.int64)
|
||
|
||
|
||
def read_texture_rgb(path: str) -> np.ndarray:
|
||
bgr = cv2.imread(path, cv2.IMREAD_COLOR)
|
||
if bgr is None:
|
||
raise FileNotFoundError(f"无法读取贴图: {path}")
|
||
return cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
|
||
|
||
|
||
def write_texture_rgb(path: str, rgb: np.ndarray):
|
||
os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
|
||
bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
|
||
ok = cv2.imwrite(path, bgr, [cv2.IMWRITE_PNG_COMPRESSION, 3])
|
||
if not ok:
|
||
raise IOError(f"写入贴图失败: {path}")
|
||
|
||
|
||
def write_obj_with_uv_coordinates(
|
||
filename: str,
|
||
vertices: np.ndarray,
|
||
uvs: np.ndarray,
|
||
vertex_indices: np.ndarray,
|
||
uv_indices: np.ndarray,
|
||
mtllib_lines: List[str] = None,
|
||
usemtl_name: str = "material_0",
|
||
):
|
||
os.makedirs(os.path.dirname(os.path.abspath(filename)), exist_ok=True)
|
||
|
||
if not mtllib_lines:
|
||
mtllib_lines = ["mtllib mesh.mtl"]
|
||
|
||
estimated_size = len(vertices) * 48 + len(uvs) * 32 + len(vertex_indices) * 64
|
||
buffer_size = min(max(int(estimated_size * 1.2), 64 * 1024 * 1024), 1024 * 1024 * 1024)
|
||
|
||
lines: List[str] = []
|
||
lines.extend(mtllib_lines)
|
||
|
||
for v in vertices:
|
||
lines.append("v %.6f %.6f %.6f" % (v[0], v[1], v[2]))
|
||
|
||
lines.append("")
|
||
for uv in uvs:
|
||
lines.append("vt %.8f %.8f" % (uv[0], uv[1]))
|
||
|
||
lines.append("")
|
||
lines.append(f"usemtl {usemtl_name}")
|
||
|
||
for v_idx, vt_idx in zip(vertex_indices, uv_indices):
|
||
lines.append(
|
||
"f %d/%d %d/%d %d/%d" % (
|
||
int(v_idx[0]) + 1, int(vt_idx[0]) + 1,
|
||
int(v_idx[1]) + 1, int(vt_idx[1]) + 1,
|
||
int(v_idx[2]) + 1, int(vt_idx[2]) + 1,
|
||
)
|
||
)
|
||
|
||
with open(filename, "w", buffering=buffer_size, encoding="utf-8") as f:
|
||
f.write("\n".join(lines))
|
||
|
||
|
||
# ============================================================
|
||
# Mesh topology
|
||
# ============================================================
|
||
|
||
def build_face_adjacency_and_vertex_faces(faces: np.ndarray, num_vertices: int):
|
||
"""
|
||
基于共享边构建 face adjacency，同时构建 vertex -> faces。
|
||
返回：face_adjacency, vertex_faces
|
||
"""
|
||
faces = np.asarray(faces, dtype=np.int64)
|
||
num_faces = len(faces)
|
||
|
||
# vertex -> faces
|
||
vertex_faces: List[List[int]] = [[] for _ in range(num_vertices)]
|
||
for fi, row in enumerate(faces):
|
||
vertex_faces[int(row[0])].append(fi)
|
||
vertex_faces[int(row[1])].append(fi)
|
||
vertex_faces[int(row[2])].append(fi)
|
||
|
||
# edge -> faces，向量化排序
|
||
edges = np.stack([
|
||
np.column_stack((faces[:, 0], faces[:, 1])),
|
||
np.column_stack((faces[:, 1], faces[:, 2])),
|
||
np.column_stack((faces[:, 2], faces[:, 0])),
|
||
], axis=1).reshape(-1, 2)
|
||
edges.sort(axis=1)
|
||
edge_faces = np.repeat(np.arange(num_faces, dtype=np.int64), 3)
|
||
|
||
key = edges[:, 0] * num_vertices + edges[:, 1]
|
||
order = np.argsort(key)
|
||
key = key[order]
|
||
edge_faces = edge_faces[order]
|
||
|
||
adj = defaultdict(list)
|
||
same = key[1:] == key[:-1]
|
||
pos = np.where(same)[0]
|
||
for p in pos:
|
||
f0 = int(edge_faces[p])
|
||
f1 = int(edge_faces[p + 1])
|
||
if f0 != f1:
|
||
adj[f0].append(f1)
|
||
adj[f1].append(f0)
|
||
|
||
return dict(adj), vertex_faces
|
||
|
||
|
||
def find_missing_regions(face_adjacency: Dict[int, List[int]], missing_faces: np.ndarray, total_faces: int):
|
||
missing_all = set(int(x) for x in missing_faces if 0 <= int(x) < total_faces)
|
||
unused = set(missing_all)
|
||
regions: List[Tuple[Set[int], Set[int]]] = []
|
||
|
||
pbar = tqdm.tqdm(total=len(unused), desc="Finding missing islands")
|
||
done = 0
|
||
|
||
while unused:
|
||
start = unused.pop()
|
||
group = {start}
|
||
boundary = set()
|
||
stack = [start]
|
||
|
||
while stack:
|
||
f = stack.pop()
|
||
for nb in face_adjacency.get(f, []):
|
||
if nb in unused:
|
||
unused.remove(nb)
|
||
group.add(nb)
|
||
stack.append(nb)
|
||
elif nb not in missing_all:
|
||
boundary.add(nb)
|
||
|
||
regions.append((group, boundary))
|
||
new_done = len(missing_all) - len(unused)
|
||
pbar.update(new_done - done)
|
||
done = new_done
|
||
|
||
pbar.close()
|
||
|
||
if regions:
|
||
sizes = [len(x[0]) for x in regions]
|
||
print(f"Total missing islands: {len(regions)}")
|
||
print(f"Island size: min={min(sizes)}, max={max(sizes)}, avg={sum(sizes)/len(sizes):.2f}")
|
||
else:
|
||
print("No missing islands found.")
|
||
|
||
return regions
|
||
|
||
|
||
# ============================================================
|
||
# Texture sampling
|
||
# ============================================================
|
||
|
||
def sample_texture_bilinear_rgb(texture_rgb: np.ndarray, uv: np.ndarray) -> np.ndarray:
|
||
"""uv: [2], v is OBJ convention, image y is flipped."""
|
||
H, W = texture_rgb.shape[:2]
|
||
u = float(np.clip(uv[0], 0.0, 1.0))
|
||
v = float(np.clip(uv[1], 0.0, 1.0))
|
||
|
||
x = u * (W - 1)
|
||
y = (1.0 - v) * (H - 1)
|
||
|
||
x0 = int(np.floor(x))
|
||
y0 = int(np.floor(y))
|
||
x1 = min(x0 + 1, W - 1)
|
||
y1 = min(y0 + 1, H - 1)
|
||
|
||
dx = x - x0
|
||
dy = y - y0
|
||
|
||
c00 = texture_rgb[y0, x0].astype(np.float32)
|
||
c10 = texture_rgb[y0, x1].astype(np.float32)
|
||
c01 = texture_rgb[y1, x0].astype(np.float32)
|
||
c11 = texture_rgb[y1, x1].astype(np.float32)
|
||
|
||
c0 = c00 * (1.0 - dx) + c10 * dx
|
||
c1 = c01 * (1.0 - dx) + c11 * dx
|
||
return c0 * (1.0 - dy) + c1 * dy
|
||
|
||
|
||
def sample_face_corner_color(
|
||
face_id: int,
|
||
corner: int,
|
||
texture_rgb: np.ndarray,
|
||
uvs: np.ndarray,
|
||
face_uv_indices: np.ndarray,
|
||
corner_inset: float = 0.90,
|
||
) -> np.ndarray:
|
||
"""
|
||
在正常 face 的某个顶点附近采样。
|
||
不直接采顶点本身，而是往三角形内部缩一点，避免 seam 边界采样不稳。
|
||
corner_inset 越接近 1，越贴近该顶点。
|
||
"""
|
||
tri_uv = uvs[face_uv_indices[face_id]] # [3,2]
|
||
w = np.full(3, (1.0 - corner_inset) / 2.0, dtype=np.float32)
|
||
w[corner] = corner_inset
|
||
uv = w @ tri_uv
|
||
return sample_texture_bilinear_rgb(texture_rgb, uv)
|
||
|
||
|
||
def median_color(colors: List[np.ndarray]) -> np.ndarray:
|
||
if len(colors) == 0:
|
||
return np.asarray([128, 128, 128], dtype=np.float32)
|
||
return np.median(np.stack(colors, axis=0).astype(np.float32), axis=0)
|
||
|
||
|
||
# ============================================================
|
||
# Harmonic vertex color diffusion
|
||
# ============================================================
|
||
|
||
def build_region_vertex_graph(region_faces: Set[int], faces: np.ndarray):
|
||
"""
|
||
对一个 missing island 构建 local vertex graph。
|
||
返回：unique_vertices, vertex_to_local, unique_edges, local_face_ids
|
||
"""
|
||
face_ids = np.asarray(sorted(region_faces), dtype=np.int64)
|
||
tri_v = faces[face_ids] # [F,3]
|
||
unique_vertices = np.unique(tri_v.reshape(-1))
|
||
v2local = {int(v): i for i, v in enumerate(unique_vertices)}
|
||
|
||
local_faces = np.vectorize(lambda x: v2local[int(x)], otypes=[np.int64])(tri_v)
|
||
|
||
edges = np.concatenate([
|
||
local_faces[:, [0, 1]],
|
||
local_faces[:, [1, 2]],
|
||
local_faces[:, [2, 0]],
|
||
], axis=0)
|
||
edges.sort(axis=1)
|
||
edges = np.unique(edges, axis=0)
|
||
|
||
return face_ids, unique_vertices, v2local, edges, local_faces
|
||
|
||
|
||
def compute_boundary_vertex_constraints(
|
||
unique_vertices: np.ndarray,
|
||
vertex_faces: List[List[int]],
|
||
faces: np.ndarray,
|
||
face_uv_indices: np.ndarray,
|
||
uvs: np.ndarray,
|
||
texture_rgb: np.ndarray,
|
||
missing_mask: np.ndarray,
|
||
corner_inset: float = 0.90,
|
||
):
|
||
"""
|
||
对 region 内的每个顶点，如果它邻接任何非 missing face，就从这些正常 face 的对应 corner 采颜色。
|
||
这比“对边界 face 求一个均值”更合理，因为边界不同位置的颜色会保留下来。
|
||
"""
|
||
n = len(unique_vertices)
|
||
constrained = np.zeros(n, dtype=bool)
|
||
constraint_colors = np.zeros((n, 3), dtype=np.float32)
|
||
|
||
for local_i, gv in enumerate(unique_vertices):
|
||
samples = []
|
||
for f in vertex_faces[int(gv)]:
|
||
if missing_mask[f]:
|
||
continue
|
||
row = faces[f]
|
||
corners = np.where(row == gv)[0]
|
||
if len(corners) == 0:
|
||
continue
|
||
for c in corners:
|
||
samples.append(sample_face_corner_color(
|
||
f,
|
||
int(c),
|
||
texture_rgb,
|
||
uvs,
|
||
face_uv_indices,
|
||
corner_inset=corner_inset,
|
||
))
|
||
|
||
if samples:
|
||
constrained[local_i] = True
|
||
constraint_colors[local_i] = median_color(samples)
|
||
|
||
return constrained, constraint_colors
|
||
|
||
|
||
def initialize_vertex_colors_by_bfs(edges: np.ndarray, constrained: np.ndarray, constraint_colors: np.ndarray) -> np.ndarray:
|
||
n = len(constrained)
|
||
colors = np.zeros((n, 3), dtype=np.float32)
|
||
|
||
if constrained.any():
|
||
colors[constrained] = constraint_colors[constrained]
|
||
fallback = np.mean(constraint_colors[constrained], axis=0)
|
||
else:
|
||
fallback = np.asarray([128, 128, 128], dtype=np.float32)
|
||
colors[:] = fallback
|
||
return colors
|
||
|
||
adj = [[] for _ in range(n)]
|
||
for a, b in edges:
|
||
adj[int(a)].append(int(b))
|
||
adj[int(b)].append(int(a))
|
||
|
||
q = deque(np.where(constrained)[0].tolist())
|
||
visited = np.zeros(n, dtype=bool)
|
||
visited[constrained] = True
|
||
|
||
while q:
|
||
cur = q.popleft()
|
||
for nb in adj[cur]:
|
||
if visited[nb]:
|
||
continue
|
||
colors[nb] = colors[cur]
|
||
visited[nb] = True
|
||
q.append(nb)
|
||
|
||
colors[~visited] = fallback
|
||
return colors
|
||
|
||
|
||
def solve_harmonic_vertex_colors(
|
||
edges: np.ndarray,
|
||
constrained: np.ndarray,
|
||
constraint_colors: np.ndarray,
|
||
smooth_iters: int = 80,
|
||
self_weight: float = 0.0,
|
||
) -> np.ndarray:
|
||
"""
|
||
对 region 顶点颜色做 Dirichlet boundary harmonic extension。
|
||
constrained 顶点保持边界颜色；非 constrained 顶点反复取邻居平均。
|
||
"""
|
||
n = len(constrained)
|
||
colors = initialize_vertex_colors_by_bfs(edges, constrained, constraint_colors)
|
||
|
||
if n == 0 or len(edges) == 0:
|
||
return colors
|
||
|
||
e0 = edges[:, 0].astype(np.int64)
|
||
e1 = edges[:, 1].astype(np.int64)
|
||
|
||
deg = np.zeros(n, dtype=np.float32)
|
||
np.add.at(deg, e0, 1.0)
|
||
np.add.at(deg, e1, 1.0)
|
||
denom = deg[:, None] + float(self_weight)
|
||
denom = np.maximum(denom, 1e-8)
|
||
|
||
for _ in range(max(0, int(smooth_iters))):
|
||
acc = np.zeros_like(colors)
|
||
np.add.at(acc, e0, colors[e1])
|
||
np.add.at(acc, e1, colors[e0])
|
||
|
||
if self_weight > 0:
|
||
acc += colors * float(self_weight)
|
||
|
||
new_colors = acc / denom
|
||
# Dirichlet boundary：边界顶点锁住，不允许被内部均化掉。
|
||
new_colors[constrained] = constraint_colors[constrained]
|
||
colors = new_colors
|
||
|
||
return np.clip(colors, 0, 255).astype(np.float32)
|
||
|
||
|
||
# ============================================================
|
||
# Per-face triangular atlas baking
|
||
# ============================================================
|
||
|
||
def precompute_triangle_tile(tile_size: int, pad: int):
|
||
"""
|
||
预计算一个标准三角 tile 的 barycentric weights 和 padding 最近邻映射。
|
||
后续每个 missing face 只需要 weights @ 三个顶点颜色。
|
||
"""
|
||
tile_size = int(tile_size)
|
||
pad = int(pad)
|
||
if tile_size < 4:
|
||
raise ValueError("tile_size 至少要 >= 4")
|
||
if pad < 1:
|
||
pad = 1
|
||
if pad * 2 + 2 >= tile_size:
|
||
pad = max(1, tile_size // 4)
|
||
|
||
# 标准右三角，三个 UV 角点都留 padding，避免双线性采样采到 tile 外。
|
||
p0 = np.asarray([pad, pad], dtype=np.float32)
|
||
p1 = np.asarray([tile_size - pad - 1, pad], dtype=np.float32)
|
||
p2 = np.asarray([pad, tile_size - pad - 1], dtype=np.float32)
|
||
|
||
yy, xx = np.meshgrid(np.arange(tile_size, dtype=np.float32),
|
||
np.arange(tile_size, dtype=np.float32), indexing="ij")
|
||
pts = np.stack([xx + 0.5, yy + 0.5], axis=-1) # [T,T,2], pixel centers
|
||
|
||
# barycentric
|
||
v0 = p1 - p0
|
||
v1 = p2 - p0
|
||
v2 = pts - p0
|
||
d00 = float(np.dot(v0, v0))
|
||
d01 = float(np.dot(v0, v1))
|
||
d11 = float(np.dot(v1, v1))
|
||
denom = d00 * d11 - d01 * d01
|
||
if abs(denom) < 1e-8:
|
||
raise ValueError("退化 tile triangle")
|
||
|
||
d20 = v2[..., 0] * v0[0] + v2[..., 1] * v0[1]
|
||
d21 = v2[..., 0] * v1[0] + v2[..., 1] * v1[1]
|
||
w1 = (d11 * d20 - d01 * d21) / denom
|
||
w2 = (d00 * d21 - d01 * d20) / denom
|
||
w0 = 1.0 - w1 - w2
|
||
weights = np.stack([w0, w1, w2], axis=-1).astype(np.float32)
|
||
|
||
inside = (weights[..., 0] >= -1e-4) & (weights[..., 1] >= -1e-4) & (weights[..., 2] >= -1e-4)
|
||
|
||
# padding：tile 外部/三角外部像素复制最近的三角内部像素颜色。
|
||
inside_coords = np.argwhere(inside)
|
||
nearest_y = np.zeros((tile_size, tile_size), dtype=np.int32)
|
||
nearest_x = np.zeros((tile_size, tile_size), dtype=np.int32)
|
||
for y in range(tile_size):
|
||
for x in range(tile_size):
|
||
if inside[y, x]:
|
||
nearest_y[y, x] = y
|
||
nearest_x[y, x] = x
|
||
else:
|
||
d2 = (inside_coords[:, 0] - y) ** 2 + (inside_coords[:, 1] - x) ** 2
|
||
k = int(np.argmin(d2))
|
||
nearest_y[y, x] = int(inside_coords[k, 0])
|
||
nearest_x[y, x] = int(inside_coords[k, 1])
|
||
|
||
padded_weights = weights[nearest_y, nearest_x]
|
||
|
||
# UV 角点使用像素中心，更稳定。
|
||
uv_points = np.stack([p0, p1, p2], axis=0).astype(np.float32)
|
||
|
||
return padded_weights, uv_points
|
||
|
||
|
||
def compute_square_atlas_layout(
|
||
width: int,
|
||
height: int,
|
||
n_tiles: int,
|
||
tile_size: int,
|
||
):
|
||
"""
|
||
计算一个严格方形输出画布，并把新增 atlas 只放到原贴图下方。
|
||
|
||
为什么不用“右侧空白 + 下方空白”的 L 形区域：
|
||
- 原始 UV 通常会贴到原贴图右边界；
|
||
- 如果右侧直接塞 atlas，线性采样 / mipmap 很容易把 atlas 颜色串到旧 UV 边缘；
|
||
- 所以右侧只做 edge padding，下方才放新增 tile。
|
||
"""
|
||
width = int(width)
|
||
height = int(height)
|
||
n_tiles = int(n_tiles)
|
||
tile_size = int(tile_size)
|
||
|
||
if tile_size <= 0:
|
||
raise ValueError("tile_size 必须为正数")
|
||
|
||
side = max(width, height)
|
||
if n_tiles <= 0:
|
||
return side, 0, 0, 0
|
||
|
||
# 至少要有一行 atlas。原贴图本身如果已经是方形，需要把方形边长一起扩大。
|
||
side = max(side, height + tile_size)
|
||
|
||
# 简单稳定地增长边长。步长用 tile_size，避免大量无效循环。
|
||
while True:
|
||
cols = side // tile_size
|
||
rows = (side - height) // tile_size
|
||
capacity = cols * rows
|
||
if cols > 0 and rows > 0 and capacity >= n_tiles:
|
||
return side, int(cols), int(rows), int(capacity)
|
||
side += tile_size
|
||
|
||
|
||
def make_square_canvas_with_edge_padding(texture_rgb: np.ndarray, side: int) -> np.ndarray:
|
||
"""
|
||
把原贴图放在方形画布左上角，并用边缘像素扩展空白区域。
|
||
|
||
这里不要用纯白初始化：
|
||
- 旧 UV remap 后仍然贴在原图矩形内；
|
||
- 但边界处的双线性采样 / mipmap 可能会读到矩形外侧；
|
||
- 用边缘扩展比白底安全，不容易出现白边。
|
||
"""
|
||
H, W = texture_rgb.shape[:2]
|
||
side = int(side)
|
||
if side < max(H, W):
|
||
raise ValueError(f"square side 太小: side={side}, texture={W}x{H}")
|
||
|
||
canvas = np.empty((side, side, 3), dtype=np.uint8)
|
||
canvas[:H, :W, :] = texture_rgb
|
||
|
||
# 右侧扩展最后一列，保护原 UV 的右边界。
|
||
if side > W:
|
||
canvas[:H, W:side, :] = texture_rgb[:, W - 1:W, :]
|
||
|
||
# 下方扩展最后一行。注意这里复制的是已经完成右侧扩展后的整行。
|
||
if side > H:
|
||
canvas[H:side, :, :] = canvas[H - 1:H, :, :]
|
||
|
||
return canvas
|
||
|
||
|
||
def remap_original_uvs_to_square_canvas(
|
||
original_uvs: np.ndarray,
|
||
old_width: int,
|
||
old_height: int,
|
||
square_side: int,
|
||
) -> np.ndarray:
|
||
"""
|
||
保持原贴图像素坐标不变，只把归一化 UV 换算到新的方形画布。
|
||
|
||
原脚本只扩高，所以只改 v；现在输出变成 square_side x square_side，
|
||
必须同时改 u 和 v，否则原始 texture 会被横向拉错。
|
||
"""
|
||
remapped = original_uvs.copy().astype(np.float32)
|
||
|
||
old_x = remapped[:, 0] * max(old_width - 1, 1)
|
||
old_y = (1.0 - remapped[:, 1]) * max(old_height - 1, 1)
|
||
|
||
remapped[:, 0] = old_x / max(square_side - 1, 1)
|
||
remapped[:, 1] = 1.0 - old_y / max(square_side - 1, 1)
|
||
return remapped
|
||
|
||
|
||
def build_harmonic_atlas_texture(
|
||
texture_rgb: np.ndarray,
|
||
original_uvs: np.ndarray,
|
||
faces: np.ndarray,
|
||
face_uv_indices: np.ndarray,
|
||
missing_faces: np.ndarray,
|
||
face_corner_colors: Dict[int, np.ndarray],
|
||
tile_size: int = 12,
|
||
tile_pad: int = 2,
|
||
):
|
||
"""
|
||
给每个 missing face 分配一个小三角 tile。
|
||
tile 内按三个顶点颜色插值，不是纯色块。
|
||
|
||
输出贴图严格为 square_side x square_side。
|
||
"""
|
||
H, W = texture_rgb.shape[:2]
|
||
missing_faces_sorted = np.asarray(sorted(set(int(x) for x in missing_faces)), dtype=np.int64)
|
||
n_missing = len(missing_faces_sorted)
|
||
|
||
square_side, cols, rows, capacity = compute_square_atlas_layout(
|
||
width=W,
|
||
height=H,
|
||
n_tiles=n_missing,
|
||
tile_size=tile_size,
|
||
)
|
||
|
||
atlas_h = rows * tile_size
|
||
print(
|
||
f"atlas tiles: {n_missing}, cols={cols}, rows={rows}, "
|
||
f"capacity={capacity}, atlas_h={atlas_h}"
|
||
)
|
||
print(f"square canvas: {square_side} x {square_side}, original texture: {W} x {H}")
|
||
|
||
new_texture = make_square_canvas_with_edge_padding(texture_rgb, square_side)
|
||
|
||
# 原始 UV remap：保持原图像素位置不变，同时重映射 u/v 到方形画布。
|
||
remapped_uvs = remap_original_uvs_to_square_canvas(
|
||
original_uvs=original_uvs,
|
||
old_width=W,
|
||
old_height=H,
|
||
square_side=square_side,
|
||
)
|
||
|
||
new_uvs: List[List[float]] = remapped_uvs.tolist()
|
||
new_face_uv_indices = face_uv_indices.copy()
|
||
|
||
weights_map, uv_points = precompute_triangle_tile(tile_size, tile_pad)
|
||
weights_flat = weights_map.reshape(-1, 3) # [T*T,3]
|
||
|
||
for i, f in enumerate(tqdm.tqdm(missing_faces_sorted, desc="Baking missing face tiles")):
|
||
col = i % cols
|
||
row = i // cols
|
||
x0 = col * tile_size
|
||
y0 = H + row * tile_size
|
||
|
||
if y0 + tile_size > square_side or x0 + tile_size > square_side:
|
||
raise RuntimeError(
|
||
"atlas layout 计算错误：tile 超出方形画布，"
|
||
f"tile=({x0},{y0},{tile_size}), side={square_side}"
|
||
)
|
||
|
||
corner_colors = face_corner_colors.get(int(f))
|
||
if corner_colors is None:
|
||
corner_colors = np.full((3, 3), 128, dtype=np.float32)
|
||
corner_colors = np.asarray(corner_colors, dtype=np.float32)
|
||
|
||
tile = weights_flat @ corner_colors # [T*T,3]
|
||
tile = np.clip(np.rint(tile), 0, 255).astype(np.uint8).reshape(tile_size, tile_size, 3)
|
||
|
||
new_texture[y0:y0 + tile_size, x0:x0 + tile_size, :] = tile
|
||
|
||
new_vt = []
|
||
for p in uv_points:
|
||
px = x0 + float(p[0])
|
||
py = y0 + float(p[1])
|
||
u = px / max(square_side - 1, 1)
|
||
v = 1.0 - py / max(square_side - 1, 1)
|
||
new_uvs.append([u, v])
|
||
new_vt.append(len(new_uvs) - 1)
|
||
|
||
new_face_uv_indices[int(f)] = np.asarray(new_vt, dtype=new_face_uv_indices.dtype)
|
||
|
||
new_uvs_np = np.asarray(new_uvs, dtype=np.float32)
|
||
|
||
if new_texture.shape[0] != new_texture.shape[1]:
|
||
raise AssertionError(f"输出贴图不是 1:1: {new_texture.shape[1]} x {new_texture.shape[0]}")
|
||
if not np.isfinite(new_uvs_np).all():
|
||
raise AssertionError("输出 UV 中出现 NaN/Inf")
|
||
if new_face_uv_indices.min() < 0 or new_face_uv_indices.max() >= len(new_uvs_np):
|
||
raise AssertionError("输出 face uv index 越界")
|
||
|
||
return new_texture, new_uvs_np, new_face_uv_indices
|
||
|
||
|
||
# ============================================================
|
||
# Main process
|
||
# ============================================================
|
||
|
||
def process(
|
||
input_obj_path: str,
|
||
input_texture_path: str,
|
||
missing_faces_path: str,
|
||
output_obj_path: str,
|
||
output_texture_path: str,
|
||
missing_index_base: int = 0,
|
||
tile_size: int = 12,
|
||
tile_pad: int = 2,
|
||
smooth_iters: int = 80,
|
||
self_weight: float = 0.0,
|
||
corner_inset: float = 0.90,
|
||
):
|
||
start = time.time()
|
||
|
||
print("Reading input files...")
|
||
vertices, uvs, faces, face_uv_indices, mtllib_lines, usemtl_name = read_obj_basic(input_obj_path)
|
||
texture_rgb = read_texture_rgb(input_texture_path)
|
||
missing_faces = read_missing_faces(missing_faces_path, index_base=missing_index_base)
|
||
|
||
total_faces = len(faces)
|
||
missing_faces = missing_faces[(missing_faces >= 0) & (missing_faces < total_faces)]
|
||
missing_faces = np.unique(missing_faces)
|
||
missing_mask = np.zeros(total_faces, dtype=bool)
|
||
missing_mask[missing_faces] = True
|
||
|
||
print(f"vertices: {len(vertices)}")
|
||
print(f"uvs: {len(uvs)}")
|
||
print(f"faces: {len(faces)}")
|
||
print(f"missing faces: {len(missing_faces)}")
|
||
print(f"texture: {texture_rgb.shape[1]} x {texture_rgb.shape[0]}")
|
||
|
||
if len(missing_faces) == 0:
|
||
# 即使没有缺色面，也保证输出贴图是 1:1；如果输入已经是方形，这里等价于原样输出。
|
||
H, W = texture_rgb.shape[:2]
|
||
square_side, _, _, _ = compute_square_atlas_layout(W, H, 0, tile_size)
|
||
square_texture = make_square_canvas_with_edge_padding(texture_rgb, square_side)
|
||
square_uvs = remap_original_uvs_to_square_canvas(uvs, W, H, square_side)
|
||
write_obj_with_uv_coordinates(output_obj_path, vertices, square_uvs, faces, face_uv_indices, mtllib_lines, usemtl_name)
|
||
write_texture_rgb(output_texture_path, square_texture)
|
||
print(f"output texture: {square_texture.shape[1]} x {square_texture.shape[0]}")
|
||
return
|
||
|
||
t0 = time.time()
|
||
print("Building adjacency...")
|
||
face_adjacency, vertex_faces = build_face_adjacency_and_vertex_faces(faces, len(vertices))
|
||
print(f"adjacency using: {time.time() - t0:.2f}s")
|
||
|
||
t0 = time.time()
|
||
regions = find_missing_regions(face_adjacency, missing_faces, total_faces)
|
||
print(f"regions using: {time.time() - t0:.2f}s")
|
||
|
||
# face -> three corner colors
|
||
face_corner_colors: Dict[int, np.ndarray] = {}
|
||
|
||
for region_idx, (region_faces, _boundary_faces) in enumerate(tqdm.tqdm(regions, desc="Solving region colors")):
|
||
face_ids, unique_vertices, v2local, edges, local_faces = build_region_vertex_graph(region_faces, faces)
|
||
|
||
constrained, constraint_colors = compute_boundary_vertex_constraints(
|
||
unique_vertices=unique_vertices,
|
||
vertex_faces=vertex_faces,
|
||
faces=faces,
|
||
face_uv_indices=face_uv_indices,
|
||
uvs=uvs,
|
||
texture_rgb=texture_rgb,
|
||
missing_mask=missing_mask,
|
||
corner_inset=corner_inset,
|
||
)
|
||
|
||
if not constrained.any():
|
||
# 极少数情况：region 没有任何正常面顶点约束。
|
||
# 这种输入本身没有颜色来源，只能给灰色，避免崩。
|
||
print(f"Warning: region {region_idx} has no boundary vertex color constraints, using gray.")
|
||
vertex_colors = np.full((len(unique_vertices), 3), 128, dtype=np.float32)
|
||
else:
|
||
vertex_colors = solve_harmonic_vertex_colors(
|
||
edges=edges,
|
||
constrained=constrained,
|
||
constraint_colors=constraint_colors,
|
||
smooth_iters=smooth_iters,
|
||
self_weight=self_weight,
|
||
)
|
||
|
||
# 写成每个 missing face 的三个角颜色。
|
||
for local_fi, f in enumerate(face_ids):
|
||
lf = local_faces[local_fi]
|
||
face_corner_colors[int(f)] = vertex_colors[lf]
|
||
|
||
print("Building harmonic atlas...")
|
||
new_texture, new_uvs, new_face_uv_indices = build_harmonic_atlas_texture(
|
||
texture_rgb=texture_rgb,
|
||
original_uvs=uvs,
|
||
faces=faces,
|
||
face_uv_indices=face_uv_indices,
|
||
missing_faces=missing_faces,
|
||
face_corner_colors=face_corner_colors,
|
||
tile_size=tile_size,
|
||
tile_pad=tile_pad,
|
||
)
|
||
|
||
print("Writing outputs...")
|
||
write_obj_with_uv_coordinates(
|
||
output_obj_path,
|
||
vertices,
|
||
new_uvs,
|
||
faces,
|
||
new_face_uv_indices,
|
||
mtllib_lines=mtllib_lines,
|
||
usemtl_name=usemtl_name,
|
||
)
|
||
write_texture_rgb(output_texture_path, new_texture)
|
||
|
||
print(f"output texture: {new_texture.shape[1]} x {new_texture.shape[0]}")
|
||
print(f"new uv count: {len(new_uvs)}")
|
||
print(f"Total using: {time.time() - start:.2f}s")
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description="Fill missing color faces with harmonic vertex-color atlas baking.")
|
||
|
||
# 原始接口保持一致
|
||
parser.add_argument("--input_obj", type=str, required=True, help="Path to the input OBJ file")
|
||
parser.add_argument("--input_texture", type=str, required=True, help="Path to the input texture file")
|
||
parser.add_argument("--missing_faces", type=str, required=True, help="Path to missing face index file")
|
||
parser.add_argument("--output_obj", type=str, required=True, help="Path to the output OBJ file")
|
||
parser.add_argument("--output_texture", type=str, required=True, help="Path to the output texture file")
|
||
|
||
# 可选参数
|
||
parser.add_argument("--missing_index_base", type=int, default=0,
|
||
help="missing_faces.txt 的索引基准。原脚本默认 0-based；如果文件是 1-based，传 1。")
|
||
parser.add_argument("--tile_size", type=int, default=12,
|
||
help="每个 missing face 在新增 atlas 中的 tile 尺寸。越大越细，贴图越大。推荐 8~14。")
|
||
parser.add_argument("--tile_pad", type=int, default=2,
|
||
help="每个 tile 内三角形到边界的 padding。用于减少纹理过滤串色。推荐 2。")
|
||
parser.add_argument("--smooth_iters", type=int, default=80,
|
||
help="顶点颜色 harmonic 扩散迭代次数。越大越平滑，稍慢。推荐 50~120。")
|
||
parser.add_argument("--self_weight", type=float, default=0.0,
|
||
help="扩散时保留当前颜色的权重。一般用 0。")
|
||
parser.add_argument("--corner_inset", type=float, default=0.90,
|
||
help="边界顶点采样时往正常 face 内部缩进的比例。0.85~0.95 比较稳。")
|
||
|
||
args = parser.parse_args()
|
||
|
||
process(
|
||
input_obj_path=args.input_obj,
|
||
input_texture_path=args.input_texture,
|
||
missing_faces_path=args.missing_faces,
|
||
output_obj_path=args.output_obj,
|
||
output_texture_path=args.output_texture,
|
||
missing_index_base=args.missing_index_base,
|
||
tile_size=args.tile_size,
|
||
tile_pad=args.tile_pad,
|
||
smooth_iters=args.smooth_iters,
|
||
self_weight=args.self_weight,
|
||
corner_inset=args.corner_inset,
|
||
)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|