Skip to content

sdfest.py

cpas_toolbox.cpas_methods.sdfest

This module defines SDFEst interface.

Method is described in SDFEst: Categorical Pose and Shape Estimation of Objects From RGB-D Using Signed Distance Fields, Bruns, 2022.

Implementation based on https://github.com/roym899/sdfest/

SDFEst

Bases: CPASMethod

Wrapper class for SDFEst.

Source code in cpas_toolbox/cpas_methods/sdfest.py 
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
class SDFEst(CPASMethod):
    """Wrapper class for SDFEst."""

    class Config(TypedDict):
        """Configuration dictionary for SDFEst.

        All keys supported by SDFPipeline are supported and will overwrite config
        contained in sdfest_... files. The keys specified here are used by this
        script only.

        The two keys sdfest_..._config_files  will be parsed with SDFEst install
        directory as part of the search paths. This allows to use the default config
        that comes with SDFEst installation.

        Attributes:
            sdfest_default_config_file: Default configuration file loaded first.
            sdfest_category_config_files: Per-category configuration file loaded second.
            device: Device used for computation.
            num_points: Numbner of points extracted from mesh.
            prior: Prior distribution to modify orientation distribution.
            visualize_optimization:
                Whether to show additional optimization visualization.
        """

    default_config: Config = {
        "sdfest_default_config_file": "estimation/configs/default.yaml",
        "sdfest_category_config_files": {
            "bottle": "estimation/configs/models/bottle.yaml",
            "bowl": "estimation/configs/models/bowl.yaml",
            "laptop": "estimation/configs/models/laptop.yaml",
            "can": "estimation/configs/models/can.yaml",
            "camera": "estimation/configs/models/camera.yaml",
            "mug": "estimation/configs/models/mug.yaml",
        },
    }

    def __init__(self, config: Config, camera: camera_utils.Camera) -> None:
        """Initialize and load SDFEst models.

        Configuration loaded in following order
            sdfest_default_config_file -> sdfest_category_config_files -> all other keys
        I.e., keys specified directly will take precedence over keys specified in
        default file.
        """
        default_config = yoco.load_config_from_file(
            config["sdfest_default_config_file"],
            search_paths=[".", "~/.sdfest/", sdfest.__path__[0]],
        )

        self._pipeline_dict = {}  # maps category to category-specific pipeline
        self._device = config["device"]
        self._visualize_optimization = config["visualize_optimization"]
        self._num_points = config["num_points"]
        self._prior = config["prior"] if "prior" in config else None

        # create per-categry models
        for category_str in config["sdfest_category_config_files"].keys():
            category_config = yoco.load_config_from_file(
                config["sdfest_category_config_files"][category_str],
                current_dict=default_config,
                search_paths=[".", "~/.sdfest/", sdfest.__path__[0]],
            )
            category_config = yoco.load_config(config, category_config)
            self._pipeline_dict[category_str] = SDFPipeline(category_config)
            self._pipeline_dict[category_str].cam = camera

    def inference(
        self,
        color_image: torch.Tensor,
        depth_image: torch.Tensor,
        instance_mask: torch.Tensor,
        category_str: str,
    ) -> PredictionDict:
        """See CPASMethod.inference."""
        # skip unsupported category
        if category_str not in self._pipeline_dict:
            return {
                "position": torch.tensor([0, 0, 0]),
                "orientation": torch.tensor([0, 0, 0, 1]),
                "extents": torch.tensor([1, 1, 1]),
                "reconstructed_pointcloud": torch.tensor([[0, 0, 0]]),
                "reconstructed_mesh": None,
            }

        pipeline = self._pipeline_dict[category_str]

        # move inputs to device
        color_image = color_image.to(self._device)
        depth_image = depth_image.to(self._device, copy=True)
        instance_mask = instance_mask.to(self._device)

        if self._prior is not None:
            prior = torch.tensor(self._prior[category_str], device=self._device)
            prior /= torch.sum(prior)
        else:
            prior = None

        position, orientation, scale, shape = pipeline(
            depth_image,
            instance_mask,
            color_image,
            visualize=self._visualize_optimization,
            prior_orientation_distribution=prior,
        )

        # outputs of SDFEst are OpenGL camera, ShapeNet object convention
        position_cv = pointset_utils.change_position_camera_convention(
            position[0], "opengl", "opencv"
        )
        orientation_cv = pointset_utils.change_orientation_camera_convention(
            orientation[0], "opengl", "opencv"
        )

        # reconstruction + extent
        mesh = pipeline.generate_mesh(shape, scale, True).get_transformed_o3d_geometry()
        reconstructed_points = torch.from_numpy(
            np.asarray(mesh.sample_points_uniformly(self._num_points).points)
        )
        extents, _ = reconstructed_points.abs().max(dim=0)
        extents *= 2.0

        return {
            "position": position_cv.detach().cpu(),
            "orientation": orientation_cv.detach().cpu(),
            "extents": extents,
            "reconstructed_pointcloud": reconstructed_points,
            "reconstructed_mesh": mesh,
        }

Config

Bases: TypedDict

Configuration dictionary for SDFEst.

All keys supported by SDFPipeline are supported and will overwrite config contained in sdfest_... files. The keys specified here are used by this script only.

