def test_load_obj_error_invalid_texture_indices(self): obj_file = "\n".join( [ "v 0.1 0.2 0.3", "v 0.1 0.2 0.3", "v 0.1 0.2 0.3", "vt 0.1 0.2", "vt 0.1 0.2", "vt 0.1 0.2", "f -2//2 2//6 1//1", ] ) obj_file = StringIO(obj_file) with self.assertRaises(ValueError) as err: load_obj(obj_file) self.assertTrue("Faces have invalid indices." in str(err.exception))
def test_load_mtl_noload(self): DATA_DIR = Path(__file__).resolve().parent.parent / "docs/tutorials/data" obj_filename = "cow_mesh/cow.obj" filename = os.path.join(DATA_DIR, obj_filename) verts, faces, aux = load_obj(filename, load_textures=False) self.assertTrue(aux.material_colors is None) self.assertTrue(aux.texture_images is None)
def test_load_obj_error_invalid_normal_indices(self): obj_file = "\n".join([ "v 0.1 0.2 0.3", "v 0.1 0.2 0.3", "v 0.1 0.2 0.3", "vn 0.1 0.2 0.3", "vn 0.1 0.2 0.3", "vn 0.1 0.2 0.3", "f -2/2 2/4 1/1", ]) with NamedTemporaryFile(mode="w", suffix=".obj") as f: f.write(obj_file) f.flush() with self.assertWarnsRegex(UserWarning, "Faces have invalid indices"): load_obj(Path(f.name))
def __init__(self, image_size, obj_filename, uv_size=256): super(ImageRenderer, self).__init__() self.image_size = image_size self.uv_size = uv_size verts, faces, aux = load_obj(obj_filename) faces = faces.verts_idx[None, ...] self.rasterizer = Pytorch3dRasterizer(image_size) # faces self.register_buffer('faces', faces)
def test_load_no_usemtl(self): obj_filename = "missing_usemtl/cow.obj" # obj_filename has no "usemtl material_1" line filename = os.path.join(DATA_DIR, obj_filename) # TexturesUV type mesh = IO().load_mesh(filename) self.assertIsNotNone(mesh.textures) verts, faces, aux = load_obj(filename) self.assertTrue("material_1" in aux.material_colors) self.assertTrue("material_1" in aux.texture_images)
def load_untextured_mesh(mesh_path, device): mesh = load_objs_as_meshes([mesh_path], device=device, load_textures = False) verts, faces_idx, _ = load_obj(mesh_path) faces = faces_idx.verts_idx verts_rgb = torch.ones_like(verts)[None] # (1, V, 3) textures = Textures(verts_rgb=verts_rgb.to(device)) mesh_no_texture = Meshes( verts=[verts.to(device)], faces=[faces.to(device)], textures=textures ) return mesh_no_texture
def _test_save_load(self, verts, faces): f = StringIO() save_obj(f, verts, faces) f.seek(0) expected_verts, expected_faces = verts, faces if not len(expected_verts): # Always compare with a (V, 3) tensor expected_verts = torch.zeros(size=(0, 3), dtype=torch.float32) if not len(expected_faces): # Always compare with an (F, 3) tensor expected_faces = torch.zeros(size=(0, 3), dtype=torch.int64) actual_verts, actual_faces, _ = load_obj(f) self.assertClose(expected_verts, actual_verts) self.assertClose(expected_faces, actual_faces.verts_idx)
def load_mesh_from_obj(file, device): verts, faces_idx, _ = load_obj(file) faces = faces_idx.verts_idx # Initialize each vertex to be white in color. verts_rgb = torch.ones_like(verts)[None] # (1, V, 3) textures = Textures(verts_rgb=verts_rgb.to(device)) # Create a Meshes object for the teapot. Here we have only one mesh in the batch. teapot_mesh = Meshes(verts=[verts.to(device)], faces=[faces.to(device)], textures=textures) return teapot_mesh
def test_load_obj_missing_mtl(self): obj_filename = "missing_files_obj/model2.obj" filename = os.path.join(DATA_DIR, obj_filename) with self.assertWarnsRegex(UserWarning, "Mtl file does not exist"): verts, faces, aux = load_obj(filename) expected_verts = torch.tensor( [[0.1, 0.2, 0.3], [0.2, 0.3, 0.4], [0.3, 0.4, 0.5], [0.4, 0.5, 0.6]], dtype=torch.float32, ) expected_faces = torch.tensor([[0, 1, 2], [0, 1, 3]], dtype=torch.int64) self.assertTrue(torch.allclose(verts, expected_verts)) self.assertTrue(torch.allclose(faces.verts_idx, expected_faces))
def test_load_obj_missing_mtl_noload(self): obj_filename = "missing_files_obj/model2.obj" filename = os.path.join(DATA_DIR, obj_filename) verts, faces, aux = load_obj(filename, load_textures=False) expected_verts = torch.tensor( [[0.1, 0.2, 0.3], [0.2, 0.3, 0.4], [0.3, 0.4, 0.5], [0.4, 0.5, 0.6]], dtype=torch.float32, ) expected_faces = torch.tensor([[0, 1, 2], [0, 1, 3]], dtype=torch.int64) self.assertTrue(torch.allclose(verts, expected_verts)) self.assertTrue(torch.allclose(faces.verts_idx, expected_faces)) self.assertTrue(aux.material_colors is None) self.assertTrue(aux.texture_images is None)
def _test_save_load(self, verts, faces): with NamedTemporaryFile(mode="w", suffix=".obj") as f: file_path = Path(f.name) save_obj(file_path, verts, faces) f.flush() expected_verts, expected_faces = verts, faces if not len(expected_verts): # Always compare with a (V, 3) tensor expected_verts = torch.zeros(size=(0, 3), dtype=torch.float32) if not len(expected_faces): # Always compare with an (F, 3) tensor expected_faces = torch.zeros(size=(0, 3), dtype=torch.int64) actual_verts, actual_faces, _ = load_obj(file_path) self.assertClose(expected_verts, actual_verts) self.assertClose(expected_faces, actual_faces.verts_idx)
def __init__(self, image_size, obj_filename, uv_size=256, rasterizer_type='pytorch3d'): super(SRenderY, self).__init__() self.image_size = image_size self.uv_size = uv_size verts, faces, aux = load_obj(obj_filename) uvcoords = aux.verts_uvs[None, ...] # (N, V, 2) uvfaces = faces.textures_idx[None, ...] # (N, F, 3) faces = faces.verts_idx[None, ...] if rasterizer_type == 'pytorch3d': self.rasterizer = Pytorch3dRasterizer(image_size) self.uv_rasterizer = Pytorch3dRasterizer(uv_size) # faces dense_triangles = util.generate_triangles(uv_size, uv_size) self.register_buffer( 'dense_faces', torch.from_numpy(dense_triangles).long()[None, :, :]) self.register_buffer('faces', faces) self.register_buffer('raw_uvcoords', uvcoords) # uv coords uvcoords = torch.cat([uvcoords, uvcoords[:, :, 0:1] * 0. + 1.], -1) #[bz, ntv, 3] uvcoords = uvcoords * 2 - 1 uvcoords[..., 1] = -uvcoords[..., 1] face_uvcoords = util.face_vertices(uvcoords, uvfaces) self.register_buffer('uvcoords', uvcoords) self.register_buffer('uvfaces', uvfaces) self.register_buffer('face_uvcoords', face_uvcoords) # shape colors, for rendering shape overlay colors = torch.tensor([180, 180, 180])[None, None, :].repeat( 1, faces.max() + 1, 1).float() / 255. face_colors = util.face_vertices(colors, faces) self.register_buffer('face_colors', face_colors) ## SH factors for lighting pi = np.pi constant_factor = torch.tensor([1/np.sqrt(4*pi), ((2*pi)/3)*(np.sqrt(3/(4*pi))), ((2*pi)/3)*(np.sqrt(3/(4*pi))),\ ((2*pi)/3)*(np.sqrt(3/(4*pi))), (pi/4)*(3)*(np.sqrt(5/(12*pi))), (pi/4)*(3)*(np.sqrt(5/(12*pi))),\ (pi/4)*(3)*(np.sqrt(5/(12*pi))), (pi/4)*(3/2)*(np.sqrt(5/(12*pi))), (pi/4)*(1/2)*(np.sqrt(5/(4*pi)))]).float() self.register_buffer('constant_factor', constant_factor)
def load_unique_meshes(json_file, model_root): with open(json_file, "r") as f: anns = json.load(f)["annotations"] # find unique models unique_models = [] for obj in anns: model_type = obj["model"] if model_type not in unique_models: unique_models.append(model_type) # read unique models object_models = {} for model in unique_models: mesh = load_obj(os.path.join(model_root, model)) object_models[model] = [mesh[0], mesh[1].verts_idx] return object_models
def test_load_obj_simple(self): obj_file = "\n".join( [ "# this is a comment", # Comments should be ignored. "v 0.1 0.2 0.3", "v 0.2 0.3 0.4", "v 0.3 0.4 0.5", "v 0.4 0.5 0.6", # some obj files have multiple spaces after v "f 1 2 3", "f 1 2 4 3 1", # Polygons should be split into triangles ] ) obj_file = StringIO(obj_file) verts, faces, aux = load_obj(obj_file) normals = aux.normals textures = aux.verts_uvs materials = aux.material_colors tex_maps = aux.texture_images expected_verts = torch.tensor( [ [0.1, 0.2, 0.3], [0.2, 0.3, 0.4], [0.3, 0.4, 0.5], [0.4, 0.5, 0.6], ], dtype=torch.float32, ) expected_faces = torch.tensor( [ [0, 1, 2], # First face [0, 1, 3], # Second face (polygon) [0, 3, 2], # Second face (polygon) [0, 2, 0], # Second face (polygon) ], dtype=torch.int64, ) self.assertTrue(torch.all(verts == expected_verts)) self.assertTrue(torch.all(faces.verts_idx == expected_faces)) self.assertTrue(faces.normals_idx == []) self.assertTrue(faces.textures_idx == []) self.assertTrue( torch.all(faces.materials_idx == -torch.ones(len(expected_faces))) ) self.assertTrue(normals is None) self.assertTrue(textures is None) self.assertTrue(materials is None) self.assertTrue(tex_maps is None)
def __init__(self, image_size, obj_filename, uv_size=256, config=None): super(Renderer, self).__init__() self.image_size = image_size self.uv_size = uv_size verts, faces, aux = load_obj(obj_filename) uvcoords = aux.verts_uvs[None, ...] # (N, V, 2) uvfaces = faces.textures_idx[None, ...] # (N, F, 3) faces = faces.verts_idx[None, ...] if config is not None: _tp = np.load(config.trim_path, allow_pickle=True) tp = {} for k in _tp.keys(): tp[k] = torch.LongTensor(_tp[k]) verts = verts[tp['idx_verts']] #uvcoords = uvcoords[:,tp['idx_verts']] uvfaces = uvfaces[:,tp['idx_faces']] faces = tp['map_verts'][faces[:,tp['idx_faces']]] self.rasterizer = Pytorch3dRasterizer(image_size) self.uv_rasterizer = Pytorch3dRasterizer(uv_size) # faces self.register_buffer('faces', faces) self.register_buffer('raw_uvcoords', uvcoords) # uv coordsw uvcoords = torch.cat([uvcoords, uvcoords[:, :, 0:1] * 0. + 1.], -1) # [bz, ntv, 3] uvcoords = uvcoords * 2 - 1 uvcoords[..., 1] = -uvcoords[..., 1] face_uvcoords = util.face_vertices(uvcoords, uvfaces) self.register_buffer('uvcoords', uvcoords) self.register_buffer('uvfaces', uvfaces) self.register_buffer('face_uvcoords', face_uvcoords) # shape colors colors = torch.tensor([74, 120, 168])[None, None, :].repeat(1, faces.max() + 1, 1).float() / 255. face_colors = util.face_vertices(colors, faces) self.register_buffer('face_colors', face_colors) ## lighting pi = np.pi constant_factor = torch.tensor( [1 / np.sqrt(4 * pi), ((2 * pi) / 3) * (np.sqrt(3 / (4 * pi))), ((2 * pi) / 3) * (np.sqrt(3 / (4 * pi))), \ ((2 * pi) / 3) * (np.sqrt(3 / (4 * pi))), (pi / 4) * (3) * (np.sqrt(5 / (12 * pi))), (pi / 4) * (3) * (np.sqrt(5 / (12 * pi))), \ (pi / 4) * (3) * (np.sqrt(5 / (12 * pi))), (pi / 4) * (3 / 2) * (np.sqrt(5 / (12 * pi))), (pi / 4) * (1 / 2) * (np.sqrt(5 / (4 * pi)))]) self.register_buffer('constant_factor', constant_factor)
def load_models(anns, model_root): # find unique models unique_models = [] for anno in anns: for obj in anno: model_type = obj["model"] if model_type not in unique_models: unique_models.append(model_type) # read unique models object_models = {} logger.info("Unique objects {}".format(len(unique_models))) for model in unique_models: mesh = load_obj(os.path.join(model_root, model)) object_models[model] = [mesh[0], mesh[1].verts_idx] logger.info("Done loading models") return object_models
def load_mesh(filename): device = torch.device("cuda:0") trg_obj = os.path.join(filename) verts, faces, aux = load_obj(trg_obj) faces_idx = faces.verts_idx.to(device) verts = verts.to(device) center = verts.mean(0) verts = verts - center scale = max(verts.abs().max(0)[0]) verts = verts / scale return Meshes(verts=[verts], faces=[faces_idx])
def test_load_obj_normals_only(self): obj_file = "\n".join( [ "v 0.1 0.2 0.3", "v 0.2 0.3 0.4", "v 0.3 0.4 0.5", "v 0.4 0.5 0.6", "vn 0.000000 0.000000 -1.000000", "vn -1.000000 -0.000000 -0.000000", "f 2//1 3//1 4//2", ] ) obj_file = StringIO(obj_file) expected_faces_normals_idx = torch.tensor( [[0, 0, 1]], dtype=torch.int64 ) expected_normals = torch.tensor( [ [0.000000, 0.000000, -1.000000], [-1.000000, -0.000000, -0.000000], ], dtype=torch.float32, ) expected_verts = torch.tensor( [ [0.1, 0.2, 0.3], [0.2, 0.3, 0.4], [0.3, 0.4, 0.5], [0.4, 0.5, 0.6], ], dtype=torch.float32, ) verts, faces, aux = load_obj(obj_file) normals = aux.normals textures = aux.verts_uvs materials = aux.material_colors tex_maps = aux.texture_images self.assertTrue( torch.allclose(faces.normals_idx, expected_faces_normals_idx) ) self.assertTrue(torch.allclose(normals, expected_normals)) self.assertTrue(torch.allclose(verts, expected_verts)) self.assertTrue(faces.textures_idx == []) self.assertTrue(textures is None) self.assertTrue(materials is None) self.assertTrue(tex_maps is None)
def load_mesh(obj_path): device = torch.device(config.cuda.device) if not os.path.exists(obj_path): raise FileNotFoundError('Cannot find moon object from \'%s\'' % obj_path) vertices, faces, aux = load_obj(obj_path) textures = load_texture(aux, faces) vertices = vertices.to(device) faces = faces.verts_idx.to(device) mesh = Meshes(verts=[vertices], faces=[faces], textures=textures) return mesh
def test_load_obj_simple(self): obj_file = "\n".join([ "# this is a comment", # Comments should be ignored. "v 0.1 0.2 0.3", "v 0.2 0.3 0.4", "v 0.3 0.4 0.5", "v 0.4 0.5 0.6", # some obj files have multiple spaces after v "f 1 2 3", "f 1 2 4 3 1", # Polygons should be split into triangles ]) with NamedTemporaryFile(mode="w", suffix=".obj") as f: f.write(obj_file) f.flush() verts, faces, aux = load_obj(Path(f.name)) normals = aux.normals textures = aux.verts_uvs materials = aux.material_colors tex_maps = aux.texture_images expected_verts = torch.tensor( [[0.1, 0.2, 0.3], [0.2, 0.3, 0.4], [0.3, 0.4, 0.5], [0.4, 0.5, 0.6]], dtype=torch.float32, ) expected_faces = torch.tensor( [ [0, 1, 2], # First face [0, 1, 3], # Second face (polygon) [0, 3, 2], # Second face (polygon) [0, 2, 0], # Second face (polygon) ], dtype=torch.int64, ) self.assertTrue(torch.all(verts == expected_verts)) self.assertTrue(torch.all(faces.verts_idx == expected_faces)) padded_vals = -(torch.ones_like(faces.verts_idx)) self.assertTrue(torch.all(faces.normals_idx == padded_vals)) self.assertTrue(torch.all(faces.textures_idx == padded_vals)) self.assertTrue( torch.all( faces.materials_idx == -(torch.ones(len(expected_faces))))) self.assertTrue(normals is None) self.assertTrue(textures is None) self.assertTrue(materials is None) self.assertTrue(tex_maps is None)
def create_mesh(self, pdb_code=None, pdb_file=None, out_dir=None): """ Creates a PyTorch3D Mesh from an .Obj file. pdb_code and pdb_file are optional arguments. Use one as suits your purposes :param pdb_code: 4-character PDB accession code :type pdb_code: str :param pdb_file: Path to local .PDB file :type pdb_file: str :param out_dir: Path to output directory :type out_dir: str :return: verts, faces, aux """ obj_file = self.get_obj_file(pdb_code=pdb_code, pdb_file=pdb_file, out_dir=out_dir) verts, faces, aux = load_obj(obj_file) return verts, faces, aux
def load_moon_mesh(_obj_filename): # Load the object verts, faces, aux = load_obj(_obj_filename) faces_idx = faces.verts_idx.to(DEVICE) verts = verts.to(DEVICE) verts_uvs = aux.verts_uvs[None, ...].to(DEVICE) faces_uvs = faces.textures_idx[None, ...].to(DEVICE) tex_maps = aux.texture_images texture_image = list(tex_maps.values())[0] texture_image = texture_image[None, ...].to(DEVICE) tex = Textures(verts_uvs=verts_uvs, faces_uvs=faces_uvs, maps=texture_image) moon_mesh = Meshes(verts=[verts], faces=[faces_idx], textures=tex) return moon_mesh
def set_moon_mesh(self): vertices, faces, aux = load_obj(OBJECT_PATH) faces_idx = faces.verts_idx.to(self.cuda_device) vertices = vertices.to(self.cuda_device) verts_uvs = aux.verts_uvs[None, ...].to(self.cuda_device) faces_uvs = faces.textures_idx[None, ...].to(self.cuda_device) tex_maps = aux.texture_images texture_image = list(tex_maps.values())[0] texture_image = texture_image[None, ...].to(self.cuda_device) tex = Textures(verts_uvs=verts_uvs, faces_uvs=faces_uvs, maps=texture_image) self.moon_mesh = Meshes(verts=[vertices], faces=[faces_idx], textures=tex)
def test_load_obj_textures_only(self): obj_file = "\n".join( [ "v 0.1 0.2 0.3", "v 0.2 0.3 0.4", "v 0.3 0.4 0.5", "v 0.4 0.5 0.6", "vt 0.999110 0.501077", "vt 0.999455 0.750380", "f 2/1 3/1 4/2", ] ) obj_file = StringIO(obj_file) expected_faces_textures_idx = torch.tensor( [[0, 0, 1]], dtype=torch.int64 ) expected_textures = torch.tensor( [[0.999110, 0.501077], [0.999455, 0.750380]], dtype=torch.float32 ) expected_verts = torch.tensor( [ [0.1, 0.2, 0.3], [0.2, 0.3, 0.4], [0.3, 0.4, 0.5], [0.4, 0.5, 0.6], ], dtype=torch.float32, ) verts, faces, aux = load_obj(obj_file) normals = aux.normals textures = aux.verts_uvs materials = aux.material_colors tex_maps = aux.texture_images self.assertTrue( torch.allclose(faces.textures_idx, expected_faces_textures_idx) ) self.assertTrue(torch.allclose(expected_textures, textures)) self.assertTrue(torch.allclose(expected_verts, verts)) self.assertTrue(faces.normals_idx == []) self.assertTrue(normals is None) self.assertTrue(materials is None) self.assertTrue(tex_maps is None)
def test_load_obj_normals_only(self): obj_file = '\n'.join([ 'v 0.1 0.2 0.3', 'v 0.2 0.3 0.4', 'v 0.3 0.4 0.5', 'v 0.4 0.5 0.6', 'vn 0.000000 0.000000 -1.000000', 'vn -1.000000 -0.000000 -0.000000', 'f 2//1 3//1 4//2', ]) obj_file = StringIO(obj_file) expected_faces_normals_idx = torch.tensor([[0, 0, 1]], dtype=torch.int64) expected_normals = torch.tensor( [ [0.000000, 0.000000, -1.000000], [-1.000000, -0.000000, -0.000000], ], dtype=torch.float32, ) expected_verts = torch.tensor( [ [0.1, 0.2, 0.3], [0.2, 0.3, 0.4], [0.3, 0.4, 0.5], [0.4, 0.5, 0.6], ], dtype=torch.float32, ) verts, faces, aux = load_obj(obj_file) normals = aux.normals textures = aux.verts_uvs materials = aux.material_colors tex_maps = aux.texture_images self.assertClose(faces.normals_idx, expected_faces_normals_idx) self.assertClose(normals, expected_normals) self.assertClose(verts, expected_verts) # Textures idx padded with -1. self.assertClose(faces.textures_idx, torch.ones_like(faces.verts_idx) * -1) self.assertTrue(textures is None) self.assertTrue(materials is None) self.assertTrue(tex_maps is None)
def test_load_obj_missing_texture(self): DATA_DIR = Path(__file__).resolve().parent / "data" obj_filename = "missing_files_obj/model.obj" filename = os.path.join(DATA_DIR, obj_filename) with self.assertWarnsRegex(Warning, "Texture file does not exist"): verts, faces, aux = load_obj(filename) expected_verts = torch.tensor( [ [0.1, 0.2, 0.3], [0.2, 0.3, 0.4], [0.3, 0.4, 0.5], [0.4, 0.5, 0.6], ], dtype=torch.float32, ) expected_faces = torch.tensor([[0, 1, 2], [0, 1, 3]], dtype=torch.int64) self.assertTrue(torch.allclose(verts, expected_verts)) self.assertTrue(torch.allclose(faces.verts_idx, expected_faces))
def test_load_obj_normals_only(self): obj_file = "\n".join([ "v 0.1 0.2 0.3", "v 0.2 0.3 0.4", "v 0.3 0.4 0.5", "v 0.4 0.5 0.6", "vn 0.000000 0.000000 -1.000000", "vn -1.000000 -0.000000 -0.000000", "f 2//1 3//1 4//2", ]) expected_faces_normals_idx = torch.tensor([[0, 0, 1]], dtype=torch.int64) expected_normals = torch.tensor( [[0.000000, 0.000000, -1.000000], [-1.000000, -0.000000, -0.000000]], dtype=torch.float32, ) expected_verts = torch.tensor( [[0.1, 0.2, 0.3], [0.2, 0.3, 0.4], [0.3, 0.4, 0.5], [0.4, 0.5, 0.6]], dtype=torch.float32, ) with NamedTemporaryFile(mode="w", suffix=".obj") as f: f.write(obj_file) f.flush() verts, faces, aux = load_obj(Path(f.name)) normals = aux.normals textures = aux.verts_uvs materials = aux.material_colors tex_maps = aux.texture_images self.assertClose(faces.normals_idx, expected_faces_normals_idx) self.assertClose(normals, expected_normals) self.assertClose(verts, expected_verts) # Textures idx padded with -1. self.assertClose(faces.textures_idx, torch.ones_like(faces.verts_idx) * -1) self.assertTrue(textures is None) self.assertTrue(materials is None) self.assertTrue(tex_maps is None)
def __init__(self, image_size, obj_filename, uv_size=256): super(Renderer, self).__init__() self.image_size = image_size self.uv_size = uv_size verts, faces, aux = load_obj(obj_filename) uvcoords = aux.verts_uvs[None, ...] # (N, V, 2) uvfaces = faces.textures_idx[None, ...] # (N, F, 3) faces = faces.verts_idx[None, ...] self.rasterizer = Pytorch3dRasterizer(image_size) self.uv_rasterizer = Pytorch3dRasterizer(uv_size) # faces self.register_buffer('faces', faces) self.register_buffer('raw_uvcoords', uvcoords) # uv coordsw uvcoords = torch.cat([uvcoords, uvcoords[:, :, 0:1] * 0. + 1.], -1) # [bz, ntv, 3] uvcoords = uvcoords * 2 - 1 uvcoords[..., 1] = -uvcoords[..., 1] face_uvcoords = util.face_vertices(uvcoords, uvfaces) self.register_buffer('uvcoords', uvcoords) self.register_buffer('uvfaces', uvfaces) self.register_buffer('face_uvcoords', face_uvcoords) # shape colors colors = torch.tensor([74, 120, 168])[None, None, :].repeat( 1, faces.max() + 1, 1).float() / 255. face_colors = util.face_vertices(colors, faces) self.register_buffer('face_colors', face_colors) ## lighting pi = np.pi constant_factor = torch.tensor( [1 / np.sqrt(4 * pi), ((2 * pi) / 3) * (np.sqrt(3 / (4 * pi))), ((2 * pi) / 3) * (np.sqrt(3 / (4 * pi))), \ ((2 * pi) / 3) * (np.sqrt(3 / (4 * pi))), (pi / 4) * (3) * (np.sqrt(5 / (12 * pi))), (pi / 4) * (3) * (np.sqrt(5 / (12 * pi))), \ (pi / 4) * (3) * (np.sqrt(5 / (12 * pi))), (pi / 4) * (3 / 2) * (np.sqrt(5 / (12 * pi))), (pi / 4) * (1 / 2) * (np.sqrt(5 / (4 * pi)))]) self.register_buffer('constant_factor', constant_factor)
def __getitem__(self, idx): """ Read a model by the given index. Returns: dictionary with following keys: - verts: FloatTensor of shape (V, 3). - faces: LongTensor of shape (F, 3) which indexes into the verts tensor. - synset_id (str): synset id - model_id (str): model id - label (str): synset label. """ model = {} model['synset_id'] = self.synset_ids[idx] model['model_id'] = self.model_ids[idx] model_path = path.join(self.data_dir, model['synset_id'], model['model_id'], self.model_dir) model['verts'], faces, _ = load_obj(model_path) model['faces'] = faces.verts_idx model['label'] = self.synset_dict[model['synset_id']] return model
def _load_mesh(self, model_path) -> Tuple: verts, faces, aux = load_obj( model_path, create_texture_atlas=self.load_textures, load_textures=self.load_textures, texture_atlas_size=self.texture_resolution, ) if self.load_textures: textures = aux.texture_atlas # Some meshes don't have textures. In this case # create a white texture map if textures is None: textures = verts.new_ones( faces.verts_idx.shape[0], self.texture_resolution, self.texture_resolution, 3, ) return verts, faces.verts_idx, textures