mirror of
https://github.com/Comfy-Org/ComfyUI.git
synced 2026-09-08 18:16:30 +08:00
feat(3d): File3DToMesh node — parse GLB/GLTF/OBJ/STL into MESH (#15919)
This commit is contained in:
Terry Jia
@@ -0,0 +1,410 @@
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
import struct
|
||||
import urllib.parse
|
||||
from io import BytesIO
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
_COMPONENT_DTYPES = {
|
||||
5120: np.int8,
|
||||
5121: np.uint8,
|
||||
5122: np.int16,
|
||||
5123: np.uint16,
|
||||
5125: np.uint32,
|
||||
5126: np.float32,
|
||||
}
|
||||
_TYPE_SIZES = {"SCALAR": 1, "VEC2": 2, "VEC3": 3, "VEC4": 4, "MAT2": 4, "MAT3": 9, "MAT4": 16}
|
||||
|
||||
_SUPPORTED_REQUIRED = {
|
||||
"EXT_mesh_gpu_instancing",
|
||||
"EXT_texture_webp",
|
||||
"KHR_materials_emissive_strength",
|
||||
"KHR_materials_unlit",
|
||||
}
|
||||
|
||||
_JSON_CHUNK = 0x4E4F534A
|
||||
_BIN_CHUNK = 0x004E4942
|
||||
|
||||
|
||||
def parse_container(data: bytes):
|
||||
if data[:4] == b"glTF":
|
||||
if len(data) < 12:
|
||||
raise ValueError("GLB file truncated (missing 12-byte header)")
|
||||
_, version, _ = struct.unpack_from("<4sII", data, 0)
|
||||
if version != 2:
|
||||
raise ValueError(f"unsupported GLB container version {version}")
|
||||
json_chunk = None
|
||||
bin_chunk = None
|
||||
offset = 12
|
||||
while offset + 8 <= len(data):
|
||||
chunk_len, chunk_type = struct.unpack_from("<II", data, offset)
|
||||
offset += 8
|
||||
chunk = data[offset:offset + chunk_len]
|
||||
offset += chunk_len
|
||||
if chunk_type == _JSON_CHUNK and json_chunk is None:
|
||||
json_chunk = chunk
|
||||
elif chunk_type == _BIN_CHUNK and bin_chunk is None:
|
||||
bin_chunk = chunk
|
||||
if json_chunk is None:
|
||||
raise ValueError("GLB file has no JSON chunk")
|
||||
return json.loads(json_chunk), bin_chunk
|
||||
return json.loads(data), None
|
||||
|
||||
|
||||
def _resolve_uri(uri: str, base_dir: str | None) -> bytes:
|
||||
if uri.startswith("data:"):
|
||||
header, _, payload = uri.partition(",")
|
||||
if ";base64" in header:
|
||||
return base64.b64decode(payload)
|
||||
return urllib.parse.unquote_to_bytes(payload)
|
||||
scheme = urllib.parse.urlparse(uri).scheme
|
||||
if scheme:
|
||||
raise ValueError(f"glTF references URI scheme {scheme!r}; only relative file paths are allowed")
|
||||
if base_dir is None:
|
||||
raise ValueError(
|
||||
f"glTF references external file {uri!r} but the source is an in-memory stream; "
|
||||
"use .glb (self-contained) or a disk-backed file"
|
||||
)
|
||||
relative = urllib.parse.unquote(uri)
|
||||
root = os.path.realpath(base_dir)
|
||||
path = os.path.realpath(os.path.join(root, relative))
|
||||
try:
|
||||
contained = not os.path.isabs(relative) and os.path.commonpath([root, path]) == root
|
||||
except ValueError:
|
||||
contained = False
|
||||
if not contained:
|
||||
raise ValueError(f"glTF references file {uri!r} outside the model directory")
|
||||
with open(path, "rb") as f:
|
||||
return f.read()
|
||||
|
||||
|
||||
def load_buffers(gltf: dict, bin_chunk: bytes | None, base_dir: str | None) -> list[bytes]:
|
||||
buffers = []
|
||||
for buf in gltf.get("buffers", []):
|
||||
if "uri" in buf:
|
||||
buffers.append(_resolve_uri(buf["uri"], base_dir))
|
||||
else:
|
||||
if bin_chunk is None:
|
||||
raise ValueError("glTF buffer has no uri and there is no GLB BIN chunk")
|
||||
buffers.append(bin_chunk)
|
||||
return buffers
|
||||
|
||||
|
||||
def _view_data(gltf: dict, buffers: list[bytes], view_index: int) -> bytes:
|
||||
view = gltf["bufferViews"][view_index]
|
||||
buf = buffers[view.get("buffer", 0)]
|
||||
start = view.get("byteOffset", 0)
|
||||
return buf[start:start + view["byteLength"]]
|
||||
|
||||
|
||||
def read_accessor(gltf: dict, buffers: list[bytes], index: int):
|
||||
acc = gltf["accessors"][index]
|
||||
count = acc["count"]
|
||||
ncomp = _TYPE_SIZES[acc["type"]]
|
||||
dtype = np.dtype(_COMPONENT_DTYPES[acc["componentType"]])
|
||||
elem = dtype.itemsize * ncomp
|
||||
|
||||
if "bufferView" in acc:
|
||||
view = gltf["bufferViews"][acc["bufferView"]]
|
||||
buf = buffers[view.get("buffer", 0)]
|
||||
start = view.get("byteOffset", 0) + acc.get("byteOffset", 0)
|
||||
stride = view.get("byteStride") or elem
|
||||
if stride == elem:
|
||||
arr = np.frombuffer(buf, dtype, count * ncomp, start).reshape(count, ncomp).copy()
|
||||
else:
|
||||
raw = np.frombuffer(buf, np.uint8, stride * (count - 1) + elem, start)
|
||||
rows = np.lib.stride_tricks.as_strided(raw, (count, elem), (stride, 1))
|
||||
arr = np.ascontiguousarray(rows).view(dtype).reshape(count, ncomp)
|
||||
else:
|
||||
arr = np.zeros((count, ncomp), dtype)
|
||||
|
||||
sparse = acc.get("sparse")
|
||||
if sparse:
|
||||
n = sparse["count"]
|
||||
idx_def = sparse["indices"]
|
||||
val_def = sparse["values"]
|
||||
idx_dtype = np.dtype(_COMPONENT_DTYPES[idx_def["componentType"]])
|
||||
iview = gltf["bufferViews"][idx_def["bufferView"]]
|
||||
ibuf = buffers[iview.get("buffer", 0)]
|
||||
sidx = np.frombuffer(ibuf, idx_dtype, n,
|
||||
iview.get("byteOffset", 0) + idx_def.get("byteOffset", 0)).astype(np.int64)
|
||||
vview = gltf["bufferViews"][val_def["bufferView"]]
|
||||
vbuf = buffers[vview.get("buffer", 0)]
|
||||
svals = np.frombuffer(vbuf, dtype, n * ncomp,
|
||||
vview.get("byteOffset", 0) + val_def.get("byteOffset", 0)).reshape(n, ncomp)
|
||||
arr[sidx] = svals
|
||||
return arr, bool(acc.get("normalized"))
|
||||
|
||||
|
||||
def to_float(arr: np.ndarray, normalized: bool) -> np.ndarray:
|
||||
if arr.dtype == np.float32:
|
||||
return arr
|
||||
out = arr.astype(np.float32)
|
||||
if not normalized:
|
||||
return out
|
||||
if arr.dtype == np.uint8:
|
||||
return out / 255.0
|
||||
if arr.dtype == np.uint16:
|
||||
return out / 65535.0
|
||||
if arr.dtype == np.int8:
|
||||
return np.maximum(out / 127.0, -1.0)
|
||||
if arr.dtype == np.int16:
|
||||
return np.maximum(out / 32767.0, -1.0)
|
||||
return out
|
||||
|
||||
|
||||
def _quat_to_matrix(x: float, y: float, z: float, w: float) -> np.ndarray:
|
||||
return np.array([
|
||||
[1 - 2 * (y * y + z * z), 2 * (x * y - z * w), 2 * (x * z + y * w)],
|
||||
[2 * (x * y + z * w), 1 - 2 * (x * x + z * z), 2 * (y * z - x * w)],
|
||||
[2 * (x * z - y * w), 2 * (y * z + x * w), 1 - 2 * (x * x + y * y)],
|
||||
], dtype=np.float32)
|
||||
|
||||
|
||||
def _node_local_matrix(node: dict) -> np.ndarray:
|
||||
if "matrix" in node:
|
||||
return np.array(node["matrix"], np.float32).reshape(4, 4).T # glTF stores column-major
|
||||
m = np.eye(4, dtype=np.float32)
|
||||
rot = _quat_to_matrix(*node.get("rotation", (0.0, 0.0, 0.0, 1.0)))
|
||||
scale = np.array(node.get("scale", (1.0, 1.0, 1.0)), np.float32)
|
||||
m[:3, :3] = rot * scale
|
||||
m[:3, 3] = node.get("translation", (0.0, 0.0, 0.0))
|
||||
return m
|
||||
|
||||
|
||||
def _instance_matrices(gltf: dict, buffers: list[bytes], node: dict, warn):
|
||||
attrs = node.get("extensions", {}).get("EXT_mesh_gpu_instancing", {}).get("attributes", {})
|
||||
if not attrs:
|
||||
return None
|
||||
counts = {gltf["accessors"][a]["count"] for a in attrs.values()}
|
||||
if len(counts) > 1:
|
||||
raise ValueError(f"EXT_mesh_gpu_instancing attribute accessors disagree on instance count: {sorted(counts)}")
|
||||
if not any(k in attrs for k in ("TRANSLATION", "ROTATION", "SCALE")):
|
||||
warn("instancing-custom", "EXT_mesh_gpu_instancing has only custom attributes; importing a single copy")
|
||||
return None
|
||||
translation = to_float(*read_accessor(gltf, buffers, attrs["TRANSLATION"])) if "TRANSLATION" in attrs else None
|
||||
rotation = to_float(*read_accessor(gltf, buffers, attrs["ROTATION"])) if "ROTATION" in attrs else None
|
||||
scale = to_float(*read_accessor(gltf, buffers, attrs["SCALE"])) if "SCALE" in attrs else None
|
||||
count = counts.pop()
|
||||
matrices = []
|
||||
for i in range(count):
|
||||
m = np.eye(4, dtype=np.float32)
|
||||
rot = _quat_to_matrix(*rotation[i]) if rotation is not None else np.eye(3, dtype=np.float32)
|
||||
m[:3, :3] = rot * (scale[i] if scale is not None else 1.0)
|
||||
if translation is not None:
|
||||
m[:3, 3] = translation[i]
|
||||
matrices.append(m)
|
||||
return matrices
|
||||
|
||||
|
||||
def _iter_mesh_nodes(gltf: dict, buffers: list[bytes], warn):
|
||||
nodes = gltf.get("nodes", [])
|
||||
scenes = gltf.get("scenes")
|
||||
if scenes:
|
||||
roots = scenes[gltf.get("scene", 0)].get("nodes", [])
|
||||
elif nodes:
|
||||
children = {c for n in nodes for c in n.get("children", [])}
|
||||
roots = [i for i in range(len(nodes)) if i not in children]
|
||||
else:
|
||||
for i in range(len(gltf.get("meshes", []))):
|
||||
yield {"mesh": i}, np.eye(4, dtype=np.float32)
|
||||
return
|
||||
seen = set()
|
||||
stack = [(i, np.eye(4, dtype=np.float32)) for i in roots]
|
||||
while stack:
|
||||
index, parent = stack.pop()
|
||||
if index in seen:
|
||||
continue
|
||||
seen.add(index)
|
||||
node = nodes[index]
|
||||
world = parent @ _node_local_matrix(node)
|
||||
if "mesh" in node:
|
||||
instances = _instance_matrices(gltf, buffers, node, warn)
|
||||
if instances is None:
|
||||
yield node, world
|
||||
else:
|
||||
warn("instancing", f"EXT_mesh_gpu_instancing: expanding {len(instances)} instances into merged geometry")
|
||||
for matrix in instances:
|
||||
yield node, world @ matrix
|
||||
for child in node.get("children", []):
|
||||
stack.append((child, world))
|
||||
|
||||
|
||||
def _to_triangles(indices: np.ndarray, mode: int) -> np.ndarray:
|
||||
if mode == 4:
|
||||
if len(indices) % 3:
|
||||
raise ValueError("TRIANGLES primitive index count must be divisible by 3")
|
||||
return indices.reshape(-1, 3)
|
||||
if len(indices) < 3:
|
||||
return np.zeros((0, 3), np.int64)
|
||||
if mode == 6:
|
||||
first = np.full(len(indices) - 2, indices[0], dtype=np.int64)
|
||||
return np.stack([first, indices[1:-1], indices[2:]], axis=1)
|
||||
tris = np.stack([indices[:-2], indices[1:-1], indices[2:]], axis=1)
|
||||
tris[1::2] = tris[1::2, ::-1]
|
||||
return tris
|
||||
|
||||
|
||||
def _vertex_attr(gltf, buffers, accessor_index, n_verts, warn, label):
|
||||
arr = to_float(*read_accessor(gltf, buffers, accessor_index))
|
||||
if arr.shape[0] < n_verts:
|
||||
warn(f"count:{label}", f"{label} has {arr.shape[0]} entries for {n_verts} vertices; attribute dropped")
|
||||
return None
|
||||
return arr[:n_verts]
|
||||
|
||||
|
||||
def load_scene_geometry(gltf: dict, buffers: list[bytes], warn) -> list[dict]:
|
||||
required = set(gltf.get("extensionsRequired", []))
|
||||
unsupported = required - _SUPPORTED_REQUIRED
|
||||
if unsupported:
|
||||
raise ValueError(f"glTF requires extensions this loader does not support: {sorted(unsupported)}")
|
||||
|
||||
prims = []
|
||||
for node, world in _iter_mesh_nodes(gltf, buffers, warn):
|
||||
if "skin" in node:
|
||||
warn("skin", "skinned mesh: joints/weights ignored, geometry imported in bind pose")
|
||||
linear = world[:3, :3]
|
||||
det = float(np.linalg.det(linear))
|
||||
normal_mat = np.linalg.inv(linear).T if abs(det) > 1e-12 else linear
|
||||
flip_winding = det < 0.0
|
||||
|
||||
for prim in gltf["meshes"][node["mesh"]].get("primitives", []):
|
||||
mode = prim.get("mode", 4)
|
||||
if mode not in (4, 5, 6):
|
||||
warn(f"mode{mode}", f"skipping non-triangle primitive (mode {mode})")
|
||||
continue
|
||||
attrs = prim.get("attributes", {})
|
||||
if "POSITION" not in attrs:
|
||||
continue
|
||||
pos = to_float(*read_accessor(gltf, buffers, attrs["POSITION"]))[:, :3]
|
||||
n_verts = pos.shape[0]
|
||||
if n_verts == 0:
|
||||
continue
|
||||
pos = pos @ linear.T + world[:3, 3]
|
||||
|
||||
if "indices" in prim:
|
||||
indices = read_accessor(gltf, buffers, prim["indices"])[0].reshape(-1).astype(np.int64)
|
||||
else:
|
||||
indices = np.arange(n_verts, dtype=np.int64)
|
||||
faces = _to_triangles(indices, mode)
|
||||
if faces.shape[0] == 0:
|
||||
continue
|
||||
if faces.min() < 0 or faces.max() >= n_verts:
|
||||
raise ValueError("primitive contains a face index outside its POSITION accessor")
|
||||
if flip_winding:
|
||||
faces = np.ascontiguousarray(faces[:, ::-1])
|
||||
if prim.get("targets"):
|
||||
warn("morph", "morph targets ignored; base geometry imported")
|
||||
|
||||
out = {"positions": np.ascontiguousarray(pos, np.float32), "faces": faces,
|
||||
"uvs": None, "colors": None, "normals": None, "tangents": None,
|
||||
"material": prim.get("material")}
|
||||
if "TEXCOORD_0" in attrs:
|
||||
uv = _vertex_attr(gltf, buffers, attrs["TEXCOORD_0"], n_verts, warn, "TEXCOORD_0")
|
||||
out["uvs"] = uv[:, :2] if uv is not None else None
|
||||
if "COLOR_0" in attrs:
|
||||
arr, normalized = read_accessor(gltf, buffers, attrs["COLOR_0"])
|
||||
arr = to_float(arr, normalized or arr.dtype != np.float32)
|
||||
if arr.shape[0] >= n_verts:
|
||||
out["colors"] = np.clip(arr[:n_verts], 0.0, 1.0)
|
||||
if "NORMAL" in attrs:
|
||||
nrm = _vertex_attr(gltf, buffers, attrs["NORMAL"], n_verts, warn, "NORMAL")
|
||||
if nrm is not None:
|
||||
nrm = nrm[:, :3] @ normal_mat.T
|
||||
out["normals"] = np.ascontiguousarray(
|
||||
nrm / np.maximum(np.linalg.norm(nrm, axis=1, keepdims=True), 1e-12), np.float32)
|
||||
if "TANGENT" in attrs:
|
||||
tan = _vertex_attr(gltf, buffers, attrs["TANGENT"], n_verts, warn, "TANGENT")
|
||||
if tan is not None and tan.shape[1] == 4:
|
||||
txyz = tan[:, :3] @ linear.T
|
||||
txyz /= np.maximum(np.linalg.norm(txyz, axis=1, keepdims=True), 1e-12)
|
||||
tw = tan[:, 3:4] * (-1.0 if flip_winding else 1.0)
|
||||
out["tangents"] = np.ascontiguousarray(np.concatenate([txyz, tw], axis=1), np.float32)
|
||||
prims.append(out)
|
||||
return prims
|
||||
|
||||
|
||||
def _decode_texture(gltf, buffers, base_dir, tex_info, warn, label):
|
||||
if tex_info is None:
|
||||
return None
|
||||
if tex_info.get("texCoord", 0) != 0:
|
||||
warn(f"texcoord:{label}", f"{label} uses TEXCOORD_{tex_info['texCoord']}; MESH only carries TEXCOORD_0")
|
||||
if "KHR_texture_transform" in tex_info.get("extensions", {}):
|
||||
warn("textransform", "KHR_texture_transform ignored; UVs used as-is")
|
||||
tex_def = gltf["textures"][tex_info["index"]]
|
||||
source = tex_def.get("source")
|
||||
if source is None:
|
||||
webp = tex_def.get("extensions", {}).get("EXT_texture_webp")
|
||||
source = webp.get("source") if webp is not None else None
|
||||
if source is None:
|
||||
warn(f"compressed:{label}", f"{label}: compressed texture (basisu/ktx2) without fallback; skipped")
|
||||
return None
|
||||
image_def = gltf["images"][source]
|
||||
if "bufferView" in image_def:
|
||||
raw = _view_data(gltf, buffers, image_def["bufferView"])
|
||||
elif "uri" in image_def:
|
||||
raw = _resolve_uri(image_def["uri"], base_dir)
|
||||
else:
|
||||
return None
|
||||
img = Image.open(BytesIO(bytes(raw)))
|
||||
return np.asarray(img.convert("RGB"), dtype=np.float32) / 255.0
|
||||
|
||||
|
||||
def extract_material(gltf: dict, buffers: list[bytes], base_dir: str | None, material_index, warn) -> dict:
|
||||
result = {"texture": None, "metallic_roughness": None, "normal_map": None, "emissive": None,
|
||||
"occlusion_in_mr": False, "unlit": False, "material": None}
|
||||
if material_index is None:
|
||||
return result
|
||||
mat = gltf["materials"][material_index]
|
||||
extensions = mat.get("extensions", {})
|
||||
result["unlit"] = "KHR_materials_unlit" in extensions
|
||||
|
||||
pbr = mat.get("pbrMetallicRoughness", {})
|
||||
result["texture"] = _decode_texture(gltf, buffers, base_dir, pbr.get("baseColorTexture"), warn, "baseColorTexture")
|
||||
mr_info = pbr.get("metallicRoughnessTexture")
|
||||
result["metallic_roughness"] = _decode_texture(gltf, buffers, base_dir, mr_info, warn, "metallicRoughnessTexture")
|
||||
normal_info = mat.get("normalTexture")
|
||||
result["normal_map"] = _decode_texture(gltf, buffers, base_dir, normal_info, warn, "normalTexture")
|
||||
result["emissive"] = _decode_texture(gltf, buffers, base_dir, mat.get("emissiveTexture"), warn, "emissiveTexture")
|
||||
|
||||
occlusion = mat.get("occlusionTexture")
|
||||
if occlusion is not None:
|
||||
if mr_info is not None and occlusion["index"] == mr_info["index"]:
|
||||
result["occlusion_in_mr"] = True
|
||||
else:
|
||||
warn("occlusion", "standalone occlusionTexture not representable in MESH (ORM packing only); skipped")
|
||||
|
||||
overrides = {}
|
||||
base_color = pbr.get("baseColorFactor")
|
||||
if base_color is not None and list(base_color) != [1.0, 1.0, 1.0, 1.0]:
|
||||
overrides["base_color_factor"] = [float(c) for c in base_color]
|
||||
overrides["metallic_factor"] = float(pbr.get("metallicFactor", 1.0))
|
||||
overrides["roughness_factor"] = float(pbr.get("roughnessFactor", 1.0))
|
||||
overrides["double_sided"] = bool(mat.get("doubleSided", False))
|
||||
if normal_info is not None and normal_info.get("scale", 1.0) != 1.0:
|
||||
overrides["normal_scale"] = float(normal_info["scale"])
|
||||
if occlusion is not None and occlusion.get("strength", 1.0) != 1.0:
|
||||
overrides["occlusion_strength"] = float(occlusion["strength"])
|
||||
emissive_factor = mat.get("emissiveFactor", (0.0, 0.0, 0.0))
|
||||
if any(c > 0.0 for c in emissive_factor):
|
||||
overrides["emissive_factor"] = [float(c) for c in emissive_factor]
|
||||
strength = extensions.get("KHR_materials_emissive_strength", {}).get("emissiveStrength")
|
||||
if strength is not None:
|
||||
overrides["emissive_strength"] = float(strength)
|
||||
result["material"] = overrides
|
||||
return result
|
||||
|
||||
|
||||
def load_gltf(data: bytes, base_dir: str | None, warn):
|
||||
gltf, bin_chunk = parse_container(data)
|
||||
version = str(gltf.get("asset", {}).get("version", ""))
|
||||
if version.partition(".")[0] != "2":
|
||||
raise ValueError(f"unsupported glTF asset version {version or 'unknown'}; only glTF 2.x is supported")
|
||||
buffers = load_buffers(gltf, bin_chunk, base_dir)
|
||||
return gltf, buffers, load_scene_geometry(gltf, buffers, warn)
|
||||
|
||||
|
||||
__all__ = ["load_gltf", "extract_material", "parse_container", "read_accessor", "to_float"]
|
||||
@@ -0,0 +1,172 @@
|
||||
import logging
|
||||
import re
|
||||
import struct
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def _srgb_to_linear(c: np.ndarray) -> np.ndarray:
|
||||
return np.where(c <= 0.04045, c / 12.92, ((c + 0.055) / 1.055) ** 2.4).astype(np.float32)
|
||||
|
||||
|
||||
def load_obj(data: bytes) -> dict:
|
||||
text = data.decode("utf-8", errors="replace")
|
||||
text = text.replace("\r\n", "\n").replace("\\\n", "")
|
||||
positions: list[tuple] = []
|
||||
pos_colors: list[tuple] = []
|
||||
uvs: list[tuple] = []
|
||||
normals: list[tuple] = []
|
||||
|
||||
corner_map: dict[tuple, int] = {}
|
||||
out_pos: list[tuple] = []
|
||||
out_col: list[tuple] = []
|
||||
out_uv: list[tuple] = []
|
||||
out_nrm: list[tuple] = []
|
||||
faces: list[tuple] = []
|
||||
has_color = False
|
||||
any_uv = False
|
||||
any_normal = False
|
||||
missing_normal = False
|
||||
warned_mtl = False
|
||||
|
||||
def resolve(index_str: str, count: int) -> int:
|
||||
i = int(index_str)
|
||||
resolved = i - 1 if i > 0 else count + i
|
||||
if i == 0 or not 0 <= resolved < count:
|
||||
raise ValueError(f"OBJ index {i} out of range for {count} entries")
|
||||
return resolved
|
||||
|
||||
def corner(spec: str) -> int:
|
||||
nonlocal any_uv, any_normal, missing_normal
|
||||
parts = spec.split("/")
|
||||
vi = resolve(parts[0], len(positions))
|
||||
ti = resolve(parts[1], len(uvs)) if len(parts) > 1 and parts[1] else None
|
||||
ni = resolve(parts[2], len(normals)) if len(parts) > 2 and parts[2] else None
|
||||
key = (vi, ti, ni)
|
||||
cached = corner_map.get(key)
|
||||
if cached is not None:
|
||||
return cached
|
||||
index = len(out_pos)
|
||||
out_pos.append(positions[vi])
|
||||
out_col.append(pos_colors[vi])
|
||||
if ti is not None:
|
||||
any_uv = True
|
||||
u, v = uvs[ti]
|
||||
out_uv.append((u, 1.0 - v))
|
||||
else:
|
||||
out_uv.append((0.0, 0.0))
|
||||
if ni is not None:
|
||||
any_normal = True
|
||||
out_nrm.append(normals[ni])
|
||||
else:
|
||||
missing_normal = True
|
||||
out_nrm.append((0.0, 0.0, 0.0))
|
||||
corner_map[key] = index
|
||||
return index
|
||||
|
||||
for raw_line in text.splitlines():
|
||||
line = raw_line.strip()
|
||||
if not line or line.startswith("#"):
|
||||
continue
|
||||
parts = line.split()
|
||||
tag = parts[0]
|
||||
if tag == "v":
|
||||
positions.append(tuple(float(x) for x in parts[1:4]))
|
||||
if len(parts) >= 7:
|
||||
pos_colors.append(tuple(float(x) for x in parts[4:7]))
|
||||
has_color = True
|
||||
else:
|
||||
pos_colors.append((1.0, 1.0, 1.0))
|
||||
elif tag == "vt":
|
||||
uvs.append((float(parts[1]), float(parts[2]) if len(parts) > 2 else 0.0))
|
||||
elif tag == "vn":
|
||||
normals.append(tuple(float(x) for x in parts[1:4]))
|
||||
elif tag == "f":
|
||||
specs = parts[1:]
|
||||
if len(specs) < 3:
|
||||
continue
|
||||
indices = [corner(s) for s in specs]
|
||||
for i in range(1, len(indices) - 1):
|
||||
faces.append((indices[0], indices[i], indices[i + 1]))
|
||||
elif tag in ("mtllib", "usemtl") and not warned_mtl:
|
||||
warned_mtl = True
|
||||
logging.warning("Get3DComponents: OBJ materials (.mtl) are not loaded; geometry only")
|
||||
|
||||
if not faces:
|
||||
raise ValueError("OBJ contains no faces")
|
||||
|
||||
prim = {
|
||||
"positions": np.array(out_pos, np.float32),
|
||||
"faces": np.array(faces, np.int64),
|
||||
"uvs": np.array(out_uv, np.float32) if any_uv else None,
|
||||
"colors": _srgb_to_linear(np.clip(np.array(out_col, np.float32), 0.0, 1.0)) if has_color else None,
|
||||
"normals": None,
|
||||
"tangents": None,
|
||||
"material": None,
|
||||
}
|
||||
if any_normal and missing_normal:
|
||||
logging.warning("Get3DComponents: OBJ has faces without vn indices; normals dropped")
|
||||
elif any_normal:
|
||||
nrm = np.array(out_nrm, np.float32)
|
||||
lengths = np.linalg.norm(nrm, axis=1, keepdims=True)
|
||||
if float(lengths.max()) > 1e-6:
|
||||
prim["normals"] = nrm / np.maximum(lengths, 1e-12)
|
||||
return prim
|
||||
|
||||
|
||||
_STL_RECORD = np.dtype([("normal", "<f4", (3,)), ("verts", "<f4", (3, 3)), ("attr", "<u2")])
|
||||
|
||||
_STL_ASCII_FACET = re.compile(
|
||||
rb"facet\s+normal\s+(\S+)\s+(\S+)\s+(\S+).*?"
|
||||
rb"vertex\s+(\S+)\s+(\S+)\s+(\S+).*?"
|
||||
rb"vertex\s+(\S+)\s+(\S+)\s+(\S+).*?"
|
||||
rb"vertex\s+(\S+)\s+(\S+)\s+(\S+)",
|
||||
re.DOTALL,
|
||||
)
|
||||
|
||||
|
||||
def load_stl(data: bytes) -> dict:
|
||||
n_faces = struct.unpack_from("<I", data, 80)[0] if len(data) >= 84 else 0
|
||||
if n_faces > 0 and 84 + n_faces * _STL_RECORD.itemsize == len(data):
|
||||
return _load_stl_binary(data, n_faces)
|
||||
if b"solid" in data[:9]:
|
||||
facets = _STL_ASCII_FACET.findall(data)
|
||||
if not facets:
|
||||
raise ValueError("ASCII STL contains no facets")
|
||||
values = np.array(facets).astype(np.float32)
|
||||
positions = values[:, 3:12].reshape(-1, 3)
|
||||
file_normals = np.repeat(values[:, 0:3], 3, axis=0)
|
||||
return _stl_prim(positions, file_normals)
|
||||
available = (len(data) - 84) // _STL_RECORD.itemsize if len(data) >= 84 else 0
|
||||
n_faces = min(n_faces, available)
|
||||
if n_faces <= 0:
|
||||
raise ValueError("not a valid STL file (neither binary layout nor ASCII 'solid')")
|
||||
return _load_stl_binary(data, n_faces)
|
||||
|
||||
|
||||
def _load_stl_binary(data: bytes, n_faces: int) -> dict:
|
||||
records = np.frombuffer(data, _STL_RECORD, count=n_faces, offset=84)
|
||||
positions = records["verts"].reshape(-1, 3).astype(np.float32)
|
||||
file_normals = np.repeat(records["normal"], 3, axis=0).astype(np.float32)
|
||||
return _stl_prim(positions, file_normals, _stl_binary_colors(data[:80], records["attr"]))
|
||||
|
||||
|
||||
def _stl_binary_colors(header: bytes, attr: np.ndarray):
|
||||
pos = header.find(b"COLOR=")
|
||||
if pos < 0 or pos + 10 > len(header):
|
||||
return None
|
||||
default = np.frombuffer(header, np.uint8, 3, pos + 6).astype(np.float32) / 255.0
|
||||
a = attr.astype(np.uint16)
|
||||
per_face = np.stack([a & 0x1F, (a >> 5) & 0x1F, (a >> 10) & 0x1F], axis=1).astype(np.float32) / 31.0
|
||||
face_colors = np.where(((a & 0x8000) != 0)[:, None], default[None, :], per_face)
|
||||
return _srgb_to_linear(np.repeat(face_colors, 3, axis=0))
|
||||
|
||||
|
||||
def _stl_prim(positions: np.ndarray, normals: np.ndarray, colors=None) -> dict:
|
||||
if positions.shape[0] == 0:
|
||||
raise ValueError("STL contains no facets")
|
||||
faces = np.arange(positions.shape[0], dtype=np.int64).reshape(-1, 3)
|
||||
lengths = np.linalg.norm(normals, axis=1, keepdims=True)
|
||||
normals = normals / np.maximum(lengths, 1e-12) if float(lengths.max()) > 1e-6 else None
|
||||
return {"positions": np.ascontiguousarray(positions), "faces": faces, "uvs": None,
|
||||
"colors": colors, "normals": normals, "tangents": None, "material": None}
|
||||
@@ -0,0 +1,163 @@
|
||||
import logging
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from typing_extensions import override
|
||||
|
||||
from comfy_api.latest import ComfyExtension, IO, Types
|
||||
from comfy_extras.mesh3d.fileio import gltf_read, mesh_file_read
|
||||
|
||||
|
||||
def _sniff_format(data: bytes) -> str:
|
||||
if data[:4] == b"glTF":
|
||||
return "glb"
|
||||
head = data[:512].lstrip()
|
||||
if head[:1] == b"{":
|
||||
return "gltf"
|
||||
if head[:5].lower() == b"solid":
|
||||
return "stl"
|
||||
return ""
|
||||
|
||||
|
||||
def _merge_primitives(prims: list[dict]) -> dict:
|
||||
any_uv = any(p["uvs"] is not None for p in prims)
|
||||
any_color = any(p["colors"] is not None for p in prims)
|
||||
all_normals = all(p["normals"] is not None for p in prims)
|
||||
all_tangents = all_normals and all(p["tangents"] is not None for p in prims)
|
||||
color_channels = max((p["colors"].shape[1] for p in prims if p["colors"] is not None), default=3)
|
||||
if not all_normals and any(p["normals"] is not None for p in prims):
|
||||
logging.warning("Get3DComponents: some primitives lack normals; normals dropped "
|
||||
"(MeshSmoothNormals can regenerate them)")
|
||||
|
||||
verts, faces, uvs, colors, normals, tangents = [], [], [], [], [], []
|
||||
offset = 0
|
||||
for p in prims:
|
||||
v = p["positions"]
|
||||
n = v.shape[0]
|
||||
verts.append(v)
|
||||
faces.append(p["faces"] + offset)
|
||||
offset += n
|
||||
if any_uv:
|
||||
uvs.append(p["uvs"] if p["uvs"] is not None else np.zeros((n, 2), np.float32))
|
||||
if any_color:
|
||||
c = p["colors"] if p["colors"] is not None else np.ones((n, color_channels), np.float32)
|
||||
if c.shape[1] < color_channels:
|
||||
c = np.concatenate([c, np.ones((n, color_channels - c.shape[1]), np.float32)], axis=1)
|
||||
colors.append(c)
|
||||
if all_normals:
|
||||
normals.append(p["normals"])
|
||||
if all_tangents:
|
||||
tangents.append(p["tangents"])
|
||||
|
||||
return {
|
||||
"vertices": np.concatenate(verts, axis=0),
|
||||
"faces": np.concatenate(faces, axis=0),
|
||||
"uvs": np.concatenate(uvs, axis=0) if any_uv else None,
|
||||
"colors": np.concatenate(colors, axis=0) if any_color else None,
|
||||
"normals": np.concatenate(normals, axis=0) if all_normals else None,
|
||||
"tangents": np.concatenate(tangents, axis=0) if all_tangents else None,
|
||||
}
|
||||
|
||||
|
||||
def _batch(arr):
|
||||
return torch.from_numpy(arr)[None] if arr is not None else None
|
||||
|
||||
|
||||
class Get3DComponents(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="Get3DComponents",
|
||||
display_name="Get 3D Components",
|
||||
category="3d",
|
||||
description=(
|
||||
"Parse a 3D model file (GLB, GLTF, OBJ, STL) into an editable MESH for the "
|
||||
"mesh-processing nodes (decimate, remesh, UV unwrap, bake, ...). All scene "
|
||||
"nodes/primitives are merged into one mesh with their transforms applied; "
|
||||
"textures and material factors come from the first material. "
|
||||
"Counterpart of MeshToFile3D."
|
||||
),
|
||||
search_aliases=["file 3d to mesh", "extract mesh", "convert 3d", "parse glb", "import mesh",
|
||||
"load mesh from file", "file to mesh"],
|
||||
is_experimental=True,
|
||||
inputs=[
|
||||
IO.MultiType.Input(
|
||||
"model_3d",
|
||||
types=[IO.File3DGLB, IO.File3DGLTF, IO.File3DOBJ, IO.File3DSTL, IO.File3DAny],
|
||||
tooltip="3D model file from Load 3D or another 3D node. "
|
||||
"FBX/USDZ are not supported - convert to GLB first.",
|
||||
),
|
||||
],
|
||||
outputs=[IO.Mesh.Output(display_name="mesh")],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, model_3d: Types.File3D) -> IO.NodeOutput:
|
||||
data = model_3d.get_bytes()
|
||||
fmt = (model_3d.format or _sniff_format(data)).lower()
|
||||
if fmt in ("fbx", "usdz"):
|
||||
raise ValueError(f"Get3DComponents: .{fmt} parsing is not supported; convert the model to GLB/GLTF first")
|
||||
|
||||
warned = set()
|
||||
|
||||
def warn_once(key, message):
|
||||
if key not in warned:
|
||||
warned.add(key)
|
||||
logging.warning("Get3DComponents: %s", message)
|
||||
|
||||
material_info = None
|
||||
if fmt in ("glb", "gltf"):
|
||||
base_dir = os.path.dirname(model_3d.get_source()) if model_3d.is_disk_backed else None
|
||||
gltf, buffers, prims = gltf_read.load_gltf(data, base_dir, warn_once)
|
||||
if not prims:
|
||||
raise ValueError("Get3DComponents: no triangle geometry found in the glTF scene")
|
||||
material_indices = [p["material"] for p in prims if p["material"] is not None]
|
||||
if len(set(material_indices)) > 1:
|
||||
warn_once("multimat", f"{len(set(material_indices))} materials found; "
|
||||
"keeping textures/factors of the first only")
|
||||
first_material = material_indices[0] if material_indices else None
|
||||
material_info = gltf_read.extract_material(gltf, buffers, base_dir, first_material, warn_once)
|
||||
elif fmt == "obj":
|
||||
prims = [mesh_file_read.load_obj(data)]
|
||||
elif fmt == "stl":
|
||||
prims = [mesh_file_read.load_stl(data)]
|
||||
else:
|
||||
raise ValueError(f"Get3DComponents: unsupported or unrecognized format {fmt!r} "
|
||||
"(supported: glb, gltf, obj, stl)")
|
||||
|
||||
merged = _merge_primitives(prims)
|
||||
n_verts = merged["vertices"].shape[0]
|
||||
max_face = int(merged["faces"].max())
|
||||
if max_face >= n_verts:
|
||||
raise ValueError(f"Get3DComponents: face index {max_face} out of range for {n_verts} vertices (corrupt file?)")
|
||||
|
||||
material_info = material_info or {}
|
||||
mesh = Types.MESH(
|
||||
vertices=_batch(merged["vertices"]),
|
||||
faces=_batch(merged["faces"]),
|
||||
uvs=_batch(merged["uvs"]),
|
||||
vertex_colors=_batch(merged["colors"]),
|
||||
normals=_batch(merged["normals"]),
|
||||
tangents=_batch(merged["tangents"]),
|
||||
texture=_batch(material_info.get("texture")),
|
||||
metallic_roughness=_batch(material_info.get("metallic_roughness")),
|
||||
normal_map=_batch(material_info.get("normal_map")),
|
||||
emissive=_batch(material_info.get("emissive")),
|
||||
unlit=bool(material_info.get("unlit", False)),
|
||||
occlusion_in_mr=bool(material_info.get("occlusion_in_mr", False)),
|
||||
material=material_info.get("material") or None,
|
||||
)
|
||||
logging.info("Get3DComponents: %s -> %d vertices, %d faces (%d primitives)",
|
||||
fmt, n_verts, merged["faces"].shape[0], len(prims))
|
||||
return IO.NodeOutput(mesh)
|
||||
|
||||
|
||||
class MeshIOExtension(ComfyExtension):
|
||||
@override
|
||||
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
|
||||
return [Get3DComponents]
|
||||
|
||||
|
||||
async def comfy_entrypoint() -> MeshIOExtension:
|
||||
return MeshIOExtension()
|
||||
Reference in New Issue
Block a user