Allow multichannel raw and single-channel mask to pass assertion

model_server/base/models.py

@@ -85,7 +85,7 @@ class InstanceSegmentationModel(ImageToImageModel):
         """
         if not mask.is_mask():
             raise InvalidInputImageError('Expecting a binary mask')
-        if not img.shape == mask.shape:
+        if img.hw != mask.hw or img.nz != mask.nz:
             raise InvalidInputImageError('Expect input image and mask to be the same shape')
 
     def label_patch_stack(self, img: PatchStack, mask: PatchStack, **kwargs):