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
