def __getitem__(self, index):
index = 0
img_id = self.img_ids[index]
img, kpt_2d, mask = self.read_data(img_id)
img, kpt_2d, mask, bbox = self.get_training_img(img, kpt_2d, mask)
orig_img, inp, trans_input, center, scale, inp_hw = \
tless_pvnet_utils.augment(img, bbox, 'train')
kpt_2d = data_utils.affine_transform(kpt_2d, trans_input)
mask = cv2.warpAffine(mask, trans_input, (inp_hw[1], inp_hw[0]), flags=cv2.INTER_NEAREST)
vertex = pvnet_data_utils.compute_vertex(mask, kpt_2d).transpose(2, 0, 1)
ret = {'inp': inp, 'mask': mask.astype(np.uint8), 'vertex': vertex}
# visualize_utils.visualize_ann(orig_img, kpt_2d, mask, False)
return ret
def __getitem__(self, index_tuple):
index, height, width = index_tuple
img_id = self.img_ids[index]
img, kpt_2d, mask = self.read_data(img_id)
if self.split == 'train':
inp, kpt_2d, mask = self.augment(img, mask, kpt_2d, height, width)
else:
inp = img
if self._transforms is not None:
inp, kpt_2d, mask = self._transforms(inp, kpt_2d, mask)
vertex = pvnet_data_utils.compute_vertex(mask, kpt_2d).transpose(2, 0, 1)
ret = {'inp': inp, 'mask': mask.astype(np.uint8), 'vertex': vertex, 'img_id': img_id, 'meta': ''}
# visualize_utils.visualize_linemod_ann(torch.tensor(inp), kpt_2d, mask, True)
return ret
def __getitem__(self, idx):
db_rec = copy.deepcopy(self.db[idx])
img_id = int(db_rec['image'].split('/')[-1].split('.')[0])
image_file = db_rec['image']
filename = db_rec['filename'] if 'filename' in db_rec else ''
imgnum = db_rec['imgnum'] if 'imgnum' in db_rec else ''
data_numpy = cv2.imread(
image_file, cv2.IMREAD_COLOR | cv2.IMREAD_IGNORE_ORIENTATION)
if data_numpy is None:
raise ValueError('Fail to read {}'.format(image_file))
joints = db_rec['joints_3d']
joints_vis = db_rec['joints_3d_vis']
c = db_rec['center']
s = db_rec['scale']
score = db_rec['score'] if 'score' in db_rec else 1
r = 0
mask = self.annToMask(db_rec)
mask = np.expand_dims(mask, axis=-1)
if self.split == 'train':
sf = self.scale_factor
rf = self.rotation_factor
s = s * np.clip(np.random.randn()*sf + 1, 1 - sf, 1 + sf)
r = np.clip(np.random.randn()*rf, -rf*2, rf*2) \
if random.random() <= 0.6 else 0
if self.flip and random.random() <= 0.5:
data_numpy = data_numpy[:, ::-1, :]
mask = mask[:, ::-1, :]
joints, joints_vis = fliplr_joints(
joints, joints_vis, data_numpy.shape[1], self.flip_pairs)
c[0] = data_numpy.shape[1] - c[0] - 1
trans = get_affine_transform(c, s, r, np.array([self.image_width, self.image_height]))
input = cv2.warpAffine(
data_numpy,
trans,
(int(self.image_width), int(self.image_height)),
flags=cv2.INTER_LINEAR)
mask = cv2.warpAffine(
mask,
trans,
(int(self.image_width), int(self.image_height)),
flags=cv2.INTER_NEAREST)
if self.transforms:
input, joints[:, :2], mask = self.transforms(input, joints[:, :2], mask)
for i in range(self.num_joints):
if joints_vis[i, 0] > 0.0:
joints[i, 0:2] = affine_transform(joints[i, 0:2], trans)
vertex = pvnet_data_utils.compute_vertex(mask.squeeze(), joints[:,:2]).transpose(2, 0, 1)
ret = {'inp': input, 'mask': mask.astype(np.uint8), 'vertex': vertex, 'img_id': img_id,
'joints_vis': torch.from_numpy(joints_vis[:, 0].astype('float32'))}
meta = {
'image': image_file,
'filename': filename,
'imgnum': imgnum,
'joints': joints,
'joints_vis': joints_vis,
'center': c,
'scale': s,
'rotation': r,
'score': score,
'mask': mask
}
if 'Coco' in cfg.test.dataset and self.split == 'test':
return ret, meta
else:
return ret
def __getitem__(self, index_tuple):
index, height, width = index_tuple
height, width = 480, 640
img_id = self.img_ids[index]
img, kpt_2d, kpt_3d, mask, pose, K = self.read_data(img_id)
if self.split == 'train':
# inp = np.asarray(img).astype(np.uint8)
inp, kpt_2d, mask, pose = self.augment(img, mask, kpt_2d, pose, K,
height, width)
else:
inp = img
K = np.array(K)
dist_coeffs = np.zeros((4, 1))
(_, rvec, tvec) = cv2.solvePnP(kpt_3d, kpt_2d.astype(np.float32), K,
dist_coeffs)
rot = cv2.Rodrigues(rvec)[0]
pnp_pose = np.concatenate([rot, tvec], axis=1)
pose_old = pnp_pose.copy()
# for symmetric objects
if self.cls in ['eggbox', 'glue']:
rot = R.from_matrix(pnp_pose[:3, :3])
euler_angles = rot.as_euler('xyz', degrees=True)
if euler_angles[2] >= 0:
#corrected pose
pnp_pose[:3, :3] = np.dot(
pnp_pose[:3, :3],
np.array([[-1, 0, 0], [0, -1, 0], [0, 0, 1]]))
# correct keypoints
kpt_2d = project(kpt_3d, K, pnp_pose)
# print('computed pose', pose)
# print('old pose', pose_old)
# print('pnp pose ', pnp_pose)
# data = {}
# data['img'] = inp
# data['img_old'] = img
# data['mask'] = mask
# data['pose_old'] = pose_old
# data['pose'] = pnp_pose
# data['kpt_2d'] = kpt_2d
# data['K'] = K
# with open('/mbrdi/sqnap1_colomirror/gupansh/data.pkl','wb') as fp:
# pickle.dump(data, fp)
# print('saved..')
# input()
if self._transforms is not None:
inp, kpt_2d, mask = self._transforms(inp, kpt_2d, mask)
# shape is (18, H, W) y1,x1,y2,x2 ...y9,x9
vertex = pvnet_data_utils.compute_vertex(mask,
kpt_2d).transpose(2, 0, 1)
ret = {
'inp': inp,
'mask': mask.astype(np.uint8),
'vertex': vertex,
'img_id': img_id,
'pose': pnp_pose,
'pose_pnp': pnp_pose,
'K': K,
'xy': kpt_2d,
'meta': {}
}
# visualize_utils.visualize_linemod_ann(torch.tensor(inp), kpt_2d, mask, True)
return ret