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]