[ { "path": ".gitignore", "content": "# Byte-compiled / optimized / DLL files\n__pycache__/\n*.py[cod]\n*$py.class\n\n# C extensions\n*.so\n\n# Distribution / packaging\n.Python\nbuild/\ndevelop-eggs/\ndist/\ndownloads/\neggs/\n.eggs/\nlib/\nlib64/\nparts/\nsdist/\nvar/\nwheels/\n*.egg-info/\n.installed.cfg\n*.egg\nMANIFEST\n\n# PyInstaller\n# Usually these files are written by a python script from a template\n# before PyInstaller builds the exe, so as to inject date/other infos into it.\n*.manifest\n*.spec\n\n# Installer logs\npip-log.txt\npip-delete-this-directory.txt\n\n# Unit test / coverage reports\nhtmlcov/\n.tox/\n.coverage\n.coverage.*\n.cache\nnosetests.xml\ncoverage.xml\n*.cover\n.hypothesis/\n.pytest_cache/\n\n# Translations\n*.mo\n*.pot\n\n# Django stuff:\n*.log\nlocal_settings.py\ndb.sqlite3\n\n# Flask stuff:\ninstance/\n.webassets-cache\n\n# Scrapy stuff:\n.scrapy\n\n# Sphinx documentation\ndocs/_build/\n\n# PyBuilder\ntarget/\n\n# Jupyter Notebook\n.ipynb_checkpoints\n\n# pyenv\n.python-version\n\n# celery beat schedule file\ncelerybeat-schedule\n\n# SageMath parsed files\n*.sage.py\n\n# Environments\n.env\n.venv\nenv/\nvenv/\nENV/\nenv.bak/\nvenv.bak/\n\n# Spyder project settings\n.spyderproject\n.spyproject\n\n# Rope project settings\n.ropeproject\n\n# mkdocs documentation\n/site\n\n# mypy\n.mypy_cache/\n" }, { "path": "LICENSE", "content": "MIT License\n\nCopyright (c) 2019 Larry\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n" }, { "path": "README.md", "content": "# MobileNetV3\nA Keras implementation of MobileNetV3 and Lite R-ASPP Semantic Segmentation (Under Development).\n\nAccording to the paper: [Searching for MobileNetV3](https://arxiv.org/abs/1905.02244?context=cs)\n\n## Requirement\n- Python 3.6\n- Tensorflow-gpu 1.10.0 \n- Keras 2.2.4\n\n## Train the model\n\n The ```config/config.json``` file provide a config for training.\n\n### Train the classification\n\n**The dataset folder structure is as follows:**\n\n\t| - data/\n\t\t| - train/\n\t \t\t| - class 0/\n\t\t\t\t| - image.jpg\n\t\t\t\t\t....\n\t\t\t| - class 1/\n\t\t\t ....\n\t\t\t| - class n/\n\t\t| - validation/\n\t \t\t| - class 0/\n\t\t\t| - class 1/\n\t\t\t ....\n\t\t\t| - class n/\n\n**Run command below to train the model:**\n\n```\npython train_cls.py\n```\n\n## Acknowledgement\n\nThank [@fzyzcjy](https://github.com/fzyzcjy) for your help in this project.\n\n## Reference\n\n\t@article{MobileNetv3, \n\t title={Searching for MobileNetV3}, \n\t author={Andrew Howard, Mark Sandler, Grace Chu, Liang-Chieh Chen, Bo Chen, Mingxing Tan, Weijun Wang, Yukun Zhu, Ruoming Pang Vijay Vasudevan, Quoc V. Le, Hartwig Adam},\n\t journal={arXiv preprint arXiv:1905.02244},\n\t year={2019}\n\t}\n\n\n## Copyright\nSee [LICENSE](LICENSE) for details.\n" }, { "path": "config/config.json", "content": "{\n \"model\": \"small\",\n \"height\": 224,\n \"width\": 224,\n \"class_number\": 100,\n \"learning_rate\": 0.001,\n \"batch\": 2,\n \"epochs\": 20,\n \"train_dir\": \"data/train\",\n \"eval_dir\": \"data/eval\",\n \"save_dir\": \"save\",\n \"weights\": \"\"\n}\n" }, { "path": "model/LR_ASPP.py", "content": "\"\"\"Lite R-ASPP Semantic Segmentation based on MobileNetV3.\n\"\"\"\n\n\nfrom keras.models import Model\nfrom keras.layers import Conv2D, AveragePooling2D, BatchNormalization, Activation, Multiply, Add\nfrom keras.utils.vis_utils import plot_model\nfrom model.layers.bilinear_upsampling import BilinearUpSampling2D\n\n\nclass LiteRASSP:\n def __init__(self, input_shape, n_class=19, alpha=1.0, weights=None, backbone='small'):\n \"\"\"Init.\n\n # Arguments\n input_shape: An integer or tuple/list of 3 integers, shape\n of input tensor (should be 1024 × 2048 or 512 × 1024 according \n to the paper).\n n_class: Integer, number of classes.\n alpha: Integer, width multiplier for mobilenetV3.\n weights: String, weights for mobilenetv3.\n backbone: String, name of backbone (must be small or large).\n \"\"\"\n self.shape = input_shape\n self.n_class = n_class\n self.alpha = alpha\n self.weights = weights\n self.backbone = backbone\n\n def _extract_backbone(self):\n \"\"\"extract feature map from backbone.\n \"\"\"\n if self.backbone == 'large':\n from model.mobilenet_v3_large import MobileNetV3_Large\n\n model = MobileNetV3_Large(self.shape, self.n_class, alpha=self.alpha, include_top=False).build()\n layer_name8 = 'batch_normalization_13'\n layer_name16 = 'add_5'\n elif self.backbone == 'small':\n from model.mobilenet_v3_small import MobileNetV3_Small\n\n model = MobileNetV3_Small(self.shape, self.n_class, alpha=self.alpha, include_top=False).build()\n layer_name8 = 'batch_normalization_7'\n layer_name16 = 'add_2'\n else:\n raise Exception('Invalid backbone: {}'.format(self.backbone))\n\n if self.weights is not None:\n model.load_weights(self.weights, by_name=True)\n\n inputs= model.input\n # 1/8 feature map.\n out_feature8 = model.get_layer(layer_name8).output\n # 1/16 feature map.\n out_feature16 = model.get_layer(layer_name16).output\n\n return inputs, out_feature8, out_feature16\n\n def build(self, plot=False):\n \"\"\"build Lite R-ASPP.\n\n # Arguments\n plot: Boolean, weather to plot model.\n\n # Returns\n model: Model, model.\n \"\"\"\n inputs, out_feature8, out_feature16 = self._extract_backbone()\n\n # branch1\n x1 = Conv2D(128, (1, 1))(out_feature16)\n x1 = BatchNormalization()(x1)\n x1 = Activation('relu')(x1)\n\n # branch2\n s = x1.shape\n\n x2 = AveragePooling2D(pool_size=(49, 49), strides=(16, 20))(out_feature16)\n x2 = Conv2D(128, (1, 1))(x2)\n x2 = Activation('sigmoid')(x2)\n x2 = BilinearUpSampling2D(target_size=(int(s[1]), int(s[2])))(x2)\n\n # branch3\n x3 = Conv2D(self.n_class, (1, 1))(out_feature8)\n\n # merge1\n x = Multiply()([x1, x2])\n x = BilinearUpSampling2D(size=(2, 2))(x)\n x = Conv2D(self.n_class, (1, 1))(x)\n\n # merge2\n x = Add()([x, x3])\n\n # out\n x = Activation('softmax')(x)\n\n model = Model(inputs=inputs, outputs=x)\n\n if plot:\n plot_model(model, to_file='images/LR_ASPP.png', show_shapes=True)\n\n return model\n" }, { "path": "model/layers/bilinear_upsampling.py", "content": "\"\"\"Keras BilinearUpSampling2D Layer.\n\"\"\"\nimport numpy as np\nimport tensorflow as tf\nimport keras.backend as K\nfrom keras.engine.topology import Layer, InputSpec\n\n\ndef resize_images_bilinear(X, height_factor=1, width_factor=1, target_height=None, target_width=None, data_format='default'):\n '''Resizes the images contained in a 4D tensor of shape\n\n - [batch, channels, height, width] (for 'channels_first' data_format)\n - [batch, height, width, channels] (for 'channels_last' data_format)\n by a factor of (height_factor, width_factor). Both factors should be\n positive integers.\n '''\n if data_format == 'default':\n data_format = K.image_data_format()\n\n if data_format == 'channels_first':\n original_shape = K.int_shape(X)\n\n if target_height and target_width:\n new_shape = tf.constant(np.array((target_height, target_width)).astype('int32'))\n else:\n new_shape = tf.shape(X)[2:]\n new_shape *= tf.constant(np.array([height_factor, width_factor]).astype('int32'))\n\n X = K.permute_dimensions(X, [0, 2, 3, 1])\n X = tf.image.resize_bilinear(X, new_shape)\n X = K.permute_dimensions(X, [0, 3, 1, 2])\n\n if target_height and target_width:\n X.set_shape((None, None, target_height, target_width))\n else:\n X.set_shape((None, None, original_shape[2] * height_factor, original_shape[3] * width_factor))\n\n return X\n elif data_format == 'channels_last':\n original_shape = K.int_shape(X)\n\n if target_height and target_width:\n new_shape = tf.constant(np.array((target_height, target_width)).astype('int32'))\n else:\n new_shape = tf.shape(X)[1:3]\n new_shape *= tf.constant(np.array([height_factor, width_factor]).astype('int32'))\n\n X = tf.image.resize_bilinear(X, new_shape)\n\n if target_height and target_width:\n X.set_shape((None, target_height, target_width, None))\n else:\n X.set_shape((None, original_shape[1] * height_factor, original_shape[2] * width_factor, None))\n\n return X\n else:\n raise Exception('Invalid data_format: ' + data_format)\n\n\nclass BilinearUpSampling2D(Layer):\n def __init__(self, size=(1, 1), target_size=None, data_format='default', **kwargs):\n \"\"\"Init.\n size: factor to original shape (ie. original-> size * original).\n target size: target size (ie. original->target).\n \"\"\"\n if data_format == 'default':\n data_format = K.image_data_format()\n self.size = tuple(size)\n\n if target_size is not None:\n self.target_size = tuple(target_size)\n else:\n self.target_size = None\n assert data_format in {'channels_last', 'channels_first'}, 'data_format must be in {tf, th}'\n\n self.data_format = data_format\n self.input_spec = [InputSpec(ndim=4)]\n\n super(BilinearUpSampling2D, self).__init__(**kwargs)\n\n def compute_output_shape(self, input_shape):\n if self.data_format == 'channels_first':\n width = int(self.size[0] * input_shape[2] if input_shape[2] is not None else None)\n height = int(self.size[1] * input_shape[3] if input_shape[3] is not None else None)\n\n if self.target_size is not None:\n width = self.target_size[0]\n height = self.target_size[1]\n return (input_shape[0],\n input_shape[1],\n width,\n height)\n elif self.data_format == 'channels_last':\n width = int(self.size[0] * input_shape[1] if input_shape[1] is not None else None)\n height = int(self.size[1] * input_shape[2] if input_shape[2] is not None else None)\n\n if self.target_size is not None:\n width = self.target_size[0]\n height = self.target_size[1]\n return (input_shape[0],\n width,\n height,\n input_shape[3])\n else:\n raise Exception('Invalid data_format: ' + self.data_format)\n\n def call(self, x, mask=None):\n if self.target_size is not None:\n return resize_images_bilinear(x, target_height=self.target_size[0], target_width=self.target_size[1], data_format=self.data_format)\n else:\n return resize_images_bilinear(x, height_factor=self.size[0], width_factor=self.size[1], data_format=self.data_format)\n\n def get_config(self):\n config = {'size': self.size, 'target_size': self.target_size}\n base_config = super(BilinearUpSampling2D, self).get_config()\n\n return dict(list(base_config.items()) + list(config.items()))\n" }, { "path": "model/mobilenet_base.py", "content": "\"\"\"MobileNet v3 models for Keras.\n# Reference\n [Searching for MobileNetV3](https://arxiv.org/abs/1905.02244?context=cs)\n\"\"\"\n\n\nfrom keras.layers import Conv2D, DepthwiseConv2D, Dense, GlobalAveragePooling2D\nfrom keras.layers import Activation, BatchNormalization, Add, Multiply, Reshape\n\nfrom keras import backend as K\n\n\nclass MobileNetBase:\n def __init__(self, shape, n_class, alpha=1.0):\n \"\"\"Init\n \n # Arguments\n input_shape: An integer or tuple/list of 3 integers, shape\n of input tensor.\n n_class: Integer, number of classes.\n alpha: Integer, width multiplier.\n \"\"\"\n self.shape = shape\n self.n_class = n_class\n self.alpha = alpha\n\n def _relu6(self, x):\n \"\"\"Relu 6\n \"\"\"\n return K.relu(x, max_value=6.0)\n\n def _hard_swish(self, x):\n \"\"\"Hard swish\n \"\"\"\n return x * K.relu(x + 3.0, max_value=6.0) / 6.0\n\n def _return_activation(self, x, nl):\n \"\"\"Convolution Block\n This function defines a activation choice.\n\n # Arguments\n x: Tensor, input tensor of conv layer.\n nl: String, nonlinearity activation type.\n\n # Returns\n Output tensor.\n \"\"\"\n if nl == 'HS':\n x = Activation(self._hard_swish)(x)\n if nl == 'RE':\n x = Activation(self._relu6)(x)\n\n return x\n\n def _conv_block(self, inputs, filters, kernel, strides, nl):\n \"\"\"Convolution Block\n This function defines a 2D convolution operation with BN and activation.\n\n # Arguments\n inputs: Tensor, input tensor of conv layer.\n filters: Integer, the dimensionality of the output space.\n kernel: An integer or tuple/list of 2 integers, specifying the\n width and height of the 2D convolution window.\n strides: An integer or tuple/list of 2 integers,\n specifying the strides of the convolution along the width and height.\n Can be a single integer to specify the same value for\n all spatial dimensions.\n nl: String, nonlinearity activation type.\n\n # Returns\n Output tensor.\n \"\"\"\n\n channel_axis = 1 if K.image_data_format() == 'channels_first' else -1\n\n x = Conv2D(filters, kernel, padding='same', strides=strides)(inputs)\n x = BatchNormalization(axis=channel_axis)(x)\n\n return self._return_activation(x, nl)\n\n def _squeeze(self, inputs):\n \"\"\"Squeeze and Excitation.\n This function defines a squeeze structure.\n\n # Arguments\n inputs: Tensor, input tensor of conv layer.\n \"\"\"\n input_channels = int(inputs.shape[-1])\n\n x = GlobalAveragePooling2D()(inputs)\n x = Dense(input_channels, activation='relu')(x)\n x = Dense(input_channels, activation='hard_sigmoid')(x)\n x = Reshape((1, 1, input_channels))(x)\n x = Multiply()([inputs, x])\n\n return x\n\n def _bottleneck(self, inputs, filters, kernel, e, s, squeeze, nl):\n \"\"\"Bottleneck\n This function defines a basic bottleneck structure.\n\n # Arguments\n inputs: Tensor, input tensor of conv layer.\n filters: Integer, the dimensionality of the output space.\n kernel: An integer or tuple/list of 2 integers, specifying the\n width and height of the 2D convolution window.\n e: Integer, expansion factor.\n t is always applied to the input size.\n s: An integer or tuple/list of 2 integers,specifying the strides\n of the convolution along the width and height.Can be a single\n integer to specify the same value for all spatial dimensions.\n squeeze: Boolean, Whether to use the squeeze.\n nl: String, nonlinearity activation type.\n\n # Returns\n Output tensor.\n \"\"\"\n\n channel_axis = 1 if K.image_data_format() == 'channels_first' else -1\n input_shape = K.int_shape(inputs)\n\n tchannel = int(e)\n cchannel = int(self.alpha * filters)\n\n r = s == 1 and input_shape[3] == filters\n\n x = self._conv_block(inputs, tchannel, (1, 1), (1, 1), nl)\n\n x = DepthwiseConv2D(kernel, strides=(s, s), depth_multiplier=1, padding='same')(x)\n x = BatchNormalization(axis=channel_axis)(x)\n x = self._return_activation(x, nl)\n\n if squeeze:\n x = self._squeeze(x)\n\n x = Conv2D(cchannel, (1, 1), strides=(1, 1), padding='same')(x)\n x = BatchNormalization(axis=channel_axis)(x)\n\n if r:\n x = Add()([x, inputs])\n\n return x\n\n def build(self):\n pass\n" }, { "path": "model/mobilenet_v3_large.py", "content": "\"\"\"MobileNet v3 Large models for Keras.\n# Reference\n [Searching for MobileNetV3](https://arxiv.org/abs/1905.02244?context=cs)\n\"\"\"\n\n\nfrom keras.models import Model\nfrom keras.layers import Input, Conv2D, GlobalAveragePooling2D, Reshape\nfrom keras.utils.vis_utils import plot_model\n\nfrom model.mobilenet_base import MobileNetBase\n\n\nclass MobileNetV3_Large(MobileNetBase):\n def __init__(self, shape, n_class, alpha=1.0, include_top=True):\n \"\"\"Init.\n\n # Arguments\n input_shape: An integer or tuple/list of 3 integers, shape\n of input tensor.\n n_class: Integer, number of classes.\n alpha: Integer, width multiplier.\n include_top: if inculde classification layer.\n\n # Returns\n MobileNetv3 model.\n \"\"\"\n super(MobileNetV3_Large, self).__init__(shape, n_class, alpha)\n self.include_top = include_top\n\n def build(self, plot=False):\n \"\"\"build MobileNetV3 Large.\n\n # Arguments\n plot: Boolean, weather to plot model.\n\n # Returns\n model: Model, model.\n \"\"\"\n inputs = Input(shape=self.shape)\n\n x = self._conv_block(inputs, 16, (3, 3), strides=(2, 2), nl='HS')\n\n x = self._bottleneck(x, 16, (3, 3), e=16, s=1, squeeze=False, nl='RE')\n x = self._bottleneck(x, 24, (3, 3), e=64, s=2, squeeze=False, nl='RE')\n x = self._bottleneck(x, 24, (3, 3), e=72, s=1, squeeze=False, nl='RE')\n x = self._bottleneck(x, 40, (5, 5), e=72, s=2, squeeze=True, nl='RE')\n x = self._bottleneck(x, 40, (5, 5), e=120, s=1, squeeze=True, nl='RE')\n x = self._bottleneck(x, 40, (5, 5), e=120, s=1, squeeze=True, nl='RE')\n x = self._bottleneck(x, 80, (3, 3), e=240, s=2, squeeze=False, nl='HS')\n x = self._bottleneck(x, 80, (3, 3), e=200, s=1, squeeze=False, nl='HS')\n x = self._bottleneck(x, 80, (3, 3), e=184, s=1, squeeze=False, nl='HS')\n x = self._bottleneck(x, 80, (3, 3), e=184, s=1, squeeze=False, nl='HS')\n x = self._bottleneck(x, 112, (3, 3), e=480, s=1, squeeze=True, nl='HS')\n x = self._bottleneck(x, 112, (3, 3), e=672, s=1, squeeze=True, nl='HS')\n x = self._bottleneck(x, 160, (5, 5), e=672, s=2, squeeze=True, nl='HS')\n x = self._bottleneck(x, 160, (5, 5), e=960, s=1, squeeze=True, nl='HS')\n x = self._bottleneck(x, 160, (5, 5), e=960, s=1, squeeze=True, nl='HS')\n\n x = self._conv_block(x, 960, (1, 1), strides=(1, 1), nl='HS')\n x = GlobalAveragePooling2D()(x)\n x = Reshape((1, 1, 960))(x)\n\n x = Conv2D(1280, (1, 1), padding='same')(x)\n x = self._return_activation(x, 'HS')\n\n if self.include_top:\n x = Conv2D(self.n_class, (1, 1), padding='same', activation='softmax')(x)\n x = Reshape((self.n_class,))(x)\n\n model = Model(inputs, x)\n\n if plot:\n plot_model(model, to_file='images/MobileNetv3_large.png', show_shapes=True)\n\n return model\n" }, { "path": "model/mobilenet_v3_small.py", "content": "\"\"\"MobileNet v3 small models for Keras.\n# Reference\n [Searching for MobileNetV3](https://arxiv.org/abs/1905.02244?context=cs)\n\"\"\"\n\n\nfrom keras.models import Model\nfrom keras.layers import Input, Conv2D, GlobalAveragePooling2D, Reshape\nfrom keras.utils.vis_utils import plot_model\n\nfrom model.mobilenet_base import MobileNetBase\n\n\nclass MobileNetV3_Small(MobileNetBase):\n def __init__(self, shape, n_class, alpha=1.0, include_top=True):\n \"\"\"Init.\n\n # Arguments\n input_shape: An integer or tuple/list of 3 integers, shape\n of input tensor.\n n_class: Integer, number of classes.\n alpha: Integer, width multiplier.\n include_top: if inculde classification layer.\n\n # Returns\n MobileNetv3 model.\n \"\"\"\n super(MobileNetV3_Small, self).__init__(shape, n_class, alpha)\n self.include_top = include_top\n\n def build(self, plot=False):\n \"\"\"build MobileNetV3 Small.\n\n # Arguments\n plot: Boolean, weather to plot model.\n\n # Returns\n model: Model, model.\n \"\"\"\n inputs = Input(shape=self.shape)\n\n x = self._conv_block(inputs, 16, (3, 3), strides=(2, 2), nl='HS')\n\n x = self._bottleneck(x, 16, (3, 3), e=16, s=2, squeeze=True, nl='RE')\n x = self._bottleneck(x, 24, (3, 3), e=72, s=2, squeeze=False, nl='RE')\n x = self._bottleneck(x, 24, (3, 3), e=88, s=1, squeeze=False, nl='RE')\n x = self._bottleneck(x, 40, (5, 5), e=96, s=2, squeeze=True, nl='HS')\n x = self._bottleneck(x, 40, (5, 5), e=240, s=1, squeeze=True, nl='HS')\n x = self._bottleneck(x, 40, (5, 5), e=240, s=1, squeeze=True, nl='HS')\n x = self._bottleneck(x, 48, (5, 5), e=120, s=1, squeeze=True, nl='HS')\n x = self._bottleneck(x, 48, (5, 5), e=144, s=1, squeeze=True, nl='HS')\n x = self._bottleneck(x, 96, (5, 5), e=288, s=2, squeeze=True, nl='HS')\n x = self._bottleneck(x, 96, (5, 5), e=576, s=1, squeeze=True, nl='HS')\n x = self._bottleneck(x, 96, (5, 5), e=576, s=1, squeeze=True, nl='HS')\n\n x = self._conv_block(x, 576, (1, 1), strides=(1, 1), nl='HS')\n x = GlobalAveragePooling2D()(x)\n x = Reshape((1, 1, 576))(x)\n\n x = Conv2D(1280, (1, 1), padding='same')(x)\n x = self._return_activation(x, 'HS')\n\n if self.include_top:\n x = Conv2D(self.n_class, (1, 1), padding='same', activation='softmax')(x)\n x = Reshape((self.n_class,))(x)\n\n model = Model(inputs, x)\n\n if plot:\n plot_model(model, to_file='images/MobileNetv3_small.png', show_shapes=True)\n\n return model\n" }, { "path": "train_cls.py", "content": "import os\nimport json\nimport pandas as pd\n\nfrom keras.optimizers import Adam\nfrom keras.preprocessing.image import ImageDataGenerator\nfrom keras.callbacks import EarlyStopping, ModelCheckpoint\n\ndef generate(batch, shape, ptrain, pval):\n \"\"\"Data generation and augmentation\n\n # Arguments\n batch: Integer, batch size.\n size: Integer, image size.\n ptrain: train dir.\n pval: eval dir.\n\n # Returns\n train_generator: train set generator\n validation_generator: validation set generator\n count1: Integer, number of train set.\n count2: Integer, number of test set.\n \"\"\"\n\n # Using the data Augmentation in traning data\n datagen1 = ImageDataGenerator(\n rescale=1. / 255,\n shear_range=0.2,\n zoom_range=0.2,\n rotation_range=90,\n width_shift_range=0.2,\n height_shift_range=0.2,\n horizontal_flip=True)\n\n datagen2 = ImageDataGenerator(rescale=1. / 255)\n\n train_generator = datagen1.flow_from_directory(\n ptrain,\n target_size=shape,\n batch_size=batch,\n class_mode='categorical')\n\n validation_generator = datagen2.flow_from_directory(\n pval,\n target_size=shape,\n batch_size=batch,\n class_mode='categorical')\n\n count1 = 0\n for root, dirs, files in os.walk(ptrain):\n for each in files:\n count1 += 1\n\n count2 = 0\n for root, dirs, files in os.walk(pval):\n for each in files:\n count2 += 1\n\n return train_generator, validation_generator, count1, count2\n\n\ndef train():\n with open('config/config.json', 'r') as f:\n cfg = json.load(f)\n\n save_dir = cfg['save_dir']\n shape = (int(cfg['height']), int(cfg['width']), 3)\n n_class = int(cfg['class_number'])\n batch = int(cfg['batch'])\n\n if not os.path.exists(save_dir):\n os.mkdir(save_dir)\n\n if cfg['model'] == 'large':\n from model.mobilenet_v3_large import MobileNetV3_Large\n model = MobileNetV3_Large(shape, n_class).build()\n if cfg['model'] == 'small':\n from model.mobilenet_v3_small import MobileNetV3_Small\n model = MobileNetV3_Small(shape, n_class).build()\n\n pre_weights = cfg['weights']\n if pre_weights and os.path.exists(pre_weights):\n model.load_weights(pre_weights, by_name=True)\n\n opt = Adam(lr=float(cfg['learning_rate']))\n earlystop = EarlyStopping(monitor='val_acc', patience=5, verbose=0, mode='auto')\n checkpoint = ModelCheckpoint(filepath=os.path.join(save_dir, '{}_weights.h5'.format(cfg['model'])),\n monitor='val_acc', save_best_only=True, save_weights_only=True)\n \n model.compile(loss='categorical_crossentropy', optimizer=opt, metrics=['accuracy'])\n\n train_generator, validation_generator, count1, count2 = generate(batch, shape[:2], cfg['train_dir'], cfg['eval_dir'])\n\n hist = model.fit_generator(\n train_generator,\n validation_data=validation_generator,\n steps_per_epoch=count1 // batch,\n validation_steps=count2 // batch,\n epochs=cfg['epochs'],\n callbacks=[earlystop,checkpoint])\n\n df = pd.DataFrame.from_dict(hist.history)\n df.to_csv(os.path.join(save_dir, 'hist.csv'), encoding='utf-8', index=False)\n #model.save_weights(os.path.join(save_dir, '{}_weights.h5'.format(cfg['model'])))\n\n\nif __name__ == '__main__':\n train()\n" } ]