Ankle/yolov5_ZWgra
1
0 Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

Error in utils/segment/general masks2segments() (#9724)

Browse Source 
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>
This commit is contained in:
Paul Guerrie authored and GitHub committed 2022-10-06 22:55:15 +02:00
1 parent 7f097ddb6c
commit 5ef69ef3e6
1 file changed
+1 -1
Download Patch File Download Diff File Expand all files Collapse all files
utils/segment/general.py
+1 -1
View File
@@ -124,7 +124,7 @@ def masks_iou(mask1, mask2, eps=1e-7):
def masks2segments(masks, strategy='largest'): def masks2segments(masks, strategy='largest'):
# Convert masks(n,160,160) into segments(n,xy) # Convert masks(n,160,160) into segments(n,xy)
segments = [] 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] c = cv2.findContours(x, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0]
if strategy == 'concat': # concatenate all segments if strategy == 'concat': # concatenate all segments
c = np.concatenate([x.reshape(-1, 2) for x in c]) c = np.concatenate([x.reshape(-1, 2) for x in c])