\u6d4b\u8bd5\u8fc7\u7a0b\u6d41\u7a0b\u53ca\u5b9e\u73b0\uff1a<\/p>\n\n\n\n
- \u89e3\u6790\u811a\u672c\u8f93\u5165\u53c2\u6570\uff1a\u4f7f\u7528argparse\u89e3\u6790\uff0c\u7531args\u53d8\u91cf\u6301\u6709<\/li>
- \u521b\u5efa\u6a21\u578b\uff1a\u81ea\u5b9a\u4e49\u51fd\u6570create_model()<\/li>
- \u6a21\u578b\u52a0\u8f7d\uff1a\u4f7f\u7528keras.models.load_model()<\/li>
- \u6570\u636e\u8bfb\u53d6\uff1a\u81ea\u5b9a\u4e49\u51fd\u6570create_image_generator()<\/li>
- \u9884\u6d4b\u4e0e\u8bc4\u4f30\uff1a\u81ea\u5b9a\u4e49\u51fd\u6570test()<\/li><\/ul>\n\n\n\n
\u73af\u5883\u642d\u5efa<\/h2>\n\n\n\n
\u5b89\u88c5TensorFlow<\/h4>\n\n\n\n
\u8f93\u5165\u4e0b\u8ff0\u547d\u4ee4\u5347\u7ea7pip\u5e76\u5b89\u88c5TensorFlow<\/p>\n\n\n\n
python -m pip install --upgrade pip && pip install tensorflow==1.14<\/pre>\n\n\n\n
\u5b89\u88c5Keras<\/h4>\n\n\n\n
\u8f93\u5165\u4e0b\u8ff0\u547d\u4ee4\u5b89\u88c5Keras<\/p>\n\n\n\n
pip install keras==2.3.1<\/pre>\n\n\n\n
\u5b89\u88c5opencv-python<\/h4>\n\n\n\n
\u8f93\u5165\u4e0b\u8ff0\u547d\u4ee4\u5b89\u88c5opencv-python<\/p>\n\n\n\n
pip install opencv-python<\/pre>\n\n\n\n
\u5b89\u88c5numpy<\/h4>\n\n\n\n
\u8f93\u5165\u4e0b\u8ff0\u547d\u4ee4\u5b89\u88c5numpy<\/p>\n\n\n\n
pip install numpy==1.19<\/pre>\n\n\n\n
\u7f16\u5199\u8bad\u7ec3\u4ee3\u7801<\/h2>\n\n\n\n
\u521b\u5efa\u6587\u4ef6<\/h4>\n\n\n\n
\u8fdb\u5165\u5de5\u7a0b\u76ee\u5f55<\/p>\n\n\n\n
cd \/traffic_symbol<\/pre>\n\n\n\n
\u521b\u5efatrain.py\u6587\u4ef6\uff0c\u672c\u5b9e\u9a8c\u7684\u540e\u7eed\u4ee3\u7801\u90fd\u5c06\u5728\u6b64\u6587\u4ef6\u4e2d\u5b8c\u6210<\/p>\n\n\n\n
touch train.py<\/pre>\n\n\n\n
\u5f15\u7528\u6587\u4ef6<\/h4>\n\n\n\n
\u70b9\u51fb\u6253\u5f00 train.py<\/a> \u6587\u4ef6\uff0c\u8f93\u5165\u4e0b\u8ff0\u5185\u5bb9\uff1a<\/p>\n\n\n\n
\u5728\u6587\u4ef6\u9876\u90e8\u8f93\u5165\u4e0b\u8ff0\u5185\u5bb9<\/p>\n\n\n\n
import os
import argparse
import shutil
import cv2
import random
import numpy as np
import keras
from keras.applications.mobilenet import preprocess_input, MobileNet
from keras.preprocessing.image import ImageDataGenerator
from keras.models import Model, load_model
from keras.layers import *
from keras.callbacks import ModelCheckpoint<\/pre>\n\n\n\n\u8bb0\u5f97\u4fdd\u5b58\uff01<\/strong><\/p>\n\n\n\n
\u4fdd\u5b58\u65b9\u6cd5\uff1aWindows<\/strong> \u7cfb\u7edf\u70b9\u51fb ctrl+s<\/strong>\uff0cMac OS<\/strong> \u70b9\u51fb command+s<\/strong> \u4fdd\u5b58<\/p>\n\n\n\n
\u5b8c\u6210\u6a21\u578b\u6784\u5efa\u4ee3\u7801<\/h4>\n\n\n\n
\u7ee7\u7eed\u5728 train.py \u4e2d\u7ee7\u7eed\u8f93\u5165\u4e0b\u9762\u7684\u5185\u5bb9\uff0c\u7136\u540e\u4fdd\u5b58<\/p>\n\n\n\n
# \u521b\u5efa\u6a21\u578b
def create_model(height, width, channel, num_class):
# \u52a0\u8f7d\u9884\u8bad\u7ec3\u6a21\u578b
base_model = MobileNet(input_shape=(height, width, channel), weights='imagenet', include_top=False)
\u200b
# \u628a\u57fa\u7840\u6a21\u578b\u540e\u90e8\u66ff\u6362\u6210GAP + FC
x = base_model.output
x = GlobalAveragePooling2D()(x)
x = Dropout(0.5)(x)
x = Dense(256, activation='relu')(x)
predictions = Dense(num_class, activation='softmax')(x)
\u200b
# \u521b\u5efa\u6a21\u578b
return Model(inputs=base_model.input, outputs=predictions)<\/pre>\n\n\n\n\u5b8c\u6210\u6570\u636e\u5904\u7406\u4ee3\u7801<\/h4>\n\n\n\n
\u7ee7\u7eed\u6dfb\u52a0\u4e0b\u5217\u4ee3\u7801\uff0c\u5e76\u4fdd\u5b58\uff1a<\/p>\n\n\n\n
# \u56fe\u7247\u589e\u5f3a
def preprocess(image):
# \u56fe\u7247\u683c\u5f0f\u8f6c\u6362\u4e3aHSV
image = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)
\u200b
# \u65e0\u5173\u4fe1\u606f\u589e\u5f3a\uff1a\u968f\u673a\u6539\u53d8\u8272\u8c03\u3001\u9971\u548c\u5ea6\u3001\u660e\u5ea6
h, s, v = cv2.split(image)
shift_h = cv2.addWeighted(h, 1, h, 0, random.randint(-10, 10))
shift_hsv = cv2.merge([shift_h, s, v])
image = cv2.cvtColor(shift_hsv, cv2.COLOR_HSV2RGB)
brightness = random.randint(-50, 50)
contrast = random.uniform(0.8, 1.2)
image = cv2.addWeighted(image, contrast, image, 0, brightness)
\u200b
# mobilenet\u7684\u56fe\u7247\u52a0\u901f\u5904\u7406\u65b9\u5f0f
image = preprocess_input(image)
return image
\u200b
\u200b
# \u521b\u5efaImageDataGenerator
def create_image_generator(args, type):
classes = [str(i) for i in range(args.n_classes)]
if type == 0:
# \u8bad\u7ec3\u96c6\u6570\u636e\u5904\u7406
generator = ImageDataGenerator(
preprocessing_function=preprocess,
shear_range=0.2, # \u88c1\u526a
zoom_range=0.2, # \u7f29\u653e
rotation_range=20, # \u65cb\u8f6c
vertical_flip=False,# \u7eb5\u5411\u5bf9\u79f0\u53d8\u6362
horizontal_flip=True# \u6a2a\u5411\u5bf9\u79f0\u53d8\u6362
).flow_from_directory(
args.dataset + '\/train',
target_size=(model.inputs[0].shape[1], model.inputs[0].shape[2]),
batch_size=args.batch_size,
classes=classes,
class_mode='categorical'
)
elif type == 1:
# \u9a8c\u8bc1\u96c6\u96c6\u6570\u636e\u5904\u7406
generator = ImageDataGenerator(
preprocessing_function=preprocess_input # \u6570\u636e\u4e0d\u505a\u4efb\u4f55\u589e\u5f3a
).flow_from_directory(
args.dataset + '\/validation',
target_size=(model.inputs[0].shape[1], model.inputs[0].shape[2]),
batch_size=args.batch_size,
classes=classes,
class_mode='categorical',
shuffle=False,
seed=0
)
return generator<\/pre>\n\n\n\n\u5b8c\u6210\u8bad\u7ec3\u6d41\u7a0b\u4ee3\u7801<\/h4>\n\n\n\n
\u7ee7\u7eed\u6dfb\u52a0\u4e0b\u5217\u4ee3\u7801\uff0c\u5e76\u4fdd\u5b58\uff1a<\/p>\n\n\n\n
# \u8bad\u7ec3\u6a21\u578b
def train(args, model):
# \u8bad\u7ec3\u6570\u636e\u5904\u7406
train_generator = create_image_generator(args, 0)
\u200b
# \u9a8c\u8bc1\u96c6\u6570\u636e\u5904\u7406
validation_generator = create_image_generator(args, 1)
\u200b
# \u7f16\u8bd1\u6a21\u578b
model.compile(optimizer=keras.optimizers.Adam(),
metrics=['accuracy'],
loss='categorical_crossentropy')
\u200b
# \u6a21\u578b\u4fdd\u5b58
model_save_path = os.path.join(args.output, 'model.h5')
saver = ModelCheckpoint(model_save_path, monitor='val_loss', verbose=1, save_best_only=True)
\u200b
# \u5f00\u59cb\u8bad\u7ec3
model.fit_generator(
generator=train_generator,
epochs=args.epochs,
validation_data=validation_generator,
callbacks=[saver])<\/pre>\n\n\n\n\u5b8c\u6210\u53c2\u6570\u5904\u7406\u4ee3\u7801<\/h4>\n\n\n\n
\u7ee7\u7eed\u6dfb\u52a0\u4e0b\u5217\u4ee3\u7801\uff0c\u5e76\u4fdd\u5b58\uff1a<\/p>\n\n\n\n
if __name__ == \"__main__\":\n \u00a0 \u00a0# \u53c2\u6570\u89e3\u6790\n \u00a0 \u00a0parser = argparse.ArgumentParser()\n \u00a0 \u00a0# \u5b9a\u4e49\u8def\u5f84\n \u00a0 \u00a0parser.add_argument(\"--dataset\", type=str, default=\".\/data\")\n \u00a0 \u00a0parser.add_argument(\"--output\", type=str, default=\".\/results\")\n \u00a0 \u00a0parser.add_argument(\"--resume_model_path\", type=str, default=\"\")\n \u00a0 \u00a0# \u5b9a\u4e49\u6a21\u578b\u53c2\u6570\n \u00a0 \u00a0parser.add_argument(\"--n_classes\", type=int, default=2)\n \u00a0 \u00a0parser.add_argument(\"--input_width\", type=int, default=128)\n \u00a0 \u00a0parser.add_argument(\"--input_height\", type=int, default=128)\n \u00a0 \u00a0parser.add_argument(\"--input_channel\", type=int, default=3)\n \u00a0 \u00a0# \u5b9a\u4e49\u8d85\u53c2\u6570\n \u00a0 \u00a0parser.add_argument(\"--epochs\", type=int, default=5)\n \u00a0 \u00a0parser.add_argument(\"--batch_size\", type=int, default=4)\n \u00a0 \u00a0# \u8fc7\u7a0b\u63a7\u5236\u3002test\u4e3a0\u8868\u793a\u8bad\u7ec3\uff0ctest\u4e3a1\u8868\u793a\u6d4b\u8bd5\n \u00a0 \u00a0parser.add_argument(\"--test\", type=int, default=0)\n \u00a0 \u00a0args = parser.parse_args()\n \u00a0 \u00a0print(\"args: \", args)\n\u200b\n \u00a0 \u00a0# \u521b\u5efa\u8f93\u51fa\u8def\u5f84\u6240\u6307\u6587\u4ef6\u5939\n \u00a0 \u00a0os.makedirs(args.output, exist_ok=True)\n\u200b\n \u00a0 \u00a0# \u521b\u5efa\u6a21\u578b\n \u00a0 \u00a0model = create_model(args.input_height, args.input_width, args.input_channel, args.n_classes)\n\u200b\n \u00a0 \u00a0# \u6253\u5370\u6a21\u578b\u7ed3\u6784\n \u00a0 \u00a0model.summary()\n\u200b\n \u00a0 \u00a0# \u52a0\u8f7d\u6a21\u578b\n \u00a0 \u00a0if args.resume_model_path != \"\":\n \u00a0 \u00a0 \u00a0 \u00a0try:\n \u00a0 \u00a0 \u00a0 \u00a0 \u00a0 \u00a0model = load_model(args.resume_model_path)\n \u00a0 \u00a0 \u00a0 \u00a0except Exception as e:\n \u00a0 \u00a0 \u00a0 \u00a0 \u00a0 \u00a0print('No saved model, using init weights!')\n\u200b\n \u00a0 \u00a0if args.test:\n \u00a0 \u00a0 \u00a0 \u00a0# \u9884\u6d4b\u6d4b\u8bd5\u56fe\u7247\n \u00a0 \u00a0 \u00a0 \u00a0test(args, model)\n \u00a0 \u00a0else:\n \u00a0 \u00a0 \u00a0 \u00a0# \u5f00\u59cb\u8bad\u7ec3\n \u00a0 \u00a0 \u00a0 \u00a0train(args, model)<\/pre>\n\n\n\n\u5f00\u59cb\u8bad\u7ec3<\/h2>\n\n\n\n
\u4f7f\u7528\u811a\u672c\u8fdb\u884c\u8bad\u7ec3<\/p>\n\n\n\n
\u8f93\u5165\u4e0b\u8ff0\u547d\u4ee4\u6267\u884c\u811a\u672c\u8bad\u7ec3\u8fc7\u7a0b<\/p>\n\n\n\n
python train.py<\/pre>\n\n\n\n
\u6d4b\u8bd5\u6a21\u578b<\/h2>\n\n\n\n
\u8f93\u5165\u6d4b\u8bd5\u6570\u636e\u5904\u7406\u51fd\u6570<\/h4>\n\n\n\n
![\"image\"\/](\"https:\/\/academy-lab-prd-pub-1258344699.cos.ap-guangzhou.myqcloud.com\/lab-prd\/contribution-attachment\/100021000061\/1635424577166-%E5%9B%BE%E7%89%87%201.jpg\")
<\/figure>\n\n\n\n