def transform_to_onnx(weight_file, batch_size, IN_IMAGE_H, IN_IMAGE_W):
model = Build_Model()
pretrained_dict = torch.load(weight_file, map_location=torch.device('cpu'))
model.load_state_dict(pretrained_dict)
input_names = ["input"]
output_names = ['boxes', 'confs']
dynamic = False
if batch_size <= 0:
dynamic = True
if dynamic:
x = torch.randn((1, 3, IN_IMAGE_H, IN_IMAGE_W), requires_grad=True)
onnx_file_name = "yolov4_-1_3_{}_{}_dynamic.onnx".format(IN_IMAGE_H, IN_IMAGE_W)
dynamic_axes = {"input": {0: "batch_size"}, "boxes": {0: "batch_size"}, "confs": {0: "batch_size"}}
# Export the model
print('Export the onnx model ...')
torch.onnx.export(model,
x,
onnx_file_name,
export_params=True,
opset_version=11,
do_constant_folding=True,
input_names=input_names, output_names=output_names,
dynamic_axes=dynamic_axes)
print('Onnx model exporting done')
return onnx_file_name
else:
x = torch.randn((batch_size, 3, IN_IMAGE_H, IN_IMAGE_W), requires_grad=True)
onnx_file_name = "yolov4_{}_3_{}_{}_static.onnx".format(batch_size, IN_IMAGE_H, IN_IMAGE_W)
# Export the model
print('Export the onnx model ...')
torch.onnx.export(model,
x,
onnx_file_name,
export_params=True,
do_constant_folding=True,
input_names=input_names, output_names=output_names,
)
print('Onnx model exporting done')
return onnx_file_name
class Trainer(object):
def __init__(self, weight_path, resume, gpu_id, accumulate):
init_seeds(0)
self.device = gpu.select_device(gpu_id)
self.start_epoch = 0
self.best_mAP = 0.
self.accumulate = accumulate
self.epochs = cfg.TRAIN["EPOCHS"]
self.weight_path = weight_path
self.multi_scale_train = cfg.TRAIN["MULTI_SCALE_TRAIN"]
self.train_dataset = data.Build_Dataset(
anno_file_type="train", img_size=cfg.TRAIN["TRAIN_IMG_SIZE"])
print('train img size is {}'.format(cfg.TRAIN["TRAIN_IMG_SIZE"]))
self.train_dataloader = DataLoader(
self.train_dataset,
batch_size=cfg.TRAIN["BATCH_SIZE"],
num_workers=cfg.TRAIN["NUMBER_WORKERS"],
shuffle=True,
pin_memory=True)
self.yolov4 = Build_Model().to(self.device)
self.optimizer = optim.SGD(self.yolov4.parameters(),
lr=cfg.TRAIN["LR_INIT"],
momentum=cfg.TRAIN["MOMENTUM"],
weight_decay=cfg.TRAIN["WEIGHT_DECAY"])
self.criterion = YoloV4Loss(
anchors=cfg.MODEL["ANCHORS"],
strides=cfg.MODEL["STRIDES"],
iou_threshold_loss=cfg.TRAIN["IOU_THRESHOLD_LOSS"])
self.__load_model_weights(weight_path, resume)
self.scheduler = cosine_lr_scheduler.CosineDecayLR(
self.optimizer,
T_max=self.epochs * len(self.train_dataloader),
lr_init=cfg.TRAIN["LR_INIT"],
lr_min=cfg.TRAIN["LR_END"],
warmup=cfg.TRAIN["WARMUP_EPOCHS"] * len(self.train_dataloader))
def __load_model_weights(self, weight_path, resume):
if resume:
last_weight = os.path.join(
os.path.split(weight_path)[0], "last.pt")
chkpt = torch.load(last_weight, map_location=self.device)
self.yolov4.load_state_dict(chkpt['model'])
self.start_epoch = chkpt['epoch'] + 1
if chkpt['optimizer'] is not None:
self.optimizer.load_state_dict(chkpt['optimizer'])
self.best_mAP = chkpt['best_mAP']
del chkpt
else:
self.yolov4.load_darknet_weights(weight_path)
# def __save_model_weights(self, epoch, mAP):
# if mAP > self.best_mAP:
# self.best_mAP = mAP
# best_weight = os.path.join(os.path.split(self.weight_path)[0], "best.pt")
# last_weight = os.path.join(os.path.split(self.weight_path)[0], "last.pt")
# chkpt = {'epoch': epoch,
# 'best_mAP': self.best_mAP,
# 'model': self.yolov3.state_dict(),
# 'optimizer': self.optimizer.state_dict()}
# torch.save(chkpt, last_weight)
#
# if self.best_mAP == mAP:
# torch.save(chkpt['model'], best_weight)
#
# if epoch > 0 and epoch % 10 == 0:
# torch.save(chkpt, os.path.join(os.path.split(self.weight_path)[0], 'backup_epoch%g.pt'%epoch))
# del chkpt
def __save_model_weights(self, epoch, loss):
weight = os.path.join(
os.path.split(self.weight_path)[0],
"epoch{}, loss{}.pt".format(epoch, loss))
chkpt = {
'epoch': epoch,
'model': self.yolov4.state_dict(),
'optimizer': self.optimizer.state_dict()
}
torch.save(chkpt['model'], weight)
def train(self):
global writer # 创建一个SummaryWriter实例
logger.info(
"Training start,img size is: {:d},batchsize is: {:d},work number is {:d}"
.format(cfg.TRAIN["TRAIN_IMG_SIZE"], cfg.TRAIN["BATCH_SIZE"],
cfg.TRAIN["NUMBER_WORKERS"]))
logger.info(self.yolov4)
logger.info("Train datasets number is : {}".format(
len(self.train_dataset)))
self.yolov4, self.optimizer = amp.initialize(self.yolov4,
self.optimizer,
opt_level='O1',
verbosity=0)
logger.info(" ======= start training ====== ")
for epoch in range(self.start_epoch, self.epochs):
start = time.time()
self.yolov4.train()
mloss = torch.zeros(4)
logger.info("===Epoch:[{}/{}]===".format(epoch, self.epochs))
for i, (imgs, label_sbbox, label_mbbox, label_lbbox, sbboxes,
mbboxes, lbboxes) in enumerate(self.train_dataloader):
self.scheduler.step(len(self.train_dataloader) * epoch + i)
imgs = imgs.to(self.device)
label_sbbox = label_sbbox.to(self.device)
label_mbbox = label_mbbox.to(self.device)
label_lbbox = label_lbbox.to(self.device)
sbboxes = sbboxes.to(self.device)
mbboxes = mbboxes.to(self.device)
lbboxes = lbboxes.to(self.device)
p, p_d = self.yolov4(imgs)
loss, loss_giou, loss_conf, loss_cls = self.criterion(
p, p_d, label_sbbox, label_mbbox, label_lbbox, sbboxes,
mbboxes, lbboxes)
with amp.scale_loss(loss, self.optimizer) as scaled_loss:
scaled_loss.backward()
# Accumulate gradient for x batches before optimizing
if i % self.accumulate == 0:
self.optimizer.step()
self.optimizer.zero_grad()
# Update running mean of tracked metrics
loss_items = torch.tensor(
[loss_giou, loss_conf, loss_cls, loss])
mloss = (mloss * i + loss_items) / (i + 1)
# Print batch results
if i % 10 == 0:
logger.info(
" === Epoch:[{:3}/{}],step:[{:3}/{}],img_size:[{:3}],total_loss:{:.4f}|loss_giou:{:.4f}|loss_conf:{:.4f}|loss_cls:{:.4f}|lr:{:.4f}"
.format(epoch, self.epochs, i,
len(self.train_dataloader) - 1,
self.train_dataset.img_size, mloss[3],
mloss[0], mloss[1], mloss[2],
self.optimizer.param_groups[0]['lr']))
writer.add_scalar('train_loss', mloss[3], i)
# multi-sclae training (320-608 pixels) every 10 batches
if self.multi_scale_train and (i + 1) % 10 == 0:
# self.train_dataset.img_size = random.choice(range(5, 15)) * 32 # for imgsize 320
# self.train_dataset.img_size = random.choice(range(12, 22)) * 32 # for imgsize 544
self.train_dataset.img_size = random.choice(range(10,
20)) * 32
if epoch >= 0 and cfg.TRAIN["DATA_TYPE"] == 'VOC':
self.__save_model_weights(epoch, mloss[3])
print('save weights done')
# mAP = 0
# if epoch >= 0:
# logger.info("===== Validate =====".format(epoch, self.epochs))
# with torch.no_grad():
# APs = Evaluator(self.yolov3,epoch,showatt=False).APs_voc()
# for i in APs:
# print("{} --> mAP : {}".format(i, APs[i]))
# mAP += APs[i]
# mAP = mAP / self.train_dataset.num_classes
# writer.add_scalar('mAP', mAP, epoch)
# logger.info(" ===test mAP:{:.3f}".format(mAP))
# writer.add_scalar('mAP', mAP, epoch)
elif epoch >= 0 and cfg.TRAIN["DATA_TYPE"] == 'COCO':
evaluator = COCOAPIEvaluator(model_type='YOLOv4',
data_dir=cfg.DATA_PATH,
img_size=cfg.VAL["TEST_IMG_SIZE"],
confthre=0.08,
nmsthre=cfg.VAL["NMS_THRESH"])
ap50_95, ap50 = evaluator.evaluate(self.yolov4)
# logger.info('ap50_95:{}|ap50:{}'.format(ap50_95, ap50))
writer.add_scalar('val/COCOAP50', ap50, epoch)
writer.add_scalar('val/COCOAP50_95', ap50_95, epoch)
self.__save_model_weights(epoch, ap50)
else:
assert print('dataset must be VOC or COCO')
end = time.time()
logger.info(" ===cost time:{:.4f}s".format(end - start))
logger.info(
"=====Training Finished. best_test_mAP:{:.3f}%====".format(
self.best_mAP))
class Trainer(object):
_resume = False
_fine_tune = False
def __init__(self, log_dir, resume=False, fine_tune=False):
init_seeds(0)
self._fine_tune = fine_tune
self._resume = resume
if self._fine_tune:
self.__prepare_fine_tune()
self.device = gpu.select_device()
self.start_epoch = 0
self.best_mAP = 0.
self.accumulate = cfg.TRAIN.ACCUMULATE
self.log_dir = log_dir
self.weight_path = "yolov4.weights"
self.multi_scale_train = cfg.TRAIN.MULTI_SCALE_TRAIN
if self.multi_scale_train:
print('Using multi scales training')
else:
print('train img size is {}'.format(cfg.TRAIN.TRAIN_IMG_SIZE))
self.train_dataset = data.Build_Train_Dataset(anno_file=cfg.TRAIN.ANNO_FILE, anno_file_type="train", img_size=cfg.TRAIN.TRAIN_IMG_SIZE)
self.epochs = cfg.TRAIN.YOLO_EPOCHS if cfg.MODEL.MODEL_TYPE == 'YOLOv4' else cfg.TRAIN.Mobilenet_YOLO_EPOCHS
self.train_dataloader = DataLoader(self.train_dataset,
batch_size=cfg.TRAIN.BATCH_SIZE//cfg.TRAIN.ACCUMULATE,
num_workers=cfg.TRAIN.NUMBER_WORKERS,
shuffle=True, pin_memory=True)
self.yolov4 = Build_Model(weight_path="yolov4.weights", resume=resume)
self.yolov4 = self.yolov4.to(self.device)
self.optimizer = optim.SGD(self.yolov4.parameters(), lr=cfg.TRAIN.LR_INIT,
momentum=cfg.TRAIN.MOMENTUM, weight_decay=cfg.TRAIN.WEIGHT_DECAY)
self.criterion = YoloV4Loss(anchors=cfg.MODEL.ANCHORS, strides=cfg.MODEL.STRIDES,
iou_threshold_loss=cfg.TRAIN.IOU_THRESHOLD_LOSS)
self.scheduler = cosine_lr_scheduler.CosineDecayLR(self.optimizer,
T_max=self.epochs*len(self.train_dataloader),
lr_init=cfg.TRAIN.LR_INIT,
lr_min=cfg.TRAIN.LR_END,
warmup=cfg.TRAIN.WARMUP_EPOCHS*len(self.train_dataloader))
if resume: self.__load_resume_weights()
if self._fine_tune: self.__load_best_weights()
def __prepare_fine_tune(self):
cfg.defrost()
cfg.TRAIN.ANNO_FILE = cfg.FINE_TUNE.ANNO_FILE
cfg.TRAIN.YOLO_EPOCHS = cfg.FINE_TUNE.YOLO_EPOCHS
cfg.TRAIN.LR_INIT = cfg.FINE_TUNE.LR_INIT
cfg.TRAIN.LR_END = cfg.FINE_TUNE.LR_END
cfg.TRAIN.WARMUP_EPOCHS = cfg.FINE_TUNE.WARMUP_EPOCHS
cfg.freeze()
def __load_best_weights(self):
best_weight = os.path.join(log_dir,"checkpoints", "best.pt")
last_weight = os.path.join(log_dir,"checkpoints", "last.pt")
shutil.copy2(best_weight,
best_weight.replace("best.pt","best_before_fine_tune.pt"))
shutil.copy2(last_weight,
last_weight.replace("last.pt","last_before_fine_tune.pt"))
last_chkpt = torch.load(last_weight, map_location=self.device)
best_chkpt = torch.load(best_weight, map_location=self.device)
self.yolov4.load_state_dict(best_chkpt)
self.best_mAP = 0
del last_chkpt, best_chkpt
def __load_resume_weights(self):
last_weight = os.path.join(log_dir,"checkpoints", "last.pt")
chkpt = torch.load(last_weight, map_location=self.device)
self.yolov4.load_state_dict(chkpt['model'])
self.start_epoch = chkpt['epoch'] + 1
if chkpt['optimizer'] is not None:
self.optimizer.load_state_dict(chkpt['optimizer'])
self.best_mAP = chkpt['best_mAP']
del chkpt
def __save_model_weights(self, epoch, mAP):
if mAP > self.best_mAP:
self.best_mAP = mAP
best_weight = os.path.join(log_dir,"checkpoints", "best.pt")
last_weight = os.path.join(log_dir,"checkpoints", "last.pt")
chkpt = {'epoch': epoch,
'best_mAP': self.best_mAP,
'model': self.yolov4.module.state_dict() if torch.cuda.device_count()>1 else self.yolov4.state_dict(),
'optimizer': self.optimizer.state_dict()}
torch.save(chkpt, last_weight)
if self.best_mAP == mAP:
torch.save(chkpt['model'], best_weight)
if self._fine_tune and epoch % 5 == 0:
torch.save(chkpt['model'], os.path.join(log_dir,"checkpoints", 'backup_fine_tune_epoch_{:02d}.pt'.format(epoch)))
del chkpt
def train(self):
global writer
logger.info("Training start,img size is: {:d},batchsize is: {:d}, subdivision: {:d}, worker number is {:d}".format(cfg.TRAIN.TRAIN_IMG_SIZE, cfg.TRAIN.BATCH_SIZE, cfg.TRAIN.ACCUMULATE, cfg.TRAIN.NUMBER_WORKERS))
logger.info(self.yolov4)
n_train = len(self.train_dataset)
n_step = n_train // (cfg.TRAIN.BATCH_SIZE//cfg.TRAIN.ACCUMULATE) + 1
logger.info("Train datasets number is : {}".format(n_train))
evaluator = COCOAPIEvaluator(cfg=cfg,
img_size=cfg.VAL.TEST_IMG_SIZE,
confthre=cfg.VAL.CONF_THRESH,
nmsthre=cfg.VAL.NMS_THRESH)
if torch.cuda.device_count() > 1: self.yolov4 = torch.nn.DataParallel(self.yolov4)
logger.info("\n=============== start training ===============")
for epoch in range(self.start_epoch, self.epochs):
start = time.time()
self.yolov4.train()
with tqdm(total=n_train, unit="imgs", desc=f'Epoch {epoch}/{self.epochs}', ncols=30) as pbar:
for i, (imgs, label_sbbox, label_mbbox, label_lbbox, sbboxes, mbboxes, lbboxes) in enumerate(self.train_dataloader):
imgs = imgs.to(self.device)
label_sbbox = label_sbbox.to(self.device)
label_mbbox = label_mbbox.to(self.device)
label_lbbox = label_lbbox.to(self.device)
sbboxes = sbboxes.to(self.device)
mbboxes = mbboxes.to(self.device)
lbboxes = lbboxes.to(self.device)
p, p_d = self.yolov4(imgs)
loss, loss_ciou, loss_conf, loss_cls = self.criterion(p, p_d, label_sbbox, label_mbbox,
label_lbbox, sbboxes, mbboxes, lbboxes)
loss.backward()
# Accumulate gradient for x batches before optimizing
if i % self.accumulate == 0:
self.scheduler.step(n_step*epoch + i)
self.optimizer.step()
self.optimizer.zero_grad()
# Print batch results
if i % (5*self.accumulate) == 0:
logger.info("{:3}: total_loss:{:.4f} | loss_ciou:{:.4f} | loss_conf:{:.4f} | loss_cls:{:.4f} | lr:{:.6f}".format(
self.train_dataset.img_size, loss, loss_ciou, loss_conf, loss_cls, self.optimizer.param_groups[0]['lr']
))
writer.add_scalar('train/loss_ciou', loss_ciou, n_step * epoch + i)
writer.add_scalar('train/loss_conf', loss_conf, n_step * epoch + i)
writer.add_scalar('train/loss_cls', loss_cls, n_step * epoch + i)
writer.add_scalar('train/train_loss', loss, n_step * epoch + i)
writer.add_scalar('train/lr', self.optimizer.param_groups[0]['lr'], n_step * epoch + i)
# multi-sclae training (320-608 pixels) every 10 batches
if self.multi_scale_train and (i+1) % (5*self.accumulate) == 0:
self.train_dataset.img_size = random.choice(range(10, 20)) * 32
pbar.update(imgs.shape[0])
mAP = 0.
# evaluator = COCOAPIEvaluator(cfg=cfg,
# img_size=cfg.VAL.TEST_IMG_SIZE,
# confthre=cfg.VAL.CONF_THRESH,
# nmsthre=cfg.VAL.NMS_THRESH)
coco_stat = evaluator.evaluate(self.yolov4)
logger.info("Average Precision (AP) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = {:.04f}".format(coco_stat[0]))
logger.info("Average Precision (AP) @[ IoU=0.50 | area= all | maxDets=100 ] = {:.04f}".format(coco_stat[1]))
logger.info("Average Precision (AP) @[ IoU=0.75 | area= all | maxDets=100 ] = {:.04f}".format(coco_stat[2]))
logger.info("Average Precision (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = {:.04f}".format(coco_stat[3]))
logger.info("Average Precision (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = {:.04f}".format(coco_stat[4]))
logger.info("Average Precision (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = {:.04f}".format(coco_stat[5]))
logger.info("Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 1 ] = {:.04f}".format(coco_stat[6]))
logger.info("Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 10 ] = {:.04f}".format(coco_stat[7]))
logger.info("Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = {:.04f}".format(coco_stat[8]))
logger.info("Average Recall (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = {:.04f}".format(coco_stat[9]))
logger.info("Average Recall (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = {:.04f}".format(coco_stat[10]))
logger.info("Average Recall (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = {:.04f}".format(coco_stat[11]))
writer.add_scalar('val/mAP_50_95', coco_stat[0], epoch)
writer.add_scalar('val/mAP_50', coco_stat[1], epoch)
writer.add_scalar('val/mAP_75', coco_stat[2], epoch)
writer.add_scalar('val/mAP_small', coco_stat[3], epoch)
writer.add_scalar('val/mAP_medium', coco_stat[4], epoch)
writer.add_scalar('val/mAP_large', coco_stat[5], epoch)
writer.add_scalar('val/mAR_max_1', coco_stat[6], epoch)
writer.add_scalar('val/mAR_max_10', coco_stat[7], epoch)
writer.add_scalar('val/mAR_max_100',coco_stat[8], epoch)
writer.add_scalar('val/mAR_small', coco_stat[9], epoch)
writer.add_scalar('val/mAR_medium', coco_stat[10], epoch)
writer.add_scalar('val/mAR_large', coco_stat[11], epoch)
self.__save_model_weights(epoch, coco_stat[0])
logger.info('save weights done')
end = time.time()
logger.info("cost time:{:.4f}s".format(end - start))
logger.info("=====Training Finished. best_test_mAP:{:.3f}%====".format(self.best_mAP))
class Evaluation(object):
def __init__(self,
gpu_id=0,
weight_path=None,
visiual=None,
eval=False,
showatt=False,
mode=None):
self.__num_class = cfg.VOC_DATA["NUM"]
self.__conf_threshold = cfg.VAL["CONF_THRESH"]
self.__nms_threshold = cfg.VAL["NMS_THRESH"]
self.__device = gpu.select_device(gpu_id)
self.__multi_scale_val = cfg.VAL["MULTI_SCALE_VAL"]
self.__flip_val = cfg.VAL["FLIP_VAL"]
self.__showatt = showatt
self.__visiual = visiual
self.__eval = eval
self.__mode = mode
self.__classes = cfg.VOC_DATA["CLASSES"]
self.__model = Build_Model(showatt=self.__showatt).to(self.__device)
self.__load_model_weights(weight_path)
self.__evalter = Evaluator(self.__model, showatt=self.__showatt)
def __load_model_weights(self, weight_path):
print("loading weight file from : {}".format(weight_path))
weight = os.path.join(weight_path)
chkpt = torch.load(weight, map_location=self.__device)
self.__model.load_state_dict(chkpt["model"])
print("loading weight file is done")
del chkpt
def val(self):
global logger
if self.__eval:
logger.info("***********Start Evaluation****************")
start = time.time()
mAP = 0
with torch.no_grad():
APs, inference_time = Evaluator(self.__model,
showatt=False).APs_voc(
self.__multi_scale_val,
self.__flip_val)
for i in APs:
logger.info("{} --> mAP : {}".format(i, APs[i]))
mAP += APs[i]
mAP = mAP / self.__num_class
logger.info("mAP:{}".format(mAP))
logger.info("inference time: {:.2f} ms".format(inference_time))
end = time.time()
logger.info(" ===val cost time:{:.4f}s".format(end - start))
def detection(self):
global logger
if self.__visiual:
imgs = os.listdir(self.__visiual)
logger.info("***********Start Detection****************")
for v in imgs:
path = os.path.join(self.__visiual, v)
logger.info("val images : {}".format(path))
img = cv2.imread(path)
assert img is not None
bboxes_prd = self.__evalter.get_bbox(img, v, mode=self.__mode)
if bboxes_prd.shape[0] != 0:
boxes = bboxes_prd[..., :4]
class_inds = bboxes_prd[..., 5].astype(np.int32)
scores = bboxes_prd[..., 4]
visualize_boxes(
image=img,
boxes=boxes,
labels=class_inds,
probs=scores,
class_labels=self.__classes,
)
path = os.path.join(cfg.PROJECT_PATH,
"detection_result/{}".format(v))
cv2.imwrite(path, img)
logger.info("saved images : {}".format(path))
class Trainer(object):
def __init__(self,
weight_path=None,
resume=False,
gpu_id=0,
accumulate=1,
fp_16=False):
init_seeds(0)
self.fp_16 = fp_16
self.device = gpu.select_device(gpu_id)
self.start_epoch = 0
self.best_mAP = 0.0
self.accumulate = accumulate
self.weight_path = weight_path
self.multi_scale_train = cfg.TRAIN["MULTI_SCALE_TRAIN"]
self.showatt = cfg.TRAIN["showatt"]
if self.multi_scale_train:
print("Using multi scales training")
else:
print("train img size is {}".format(cfg.TRAIN["TRAIN_IMG_SIZE"]))
self.train_dataset = data.Build_Dataset(
anno_file_type="train", img_size=cfg.TRAIN["TRAIN_IMG_SIZE"])
self.epochs = (cfg.TRAIN["YOLO_EPOCHS"] if cfg.MODEL_TYPE["TYPE"]
== "YOLOv4" else cfg.TRAIN["Mobilenet_YOLO_EPOCHS"])
self.eval_epoch = (30 if cfg.MODEL_TYPE["TYPE"] == "YOLOv4" else 50)
self.train_dataloader = DataLoader(
self.train_dataset,
batch_size=cfg.TRAIN["BATCH_SIZE"],
num_workers=cfg.TRAIN["NUMBER_WORKERS"],
shuffle=True,
pin_memory=True,
)
self.yolov4 = Build_Model(weight_path=weight_path,
resume=resume,
showatt=self.showatt).to(self.device)
self.optimizer = optim.SGD(
self.yolov4.parameters(),
lr=cfg.TRAIN["LR_INIT"],
momentum=cfg.TRAIN["MOMENTUM"],
weight_decay=cfg.TRAIN["WEIGHT_DECAY"],
)
self.criterion = YoloV4Loss(
anchors=cfg.MODEL["ANCHORS"],
strides=cfg.MODEL["STRIDES"],
iou_threshold_loss=cfg.TRAIN["IOU_THRESHOLD_LOSS"],
)
self.scheduler = cosine_lr_scheduler.CosineDecayLR(
self.optimizer,
T_max=self.epochs * len(self.train_dataloader),
lr_init=cfg.TRAIN["LR_INIT"],
lr_min=cfg.TRAIN["LR_END"],
warmup=cfg.TRAIN["WARMUP_EPOCHS"] * len(self.train_dataloader),
)
if resume:
self.__load_resume_weights(weight_path)
def __load_resume_weights(self, weight_path):
last_weight = os.path.join(os.path.split(weight_path)[0], "last.pt")
chkpt = torch.load(last_weight, map_location=self.device)
self.yolov4.load_state_dict(chkpt["model"])
self.start_epoch = chkpt["epoch"] + 1
if chkpt["optimizer"] is not None:
self.optimizer.load_state_dict(chkpt["optimizer"])
self.best_mAP = chkpt["best_mAP"]
del chkpt
def __save_model_weights(self, epoch, mAP):
if mAP > self.best_mAP:
self.best_mAP = mAP
best_weight = os.path.join(
os.path.split(self.weight_path)[0], "best.pt")
last_weight = os.path.join(
os.path.split(self.weight_path)[0], "last.pt")
chkpt = {
"epoch": epoch,
"best_mAP": self.best_mAP,
"model": self.yolov4.state_dict(),
"optimizer": self.optimizer.state_dict(),
}
torch.save(chkpt, last_weight)
if self.best_mAP == mAP:
torch.save(chkpt["model"], best_weight)
if epoch > 0 and epoch % 10 == 0:
torch.save(
chkpt,
os.path.join(
os.path.split(self.weight_path)[0],
"backup_epoch%g.pt" % epoch,
),
)
del chkpt
def train(self):
global writer
logger.info(
"Training start,img size is: {:d},batchsize is: {:d},work number is {:d}"
.format(
cfg.TRAIN["TRAIN_IMG_SIZE"],
cfg.TRAIN["BATCH_SIZE"],
cfg.TRAIN["NUMBER_WORKERS"],
))
logger.info(self.yolov4)
logger.info("Train datasets number is : {}".format(
len(self.train_dataset)))
def is_valid_number(x):
return not (math.isnan(x) or math.isinf(x) or x > 1e4)
if self.fp_16:
self.yolov4, self.optimizer = amp.initialize(self.yolov4,
self.optimizer,
opt_level="O1",
verbosity=0)
logger.info(" ======= start training ====== ")
for epoch in range(self.start_epoch, self.epochs):
start = time.time()
self.yolov4.train()
mloss = torch.zeros(4)
logger.info("===Epoch:[{}/{}]===".format(epoch, self.epochs))
for i, (
imgs,
label_sbbox,
label_mbbox,
label_lbbox,
sbboxes,
mbboxes,
lbboxes,
) in enumerate(self.train_dataloader):
self.scheduler.step(
len(self.train_dataloader) /
(cfg.TRAIN["BATCH_SIZE"]) * epoch + i)
imgs = imgs.to(self.device)
label_sbbox = label_sbbox.to(self.device)
label_mbbox = label_mbbox.to(self.device)
label_lbbox = label_lbbox.to(self.device)
sbboxes = sbboxes.to(self.device)
mbboxes = mbboxes.to(self.device)
lbboxes = lbboxes.to(self.device)
p, p_d = self.yolov4(imgs)
loss, loss_ciou, loss_conf, loss_cls = self.criterion(
p,
p_d,
label_sbbox,
label_mbbox,
label_lbbox,
sbboxes,
mbboxes,
lbboxes,
)
if is_valid_number(loss.item()):
if self.fp_16:
with amp.scale_loss(loss,
self.optimizer) as scaled_loss:
scaled_loss.backward()
else:
loss.backward()
# Accumulate gradient for x batches before optimizing
if i % self.accumulate == 0:
self.optimizer.step()
self.optimizer.zero_grad()
# Update running mean of tracked metrics
loss_items = torch.tensor(
[loss_ciou, loss_conf, loss_cls, loss])
mloss = (mloss * i + loss_items) / (i + 1)
# Print batch results
if i % 10 == 0:
logger.info(
" === Epoch:[{:3}/{}],step:[{:3}/{}],img_size:[{:3}],total_loss:{:.4f}|loss_ciou:{:.4f}|loss_conf:{:.4f}|loss_cls:{:.4f}|lr:{:.4f}"
.format(
epoch,
self.epochs,
i,
len(self.train_dataloader) - 1,
self.train_dataset.img_size,
mloss[3],
mloss[0],
mloss[1],
mloss[2],
self.optimizer.param_groups[0]["lr"],
))
writer.add_scalar(
"loss_ciou",
mloss[0],
len(self.train_dataloader) * epoch + i,
)
writer.add_scalar(
"loss_conf",
mloss[1],
len(self.train_dataloader) * epoch + i,
)
writer.add_scalar(
"loss_cls",
mloss[2],
len(self.train_dataloader) * epoch + i,
)
writer.add_scalar(
"train_loss",
mloss[3],
len(self.train_dataloader) * epoch + i,
)
# multi-sclae training (320-608 pixels) every 10 batches
if self.multi_scale_train and (i + 1) % 10 == 0:
self.train_dataset.img_size = (
random.choice(range(10, 20)) * 32)
if (cfg.TRAIN["DATA_TYPE"] == "VOC"
or cfg.TRAIN["DATA_TYPE"] == "Customer"):
mAP = 0.0
if epoch >= self.eval_epoch:
logger.info("===== Validate =====".format(
epoch, self.epochs))
logger.info("val img size is {}".format(
cfg.VAL["TEST_IMG_SIZE"]))
with torch.no_grad():
APs, inference_time = Evaluator(
self.yolov4, showatt=self.showatt).APs_voc()
for i in APs:
logger.info("{} --> mAP : {}".format(i, APs[i]))
mAP += APs[i]
mAP = mAP / self.train_dataset.num_classes
logger.info("mAP : {}".format(mAP))
logger.info(
"inference time: {:.2f} ms".format(inference_time))
writer.add_scalar("mAP", mAP, epoch)
self.__save_model_weights(epoch, mAP)
logger.info("save weights done")
logger.info(" ===test mAP:{:.3f}".format(mAP))
elif epoch >= 0 and cfg.TRAIN["DATA_TYPE"] == "COCO":
evaluator = COCOAPIEvaluator(
model_type="YOLOv4",
data_dir=cfg.DATA_PATH,
img_size=cfg.VAL["TEST_IMG_SIZE"],
confthre=0.08,
nmsthre=cfg.VAL["NMS_THRESH"],
)
ap50_95, ap50 = evaluator.evaluate(self.yolov4)
logger.info("ap50_95:{}|ap50:{}".format(ap50_95, ap50))
writer.add_scalar("val/COCOAP50", ap50, epoch)
writer.add_scalar("val/COCOAP50_95", ap50_95, epoch)
self.__save_model_weights(epoch, ap50)
print("save weights done")
end = time.time()
logger.info(" ===cost time:{:.4f}s".format(end - start))
logger.info(
"=====Training Finished. best_test_mAP:{:.3f}%====".format(
self.best_mAP))
class Tester(object):
def __init__(self, log_dir):
init_seeds(0)
self.device = gpu.select_device()
self.log_dir = log_dir
self.yolov4 = Build_Model(weight_path=None, resume=False)
self.yolov4 = self.yolov4.to(self.device)
self.__load_best_weights()
def __load_best_weights(self):
best_weight = os.path.join(log_dir, "checkpoints", "best.pt")
chkpt = torch.load(best_weight, map_location=self.device)
self.yolov4.load_state_dict(chkpt)
del chkpt
def test(self):
logger.info(self.yolov4)
evaluator = COCOAPIEvaluator(cfg=cfg,
img_size=cfg.VAL.TEST_IMG_SIZE,
confthre=cfg.VAL.CONF_THRESH,
nmsthre=cfg.VAL.NMS_THRESH)
logger.info("\n=============== start testing ===============")
start = time.time()
coco_stat = evaluator.evaluate(self.yolov4)
end = time.time()
logger.info(
"Average Precision (AP) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = {:.04f}"
.format(coco_stat[0]))
logger.info(
"Average Precision (AP) @[ IoU=0.50 | area= all | maxDets=100 ] = {:.04f}"
.format(coco_stat[1]))
logger.info(
"Average Precision (AP) @[ IoU=0.75 | area= all | maxDets=100 ] = {:.04f}"
.format(coco_stat[2]))
logger.info(
"Average Precision (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = {:.04f}"
.format(coco_stat[3]))
logger.info(
"Average Precision (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = {:.04f}"
.format(coco_stat[4]))
logger.info(
"Average Precision (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = {:.04f}"
.format(coco_stat[5]))
logger.info(
"Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 1 ] = {:.04f}"
.format(coco_stat[6]))
logger.info(
"Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 10 ] = {:.04f}"
.format(coco_stat[7]))
logger.info(
"Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = {:.04f}"
.format(coco_stat[8]))
logger.info(
"Average Recall (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = {:.04f}"
.format(coco_stat[9]))
logger.info(
"Average Recall (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = {:.04f}"
.format(coco_stat[10]))
logger.info(
"Average Recall (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = {:.04f}"
.format(coco_stat[11]))
logger.info("cost time:{:.4f}s".format(end - start))
logger.info("FPS:{:.4f}".format(
len(evaluator.dataloader) / (end - start)))
class Evaluation(object):
def __init__(self,
gpu_id=0,
weight_path=None,
visiual=None,
heatmap=False):
self.__num_class = cfg.COCO_DATA["NUM"]
self.__conf_threshold = cfg.VAL["CONF_THRESH"]
self.__nms_threshold = cfg.VAL["NMS_THRESH"]
self.__device = gpu.select_device(gpu_id)
self.__multi_scale_val = cfg.VAL["MULTI_SCALE_VAL"]
self.__flip_val = cfg.VAL["FLIP_VAL"]
self.__visiual = visiual
self.__eval = eval
self.__classes = cfg.COCO_DATA["CLASSES"]
self.__model = Build_Model().to(self.__device)
self.__load_model_weights(weight_path)
self.__evalter = Evaluator(self.__model, showatt=heatmap)
def __load_model_weights(self, weight_path):
print("loading weight file from : {}".format(weight_path))
weight = os.path.join(weight_path)
chkpt = torch.load(weight, map_location=self.__device)
self.__model.load_state_dict(chkpt)
print("loading weight file is done")
del chkpt
def reset(self):
path1 = os.path.join(cfg.DETECTION_PATH, "detection_result/")
path2 = os.path.join(cfg.DETECTION_PATH, "ShowAtt/")
for i in os.listdir(path1):
path_file = os.path.join(path1, i)
if os.path.isfile(path_file):
os.remove(path_file)
else:
for f in os.listdir(path_file):
path_file2 = os.path.join(path_file, f)
if os.path.isfile(path_file2):
os.remove(path_file2)
for i in os.listdir(path2):
path_file = os.path.join(path2, i)
if os.path.isfile(path_file):
os.remove(path_file)
else:
for f in os.listdir(path_file):
path_file2 = os.path.join(path_file, f)
if os.path.isfile(path_file2):
os.remove(path_file2)
def study(self):
# Parameter study
y = []
for i in [0.08, 0.07, 0.06]:
t = time.time()
evaluator = COCOAPIEvaluator(model_type='YOLOv3',
data_dir=cfg.DATA_PATH,
img_size=cfg.VAL["TEST_IMG_SIZE"],
confthre=i,
nmsthre=cfg.VAL["NMS_THRESH"])
_, r = evaluator.evaluate(self.__model)
y.append(
str(i) + str(' ') + str(r) + str(' ') +
str(time.time() - t, ))
np.savetxt('study.txt', y, fmt='%s') # y = np.loadtxt('study.txt')
def val(self):
global logger
logger.info("***********Start Evaluation****************")
start = time.time()
evaluator = COCOAPIEvaluator(model_type='YOLOv4',
data_dir=cfg.DATA_PATH,
img_size=cfg.VAL["TEST_IMG_SIZE"],
confthre=cfg.VAL["CONF_THRESH"],
nmsthre=cfg.VAL["NMS_THRESH"])
ap50_95, ap50 = evaluator.evaluate(self.__model)
logger.info('ap50_95:{}|ap50:{}'.format(ap50_95, ap50))
end = time.time()
logger.info(" ===val cost time:{:.4f}s".format(end - start))
def Inference(self):
global logger
# clear cache
self.reset()
logger.info("***********Start Inference****************")
imgs = os.listdir(self.__visiual)
logger.info("images path: {}".format(self.__visiual))
path = os.path.join(cfg.DETECTION_PATH, "detection_result")
logger.info("saved images at: {}".format(path))
inference_times = []
for v in imgs:
start_time = time.time()
path = os.path.join(self.__visiual, v)
img = cv2.imread(path)
assert img is not None
bboxes_prd = self.__evalter.get_bbox(img, v)
if bboxes_prd.shape[0] != 0:
boxes = bboxes_prd[..., :4]
class_inds = bboxes_prd[..., 5].astype(np.int32)
scores = bboxes_prd[..., 4]
visualize_boxes(image=img,
boxes=boxes,
labels=class_inds,
probs=scores,
class_labels=self.__classes)
path = os.path.join(cfg.DETECTION_PATH,
"detection_result/{}".format(v))
cv2.imwrite(path, img)
end_time = time.time()
inference_times.append(end_time - start_time)
inference_time = sum(inference_times) / len(inference_times)
fps = 1.0 / inference_time
logging.info("Inference_Time: {:.5f} s/image, FPS: {}".format(
inference_time, fps))
class Detection(object):
def __init__(self,
gpu_id=0,
weight_path=None,
video_path=None,
output_dir=None,
):
self.__num_class = cfg.VOC_DATA["NUM"]
self.__conf_threshold = cfg.VAL["CONF_THRESH"]
self.__nms_threshold = cfg.VAL["NMS_THRESH"]
self.__device = gpu.select_device(gpu_id)
self.__multi_scale_val = cfg.VAL["MULTI_SCALE_VAL"]
self.__flip_val = cfg.VAL["FLIP_VAL"]
self.__classes = cfg.VOC_DATA["CLASSES"]
self.__video_path = video_path
self.__output_dir = output_dir
self.__model = Build_Model().to(self.__device)
self.__load_model_weights(weight_path)
self.__evalter = Evaluator(self.__model, showatt=False)
def __load_model_weights(self, weight_path):
print("loading weight file from : {}".format(weight_path))
weight = os.path.join(weight_path)
chkpt = torch.load(weight, map_location=self.__device)
self.__model.load_state_dict(chkpt)
print("loading weight file is done")
del chkpt
def Video_detection(self):
import cv2
vid = cv2.VideoCapture(self.__video_path)
if not vid.isOpened():
raise IOError("Couldn't open webcam or video")
video_FourCC = int(vid.get(cv2.CAP_PROP_FOURCC))
video_fps = vid.get(cv2.CAP_PROP_FPS)
video_size = (int(vid.get(cv2.CAP_PROP_FRAME_WIDTH)),
int(vid.get(cv2.CAP_PROP_FRAME_HEIGHT)))
isOutput = True if self.__output_dir != "" else False
if isOutput:
print("!!! TYPE:", type(self.__output_dir), type(video_FourCC), type(video_fps), type(video_size))
out = cv2.VideoWriter(self.__output_dir, video_FourCC, video_fps, video_size)
accum_time = 0
curr_fps = 0
fps = "FPS: ??"
prev_time = timer()
while True:
return_value, frame = vid.read()
bboxes_prd = self.__evalter.get_bbox(frame)
if bboxes_prd.shape[0] != 0:
boxes = bboxes_prd[..., :4]
class_inds = bboxes_prd[..., 5].astype(np.int32)
scores = bboxes_prd[..., 4]
visualize_boxes(image=frame, boxes=boxes, labels=class_inds, probs=scores, class_labels=self.__classes)
curr_time = timer()
exec_time = curr_time - prev_time
prev_time = curr_time
accum_time = accum_time + exec_time
curr_fps = curr_fps + 1
if accum_time > 1:
accum_time = accum_time - 1
fps = "FPS: " + str(curr_fps)
curr_fps = 0
cv2.putText(frame, text=fps, org=(3, 15), fontFace=cv2.FONT_HERSHEY_SIMPLEX,
fontScale=0.50, color=(255, 0, 0), thickness=2)
cv2.namedWindow("result", cv2.WINDOW_NORMAL)
cv2.imshow("result", frame)
if isOutput:
out.write(frame)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
class Evaluation(object):
def __init__(self,
gpu_id='0',
weight_path=None,
img_size=544,
visiual=None,
eval=False,
epoch=None):
self.img_size = img_size
self.__num_class = cfg.DATA["NUM"]
self.__conf_threshold = cfg.VAL["CONF_THRESH"]
self.__nms_threshold = cfg.VAL["NMS_THRESH"]
# self.__device = select_device('0', batch_size=cfg.VAL["BATCH_SIZE"])
self.__device = gpu.select_device(gpu_id)
self.__multi_scale_val = cfg.VAL["MULTI_SCALE_VAL"]
self.__flip_val = cfg.VAL["FLIP_VAL"]
self.__visiual = visiual
self.__eval = eval
self.__classes = cfg.DATA["CLASSES"]
self.__model = Build_Model().to(self.__device)
self.__load_model_weights(weight_path)
self.__evalter = Evaluator(self.__model, epoch, showatt=False)
self.epoch = epoch
def __load_model_weights(self, weight_path):
print("loading weight file from : {}".format(weight_path))
weight = os.path.join(weight_path)
chkpt = torch.load(weight, map_location=self.__device)
self.__model.load_state_dict(chkpt)
print("loading weight file is done")
del chkpt
def val(self):
global writer, logger
if self.__eval:
logger.info("***********Start Evaluation****************")
start = time.time()
mAP = 0
with torch.no_grad():
_, _, APs = Evaluator(self.__model, self.epoch,
showatt=False).APs_voc(
self.__multi_scale_val,
self.__flip_val)
for i in APs:
logger.info("{} --> mAP : {}".format(i, APs[i]))
mAP += APs[i]
mAP = mAP / self.__num_class
logger.info('mAP:{}'.format(mAP))
writer.add_scalar('mAP', mAP, self.epoch)
end = time.time()
logger.info(" ===val cost time:{:.4f}s".format(end - start))
def detection(self):
global writer, logger
if self.__visiual:
imgs = os.listdir(self.__visiual)
logger.info("***********Start Detection****************")
start = time.clock()
for v in imgs:
path = os.path.join(self.__visiual, v)
logger.info("val images : {}".format(path))
img = cv2.imread(path)
assert img is not None
bboxes_prd = self.__evalter.get_bbox(img, v)
if bboxes_prd.shape[0] != 0:
boxes = bboxes_prd[..., :4]
class_inds = bboxes_prd[..., 5].astype(np.int32)
scores = bboxes_prd[..., 4]
visualize_boxes(image=img,
boxes=boxes,
labels=class_inds,
probs=scores,
class_labels=self.__classes)
path = os.path.join(cfg.PROJECT_PATH,
"detection_result/{}".format(v))
cv2.imwrite(path, img)
logger.info("saved images : {}".format(path))
end = time.clock()
times = end - start
FPS = len(imgs) / times
logger.info('FPS:{}'.format(FPS))
logger.info(" ===detection cost time:{:.4f}s".format(times))
class Detection(object):
def __init__(
self,
label_path,
weight_path=None,
output_dir=None,
):
self.__label_path = os.path.join("/data", label_path)
self.get_first_10_imgs()
self.__num_class = cfg.DATASET.NUM
# these should be set still
self.__conf_threshold = 0.25
self.__nms_threshold = 0.5
#######################################
self.__device = gpu.select_device()
self.__classes = cfg.DATASET.CLASSES
self.__output_dir = output_dir
self.__model = Build_Model().to(self.__device)
self.__load_model_weights(weight_path)
self.__evalter = Evaluator(self.__model, showatt=False)
def __load_model_weights(self, weight_path):
print("loading weight file from : {}".format(weight_path))
weight = os.path.join(weight_path)
chkpt = torch.load(weight, map_location=self.__device)
self.__model.load_state_dict(chkpt)
print("loading weight file is done")
del chkpt
def detect_demo(self):
import cv2
accum_time = 0
for image_path in tqdm(self.__file_dict.keys()):
# print(self.__file_dict[image_path])
frame = cv2.imread(image_path)
# prev_time = timer()
bboxes_prd = self.__evalter.get_bbox(frame)
if bboxes_prd.shape[0] != 0:
boxes = bboxes_prd[..., :4]
class_inds = bboxes_prd[..., 5].astype(np.int32)
scores = bboxes_prd[..., 4]
visualize_boxes(image=frame,
boxes=boxes,
labels=class_inds,
probs=scores,
class_labels=self.__classes)
visualize_boxes_and_labels_on_image_array(
image=frame,
boxes=self.__file_dict[image_path][:, :4],
classes=self.__file_dict[image_path][:, 4],
scores=None,
line_thickness=10,
category_index=self.__classes)
cv2.imwrite(
os.path.join(self.__output_dir, os.path.basename(image_path)),
frame)
print("FPS: {:.04f}".format(1000 * len(self.__file_dict.keys()) /
self.__evalter.inference_time))
def get_first_10_imgs(self):
fh = open(self.__label_path)
image_paths = {}
####################################
# for random select
####################################
# random.seed(1)
# lines = random.choices(fh.readlines(),k=10)
# for line in lines:
# line = line.rstrip().split()
# if len(line)>1:
# image_paths[os.path.join("/data",line[0])] = np.array([list(map(int,i.split(","))) for i in line[1:]])
# else:
# break
####################################
# for on demand plot
####################################
lines = fh.readlines()
imgs = [
"images/0021023.png", "images/0020485.png", "images/0021042.png",
"images/0021630.png", "images/0021729.png", "images/0021781.png"
]
for line in lines:
line = line.rstrip().split()
if line[0] in imgs:
image_paths[os.path.join("/data", line[0])] = np.array(
[list(map(int, i.split(","))) for i in line[1:]])
#####################################
self.__file_dict = image_paths
class Detection(object):
def __init__(self,
label_path,
weight_path=None,
output_dir=None,
):
self.__label_path = os.path.join("/data",label_path)
self.get_first_10_imgs()
self.__num_class = cfg.DATASET.NUM
self.__conf_threshold = cfg.VAL.CONF_THRESH
self.__nms_threshold = cfg.VAL.NMS_THRESH
self.__device = gpu.select_device()
self.__classes = cfg.DATASET.CLASSES
# self.__video_path = video_path
self.__output_dir = output_dir
self.__model = Build_Model().to(self.__device)
self.__load_model_weights(weight_path)
self.__evalter = Evaluator(self.__model, showatt=False)
def __load_model_weights(self, weight_path):
print("loading weight file from : {}".format(weight_path))
weight = os.path.join(weight_path)
chkpt = torch.load(weight, map_location=self.__device)
self.__model.load_state_dict(chkpt)
print("loading weight file is done")
del chkpt
def detect_demo(self):
import cv2
accum_time = 0
for image_path in self.__file_list:
frame = cv2.imread(image_path)
# prev_time = timer()
bboxes_prd = self.__evalter.get_bbox(frame)
if bboxes_prd.shape[0] != 0:
boxes = bboxes_prd[..., :4]
class_inds = bboxes_prd[..., 5].astype(np.int32)
scores = bboxes_prd[..., 4]
visualize_boxes(image=frame, boxes=boxes, labels=class_inds, probs=scores, class_labels=self.__classes)
# curr_time = timer()
# accum_time += curr_time - prev_time
cv2.imwrite(os.path.join(self.__output_dir,os.path.basename(image_path)),frame)
print("FPS: {:.04f}".format(len(self.__file_list)/self.__evalter.inference_time))
def get_first_10_imgs(self):
fh = open(self.__label_path)
image_paths = []
for line in fh.readlines():
line = line.rstrip().split()
if len(image_paths) < 100:
if len(line) > 1:
image_paths.append(os.path.join("/data",line[0]))
else:
break
print(image_paths)
self.__file_list = image_paths
def transform_to_onnx(weight_file, batch_size, n_classes, IN_IMAGE_H,
IN_IMAGE_W):
device = gpu.select_device(id=0)
model = Build_Model().to(device)
#model = Build_Model(weight_path=weight_file, resume=False).to(device)
pretrained_dict = torch.load(weight_file,
map_location=device) #torch.device('cuda')
model.load_state_dict(pretrained_dict)
evaluator = Evaluator(model, showatt=False)
input_names = ["input"]
output_names = ['boxes', 'confs']
dynamic = False
if batch_size <= 0:
dynamic = True
if dynamic:
x = torch.randn((1, 3, IN_IMAGE_H, IN_IMAGE_W), requires_grad=True)
onnx_file_name = "yolov4_-1_3_{}_{}_dynamic.onnx".format(
IN_IMAGE_H, IN_IMAGE_W)
dynamic_axes = {
"input": {
0: "batch_size"
},
"boxes": {
0: "batch_size"
},
"confs": {
0: "batch_size"
}
}
# Export the model
print('Export the onnx model ...')
torch.onnx.export(model,
x,
onnx_file_name,
export_params=True,
opset_version=11,
do_constant_folding=True,
input_names=input_names,
output_names=output_names,
dynamic_axes=dynamic_axes)
print('Onnx model exporting done')
return onnx_file_name
else:
x = torch.randn((batch_size, 3, IN_IMAGE_H, IN_IMAGE_W),
requires_grad=True)
x = x.to(device)
onnx_file_name = "yolov4_{}_3_{}_{}_static.onnx".format(
batch_size, IN_IMAGE_H, IN_IMAGE_W)
# Export the model
print('Export the onnx model ...')
torch.onnx.export(model,
x,
onnx_file_name,
export_params=True,
opset_version=11,
do_constant_folding=True,
input_names=input_names,
output_names=output_names,
dynamic_axes=None)
print('Onnx model exporting done')
return onnx_file_name, evaluator
from utils.torch_utils import *
# GPU device
# Check GPU
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print("Device: ", device)
model = Build_Model().to(device)
w_path = 'model.pt'
# chkpt = torch.load(os.path.join(w_path))
chkpt = torch.load(os.path.join(w_path), map_location=torch.device('cpu'))
model.load_state_dict(chkpt['model'])
classes = ['1', '2', '3', '4', '5', '6']
test_path = './data/test/'
pred_path = './data/pred/'
imgs_path = os.listdir(test_path)
n = len(imgs_path)
ratio = 0.1
for i, img_path in enumerate(imgs_path):
path = os.path.join(test_path, img_path)
print(path)
img = cv2.imread(path)
# resized_img = cv2.resize(img, None, fx=ratio, fy=ratio, interpolation=cv2.INTER_AREA)
# resized_img = cv2.resize(img, (832, 832), interpolation=cv2.INTER_AREA)
class LAMR_Tester(object):
id_map = ["pedestrian"]
def __init__(self, log_dir, test_images):
init_seeds(0)
self.device = gpu.select_device()
self.log_dir = log_dir
self.yolov4 = Build_Model(weight_path=None, resume=False)
self.yolov4 = self.yolov4.to(self.device)
self.dataset = Naive_Test_Dataset(test_images)
self.dataloader = torch.utils.data.DataLoader(
self.dataset,
batch_size=cfg.VAL.BATCH_SIZE,
shuffle=False,
pin_memory=True,
num_workers=cfg.VAL.NUMBER_WORKERS)
self.__load_best_weights()
def __load_best_weights(self):
best_weight = os.path.join(log_dir, "checkpoints", "best.pt")
chkpt = torch.load(best_weight, map_location=self.device)
self.yolov4.load_state_dict(chkpt)
del chkpt
def test(self):
logger.info(self.yolov4)
self.yolov4.eval()
Tensor = torch.cuda.FloatTensor if torch.cuda.is_available(
) else torch.FloatTensor
results_path = os.path.join("/data", "mock_detections",
os.path.basename(self.log_dir), "day",
"val")
if not os.path.exists(results_path):
os.makedirs(results_path)
for i, (img_path, img, info_img) in tqdm(enumerate(self.dataloader),
desc="Test to ECP... ",
unit="imgs",
total=len(self.dataloader)):
info_img = [float(info) for info in info_img]
data_dict = {
"tags": [],
"imageheight": int(info_img[0]),
"imagewidth": int(info_img[1]),
"children": [],
"identity": "frame"
}
city_name = os.path.basename(os.path.dirname(img_path[0]))
os.makedirs(os.path.join(results_path, city_name), exist_ok=True)
result_json_path = os.path.join(
results_path,
os.path.basename(img_path[0]).replace("png", "json"))
with torch.no_grad():
img = Variable(img.type(Tensor))
_, outputs = self.yolov4(img)
outputs = outputs.unsqueeze(0)
outputs = postprocess(outputs, len(cfg.DATASET.CLASSES), 0.01,
cfg.VAL.NMS_THRESH)
if outputs[0] is None:
with open(result_json_path, "w") as json_fh:
json.dump(data_dict, json_fh, indent=4)
continue
outputs = outputs[0].cpu().data
for output in outputs:
x1 = float(output[0])
y1 = float(output[1])
x2 = float(output[2])
y2 = float(output[3])
box = yolobox2label((y1, x1, y2, x2), info_img)
data_dict["children"].append({
"tags": [],
"children": [],
"identity":
self.id_map[int(output[6])],
"score":
float(output[4]) * float(output[5]),
"y0":
float(box[0]),
"x0":
float(box[1]),
"y1":
float(box[2]),
"x1":
float(box[3])
}) # ECP Formats
with open(result_json_path, "w") as json_fh:
json.dump(data_dict, json_fh, indent=4)
class Evaluation(object):
def __init__(
self,
gpu_id=0,
model1_path=None,
model2_path=None,
data_dir=None,
# result_dir=None,
mnist=False,
):
self.__num_class = cfg.Customer_DATA["NUM"]
self.__conf_threshold = cfg.VAL["CONF_THRESH"]
self.__nms_threshold = cfg.VAL["NMS_THRESH"]
self.__device = gpu.select_device(gpu_id)
self.__multi_scale_val = cfg.VAL["MULTI_SCALE_VAL"]
self.__flip_val = cfg.VAL["FLIP_VAL"]
self.__data_dir = data_dir
print(self.__data_dir)
self.__classes = cfg.Customer_DATA["CLASSES"]
self.__mnist = mnist
self.__model1 = Build_Model().to(self.__device)
if mnist:
self.__model2 = torch.load(model2_path).double().cuda()
else:
self.__model2 = torch.load(model2_path).cuda()
self.__load_model_weights(model1_path)
self.__evalter = Evaluator(self.__model1, showatt=False)
# self.__result_dir = result_dir
def __load_model_weights(self, weight_path):
print("loading weight file from : {}".format(weight_path))
weight = os.path.join(weight_path)
chkpt = torch.load(weight, map_location=self.__device)
self.__model1.load_state_dict(chkpt)
# print("loading weight file is done")
del chkpt
def detection(self):
if os.path.isdir(self.__data_dir):
imgs = os.listdir(self.__data_dir)
print("***********Start Detection****************")
for v in imgs:
path = self.__data_dir + "/" + v
print("val images : {}".format(path))
img = cv2.imread(path)
assert img is not None
bboxes_prd = self.__evalter.get_bbox(img, v)
points, max_l = get_points(img, bboxes_prd=bboxes_prd, model=self.__model2, mnist=self.__mnist)
if points is None:
return
# print(points)
matrix_calculator = calculator(self.__classes)
matrix_calculator.get_from_points(points, max_l)
print(matrix_calculator())
else:
raise NotImplementedError("The data directory is not exist!")