Merge branch 'bugfix'

Avoid custom handling of PydanticCoordinate/Roi tuples. These are now directly cast to the appropriate types.
Handle writing to datasets that have lazy ops associated with them. We now write directly to the source array.

Author: pattonw

src/volara_torch/blockwise/predict.py

@@ -65,7 +65,6 @@ def process(self, batch, request):
 
 class Predict(BlockwiseTask):
     task_type: Literal["predict"] = "predict"
-    roi: tuple[PydanticCoordinate, PydanticCoordinate] | None = None
     checkpoint: Annotated[
         TorchModel,
         Field(discriminator="model_type"),
@@ -75,6 +74,7 @@ class Predict(BlockwiseTask):
 
     fit: Literal["overhang"] = "overhang"
     read_write_conflict: Literal[False] = False
+    context_override: tuple[PydanticCoordinate, PydanticCoordinate] | None = None
     _out_array_dtype: np.dtype = np.dtype(np.uint8)
 
     @property
@@ -85,7 +85,7 @@ def checkpoint_config(self) -> Model:
     def write_roi(self) -> Roi:
         in_data_roi = self.in_data.array("r").roi
         if self.roi is not None:
-            return in_data_roi.intersect(Roi(self.roi[0], self.roi[1]))
+            return in_data_roi.intersect(self.roi)
         else:
             return in_data_roi
 
@@ -98,8 +98,16 @@ def write_size(self) -> Coordinate:
         return self.checkpoint_config.eval_output_shape * self.voxel_size
 
     @property
-    def context_size(self) -> Coordinate:
-        return self.checkpoint_config.context * self.voxel_size
+    def context_size(self) -> Coordinate | tuple[Coordinate, Coordinate]:
+        context = self.checkpoint_config.context
+        if isinstance(context, Coordinate):
+            return self.checkpoint_config.context * self.voxel_size
+        elif isinstance(context[0], Coordinate) and isinstance(context[1], Coordinate):
+            return (context[0] * self.voxel_size, context[1] * self.voxel_size)
+        else:
+            raise NotImplementedError(
+                f"Unsupported context {context} type: {type(context)}. Expected Coordinate or tuple of Coordinates."
+            )
 
     @property
     def task_name(self) -> str:
@@ -203,7 +211,7 @@ def process_block_func(self):
         for output_key, out_data in zip(output_keys, self.out_data):
             if out_data is not None:
                 pipeline += ArrayWrite(
-                    output_key, out_data.array("a"), self.checkpoint.to_uint8
+                    output_key, Dataset.array(out_data, "a"), self.checkpoint.to_uint8
                 )
 
         print("Starting prediction...")

src/volara_torch/models.py

@@ -44,14 +44,32 @@ class Model(StrictBaseModel, ABC):
     The range of the output data. This is used to convert between `np.uint8` and
     `np.float32` data types for efficient reading/writing of model outputs.
     """
+    context_override: tuple[PydanticCoordinate, PydanticCoordinate] | None = None
+    """
+    An optional override for asymetrical context sizes. This lets you define
+    a specific lower and upper context size for the model which must equal
+    the expected context size of input_shape - output_shape
+    """
 
     @property
-    def context(self) -> Coordinate:
+    def context(self) -> Coordinate | tuple[Coordinate, Coordinate]:
         """
         The context required to make tile artifact free predictions
         with this model.
         """
-        return (self.eval_input_shape - self.eval_output_shape) // 2
+        expected_context = self.min_input_shape - self.min_output_shape
+        if self.context_override is not None:
+            assert (
+                self.context_override[0] + self.context_override[1]
+            ) == expected_context, (
+                f"Expected context override {self.context_override} to sum to {expected_context}, "
+                f"but got {self.context_override[0] + self.context_override[1]}"
+            )
+
+        if self.context_override is not None:
+            return self.context_override
+        else:
+            return expected_context // 2
 
     @abstractmethod
     def model(self) -> torch.nn.Module: