Predictions are cut off for larger images on Apple M1/macOS11/tensorflow-metal 2.6

[maweigert]

Describe the bug

On Apple M1 using tensorflow-metal 2.6, predicting larger images (larger than ~1446 height/width) leads to results being cut off.

import numpy as np
import tensorflow as tf
from stardist.data import test_image_nuclei_2d
from stardist.models import StarDist2D

x = np.tile(test_image_nuclei_2d(),(4,4))/200

model = StarDist2D.from_pretrained('2D_versatile_fluo')

y, _ = model.predict(x)

Actually, this seems to have nothing to to with stardist but appears to be a bug probably specific to tensorflow-metal 2.6 for Apple Silicon/M1/maxOS11, so probably nothing we can do about it (update to macOS 12 might help in the future)

x = np.ones((1,1536,1536,1))

inp = tf.keras.layers.Input((None,None,1))
out = tf.keras.layers.Conv2D(1,3,padding='same')(inp)
model = tf.keras.models.Model(inp, out)
y = model.predict(x)

[psobolewskiPhD]

I'm still on tensorflow-macos 2.5, but on Monterey.

x = np.tile(test_image_nuclei_2d(),(4,4))
model = StarDist2D.from_pretrained("2D_versatile_fluo")
y, _ = model.predict_instances(normalize(x))

works:

Interestingly, on Big Sur I had issues with larger images and Stardist not labeling the bottom part of the image. I ended up using tiling, assuming it was a memory issue. Now on Monterey the same code doesn't require tiles, even for (3899, 6811)

[maweigert]

Hi @psobolewskiPhD ,

thanks a lot for reporting/testing this!

Yes, it seems that this is a version/platform depending bug, although the respective thread in the image.sc forum indicates it affects equally windows 10. I never saw this happening on linux (which I use for large scale predictions).

[psobolewskiPhD]

Re: Windows, I have no idea. ...but officially from Apple, the TF-macos and TF-metal are only for macOS 12... I used them both on Big Sur a lot, but maybe there is something that is needed in macOS 12. I will try to update to TF 2.6 this weekend and will report back.

[psobolewskiPhD]

Ok, so I tested using the code you posted on image.sc and using my own large image (a stitched mosaic 3899, 6811). On macOS 12 Monterey, both TF 2.5 and stardist 0.7.1 and TF 2.6 and stardist 0.7.3 work fine, fully labeling everything and at the same speed, ~4.8 s according to %%timeit. I was not able to run these types of image without tiling under Big Sur due to the incomplete labeling. Edit: btw napari-stardist plugin also works perfectly, while previously I always used tiling.

[maweigert]

Ok, can confirm that on macOS 12 with tf 2.6 everything is fine.

platform: macOS-12.0.1-arm64-arm-64bit
tf:       2.6.0

[psobolewskiPhD]

Perhaps the README just needs to be updated that only macOS 12 is supported on arm64? I'm not sure how this behaves on x86, but Apple is pretty clear about tensorflow arm64 being macOS 12 and up.