def _get_model_points(self, dataset_name): """convert to label based keys.""" if dataset_name in self.model_points: return self.model_points[dataset_name] dset_meta = MetadataCatalog.get(dataset_name) ref_key = dset_meta.ref_key data_ref = ref.__dict__[ref_key] objs = dset_meta.objs cfg = self.cfg cur_model_points = {} num = np.inf for i, obj_name in enumerate(objs): obj_id = data_ref.obj2id[obj_name] model_path = osp.join(data_ref.model_dir, f"obj_{obj_id:06d}.ply") model = inout.load_ply(model_path, vertex_scale=data_ref.vertex_scale) cur_model_points[i] = pts = model["pts"] if pts.shape[0] < num: num = pts.shape[0] num = min(num, cfg.MODEL.CDPN.PNP_NET.NUM_PM_POINTS) for i in range(len(cur_model_points)): keep_idx = np.arange(num) np.random.shuffle(keep_idx) # random sampling cur_model_points[i] = cur_model_points[i][keep_idx, :] self.model_points[dataset_name] = cur_model_points return self.model_points[dataset_name]
def load_meshes_sixd(obj_files, vertex_tmp_store_folder, recalculate_normals=False): hashed_file_name = ( hashlib.md5(("".join(obj_files) + "load_meshes_sixd" + str(recalculate_normals)).encode("utf-8")).hexdigest() + ".npy" ) out_file = os.path.join(vertex_tmp_store_folder, hashed_file_name) if os.path.exists(out_file): return np.load(out_file, allow_pickle=True) else: attributes = [] for model_path in tqdm(obj_files): model = inout.load_ply(model_path) vertices = np.array(model["pts"]).astype(np.float32) if recalculate_normals: normals = calc_normals(vertices) else: normals = np.array(model["normals"]).astype(np.float32) faces = np.array(model["faces"]).astype(np.uint32) if "colors" in model: colors = np.array(model["colors"]).astype(np.uint32) attributes.append((vertices, normals, colors, faces)) else: attributes.append((vertices, normals, faces)) os.makedirs(vertex_tmp_store_folder, exist_ok=True) np.save(out_file, attributes) return attributes
def models(self): """Load models into a list.""" cache_path = osp.join(self.cache_dir, "models_{}.pkl".format("_".join(self.objs))) if osp.exists(cache_path) and self.use_cache: # dprint("{}: load cached object models from {}".format(self.name, cache_path)) return mmcv.load(cache_path) models = [] for obj_name in self.objs: model = inout.load_ply( osp.join(self.models_root, f"obj_{ref.lm_full.obj2id[obj_name]:06d}.ply"), vertex_scale=self.scale_to_meter, ) # NOTE: the bbox3d_and_center is not obtained from centered vertices # for BOP models, not a big problem since they had been centered model["bbox3d_and_center"] = misc.get_bbox3d_and_center( model["pts"]) models.append(model) logger.info("cache models to {}".format(cache_path)) mmcv.mkdir_or_exist(osp.dirname(cache_path)) mmcv.dump(models, cache_path, protocol=4) return models
def _get_extents(self, dataset_name): """label based keys.""" if dataset_name in self.extents: return self.extents[dataset_name] dset_meta = MetadataCatalog.get(dataset_name) try: ref_key = dset_meta.ref_key except: # FIXME: for some reason, in distributed training, this need to be re-registered register_datasets([dataset_name]) dset_meta = MetadataCatalog.get(dataset_name) ref_key = dset_meta.ref_key data_ref = ref.__dict__[ref_key] objs = dset_meta.objs cfg = self.cfg cur_extents = {} for i, obj_name in enumerate(objs): obj_id = data_ref.obj2id[obj_name] model_path = osp.join(data_ref.model_dir, f"obj_{obj_id:06d}.ply") model = inout.load_ply(model_path, vertex_scale=data_ref.vertex_scale) pts = model["pts"] xmin, xmax = np.amin(pts[:, 0]), np.amax(pts[:, 0]) ymin, ymax = np.amin(pts[:, 1]), np.amax(pts[:, 1]) zmin, zmax = np.amin(pts[:, 2]), np.amax(pts[:, 2]) size_x = xmax - xmin size_y = ymax - ymin size_z = zmax - zmin cur_extents[i] = np.array([size_x, size_y, size_z], dtype="float32") self.extents[dataset_name] = cur_extents return self.extents[dataset_name]
def calc_texture_uv_emb_proj(uv_model_path_or_model, R, T, K, height=480, width=640): """calculate uv map emb via projection it seems to be better not to use expand the models are generated by blender, where texture_u, texture_v are provided """ if isinstance(uv_model_path_or_model, str): model = load_ply(uv_model_path_or_model) else: model = uv_model_path_or_model points = model["pts"] uv_gb = model["texture_uv"] points_2d, z = points_to_2D(points, R, T, K) image_points = np.round(points_2d).astype(np.int32) # image_points = (points_2d + 0.5).astype(np.int32) uv_ProjEmb = np.zeros((height, width, 2)).astype(np.float32) depth = np.zeros((height, width, 1)).astype(np.float32) for i, (x, y) in enumerate(image_points): if x >= width or y >= height or x < 0 or y < 0: continue if depth[y, x, 0] == 0: depth[y, x, 0] = z[i] uv_ProjEmb[y, x] = uv_gb[i] elif z[i] < depth[y, x, 0]: depth[y, x, 0] = z[i] uv_ProjEmb[y, x] = uv_gb[i] else: pass # print("ProjEmb: min {} max {}".format(ProjEmb.min(), ProjEmb.max())) return uv_ProjEmb
def test_draw_3d_bbox(): cur_dir = os.path.abspath(os.path.dirname(__file__)) K = np.array([[572.4114, 0, 325.2611], [0, 573.57043, 242.04899], [0, 0, 1]]) # LM6d model_dir = os.path.join(cur_dir, "../../datasets/BOP_DATASETS/lm/models") class_name = "ape" cls_idx = 1 model_path = os.path.join(model_dir, "obj_{:06d}.ply".format(cls_idx)) pts_3d = load_ply(model_path, vertex_scale=0.001)['pts'] corners_3d = get_3D_corners(pts_3d) image_path = os.path.join( cur_dir, "../../datasets/BOP_DATASETS/lm/test/000001/rgb/000011.png") image = cv2.imread(image_path, cv2.IMREAD_COLOR) gt_dict = mmcv.load( os.path.join( cur_dir, '../../datasets/BOP_DATASETS/lm/test/000001/scene_gt.json')) R = np.array(gt_dict['11'][0]['cam_R_m2c']).reshape(3, 3) t = np.array(gt_dict['11'][0]['cam_t_m2c']) / 1000.0 corners_2d, _ = points_to_2D(corners_3d, R, t, K) # print(pts_2d.shape) image_3dbb = draw_projected_box3d(image, corners_2d, thickness=1) cv2.imshow("image with 3d bbox", image_3dbb) cv2.waitKey()
def calc_uv_emb_proj(uv_model_path_or_model, R, T, K, height=480, width=640, expand=False): """calculate uv map emb via projection it seems to be better not to use expand.""" if isinstance(uv_model_path_or_model, str): model = load_ply(uv_model_path_or_model) if expand: model = ply_vtx_color_expand(model) else: model = uv_model_path_or_model if expand: points = model["pts_expand"] uv_gb = model["colors_expand"][:, [1, 2]] else: points = model["pts"] uv_gb = model["colors"][:, [1, 2]] points_2d, z = points_to_2D(points, R, T, K) image_points = np.round(points_2d).astype(np.int32) # image_points = (points_2d + 0.5).astype(np.int32) uv_ProjEmb = np.zeros((height, width, 2)).astype(np.float32) depth = np.zeros((height, width, 1)).astype(np.float32) for i, (x, y) in enumerate(image_points): if x >= width or y >= height or x < 0 or y < 0: continue if depth[y, x, 0] == 0: depth[y, x, 0] = z[i] uv_ProjEmb[y, x] = uv_gb[i] elif z[i] < depth[y, x, 0]: depth[y, x, 0] = z[i] uv_ProjEmb[y, x] = uv_gb[i] else: pass # print("ProjEmb: min {} max {}".format(ProjEmb.min(), ProjEmb.max())) return uv_ProjEmb
def main(): vertex_scale = 0.001 fps_dict = {} for obj_id in tqdm(id2obj): print(obj_id) model_path = osp.join(model_dir, f"obj_{obj_id:06d}.ply") model = inout.load_ply(model_path, vertex_scale=vertex_scale) fps_dict[str(obj_id)] = {} fps_dict[str(obj_id)]["fps4_and_center"] = get_fps_and_center( model["pts"], num_fps=4, init_center=True) fps_dict[str(obj_id)]["fps8_and_center"] = get_fps_and_center( model["pts"], num_fps=8, init_center=True) fps_dict[str(obj_id)]["fps12_and_center"] = get_fps_and_center( model["pts"], num_fps=12, init_center=True) fps_dict[str(obj_id)]["fps16_and_center"] = get_fps_and_center( model["pts"], num_fps=16, init_center=True) fps_dict[str(obj_id)]["fps20_and_center"] = get_fps_and_center( model["pts"], num_fps=20, init_center=True) fps_dict[str(obj_id)]["fps32_and_center"] = get_fps_and_center( model["pts"], num_fps=32, init_center=True) fps_dict[str(obj_id)]["fps64_and_center"] = get_fps_and_center( model["pts"], num_fps=64, init_center=True) fps_dict[str(obj_id)]["fps128_and_center"] = get_fps_and_center( model["pts"], num_fps=128, init_center=True) fps_dict[str(obj_id)]["fps256_and_center"] = get_fps_and_center( model["pts"], num_fps=256, init_center=True) save_path = osp.join(model_dir, "fps_points.pkl") mmcv.dump(fps_dict, save_path) print(f"saved to {save_path}")
def __init__(self, cfg, dataset_name, distributed, output_dir, train_objs=None): self.cfg = cfg self._distributed = distributed self._output_dir = output_dir self._cpu_device = torch.device("cpu") self._logger = logging.getLogger(__name__) # if test objs are just a subset of train objs self.train_objs = train_objs self._metadata = MetadataCatalog.get(dataset_name) self.data_ref = ref.__dict__[self._metadata.ref_key] self.obj_names = self._metadata.objs self.obj_ids = [self.data_ref.obj2id[obj_name] for obj_name in self.obj_names] # with contextlib.redirect_stdout(io.StringIO()): # self._coco_api = COCO(self._metadata.json_file) self.model_paths = [ osp.join(self.data_ref.model_dir, "obj_{:06d}.ply".format(obj_id)) for obj_id in self.obj_ids ] self.models_3d = [ inout.load_ply(model_path, vertex_scale=self.data_ref.vertex_scale) for model_path in self.model_paths ] # eval cached if cfg.VAL.EVAL_CACHED or cfg.VAL.EVAL_PRINT_ONLY: eval_cached_results(self.cfg, self._output_dir, obj_ids=self.obj_ids)
def __init__(self, cfg, dataset_name, distributed, output_dir, train_objs=None): self.cfg = cfg self._distributed = distributed self._output_dir = output_dir self._cpu_device = torch.device("cpu") self._logger = logging.getLogger(__name__) # if test objs are just a subset of train objs self.train_objs = train_objs self._metadata = MetadataCatalog.get(dataset_name) self.data_ref = ref.__dict__[self._metadata.ref_key] self.obj_names = self._metadata.objs self.obj_ids = [ self.data_ref.obj2id[obj_name] for obj_name in self.obj_names ] # with contextlib.redirect_stdout(io.StringIO()): # self._coco_api = COCO(self._metadata.json_file) self.model_paths = [ osp.join(self.data_ref.model_dir, "obj_{:06d}.ply".format(obj_id)) for obj_id in self.obj_ids ] self.models_3d = [ inout.load_ply(model_path, vertex_scale=self.data_ref.vertex_scale) for model_path in self.model_paths ] if cfg.DEBUG: from lib.render_vispy.model3d import load_models from lib.render_vispy.renderer import Renderer self.ren = Renderer(size=(self.data_ref.width, self.data_ref.height), cam=self.data_ref.camera_matrix) self.ren_models = load_models( model_paths=self.data_ref.model_paths, scale_to_meter=0.001, cache_dir=".cache", texture_paths=self.data_ref.texture_paths, center=False, use_cache=True, ) # eval cached if cfg.VAL.EVAL_CACHED or cfg.VAL.EVAL_PRINT_ONLY: eval_cached_results(self.cfg, self._output_dir, obj_ids=self.obj_ids)
def __init__(self, cfg, dataset_name, distributed, output_dir, train_objs=None): self.cfg = cfg self._distributed = distributed self._output_dir = output_dir mmcv.mkdir_or_exist(output_dir) self._cpu_device = torch.device("cpu") self._logger = logging.getLogger(__name__) # if test objs are just a subset of train objs self.train_objs = train_objs self.dataset_name = dataset_name self._metadata = MetadataCatalog.get(dataset_name) self.data_ref = ref.__dict__[self._metadata.ref_key] self.obj_names = self._metadata.objs self.obj_ids = [ self.data_ref.obj2id[obj_name] for obj_name in self.obj_names ] # with contextlib.redirect_stdout(io.StringIO()): # self._coco_api = COCO(self._metadata.json_file) self.model_paths = [ osp.join(self.data_ref.model_eval_dir, "obj_{:06d}.ply".format(obj_id)) for obj_id in self.obj_ids ] self.diameters = [ self.data_ref.diameters[self.data_ref.objects.index(obj_name)] for obj_name in self.obj_names ] self.models_3d = [ inout.load_ply(model_path, vertex_scale=self.data_ref.vertex_scale) for model_path in self.model_paths ] self.eval_precision = cfg.VAL.get("EVAL_PRECISION", False) self._logger.info(f"eval precision: {self.eval_precision}") # eval cached self.use_cache = False if cfg.VAL.EVAL_CACHED or cfg.VAL.EVAL_PRINT_ONLY: self.use_cache = True if self.eval_precision: self._eval_predictions_precision() else: self._eval_predictions() # recall exit(0)
def test_qrot_points(): from lib.pysixd.inout import load_ply data_root = osp.normpath(osp.join(cur_dir, "../../datasets/BOP_DATASETS/lm/")) models_cad_files = [osp.join(data_root, "models/obj_{:06d}.ply".format(i)) for i in range(1, 15 + 1)] obj_id = 0 points = load_ply(models_cad_files[obj_id])["pts"] axis = np.array([1, 2, 0]) rot = axangle2mat(axis, -pi / 3) quat = mat2quat(rot) points_q = qrot_points_th(torch.from_numpy(quat), torch.from_numpy(points)) # N = points.shape[0] # points_q = qrot_torch(torch.from_numpy(quat).expand(N, 4), torch.from_numpy(points)) points_r = rot.dot(points.T).T print(np.allclose(points_q.numpy(), points_r))
def __init__(self, cfg, dataset_name, distributed, output_dir, train_objs=None): self.cfg = cfg self._distributed = distributed self._output_dir = output_dir mmcv.mkdir_or_exist(output_dir) self._cpu_device = torch.device("cpu") self._logger = logging.getLogger(__name__) # if test objs are just a subset of train objs self.train_objs = train_objs self.dataset_name = dataset_name self._metadata = MetadataCatalog.get(dataset_name) self.data_ref = ref.__dict__[self._metadata.ref_key] self.obj_names = self._metadata.objs self.obj_ids = [ self.data_ref.obj2id[obj_name] for obj_name in self.obj_names ] # with contextlib.redirect_stdout(io.StringIO()): # self._coco_api = COCO(self._metadata.json_file) self.model_paths = [ osp.join(self.data_ref.model_eval_dir, "obj_{:06d}.ply".format(obj_id)) for obj_id in self.obj_ids ] self.diameters = [ self.data_ref.diameters[self.data_ref.objects.index(obj_name)] for obj_name in self.obj_names ] self.models_3d = [ inout.load_ply(model_path, vertex_scale=self.data_ref.vertex_scale) for model_path in self.model_paths ] if cfg.DEBUG: from lib.render_vispy.model3d import load_models from lib.render_vispy.renderer import Renderer self.get_gts() self.kpts_3d = [ misc.get_bbox3d_and_center(m["pts"]) for m in self.models_3d ] self.kpts_axis_3d = [ misc.get_axis3d_and_center(m["pts"], scale=0.5) for m in self.models_3d ] self.ren = Renderer(size=(self.data_ref.width, self.data_ref.height), cam=self.data_ref.camera_matrix) self.ren_models = load_models( model_paths=self.data_ref.model_paths, scale_to_meter=0.001, cache_dir=".cache", texture_paths=self.data_ref.texture_paths, center=False, use_cache=True, ) self.eval_precision = cfg.VAL.get("EVAL_PRECISION", False) self._logger.info(f"eval precision: {self.eval_precision}") # eval cached self.use_cache = False if cfg.VAL.EVAL_CACHED or cfg.VAL.EVAL_PRINT_ONLY: self.use_cache = True if self.eval_precision: self._eval_predictions_precision() else: self._eval_predictions() # recall exit(0)
def add_object(self, obj_id, model_path, **kwargs): """See base class.""" # Color of the object model (the original color saved with the object model # will be used if None). surf_color = None if "surf_color" in kwargs: surf_color = kwargs["surf_color"] # Load the object model. model = inout.load_ply(model_path) self.models[obj_id] = model # Calculate the 3D bounding box of the model (will be used to set the near # and far clipping plane). bb = misc.calc_3d_bbox(model["pts"][:, 0], model["pts"][:, 1], model["pts"][:, 2]) self.model_bbox_corners[obj_id] = np.array([ [bb[0], bb[1], bb[2]], [bb[0], bb[1], bb[2] + bb[5]], [bb[0], bb[1] + bb[4], bb[2]], [bb[0], bb[1] + bb[4], bb[2] + bb[5]], [bb[0] + bb[3], bb[1], bb[2]], [bb[0] + bb[3], bb[1], bb[2] + bb[5]], [bb[0] + bb[3], bb[1] + bb[4], bb[2]], [bb[0] + bb[3], bb[1] + bb[4], bb[2] + bb[5]], ]) # Set texture/color of vertices. self.model_textures[obj_id] = None # Use the specified uniform surface color. if surf_color is not None: colors = np.tile( list(surf_color) + [1.0], [model["pts"].shape[0], 1]) # Set UV texture coordinates to dummy values. texture_uv = np.zeros((model["pts"].shape[0], 2), np.float32) # Use the model texture. elif "texture_file" in self.models[obj_id].keys(): model_texture_path = os.path.join( os.path.dirname(model_path), self.models[obj_id]["texture_file"]) model_texture = inout.load_im(model_texture_path) # Normalize the texture image. if model_texture.max() > 1.0: model_texture = model_texture.astype(np.float32) / 255.0 model_texture = np.flipud(model_texture) self.model_textures[obj_id] = model_texture # UV texture coordinates. texture_uv = model["texture_uv"] # Set the per-vertex color to dummy values. colors = np.zeros((model["pts"].shape[0], 3), np.float32) # Use the original model color. elif "colors" in model.keys(): assert model["pts"].shape[0] == model["colors"].shape[0] colors = model["colors"] if colors.max() > 1.0: colors /= 255.0 # Color values are expected in range [0, 1]. # Set UV texture coordinates to dummy values. texture_uv = np.zeros((model["pts"].shape[0], 2), np.float32) # Set the model color to gray. else: colors = np.ones((model["pts"].shape[0], 3), np.float32) * 0.5 # Set UV texture coordinates to dummy values. texture_uv = np.zeros((model["pts"].shape[0], 2), np.float32) # Set the vertex data. if self.mode == "depth": vertices_type = [("a_position", np.float32, 3), ("a_color", np.float32, colors.shape[1])] vertices = np.array(list(zip(model["pts"], colors)), vertices_type) else: if self.shading == "flat": vertices_type = [ ("a_position", np.float32, 3), ("a_color", np.float32, colors.shape[1]), ("a_texcoord", np.float32, 2), ] vertices = np.array( list(zip(model["pts"], colors, texture_uv)), vertices_type) elif self.shading == "phong": vertices_type = [ ("a_position", np.float32, 3), ("a_normal", np.float32, 3), ("a_color", np.float32, colors.shape[1]), ("a_texcoord", np.float32, 2), ] vertices = np.array( list( zip(model["pts"], model["normals"], colors, texture_uv)), vertices_type) else: raise ValueError("Unknown shading type.") # Create vertex and index buffer for the loaded object model. self.vertex_buffers[obj_id] = vertices.view(gloo.VertexBuffer) self.index_buffers[obj_id] = model["faces"].flatten().astype( np.uint32).view(gloo.IndexBuffer) # Set shader for the selected shading. if self.shading == "flat": rgb_fragment_code = _rgb_fragment_flat_code elif self.shading == "phong": rgb_fragment_code = _rgb_fragment_phong_code else: raise ValueError("Unknown shading type.") # Prepare the RGB OpenGL program. rgb_program = gloo.Program(_rgb_vertex_code, rgb_fragment_code) rgb_program.bind(self.vertex_buffers[obj_id]) if self.model_textures[obj_id] is not None: rgb_program["u_use_texture"] = int(True) rgb_program["u_texture"] = self.model_textures[obj_id] else: rgb_program["u_use_texture"] = int(False) rgb_program["u_texture"] = np.zeros((1, 1, 4), np.float32) self.rgb_programs[obj_id] = rgb_program # Prepare the depth OpenGL program. depth_program = gloo.Program(_depth_vertex_code, _depth_fragment_code) depth_program.bind(self.vertex_buffers[obj_id]) self.depth_programs[obj_id] = depth_program
"add", "adi", "mssd", "mspd", "proj", "projS", "ABSadd", "ABSadi", "ABSad", "AUCadd", "AUCadi", "AUCad", ]: misc.log("Loading object models...") for obj_id in dp_model["obj_ids"]: models[obj_id] = inout.load_ply(dp_model["model_tpath"].format(obj_id=obj_id)) # Load models info. models_info = None if p["error_type"] in ["ad", "add", "adi", "vsd", "mssd", "mspd", "cus", "reteS", "reS", "teS", "projS"]: models_info = inout.load_json(dp_model["models_info_path"], keys_to_int=True) # Get sets of symmetry transformations for the object models. models_sym = None if p["error_type"] in ["mssd", "mspd", "reteS", "reS", "teS", "projS"]: models_sym = {} for obj_id in dp_model["obj_ids"]: models_sym[obj_id] = misc.get_symmetry_transformations(models_info[obj_id], p["max_sym_disc_step"]) # Initialize a renderer. ren = None