Fix permute bug in order to have several output masks and images

2023-11-29 16:30:23 +01:00
parent 906a51f9ef
commit 52215ae067
1 changed files with 1 additions and 1 deletions
@@ -231,7 +231,7 @@ def sam_segment(
        point_labels=None,
        boxes=transformed_boxes.to(sam_device),
        multimask_output=False)
-    masks = masks.permute(1, 0, 2, 3).cpu().numpy()
+    masks = masks.cpu().numpy()
    return create_tensor_output(image_np, masks, boxes)