Commit b0ca8b5c authored by Stanislav Pidhorskyi's avatar Stanislav Pidhorskyi Committed by Facebook GitHub Bot
Browse files

Make torch script compatible

Summary:
Sequence is not TorchScript compatible, replacing with a List.

Also the original annotation `Optional[Sequence[str]]` was not correct, as implementation was actially allowing usage of `str` type.

Wrapped `project_points` with decorator `torch.jit.ignore` as it violates TorchScript too much (uses sets, changes types of variables, uses None with other types aggressively)

Differential Revision: D63444533

fbshipit-source-id: 46604cbd239ed8b9051ad371779fb81106987dc5
parent e2335d06
...@@ -3,7 +3,7 @@ ...@@ -3,7 +3,7 @@
# This source code is licensed under the MIT license found in the # This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree. # LICENSE file in the root directory of this source tree.
from typing import Optional, Sequence, Tuple from typing import List, Optional, Tuple, Union
import torch as th import torch as th
from drtk.utils import project_points from drtk.utils import project_points
...@@ -17,7 +17,7 @@ def transform( ...@@ -17,7 +17,7 @@ def transform(
princpt: Optional[th.Tensor] = None, princpt: Optional[th.Tensor] = None,
K: Optional[th.Tensor] = None, K: Optional[th.Tensor] = None,
Rt: Optional[th.Tensor] = None, Rt: Optional[th.Tensor] = None,
distortion_mode: Optional[Sequence[str]] = None, distortion_mode: Optional[Union[List[str], str]] = None,
distortion_coeff: Optional[th.Tensor] = None, distortion_coeff: Optional[th.Tensor] = None,
fov: Optional[th.Tensor] = None, fov: Optional[th.Tensor] = None,
) -> th.Tensor: ) -> th.Tensor:
...@@ -72,7 +72,7 @@ def transform_with_v_cam( ...@@ -72,7 +72,7 @@ def transform_with_v_cam(
princpt: Optional[th.Tensor] = None, princpt: Optional[th.Tensor] = None,
K: Optional[th.Tensor] = None, K: Optional[th.Tensor] = None,
Rt: Optional[th.Tensor] = None, Rt: Optional[th.Tensor] = None,
distortion_mode: Optional[Sequence[str]] = None, distortion_mode: Optional[Union[List[str], str]] = None,
distortion_coeff: Optional[th.Tensor] = None, distortion_coeff: Optional[th.Tensor] = None,
fov: Optional[th.Tensor] = None, fov: Optional[th.Tensor] = None,
) -> Tuple[th.Tensor, th.Tensor]: ) -> Tuple[th.Tensor, th.Tensor]:
......
...@@ -3,7 +3,7 @@ ...@@ -3,7 +3,7 @@
# This source code is licensed under the MIT license found in the # This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree. # LICENSE file in the root directory of this source tree.
from typing import Optional, Sequence, Set, Tuple, Union from typing import List, Optional, Set, Tuple, Union
import numpy as np import numpy as np
import torch as th import torch as th
...@@ -321,13 +321,14 @@ def estimate_fisheye_fov(D: Union[np.ndarray, th.Tensor]) -> th.Tensor: ...@@ -321,13 +321,14 @@ def estimate_fisheye_fov(D: Union[np.ndarray, th.Tensor]) -> th.Tensor:
return fov return fov
@th.jit.ignore
def project_points( def project_points(
v: th.Tensor, v: th.Tensor,
campos: th.Tensor, campos: th.Tensor,
camrot: th.Tensor, camrot: th.Tensor,
focal: th.Tensor, focal: th.Tensor,
princpt: th.Tensor, princpt: th.Tensor,
distortion_mode: Optional[Sequence[str]] = None, distortion_mode: Optional[Union[List[str], str]] = None,
distortion_coeff: Optional[th.Tensor] = None, distortion_coeff: Optional[th.Tensor] = None,
fov: Optional[th.Tensor] = None, fov: Optional[th.Tensor] = None,
) -> Tuple[th.Tensor, th.Tensor]: ) -> Tuple[th.Tensor, th.Tensor]:
...@@ -430,13 +431,14 @@ def project_points( ...@@ -430,13 +431,14 @@ def project_points(
return v_pix, v_cam return v_pix, v_cam
@th.jit.ignore
def project_points_grad( def project_points_grad(
v_grad: th.Tensor, v_grad: th.Tensor,
v: th.Tensor, v: th.Tensor,
campos: th.Tensor, campos: th.Tensor,
camrot: th.Tensor, camrot: th.Tensor,
focal: th.Tensor, focal: th.Tensor,
distortion_mode: Optional[Sequence[str]] = None, distortion_mode: Optional[Union[List[str], str]] = None,
distortion_coeff: Optional[th.Tensor] = None, distortion_coeff: Optional[th.Tensor] = None,
) -> th.Tensor: ) -> th.Tensor:
"""Computes the gradient of projected (pixel-space) vertex positions with """Computes the gradient of projected (pixel-space) vertex positions with
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment