diff --git a/ocr4all_pixel_classifier/lib/architecture.py b/ocr4all_pixel_classifier/lib/architecture.py new file mode 100644 index 0000000..1af7edf --- /dev/null +++ b/ocr4all_pixel_classifier/lib/architecture.py @@ -0,0 +1,90 @@ +import enum +from functools import partial + + +class Architecture(enum.Enum): + FCN_SKIP = 'fcn_skip' + FCN = 'fcn' + RES_NET = 'image_res_net' + RES_UNET = 'res_unet' + MOBILE_NET = 'mobile_net' + UNET = 'unet' + EFFNETB0 = 'effb0' + EFFNETB1 = 'effb1' + EFFNETB2 = 'effb2' + EFFNETB3 = 'effb3' + EFFNETB4 = 'effb4' + EFFNETB5 = 'effb5' + EFFNETB6 = 'effb6' + EFFNETB7 = 'effb7' + + def __call__(self, *args, **kwargs): + return self.model() + + def model(self): + from ocr4all_pixel_classifier.lib.model import model_fcn_skip, model_fcn, res_net_fine_tuning, res_unet, \ + unet_with_mobile_net_encoder, unet, eff_net_fine_tuning + from efficientnet import tfkeras as efn + return { + Architecture.FCN_SKIP: model_fcn_skip, + Architecture.FCN: model_fcn, + Architecture.RES_NET: res_net_fine_tuning, + Architecture.RES_UNET: res_unet, + Architecture.MOBILE_NET: unet_with_mobile_net_encoder, + Architecture.UNET: unet, + Architecture.EFFNETB0: partial(eff_net_fine_tuning, efnet=efn.EfficientNetB0), + Architecture.EFFNETB1: partial(eff_net_fine_tuning, efnet=efn.EfficientNetB1), + Architecture.EFFNETB2: partial(eff_net_fine_tuning, efnet=efn.EfficientNetB2), + Architecture.EFFNETB3: partial(eff_net_fine_tuning, efnet=efn.EfficientNetB3), + Architecture.EFFNETB4: partial(eff_net_fine_tuning, efnet=efn.EfficientNetB4), + Architecture.EFFNETB5: partial(eff_net_fine_tuning, efnet=efn.EfficientNetB5), + Architecture.EFFNETB6: partial(eff_net_fine_tuning, efnet=efn.EfficientNetB6), + Architecture.EFFNETB7: partial(eff_net_fine_tuning, efnet=efn.EfficientNetB7), + }[self] + + def preprocess(self): + from efficientnet import tfkeras as efn + import tensorflow as tf + + return { + Architecture.FCN_SKIP: (default_preprocess, False), + Architecture.FCN: (default_preprocess, False), + Architecture.RES_NET: (tf.keras.applications.resnet50.preprocess_input, True), + Architecture.RES_UNET: (default_preprocess, False), + Architecture.MOBILE_NET: (tf.keras.applications.mobilenet_v2.preprocess_input, True), + Architecture.UNET: (default_preprocess, False), + Architecture.EFFNETB0: (efn.preprocess_input, True), + Architecture.EFFNETB1: (efn.preprocess_input, True), + Architecture.EFFNETB2: (efn.preprocess_input, True), + Architecture.EFFNETB3: (efn.preprocess_input, True), + Architecture.EFFNETB4: (efn.preprocess_input, True), + Architecture.EFFNETB5: (efn.preprocess_input, True), + Architecture.EFFNETB6: (efn.preprocess_input, True), + Architecture.EFFNETB7: (efn.preprocess_input, True), + }[self] + + +def default_preprocess(x): + return x / 255.0 + + +class Optimizers(enum.Enum): + ADAM = 'adam' + ADAMAX = 'adamax' + ADADELTA = 'adadelta' + ADAGRAD = 'adagrad' + RMSPROP = 'rmsprop' + SGD = 'sgd' + NADAM = 'nadam' + + def __call__(self, *args, **kwargs): + import tensorflow.keras.optimizers as tfo + return { + Optimizers.ADAM: tfo.Adam, + Optimizers.ADAMAX: tfo.Adamax, + Optimizers.ADADELTA: tfo.Adadelta, + Optimizers.ADAGRAD: tfo.Adagrad, + Optimizers.RMSPROP: tfo.RMSprop, + Optimizers.SGD: tfo.SGD, + Optimizers.NADAM: tfo.Nadam, + }[self] diff --git a/ocr4all_pixel_classifier/lib/model.py b/ocr4all_pixel_classifier/lib/model.py index 0cefbdb..2bc347d 100644 --- a/ocr4all_pixel_classifier/lib/model.py +++ b/ocr4all_pixel_classifier/lib/model.py @@ -1,5 +1,3 @@ -import enum -from functools import partial from typing import Union, List import efficientnet.tfkeras as efn @@ -43,6 +41,7 @@ def crop(input_tensors): output = tf.image.crop_to_bounding_box(input, 0, 0, tf.gather(shape, 1) - px, tf.gather(shape, 2) - py) return output + def model_fcn_skip(input: Tensors, n_classes: int): input_image = input[0] padding = tf.keras.layers.Lambda(lambda x: calculate_padding(x))(input_image) @@ -308,16 +307,17 @@ def bn_act(x, act=True, batch_normailzation=False): return model -def res_net_fine_tuning(input: Tensors, n_classes: int): - def conv_block_simple(prevlayer, filters, prefix, strides=(1, 1), batch_nm=False): - conv = tf.keras.layers.Conv2D(filters, (3, 3), - padding="same", kernel_initializer="he_normal", - strides=strides, name=prefix + "_conv")(prevlayer) - if batch_nm: - conv = tf.keras.layers.BatchNormalization(name=prefix + "_bn")(conv) - conv = tf.keras.layers.Activation('relu', name=prefix + "_activation")(conv) - return conv +def conv_block_simple(prevlayer, filters, prefix, strides=(1, 1), batch_nm=False): + conv = tf.keras.layers.Conv2D(filters, (3, 3), + padding="same", kernel_initializer="he_normal", + strides=strides, name=prefix + "_conv")(prevlayer) + if batch_nm: + conv = tf.keras.layers.BatchNormalization(name=prefix + "_bn")(conv) + conv = tf.keras.layers.Activation('relu', name=prefix + "_activation")(conv) + return conv + +def res_net_fine_tuning(input: Tensors, n_classes: int): input_image = input[0] padding = tf.keras.layers.Lambda(lambda x: calculate_padding(x), name='lambda_calc_pad')(input_image) @@ -366,15 +366,6 @@ def conv_block_simple(prevlayer, filters, prefix, strides=(1, 1), batch_nm=False def eff_net_fine_tuning(input: Tensors, n_classes: int, efnet=efn.EfficientNetB1): - def conv_block_simple(prevlayer, filters, prefix, strides=(1, 1), batch_nm=False): - conv = tf.keras.layers.Conv2D(filters, (3, 3), - padding="same", kernel_initializer="he_normal", - strides=strides, name=prefix + "_conv")(prevlayer) - if batch_nm: - conv = tf.keras.layers.BatchNormalization(name=prefix + "_bn")(conv) - conv = tf.keras.layers.Activation('relu', name=prefix + "_activation")(conv) - return conv - input_image = input[0] padding = tf.keras.layers.Lambda(lambda x: calculate_padding(x))(input_image) padded = tf.keras.layers.Lambda(pad)([input_image, padding]) @@ -414,90 +405,3 @@ def conv_block_simple(prevlayer, filters, prefix, strides=(1, 1), batch_nm=False model = tf.keras.Model(input, out, name='effb0') return model - - -def default_preprocess(x): - return x / 255.0 - - -class Architecture(enum.Enum): - FCN_SKIP = 'fcn_skip' - FCN = 'fcn' - RES_NET = 'image_res_net' - RES_UNET = 'res_unet' - MOBILE_NET = 'mobile_net' - UNET = 'unet' - EFFNETB0 = 'effb0' - EFFNETB1 = 'effb1' - EFFNETB2 = 'effb2' - EFFNETB3 = 'effb3' - EFFNETB4 = 'effb4' - EFFNETB5 = 'effb5' - EFFNETB6 = 'effb6' - EFFNETB7 = 'effb7' - - def __call__(self, *args, **kwargs): - return self.model() - - def model(self): - return { - Architecture.FCN_SKIP: model_fcn_skip, - Architecture.FCN: model_fcn, - Architecture.RES_NET: res_net_fine_tuning, - Architecture.RES_UNET: res_unet, - Architecture.MOBILE_NET: unet_with_mobile_net_encoder, - Architecture.UNET: unet, - Architecture.EFFNETB0: partial(eff_net_fine_tuning, efnet=efn.EfficientNetB0), - Architecture.EFFNETB1: partial(eff_net_fine_tuning, efnet=efn.EfficientNetB1), - Architecture.EFFNETB2: partial(eff_net_fine_tuning, efnet=efn.EfficientNetB2), - Architecture.EFFNETB3: partial(eff_net_fine_tuning, efnet=efn.EfficientNetB3), - Architecture.EFFNETB4: partial(eff_net_fine_tuning, efnet=efn.EfficientNetB4), - Architecture.EFFNETB5: partial(eff_net_fine_tuning, efnet=efn.EfficientNetB5), - Architecture.EFFNETB6: partial(eff_net_fine_tuning, efnet=efn.EfficientNetB6), - Architecture.EFFNETB7: partial(eff_net_fine_tuning, efnet=efn.EfficientNetB7), - }[self] - - def preprocess(self): - return { - Architecture.FCN_SKIP: (default_preprocess, False), - Architecture.FCN: (default_preprocess, False), - Architecture.RES_NET: (tf.keras.applications.resnet50.preprocess_input, True), - Architecture.RES_UNET: (default_preprocess, False), - Architecture.MOBILE_NET: (tf.keras.applications.mobilenet_v2.preprocess_input, True), - Architecture.UNET: (default_preprocess, False), - Architecture.EFFNETB0: (efn.preprocess_input, True), - Architecture.EFFNETB1: (efn.preprocess_input, True), - Architecture.EFFNETB2: (efn.preprocess_input, True), - Architecture.EFFNETB3: (efn.preprocess_input, True), - Architecture.EFFNETB4: (efn.preprocess_input, True), - Architecture.EFFNETB5: (efn.preprocess_input, True), - Architecture.EFFNETB6: (efn.preprocess_input, True), - Architecture.EFFNETB7: (efn.preprocess_input, True), - }[self] - - -class Optimizers(enum.Enum): - ADAM = 'adam' - ADAMAX = 'adamax' - ADADELTA = 'adadelta' - ADAGRAD = 'adagrad' - RMSPROP = 'rmsprop' - SGD = 'sgd' - NADAM = 'nadam' - - def __call__(self, *args, **kwargs): - return { - Optimizers.ADAM: tf.keras.optimizers.Adam, - Optimizers.ADAMAX: tf.keras.optimizers.Adamax, - Optimizers.ADADELTA: tf.keras.optimizers.Adadelta, - Optimizers.ADAGRAD: tf.keras.optimizers.Adagrad, - Optimizers.RMSPROP: tf.keras.optimizers.RMSprop, - Optimizers.SGD: tf.keras.optimizers.SGD, - Optimizers.NADAM: tf.keras.optimizers.Nadam, - }[self] - - -if __name__ == '__main__': - #effnet_base = efn.EfficientNetB1(weights='imagenet', include_top=False, input_shape=(256, 256, 3)) - #print(effnet_base.summary()) - pass \ No newline at end of file diff --git a/ocr4all_pixel_classifier/lib/network.py b/ocr4all_pixel_classifier/lib/network.py index 3236e64..7bf331f 100644 --- a/ocr4all_pixel_classifier/lib/network.py +++ b/ocr4all_pixel_classifier/lib/network.py @@ -6,7 +6,7 @@ import tensorflow as tf from ocr4all_pixel_classifier.lib.callback import TrainProgressCallback, TrainProgressCallbackWrapper -from ocr4all_pixel_classifier.lib.model import Optimizers, Architecture +from .architecture import Architecture, Optimizers from ocr4all_pixel_classifier.lib.trainer import TrainSettings, AugmentationSettings from ocr4all_pixel_classifier.lib.util import image_to_batch, gray_to_rgb from .dataset import Dataset, SingleData diff --git a/ocr4all_pixel_classifier/lib/trainer.py b/ocr4all_pixel_classifier/lib/trainer.py index c660833..140a00d 100644 --- a/ocr4all_pixel_classifier/lib/trainer.py +++ b/ocr4all_pixel_classifier/lib/trainer.py @@ -2,7 +2,7 @@ from ocr4all_pixel_classifier.lib.callback import TrainProgressCallback from typing import NamedTuple, Optional, List from ocr4all_pixel_classifier.lib.metrics import Loss, Monitor -from ocr4all_pixel_classifier.lib.model import Architecture, Optimizers +from ocr4all_pixel_classifier.lib.architecture import Architecture, Optimizers import numpy as np import logging import tensorflow as tf