Prediction with one channel not working on format version v0.4

[IvanHCenalmor]

Hi @FynnBe!

I am working in the notebook to load, fine-tune and export models from the BMZ.

I manage it to work for the affable-shark model :hugs: (which has the format version v0.5), but now I was trying with wild-whale (which has the format version v0.4) and I am getting a different performance on the prediction function. I found what was causing this difference, but I don't know if my proposal is a good solution.

To put you on context, the image I'm passing for prediction has the shape (1024,1024), without any batch or channel axis on the original .tif file. The input axes for affable-shark are:

model.inputs[0].axes = [BatchAxis(id='batch', description='', type='batch', size=None),
 ChannelAxis(id='channel', description='', type='channel', channel_names=['channel0']),
 SpaceInputAxis(size=ParameterizedSize(min=64, step=16), id='y', description='', type='space', unit=None, scale=1.0, concatenable=False),
 SpaceInputAxis(size=ParameterizedSize(min=64, step=16), id='x', description='', type='space', unit=None, scale=1.0, concatenable=False)]

While the axes on wild-shale are:

model.inputs[0].axes='bcyx'

The code that I'm using for the prediction is:

#@markdown ### If you have an image in a folder to segment, copy the path to it here:
path2image = "/DeepBacs/Multilabel_U-Net_dataset_B.subtilis/test/source/test_9.tif"  #@param {type:"string"}

# Load the paths to the input images
input_paths = {"input0": Path(path2image)}

# The prediction pipeline expects a Sample object from bioimageio.core
input_sample = create_sample_for_model(
    model=model, inputs=input_paths, sample_id="my_demo_sample"
)

# Use the predict function with the defined Samples
prediction = predict(model=model, inputs=input_sample)

I now that this difference is due to the different format versions and everything works until loading, but when making the prediction I get the following error on wild-whale:

ValueError: Array shape (1024, 1024) does not map to axes [AxisInfo(id='batch', type='batch', maybe_singleton=True), AxisInfo(id='channel', type='channel', maybe_singleton=False), AxisInfo(id='y', type='space', maybe_singleton=False), AxisInfo(id='x', type='space', maybe_singleton=False)]

After some debugging I realized that on wild-whale the maybe_singleton attribute on the channel axis was set to False, while on affable-shark is set to True. I found on the code that the function assigning this value is create on AxisInfo.

More specifically, the line which is assigning that False value is this one. As it is a string axis and channel is represented as 'c' and not 'b', maybe_singleton is set to False. But in this case, the channel axis is a singleton where there is only one channel. Would this new line work?

maybe_singleton = axis == "b" or axis == "c"

I have tried like this on my local version of the core library and like that wild-whale is capable of predicting the (1024,1024) image without any issue.

Is this a good solution or would it change a lot the performance on other old models?

[FynnBe]

unfortunately your proposed solution maybe_singleton = axis == "b" or axis == "c" would be wrong for any channel axis with more than one channel... It is possible to improve the library to detect trivial SizeReferences and know that this particular channel axis may be a singleton axis, but as move to using model spec v0_5 which allows to specify a size reference per axis, an updated model should simply hardcode the number of output channels (if, as usually is the case, they are not dependent on the input size). So I encourage you to update wild-whale instead to avoid this problem and load 2d data correctly with a singleton batch and singelton channel axis.

[IvanHCenalmor]

Aaaa okey, I thought that then it was assessed with the shape or something similar. Then, it's true that it doesn't make sense for cases with more than one axis. Okey then maybe the solution of updating the wild-whale seems the most straightforward one.