| 11 | |
| 12 | |
| 13 | class MeshSampler(object): |
| 14 | def __init__(self, path): |
| 15 | self.path = path |
| 16 | self._load(path) |
| 17 | |
| 18 | def _load(self, path): |
| 19 | # load mesh obj and mtl |
| 20 | obj = trimesh.load(path, process=False) |
| 21 | assert type(obj) == trimesh.Trimesh |
| 22 | |
| 23 | self.vs = obj.vertices |
| 24 | self.fs = obj.faces |
| 25 | self.vns = obj.vertex_normals |
| 26 | self.uvs = obj.visual.uv |
| 27 | assert self.uvs is not None |
| 28 | |
| 29 | # albedo_path = os.path.join(os.path.dirname(path), "textures/albedo.png") |
| 30 | albedo_path = glob.glob(os.path.join(os.path.dirname(path), "textures/albedo.*"))[0] |
| 31 | self.image_albedo = PIL.Image.open(albedo_path) |
| 32 | |
| 33 | # metal_rough_path = os.path.join(os.path.dirname(path), "textures/metallicRoughness.png") |
| 34 | metal_rough_path = glob.glob(os.path.join(os.path.dirname(path), "textures/metallicRoughness.*"))[0] |
| 35 | if os.path.exists(metal_rough_path): |
| 36 | self.image_metallicRoughness = PIL.Image.open(metal_rough_path) |
| 37 | else: |
| 38 | # metallic_path = os.path.join(os.path.dirname(path), "textures/metallic.png") |
| 39 | metallic_path = glob.glob(os.path.join(os.path.dirname(path), "textures/metallic.*"))[0] |
| 40 | if os.path.exists(metallic_path): |
| 41 | image_metallic = PIL.Image.open(metallic_path) |
| 42 | else: |
| 43 | image_metallic = None |
| 44 | |
| 45 | # roughness_path = os.path.join(os.path.dirname(path), "textures/roughness.png") |
| 46 | roughness_path = glob.glob(os.path.join(os.path.dirname(path), "textures/roughness.*"))[0] |
| 47 | if os.path.exists(roughness_path): |
| 48 | image_roughness = PIL.Image.open(roughness_path) |
| 49 | else: |
| 50 | image_roughness = None |
| 51 | |
| 52 | # compose metallicRoughness |
| 53 | if image_metallic is None and image_roughness is None: |
| 54 | self.image_metallicRoughness = PIL.Image.fromarray(np.zeros_like(self.image_albedo, dtype=np.uint8)) |
| 55 | elif image_metallic is None: |
| 56 | image_roughness = np.asarray(image_roughness) |
| 57 | image_metallic = np.zeros_like(image_roughness, dtype=np.uint8) |
| 58 | self.image_metallicRoughness = PIL.Image.fromarray( |
| 59 | np.stack([image_metallic, image_roughness, np.zeros_like(image_roughness, dtype=np.uint8)], axis=-1) |
| 60 | ) |
| 61 | elif image_roughness is None: |
| 62 | image_metallic = np.asarray(image_metallic) |
| 63 | image_roughness = np.zeros_like(image_metallic, dtype=np.uint8) |
| 64 | self.image_metallicRoughness = PIL.Image.fromarray( |
| 65 | np.stack([image_metallic, image_roughness, np.zeros_like(image_metallic, dtype=np.uint8)], axis=-1) |
| 66 | ) |
| 67 | else: |
| 68 | image_metallic = np.asarray(image_metallic) |
| 69 | image_roughness = np.asarray(image_roughness) |
| 70 | self.image_metallicRoughness = PIL.Image.fromarray( |