pydrake_custom_pickling_example.py

"""
We (TRI) used this pattern in Anzu (our internal app codebase) before some of
our classes supported pickling.

Here is how we support pickling an object that contained two classes that
were unpicklable:
- RigidTransform - now pickable, as of v0.10.0:
    https://github.com/RobotLocomotion/drake/commit/3b559846
- Certain `pyassimp` objects (e.g. `pyassimp.structs.Scene` from
`pyassimp.load`).

At time of this writing, this code was using upstream deps:
- drake@69c771f
- pyassimp (v4.1.0, installed from `apt` on Bionic)
"""

import numpy as np
import pyassimp

from pydrake.math import RigidTransform


class ObjectModel(object):
    """This is only a subset of the real implementation."""
    def __init__(self, mesh_file, X_OOm):
        self.mesh_file = mesh_file
        # Pose of object mesh frame (Om) w.r.t. object frame (O).
        assert X_OOm is not None
        self.X_OOm = RigidTransform(X_OOm)
        # To be lazy-loaded using `self._load_mesh_if_necessary()`.
        self._mesh = None

    def __getstate__(self):
        # TODO(eric): Remove this once `RigidTransform` supports pickling
        # (drake#11957).
        d = self.__dict__.copy()
        # WARNING: Mesh path will be invalid. Consider normalizing?
        # Use numpy value from matrix.
        d["X_OOm"] = self.X_OOm.matrix()
        d["_mesh"] = None  # Cannot pickle pyassimp object
        return d

    def __setstate__(self, d):
        self.__dict__.update(d)
        # Reconstruct with numpy value that was stored.
        self.X_OOm = RigidTransform(self.X_OOm)

    def _load_mesh_if_necessary(self):
        if self._mesh is None:
            print(f"loading {self.mesh_file}")
            self._mesh = pyassimp.load(self.mesh_file)
        assert self._mesh is not None