Skip to content

Commit

Permalink
update eval when train
Browse files Browse the repository at this point in the history
  • Loading branch information
@bubbliiiing
bubbliiiing committed May 11, 2022
1 parent bf9bbf6 commit 0834bbc
Show file tree
Hide file tree
Showing 7 changed files with 275 additions and 37 deletions.
1 change: 1 addition & 0 deletions .gitignore
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ Medical_Datasets/
lfw/
logs/
model_data/
.temp_map_out/

# Byte-compiled / optimized / DLL files
__pycache__/
Expand Down
47 changes: 36 additions & 11 deletions get_map.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,11 +15,11 @@

if __name__ == "__main__":
'''
Recall和Precision不像AP是一个面积的概念，在门限值不同时，网络的Recall和Precision值是不同的。
map计算结果中的Recall和Precision代表的是当预测时，门限置信度为0.5时，所对应的Recall和Precision值。
Recall和Precision不像AP是一个面积的概念，因此在门限值（Confidence）不同时，网络的Recall和Precision值是不同的。
默认情况下，本代码计算的Recall和Precision代表的是当门限值（Confidence）为0.5时，所对应的Recall和Precision值。
此处获得的./map_out/detection-results/里面的txt的框的数量会比直接predict多一些，这是因为这里的门限低，
目的是为了计算不同门限条件下的Recall和Precision值，从而实现map的计算。
受到mAP计算原理的限制，网络在计算mAP时需要获得近乎所有的预测框，这样才可以计算不同门限条件下的Recall和Precision值
因此，本代码获得的map_out/detection-results/里面的txt的框的数量一般会比直接predict多一些，目的是列出所有可能的预测框，
'''
#------------------------------------------------------------------------------------------------------------------#
# map_mode用于指定该文件运行时计算的内容
Expand All @@ -30,16 +30,41 @@
# map_mode为4代表利用COCO工具箱计算当前数据集的0.50:0.95map。需要获得预测结果、获得真实框后并安装pycocotools才行
#-------------------------------------------------------------------------------------------------------------------#
map_mode = 0
#-------------------------------------------------------#
#--------------------------------------------------------------------------------------#
# 此处的classes_path用于指定需要测量VOC_map的类别
# 一般情况下与训练和预测所用的classes_path一致即可
#-------------------------------------------------------#
#--------------------------------------------------------------------------------------#
classes_path = 'model_data/voc_classes.txt'
#-------------------------------------------------------#
# MINOVERLAP用于指定想要获得的mAP0.x
#--------------------------------------------------------------------------------------#
# MINOVERLAP用于指定想要获得的mAP0.x，mAP0.x的意义是什么请同学们百度一下。
# 比如计算mAP0.75，可以设定MINOVERLAP = 0.75。
#-------------------------------------------------------#
#
# 当某一预测框与真实框重合度大于MINOVERLAP时，该预测框被认为是正样本，否则为负样本。
# 因此MINOVERLAP的值越大，预测框要预测的越准确才能被认为是正样本，此时算出来的mAP值越低，
#--------------------------------------------------------------------------------------#
MINOVERLAP = 0.5
#--------------------------------------------------------------------------------------#
# 受到mAP计算原理的限制，网络在计算mAP时需要获得近乎所有的预测框，这样才可以计算mAP
# 因此，confidence的值应当设置的尽量小进而获得全部可能的预测框。
#
# 该值一般不调整。因为计算mAP需要获得近乎所有的预测框，此处的confidence不能随便更改。
# 想要获得不同门限值下的Recall和Precision值，请修改下方的score_threhold。
#--------------------------------------------------------------------------------------#
confidence = 0.001
#--------------------------------------------------------------------------------------#
# 预测时使用到的非极大抑制值的大小，越大表示非极大抑制越不严格。
#
# 该值一般不调整。
#--------------------------------------------------------------------------------------#
nms_iou = 0.5
#---------------------------------------------------------------------------------------------------------------#
# Recall和Precision不像AP是一个面积的概念，因此在门限值不同时，网络的Recall和Precision值是不同的。
#
# 默认情况下，本代码计算的Recall和Precision代表的是当门限值为0.5（此处定义为score_threhold）时所对应的Recall和Precision值。
# 因为计算mAP需要获得近乎所有的预测框，上面定义的confidence不能随便更改。
# 这里专门定义一个score_threhold用于代表门限值，进而在计算mAP时找到门限值对应的Recall和Precision值。
#---------------------------------------------------------------------------------------------------------------#
score_threhold = 0.5
#-------------------------------------------------------#
# map_vis用于指定是否开启VOC_map计算的可视化
#-------------------------------------------------------#
Expand Down Expand Up @@ -69,7 +94,7 @@

