Add files via upload

FPS_test.py

-import os
 import time
 
 import numpy as np

get_dr_txt.py

@@ -8,8 +8,8 @@ import os
 
 import numpy as np
 from keras import backend as K
-from keras.applications.imagenet_utils import preprocess_input
 from keras.layers import Input
+from keras.models import load_model
 from PIL import Image
 from tqdm import tqdm
 
@@ -17,7 +17,6 @@ from nets.yolo4 import yolo_body, yolo_eval
 from utils.utils import letterbox_image
 from yolo import YOLO
 
-
 '''
 这里设置的门限值较低是因为计算map需要用到不同门限条件下的Recall和Precision值。
 所以只有保留的框足够多，计算的map才会更精确，详情可以了解map的原理。

nets/CSPdarknet53.py

 from functools import wraps
 
 from keras import backend as K
-from keras.layers import (Add, Concatenate, Conv2D, Layer, MaxPooling2D,
-                          UpSampling2D, ZeroPadding2D)
-from keras.layers.advanced_activations import LeakyReLU
+from keras.initializers import random_normal
+from keras.layers import Add, Concatenate, Conv2D, Layer, ZeroPadding2D
 from keras.layers.normalization import BatchNormalization
 from keras.regularizers import l2
 from utils.utils import compose
@@ -32,7 +31,7 @@ class Mish(Layer):
 @wraps(Conv2D)
 def DarknetConv2D(*args, **kwargs):
     # darknet_conv_kwargs = {'kernel_regularizer': l2(5e-4)}
-    darknet_conv_kwargs = {}
+    darknet_conv_kwargs = {'kernel_initializer' : random_normal(stddev=0.02)}
     darknet_conv_kwargs['padding'] = 'valid' if kwargs.get('strides')==(2,2) else 'same'
     darknet_conv_kwargs.update(kwargs)
     return Conv2D(*args, **darknet_conv_kwargs)

nets/yolo4.py

 from functools import wraps
 
 import keras
-import numpy as np
 import tensorflow as tf
 from keras import backend as K
-from keras.layers import (Add, Concatenate, Conv2D, MaxPooling2D, UpSampling2D,
+from keras.initializers import random_normal
+from keras.layers import (Concatenate, Conv2D, MaxPooling2D, UpSampling2D,
                           ZeroPadding2D)
 from keras.layers.advanced_activations import LeakyReLU
 from keras.layers.normalization import BatchNormalization
@@ -23,7 +23,7 @@ from nets.CSPdarknet53 import darknet_body
 @wraps(Conv2D)
 def DarknetConv2D(*args, **kwargs):
     # darknet_conv_kwargs = {'kernel_regularizer': l2(5e-4)}
-    darknet_conv_kwargs = {}
+    darknet_conv_kwargs = {'kernel_initializer' : random_normal(stddev=0.02)}
     darknet_conv_kwargs['padding'] = 'valid' if kwargs.get('strides')==(2,2) else 'same'
     darknet_conv_kwargs.update(kwargs)
     return Conv2D(*args, **darknet_conv_kwargs)

train.py

 import keras.backend as K
 import numpy as np
-import tensorflow as tf
-from keras.backend.tensorflow_backend import set_session
 from keras.callbacks import (EarlyStopping, ModelCheckpoint, ReduceLROnPlateau,
                              TensorBoard)
 from keras.layers import Input, Lambda
@@ -10,8 +8,8 @@ from keras.optimizers import Adam
 
 from nets.loss import yolo_loss
 from nets.yolo4 import yolo_body
-from utils.utils import (WarmUpCosineDecayScheduler, get_random_data,
-                         get_random_data_with_Mosaic, rand)
+from utils.utils import (LossHistory, WarmUpCosineDecayScheduler,
+                         get_random_data, get_random_data_with_Mosaic)
 
 
 #---------------------------------------------------#
@@ -216,7 +214,7 @@ if __name__ == "__main__":
     #   获取classes和anchor
     #----------------------------------------------------#
     class_names = get_classes(classes_path)
-    anchors = get_anchors(anchors_path)
+    anchors     = get_anchors(anchors_path)
     #------------------------------------------------------#
     #   一共有多少类和多少先验框
     #------------------------------------------------------#
@@ -272,6 +270,7 @@ if __name__ == "__main__":
     checkpoint = ModelCheckpoint(log_dir + 'ep{epoch:03d}-loss{loss:.3f}-val_loss{val_loss:.3f}.h5',
         monitor='val_loss', save_weights_only=True, save_best_only=False, period=1)
     early_stopping = EarlyStopping(monitor='val_loss', min_delta=0, patience=10, verbose=1)
+    loss_history = LossHistory(log_dir)
 
     #----------------------------------------------------------------------#
     #   验证集的划分在train.py代码里面进行
@@ -300,20 +299,20 @@ if __name__ == "__main__":
     #   提示OOM或者显存不足请调小Batch_size
     #------------------------------------------------------#
     if True:
-        Init_epoch = 0
-        Freeze_epoch = 50
-        batch_size = 8
-        learning_rate_base = 1e-3
+        Init_epoch          = 0
+        Freeze_epoch        = 50
+        batch_size          = 8
+        learning_rate_base  = 1e-3
 
         if Cosine_scheduler:
             # 预热期
-            warmup_epoch = int((Freeze_epoch-Init_epoch)*0.2)
+            warmup_epoch    = int((Freeze_epoch-Init_epoch)*0.2)
             # 总共的步长
-            total_steps = int((Freeze_epoch-Init_epoch) * num_train / batch_size)
+            total_steps     = int((Freeze_epoch-Init_epoch) * num_train / batch_size)
             # 预热步长
-            warmup_steps = int(warmup_epoch * num_train / batch_size)
+            warmup_steps    = int(warmup_epoch * num_train / batch_size)
             # 学习率
-            reduce_lr = WarmUpCosineDecayScheduler(learning_rate_base=learning_rate_base,
+            reduce_lr       = WarmUpCosineDecayScheduler(learning_rate_base=learning_rate_base,
                                                         total_steps=total_steps,
                                                         warmup_learning_rate=1e-4,
                                                         warmup_steps=warmup_steps,
@@ -325,23 +324,29 @@ if __name__ == "__main__":
             reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.5, patience=3, verbose=1)
             model.compile(optimizer=Adam(learning_rate_base), loss={'yolo_loss': lambda y_true, y_pred: y_pred})
 
+        epoch_size          = num_train // batch_size
+        epoch_size_val      = num_val // batch_size
+
+        if epoch_size == 0 or epoch_size_val == 0:
+            raise ValueError("数据集过小，无法进行训练，请扩充数据集。")
+
         print('Train on {} samples, val on {} samples, with batch size {}.'.format(num_train, num_val, batch_size))
         model.fit_generator(data_generator(lines[:num_train], batch_size, input_shape, anchors, num_classes, mosaic=mosaic, random=True),
-                steps_per_epoch=max(1, num_train//batch_size),
+                steps_per_epoch=epoch_size,
                 validation_data=data_generator(lines[num_train:], batch_size, input_shape, anchors, num_classes, mosaic=False, random=False),
-                validation_steps=max(1, num_val//batch_size),
+                validation_steps=epoch_size_val,
                 epochs=Freeze_epoch,
                 initial_epoch=Init_epoch,
-                callbacks=[logging, checkpoint, reduce_lr, early_stopping])
+                callbacks=[logging, checkpoint, reduce_lr, early_stopping, loss_history])
         model.save_weights(log_dir + 'trained_weights_stage_1.h5')
 
     for i in range(freeze_layers): model_body.layers[i].trainable = True
 
     if True:
-        Freeze_epoch = 50
-        Epoch = 100
-        batch_size = 2
-        learning_rate_base = 1e-4
+        Freeze_epoch        = 50
+        Epoch               = 100
+        batch_size          = 2
+        learning_rate_base  = 1e-4
 
         if Cosine_scheduler:
             # 预热期
@@ -363,12 +368,18 @@ if __name__ == "__main__":
             reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.5, patience=3, verbose=1)
             model.compile(optimizer=Adam(learning_rate_base), loss={'yolo_loss': lambda y_true, y_pred: y_pred})
 
+        epoch_size          = num_train // batch_size
+        epoch_size_val      = num_val // batch_size
+
+        if epoch_size == 0 or epoch_size_val == 0:
+            raise ValueError("数据集过小，无法进行训练，请扩充数据集。")
+
         print('Train on {} samples, val on {} samples, with batch size {}.'.format(num_train, num_val, batch_size))
         model.fit_generator(data_generator(lines[:num_train], batch_size, input_shape, anchors, num_classes, mosaic=mosaic, random=True),
-                steps_per_epoch=max(1, num_train//batch_size),
+                steps_per_epoch=epoch_size,
                 validation_data=data_generator(lines[num_train:], batch_size, input_shape, anchors, num_classes, mosaic=False, random=False),
-                validation_steps=max(1, num_val//batch_size),
+                validation_steps=epoch_size_val,
                 epochs=Epoch,
                 initial_epoch=Freeze_epoch,
-                callbacks=[logging, checkpoint, reduce_lr, early_stopping])
+                callbacks=[logging, checkpoint, reduce_lr, early_stopping, loss_history])
         model.save_weights(log_dir + 'last1.h5')

utils/utils.py

+import os
 from functools import reduce
 
 import cv2
 import keras
 import keras.backend as K
+import matplotlib.pyplot as plt
 import numpy as np
-from matplotlib.colors import hsv_to_rgb, rgb_to_hsv
+import scipy.signal
 from PIL import Image
 
 
@@ -298,6 +300,58 @@ def get_random_data(annotation_line, input_shape, max_boxes=100, jitter=.3, hue=
     return image_data, box_data
 
 
+class LossHistory(keras.callbacks.Callback):
+    def __init__(self, log_dir):
+        import datetime
+        curr_time = datetime.datetime.now()
+        time_str = datetime.datetime.strftime(curr_time,'%Y_%m_%d_%H_%M_%S')
+        self.log_dir    = log_dir
+        self.time_str   = time_str
+        self.save_path  = os.path.join(self.log_dir, "loss_" + str(self.time_str))  
+        self.losses     = []
+        self.val_loss   = []
+        
+        os.makedirs(self.save_path)
+
+    def on_epoch_end(self, batch, logs={}):
+        self.losses.append(logs.get('loss'))
+        self.val_loss.append(logs.get('val_loss'))
+        with open(os.path.join(self.save_path, "epoch_loss_" + str(self.time_str) + ".txt"), 'a') as f:
+            f.write(str(logs.get('loss')))
+            f.write("\n")
+        with open(os.path.join(self.save_path, "epoch_val_loss_" + str(self.time_str) + ".txt"), 'a') as f:
+            f.write(str(logs.get('val_loss')))
+            f.write("\n")
+        self.loss_plot()
+
+    def loss_plot(self):
+        iters = range(len(self.losses))
+
+        plt.figure()
+        plt.plot(iters, self.losses, 'red', linewidth = 2, label='train loss')
+        plt.plot(iters, self.val_loss, 'coral', linewidth = 2, label='val loss')
+        try:
+            if len(self.losses) < 25:
+                num = 5
+            else:
+                num = 15
+            
+            plt.plot(iters, scipy.signal.savgol_filter(self.losses, num, 3), 'green', linestyle = '--', linewidth = 2, label='smooth train loss')
+            plt.plot(iters, scipy.signal.savgol_filter(self.val_loss, num, 3), '#8B4513', linestyle = '--', linewidth = 2, label='smooth val loss')
+        except:
+            pass
+
+        plt.grid(True)
+        plt.xlabel('Epoch')
+        plt.ylabel('Loss')
+        plt.title('A Loss Curve')
+        plt.legend(loc="upper right")
+
+        plt.savefig(os.path.join(self.save_path, "epoch_loss_" + str(self.time_str) + ".png"))
+
+        plt.cla()
+        plt.close("all")
+
 def cosine_decay_with_warmup(global_step,
                              learning_rate_base,
                              total_steps,
@@ -339,7 +393,6 @@ def cosine_decay_with_warmup(global_step,
     learning_rate = max(learning_rate,min_learn_rate)
     return learning_rate
 
-
 class WarmUpCosineDecayScheduler(keras.callbacks.Callback):
     """
     继承Callback，实现对学习率的调度

voc_annotation.py

@@ -7,7 +7,9 @@ import xml.etree.ElementTree as ET
 from os import getcwd
 
 sets=[('2007', 'train'), ('2007', 'val'), ('2007', 'test')]
-
+#-----------------------------------------------------#
+#   这里设定的classes顺序要和model_data里的txt一样
+#-----------------------------------------------------#
 classes = ["aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "diningtable", "dog", "horse", "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]
 
 def convert_annotation(year, image_id, list_file):
@@ -19,20 +21,19 @@ def convert_annotation(year, image_id, list_file):
         difficult = 0 
         if obj.find('difficult')!=None:
             difficult = obj.find('difficult').text
-            
         cls = obj.find('name').text
         if cls not in classes or int(difficult)==1:
             continue
         cls_id = classes.index(cls)
         xmlbox = obj.find('bndbox')
-        b = (int(xmlbox.find('xmin').text), int(xmlbox.find('ymin').text), int(xmlbox.find('xmax').text), int(xmlbox.find('ymax').text))
+        b = (int(float(xmlbox.find('xmin').text)), int(float(xmlbox.find('ymin').text)), int(float(xmlbox.find('xmax').text)), int(float(xmlbox.find('ymax').text)))
         list_file.write(" " + ",".join([str(a) for a in b]) + ',' + str(cls_id))
 
 wd = getcwd()
 
 for year, image_set in sets:
-    image_ids = open('VOCdevkit/VOC%s/ImageSets/Main/%s.txt'%(year, image_set)).read().strip().split()
-    list_file = open('%s_%s.txt'%(year, image_set), 'w')
+    image_ids = open('VOCdevkit/VOC%s/ImageSets/Main/%s.txt'%(year, image_set), encoding='utf-8').read().strip().split()
+    list_file = open('%s_%s.txt'%(year, image_set), 'w', encoding='utf-8')
     for image_id in image_ids:
         list_file.write('%s/VOCdevkit/VOC%s/JPEGImages/%s.jpg'%(wd, year, image_id))
         convert_annotation(year, image_id, list_file)

yolo.py

 import colorsys
-import copy
 import os
-from timeit import default_timer as timer
 
 import numpy as np
 from keras import backend as K
@@ -131,6 +129,11 @@ class YOLO(object):
     #   检测图片
     #---------------------------------------------------#
     def detect_image(self, image):
+        #---------------------------------------------------------#
+        #   在这里将图像转换成RGB图像，防止灰度图在预测时报错。
+        #---------------------------------------------------------#
+        image = image.convert('RGB')
+
         #---------------------------------------------------------#
         #   给图像增加灰条，实现不失真的resize
         #   也可以直接resize进行识别
@@ -138,8 +141,7 @@ class YOLO(object):
         if self.letterbox_image:
             boxed_image = letterbox_image(image, (self.model_image_size[1],self.model_image_size[0]))
         else:
-            boxed_image = image.convert('RGB')
-            boxed_image = boxed_image.resize((self.model_image_size[1],self.model_image_size[0]), Image.BICUBIC)
+            boxed_image = image.resize((self.model_image_size[1],self.model_image_size[0]), Image.BICUBIC)
         image_data = np.array(boxed_image, dtype='float32')
         image_data /= 255.
         #---------------------------------------------------------#
@@ -162,7 +164,7 @@ class YOLO(object):
         #---------------------------------------------------------#
         #   设置字体
         #---------------------------------------------------------#
-        font = ImageFont.truetype(font='font/simhei.ttf',
+        font = ImageFont.truetype(font='model_data/simhei.ttf',
                     size=np.floor(3e-2 * image.size[1] + 0.5).astype('int32'))
 
         thickness = max((image.size[0] + image.size[1]) // 300, 1)