Frame with labels

aloscene/frame.py

@@ -7,7 +7,7 @@
 import aloscene
 from aloscene.renderer import View
 from aloscene.disparity import Disparity
-from aloscene import BoundingBoxes2D, BoundingBoxes3D, Flow, Mask
+from aloscene import BoundingBoxes2D, BoundingBoxes3D, Flow, Mask, Labels
 
 # from aloscene.camera_calib import CameraExtrinsic, CameraIntrinsic
 from aloscene.io.image import load_image
@@ -24,6 +24,7 @@ def __new__(
         x,
         boxes2d: Union[dict, BoundingBoxes2D] = None,
         boxes3d: Union[dict, BoundingBoxes3D] = None,
+        labels: Union[dict, Labels] = None,
         flow: Flow = None,
         segmentation: Mask = None,
         disparity: Disparity = None,
@@ -46,6 +47,7 @@ def __new__(
         tensor.add_label("flow", flow, align_dim=["B", "T"], mergeable=False)
         tensor.add_label("disparity", disparity, align_dim=["B", "T"], mergeable=True)
         tensor.add_label("segmentation", segmentation, align_dim=["B", "T"], mergeable=False)
+        tensor.add_label("labels", labels, align_dim=["B", "T"], mergeable=True)
 
         # Add other tensor property
         tensor.add_property("normalization", normalization)
@@ -71,6 +73,20 @@ def save(self, tgt_path: str):
         """
         torchvision.utils.save_image(self.cpu().norm01().as_tensor(), tgt_path)
 
+    def append_labels(self, labels: Labels, name: str = None):
+        """Attach a set of labels to the frame. This can be usefull for classification
+        or multi label classification. The rank of the label must be >= 1
+
+        Parameters
+        ----------
+        labels: aloscene.Labels
+            Set of labels to attached to the frame
+        name: str
+            If none, the label will be attached without name (if possible). Otherwise if no other unnamed
+            labels are attached to the frame, the labels will be added to the set of labels.
+        """
+        self._append_label("labels", labels, name)
+
     def append_boxes2d(self, boxes: BoundingBoxes2D, name: str = None):
         """Attach a set of boxes to the frame.
 

unittest/test_frame.py

@@ -1,3 +1,4 @@
+from aloscene.labels import Labels
 from aloscene.bounding_boxes_2d import BoundingBoxes2D
 
 # from aloscene.renderer import View
@@ -523,7 +524,33 @@ def test_device_propagation():
     pass
 
 
+def test_frame_label():
+    frame = aloscene.Frame(np.random.uniform(0, 1, (3, 600, 600)), names=("C", "H", "W"))
+
+    label = aloscene.Labels([1], encoding="id")
+    frame.append_labels(label)
+
+    assert frame.labels == 1
+    frame = frame.batch()
+    assert len(frame.labels) == 1
+    frame = frame.temporal()
+    assert len(frame.labels) == 1 and len(frame.labels[0]) == 1
+    frame1 = frame.clone()
+    frame2 = frame.clone()
+    frames_batch = torch.cat([frame1, frame2], dim=1)
+    assert len(frames_batch.labels) == 1 and len(frames_batch.labels[0]) == 2
+    frames_temporal = torch.cat([frame1, frame2], dim=0)
+    assert len(frames_temporal.labels) == 2 and len(frames_temporal.labels[0]) == 1
+    n_frame = frames_batch[:, 0]
+    assert len(n_frame.labels) == 1 and len(n_frame.labels[0]) == 1
+    n_frame = frames_temporal[0]
+    assert len(n_frame.labels) == 1 and len(n_frame.labels[0]) == 1
+    frame = n_frame[0]
+    assert len(frame.labels) == 1 and len(frame.labels.shape) == 1
+
+
 if __name__ == "__main__":
+    test_frame_label()
     test_frame_from_dt()
     test_frame_01()
     test_frame_255()