faustomorales / keras-ocr

A packaged and flexible version of the CRAFT text detector and Keras CRNN recognition model.
https://keras-ocr.readthedocs.io/
MIT License
1.38k stars 355 forks source link

Saving and Loading custom trained model #111

Closed srivatsan88 closed 4 years ago

srivatsan88 commented 4 years ago

@faustomorales .. This is really an amazing package and simple to use as well

I am facing issue in loading a fine tuned model and using it in keras_ocr.pipeline.Pipeline

It works when I give detector and recognizer of trained model in same session but when I save and load custom trained recognizer the load is not helping much. Below is what I did

recognizer.model.save('custom_words.hdf5') and while loading used

recognizer_l = keras_ocr.recognition.keras.models.load_model('custom_words.hdf5')

later I am using it in pipeline as below

pipeline = keras_ocr.pipeline.Pipeline(detector=detector,recognizer=recognizer_l) and it is failing to initialize recognizer

I see that model is downloaded from PRETRAINED_WEIGHTS variable parameters in recognition.py file and it defaults to Kurapan

Is there a better way of loading my custom weights or do i need to override PRETRAINED_WEIGHTS variable?

Not sure if I am missing anything here or something I missed. I did loaded model.summary as well and do see weights associated with it

faustomorales commented 4 years ago

One could think about this in terms of four different objects.

We use these wrappers in order to abstract out pre- and post-processing steps that would result in a lot of boilerplate if you just used the raw keras.models.Model all the time. For example, the keras_ocr.recognition.Recognizer keeps track of your alphabet and handles converting numeric model outputs to and from string representations.

The problem in your code is that recognizer_l = keras_ocr.recognition.keras.models.load_model('custom_words.hdf5') will only load the keras.models.Model instance but not any of the wrapping provided by keras_ocr.recognition.Recognizer. When you pass recognizer_l to keras_ocr.pipeline.Pipeline, it doesn't work properly because a keras_ocr.recognition.Recognizer is expected, not a raw keras.models.Model.

The following example I believe is more like what you're after.

# Previous to this assume you've trained a recognizer.
recognizer.model.save_weights('custom_words.h5')

# Let's reload the recognizer. First, instantiate a recognizer with the *same*
# arguments as those you used earlier (e.g., alphabet).
reloaded_recognizer = keras_ocr.recognition.Recognizer(...)

# Now load the weights you saved earlier.
reloaded_recognizer.model.load_weights('custom_words.h5')

The reason why you have to do this in two steps is fundamentally because keras_ocr doesn't have an end-to-end serialization strategy (e.g., for alphabets and other parameters). So you have to instantiate the "outer wrapper" first and then load the model weights into it (which uses Keras' serialization functionality for the actual model component).

I hope this has been helpful! Since this is not a bug, I'm closing this issue for the moment. Thanks.

srivatsan88 commented 4 years ago

Thanks @faustomorales .. Makes sense.. I was trying to pickle entire recognition object to overcome what you said but later ended up forking the code to add my new weights to PRETRAINED_WEIGHTS in recognition.py object which worked. What you mentioned seems to me a neat option. Will try and let you know in case if I face any issue. Thank for your help on this