def record_ann(model_meta, img_id, ann_id, images, annotations): data_root = model_meta['data_root'] corner_3d = model_meta['corner_3d'] center_3d = model_meta['center_3d'] fps_3d = model_meta['fps_3d'] K = model_meta['K'] pose_dir = os.path.join(data_root, 'pose') rgb_dir = os.path.join(data_root, 'rgb') mask_dir = os.path.join(data_root, 'mask') inds = range(len(os.listdir(rgb_dir))) for ind in tqdm.tqdm(inds): rgb_path = os.path.join(rgb_dir, '{}.jpg'.format(ind)) rgb = Image.open(rgb_path) img_size = rgb.size img_id += 1 info = { 'file_name': rgb_path, 'height': img_size[1], 'width': img_size[0], 'id': img_id } images.append(info) pose_path = os.path.join(pose_dir, 'pose{}.npy'.format(ind)) pose = np.load(pose_path) corner_2d = base_utils.project(corner_3d, K, pose) center_2d = base_utils.project(center_3d[None], K, pose)[0] fps_2d = base_utils.project(fps_3d, K, pose) mask_path = os.path.join(mask_dir, '{}.png'.format(ind)) ann_id += 1 anno = { 'mask_path': mask_path, 'image_id': img_id, 'category_id': 1, 'id': ann_id } anno.update({ 'corner_3d': corner_3d.tolist(), 'corner_2d': corner_2d.tolist() }) anno.update({ 'center_3d': center_3d.tolist(), 'center_2d': center_2d.tolist() }) anno.update({'fps_3d': fps_3d.tolist(), 'fps_2d': fps_2d.tolist()}) anno.update({'K': K.tolist(), 'pose': pose.tolist()}) anno.update({'data_root': rgb_dir}) anno.update({'type': 'real', 'cls': 'cat'}) annotations.append(anno) return img_id, ann_id
def get_bound_2d_mask(bounds, K, pose, H, W): corners_3d = get_bound_corners(bounds) corners_2d = base_utils.project(corners_3d, K, pose) corners_2d = np.round(corners_2d).astype(int) mask = np.zeros((H, W), dtype=np.uint8) cv2.fillPoly(mask, [corners_2d[[0, 1, 3, 2, 0]]], 1) cv2.fillPoly(mask, [corners_2d[[4, 5, 7, 6, 5]]], 1) cv2.fillPoly(mask, [corners_2d[[0, 1, 5, 4, 0]]], 1) cv2.fillPoly(mask, [corners_2d[[2, 3, 7, 6, 2]]], 1) cv2.fillPoly(mask, [corners_2d[[0, 2, 6, 4, 0]]], 1) cv2.fillPoly(mask, [corners_2d[[1, 3, 7, 5, 1]]], 1) return mask
def prepare_inside_pts(self, pts, msk, K, R, T): sh = pts.shape pts3d = pts.reshape(-1, 3) RT = np.concatenate([R, T], axis=1) pts2d = base_utils.project(pts3d, K, RT) H, W = msk.shape pts2d = np.round(pts2d).astype(np.int32) pts2d[:, 0] = np.clip(pts2d[:, 0], 0, W - 1) pts2d[:, 1] = np.clip(pts2d[:, 1], 0, H - 1) inside = msk[pts2d[:, 1], pts2d[:, 0]] inside = inside.reshape(*sh[:-1]) return inside
def prepare_inside_pts(self, pts, i): sh = pts.shape pts3d = pts.reshape(-1, 3) inside = np.ones([len(pts3d)]).astype(np.uint8) for nv in range(self.ims.shape[1]): ind = inside == 1 pts3d_ = pts3d[ind] RT = np.concatenate([self.Rs[nv], self.Ts[nv]], axis=1) pts2d = base_utils.project(pts3d_, self.Ks[nv], RT) msk = self.get_mask(i, nv) H, W = msk.shape pts2d = np.round(pts2d).astype(np.int32) pts2d[:, 0] = np.clip(pts2d[:, 0], 0, W - 1) pts2d[:, 1] = np.clip(pts2d[:, 1], 0, H - 1) msk_ = msk[pts2d[:, 1], pts2d[:, 0]] inside[ind] = msk_ inside = inside.reshape(*sh[:-1]) return inside
def record_ann(model_meta, img_id, ann_id, images, annotations, cls_type): data_root = model_meta['data_root'] corner_3d = model_meta['corner_3d'] center_3d = model_meta['center_3d'] fps_3d = model_meta['fps_3d'] K = model_meta['K'] length = int(input("How many pictures do you want to use? ")) inds = range(length) for ind in tqdm.tqdm(inds): desired_class = cls_type # name of the objectclass, on which you want to train number = str(ind) number = number.zfill(6) # getting rgb datei = number + '.png' rgb_path = os.path.join(data_root, datei) rgb = Image.open(rgb_path).convert('RGB') img_size = rgb.size img_id += 1 info = { 'file_name': rgb_path, 'height': img_size[1], 'width': img_size[0], 'id': img_id } images.append(info) # path to annotations from ndds datei = number + '.json' pose_path = os.path.join(data_root, datei) # getting pose of annotations from ndds with open(pose_path, 'r') as file: annotation = json.loads(file.read()) object_from_annotation = annotation['objects'] object_class = object_from_annotation[0]["class"] if desired_class in object_class: # translation translation = np.array(object_from_annotation[0]['location']) * 10 # rotation rotation = np.asarray( object_from_annotation[0]['pose_transform'])[0:3, 0:3] rotation = np.dot(rotation, np.array([[-1, 0, 0], [0, -1, 0], [0, 0, -1]])) rotation = np.dot(rotation.T, np.array([[1, 0, 0], [0, 1, 0], [0, 0, -1]])) # pose pose = np.column_stack((rotation, translation)) else: print("No such class in annotations!") pass corner_2d = base_utils.project(corner_3d, K, pose) center_2d = base_utils.project(center_3d[None], K, pose)[0] fps_2d = base_utils.project(fps_3d, K, pose) # getting segmentation-mask datei = number + '.cs.png' mask_path = os.path.join(data_root, datei) ann_id += 1 anno = { 'mask_path': mask_path, 'image_id': img_id, 'category_id': 1, 'id': ann_id } anno.update({ 'corner_3d': corner_3d.tolist(), 'corner_2d': corner_2d.tolist() }) anno.update({ 'center_3d': center_3d.tolist(), 'center_2d': center_2d.tolist() }) anno.update({'fps_3d': fps_3d.tolist(), 'fps_2d': fps_2d.tolist()}) anno.update({'K': K.tolist(), 'pose': pose.tolist()}) anno.update({'data_root': data_root}) anno.update({'type': 'real', 'cls': cls_type}) annotations.append(anno) return img_id, ann_id
def record_scene_ann(model_meta, pose_meta, img_id, ann_id, images, annotations): corner_3d_dict = model_meta['corner_3d'] center_3d_dict = model_meta['center_3d'] fps_3d_dict = model_meta['fps_3d'] rgb_paths = pose_meta['rgb_paths'] gt = pose_meta['gt'] K = pose_meta['K'] a_pixel_num_dict = pose_meta['a_pixel_num'] for rgb_path in tqdm.tqdm(rgb_paths): rgb = Image.open(rgb_path) img_size = rgb.size img_id += 1 depth_path = rgb_path.replace('rgb', 'depth') info = { 'file_name': rgb_path, 'depth_path': depth_path, 'height': img_size[1], 'width': img_size[0], 'id': img_id } images.append(info) ind = int(os.path.basename(rgb_path).replace('.png', '')) K_ = K[ind]['cam_K'] K_ = np.array(K_).reshape(3, 3) gt_ = gt[ind] mask_path = rgb_path.replace('rgb', 'mask') mask = np.array(Image.open(mask_path)) a_pixel_nums = a_pixel_num_dict[ind] for instance_id, instance_gt in enumerate(gt_): obj_id = instance_gt['obj_id'] mask_id = obj_id * 1000 + instance_id mask_ = (mask == mask_id).astype(np.uint8) pixel_num = np.sum(mask_) a_pixel_num = a_pixel_nums[instance_id] if pixel_num / a_pixel_num < 0.1: continue R = np.array(instance_gt['cam_R_m2c']).reshape(3, 3) t = np.array(instance_gt['cam_t_m2c']) * 0.001 pose = np.concatenate([R, t[:, None]], axis=1) corner_3d = corner_3d_dict[obj_id] center_3d = center_3d_dict[obj_id] fps_3d = fps_3d_dict[obj_id] corner_2d = base_utils.project(corner_3d, K_, pose) center_2d = base_utils.project(center_3d[None], K_, pose)[0] fps_2d = base_utils.project(fps_3d, K_, pose) ann_id += 1 anno = {} anno.update({ 'corner_3d': corner_3d.tolist(), 'corner_2d': corner_2d.tolist() }) anno.update({ 'center_3d': center_3d.tolist(), 'center_2d': center_2d.tolist() }) anno.update({'fps_3d': fps_3d.tolist(), 'fps_2d': fps_2d.tolist()}) anno.update({'K': K_.tolist(), 'pose': pose.tolist()}) anno.update({ 'category_id': obj_id, 'image_id': img_id, 'id': ann_id }) annotations.append(anno) return img_id, ann_id