From 4fac5e74218c2496f49f276ce38ffeceaace525e Mon Sep 17 00:00:00 2001 From: Lukasz Migas Date: Mon, 9 Dec 2024 18:05:24 +0100 Subject: [PATCH] Update image.py --- src/koyo/image.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/src/koyo/image.py b/src/koyo/image.py index adf7151..634793d 100644 --- a/src/koyo/image.py +++ b/src/koyo/image.py @@ -93,10 +93,19 @@ def reshape_array_from_coordinates( def flatten_array_from_coordinates(array: np.ndarray, coordinates: np.ndarray, offset: int = 0) -> np.ndarray: """Flatten array based on xy coordinates.""" - try: - return array[coordinates[:, 1] - offset, coordinates[:, 0] - offset] - except IndexError: - return array[coordinates[:, 0] - offset, coordinates[:, 1] - offset] + if array.ndim == 2: + try: + return array[coordinates[:, 1] - offset, coordinates[:, 0] - offset] + except IndexError: + return array[coordinates[:, 0] - offset, coordinates[:, 1] - offset] + else: + try: + res = array[:, coordinates[:, 1] - offset, coordinates[:, 0] - offset] + except IndexError: + res = array[:, coordinates[:, 0] - offset, coordinates[:, 1] - offset] + # need to swap axes + res = np.swapaxes(res, 0, 1) + return res def reshape_array_batch(