Allow index-based access to OnImage instances

changelogs/master/added/20200102_cbasoi_getitem.md

@@ -0,0 +1,11 @@
+# Index-based Access to Coordinate-based `*OnImage` Instances #547
+
+Enabled index-based access to coordinate-based `*OnImage` instances, i.e. to
+`KeypointsOnImage`, `BoundingBoxesOnImage`, `LineStringsOnImage` and
+`PolygonsOnImage`. This allows to do things like
+`bbsoi = BoundingBoxesOnImage(...); bbs = bbsoi[0:2];`.
+
+* Added `imgaug.augmentables.kps.KeypointsOnImage.__getitem__()`.
+* Added `imgaug.augmentables.bbs.BoundingBoxesOnImage.__getitem__()`.
+* Added `imgaug.augmentables.lines.LineStringsOnImage.__getitem__()`.
+* Added `imgaug.augmentables.polys.PolygonsOnImage.__getitem__()`.

imgaug/augmentables/bbs.py

@@ -1890,6 +1890,17 @@ def deepcopy(self, bounding_boxes=None, shape=None):
 
         return BoundingBoxesOnImage(bounding_boxes, shape)
 
+    def __getitem__(self, indices):
+        """Get the bounding box(es) with given indices.
+
+        Returns
+        -------
+        list of imgaug.augmentables.bbs.BoundingBoxes
+            Bounding box(es) with given indices.
+
+        """
+        return self.bounding_boxes[indices]
+
     def __iter__(self):
         """Iterate over the bounding boxes in this container.
 

imgaug/augmentables/kps.py

@@ -1398,6 +1398,17 @@ def deepcopy(self, keypoints=None, shape=None):
 
         return KeypointsOnImage(keypoints, shape)
 
+    def __getitem__(self, indices):
+        """Get the keypoint(s) with given indices.
+
+        Returns
+        -------
+        list of imgaug.augmentables.kps.Keypoint
+            Keypoint(s) with given indices.
+
+        """
+        return self.keypoints[indices]
+
     def __iter__(self):
         """Iterate over the keypoints in this container.
 

imgaug/augmentables/lines.py

@@ -2286,6 +2286,17 @@ def deepcopy(self, line_strings=None, shape=None):
 
         return LineStringsOnImage(line_strings, shape)
 
+    def __getitem__(self, indices):
+        """Get the line string(s) with given indices.
+
+        Returns
+        -------
+        list of imgaug.augmentables.lines.LineString
+            Line string(s) with given indices.
+
+        """
+        return self.line_strings[indices]
+
     def __iter__(self):
         """Iterate over the line strings in this container.
 

imgaug/augmentables/polys.py

@@ -2046,6 +2046,17 @@ def deepcopy(self, polygons=None, shape=None):
 
         return PolygonsOnImage(polygons, shape)
 
+    def __getitem__(self, indices):
+        """Get the polygon(s) with given indices.
+
+        Returns
+        -------
+        list of imgaug.augmentables.polys.Polygon
+            Polygon(s) with given indices.
+
+        """
+        return self.polygons[indices]
+
     def __iter__(self):
         """Iterate over the polygons in this container.
 

test/augmentables/test_bbs.py

@@ -2009,6 +2009,17 @@ def test_deepcopy_shape_set(self):
         assert bbsoi_copy.bounding_boxes[0].coords_almost_equals(bb1)
         assert bbsoi_copy.bounding_boxes[1].coords_almost_equals(bb2)
 
+    def test___getitem__(self):
+        cbas = [
+            ia.BoundingBox(x1=1, y1=2, x2=3, y2=4),
+            ia.BoundingBox(x1=2, y1=3, x2=4, y2=5)
+        ]
+        cbasoi = ia.BoundingBoxesOnImage(cbas, shape=(3, 4, 3))
+
+        assert cbasoi[0] is cbas[0]
+        assert cbasoi[1] is cbas[1]
+        assert cbasoi[0:2] == cbas
+
     def test___iter__(self):
         cbas = [ia.BoundingBox(x1=0, y1=0, x2=2, y2=2),
                 ia.BoundingBox(x1=1, y1=2, x2=3, y2=4)]

test/augmentables/test_kps.py

@@ -1330,6 +1330,17 @@ def test_deepcopy_shape_set(self):
         assert kpsoi_copy.keypoints[0].coords_almost_equals(kp1)
         assert kpsoi_copy.keypoints[1].coords_almost_equals(kp2)
 
+    def test___getitem__(self):
+        cbas = [
+            ia.Keypoint(x=1, y=2),
+            ia.Keypoint(x=2, y=3)
+        ]
+        cbasoi = ia.KeypointsOnImage(cbas, shape=(3, 4, 3))
+
+        assert cbasoi[0] is cbas[0]
+        assert cbasoi[1] is cbas[1]
+        assert cbasoi[0:2] == cbas
+
     def test___iter__(self):
         cbas = [ia.Keypoint(x=1, y=2),
                 ia.Keypoint(x=3, y=4)]

test/augmentables/test_lines.py

@@ -2627,6 +2627,17 @@ def test_deepcopy_and_replace_line_strings_attribute_with_empty_list(self):
         assert observed.line_strings == []
         assert observed.shape == (200, 201, 3)
 
+    def test___getitem__(self):
+        cbas = [
+            ia.LineString([(0, 0), (1, 0), (1, 1)]),
+            ia.LineString([(1, 1), (2, 1), (2, 2)])
+        ]
+        cbasoi = ia.LineStringsOnImage(cbas, shape=(3, 4, 3))
+
+        assert cbasoi[0] is cbas[0]
+        assert cbasoi[1] is cbas[1]
+        assert cbasoi[0:2] == cbas
+
     def test___iter__(self):
         cbas = [ia.LineString([(0, 0), (1, 1)]),
                 ia.LineString([(1, 2), (3, 4)])]

test/augmentables/test_polys.py

@@ -3128,6 +3128,19 @@ def test_shape_parameter_set(self):
         assert psoi_copy.polygons[1].coords_almost_equals(poly2)
 
 
+class TestPolygonsOnImage___getitem__(unittest.TestCase):
+    def test_with_two_polygons(self):
+        cbas = [
+            ia.Polygon([(0, 0), (1, 0), (1, 1)]),
+            ia.Polygon([(1, 1), (2, 1), (2, 2)])
+        ]
+        cbasoi = ia.PolygonsOnImage(cbas, shape=(3, 4, 3))
+
+        assert cbasoi[0] is cbas[0]
+        assert cbasoi[1] is cbas[1]
+        assert cbasoi[0:2] == cbas
+
+
 class TestPolygonsOnImage___iter__(unittest.TestCase):
     def test_with_two_polygons(self):
         cbas = [ia.Polygon([(0, 0), (1, 0), (1, 1)]),