if map_mode == 0 or map_mode == 1:
print("Load model.")
yolo = YOLO(confidence = 0.001, nms_iou = 0.5)
yolo = YOLO(confidence = confidence, nms_iou = nms_iou)
print("Load model done.")

print("Get predict result.")
Expand Down Expand Up @@ -109,7 +134,7 @@

if map_mode == 0 or map_mode == 3:
print("Get map.")
get_map(MINOVERLAP, True, path = map_out_path)
get_map(MINOVERLAP, True, score_threhold = score_threhold, path = map_out_path)
print("Get map done.")

if map_mode == 4:
Expand Down
23 changes: 19 additions & 4 deletions train.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@

from nets.yolo import get_train_model, yolo_body
from nets.yolo_training import get_lr_scheduler
from utils.callbacks import LossHistory, ModelCheckpoint
from utils.callbacks import EvalCallback, LossHistory, ModelCheckpoint
from utils.dataloader import YoloDatasets
from utils.utils import get_anchors, get_classes, show_config
from utils.utils_fit import fit_one_epoch
Expand Down Expand Up @@ -173,6 +173,17 @@
#------------------------------------------------------------------#
save_dir = 'logs'
#------------------------------------------------------------------#
# eval_flag 是否在训练时进行评估，评估对象为验证集
# 安装pycocotools库后，评估体验更佳。
# eval_period 代表多少个epoch评估一次，不建议频繁的评估
# 评估需要消耗较多的时间，频繁评估会导致训练非常慢
# 此处获得的mAP会与get_map.py获得的会有所不同，原因有二：
# （一）此处获得的mAP为验证集的mAP。
# （二）此处设置评估参数较为保守，目的是加快评估速度。
#------------------------------------------------------------------#
eval_flag = True
eval_period = 10
#------------------------------------------------------------------#
# num_workers 用于设置是否使用多线程读取数据，1代表关闭多线程
# 开启后会加快数据读取速度，但是会占用更多内存
# keras里开启多线程有些时候速度反而慢了许多
Expand Down Expand Up @@ -335,6 +346,8 @@
time_str = datetime.datetime.strftime(datetime.datetime.now(),'%Y_%m_%d_%H_%M_%S')
log_dir = os.path.join(save_dir, "loss_" + str(time_str))
loss_history = LossHistory(log_dir)
eval_callback = EvalCallback(model_body, input_shape, anchors, anchors_mask, class_names, num_classes, val_lines, log_dir, \
eval_flag=eval_flag, period=eval_period)
#---------------------------------------#
# 开始模型训练
#---------------------------------------#
Expand Down Expand Up @@ -386,7 +399,7 @@
lr = lr_scheduler_func(epoch)
K.set_value(optimizer.lr, lr)

fit_one_epoch(model_body, loss_history, optimizer, epoch, epoch_step, epoch_step_val, gen, gen_val,
fit_one_epoch(model_body, loss_history, eval_callback, optimizer, epoch, epoch_step, epoch_step_val, gen, gen_val,
end_epoch, input_shape, anchors, anchors_mask, num_classes, save_period, save_dir, strategy)

train_dataloader.on_epoch_end()
Expand Down Expand Up @@ -419,7 +432,9 @@
monitor = 'val_loss', save_weights_only = True, save_best_only = True, period = 1)
early_stopping = EarlyStopping(monitor='val_loss', min_delta = 0, patience = 10, verbose = 1)
lr_scheduler = LearningRateScheduler(lr_scheduler_func, verbose = 1)
callbacks = [logging, loss_history, checkpoint, checkpoint_last, checkpoint_best, lr_scheduler]
eval_callback = EvalCallback(model_body, input_shape, anchors, anchors_mask, class_names, num_classes, val_lines, log_dir, \
eval_flag=eval_flag, period=eval_period)
callbacks = [logging, loss_history, checkpoint, checkpoint_last, checkpoint_best, lr_scheduler, eval_callback]

