From 50cccf89be88a90f1def42a22f1b7dcf7c955de7 Mon Sep 17 00:00:00 2001 From: jevans Date: Sat, 4 Oct 2014 22:16:09 -0400 Subject: [PATCH] additional ellipsis cases --- glymur/jp2k.py | 30 ++++++++++++++++++++++++++---- glymur/test/test_jp2k.py | 10 ++++++++++ 2 files changed, 36 insertions(+), 4 deletions(-) diff --git a/glymur/jp2k.py b/glymur/jp2k.py index ea8e955..751d378 100644 --- a/glymur/jp2k.py +++ b/glymur/jp2k.py @@ -18,6 +18,7 @@ else: from collections import Counter import ctypes +import itertools import math import os import re @@ -807,19 +808,40 @@ class Jp2k(Jp2kBox): return pixel[pargs[2]] if isinstance(pargs, tuple) and any(x is ... for x in pargs): - nrows = codestream.segment[1].ysiz ncols = codestream.segment[1].xsiz nbands = codestream.segment[1].Csiz # Reformulate without the ellipsis. if pargs[0] is ...: - newindex = (slice(0, nrows), slice(0, ncols), pargs[1]) + if len(pargs) == 2: + newindex = (slice(0, nrows), slice(0, ncols), pargs[1]) + else: + newindex = (slice(0, nrows), pargs[1], pargs[2]) elif pargs[1] is ...: - # Assume we have something like (r,...) where r is a scalar. - newindex = pargs[0] + if len(pargs) == 2: + newindex = (pargs[0], slice(0, ncols), slice(0, nbands)) + else: + newindex = (pargs[0], slice(0, ncols), pargs[2]) + else: + newindex = (pargs[0], pargs[1], slice(0, nbands)) return self.__getitem__(newindex) + if isinstance(pargs, tuple) and not all(isinstance(x, slice) for x in pargs): + # Search out any remaining non-slices and turn them into slices. + lst = list(pargs) + predicate = lambda x: not isinstance(x[1], int) + g = itertools.filterfalse(predicate, enumerate(pargs)) + idx = list(g)[0][0] + lst[idx] = slice(pargs[idx], pargs[idx] + 1) + newindex = tuple(lst) + + data = self.__getitem__(newindex) + + # Reduce dimensionality in the scalar dimension. + return np.squeeze(data, axis=idx) + + # Assuming pargs is a tuple of slices from now on. rows = pargs[0] cols = pargs[1] diff --git a/glymur/test/test_jp2k.py b/glymur/test/test_jp2k.py index ef86744..77d796d 100644 --- a/glymur/test/test_jp2k.py +++ b/glymur/test/test_jp2k.py @@ -261,6 +261,16 @@ class TestSliceProtocolRead(SliceProtocolBase): expected = self.j2k_data[0, ...] np.testing.assert_array_equal(actual, expected) + def test_two_ellipsis_band_select(self): + actual = self.j2k[..., ..., 1] + expected = self.j2k_data[:, :, 1] + np.testing.assert_array_equal(actual, expected) + + def test_two_ellipsis_row_select(self): + actual = self.j2k[1, ..., ...] + expected = self.j2k_data[1, :, :] + np.testing.assert_array_equal(actual, expected) + def test_slice_protocol_2d_reduce_resolution(self): d = self.j2k[:] self.assertEqual(d.shape, (800, 480, 3))