Add LUT to project_fisheye_distort in drtk

Summary:
The purpos of this stack is to enable LUT in rosetta2. To do this:
1. extend project_fisheye_distort_62 to support LUT
2. propagate lut in DRTK rendering functions
3. load the LUT in HMCCalibrationSCIOVRS
4. register LUT as an asset for loading
5. extend the rendering pipeline in rosetta2 with  LUT
6. hook the care cmd string generated in fblearner pipeline with assets defined in 4.

This diff implements 1. The look up table contains two parts:
1. lut_vector_field: N x 2 X H_lut x W_lut stores the offset for each camera, e.g. 2 x 41 x 41 for FETA camera
2. lut_spacing: N x 2 stores the spacing of each grid in lut_vector_field, e.g. 10.0 for FETA 400 x 400 camera. Usually spacing is the same for both with and height directions in calibration file but might be different if we resize the field.

The offset for each pixel is found by grid sampling

Reviewed By: patricksnape

Differential Revision: D62097271

fbshipit-source-id: 1655cb1dfc8d5e34bc97c4bab16b07a9ee2cf636

drtk/utils/projection.py

@@ -160,6 +160,8 @@ def project_fisheye_distort_62(
     princpt: th.Tensor,
     D: th.Tensor,
     fov: Optional[th.Tensor] = None,
+    lut_vector_field: Optional[th.Tensor] = None,
+    lut_spacing: Optional[th.Tensor] = None,
 ) -> th.Tensor:
     """Project camera-space points to distort pixel-space points using the
     Fisheye62 distortion model. See the perception camera model implementation
@@ -170,6 +172,8 @@ def project_fisheye_distort_62(
     princpt:            N x 2
     D:                  N x 8
     fov:                N x 1
+    lut_vector_field:   N x 2 x H_lut x W_lut
+    lut_spacing:        N x 2
     """
     assert (
         D.shape[1] == 8
@@ -230,6 +234,38 @@ def project_fisheye_distort_62(
 
     v_pix_dist = (focal[:, None] @ v_proj_dist[..., None])[..., 0] + princpt[:, None]
 
+    if lut_vector_field is not None:
+        assert (
+            lut_spacing is not None
+        ), "lookup table spacing must be provided along with vector field"
+        # Normalize pixel position to [-1, 1] range for torch.nn.functional.grid_sample
+        normalized_pixel_pos = v_pix_dist / lut_spacing[:, None, :]
+        lut_col, lut_row = lut_vector_field.shape[2:4]
+        normalized_pixel_pos[..., 0] = (
+            normalized_pixel_pos[..., 0] / (lut_col - 1) * 2.0 - 1.0
+        )
+        normalized_pixel_pos[..., 1] = (
+            normalized_pixel_pos[..., 1] / (lut_row - 1) * 2.0 - 1.0
+        )
+        # do grid sample
+        pixel_offset = th.nn.functional.grid_sample(
+            input=lut_vector_field,
+            grid=normalized_pixel_pos.unsqueeze(1),
+            align_corners=True,
+        ).permute(0, 2, 3, 1)[:, 0, ...]
+
+        # mask for pixels out of the lookup table, whose offset need to be set to 0.0
+        out_bound_mask = (
+            (normalized_pixel_pos[..., 0] < -1.0)
+            | (normalized_pixel_pos[..., 0] > 1.0)
+            | (normalized_pixel_pos[..., 1] < -1.0)
+            | (normalized_pixel_pos[..., 1] > 1.0)
+        )
+
+        pixel_offset[out_bound_mask] = 0.0
+
+        v_pix_dist = v_pix_dist + pixel_offset
+
     return v_pix_dist