Allow index-based acces to OnImage instances

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];`.

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

@@ -1633,6 +1633,17 @@ def deepcopy(self):
         bbs = [bb.deepcopy() for bb in self.bounding_boxes]
         return BoundingBoxesOnImage(bbs, tuple(self.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

@@ -1250,6 +1250,17 @@ def deepcopy(self, keypoints=None, shape=None):
             shape = tuple(self.shape)
         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

@@ -2081,6 +2081,17 @@ def deepcopy(self, line_strings=None, shape=None):
             line_strings=[ls.deepcopy() for ls in lss],
             shape=tuple(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

@@ -1784,6 +1784,17 @@ def deepcopy(self):
         polys = [poly.deepcopy() for poly in self.polygons]
         return PolygonsOnImage(polys, tuple(self.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

@@ -1726,6 +1726,17 @@ def test_deepcopy(self):
         assert bbsoi.bounding_boxes[0].y1 == 10
         assert bbsoi_copy.bounding_boxes[0].y1 == 0
 
+    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

@@ -1141,6 +1141,17 @@ def test_deepcopy(self):
         assert kpi2.keypoints[1].x == 3
         assert kpi2.keypoints[1].y == 4
 
+    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

@@ -2426,6 +2426,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

@@ -2822,6 +2822,19 @@ def test_with_two_polygons(self):
                            rtol=0, atol=1e-4)
 
 
+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)]),