Add a padding logic to build_mask for the -1 size offset.

PiperOrigin-RevId: 715385760

connectomics/volume/mask.py

@@ -222,6 +222,17 @@ def build_mask(
                 size=fov.size[::-1],
                 origin=fov.size[::-1] // 2 + fov.start[::-1],
             )
+        # When using an upsampled mask, the mask could be not the exactly same
+        # size as the input image (i.e., -1 offset), we need to pad the mask to
+        # match the image size.
+        if curr_mask.shape != bool_mask.shape:
+          padding = [
+              (0, curr_mask.shape[dim] - bool_mask.shape[dim])
+              for dim in range(len(curr_mask.shape))
+          ]
+          bool_mask = np.pad(
+              bool_mask, padding, mode='constant', constant_values=0
+          )
         curr_mask |= bool_mask
 
     if config.invert:

connectomics/volume/mask_test.py

@@ -65,5 +65,48 @@ def test_build_mask(self):
     )
 
 
+def test_build_mask_with_padding(self):
+  subvol_size = (150, 50, 75)
+  smaller_size = (149, 49, 75)
+
+  smaller_mask = np.random.choice([True, False], size=smaller_size)
+  image = np.random.choice(
+      [0, 1], size=subvol_size
+  )
+
+  mask_config = m.MaskConfig.from_dict({
+      'image': {
+          'channels': [{
+              'channel': 0,
+              'min_value': 0,
+              'max_value': 1,
+          }],
+      }
+  })
+
+  box = bounding_box.BoundingBox(start=(0, 0, 0), size=subvol_size[::-1])
+
+  def mock_volume_loader(volume):
+    del volume  # Unused
+    return smaller_mask
+
+  mask = m.build_mask(
+      [mask_config],
+      box,
+      decorated_volume_loader=mock_volume_loader,
+      image=image,
+  )
+
+  # Ensure the resulting mask matches the target dimensions
+  self.assertEqual(mask.shape, subvol_size)
+
+  # Verify that the padded regions are zeros in the final mask
+  padded_image = np.zeros(subvol_size, dtype=bool)
+  padded_image[: smaller_size[0], : smaller_size[1], : smaller_size[2]] = (
+      smaller_mask
+  )
+  np.testing.assert_array_equal(mask, padded_image)
+
+
 if __name__ == '__main__':
   absltest.main()