diff --git a/train.py b/train.py index 7b5a66d..3e21526 100755 --- a/train.py +++ b/train.py @@ -4,7 +4,7 @@ import tensorflow as tf import datetime import zipfile -from yolo3.model import darknet_yolo_body, YoloLoss, mobilenetv2_yolo_body, efficientnet_yolo_body +from yolo3.model import darknet_yolo_body, YoloLoss, mobilenetv2_yolo_body, efficientnet_yolo_body, mobilenetv2_yolo_body_r13 from yolo3.data import Dataset from yolo3.enum import OPT, BACKBONE, DATASET_MODE from yolo3.map import MAPCallback @@ -115,6 +115,12 @@ def train(FLAGS): len(anchors) // 3, num_classes, alpha=FLAGS['alpha']) + elif backbone == BACKBONE.MOBILENETV2_R13: + model = factory.build(mobilenetv2_yolo_body_r13, + 155, + len(anchors) // 3, + num_classes, + alpha=FLAGS['alpha']) elif backbone == BACKBONE.DARKNET53: model = factory.build(darknet_yolo_body, 185, len(anchors) // 3, num_classes) diff --git a/train_backbone.py b/train_backbone.py index 73fba15..d2b5cef 100755 --- a/train_backbone.py +++ b/train_backbone.py @@ -76,7 +76,7 @@ def train(FLAGS): batch_size = batch_size * strategy.num_replicas_in_sync with strategy.scope(): factory = ModelFactory(weights_path=model_path) - if backbone == BACKBONE.MOBILENETV2: + if backbone == BACKBONE.MOBILENETV2 or backbone == BACKBONE.MOBILENETV2_R13: model = factory.build(mobilenetv2, 0, alpha=1.4, diff --git a/yolo.py b/yolo.py index a658aea..2e7bce5 100755 --- a/yolo.py +++ b/yolo.py @@ -9,7 +9,7 @@ from PIL import Image, ImageFont, ImageDraw import cv2 import tensorflow as tf -from yolo3.model import yolo_eval, darknet_yolo_body, mobilenetv2_yolo_body, efficientnet_yolo_body, YoloEval +from yolo3.model import yolo_eval, darknet_yolo_body, mobilenetv2_yolo_body, efficientnet_yolo_body, YoloEval, mobilenetv2_yolo_body_r13 from yolo3.utils import letterbox_image, get_anchors, get_classes from yolo3.enum import OPT, BACKBONE from yolo3.map import MAPCallback @@ -79,6 +79,9 @@ def generate(self, FLAGS): if self.backbone == BACKBONE.MOBILENETV2: model_body = partial(mobilenetv2_yolo_body, alpha=FLAGS['alpha']) + elif self.backbone == BACKBONE.MOBILENETV2_R13: + model_body = partial(mobilenetv2_yolo_body_r13, + alpha=FLAGS['alpha']) elif self.backbone == BACKBONE.DARKNET53: model_body = darknet_yolo_body elif self.backbone == BACKBONE.EFFICIENTNET: diff --git a/yolo3/enum.py b/yolo3/enum.py index 527afc5..ef2b391 100755 --- a/yolo3/enum.py +++ b/yolo3/enum.py @@ -26,6 +26,7 @@ class BACKBONE(Enum): MOBILENETV2 = 0 EFFICIENTNET = 1 DARKNET53 = 2 + MOBILENETV2_R13 = 3 @unique diff --git a/yolo3/model.py b/yolo3/model.py index bd97c9f..b1e501d 100755 --- a/yolo3/model.py +++ b/yolo3/model.py @@ -195,6 +195,48 @@ def mobilenetv2_yolo_body(inputs, num_anchors, num_classes, alpha=1.0): return tf.keras.models.Model(inputs, [y1, y2, y3]) +# the mobkenetv2-yolo model befroe r1.14 +def mobilenetv2_yolo_body_r13(inputs, num_anchors, num_classes, alpha=1.0): + MobilenetConv2D_BN_Relu = MobilenetConv2D + mobilenetv2 = tf.keras.applications.MobileNetV2(alpha=alpha, input_tensor=inputs, include_top=False, + weights='imagenet') + x=mobilenetv2.output + y1 = MobilenetConv2D_BN_Relu((1,1),alpha, 1280)(x) + y1 = DarknetConv2D(num_anchors * (num_classes + 5), (1, 1),padding='same',kernel_regularizer=tf.keras.regularizers.l2(5e-4))(y1) + x = compose( + MobilenetConv2D_BN_Relu((1,1),alpha, 640), + tf.keras.layers.UpSampling2D(2))(x) + x = tf.keras.layers.Concatenate()( + [x, MobilenetConv2D_BN_Relu((1,1),alpha, 640)(mobilenetv2.get_layer('block_12_project_BN').output)]) + y2 = MobilenetConv2D_BN_Relu((1,1),alpha, 640)(x) + y2 = DarknetConv2D(num_anchors * (num_classes + 5), (1, 1),padding='same',kernel_regularizer=tf.keras.regularizers.l2(5e-4))(y2) + x = compose( + MobilenetConv2D_BN_Relu((1,1),alpha, 320), + tf.keras.layers.UpSampling2D(2))(x) + x = tf.keras.layers.Concatenate()( + [x, MobilenetConv2D_BN_Relu((1,1),alpha, 320)(mobilenetv2.get_layer('block_5_project_BN').output)]) + y3 = MobilenetConv2D_BN_Relu((1,1),alpha, 320)(x) + y3 = DarknetConv2D(num_anchors * (num_classes + 5), (1, 1),padding='same',kernel_regularizer=tf.keras.regularizers.l2(5e-4))(y3) + + # reshape to make compatible with r1.14 + y1 = tf.keras.layers.Lambda(lambda y: tf.reshape(y, [ + -1, tf.shape(y)[1], + tf.shape(y)[2], num_anchors, num_classes + 5 + ]), + name='y1')(y1) + y2 = tf.keras.layers.Lambda(lambda y: tf.reshape(y, [ + -1, tf.shape(y)[1], + tf.shape(y)[2], num_anchors, num_classes + 5 + ]), + name='y2')(y2) + y3 = tf.keras.layers.Lambda(lambda y: tf.reshape(y, [ + -1, tf.shape(y)[1], + tf.shape(y)[2], num_anchors, num_classes + 5 + ]), + name='y3')(y3) + return tf.keras.models.Model(inputs, [y1, y2, y3]) + + def make_last_layers_efficientnet(x, block_args, global_params): if global_params.data_format == 'channels_first': channel_axis = 1