From 7d42ed7b3b0d17ac6e90fd05f0ecd5ec918687bb Mon Sep 17 00:00:00 2001 From: Howcanoe WANG Date: Thu, 24 Oct 2024 10:36:54 +0900 Subject: [PATCH] cvtools imarray_crop return mask --- easyidp/cvtools.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/easyidp/cvtools.py b/easyidp/cvtools.py index ea9cd98..ad1b781 100644 --- a/easyidp/cvtools.py +++ b/easyidp/cvtools.py @@ -7,7 +7,7 @@ warnings.filterwarnings("ignore", message="The array interface is deprecated and will no longer work in Shapely 2.0") -def imarray_crop(imarray, polygon_hv, nodata_value=0, transparent_layer=None): +def imarray_crop(imarray, polygon_hv, nodata_value=0, transparent_layer=None, return_mask=False): """crop a given ndarray image by given polygon pixel positions Parameters @@ -60,12 +60,17 @@ def imarray_crop(imarray, polygon_hv, nodata_value=0, transparent_layer=None): | 0-x : specific alpha layer | -1 : the last layer + return_mask: bool + | return the cooresponding binary mask for further usage, default False + returns ------- imarray_out : ndarray the (m,n,d) ndrray to store pixel info roi_top_left_offset : ndarray the (h, v) pixel index that represent the polygon bbox left top corner + mask: ndarray + the (m, n) binary img ndarray of polgyon masks """ # check if the imarray is correct imarray @@ -214,7 +219,10 @@ def imarray_crop(imarray, polygon_hv, nodata_value=0, transparent_layer=None): # imarray_out[~mask, i] = nodata_value imarray_out[~mask] = nodata_value - return imarray_out, roi_top_left_offset + if return_mask: + return imarray_out, roi_top_left_offset, mask + else: + return imarray_out, roi_top_left_offset def poly2mask(image_shape, poly_coord, engine="skimage"):