def create_scene(self, hws, alphas, N):
batch = []
for i in range(N):
scene = []
for mesh_name in self.models:
hw = hws[mesh_name]
alpha = alphas[mesh_name]
N, K, _ = hw.shape
for k in range(K):
c = self.colors[mesh_name].clone()
c[..., 3] = alpha[i, k]
textures = TexturesVertex(verts_features=[c])
m = Meshes(verts=[self.verts[mesh_name].clone()],
faces=[self.faces[mesh_name].clone()],
textures=textures)
t = Translate(y=hw[i, k, 0],
x=hw[i, k, 1],
z=torch.zeros(1, device=self.device),
device=str(self.device))
m = m.update_padded(t.transform_points(m.verts_padded()))
scene += [m]
batch += [join_meshes_as_scene(scene)]
batch = join_meshes_as_batch(batch)
return batch
def load_objs_as_meshes(files: list, device=None, load_textures: bool = True):
"""
Load meshes from a list of .obj files using the load_obj function, and
return them as a Meshes object. This only works for meshes which have a
single texture image for the whole mesh. See the load_obj function for more
details. material_colors and normals are not stored.
Args:
f: A list of file-like objects (with methods read, readline, tell,
and seek), pathlib paths or strings containing file names.
device: Desired device of returned Meshes. Default:
uses the current device for the default tensor type.
load_textures: Boolean indicating whether material files are loaded
Returns:
New Meshes object.
"""
mesh_list = []
for f_obj in files:
# TODO: update this function to support the two texturing options.
verts, faces, aux = load_obj(f_obj, load_textures=load_textures)
verts = verts.to(device)
tex = None
tex_maps = aux.texture_images
if tex_maps is not None and len(tex_maps) > 0:
verts_uvs = aux.verts_uvs[None, ...].to(device) # (1, V, 2)
faces_uvs = faces.textures_idx[None, ...].to(device) # (1, F, 3)
image = list(tex_maps.values())[0].to(device)[None]
tex = Textures(verts_uvs=verts_uvs, faces_uvs=faces_uvs, maps=image)
mesh = Meshes(verts=[verts], faces=[faces.verts_idx.to(device)], textures=tex)
mesh_list.append(mesh)
if len(mesh_list) == 1:
return mesh_list[0]
return join_meshes_as_batch(mesh_list)
def load_objs_as_meshes(
files: list,
device=None,
load_textures: bool = True,
create_texture_atlas: bool = False,
texture_atlas_size: int = 4,
texture_wrap: Optional[str] = "repeat",
path_manager: Optional[PathManager] = None,
):
"""
Load meshes from a list of .obj files using the load_obj function, and
return them as a Meshes object. This only works for meshes which have a
single texture image for the whole mesh. See the load_obj function for more
details. material_colors and normals are not stored.
Args:
files: A list of file-like objects (with methods read, readline, tell,
and seek), pathlib paths or strings containing file names.
device: Desired device of returned Meshes. Default:
uses the current device for the default tensor type.
load_textures: Boolean indicating whether material files are loaded
create_texture_atlas, texture_atlas_size, texture_wrap: as for load_obj.
path_manager: optionally a PathManager object to interpret paths.
Returns:
New Meshes object.
"""
mesh_list = []
for f_obj in files:
verts, faces, aux = load_obj(
f_obj,
load_textures=load_textures,
create_texture_atlas=create_texture_atlas,
texture_atlas_size=texture_atlas_size,
texture_wrap=texture_wrap,
path_manager=path_manager,
)
tex = None
if create_texture_atlas:
# TexturesAtlas type
tex = TexturesAtlas(atlas=[aux.texture_atlas.to(device)])
else:
# TexturesUV type
tex_maps = aux.texture_images
if tex_maps is not None and len(tex_maps) > 0:
verts_uvs = aux.verts_uvs.to(device) # (V, 2)
faces_uvs = faces.textures_idx.to(device) # (F, 3)
image = list(tex_maps.values())[0].to(device)[None]
tex = TexturesUV(verts_uvs=[verts_uvs],
faces_uvs=[faces_uvs],
maps=image)
mesh = Meshes(verts=[verts.to(device)],
faces=[faces.verts_idx.to(device)],
textures=tex)
mesh_list.append(mesh)
if len(mesh_list) == 1:
return mesh_list[0]
return join_meshes_as_batch(mesh_list)
def test_join_meshes_as_batch(self):
"""
Test that join_meshes_as_batch and load_objs_as_meshes are consistent
with single meshes.
"""
def check_triple(mesh, mesh3):
"""
Verify that mesh3 is three copies of mesh.
"""
def check_item(x, y):
self.assertEqual(x is None, y is None)
if x is not None:
self.assertClose(torch.cat([x, x, x]), y)
check_item(mesh.verts_padded(), mesh3.verts_padded())
check_item(mesh.faces_padded(), mesh3.faces_padded())
if mesh.textures is not None:
if isinstance(mesh.textures, TexturesUV):
check_item(
mesh.textures.faces_uvs_padded(),
mesh3.textures.faces_uvs_padded(),
)
check_item(
mesh.textures.verts_uvs_padded(),
mesh3.textures.verts_uvs_padded(),
)
check_item(mesh.textures.maps_padded(),
mesh3.textures.maps_padded())
elif isinstance(mesh.textures, TexturesVertex):
check_item(
mesh.textures.verts_features_padded(),
mesh3.textures.verts_features_padded(),
)
elif isinstance(mesh.textures, TexturesAtlas):
check_item(mesh.textures.atlas_padded(),
mesh3.textures.atlas_padded())
DATA_DIR = Path(
__file__).resolve().parent.parent / "docs/tutorials/data"
obj_filename = DATA_DIR / "cow_mesh/cow.obj"
mesh = load_objs_as_meshes([obj_filename])
mesh3 = load_objs_as_meshes([obj_filename, obj_filename, obj_filename])
check_triple(mesh, mesh3)
self.assertTupleEqual(mesh.textures.maps_padded().shape,
(1, 1024, 1024, 3))
# Try mismatched texture map sizes, which needs a call to interpolate()
mesh2048 = mesh.clone()
maps = mesh.textures.maps_padded()
mesh2048.textures._maps_padded = torch.cat([maps, maps], dim=1)
join_meshes_as_batch([mesh.to("cuda:0"), mesh2048.to("cuda:0")])
mesh_notex = load_objs_as_meshes([obj_filename], load_textures=False)
mesh3_notex = load_objs_as_meshes(
[obj_filename, obj_filename, obj_filename], load_textures=False)
check_triple(mesh_notex, mesh3_notex)
self.assertIsNone(mesh_notex.textures)
# meshes with vertex texture, join into a batch.
verts = torch.randn((4, 3), dtype=torch.float32)
faces = torch.tensor([[2, 1, 0], [3, 1, 0]], dtype=torch.int64)
vert_tex = torch.ones_like(verts)
rgb_tex = TexturesVertex(verts_features=[vert_tex])
mesh_rgb = Meshes(verts=[verts], faces=[faces], textures=rgb_tex)
mesh_rgb3 = join_meshes_as_batch([mesh_rgb, mesh_rgb, mesh_rgb])
check_triple(mesh_rgb, mesh_rgb3)
# meshes with texture atlas, join into a batch.
device = "cuda:0"
atlas = torch.rand((2, 4, 4, 3), dtype=torch.float32, device=device)
atlas_tex = TexturesAtlas(atlas=[atlas])
mesh_atlas = Meshes(verts=[verts], faces=[faces], textures=atlas_tex)
mesh_atlas3 = join_meshes_as_batch(
[mesh_atlas, mesh_atlas, mesh_atlas])
check_triple(mesh_atlas, mesh_atlas3)
# Test load multiple meshes with textures into a batch.
teapot_obj = DATA_DIR / "teapot.obj"
mesh_teapot = load_objs_as_meshes([teapot_obj])
teapot_verts, teapot_faces = mesh_teapot.get_mesh_verts_faces(0)
mix_mesh = load_objs_as_meshes([obj_filename, teapot_obj],
load_textures=False)
self.assertEqual(len(mix_mesh), 2)
self.assertClose(mix_mesh.verts_list()[0], mesh.verts_list()[0])
self.assertClose(mix_mesh.faces_list()[0], mesh.faces_list()[0])
self.assertClose(mix_mesh.verts_list()[1], teapot_verts)
self.assertClose(mix_mesh.faces_list()[1], teapot_faces)
cow3_tea = join_meshes_as_batch([mesh3, mesh_teapot],
include_textures=False)
self.assertEqual(len(cow3_tea), 4)
check_triple(mesh_notex, cow3_tea[:3])
self.assertClose(cow3_tea.verts_list()[3], mesh_teapot.verts_list()[0])
self.assertClose(cow3_tea.faces_list()[3], mesh_teapot.faces_list()[0])
# Check error raised if all meshes in the batch don't have the same texture type
with self.assertRaisesRegex(ValueError, "same type of texture"):
join_meshes_as_batch([mesh_atlas, mesh_rgb, mesh_atlas])
def save_pair_objects(
img_file1,
img_file2,
p_instances,
output_dir,
prefix="",
pred_camera=None,
plane_param_override=None,
show_camera=True,
corr_list=[],
webvis=False,
):
"""
if tran_topk == -2 and rot_topk == -2, then pred_camera should not be None, this is used for non-binned camera.
if exclude is not None, exclude some instances to make fig 2.
idx=7867
exclude = {
'0': [2,3,4,5,6,7],
'1': [0,1,2,4,5,6,7],
}
"""
image_paths = {"0": img_file1, "1": img_file2}
meshes_list = []
# map_files = []
uv_maps = []
cam_list = []
# get plane parameters
plane_locals = {}
for i in range(2):
if plane_param_override is None:
plane_locals[str(i)] = p_instances[str(i)].pred_planes
else:
plane_locals[str(i)] = plane_param_override[str(i)]
# get camera 1 to 2
camera1to2 = {
"position": np.array(pred_camera["position"]),
"rotation": quaternion.from_float_array(pred_camera["rotation"]),
}
# Merge planes if they are in correspondence
if len(corr_list) != 0:
plane_locals = merge_plane_params_from_local_params(
plane_locals, corr_list, camera1to2)
os.makedirs(output_dir, exist_ok=True)
for i in range(2):
if i == 0:
camera_info = camera1to2
else:
camera_info = {
"position": np.array([0, 0, 0]),
"rotation": np.quaternion(1, 0, 0, 0),
}
p_instance = p_instances[str(i)]
plane_params = plane_locals[str(i)]
segmentations = p_instance.pred_masks
meshes, uv_map = get_single_image_mesh_plane(
plane_params,
segmentations,
img_file=image_paths[str(i)],
height=480,
width=640,
webvis=False,
)
uv_maps.extend(uv_map)
meshes = transform_meshes(meshes, camera_info)
meshes_list.append(meshes)
cam_list.append(camera_info)
joint_mesh = join_meshes_as_batch(meshes_list)
if webvis:
joint_mesh = rotate_mesh_for_webview(joint_mesh)
# add camera into the mesh
if show_camera:
cam_meshes = get_camera_meshes(cam_list)
if webvis:
cam_meshes = rotate_mesh_for_webview(cam_meshes)
else:
cam_meshes = None
# save obj
if len(prefix) == 0:
prefix = "pred"
save_obj(
folder=output_dir,
prefix=prefix,
meshes=joint_mesh,
cam_meshes=cam_meshes,
decimal_places=10,
blend_flag=True,
map_files=None,
uv_maps=uv_maps,
)
def test_join_meshes_as_batch(self):
"""
Test that join_meshes_as_batch and load_objs_as_meshes are consistent
with single meshes.
"""
def check_triple(mesh, mesh3):
"""
Verify that mesh3 is three copies of mesh.
"""
def check_item(x, y):
self.assertEqual(x is None, y is None)
if x is not None:
self.assertClose(torch.cat([x, x, x]), y)
check_item(mesh.verts_padded(), mesh3.verts_padded())
check_item(mesh.faces_padded(), mesh3.faces_padded())
if mesh.textures is not None:
check_item(mesh.textures.maps_padded(),
mesh3.textures.maps_padded())
check_item(
mesh.textures.faces_uvs_padded(),
mesh3.textures.faces_uvs_padded(),
)
check_item(
mesh.textures.verts_uvs_padded(),
mesh3.textures.verts_uvs_padded(),
)
check_item(
mesh.textures.verts_rgb_padded(),
mesh3.textures.verts_rgb_padded(),
)
DATA_DIR = (Path(__file__).resolve().parent.parent /
'docs/tutorials/data')
obj_filename = DATA_DIR / 'cow_mesh/cow.obj'
mesh = load_objs_as_meshes([obj_filename])
mesh3 = load_objs_as_meshes([obj_filename, obj_filename, obj_filename])
check_triple(mesh, mesh3)
self.assertTupleEqual(mesh.textures.maps_padded().shape,
(1, 1024, 1024, 3))
# Try mismatched texture map sizes, which needs a call to interpolate()
mesh2048 = mesh.clone()
maps = mesh.textures.maps_padded()
mesh2048.textures._maps_padded = torch.cat([maps, maps], dim=1)
join_meshes_as_batch([mesh.to('cuda:0'), mesh2048.to('cuda:0')])
mesh_notex = load_objs_as_meshes([obj_filename], load_textures=False)
mesh3_notex = load_objs_as_meshes(
[obj_filename, obj_filename, obj_filename], load_textures=False)
check_triple(mesh_notex, mesh3_notex)
self.assertIsNone(mesh_notex.textures)
verts = torch.randn((4, 3), dtype=torch.float32)
faces = torch.tensor([[2, 1, 0], [3, 1, 0]], dtype=torch.int64)
vert_tex = torch.tensor([[0, 1, 0], [0, 1, 1], [1, 1, 0], [1, 1, 1]],
dtype=torch.float32)
tex = Textures(verts_rgb=vert_tex[None, :])
mesh_rgb = Meshes(verts=[verts], faces=[faces], textures=tex)
mesh_rgb3 = join_meshes_as_batch([mesh_rgb, mesh_rgb, mesh_rgb])
check_triple(mesh_rgb, mesh_rgb3)
teapot_obj = DATA_DIR / 'teapot.obj'
mesh_teapot = load_objs_as_meshes([teapot_obj])
teapot_verts, teapot_faces = mesh_teapot.get_mesh_verts_faces(0)
mix_mesh = load_objs_as_meshes([obj_filename, teapot_obj],
load_textures=False)
self.assertEqual(len(mix_mesh), 2)
self.assertClose(mix_mesh.verts_list()[0], mesh.verts_list()[0])
self.assertClose(mix_mesh.faces_list()[0], mesh.faces_list()[0])
self.assertClose(mix_mesh.verts_list()[1], teapot_verts)
self.assertClose(mix_mesh.faces_list()[1], teapot_faces)
cow3_tea = join_meshes_as_batch([mesh3, mesh_teapot],
include_textures=False)
self.assertEqual(len(cow3_tea), 4)
check_triple(mesh_notex, cow3_tea[:3])
self.assertClose(cow3_tea.verts_list()[3], mesh_teapot.verts_list()[0])
self.assertClose(cow3_tea.faces_list()[3], mesh_teapot.faces_list()[0])
def my_collate(batch):
meshes, targets = zip(*batch)
meshes = join_meshes_as_batch(meshes, include_textures=True)
targets = torch.tensor(targets)
return [meshes, targets]
def batch(self):
return join_meshes_as_batch([scene.mesh for scene in self.scenes])
def load_objs_as_meshes(
files: list,
device=None,
load_textures: bool = True,
create_texture_atlas: bool = False,
texture_atlas_size: int = 4,
texture_wrap: Optional[str] = "repeat",
path_manager: Optional[PathManager] = None,
):
"""
Load meshes from a list of .obj files using the load_obj function, and
return them as a Meshes object. This only works for meshes which have a
single texture image for the whole mesh. See the load_obj function for more
details. material_colors and normals are not stored.
Args:
files: A list of file-like objects (with methods read, readline, tell,
and seek), pathlib paths or strings containing file names.
device: Desired device of returned Meshes. Default:
uses the current device for the default tensor type.
load_textures: Boolean indicating whether material files are loaded
create_texture_atlas, texture_atlas_size, texture_wrap: as for load_obj.
path_manager: optionally a PathManager object to interpret paths.
Returns:
New Meshes object.
"""
mesh_list = []
for f_obj in files:
verts, faces, aux = load_obj(
f_obj,
load_textures=load_textures,
create_texture_atlas=create_texture_atlas,
texture_atlas_size=texture_atlas_size,
texture_wrap=texture_wrap,
path_manager=path_manager,
)
tex = None
if create_texture_atlas:
# TexturesAtlas type
tex = TexturesAtlas(atlas=[aux.texture_atlas.to(device)])
else:
# TexturesUV type
tex_maps = aux.texture_images
textures = []
if tex_maps is not None and len(tex_maps) > 0:
verts_uvs = aux.verts_uvs.to(device) # (V, 2)
faces_uvs = faces.textures_idx.to(device) # (F, 3)
face_to_mat = faces.materials_idx
# code for checking
current_offset = 0
for mat_idx, (mat_name,
tex_map) in enumerate(tex_maps.items()):
image = tex_map.flip(0).unsqueeze(0).to(device)
faces_mask = face_to_mat == mat_idx
faces_verts_uvs = faces_uvs[faces_mask].unique()
tex_verts_uvs = verts_uvs[faces_verts_uvs]
tex_faces_uvs = faces_uvs[faces_mask] - current_offset
tex = TexturesUV(verts_uvs=[tex_verts_uvs],
faces_uvs=[tex_faces_uvs],
maps=image,
mat_names=[mat_name])
textures.append(tex)
current_offset += tex_verts_uvs.shape[0]
tex = textures[0]
if len(textures) > 1:
tex = tex.join_batch(textures[1:]).join_scene()
mesh = Meshes(verts=[verts.to(device)],
faces=[faces.verts_idx.to(device)],
textures=tex,
aux=aux)
mesh_list.append(mesh)
if len(mesh_list) == 1:
return mesh_list[0]
return join_meshes_as_batch(mesh_list)
def my_collate(batch):
## load unregular mesh within a batch
meshes, label = zip(*batch)
meshes = join_meshes_as_batch(meshes, include_textures=False)
label = torch.tensor(label)
return [meshes, label]