The two keys sdfest_..._config_files will be parsed with SDFEst install directory as part of the search paths. This allows to use the default config that comes with SDFEst installation.

ATTRIBUTE DESCRIPTION
sdfest_default_config_file

Default configuration file loaded first.
sdfest_category_config_files

Per-category configuration file loaded second.
device

Device used for computation.
num_points

Numbner of points extracted from mesh.
prior

Prior distribution to modify orientation distribution.
visualize_optimization

Whether to show additional optimization visualization.
Source code in cpas_toolbox/cpas_methods/sdfest.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
class Config(TypedDict):
    """Configuration dictionary for SDFEst.

    All keys supported by SDFPipeline are supported and will overwrite config
    contained in sdfest_... files. The keys specified here are used by this
    script only.

    The two keys sdfest_..._config_files  will be parsed with SDFEst install
    directory as part of the search paths. This allows to use the default config
    that comes with SDFEst installation.

    Attributes:
        sdfest_default_config_file: Default configuration file loaded first.
        sdfest_category_config_files: Per-category configuration file loaded second.
        device: Device used for computation.
        num_points: Numbner of points extracted from mesh.
        prior: Prior distribution to modify orientation distribution.
        visualize_optimization:
            Whether to show additional optimization visualization.
    """

__init__ 

__init__(config: Config, camera: camera_utils.Camera) -> None

Initialize and load SDFEst models.

Configuration loaded in following order sdfest_default_config_file -> sdfest_category_config_files -> all other keys I.e., keys specified directly will take precedence over keys specified in default file.

Source code in cpas_toolbox/cpas_methods/sdfest.py 
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
def __init__(self, config: Config, camera: camera_utils.Camera) -> None:
    """Initialize and load SDFEst models.

    Configuration loaded in following order
        sdfest_default_config_file -> sdfest_category_config_files -> all other keys
    I.e., keys specified directly will take precedence over keys specified in
    default file.
    """
    default_config = yoco.load_config_from_file(
        config["sdfest_default_config_file"],
        search_paths=[".", "~/.sdfest/", sdfest.__path__[0]],
    )

    self._pipeline_dict = {}  # maps category to category-specific pipeline
    self._device = config["device"]
    self._visualize_optimization = config["visualize_optimization"]
    self._num_points = config["num_points"]
    self._prior = config["prior"] if "prior" in config else None

    # create per-categry models
    for category_str in config["sdfest_category_config_files"].keys():
        category_config = yoco.load_config_from_file(
            config["sdfest_category_config_files"][category_str],
            current_dict=default_config,
            search_paths=[".", "~/.sdfest/", sdfest.__path__[0]],
        )
        category_config = yoco.load_config(config, category_config)
        self._pipeline_dict[category_str] = SDFPipeline(category_config)
        self._pipeline_dict[category_str].cam = camera

inference 

inference(
    color_image: torch.Tensor,
    depth_image: torch.Tensor,
    instance_mask: torch.Tensor,
    category_str: str,
) -> PredictionDict

See CPASMethod.inference.

Source code in cpas_toolbox/cpas_methods/sdfest.py 
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
def inference(
    self,
    color_image: torch.Tensor,
    depth_image: torch.Tensor,
    instance_mask: torch.Tensor,
    category_str: str,
) -> PredictionDict:
    """See CPASMethod.inference."""
    # skip unsupported category
    if category_str not in self._pipeline_dict:
        return {
            "position": torch.tensor([0, 0, 0]),
            "orientation": torch.tensor([0, 0, 0, 1]),
            "extents": torch.tensor([1, 1, 1]),
            "reconstructed_pointcloud": torch.tensor([[0, 0, 0]]),
            "reconstructed_mesh": None,
        }

    pipeline = self._pipeline_dict[category_str]

    # move inputs to device
    color_image = color_image.to(self._device)
    depth_image = depth_image.to(self._device, copy=True)
    instance_mask = instance_mask.to(self._device)

    if self._prior is not None:
        prior = torch.tensor(self._prior[category_str], device=self._device)
        prior /= torch.sum(prior)
    else:
        prior = None

    position, orientation, scale, shape = pipeline(
        depth_image,
        instance_mask,
        color_image,
        visualize=self._visualize_optimization,
        prior_orientation_distribution=prior,
    )

    # outputs of SDFEst are OpenGL camera, ShapeNet object convention
    position_cv = pointset_utils.change_position_camera_convention(
        position[0], "opengl", "opencv"
    )
    orientation_cv = pointset_utils.change_orientation_camera_convention(
        orientation[0], "opengl", "opencv"
    )

    # reconstruction + extent
    mesh = pipeline.generate_mesh(shape, scale, True).get_transformed_o3d_geometry()
    reconstructed_points = torch.from_numpy(
        np.asarray(mesh.sample_points_uniformly(self._num_points).points)
    )
    extents, _ = reconstructed_points.abs().max(dim=0)
    extents *= 2.0

    return {
        "position": position_cv.detach().cpu(),
        "orientation": orientation_cv.detach().cpu(),
        "extents": extents,
        "reconstructed_pointcloud": reconstructed_points,
        "reconstructed_mesh": mesh,
    }