feat: add VRM model loading functionality
This commit introduces the core VRM loading capabilities for the 3D rendering system. It implements a dedicated VRMLoader class that can parse VRM files (which are binary glTF format) and extract essential data including meshes, materials, textures, and VRM-specific extension metadata. The implementation handles both binary and JSON formatted VRM files while supporting standard glTF accessor data extraction for vertices, normals, UV coordinates, and indices. This enables the application to load and render 3D models with proper material and texture support from VRM files.
This commit is contained in:
243
src/rendering/vrm_loader.py
Normal file
243
src/rendering/vrm_loader.py
Normal file
@@ -0,0 +1,243 @@
|
||||
"""
|
||||
VRM Loader - Loads and parses VRM model files
|
||||
"""
|
||||
import os
|
||||
import base64
|
||||
from typing import Optional, Dict, List, Tuple
|
||||
import numpy as np
|
||||
from pygltflib import GLTF2
|
||||
from PIL import Image
|
||||
import io
|
||||
|
||||
class VRMModel:
|
||||
"""Container for loaded VRM model data"""
|
||||
|
||||
def __init__(self):
|
||||
self.meshes: List[Dict] = []
|
||||
self.materials: List[Dict] = []
|
||||
self.textures: List[Optional[Image.Image]] = []
|
||||
self.nodes: List[Dict] = []
|
||||
self.vrm_meta: Optional[Dict] = None
|
||||
self.blend_shapes: Dict[str, List] = {}
|
||||
|
||||
class VRMLoader:
|
||||
"""Loads VRM models using pygltflib"""
|
||||
|
||||
def __init__(self):
|
||||
self.model: Optional[VRMModel] = None
|
||||
self.gltf: Optional[GLTF2] = None
|
||||
|
||||
def load(self, file_path: str) -> VRMModel:
|
||||
"""Load VRM file and extract data"""
|
||||
if not os.path.exists(file_path):
|
||||
raise FileNotFoundError(f"VRM file not found: {file_path}")
|
||||
|
||||
print(f"Loading VRM: {file_path}")
|
||||
|
||||
# VRM files are binary glTF (.glb) format
|
||||
# Check file magic bytes to determine if binary
|
||||
with open(file_path, 'rb') as f:
|
||||
magic = f.read(4)
|
||||
|
||||
is_binary = magic == b'glTF'
|
||||
|
||||
# Load glTF/VRM file
|
||||
if is_binary:
|
||||
self.gltf = GLTF2().load_binary(file_path)
|
||||
else:
|
||||
self.gltf = GLTF2().load_json(file_path)
|
||||
|
||||
self.model = VRMModel()
|
||||
|
||||
# Extract VRM extension data
|
||||
self._extract_vrm_extensions()
|
||||
|
||||
# Extract meshes
|
||||
self._extract_meshes()
|
||||
|
||||
# Extract materials
|
||||
self._extract_materials()
|
||||
|
||||
# Extract textures
|
||||
self._extract_textures(file_path)
|
||||
|
||||
print(f"VRM loaded: {len(self.model.meshes)} meshes, {len(self.model.materials)} materials")
|
||||
|
||||
return self.model
|
||||
|
||||
def _extract_vrm_extensions(self):
|
||||
"""Extract VRM-specific extension data"""
|
||||
if hasattr(self.gltf, 'extensions') and self.gltf.extensions:
|
||||
if 'VRM' in self.gltf.extensions:
|
||||
self.model.vrm_meta = self.gltf.extensions['VRM']
|
||||
print("VRM extension found")
|
||||
|
||||
def _extract_meshes(self):
|
||||
"""Extract mesh data from glTF"""
|
||||
if not self.gltf.meshes:
|
||||
return
|
||||
|
||||
for mesh_idx, mesh in enumerate(self.gltf.meshes):
|
||||
mesh_data = {
|
||||
'name': mesh.name or f'Mesh_{mesh_idx}',
|
||||
'primitives': []
|
||||
}
|
||||
|
||||
for prim in mesh.primitives:
|
||||
prim_data = self._extract_primitive(prim)
|
||||
mesh_data['primitives'].append(prim_data)
|
||||
|
||||
self.model.meshes.append(mesh_data)
|
||||
|
||||
def _extract_primitive(self, primitive) -> Dict:
|
||||
"""Extract primitive (submesh) data"""
|
||||
prim_data = {
|
||||
'vertices': None,
|
||||
'normals': None,
|
||||
'uvs': None,
|
||||
'indices': None,
|
||||
'material_idx': primitive.material if primitive.material is not None else 0
|
||||
}
|
||||
|
||||
# Get vertices (POSITION)
|
||||
if primitive.attributes.POSITION is not None:
|
||||
prim_data['vertices'] = self._get_accessor_data(primitive.attributes.POSITION)
|
||||
|
||||
# Get normals
|
||||
if primitive.attributes.NORMAL is not None:
|
||||
prim_data['normals'] = self._get_accessor_data(primitive.attributes.NORMAL)
|
||||
|
||||
# Get UVs (TEXCOORD_0)
|
||||
if primitive.attributes.TEXCOORD_0 is not None:
|
||||
prim_data['uvs'] = self._get_accessor_data(primitive.attributes.TEXCOORD_0)
|
||||
|
||||
# Get indices
|
||||
if primitive.indices is not None:
|
||||
prim_data['indices'] = self._get_accessor_data(primitive.indices)
|
||||
|
||||
return prim_data
|
||||
|
||||
def _get_accessor_data(self, accessor_idx: int) -> np.ndarray:
|
||||
"""Get data from glTF accessor"""
|
||||
accessor = self.gltf.accessors[accessor_idx]
|
||||
buffer_view = self.gltf.bufferViews[accessor.bufferView]
|
||||
buffer = self.gltf.buffers[buffer_view.buffer]
|
||||
|
||||
# Get buffer data
|
||||
if buffer.uri:
|
||||
# External buffer (shouldn't happen with .vrm files)
|
||||
raise NotImplementedError("External buffers not supported")
|
||||
else:
|
||||
# Embedded buffer
|
||||
data = self.gltf.binary_blob()
|
||||
|
||||
# Calculate offset
|
||||
offset = buffer_view.byteOffset + accessor.byteOffset if accessor.byteOffset else buffer_view.byteOffset
|
||||
|
||||
# Map component type to numpy dtype
|
||||
component_types = {
|
||||
5120: np.int8,
|
||||
5121: np.uint8,
|
||||
5122: np.int16,
|
||||
5123: np.uint16,
|
||||
5125: np.uint32,
|
||||
5126: np.float32,
|
||||
}
|
||||
|
||||
dtype = component_types.get(accessor.componentType)
|
||||
if dtype is None:
|
||||
raise ValueError(f"Unknown component type: {accessor.componentType}")
|
||||
|
||||
# Map accessor type to element count
|
||||
type_sizes = {
|
||||
'SCALAR': 1,
|
||||
'VEC2': 2,
|
||||
'VEC3': 3,
|
||||
'VEC4': 4,
|
||||
'MAT2': 4,
|
||||
'MAT3': 9,
|
||||
'MAT4': 16,
|
||||
}
|
||||
|
||||
element_size = type_sizes.get(accessor.type, 1)
|
||||
|
||||
# Extract data
|
||||
count = accessor.count * element_size
|
||||
array = np.frombuffer(data, dtype=dtype, count=count, offset=offset)
|
||||
|
||||
# Reshape if needed
|
||||
if element_size > 1:
|
||||
array = array.reshape((accessor.count, element_size))
|
||||
|
||||
return array
|
||||
|
||||
def _extract_materials(self):
|
||||
"""Extract material data"""
|
||||
if not self.gltf.materials:
|
||||
# Create default material
|
||||
self.model.materials.append({
|
||||
'name': 'Default',
|
||||
'base_color': [1.0, 1.0, 1.0, 1.0],
|
||||
'texture_idx': None
|
||||
})
|
||||
return
|
||||
|
||||
for mat_idx, material in enumerate(self.gltf.materials):
|
||||
mat_data = {
|
||||
'name': material.name or f'Material_{mat_idx}',
|
||||
'base_color': [1.0, 1.0, 1.0, 1.0],
|
||||
'texture_idx': None
|
||||
}
|
||||
|
||||
# Get base color
|
||||
if material.pbrMetallicRoughness:
|
||||
if material.pbrMetallicRoughness.baseColorFactor:
|
||||
mat_data['base_color'] = material.pbrMetallicRoughness.baseColorFactor
|
||||
|
||||
# Get base color texture
|
||||
if material.pbrMetallicRoughness.baseColorTexture:
|
||||
mat_data['texture_idx'] = material.pbrMetallicRoughness.baseColorTexture.index
|
||||
|
||||
self.model.materials.append(mat_data)
|
||||
|
||||
def _extract_textures(self, vrm_path: str):
|
||||
"""Extract texture images"""
|
||||
if not self.gltf.textures:
|
||||
return
|
||||
|
||||
base_dir = os.path.dirname(vrm_path)
|
||||
|
||||
for tex_idx, texture in enumerate(self.gltf.textures):
|
||||
try:
|
||||
image = self.gltf.images[texture.source]
|
||||
|
||||
if image.uri:
|
||||
# External image file
|
||||
if image.uri.startswith('data:'):
|
||||
# Data URI
|
||||
header, data = image.uri.split(',', 1)
|
||||
img_data = base64.b64decode(data)
|
||||
img = Image.open(io.BytesIO(img_data))
|
||||
else:
|
||||
# File path
|
||||
img_path = os.path.join(base_dir, image.uri)
|
||||
img = Image.open(img_path)
|
||||
else:
|
||||
# Embedded in bufferView
|
||||
buffer_view = self.gltf.bufferViews[image.bufferView]
|
||||
data = self.gltf.binary_blob()
|
||||
offset = buffer_view.byteOffset
|
||||
length = buffer_view.byteLength
|
||||
img_data = data[offset:offset+length]
|
||||
img = Image.open(io.BytesIO(img_data))
|
||||
|
||||
# Convert to RGBA if needed
|
||||
if img.mode != 'RGBA':
|
||||
img = img.convert('RGBA')
|
||||
|
||||
self.model.textures.append(img)
|
||||
print(f"Loaded texture {tex_idx}: {img.size}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"Failed to load texture {tex_idx}: {e}")
|
||||
self.model.textures.append(None)
|
||||
Reference in New Issue
Block a user