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