diff --git a/glymur/jp2k.py b/glymur/jp2k.py index c9021be..6919555 100644 --- a/glymur/jp2k.py +++ b/glymur/jp2k.py @@ -759,60 +759,51 @@ class Jp2k(Jp2kBox): return boxes - def __getitem__(self, *pargs): + def __getitem__(self, pargs): """ Slicing protocol. """ - if isinstance(pargs[0], slice): + if isinstance(pargs, slice): + # Case of jp2[:] + # # Should have a slice object where start = stop = step = None - slc = pargs[0] + slc = pargs if slc.start is None and slc.stop is None and slc.step is None: return self.read() else: raise IndexError("Illegal syntax.") - if not isinstance(pargs[0], tuple): - msg = "Unexpected situation, slicing invoked, but not passed " - msg += "a slice or tuple." - raise RuntimeError(msg) + # Assuming pargs is a tuple from now on. + rows = pargs[0] + cols = pargs[1] + if len(pargs) == 2: + bands = slice(None, None, None) + else: + bands = pargs[2] - # Assuming tuple from now on. - ridx = pargs[0][0] - cidx = pargs[0][1] - bidx = pargs[0][2] - - if ((ridx.step is None) and (cidx.step is None)): + if ((rows.step is None) and (cols.step is None)): # Slicing with full resolution. - return self.read()[ridx, cidx, bidx] + # This can be improved to take advantage of tiling. + return self.read()[rows, cols, bands] - if ((ridx.start is not None) or - (ridx.stop is not None) or - (cidx.start is not None) or - (cidx.stop is not None)): - msg = "Only strides are supported when slicing a Jp2k object." - raise IndexError(msg) - - if ridx.step is None and cidx.step is None: - step = 1 - elif ridx.step != cidx.step: + if rows.step != cols.step: msg = "Row and column strides must be the same." raise IndexError(msg) - else: - step = ridx.step + + # Ok, reduce layer step is the same in both xy directions, so just take + # one of them. + step = rows.step if np.log2(step) != np.floor(np.log2(step)): msg = "Row and column strides must be powers of 2." raise IndexError(msg) data = self.read(rlevel=np.int(np.log2(step))) - if len(pargs[0]) == 2: + if len(pargs) == 2: return data # Ok, 3 arguments in pargs. - if isinstance(pargs[0][2], slice): - return data[:,:,pargs[0][2]] - elif isinstance(pargs[0][2], int): - return data[:,:,pargs[0][2]] + return data[:, :, bands] def read(self, **kwargs): diff --git a/glymur/test/test_jp2k.py b/glymur/test/test_jp2k.py index bbbe657..50b153a 100644 --- a/glymur/test/test_jp2k.py +++ b/glymur/test/test_jp2k.py @@ -74,24 +74,11 @@ class TestSliceProtocol(unittest.TestCase): with self.assertRaises(IndexError): self.j2k[::3, ::3] - def test_start_and_resolution_stride_not_allowed_at_same_time(self): - with self.assertRaises(IndexError): - self.j2k[2::2, 2::2] - - def test_stop_and_resolution_stride_not_allowed_at_same_time(self): - with self.assertRaises(IndexError): - self.j2k[:8:2, :8:2] - def test_integer_index_in_3d(self): - d = self.j2k[:,:,0] - np.testing.assert_array_equal(self.j2k_data[:,:,0], d) - - d = self.j2k[:,:,1] - np.testing.assert_array_equal(self.j2k_data[:,:,1], d) - - d = self.j2k[:,:,2] - np.testing.assert_array_equal(self.j2k_data[:,:,2], d) + for j in [0, 1, 2]: + band = self.j2k[:, :, j] + np.testing.assert_array_equal(self.j2k_data[:, :, j], band) def test_slice_in_third_dimension(self): actual = self.j2k[:,:,1:3] @@ -103,13 +90,15 @@ class TestSliceProtocol(unittest.TestCase): all = self.j2k.read(rlevel=1) np.testing.assert_array_equal(all[:,:,1:3], d) - def test_full_resolution_slicing_by_quarters(self): - # upper left - np.testing.assert_array_equal(self.jp2_data[:728, :1296], - self.jp2[:728, :1296]) - # lower left - np.testing.assert_array_equal(self.jp2_data[728:, :1296], - self.jp2[728:, :1296]) + def test_full_resolution_slicing_by_quarters_upper_left(self): + actual = self.jp2[:728, :1296] + expected = self.jp2_data[:728, :1296] + np.testing.assert_array_equal(actual, expected) + + def test_full_resolution_slicing_by_quarters_lower_left(self): + actual = self.jp2[728:, :1296] + expected = self.jp2_data[728:, :1296] + np.testing.assert_array_equal(actual, expected) def test_full_resolution_slicing_by_quarters_upper_right(self): actual = self.jp2[:728, 1296:]