Merge pull request #208 from Visual-Behavior/scene_flow

Scene flow in frame

Author: thibo73800

aloscene/__init__.py

@@ -12,6 +12,7 @@
 from .bounding_boxes_2d import BoundingBoxes2D
 from .bounding_boxes_3d import BoundingBoxes3D
 from .oriented_boxes_2d import OrientedBoxes2D
+from .scene_flow import SceneFlow
 from .frame import Frame
 from .tensors.spatial_augmented_tensor import SpatialAugmentedTensor
 

aloscene/frame.py

@@ -6,7 +6,19 @@
 
 import aloscene
 from aloscene.renderer import View
-from aloscene import BoundingBoxes2D, BoundingBoxes3D, Depth, Disparity, Flow, Mask, Labels, Points2D, Points3D, Pose
+from aloscene import (
+    BoundingBoxes2D,
+    BoundingBoxes3D,
+    Depth,
+    Disparity,
+    Flow,
+    Mask,
+    Labels,
+    Points2D,
+    Points3D,
+    Pose,
+    SceneFlow,
+)
 
 # from aloscene.camera_calib import CameraExtrinsic, CameraIntrinsic
 from aloscene.io.image import load_image
@@ -114,6 +126,7 @@ def __new__(
         tensor.add_child("segmentation", segmentation, align_dim=["B", "T"], mergeable=False)
         tensor.add_child("labels", labels, align_dim=["B", "T"], mergeable=True)
         tensor.add_child("pose", labels, align_dim=["B", "T"], mergeable=True)
+        tensor.add_child("scene_flow", labels, align_dim=["B", "T"], mergeable=False)
 
         # Add other tensor property
         tensor.add_property("normalization", normalization)
@@ -318,6 +331,25 @@ def append_pose(self, pose: Pose, name: str = None):
         """
         self._append_child("pose", pose, name)
 
+    def append_scene_flow(self, scene_flow: SceneFlow, name: str = None):
+        """Attach a scene flow to the frame.
+
+        Parameters
+        ----------
+        scene_flow : :mod:`SceneFlow <aloscene.SceneFlow>`
+            Depth to attach to the Frame
+        name : str
+            If none, the scene flow will be attached without name (if possible). Otherwise if no other unnamed
+            scene flow is attached to the frame, the scene flow will be added to the set of scene flows.
+
+        Examples
+        --------
+        >>> frame = aloscene.Frame("/path/to/image.jpeg")
+        >>> scene_flow = aloscene.SceneFlow(np.random.rand((3, frame.H, frame.W)))
+        >>> frame.append_scene_flow(scene_flow)
+        """
+        self._append_child("scene_flow", scene_flow, name)
+
     @staticmethod
     def _get_mean_std_tensor(shape, names, mean_std: tuple, device="cpu"):
         """Utils method to a get the mean and the std

aloscene/scene_flow.py

@@ -123,7 +123,7 @@ def from_optical_flow(
         if not has_batch:
             scene_flow_vector = scene_flow_vector.squeeze(0)
             optical_flow = optical_flow.squeeze(0)
-            occlusion = occlusion.squeeze(0)
+            occlusion = occlusion.squeeze(0) if occlusion is not None else None
 
         # Create the scene flow object
         tensor = cls(