if start_epoch < end_epoch:
print('Train on {} samples, val on {} samples, with batch size {}.'.format(num_train, num_val, batch_size))
Expand Down Expand Up @@ -456,7 +471,7 @@
#---------------------------------------#
lr_scheduler_func = get_lr_scheduler(lr_decay_type, Init_lr_fit, Min_lr_fit, UnFreeze_Epoch)
lr_scheduler = LearningRateScheduler(lr_scheduler_func, verbose = 1)
callbacks = [logging, loss_history, checkpoint, checkpoint_last, checkpoint_best, lr_scheduler]
callbacks = [logging, loss_history, checkpoint, checkpoint_last, checkpoint_best, lr_scheduler, eval_callback]

for i in range(len(model_body.layers)):
model_body.layers[i].trainable = True
Expand Down
177 changes: 176 additions & 1 deletion utils/callbacks.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,22 @@
import matplotlib
matplotlib.use('Agg')
from matplotlib import pyplot as plt
import numpy as np
import scipy.signal

import shutil
import numpy as np
import tensorflow as tf

from tensorflow import keras
from tensorflow.keras import backend as K
from tensorflow.keras.layers import Input, Lambda
from tensorflow.keras.models import Model
from PIL import Image
from tqdm import tqdm
from .utils import cvtColor, preprocess_input, resize_image
from .utils_bbox import DecodeBox
from .utils_map import get_coco_map, get_map


class LossHistory(keras.callbacks.Callback):
def __init__(self, log_dir):
Expand Down Expand Up @@ -75,6 +87,169 @@ def on_epoch_end(self, batch, logs=None):
if self.verbose > 0:
print('Setting learning rate to %s.' % (learning_rate))

class EvalCallback(keras.callbacks.Callback):
def __init__(self, model_body, input_shape, anchors, anchors_mask, class_names, num_classes, val_lines, log_dir,\
map_out_path=".temp_map_out", max_boxes=100, confidence=0.05, nms_iou=0.5, letterbox_image=True, MINOVERLAP=0.5, eval_flag=True, period=1):
super(EvalCallback, self).__init__()

self.model_body = model_body
self.input_shape = input_shape
self.anchors = anchors
self.anchors_mask = anchors_mask
self.class_names = class_names
self.num_classes = num_classes
self.val_lines = val_lines
self.log_dir = log_dir
self.map_out_path = map_out_path
self.max_boxes = max_boxes
self.confidence = confidence
self.nms_iou = nms_iou
self.letterbox_image = letterbox_image
self.MINOVERLAP = MINOVERLAP
self.eval_flag = eval_flag
self.period = period

#---------------------------------------------------------#
# 在DecodeBox函数中，我们会对预测结果进行后处理
# 后处理的内容包括，解码、非极大抑制、门限筛选等
#---------------------------------------------------------#
self.input_image_shape = Input([2,],batch_size=1)
inputs = [*self.model_body.output, self.input_image_shape]
outputs = Lambda(
DecodeBox,
output_shape = (1,),
name = 'yolo_eval',
arguments = {
'anchors' : self.anchors,
'num_classes' : self.num_classes,
'input_shape' : self.input_shape,
'anchor_mask' : self.anchors_mask,
'confidence' : self.confidence,
'nms_iou' : self.nms_iou,
'max_boxes' : self.max_boxes,
'letterbox_image' : self.letterbox_image
}
)(inputs)
self.yolo_model = Model([self.model_body.input, self.input_image_shape], outputs)

self.maps = [0]
self.epoches = [0]
if self.eval_flag:
with open(os.path.join(self.log_dir, "epoch_map.txt"), 'a') as f:
f.write(str(0))
f.write("\n")

@tf.function
def get_pred(self, image_data, input_image_shape):
out_boxes, out_scores, out_classes = self.yolo_model([image_data, input_image_shape], training=False)
return out_boxes, out_scores, out_classes

