changes eval

tools/eval_linemod.py

@@ -21,25 +21,146 @@ from lib.loss import Loss
 from lib.loss_refiner import Loss_refine
 from lib.transformations import euler_matrix, quaternion_matrix, quaternion_from_matrix
 from lib.knn.__init__ import KNearestNeighbor
+from PIL import Image
+import cv2
+import math
+
+
+def draw_axis(img, rotation_vec, t, K, scale=0.1, dist=None):
+    """
+    Draw a 6dof axis (XYZ -> RGB) in the given rotation and translation
+    :param img - rgb numpy array
+    :rotation_vec - euler rotations, numpy array of length 3,
+                    use cv2.Rodrigues(R)[0] to convert from rotation matrix
+    :t - 3d translation vector, in meters (dtype must be float)
+    :K - intrinsic calibration matrix , 3x3
+    :scale - factor to control the axis lengths
+    :dist - optional distortion coefficients, numpy array of length 4. If None distortion is ignored.
+    """
+    # img = img.astype(np.float32)
+    dist = np.zeros(4, dtype=float) if dist is None else dist
+    points = scale * np.float32([[1, 0, 0], [0, 1, 0], [0, 0, 1], [0, 0, 0]]).reshape(-1, 3)
+    axis_points, _ = cv2.projectPoints(points, rotation_vec, t, K, dist)
+    axis_points = np.asarray(axis_points, dtype='int')
+    img = cv2.line(img, tuple(axis_points[3].ravel()), tuple(axis_points[0].ravel()), (255, 0, 0), 3)
+    img = cv2.line(img, tuple(axis_points[3].ravel()), tuple(axis_points[1].ravel()), (0, 255, 0), 3)
+    img = cv2.line(img, tuple(axis_points[3].ravel()), tuple(axis_points[2].ravel()), (0, 0, 255), 3)
+    return img
+
+
+def draw_cube(img, rotation_vec, t, K, size, scale=0.1, dist=None):
+    # sizes donne un tableau avec les tailles selon chaque axe
+    # img = img.astype(np.float32)
+
+    dist = np.zeros(4, dtype=float) if dist is None else dist
+    x, y, z = size, size, size
+    points = scale * np.float32([[-x, -y, z], [-x, y, z], [x, y, z], [x, -y, z],
+                                 [-x, -y, -z], [-x, y, -z], [x, y, -z], [x, -y, -z]])
+    axis_points, _ = cv2.projectPoints(points, rotation_vec, t, K, dist)
+    axis_points = np.asarray(axis_points, dtype='int')
+
+    print(axis_points)
+    img = cv2.drawContours(img, [axis_points[:4]], -1, (0, 255, 0), 3)
+
+    for i, j in zip(range(4), range(4, 8)):
+        img = cv2.line(img, tuple(axis_points[i].ravel()), tuple(axis_points[j].ravel()), (255, 0, 0), 3)
+
+    # draw top layer in red color
+    img = cv2.drawContours(img, [axis_points[4:]], -1, (0, 0, 255), 3)
+    return img
+
+
+def generate_pose(rot, trans):
+    pose = quaternion_to_rotation_matrix(rot)
+    pose[3, :3] = trans
+    return pose
+
+
+def get_R_t(mat):
+    return mat[:3, :3], mat[:3, 3]
+
+
+def quaternion_to_rotation_matrix(quat):
+    q = quat.copy()
+    n = np.dot(q, q)
+    if n < np.finfo(q.dtype).eps:
+        return np.identity(4)
+    q = q * np.sqrt(2.0 / n)
+    q = np.outer(q, q)
+    rot_matrix = np.array(
+        [[1.0 - q[2, 2] - q[3, 3], q[1, 2] + q[3, 0], q[1, 3] - q[2, 0], 0.0],
+         [q[1, 2] - q[3, 0], 1.0 - q[1, 1] - q[3, 3], q[2, 3] + q[1, 0], 0.0],
+         [q[1, 3] + q[2, 0], q[2, 3] - q[1, 0], 1.0 - q[1, 1] - q[2, 2], 0.0],
+         [0.0, 0.0, 0.0, 1.0]],
+        dtype=q.dtype)
+    return rot_matrix
+
+
+def euler_from_quaternion(x, y, z, w):
+    """
+        Convert a quaternion into euler angles (roll, pitch, yaw)
+        roll is rotation around x in radians (counterclockwise)
+        pitch is rotation around y in radians (counterclockwise)
+        yaw is rotation around z in radians (counterclockwise)
+        """
+    t0 = +2.0 * (w * x + y * z)
+    t1 = +1.0 - 2.0 * (x * x + y * y)
+    roll_x = math.atan2(t0, t1)
+
+    t2 = +2.0 * (w * y - z * x)
+    t2 = +1.0 if t2 > +1.0 else t2
+    t2 = -1.0 if t2 < -1.0 else t2
+    pitch_y = math.asin(t2)
+
+    t3 = +2.0 * (w * z + x * y)
+    t4 = +1.0 - 2.0 * (y * y + z * z)
+    yaw_z = math.atan2(t3, t4)
+
+    return roll_x, pitch_y, yaw_z  # in radians
+
+
+flip_rot_matrix = np.array(
+    [[1, 0, 0, 0],
+     [0, -1, 0, 0],
+     [0, 0, -1, 0],
+     [0, 0, 0, 1]],
+    dtype=np.float32)
 
 parser = argparse.ArgumentParser()
