Use bioimageio.core to generate model export (WIP)

[constantinpape]

The cell is running through, but there are still several issues:

• the keras model is exported with the custom loss and cannot be loaded again (how do we get rid of this?)
• the tensorflow weights are not added yet (will work on a weight converter based on https://github.com/deepimagej/pydeepimagej/blob/5aaf0e71f9b04df591d5ca596f0af633a7e024f5/pydeepimagej/yaml/create_config.py#L236 tomorrow)
• test_model cannot be used because of the dask import error caused by mismatch with pyyaml version (fixed by installing an older pyyaml < 6, see PyYAML 6.0 load() function requires Loader argument yaml/pyyaml#576)
• there is some additional post-processing added in the deepimagej part (scale_range and scale_linear) that I don't understand and that I am not adding right now.
• the pixel size is not written as properties to the tifs (not sure if this matters)

The notebook is available at
https://colab.research.google.com/github/constantinpape/ZeroCostDL4Mic/blob/master/Colab_notebooks/BioImage.io notebooks/U-Net_2D_ZeroCostDL4Mic_BioImageModelZoo_export.ipynb

(p.s. white spaces in urls are evil ;))

cc @esgomezm

[esgomezm]

Hi @constantinpape

Let me know if you prefer a PR in your forked repo. Meanwhile:

• To get only the model architecture I did the following:

from keras.models import Model
unet = load_model(os.path.join(full_QC_model_path, 'weights_best.hdf5'), custom_objects={'_weighted_binary_crossentropy': weighted_binary_crossentropy(np.ones(2))})
# Create a model without compilation so it can be used in any other environment.
input  = unet.input
single_output = unet.output
unet = Model(input, single_output)

I've just seen that in this notebook they import keras directly rather than from TF. This could fail again in the bioimageio.core validator.

• Postprocessing:
  In the notebooks, the raw output values ("probabilities") are scaled from the [0,1] range to the [0, 255] range and then, the optimal threshold is applied in this new scale. The scale_range maybe can be removed as it doesn't really do anything (the output values are already in the [0,1] range).

• The function in the notebook that calculates the pixel size provides a value of "1 pixel" whenever this value is not given in the image. That can be a good compromise. Yet another chance is to always ask for the pixel size.

[constantinpape]

@esgomezm I continued a bit on this and have used your code snippet to export the "uncompiled" keras model.
I also updated the postprocessing a bit and translated the threshold to float, which I think makes it a bit easier to unerstand what's going on ...
I haven't switched from import keras to from keras import tensorflow yet, because the keras imports are pretty complex:

from keras import models
from keras.models import Model, load_model
from keras.layers import Input, Conv2D, MaxPooling2D, Dropout, concatenate, UpSampling2D
from keras.optimizers import Adam
# from keras.callbacks import ModelCheckpoint, LearningRateScheduler, CSVLogger # we currently don't use any other callbacks from ModelCheckpoints
from keras.callbacks import ModelCheckpoint
from keras.callbacks import ReduceLROnPlateau
from keras.preprocessing.image import ImageDataGenerator, img_to_array, load_img
from keras import backend as keras

and I also haven't implemented tensorflow weight export yet. (It would be good to implement a function in bioimageio.core for this, maybe we can discuss this now.)

[esgomezm]

I haven't switched from import keras to from keras import tensorflow yet, because the keras imports are pretty complex

Do you mean from tensorflow import keras? I'll do this part.

I also haven't implemented tensorflow weight export yet. (It would be good to implement a function in bioimageio.core for this, maybe we can discuss this now.)

For the export to TF, we can take the function in pydeepimagej. There are just few things to take into account:

• The folder in which the model is exported cannot exist before. If so, it needs to be removed and then exported. Otherwise, TF will complain as it does not replace any existing file.
• The exported model needs to be zipped according to the BMZ specs and the previous folder should be removed to avoid redundancy.

[constantinpape]

Do you mean from tensorflow import keras? I'll do this part.

Yes, sorry I always mix this up.

For the export to TF, we can take the function in pydeepimagej. There are just few things to take into account:

* The folder in which the model is exported cannot exist before. If so, it needs to be removed and then exported. Otherwise, TF will complain as it does not replace any existing file.

* The exported model needs to be zipped according to the BMZ specs and the previous folder should be removed to avoid redundancy.

Yes, it would be good if we could re-implement this function in bioimageio.core.weight_converter.keras.

[constantinpape]

I have updated the notebook now so that it also adds the tensorflow weights.

The model is available here:
https://drive.google.com/file/d/1-DBzP-YEd8JaDpuBajHDRkkBYyWVloen/view?usp=sharing
I have still exported it based on your old model, somehow I can't access the one you shared with me yesterday. So test-model fails for this model, but otherwise I can run it. I hope this issue is fixed when using a model that was trained correctly using the from tensorflow import keras.

@esgomezm could you again share the new model with me?
Could you also test if you can run the model I have exported in deepimagej?

[esgomezm]

Nice. Please, find it here: https://drive.google.com/drive/u/0/folders/1-vGzP95Po_Pu0n6K1ei5GYKlGU2aqI74

[constantinpape]

@esgomezm I had to fix a few more things, but now it's almost working:
I have a version without thresholding now that works (it's added as an option to the export field).
The version with thresholding is disagreeing for 27 pixels, so there is probably some numerical inacuracies going on ...
Maybe the post-processing I have could also be improved.

Here's the current model (without threshold): https://drive.google.com/file/d/1-LWeHUBUp_1xeqUAxph6ugby659N4y6r/view?usp=sharing

[constantinpape]

@esgomezm I have updated the thresholding now so that the tests also pass when applying it.
Also, Carlos has made some updates to deepimagej and has tested the model and it works!
(There is one minor issue we still need to fix in bioimageio.core: bioimage-io/core-bioimage-io-python#184)

[constantinpape]

@esgomezm we fixed a few more things and now everything should be working if you use:
bioimageio.core 0.4.6.post0
bioimageio.spec 0.4.0.post1
Let me know if you run into any issues with this.

At some point we should also tackle the model import. That should be pretty straightforward:

model_spec = bioimageio.core.load_resource_description(doi)  # load the model by providing the doi
weight_path = model_spec.weights["keras_hdf5"].source
model = keras.models.load_model(weight_path)
# probably do some more stuff to compile the model with the custom weighted bce loss again

[constantinpape]

@esgomezm I have updated the notebook now so that it uses the correct versions of the bioimageio.core library and I have also implemented loading a pretrained model from bioimage.io.
This is good to go from my side. Let me know if you think any more changes are necessary.

Here is the most recent model I have exported with the notebook: https://drive.google.com/file/d/1-LWeHUBUp_1xeqUAxph6ugby659N4y6r/view?usp=sharing
I have checked that I can run it locally, it would be good if you could also test that it runs in deepImageJ.

[esgomezm]

Hi @constantinpape

Is there any 2D UNet uploaded to zenodo so I can try the bioimage.io model load from the notebook? Currently the bioimage model zoo is empty and the old version is not working so I cannot access any model

[constantinpape]

@esgomezm yes, you can use this model to try: https://drive.google.com/file/d/1-LWeHUBUp_1xeqUAxph6ugby659N4y6r/view

[esgomezm]

Hey @constantinpape, Happy New Year!

I managed to make the notebook working also with a pre-trained model and the bioimage.io model also works in deepImageJ :)

I'm cleaning the output to update it in my repo so I can then accept your pull request here with the updated notebook.

Also, I tried to upload the model to the BMZ but I get the following message referring to the version number. I think it has to do more with the bioimageio.core export model functionality than with the notebook:

[esgomezm]

Solved the part of the validator. The name in the rdf.yaml needs to be the same we write in the upload formulary.

There's a PR already opened from my repository and I cannot create a new one until this is either accepted or closed. Please let me know how you want to proceed.

[constantinpape]

Hi @esgomezm and also happy new year :).
Glad you could figure out the issues. Regarding the PR: I guess you refer to this one: constantinpape#2. I just saw this now, but I am a bit confused by what's going on there...

[constantinpape]

I merged your PR @esgomezm and I think this is good to go.
Let me know if you need anything else from my side here :).