improve the robustness of tile creation

bioimage-io · Oct 25, 2023 · e0e09c4 · e0e09c4
1 parent 5d36645
commit e0e09c4
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 36 deletions.
diff --git a/src/main/java/io/bioimage/modelrunner/bioimageio/description/TensorSpec.java b/src/main/java/io/bioimage/modelrunner/bioimageio/description/TensorSpec.java
@@ -345,41 +345,43 @@ public int[] getPatchArrFromStr(String patchStr) {
 
     /**
      * Validates if a given patch array fulfills the conditions specified in the yaml file.
-     * If it is valid, it sets the value as the {@link #processingPatch}
-     * @param patch
-     * 	the patch array to validate
+     * To check whether the tile size works also for a given image, 
+     * please use {@link #setTileSizeForTensorAndImageSize(int[], int[])}
+     * @param tileSize
+     * 	the patch array size to validate
      * @throws Exception if the patch does not comply with the constraints specified
      */
-    public void validate(int[] patch) throws Exception {
+    public void validate(int[] tileSize) throws Exception {
     	// VAlidate that the minimum size and step constraints are fulfilled
-    	validateStepMin(patch);
-    	this.processingPatch = patch;
+    	validateStepMin(tileSize);
     }
 
     /**
-     * Validates if a given patch array fulfills the conditions specified in the yaml file.
-     * If it is valid, it sets the value as the {@link #processingPatch}
-     * @param patch
-     * 	the patch array to validate
-     * @param seqSize
-     * 	array containing the dimensions of the sequence that is going to be processed
-     * 	seqSize is defined following the Icy axes order (xyztc)
+     * Sets the tile size to process the given tensor regarding that the tensor
+     * has the dimensions specified by the second argument.
+     * If also validates if the tile size selected fufils the requirements
+     * specified in the bioimage.io rdf.yaml file.
+     * Both arguments must follow the same axis order of this tensor.
+     * TODO maybe allow the possibility of changing the tensor axes order
+     * If everything is correct, sets {@link #processingPatch} to the first argument.
+     * @param tileSize
+     * 	the size of the tile/patch in which the main tensor is going to be divided
+     * @param tensorImageSize
+     * 	size of the image that is used for the tensor.
      * @throws Exception if the patch size is not able to 
      * 	fulfill the requirements of the tensor
      */
-    public void validate(int[] patch, int[] seqSize) throws Exception {
-    	// Convert the Icy sequence array dims into the tensor axes order
-    	seqSize = PatchGridCalculator.icySeqAxesOrderToWantedOrder(seqSize, axes);
+    public void setTileSizeForTensorAndImageSize(int[] tileSize, int[] tensorImageSize) throws Exception {
     	// If tiling is not allowed, the patch array needs to be equal to the
     	// optimal patch
     	if (!tiling) {
-    		validateNoTiling(patch, seqSize);
+    		validateNoTiling(tileSize, tensorImageSize);
     	}
     	// VAlidate that the minimum size and step constraints are fulfilled
-    	validateStepMin(patch);
+    	validateStepMin(tileSize);
     	// Finally validate that the sequence size complies with the patch size selected
-    	validatePatchVsImage(patch, seqSize);
-    	this.processingPatch = patch;
+    	validatePatchVsImage(tileSize, tensorImageSize);
+    	this.processingPatch = tileSize;
     }
 
 

diff --git a/src/main/java/io/bioimage/modelrunner/tiling/PatchGridCalculator.java b/src/main/java/io/bioimage/modelrunner/tiling/PatchGridCalculator.java
@@ -101,6 +101,7 @@ public static <T extends RealType<T> & NativeType<T>> PatchGridCalculator<T> bui
     	try {
 	    	descriptor = 
 	    			ModelDescriptor.readFromLocalFile(modelFolder + File.separator + Constants.RDF_FNAME, false);
+	    	descriptor.getInputTensors().get(0).setTileSizeForTensorAndImageSize(new int[]{1, 256, 256, 3}, new int[]{1, 512, 512, 3});
     	} catch (Exception ex) {
     		throw new IOException("Unable to process the rf.yaml specifications file.", ex);
     	}
@@ -267,25 +268,21 @@ private PatchSpec computePatchSpecs(TensorSpec inputTensorSpec, Tensor<T> inputO
     		throws IllegalArgumentException {
     	return computePatchSpecs(inputTensorSpec, inputObject.getData());
     }
-
-    /**
-     * Compute the patch details needed to perform the tiling strategy. The calculations
-     * obtain the input patch, the padding needed at each side and the number of patches
-     * needed for every tensor.
-     * 
-     * @param spec
-     * 	specs of the tensor
-     * @param rai
-     * 	ImgLib2 rai, backend of a tensor, thta is going to be tiled
-     * 
-     * @return an object containing the specs needed to perform patching for the particular tensor
-     */
-    private PatchSpec computePatchSpecs(TensorSpec spec, RandomAccessibleInterval<T> rai)
+
+    private PatchSpec computePatchSpecs(TensorSpec spec, RandomAccessibleInterval<T> rai) throws IllegalArgumentException
     {
-    	String processingAxesOrder = spec.getAxesOrder();
     	int[] intShape = new int[rai.dimensionsAsLongArray().length];
     	for (int i = 0; i < intShape.length; i ++) intShape[i] = (int) rai.dimensionsAsLongArray()[i];
-    	return computePatchSpecs(spec, rai, spec.getOptimalPatch(intShape, processingAxesOrder));
+    	if (spec.getProcessingPatch() == null) {
+			try {
+				spec.setTileSizeForTensorAndImageSize(spec.getOptimalPatch(intShape), intShape);
+			} catch (Exception e) {
+				throw new IllegalArgumentException("Tensor dimensions of tensor named '" + spec.getName() + "' "
+						+ "are not compatible with the requirements set by the"
+						+ " rdf.yaml file for tensor '" + spec.getName() + "': " + e.getMessage());
+			}
+    	}
+    	return computePatchSpecs(spec, rai, spec.getProcessingPatch());
     }
 
     /**