Adjust schema for virtual chunked.

PiperOrigin-RevId: 589147704

Author: jan-matthis

decorators/flow.py

@@ -19,6 +19,7 @@
 from connectomics.common import jax_utils
 # pylint: disable=g-importing-member
 from connectomics.volume.decorators import adjust_schema_for_chunksize
+from connectomics.volume.decorators import adjust_schema_for_virtual_chunked
 from connectomics.volume.decorators import Decorator
 from connectomics.volume.decorators import Filter
 # pylint: enable=g-importing-member
@@ -73,6 +74,7 @@ def filt_read(domain: ts.IndexDomain, array: np.ndarray,
     schema = input_ts.schema
     if self._min_chunksize is not None:
       schema = adjust_schema_for_chunksize(schema, self._min_chunksize)
+    schema = adjust_schema_for_virtual_chunked(schema)
 
     # Remove non-spatial dimensions from `fc`-dimension (0th).
     json = schema.to_json()

decorators/maps.py

@@ -17,7 +17,10 @@
 from typing import Any, Mapping, MutableMapping, Optional, Sequence
 
 from connectomics.common import bounding_box
-from connectomics.volume.decorators import Decorator  # pylint: disable=g-importing-member
+# pylint: disable=g-importing-member
+from connectomics.volume.decorators import adjust_schema_for_virtual_chunked
+from connectomics.volume.decorators import Decorator
+# pylint: enable=g-importing-member
 import gin
 import numpy as np
 import sofima
@@ -84,7 +87,8 @@ def warp_fn(domain: ts.IndexDomain, array: np.ndarray,
         chunksize.append(dim.size)
       else:
         chunksize.append(1)
-    json = input_ts.schema.to_json()
+    schema = adjust_schema_for_virtual_chunked(input_ts.schema)
+    json = schema.to_json()
     json['chunk_layout']['read_chunk']['shape'] = chunksize
     json['chunk_layout']['write_chunk']['shape'] = chunksize
 

decorators/warp.py

@@ -18,7 +18,10 @@
 
 from connectomics.common import bounding_box
 from connectomics.common import opencv_utils
-from connectomics.volume.decorators import Decorator  # pylint: disable=g-importing-member
+# pylint: disable=g-importing-member
+from connectomics.volume.decorators import adjust_schema_for_virtual_chunked
+from connectomics.volume.decorators import Decorator
+# pylint: enable=g-importing-member
 import gin
 import numpy as np
 import scipy.ndimage
@@ -205,7 +208,8 @@ def warp_fn(domain: ts.IndexDomain, array: np.ndarray,
         chunksize.append(dim.size)
       else:
         chunksize.append(1)
-    json = input_ts.schema.to_json()
+    schema = adjust_schema_for_virtual_chunked(input_ts.schema)
+    json = schema.to_json()
     json['chunk_layout']['read_chunk']['shape'] = chunksize
     json['chunk_layout']['write_chunk']['shape'] = chunksize
 
@@ -302,7 +306,8 @@ def warp_fn(domain: ts.IndexDomain, array: np.ndarray,
         chunksize.append(dim.size)
       else:
         chunksize.append(1)
-    json = input_ts.schema.to_json()
+    schema = adjust_schema_for_virtual_chunked(input_ts.schema)
+    json = schema.to_json()
     json['chunk_layout']['read_chunk']['shape'] = chunksize
     json['chunk_layout']['write_chunk']['shape'] = chunksize