From 5ef69ef3e6180709bc292370ed314b6029ecabfc Mon Sep 17 00:00:00 2001 From: Paul Guerrie <97041392+paulguerrie@users.noreply.github.com> Date: Thu, 6 Oct 2022 14:55:15 -0600 Subject: [PATCH] Error in utils/segment/general `masks2segments()` (#9724) When running segmentation predict on gpu, the conversion from tensor to numpy fails. Calling `.cpu()` solves this problem. Signed-off-by: Paul Guerrie <97041392+paulguerrie@users.noreply.github.com> Signed-off-by: Paul Guerrie <97041392+paulguerrie@users.noreply.github.com> --- utils/segment/general.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/utils/segment/general.py b/utils/segment/general.py index 655123bdcfeb..43bdc460f928 100644 --- a/utils/segment/general.py +++ b/utils/segment/general.py @@ -124,7 +124,7 @@ def masks_iou(mask1, mask2, eps=1e-7): def masks2segments(masks, strategy='largest'): # Convert masks(n,160,160) into segments(n,xy) segments = [] - for x in masks.int().numpy().astype('uint8'): + for x in masks.int().cpu().numpy().astype('uint8'): c = cv2.findContours(x, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0] if strategy == 'concat': # concatenate all segments c = np.concatenate([x.reshape(-1, 2) for x in c])