Some refactoring

examples/efficientpose/losses_refactored.py

@@ -12,7 +12,8 @@ def __init__(self, object_id, translation_priors, data_path,
         self.object_id = object_id
         self.translation_priors = translation_priors
         self.num_pose_dims = num_pose_dims
-        self.translation_scale_norm = translation_scale_norm
+        self.tz_scale = tf.convert_to_tensor(translation_scale_norm,
+                                             dtype=tf.float32)
         self.model_path = data_path + model_path + 'obj_' + object_id + '.ply'
         self.model_points = self._load_model_file()
         self.model_points = self._filter_model_points(target_num_points)
@@ -41,11 +42,10 @@ def _filter_model_points(self, target_num_points):
         return tf.convert_to_tensor(points)
 
     def _compute_translation(self, translation_raw_pred, scale):
-        camera_parameter = self._compute_camera_parameter(
-            scale, LINEMOD_CAMERA_MATRIX)
+        camera_matrix = tf.convert_to_tensor(LINEMOD_CAMERA_MATRIX)
         translation_pred = self._regress_translation(translation_raw_pred)
-        translation_pred = self._compute_tx_ty(translation_pred,
-                                               camera_parameter)
+        translation_pred = self._compute_tx_ty(translation_pred, camera_matrix,
+                                               scale)
         return translation_pred
 
     def _regress_translation(self, translation_raw):
@@ -58,20 +58,13 @@ def _regress_translation(self, translation_raw):
         translations_predicted = tf.concat([x, y, Tz], axis=-1)
         return translations_predicted
 
-    def _compute_camera_parameter(self, image_scale, camera_matrix):
-        camera_parameter = tf.convert_to_tensor(
-            [camera_matrix[0, 0], camera_matrix[1, 1], camera_matrix[0, 2],
-             camera_matrix[1, 2], self.translation_scale_norm, image_scale])
-        return camera_parameter
+    def _compute_tx_ty(self, translation_xy_Tz, camera_matrix, scale):
+        fx, fy = camera_matrix[0, 0], camera_matrix[1, 1]
+        px, py = camera_matrix[0, 2], camera_matrix[1, 2]
 
-    def _compute_tx_ty(self, translation_xy_Tz, camera_parameter):
-        fx, fy = camera_parameter[0], camera_parameter[1],
-        px, py = camera_parameter[2], camera_parameter[3],
-        tz_scale, image_scale = camera_parameter[4], camera_parameter[5]
-
-        x = translation_xy_Tz[:, :, 0] / image_scale
-        y = translation_xy_Tz[:, :, 1] / image_scale
-        tz = translation_xy_Tz[:, :, 2] * tz_scale
+        x = translation_xy_Tz[:, :, 0] / scale
+        y = translation_xy_Tz[:, :, 1] / scale
+        tz = translation_xy_Tz[:, :, 2] * self.tz_scale
 
         x = x - px
         y = y - py
@@ -95,8 +88,10 @@ def _separate_axis_from_angle(self, axis_angle):
     def _rotate(self, point, axis, angle):
         cos_angle = tf.cos(angle)
         axis_dot_point = self._dot(axis, point)
-        return (point * cos_angle + self._cross(axis, point) *
-                tf.sin(angle) + axis * axis_dot_point * (1.0 - cos_angle))
+        points_rotated = (
+            point * cos_angle + self._cross(axis, point) * tf.sin(angle) +
+            axis * axis_dot_point * (1.0 - cos_angle))
+        return points_rotated
 
     def _dot(self, vector1, vector2, axis=-1, keepdims=True):
         return tf.reduce_sum(input_tensor=vector1 * vector2,
@@ -115,8 +110,8 @@ def _cross(self, vector1, vector2):
     def _calc_sym_distances(self, sym_points_pred, sym_points_true):
         sym_points_pred = sym_points_pred[:, :, tf.newaxis]
         sym_points_true = sym_points_true[:, tf.newaxis]
-        distances = tf.reduce_min(tf.norm(
-            sym_points_pred - sym_points_true, axis=-1), axis=-1)
+        distances = tf.reduce_min(tf.norm(sym_points_pred - sym_points_true,
+                                          axis=-1), axis=-1)
         return tf.reduce_mean(distances, axis=-1)
 
     def _calc_asym_distances(self, asym_points_pred, asym_points_target):