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
