Kornia adaption (#552)

pcdet/models/__init__.py

@@ -2,7 +2,7 @@ from collections import namedtuple
 
 import numpy as np
 import torch
-import kornia
+from kornia.utils.image import image_to_tensor
 
 from .detectors import build_detector
 
@@ -21,7 +21,7 @@ def load_data_to_gpu(batch_dict):
         elif key in ['frame_id', 'metadata', 'calib']:
             continue
         elif key in ['images']:
-            batch_dict[key] = kornia.image_to_tensor(val).float().cuda().contiguous()
+            batch_dict[key] = image_to_tensor(val).float().cuda().contiguous()
         elif key in ['image_shape']:
             batch_dict[key] = torch.from_numpy(val).int().cuda()
         else:

pcdet/models/backbones_3d/vfe/image_vfe_modules/f2v/frustum_grid_generator.py

 import torch
 import torch.nn as nn
-import kornia
+try:
+    from kornia.utils.grid import create_meshgrid3d
+    from kornia.geometry.linalg import transform_points
+except Exception as e:
+    # Note: Kornia team will fix this import issue to try to allow the usage of lower torch versions.
+    raise ImportError("It is recommended to use torch version greater than 1.2 to use kornia properly.")
 
 from pcdet.utils import transform_utils
 
@@ -30,7 +35,7 @@ class FrustumGridGenerator(nn.Module):
 
         # Create voxel grid
         self.depth, self.width, self.height = self.grid_size.int()
-        self.voxel_grid = kornia.utils.create_meshgrid3d(depth=self.depth,
+        self.voxel_grid = create_meshgrid3d(depth=self.depth,
                                                          height=self.height,
                                                          width=self.width,
                                                          normalized_coordinates=False)
@@ -85,7 +90,7 @@ class FrustumGridGenerator(nn.Module):
         voxel_grid = voxel_grid.repeat_interleave(repeats=B, dim=0)
 
         # Transform to camera frame
-        camera_grid = kornia.transform_points(trans_01=trans, points_1=voxel_grid)
+        camera_grid = transform_points(trans_01=trans, points_1=voxel_grid)
 
         # Project to image
         I_C = I_C.reshape(B, 1, 1, 3, 4)

pcdet/models/backbones_3d/vfe/image_vfe_modules/ffn/ddn/ddn_template.py

@@ -7,7 +7,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import torchvision
-import kornia
+from kornia.enhance.normalize import normalize
 
 
 class DDNTemplate(nn.Module):
@@ -151,7 +151,7 @@ class DDNTemplate(nn.Module):
             mask = torch.isnan(x)
 
             # Match ResNet pretrained preprocessing
-            x = kornia.normalize(x, mean=self.norm_mean, std=self.norm_std)
+            x = normalize(x, mean=self.norm_mean, std=self.norm_std)
 
             # Make padded pixels = 0
             x[mask] = 0

pcdet/models/backbones_3d/vfe/image_vfe_modules/ffn/ddn_loss/ddn_loss.py

 import torch
 import torch.nn as nn
-import kornia
+from kornia.losses.focal import FocalLoss
 
 from .balancer import Balancer
 from pcdet.utils import transform_utils
@@ -37,7 +37,7 @@ class DDNLoss(nn.Module):
         # Set loss function
         self.alpha = alpha
         self.gamma = gamma
-        self.loss_func = kornia.losses.FocalLoss(alpha=self.alpha, gamma=self.gamma, reduction="none")
+        self.loss_func = FocalLoss(alpha=self.alpha, gamma=self.gamma, reduction="none")
         self.weight = weight
 
     def forward(self, depth_logits, depth_maps, gt_boxes2d):

pcdet/utils/transform_utils.py

 import math
 import torch
-import kornia
+from kornia.geometry.conversions import (
+    convert_points_to_homogeneous,
+    convert_points_from_homogeneous,
+)
 
 
 def project_to_image(project, points):
@@ -14,14 +17,14 @@ def project_to_image(project, points):
         points_depth [torch.Tensor(...)]: Depth of each point
     """
     # Reshape tensors to expected shape
-    points = kornia.convert_points_to_homogeneous(points)
+    points = convert_points_to_homogeneous(points)
     points = points.unsqueeze(dim=-1)
     project = project.unsqueeze(dim=1)
 
     # Transform points to image and get depths
     points_t = project @ points
     points_t = points_t.squeeze(dim=-1)
-    points_img = kornia.convert_points_from_homogeneous(points_t)
+    points_img = convert_points_from_homogeneous(points_t)
     points_depth = points_t[..., -1] - project[..., 2, 3]
 
     return points_img, points_depth

setup.py

 import os
+import sys
 import subprocess
 
 from setuptools import find_packages, setup
+from setuptools.command.install import install
+# TODO: This is a bit buggy since it requires torch before installing torch.
 from torch.utils.cpp_extension import BuildExtension, CUDAExtension
 
 
@@ -27,6 +30,15 @@ def write_version_to_file(version, target_file):
         print('__version__ = "%s"' % version, file=f)
 
 
+class PostInstallation(install):
+    """Post-installation for installation mode."""
+    def run(self):
+        install.run(self)
+        # Note: buggy for kornia==0.5.3 and it will be fixed in the next version.
+        # Set kornia to 0.5.2 temporarily
+        subprocess.call([sys.executable, '-m', 'pip', 'install', 'kornia==0.5.2', '--no-dependencies'])
+
+
 if __name__ == '__main__':
     version = '0.3.0+%s' % get_git_commit_number()
     write_version_to_file(version, 'pcdet/version.py')
@@ -48,7 +60,12 @@ if __name__ == '__main__':
         author_email='shaoshuaics@gmail.com',
         license='Apache License 2.0',
         packages=find_packages(exclude=['tools', 'data', 'output']),
-        cmdclass={'build_ext': BuildExtension},
+        cmdclass={
+            'build_ext': BuildExtension,
+            'install': PostInstallation,
+            # Post installation cannot be done. ref: https://github.com/pypa/setuptools/issues/1936.
+            # 'develop': PostInstallation,
+        },
         ext_modules=[
             make_cuda_ext(
                 name='iou3d_nms_cuda',