def get_map_txt(self, image_id, image, class_names, map_out_path):
f = open(os.path.join(map_out_path, "detection-results/"+image_id+".txt"),"w")
#---------------------------------------------------------#
# 在这里将图像转换成RGB图像，防止灰度图在预测时报错。
#---------------------------------------------------------#
image = cvtColor(image)
#---------------------------------------------------------#
# 给图像增加灰条，实现不失真的resize
# 也可以直接resize进行识别
#---------------------------------------------------------#
image_data = resize_image(image, (self.input_shape[1], self.input_shape[0]), self.letterbox_image)
#---------------------------------------------------------#
# 添加上batch_size维度，并进行归一化
#---------------------------------------------------------#
image_data = np.expand_dims(preprocess_input(np.array(image_data, dtype='float32')), 0)

#---------------------------------------------------------#
# 将图像输入网络当中进行预测！
#---------------------------------------------------------#
input_image_shape = np.expand_dims(np.array([image.size[1], image.size[0]], dtype='float32'), 0)
outputs = self.get_pred(image_data, input_image_shape)
out_boxes, out_scores, out_classes = [out.numpy() for out in outputs]

top_100 = np.argsort(out_scores)[::-1][:self.max_boxes]
out_boxes = out_boxes[top_100]
out_scores = out_scores[top_100]
out_classes = out_classes[top_100]

for i, c in enumerate(out_classes):
predicted_class = self.class_names[int(c)]
try:
score = str(out_scores[i].numpy())
except:
score = str(out_scores[i])
top, left, bottom, right = out_boxes[i]
if predicted_class not in class_names:
continue

f.write("%s %s %s %s %s %s\n" % (predicted_class, score[:6], str(int(left)), str(int(top)), str(int(right)),str(int(bottom))))

f.close()
return

def on_epoch_end(self, epoch, logs=None):
temp_epoch = epoch + 1
if temp_epoch % self.period == 0 and self.eval_flag:
if not os.path.exists(self.map_out_path):
os.makedirs(self.map_out_path)
if not os.path.exists(os.path.join(self.map_out_path, "ground-truth")):
os.makedirs(os.path.join(self.map_out_path, "ground-truth"))
if not os.path.exists(os.path.join(self.map_out_path, "detection-results")):
os.makedirs(os.path.join(self.map_out_path, "detection-results"))
print("Get map.")
for annotation_line in tqdm(self.val_lines):
line = annotation_line.split()
image_id = os.path.basename(line[0]).split('.')[0]
#------------------------------#
# 读取图像并转换成RGB图像
#------------------------------#
image = Image.open(line[0])
#------------------------------#
# 获得预测框
#------------------------------#
gt_boxes = np.array([np.array(list(map(int,box.split(',')))) for box in line[1:]])
#------------------------------#
# 获得预测txt
#------------------------------#
self.get_map_txt(image_id, image, self.class_names, self.map_out_path)

#------------------------------#
# 获得真实框txt
#------------------------------#
with open(os.path.join(self.map_out_path, "ground-truth/"+image_id+".txt"), "w") as new_f:
for box in gt_boxes:
left, top, right, bottom, obj = box
obj_name = self.class_names[obj]
new_f.write("%s %s %s %s %s\n" % (obj_name, left, top, right, bottom))

print("Calculate Map.")
try:
temp_map = get_coco_map(class_names = self.class_names, path = self.map_out_path)[1]
except:
temp_map = get_map(self.MINOVERLAP, False, path = self.map_out_path)
self.maps.append(temp_map)
self.epoches.append(temp_epoch)

with open(os.path.join(self.log_dir, "epoch_map.txt"), 'a') as f:
f.write(str(temp_map))
f.write("\n")

plt.figure()
plt.plot(self.epoches, self.maps, 'red', linewidth = 2, label='train map')

plt.grid(True)
plt.xlabel('Epoch')
plt.ylabel('Map %s'%str(self.MINOVERLAP))
plt.title('A Map Curve')
plt.legend(loc="upper right")

plt.savefig(os.path.join(self.log_dir, "epoch_map.png"))
plt.cla()
plt.close("all")

print("Get map done.")
shutil.rmtree(self.map_out_path)

class ModelCheckpoint(keras.callbacks.Callback):
def __init__(self, filepath, monitor='val_loss', verbose=0,
save_best_only=False, save_weights_only=False,
Expand Down
Loading

0 comments on commit 0834bbc

Please sign in to comment.