deepimagej / python4deepimagej

BSD 2-Clause "Simplified" License
14 stars 5 forks source link

Exporting StarDist model to DeepImageJ model #2

Closed romainGuiet closed 4 years ago

romainGuiet commented 4 years ago

Hi @deepimagej ,

I would like to use my own StarDist model ( we had to add images from our dataset with new annotations) but I can not find a way to export it (yet).

I tried re-using your code example but got these troubles :

-I’ve loaded my model using the function from StarDist and got :

AttributeError                            Traceback (most recent call last)
<ipython-input-23-843654a17ca2> in <module>
     15 
     16 signature = tf.saved_model.signature_def_utils.predict_signature_def(
---> 17             inputs  = {'input':  model.input},
     18             outputs = {'output': model.output})
     19 

AttributeError: 'StarDist2D' object has no attribute 'input'

so I guessed that StarDist model might not have an input !

So I tried to use the keras function from your code and got :

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-27-81d916cc2deb> in <module>
      8 
      9 K.set_learning_phase(1)
---> 10 model = keras.models.load_model(model_path_str)
     11 #model = StarDist2D(None, name=model_name , basedir=model_path)
     12 

m:\2-stardist-env\stardist_0.3.2\lib\site-packages\keras\engine\saving.py in load_wrapper(*args, **kwargs)
    456                 os.remove(tmp_filepath)
    457             return res
--> 458         return load_function(*args, **kwargs)
    459 
    460     return load_wrapper

m:\2-stardist-env\stardist_0.3.2\lib\site-packages\keras\engine\saving.py in load_model(filepath, custom_objects, compile)
    548     if H5Dict.is_supported_type(filepath):
    549         with H5Dict(filepath, mode='r') as h5dict:
--> 550             model = _deserialize_model(h5dict, custom_objects, compile)
    551     elif hasattr(filepath, 'write') and callable(filepath.write):
    552         def load_function(h5file):

m:\2-stardist-env\stardist_0.3.2\lib\site-packages\keras\engine\saving.py in _deserialize_model(h5dict, custom_objects, compile)
    237         return obj
    238 
--> 239     model_config = h5dict['model_config']
    240     if model_config is None:
    241         raise ValueError('No model found in config.')

m:\2-stardist-env\stardist_0.3.2\lib\site-packages\keras\utils\io_utils.py in __getitem__(self, attr)
    316             else:
    317                 if self.read_only:
--> 318                     raise ValueError('Cannot create group in read-only mode.')
    319                 val = H5Dict(self.data.create_group(attr))
    320         return val

ValueError: Cannot create group in read-only mode.

so the StarDist model might not be compatible with keras function!

Would you have an example of code for stardist?

Best regards,

Romain

esgomezm commented 4 years ago

Hi @romainGuiet,

StarDist library works a bit different so the code needs to be adapted. There is a new notebook in the repository to store StarDist models as TensorFlow SavedModel:

https://github.com/deepimagej/python4deepimagej/blob/master/export_StarDist_to_TensorFlow_SavedModel.ipynb

It is written for StarDist version 0.3.6 (most updated version until now):

Once you read your model using StarDist2D:

model_paper = StarDist2D(None, name='your_model', basedir='/path_to_your_model/')
model_paper.load_weights('weights_best.h5')

the input and output are given as

# Input

model_paper.keras_model.input[0]

# Output 1: the object probability (normalized euclidean distance to the nearest background point)

model_paper.keras_model.output[0]

# Output 2: 32 radial distances to the boundary of the object

model_paper.keras_model.output[1]

Hence, to export it as a TensorFlow SavedModel you can now use a very similar code:

OUTPUT_DIR = "/content/drive/My Drive/the_path_where_you_want_to_save_your_model/new_folder"

builder = tf.saved_model.builder.SavedModelBuilder(OUTPUT_DIR)

# StarDist has two different outputs. DeepImageJ can only read one of them, so 
# we concatenate them as different channels to use them in ImageJ.

signature = tf.saved_model.signature_def_utils.predict_signature_def(
            inputs  = {'input':  model_paper.keras_model.input[0]},
            # concatenate the output of StarDist
            outputs = {'output': concatenate([model_paper.keras_model.output[0],model_paper.keras_model.output[1]], axis = 3)})

signature_def_map = { tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature }

builder.add_meta_graph_and_variables(K.get_session(), [tf.saved_model.tag_constants.SERVING],
                                             signature_def_map=signature_def_map)
builder.save()

I hope this helps!

Thank you for reporting the issue

romainGuiet commented 4 years ago

Hi @esgomezm, That's awesome! Thanks a lot for your efforts! Best,