def test_quaternion2matrix(self):
testing.assert_array_equal(quaternion2matrix([1, 0, 0, 0]), np.eye(3))
testing.assert_almost_equal(
quaternion2matrix([1.0 / np.sqrt(2), 1.0 / np.sqrt(2), 0, 0]),
np.array([[1., 0., 0.], [0., 0., -1.], [0., 1., 0.]]))
testing.assert_almost_equal(
quaternion2matrix(
normalize_vector(
[1.0, 1 / np.sqrt(2), 1 / np.sqrt(2), 1 / np.sqrt(2)])),
np.array([[0.2000000, -0.1656854, 0.9656854],
[0.9656854, 0.2000000, -0.1656854],
[-0.1656854, 0.9656854, 0.2000000]]))
testing.assert_almost_equal(
quaternion2matrix(
normalize_vector(
[1.0, -1 / np.sqrt(2), 1 / np.sqrt(2), -1 / np.sqrt(2)])),
np.array([[0.2000000, 0.1656854, 0.9656854],
[-0.9656854, 0.2000000, 0.1656854],
[-0.1656854, -0.9656854, 0.2000000]]))
testing.assert_almost_equal(
quaternion2matrix([0.925754, 0.151891, 0.159933, 0.307131]),
rotate_matrix(
rotate_matrix(rotate_matrix(np.eye(3), 0.2, 'x'), 0.4, 'y'),
0.6, 'z'),
decimal=5)
def rotation(self, rotation):
"""Set rotation of this coordinate
This setter checkes the given rotation and set it this coordinate.
Parameters
----------
rotation : list or numpy.ndarray
we can take 3x3 rotation matrix or
rpy angle [yaw, pitch, roll] or
quaternion [w, x, y, z] order
"""
rotation = np.array(rotation)
# Convert quaternions
if rotation.shape == (4, ):
self._q = np.array([q for q in rotation])
if np.abs(np.linalg.norm(self._q) - 1.0) > 1e-3:
raise ValueError('Invalid quaternion. Must be '
'norm 1.0, get {}'.format(
np.linalg.norm(self._q)))
rotation = quaternion2matrix(self._q)
elif rotation.shape == (3, ):
# Convert [yaw-pitch-roll] to rotation matrix
self._q = rpy2quaternion(rotation)
rotation = quaternion2matrix(self._q)
else:
self._q = matrix2quaternion(rotation)
# Convert lists and tuples
if type(rotation) in (list, tuple):
rotation = np.array(rotation).astype(np.float32)
_check_valid_rotation(rotation)
self._rotation = rotation * 1.
def test_rotation(self):
qr = normalize_vector(np.array([1, 2, 3, 4]))
x, y, z = np.array([1, 2, 3])
qd = 0.5 * quaternion_multiply(np.array([0, x, y, z]), qr)
dq = DualQuaternion(qr, qd)
testing.assert_almost_equal(dq.rotation, quaternion2matrix(qr))
def test_matrix2quaternion(self):
testing.assert_almost_equal(matrix2quaternion(np.eye(3)),
np.array([1, 0, 0, 0]))
m = rotate_matrix(
rotate_matrix(rotate_matrix(np.eye(3), 0.2, 'x'), 0.4, 'y'), 0.6,
'z')
testing.assert_almost_equal(quaternion2matrix(matrix2quaternion(m)), m)
testing.assert_almost_equal(matrix2quaternion(
np.array([[0.428571, 0.514286, -0.742857],
[-0.857143, -0.028571, -0.514286],
[-0.285714, 0.857143, 0.428571]])),
normalize_vector(np.array([4, 3, -1, -3])),
decimal=5)
def quaternion2vec(q, axis='x'):
"""Calculate axis vector from quaternion
Parameters
----------
q : np.ndarray
qauternion
axis : str, optional
axis x, y, z, by default 'x'
Returns
-------
vec : numpy.ndarray
"""
m = quaternion2matrix(q)
return matrix2vec(m, axis=axis)
def rotation(self):
"""Return rotation matrix.
Note that this property internally normalizes quaternion.
Returns
-------
quaternion2matrix(self.q) : numpy.ndarray
3x3 rotation matrix
Examples
--------
>>> from skrobot.coordinates.quaternion import Quaternion
>>> q = Quaternion()
>>> q
#<Quaternion 0x12f1aa6a0 w: 1.0 x: 0.0 y: 0.0 z: 0.0>
>>> q.rotation
array([[1., 0., 0.],
[0., 1., 0.],
[0., 0., 1.]])
>>> q.q = [0, 1, 0, 0]
>>> q
#<Quaternion 0x12f1aa6a0 w: 0.0 x: 1.0 y: 0.0 z: 0.0>
>>> q.rotation
array([[ 1., 0., 0.],
[ 0., -1., 0.],
[ 0., 0., -1.]])
>>> q.q = [1, 2, 3, 4]
>>> q
#<Quaternion 0x12f1aa6a0 w: 1.0 x: 2.0 y: 3.0 z: 4.0>
>>> q.rotation
array([[-0.66666667, 0.13333333, 0.73333333],
[ 0.66666667, -0.33333333, 0.66666667],
[ 0.33333333, 0.93333333, 0.13333333]])
"""
return quaternion2matrix(self.normalized.q)
def predict(self, rgb_img, depth, label_img, label, bboxes,
intrinsic_matrix):
num_points = self.num_points
iteration = self.iteration
batch_size = 1
lst = np.array(label.flatten(), dtype=np.int32)
img_height, img_width, _ = rgb_img.shape
cameramodel = cameramodels.PinholeCameraModel.from_intrinsic_matrix(
intrinsic_matrix, img_height, img_width)
translations = []
rotations = []
for idx in range(len(lst)):
itemid = lst[idx]
rmin, cmin, rmax, cmax = bboxes[idx]
mask_depth = ma.getmaskarray(ma.masked_not_equal(depth, 0))
mask_label = ma.getmaskarray(ma.masked_equal(label_img, itemid))
mask = mask_label * mask_depth
choose = mask[rmin:rmax, cmin:cmax].flatten().nonzero()[0]
if len(choose) == 0:
translations.append([])
rotations.append([])
continue
if len(choose) > num_points:
c_mask = np.zeros(len(choose), dtype=int)
c_mask[:num_points] = 1
np.random.shuffle(c_mask)
choose = choose[c_mask.nonzero()]
else:
choose = np.pad(choose, (0, num_points - len(choose)), 'wrap')
xmap = np.array([[j for i in range(img_width)]
for j in range(img_height)])
ymap = np.array([[i for i in range(img_width)]
for j in range(img_height)])
depth_masked = depth[
rmin:rmax,
cmin:cmax].flatten()[choose][:, np.newaxis].astype(np.float32)
xmap_masked = xmap[rmin:rmax,
cmin:cmax].flatten()[choose][:,
np.newaxis].astype(
np.float32)
ymap_masked = ymap[rmin:rmax,
cmin:cmax].flatten()[choose][:,
np.newaxis].astype(
np.float32)
choose = np.array([choose])
cloud = cameramodel.batch_project_pixel_to_3d_ray(
np.concatenate([ymap_masked, xmap_masked], axis=1),
depth_masked)
img_masked = np.array(rgb_img)[:, :, :3]
img_masked = np.transpose(img_masked, (2, 0, 1))
img_masked = img_masked[:, rmin:rmax, cmin:cmax]
with torch.no_grad():
cloud = torch.from_numpy(cloud.astype(np.float32))
choose = torch.LongTensor(choose.astype(np.int32))
img_masked = normalize(
torch.from_numpy(img_masked.astype(np.float32)))
index = torch.LongTensor([itemid - 1])
cloud = cloud.cuda()
choose = choose.cuda()
img_masked = img_masked.cuda()
index = index.cuda()
cloud = cloud.view(1, num_points, 3)
img_masked = img_masked.view(1, 3,
img_masked.size()[1],
img_masked.size()[2])
pred_rot, pred_trans, pred_score, emb = self.estimator(
img_masked, cloud, choose, index)
pred_rot = pred_rot / torch.norm(pred_rot, dim=2).view(
1, num_points, 1)
pred_score = pred_score.view(batch_size, num_points)
_, which_max = torch.max(pred_score, 1)
pred_trans = pred_trans.view(batch_size * num_points, 1, 3)
points = cloud.view(batch_size * num_points, 1, 3)
rotation = pred_rot[0][which_max[0]].view(-1).cpu().\
data.numpy()
translation = (points + pred_trans)[which_max[0]].view(-1).\
cpu().data.numpy()
for _ in range(iteration):
T = torch.from_numpy(translation.astype(np.float32)).\
cuda().view(1, 3).\
repeat(num_points, 1).contiguous().\
view(1, num_points, 3)
trans_matrix = np.eye(4)
trans_matrix[:3, :3] = quaternion2matrix(
quaternion_normalize(rotation))
R = torch.from_numpy(trans_matrix[:3, :3].astype(
np.float32)).cuda().view(1, 3, 3)
trans_matrix[0:3, 3] = translation
new_cloud = torch.bmm((cloud - T), R).contiguous()
refined_rot, refined_trans = self.refiner(
new_cloud, emb, index)
refined_rot = refined_rot.view(1, 1, -1)
refined_rot = refined_rot / (torch.norm(
refined_rot, dim=2).view(1, 1, 1))
rotation_2 = refined_rot.view(-1).cpu().data.numpy()
translation_2 = refined_trans.view(-1).cpu().data.numpy()
trans_matrix_2 = np.eye(4)
trans_matrix_2[:3, :3] = quaternion2matrix(
quaternion_normalize(rotation_2))
trans_matrix_2[0:3, 3] = translation_2
trans_matrix_final = np.dot(trans_matrix, trans_matrix_2)
rotation_final = matrix2quaternion(
trans_matrix_final[:3, :3])
translation_final = np.array([
trans_matrix_final[0][3], trans_matrix_final[1][3],
trans_matrix_final[2][3]
])
rotation = rotation_final
translation = translation_final
translations.append(translation)
rotations.append(quaternion_normalize(rotation))
return rotations, translations
def callback(self, camera_info_msg, img_raw_msg, img_msg, depth_msg):
ymin = camera_info_msg.roi.y_offset
xmin = camera_info_msg.roi.x_offset
ymax = camera_info_msg.roi.y_offset + camera_info_msg.roi.height
xmax = camera_info_msg.roi.x_offset + camera_info_msg.roi.width
self.camera_model.roi = [ymin, xmin, ymax, xmax]
self.camera_model.target_size = self.target_size
bgr_raw = self.bridge.imgmsg_to_cv2(img_raw_msg, "bgr8")
bgr = self.bridge.imgmsg_to_cv2(img_msg, "bgr8")
cv_depth = self.bridge.imgmsg_to_cv2(depth_msg, "32FC1")
if cv_depth is None or bgr is None:
return
remove_nan(cv_depth)
cv_depth[cv_depth < self.depth_range[0]] = 0
cv_depth[cv_depth > self.depth_range[1]] = 0
bgr = cv2.resize(bgr, self.target_size)
cv_depth = cv2.resize(cv_depth,
self.target_size,
interpolation=cv2.INTER_NEAREST)
depth_bgr = colorize_depth(cv_depth, ignore_value=0)
in_feature = cv_depth.copy().astype(np.float32) * 0.001
if self.config['use_bgr2gray']:
gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
gray = cv2.resize(gray, self.target_size)[..., None] / 255.
normalized_depth = normalize_depth(cv_depth, self.depth_range[0],
self.depth_range[1])[..., None]
in_feature = np.concatenate((normalized_depth, gray),
axis=2).astype(np.float32)
if self.transform:
in_feature = self.transform(in_feature)
in_feature = in_feature.to(self.device)
in_feature = in_feature.unsqueeze(0)
confidence, depth, rotation = self.model(in_feature)
confidence = confidence[0, 0:1, ...]
confidence_np = confidence.cpu().detach().numpy().copy() * 255
confidence_np = confidence_np.transpose(1, 2, 0)
confidence_np[confidence_np <= 0] = 0
confidence_np[confidence_np >= 255] = 255
confidence_img = confidence_np.astype(np.uint8)
confidence_img = cv2.resize(confidence_img, self.target_size)
heatmap = overlay_heatmap(bgr, confidence_img)
axis_pred = bgr.copy()
axis_pred_raw = bgr_raw.copy()
hanging_points_pose_array = PoseArray()
for i, (roi, roi_center) in enumerate(
zip(self.model.rois_list[0], self.model.rois_center_list[0])):
if roi.tolist() == [0, 0, 0, 0]:
continue
roi = roi.cpu().detach().numpy().copy()
hanging_point_x = roi_center[0]
hanging_point_y = roi_center[1]
v = rotation[i].cpu().detach().numpy()
v /= np.linalg.norm(v)
rot = rotation_matrix_from_axis(v, [0, 1, 0], 'xy')
q = matrix2quaternion(rot)
hanging_point = np.array(
self.camera_model.project_pixel_to_3d_ray(
[int(hanging_point_x),
int(hanging_point_y)]))
if self.predict_depth:
dep = depth[i].cpu().detach().numpy().copy()
dep = unnormalize_depth(dep, self.depth_range[0],
self.depth_range[1]) * 0.001
length = float(dep) / hanging_point[2]
else:
depth_roi = make_box(roi_center,
width=self.depth_roi_size[1],
height=self.depth_roi_size[0],
img_shape=self.target_size,
xywh=False)
depth_roi_clip = cv_depth[depth_roi[0]:depth_roi[2],
depth_roi[1]:depth_roi[3]]
dep_roi_clip = depth_roi_clip[np.where(
np.logical_and(self.depth_range[0] < depth_roi_clip,
depth_roi_clip < self.depth_range[1]))]
dep_roi_clip = np.median(dep_roi_clip) * 0.001
if dep_roi_clip == np.nan:
continue
length = float(dep_roi_clip) / hanging_point[2]
hanging_point *= length
hanging_point_pose = Pose()
hanging_point_pose.position.x = hanging_point[0]
hanging_point_pose.position.y = hanging_point[1]
hanging_point_pose.position.z = hanging_point[2]
hanging_point_pose.orientation.w = q[0]
hanging_point_pose.orientation.x = q[1]
hanging_point_pose.orientation.y = q[2]
hanging_point_pose.orientation.z = q[3]
hanging_points_pose_array.poses.append(hanging_point_pose)
axis_pred_raw = cv2.rectangle(
axis_pred_raw,
(int(roi[0] * (xmax - xmin) / self.target_size[1] + xmin),
int(roi[1] * (ymax - ymin) / self.target_size[0] + ymin)),
(int(roi[2] * (xmax - xmin) / self.target_size[1] + xmin),
int(roi[3] * (ymax - ymin) / self.target_size[0] + ymin)),
(0, 255, 0), 1)
try:
axis_pred_raw = draw_axis(axis_pred_raw, quaternion2matrix(q),
hanging_point,
self.camera_model.full_K)
except Exception:
print('Fail to draw axis')
axis_pred = self.camera_model.crop_image(axis_pred_raw,
copy=True).astype(np.uint8)
msg_out = self.bridge.cv2_to_imgmsg(heatmap, "bgr8")
msg_out.header.stamp = depth_msg.header.stamp
confidence_msg = self.bridge.cv2_to_imgmsg(confidence_img, "mono8")
confidence_msg.header.stamp = depth_msg.header.stamp
colorized_depth_msg = self.bridge.cv2_to_imgmsg(depth_bgr, "bgr8")
colorized_depth_msg.header.stamp = depth_msg.header.stamp
axis_pred_msg = self.bridge.cv2_to_imgmsg(axis_pred, "bgr8")
axis_pred_msg.header.stamp = depth_msg.header.stamp
axis_pred_raw_msg = self.bridge.cv2_to_imgmsg(axis_pred_raw, "bgr8")
axis_pred_raw_msg.header.stamp = depth_msg.header.stamp
hanging_points_pose_array.header = camera_info_msg.header
self.pub.publish(msg_out)
self.pub_confidence.publish(confidence_msg)
self.pub_depth.publish(colorized_depth_msg)
self.pub_axis.publish(axis_pred_msg)
self.pub_axis_raw.publish(axis_pred_raw_msg)
self.pub_hanging_points.publish(hanging_points_pose_array)
def step(self, dataloader, mode):
print('Start {}'.format(mode))
# self.model = self.prev_model
if mode == 'train':
self.model.train()
elif mode == 'val' or mode == 'test':
self.model.eval()
loss_sum = 0
confidence_loss_sum = 0
depth_loss_sum = 0
rotation_loss_sum = 0
rotation_loss_count = 0
for index, (hp_data, depth_image, camera_info_path, hp_data_gt,
annotation_data) in tqdm.tqdm(enumerate(dataloader),
total=len(dataloader),
desc='{} epoch={}'.format(
mode, self.epo),
leave=False):
# if index == 0:
# self.model = self.prev_model
self.cameramodel\
= cameramodels.PinholeCameraModel.from_yaml_file(
camera_info_path[0])
self.cameramodel.target_size = self.target_size
depth_image = hp_data.numpy().copy()[0, 0, ...]
depth_image = np.nan_to_num(depth_image)
depth_image = unnormalize_depth(depth_image, self.depth_range[0],
self.depth_range[1])
hp_data = hp_data.to(self.device)
depth_image_bgr = colorize_depth(depth_image,
ignore_value=self.depth_range[0])
if mode == 'train':
confidence, depth, rotation = self.model(hp_data)
elif mode == 'val' or mode == 'test':
with torch.no_grad():
confidence, depth, rotation = self.model(hp_data)
confidence_np = confidence[0, ...].cpu().detach().numpy().copy()
confidence_np[confidence_np >= 1] = 1.
confidence_np[confidence_np <= 0] = 0.
confidence_vis = cv2.cvtColor(confidence_np[0, ...] * 255,
cv2.COLOR_GRAY2BGR)
if mode != 'test':
pos_weight = hp_data_gt.detach().numpy().copy()
pos_weight = pos_weight[:, 0, ...]
zeroidx = np.where(pos_weight < 0.5)
nonzeroidx = np.where(pos_weight >= 0.5)
pos_weight[zeroidx] = 0.5
pos_weight[nonzeroidx] = 1.0
pos_weight = torch.from_numpy(pos_weight)
pos_weight = pos_weight.to(self.device)
hp_data_gt = hp_data_gt.to(self.device)
confidence_gt = hp_data_gt[:, 0:1, ...]
rois_list_gt, rois_center_list_gt = find_rois(confidence_gt)
criterion = HPNETLoss(self.use_coords).to(self.device)
if self.model.rois_list is None or rois_list_gt is None:
return None, None
annotated_rois = annotate_rois(self.model.rois_list,
rois_list_gt, annotation_data)
confidence_loss, depth_loss, rotation_loss = criterion(
confidence, hp_data_gt, pos_weight, depth, rotation,
annotated_rois)
if self.train_depth:
loss = confidence_loss + rotation_loss + depth_loss
else:
loss = confidence_loss + rotation_loss
if torch.isnan(loss):
print('loss is nan!!')
self.model = self.prev_model
self.optimizer = torch.optim.Adam(self.model.parameters(),
lr=self.lr,
betas=(0.9, 0.999),
eps=1e-10,
weight_decay=0,
amsgrad=False)
self.optimizer.load_state_dict(
self.prev_optimizer.state_dict())
continue
else:
self.prev_model = copy.deepcopy(self.model)
self.prev_optimizer = copy.deepcopy(self.optimizer)
if mode == 'train':
self.optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(self.model.parameters(), 5)
self.optimizer.step()
axis_gt = depth_image_bgr.copy()
confidence_gt_vis = cv2.cvtColor(
confidence_gt[0, 0, ...].cpu().detach().numpy().copy() *
255, cv2.COLOR_GRAY2BGR)
# Visualize gt axis and roi
for roi, roi_c in zip(rois_list_gt[0], rois_center_list_gt[0]):
if roi.tolist() == [0, 0, 0, 0]:
continue
roi = roi.cpu().detach().numpy().copy()
cx = roi_c[0]
cy = roi_c[1]
depth_and_rotation_gt = get_value_gt([cx, cy],
annotation_data[0])
rotation_gt = depth_and_rotation_gt[1:]
depth_gt_val = depth_and_rotation_gt[0]
unnormalized_depth_gt_val = unnormalize_depth(
depth_gt_val, self.depth_range[0], self.depth_range[1])
hanging_point_pose = np.array(
self.cameramodel.project_pixel_to_3d_ray(
[int(cx), int(cy)])) \
* unnormalized_depth_gt_val * 0.001
if self.use_coords:
rot = quaternion2matrix(rotation_gt),
else:
v = np.matmul(quaternion2matrix(rotation_gt),
[1, 0, 0])
rot = rotation_matrix_from_axis(v, [0, 1, 0], 'xy')
try:
draw_axis(axis_gt, rot, hanging_point_pose,
self.cameramodel.K)
except Exception:
print('Fail to draw axis')
confidence_gt_vis = draw_roi(confidence_gt_vis,
roi,
val=depth_gt_val,
gt=True)
axis_gt = draw_roi(axis_gt, roi, val=depth_gt_val, gt=True)
# Visualize pred axis and roi
axis_pred = depth_image_bgr.copy()
for i, (roi, roi_c) in enumerate(
zip(self.model.rois_list[0],
self.model.rois_center_list[0])):
if roi.tolist() == [0, 0, 0, 0]:
continue
roi = roi.cpu().detach().numpy().copy()
cx = roi_c[0]
cy = roi_c[1]
dep = depth[i].cpu().detach().numpy().copy()
normalized_dep_pred = float(dep)
dep = unnormalize_depth(dep, self.depth_range[0],
self.depth_range[1])
confidence_vis = draw_roi(confidence_vis,
roi,
val=normalized_dep_pred)
axis_pred = draw_roi(axis_pred, roi, val=normalized_dep_pred)
if mode != 'test':
if annotated_rois[i][2]:
confidence_vis = draw_roi(confidence_vis,
annotated_rois[i][0],
val=annotated_rois[i][1][0],
gt=True)
axis_pred = draw_roi(axis_pred,
annotated_rois[i][0],
val=annotated_rois[i][1][0],
gt=True)
hanging_point_pose = np.array(
self.cameramodel.project_pixel_to_3d_ray(
[int(cx), int(cy)])) * float(dep * 0.001)
if self.use_coords:
# have not check this yet
q = rotation[i].cpu().detach().numpy().copy()
q /= np.linalg.norm(q)
rot = quaternion2matrix(q)
else:
v = rotation[i].cpu().detach().numpy()
v /= np.linalg.norm(v)
rot = rotation_matrix_from_axis(v, [0, 1, 0], 'xy')
try:
draw_axis(axis_pred, rot, hanging_point_pose,
self.cameramodel.K)
except Exception:
print('Fail to draw axis')
axis_pred = cv2.cvtColor(axis_pred, cv2.COLOR_BGR2RGB)
confidence_vis = cv2.cvtColor(confidence_vis, cv2.COLOR_BGR2RGB)
if self.config['use_bgr']:
if self.config['use_bgr2gray']:
in_gray = hp_data.cpu().detach().numpy().copy()[0, 1:2,
...] * 255
in_gray = in_gray.transpose(1, 2, 0).astype(np.uint8)
in_gray = cv2.cvtColor(in_gray, cv2.COLOR_GRAY2RGB)
in_gray = in_gray.transpose(2, 0, 1)
in_img = in_gray
else:
in_bgr = hp_data.cpu().detach().numpy().copy()[
0, 3:, ...].transpose(1, 2, 0)
in_rgb = cv2.cvtColor(in_bgr, cv2.COLOR_BGR2RGB).transpose(
2, 0, 1)
in_img = in_rgb
if mode != 'test':
confidence_loss_sum += confidence_loss.item()
axis_gt = cv2.cvtColor(axis_gt, cv2.COLOR_BGR2RGB)
confidence_gt_vis = cv2.cvtColor(confidence_gt_vis,
cv2.COLOR_BGR2RGB)
if rotation_loss.item() > 0:
depth_loss_sum += depth_loss.item()
rotation_loss_sum += rotation_loss.item()
loss_sum = loss_sum \
+ confidence_loss.item() \
+ rotation_loss.item()
rotation_loss_count += 1
if np.mod(index, 1) == 0:
print(
'epoch {}, {}/{},{} loss is confidence:{} rotation:{} depth:{}'
.format( # noqa
self.epo, index, len(dataloader), mode,
confidence_loss.item(), rotation_loss.item(),
depth_loss.item()))
self.vis.images(
[axis_gt.transpose(2, 0, 1),
axis_pred.transpose(2, 0, 1)],
win='{} axis'.format(mode),
opts=dict(title='{} axis'.format(mode)))
self.vis.images(
[
confidence_gt_vis.transpose(2, 0, 1),
confidence_vis.transpose(2, 0, 1)
],
win='{}_confidence_roi'.format(mode),
opts=dict(title='{} confidence(GT, Pred)'.format(mode)))
if self.config['use_bgr']:
self.vis.images([in_img],
win='{} in_gray'.format(mode),
opts=dict(title='{} in_gray'.format(mode)))
else:
if self.config['use_bgr']:
self.vis.images(
[
in_img,
confidence_vis.transpose(2, 0, 1),
axis_pred.transpose(2, 0, 1)
],
win='{}-{}'.format(mode, index),
opts=dict(
title='{}-{} hanging_point_depth (pred)'.format(
mode, index)))
else:
self.vis.images(
[
confidence_vis.transpose(2, 0, 1),
axis_pred.transpose(2, 0, 1)
],
win='{}-{}'.format(mode, index),
opts=dict(
title='{}-{} hanging_point_depth (pred)'.format(
mode, index)))
if np.mod(index, 1000) == 0:
save_file = osp.join(
self.save_dir,
'hpnet_latestmodel_' + self.time_now + '.pt')
print('save {}'.format(save_file))
torch.save(self.model.state_dict(),
save_file,
_use_new_zipfile_serialization=False)
if mode != 'test':
if len(dataloader) > 0:
avg_confidence_loss\
= confidence_loss_sum / len(dataloader)
if rotation_loss_count > 0:
avg_rotation_loss\
= rotation_loss_sum / rotation_loss_count
avg_depth_loss\
= depth_loss_sum / rotation_loss_count
avg_loss\
= loss_sum / rotation_loss_count
else:
avg_rotation_loss = 1e10
avg_depth_loss = 1e10
avg_loss = 1e10
else:
avg_loss = loss_sum
avg_confidence_loss = confidence_loss_sum
avg_rotation_loss = rotation_loss_sum
avg_depth_loss = rotation_loss_sum
self.vis.line(X=np.array([self.epo]),
Y=np.array([avg_confidence_loss]),
opts={'title': 'confidence'},
win='confidence loss',
name='{}_confidence_loss'.format(mode),
update='append')
if rotation_loss_count > 0:
self.vis.line(X=np.array([self.epo]),
Y=np.array([avg_rotation_loss]),
opts={'title': 'rotation loss'},
win='rotation loss',
name='{}_rotation_loss'.format(mode),
update='append')
self.vis.line(X=np.array([self.epo]),
Y=np.array([avg_depth_loss]),
opts={'title': 'depth loss'},
win='depth loss',
name='{}_depth_loss'.format(mode),
update='append')
self.vis.line(X=np.array([self.epo]),
Y=np.array([avg_loss]),
opts={'title': 'loss'},
win='loss',
name='{}_loss'.format(mode),
update='append')
if mode == 'val':
if np.mod(self.epo, self.save_model_interval) == 0:
save_file = osp.join(
self.save_dir,
'hpnet_latestmodel_' + self.time_now + '.pt')
print('save {}'.format(save_file))
torch.save(self.model.state_dict(),
save_file,
_use_new_zipfile_serialization=False)
if self.best_loss > avg_loss:
print('update best model {} -> {}'.format(
self.best_loss, avg_loss))
self.best_loss = avg_loss
save_file = osp.join(
self.save_dir,
'hpnet_bestmodel_' + self.time_now + '.pt')
print('save {}'.format(save_file))
# For ros(python 2, torch 1.4)
torch.save(self.model.state_dict(),
save_file,
_use_new_zipfile_serialization=False)
'--input-dir', '-i', type=str,
help='input_dir',
# default='/home/kosuke55/catkin_ws/src/pose_annotation_tool/annotation_obj') # noqa
default='/media/kosuke55/SANDISK/meshdata/ycb_pouring_object_16/textured_urdf/annotation_obj') # noqa
args = parser.parse_args()
paths = list(sorted(Path(args.input_dir).glob('*.txt')))
for path in paths:
input_file = str(path)
print(input_file)
contact_points = {'contact_points': [],
'labels': [],
'urdf_file': str(path.with_suffix('.urdf'))}
with open(input_file)as f:
for line in f.readlines():
values = [float(v) for v in line.split()]
label = int(values[0])
pos = values[1:4]
quaternion = values[4:]
matrix = quaternion2matrix(quaternion)
contact_points['contact_points'].append(
np.vstack([pos, matrix]).tolist())
contact_points['labels'].append(label)
print(line)
save_json(str(path.with_suffix('.json')), contact_points)
# print(lines)