-parser.add_argument('--dataset_root', type=str, default = '', help='dataset root dir')
-parser.add_argument('--model', type=str, default = '',  help='resume PoseNet model')
-parser.add_argument('--refine_model', type=str, default = '',  help='resume PoseRefineNet model')
+parser.add_argument('--dataset_root', type=str, default='', help='dataset root dir')
+parser.add_argument('--model', type=str, default='', help='resume PoseNet model')
+parser.add_argument('--refine_model', type=str, default='', help='resume PoseRefineNet model')
 opt = parser.parse_args()
 
-num_objects = 13
-objlist = [1, 2, 4, 5, 6, 8, 9, 10, 11, 12, 13, 14, 15]
+num_objects = 8
+objlist = [1, 2, 3, 4, 5, 6, 7, 8]
 num_points = 500
 iteration = 4
 bs = 1
+cam_cx = 320.25  # TODO
+cam_cy = 240.33333333333331  # TODO
+cam_fx = 543.2527222420504  # TODO
+cam_fy = 724.3369629894005  # TODO
+
+# ["banana1", "kiwi1", "pear2", "strawberry1", "apricot", "orange2", "peach1", "lemon2", "apple2" ]
+map_id_obj = {
+    1: 'banana1',
+    2: 'kiwi1',
+    3: 'pear2',
+    4: 'strawberry1',
+    5: 'orange2',
+    6: 'peach1',
+    7: 'lemon2',
+    8: 'apple2',
+}
+K = np.array([[cam_fx, 0, cam_cx], [0, cam_fy, cam_cy], [0, 0, 1]])
 dataset_config_dir = 'datasets/linemod/dataset_config'
 output_result_dir = 'experiments/eval_result/linemod'
 knn = KNearestNeighbor(1)
 
-estimator = PoseNet(num_points = num_points, num_obj = num_objects)
+estimator = PoseNet(num_points=num_points, num_obj=num_objects)
 estimator.cuda()
-refiner = PoseRefineNet(num_points = num_points, num_obj = num_objects)
+refiner = PoseRefineNet(num_points=num_points, num_obj=num_objects)
 refiner.cuda()
 estimator.load_state_dict(torch.load(opt.model))
 refiner.load_state_dict(torch.load(opt.refine_model))
@@ -64,20 +185,19 @@ print(diameter)
 success_count = [0 for i in range(num_objects)]
 num_count = [0 for i in range(num_objects)]
 fw = open('{0}/eval_result_logs.txt'.format(output_result_dir), 'w')
-
 for i, data in enumerate(testdataloader, 0):
     points, choose, img, target, model_points, idx = data
+
     if len(points.size()) == 2:
         print('No.{0} NOT Pass! Lost detection!'.format(i))
         fw.write('No.{0} NOT Pass! Lost detection!\n'.format(i))
         continue
     points, choose, img, target, model_points, idx = Variable(points).cuda(), \
-                                                     Variable(choose).cuda(), \
-                                                     Variable(img).cuda(), \
-                                                     Variable(target).cuda(), \
-                                                     Variable(model_points).cuda(), \
-                                                     Variable(idx).cuda()
-
+        Variable(choose).cuda(), \
+        Variable(img).cuda(), \
+        Variable(target).cuda(), \
+        Variable(model_points).cuda(), \
+        Variable(idx).cuda()
     pred_r, pred_t, pred_c, emb = estimator(img, points, choose, idx)
     pred_r = pred_r / torch.norm(pred_r, dim=2).view(1, num_points, 1)
     pred_c = pred_c.view(bs, num_points)
@@ -89,11 +209,14 @@ for i, data in enumerate(testdataloader, 0):
     my_pred = np.append(my_r, my_t)
 
     for ite in range(0, iteration):
-        T = Variable(torch.from_numpy(my_t.astype(np.float32))).cuda().view(1, 3).repeat(num_points, 1).contiguous().view(1, num_points, 3)
+        T = Variable(torch.from_numpy(my_t.astype(np.float32))).cuda().view(1, 3).repeat(num_points,
+                                                                                         1).contiguous().view(1,
+                                                                                                              num_points,
+                                                                                                              3)
         my_mat = quaternion_matrix(my_r)
         R = Variable(torch.from_numpy(my_mat[:3, :3].astype(np.float32))).cuda().view(1, 3, 3)
         my_mat[0:3, 3] = my_t
-        
+
         new_points = torch.bmm((points - T), R).contiguous()
         pred_r, pred_t = refiner(new_points, emb, idx)
         pred_r = pred_r.view(1, 1, -1)
@@ -120,6 +243,14 @@ for i, data in enumerate(testdataloader, 0):
     pred = np.dot(model_points, my_r.T) + my_t
     target = target[0].cpu().detach().numpy()
 
+    matrix = np.zeros((3, 4))
+    matrix[:3, :3] = my_r
+    matrix[:3, 3] = my_t
+    # np.save(f"results2/result_{i}.npy", matrix)
+    """
+   with open(f"results/result_{temp_id}.txt", "w") as f:
+        f.write(f"{matrix}")
+    """
     if idx[0].item() in sym_list:
         pred = torch.from_numpy(pred.astype(np.float32)).cuda().transpose(1, 0).contiguous()
         target = torch.from_numpy(target.astype(np.float32)).cuda().transpose(1, 0).contiguous()