def __init__(self, bgr=True, gpu_device=0, **kwargs):
self.__dict__.update(self._defaults) # set up default values
# for portability between keras-yolo3/yolo.py and this
if 'model_path' in kwargs:
kwargs['weights'] = kwargs['model_path']
if 'score' in kwargs:
kwargs['thresh'] = kwargs['score']
self.__dict__.update(kwargs) # update with user overrides
self.class_names = self._get_class()
# self.model = Darknet(self.config)
# self.model.load_weights(self.weights)
self.model = Yolov4(n_classes=len(self.class_names), inference=True)
checkpoint = torch.load(self.weights, map_location=torch.device('cpu'))
# checkpoint = self._rename_checkpoint(checkpoint)
self.model.load_state_dict(checkpoint)
self.device = gpu_device
self.model.cuda(self.device)
self.model.eval()
self.bgr = bgr
if self.half:
self.model.half()
# warm up
self._detect([np.zeros((10, 10, 3), dtype=np.uint8)])
# self._detect([np.zeros((10,10,3), dtype=np.uint8), np.zeros((10,10,3), dtype=np.uint8), np.zeros((10,10,3), dtype=np.uint8), np.zeros((10,10,3), dtype=np.uint8)])
print('Warmed up!')
def transform_to_onnx(weight_file, batch_size, n_classes, IN_IMAGE_H,
IN_IMAGE_W):
model = Yolov4(n_classes=n_classes, inference=True)
pretrained_dict = torch.load(weight_file,
map_location=torch.device('cuda'))
model.load_state_dict(pretrained_dict)
x = torch.randn((batch_size, 3, IN_IMAGE_H, IN_IMAGE_W),
requires_grad=True) # .cuda()
onnx_file_name = "yolov4_{}_3_{}_{}.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'],
output_names=['boxes', 'confs'],
dynamic_axes=None)
print('Onnx model exporting done')
return onnx_file_name
def detect_cv2(namesfile,anchors,weightfile,imgfile):
import cv2
class_names = load_class_names(namesfile)
n_classes=len(class_names)
m = Yolov4(anchors=anchors,yolov4conv137weight=None, n_classes=n_classes, inference=True)
pretrained_dict = torch.load(weightfile, map_location=torch.device('cuda'))
m.load_state_dict(pretrained_dict)
print('Loading weights from %s... Done!' % (weightfile))
if use_cuda:
m.cuda()
img = cv2.imread(imgfile)
sized = cv2.resize(img, (416, 416))
sized = cv2.cvtColor(sized, cv2.COLOR_BGR2RGB)
for i in range(2):
start = time.time()
boxes = do_detect(m, sized, 0.4, 0.6, use_cuda)
finish = time.time()
if i == 1:
print('%s: Predicted in %f seconds.' % (imgfile, (finish - start)))
plot_boxes_cv2(img, boxes[0], savename='predictions.jpg', class_names=class_names)
def load_model_weight(model_path, num_classes):
model = Yolov4(n_classes=num_classes)
print("------------loading model")
model.load_state_dict(
torch.load(model_path, map_location=torch.device('cuda')))
print("------------done loading")
return model
def __init__(self, weightfile, use_cuda=True):
self.use_cuda = use_cuda
self.model = Yolov4(yolov4conv137weight=None, n_classes=80, inference=True)
pretrained_dict = torch.load(
weightfile, map_location=torch.device("cuda" if use_cuda else "cpu")
)
self.model.load_state_dict(pretrained_dict)
if self.use_cuda:
self.model.cuda()
def __init__(self, weightfile, use_cuda=True):
if use_cuda and not torch.cuda.is_available():
raise Exception(
"Selected use_cuda=True, but cuda is not available to Pytorch"
)
self.use_cuda = use_cuda
self.model = Yolov4(yolov4conv137weight=None, n_classes=80, inference=True)
pretrained_dict = torch.load(
weightfile, map_location=torch.device("cuda" if use_cuda else "cpu")
)
self.model.load_state_dict(pretrained_dict)
if self.use_cuda:
self.model.cuda()
def transform_to_onnx(weight_file, batch_size, n_classes, IN_IMAGE_H, IN_IMAGE_W):
model = Yolov4(n_classes=n_classes, inference=True)
pretrained_dict = torch.load(weight_file, map_location=torch.device('cuda'))
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,
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
def transform_to_onnx(weight_file, onnx_file_name, batch_size, n_classes,
input_h, input_w):
model = Yolov4(n_classes=n_classes, inference=True).to(device)
pretrained_dict = torch.load(weight_file)
model.load_state_dict(pretrained_dict)
input_names = ["input"]
output_names = ["boxes", "confs"]
x = torch.randn(batch_size, 3, input_h, input_w).to(device)
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')
def __init__(self,weightfile,num_classes=80,width=416,height=416,use_cuda = True,is_half=True):
self.num_classes = num_classes
self.model = Yolov4(yolov4conv137weight=None, n_classes=self.num_classes, inference=True)
pretrained_dict = torch.load(weightfile, map_location=torch.device('cuda'))
self.model.load_state_dict(pretrained_dict)
self.use_cuda = use_cuda
self.is_half=is_half
self.width=width
self.height=height
if self.use_cuda:
if self.is_half:
self.model=self.model.cuda().half()
else:
self.model=self.model.cuda()
if self.num_classes == 20:
namesfile = 'data/voc.names'
elif self.num_classes == 80:
namesfile = 'data/coco.names'
else:
namesfile = 'data/x.names'
self.class_names = load_class_names(namesfile)
print("initialize_a_pytorch")
def detect_cv2_camera(namesfile,anchors,weightfile,size=416):
import cv2
anchors=np.array(anchors)*size
anchors=anchors.tolist()
class_names = load_class_names(namesfile)
n_classes=len(class_names)
m = Yolov4(anchors=anchors,yolov4conv137weight=None, n_classes=n_classes, inference=True)
pretrained_dict = torch.load(weightfile, map_location=torch.device('cuda'))
m.load_state_dict(pretrained_dict)
print('Loading weights from %s... Done!' % (weightfile))
if use_cuda:
m.cuda()
# cap = cv2.VideoCapture(0)
cap = cv2.VideoCapture("../dataset/sample1.mov")
cap.set(3, 1280)
cap.set(4, 720)
print("Starting the YOLO loop...")
while True:
ret, img = cap.read()
if(ret==False): continue
sized = cv2.resize(img, (size,size))
sized = cv2.cvtColor(sized, cv2.COLOR_BGR2RGB)
start = time.time()
boxes = do_detect(m, sized, 0.2, 0.6, use_cuda)
finish = time.time()
print('Predicted in %f seconds.' % (finish - start))
result_img = plot_boxes_cv2(img, boxes[0], savename=None, class_names=class_names)
cv2.imshow('Yolo demo', result_img)
cv2.waitKey(1)
cap.release()
def init_net(pretrained, n_classes):
# m = Darknet(cfgfile)
# model = Yolov4(cfg.pretrained, n_classes=cfg.classes)
m = Yolov4(n_classes=n_classes)
# model = Yolov4(n_classes=n_classes)
pretrained_dict = torch.load(pretrained, map_location=torch.device('cuda'))
from collections import OrderedDict
new_state_dict = OrderedDict()
for key, value in pretrained_dict.items():
new_key = key.replace('module.', '')
new_state_dict[new_key] = value
m.load_state_dict(new_state_dict)
# m.print_network()
# m.load_weights(weightfile)
print('Loading weights from %s... Done!' % (pretrained))
if use_cuda:
m.cuda()
return m
def get_image_objects(image: Image):
models_path = os.path.join(os.getcwd(), 'modules', 'pytorch-YOLOv4')
if models_path not in sys.path:
sys.path.insert(1, models_path)
from models import Yolov4
from tool.utils import load_class_names
device = 'cpu'
model = Yolov4(yolov4conv137weight=None, n_classes=80, inference=True)
model.load_state_dict(
torch.load(os.path.join(os.getcwd(), 'models', 'yolov4', 'yolov4.pth'),
map_location=torch.device(device)))
model.eval()
model = model.to(torch.device(device))
image_tensor = ToTensor()(image)
image_tensor = image_tensor.unsqueeze(0)
image_tensor = image_tensor.to(torch.device(device))
with torch.no_grad():
output = model(image_tensor)
conf_thresh = 0.4
nms_thresh = 0.6
boxes = post_process_boxes(conf_thresh, nms_thresh, output)
class_names = load_class_names(
os.path.join(models_path, 'data', 'coco.names'))
return list(
map(
lambda obj: {
"x1": obj[0],
"y1": obj[1],
"x2": obj[2],
"y2": obj[3],
"class_name": class_names[obj[6]]
}, boxes[0]))
def detect_skimage(namesfile,anchors,weightfile, imgfile):
from skimage import io
from skimage.transform import resize
class_names = load_class_names(namesfile)
n_classes=len(class_names)
m = Yolov4(anchors=anchors,yolov4conv137weight=None, n_classes=n_classes, inference=True)
pretrained_dict = torch.load(weightfile, map_location=torch.device('cuda'))
m.load_state_dict(pretrained_dict)
print('Loading weights from %s... Done!' % (weightfile))
if use_cuda:
m.cuda()
img = io.imread(imgfile)
sized = resize(img, (m.width, m.height)) * 255
for i in range(2):
start = time.time()
boxes = do_detect(m, sized, 0.4, 0.4, use_cuda)
finish = time.time()
if i == 1:
print('%s: Predicted in %f seconds.' % (imgfile, (finish - start)))
plot_boxes_cv2(img, boxes, savename='predictions.jpg', class_names=class_names)
def train_yolov4(cfg):
logging = init_logger(log_dir=cfg.TRAIN_TENSORBOARD_DIR)
os.environ["CUDA_VISIBLE_DEVICES"] = cfg.gpu
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
logging.info(f'Using device {device}')
if cfg.use_darknet_cfg:
model = Darknet(cfg.cfgfile)
else:
model = Yolov4(cfg.pretrained, n_classes=cfg.classes)
#TODO, load checkpoints
'''
WORK_FOLDER = '/mnt/bos/modules/perception/emergency_detection'
weightfile = os.path.join(WORK_FOLDER, 'checkpoints/Yolov4_epoch291.pth')
pretrained_dict = torch.load(weightfile, map_location=torch.device('cuda'))
model.load_state_dict(pretrained_dict)
'''
if torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model)
model.to(device=device)
try:
train(model=model,
config=cfg,
epochs=cfg.TRAIN_EPOCHS,
device=device, )
except KeyboardInterrupt:
torch.save(model.state_dict(), 'INTERRUPTED.pth')
logging.info('Saved interrupt')
try:
sys.exit(0)
except SystemExit:
os._exit(0)
return logging
def _get_date_str():
now = datetime.datetime.now()
return now.strftime('%Y-%m-%d_%H-%M')
if __name__ == "__main__":
logging = init_logger(log_dir='log')
cfg = Cfg
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
logging.info(f'Using device {device}')
model = Yolov4("yolov4.conv.137.pth", n_classes=cfg.classes)
model.to(device=device)
try:
train(model=model,
config=cfg,
epochs=cfg.TRAIN_EPOCHS,
device=device, )
except KeyboardInterrupt:
torch.save(model.state_dict(), 'INTERRUPTED.pth')
logging.info('Saved interrupt')
try:
sys.exit(0)
except SystemExit:
os._exit(0)
def plot_lines(image, boxes):
width = img.shape[1]
height = img.shape[0]
angled = np.sqrt(width**2 + height**2)
for box in boxes:
x = (box[2] + box[0]) / 2 * width
y = (box[3] + box[1]) / 2 * height
print()
return image
if __name__ == '__main__':
model = Yolov4(None, n_classes=80, inference=True)
state = torch.load('16.pth')
model.load_state_dict(state['model'])
del state
model.eval()
rawimg = cv2.imread(
'/media/palm/data/coco/images/val2017/000000289343.jpg')
rawimg = cv2.cvtColor(rawimg, cv2.COLOR_BGR2RGB)
rawimg, bboxes = resize_image(rawimg, None, 608, 608)
img = rawimg.copy().transpose(2, 0, 1)
inputs = torch.from_numpy(np.expand_dims(img,
0).astype('float32')).div(255.0)
output = model(inputs)
boxes = post_processing(img, 0.4, 0.4, output)
is_first_frame = True
num_frame = 0
counter_detection = 0
person_id = 0
prev_id = 0
prev_coords = None
people_actions = dict()
person_inside = dict()
person_outside = dict()
people_actions['in'] = 0
people_actions['out'] = 0
check_gpu()
model = Yolov4(weight_path='yolo4_config/yolov4.weights',
class_name_path='yolo4_config/coco_classes.txt')
videofile = cv2.VideoCapture(video_path)
fps = videofile.get(cv2.CAP_PROP_FPS)
fourcc = cv2.VideoWriter_fourcc(*'MP4V')
output_video = cv2.VideoWriter(output_name, fourcc, fps, (640, 480))
rr, first_frame = videofile.read()
if initialize_door_by_yourself:
door_array = select_object(first_frame)[0]
else:
door_array = [361, 20, 507, 352]
while rr:
ret, frame = videofile.read()
if not ret:
logging.getLogger('').addHandler(console)
return logging
if __name__ == "__main__":
logging = init_logger(log_dir='log')
cfg = get_args(**Cfg)
os.environ["CUDA_VISIBLE_DEVICES"] = cfg.gpu
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
logging.info(f'Using device {device}')
if cfg.use_darknet_cfg:
model = Darknet(cfg.cfgfile)
else:
model = Yolov4(cfg.weight, n_classes=cfg.classes)
# pretrained_dict = torch.load(cfg.load)
# model_dict = model.state_dict()
# pretrained_dict = {k: v for k, v in pretrained_dict.items() if
# k in model_dict} # pretrained_dict只保留了model_dict中存在的键。为什么直接load它会报错。要先给model_dict更新。
# model_dict.update(pretrained_dict)
# model.load_state_dict(model_dict)
total_params = sum(p.numel() for p in model.parameters())
print(f'{total_params:,} total parameters.')
total_trainable_params = sum(p.numel() for p in model.parameters()
if p.requires_grad)
print(f'{total_trainable_params:,} training parameters.')
if torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model)
def main():
n_classes = 18
weightfile = r'D:\DeepBlueProjects\chem-lab\pytorch-YOLOv4-master\checkpoints\Yolov4_epoch6.pth'
imgfile = r'D:\Data\CMS01_single-end\val\JPEGImages\frontfront_0518.jpg'
base_dir = r'D:\Data\chem-yolov4\eval-dataset\top'
gt_path = os.path.join(base_dir, 'gt.json')
name_id_path = os.path.join(base_dir, 'name_id.json')
with open(gt_path, 'r') as f:
gt_dict = json.load(f)
with open(name_id_path, 'r') as f:
name_id_dict = json.load(f)
input_size = (960, 960)
model = Yolov4(yolov4conv137weight=None,
n_classes=n_classes,
inference=True)
pretrained_dict = torch.load(weightfile, map_location=torch.device('cuda'))
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in pretrained_dict.items():
name = k[7:] # remove `module.`，表面从第7个key值字符取到最后一个字符，正好去掉了module.
new_state_dict[name] = v # 新字典的key值对应的value为一一对应的值。
model.load_state_dict(pretrained_dict)
use_cuda = True
if use_cuda:
model.cuda()
data_txt = os.path.join(base_dir, 'data.txt')
save_dir = os.path.join(base_dir, 'JPEGImages_pred')
result_dir = os.path.join(base_dir, 'result_txt')
with open(data_txt, 'r') as f:
imgfiles = f.readlines()
box_list = []
for imgfile in imgfiles:
img = cv2.imread(imgfile.strip('\n'))
img_h, img_w, _ = img.shape
img_name = imgfile.split('\\')[-1].strip('\n')
img_id = name_id_dict[img_name]
result_txt = os.path.join(result_dir, img_name[:-4] + '.txt')
result_f = open(result_txt, 'w')
# Inference input size is 416*416 does not mean training size is the same
# Training size could be 608*608 or even other sizes
# Optional inference sizes:
# Hight in {320, 416, 512, 608, ... 320 + 96 * n}
# Width in {320, 416, 512, 608, ... 320 + 96 * m}
sized = cv2.resize(img, (input_size[1], input_size[0]))
sized = cv2.cvtColor(sized, cv2.COLOR_BGR2RGB)
# for i in range(2): # This 'for' loop is for speed check
# Because the first iteration is usually longer
boxes = do_detect(model, sized, 0.01, 0.3, use_cuda)
for box in boxes[0]:
x1 = int((box[0] - box[2] / 2.0) * img_w)
y1 = int((box[1] - box[3] / 2.0) * img_h)
x2 = int((box[0] + box[2] / 2.0) * img_w)
y2 = int((box[1] + box[3] / 2.0) * img_h)
w = x2 - x1
h = y2 - y1
if len(box) >= 7:
cls_conf = box[5]
cls_id = box[6]
box_list.append({
"image_id": img_id,
"category_id": int(cls_id),
"bbox": [x1, y1, w, h],
"score": float(cls_conf)
})
string = ','.join([
str(cls_id),
str(x1),
str(y1),
str(x2),
str(y2),
str(cls_conf)
]) + '\n'
result_f.write(string)
if cls_conf > 0.3:
cv2.rectangle(img, (int(x1), int(y1)), (int(x2), int(y2)),
(255, 0, 255), 1)
cv2.putText(img, str(cls_id), (int(x1 + 10), int(y1 + 20)),
cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 255),
1)
cv2.putText(img, str(round(cls_conf, 3)),
(int(x1 + 30), int(y1 + 20)),
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (20, 240, 0), 1)
else:
print('????')
result_f.close()
namesfile = 'data/chem.names'
class_names = load_class_names(namesfile)
save_name = os.path.join(save_dir, img_name)
plot_boxes_cv2(img, boxes[0], save_name, class_names)
# cv2.imshow('result', img)
# cv2.waitKey(0)
# cv2.destroyAllWindows()
info, map_iou0_5 = get_coco_mAP(gt_dict, box_list)
# print("---base_eval---epoch%d"%real_epoch)
print(info)
n_classes = int(sys.argv[1])
namesfile = sys.argv[2]
dataset_dir = sys.argv[3]
width = int(sys.argv[4])
height = int(sys.argv[5])
weightfile = sys.argv[6]
eval_model = True
else:
print('Usage: ')
print(
' python models.py [num_classes] [names_file] [datasetdir] [width] [height] ([pretrained])'
)
if eval_model:
model = Yolov4(yolov4conv137weight=None,
n_classes=n_classes,
inference=True)
pretrained_dict = torch.load(weightfile,
map_location=torch.device('cuda'))
model.load_state_dict(pretrained_dict)
use_device = 1 # -1 for cpu, 0,1,2,... for gpu 0,1,2,...
if use_device >= 0:
# model.cuda()
model.to("cuda:{}".format(use_device))
if not output_video:
window_name = "gt vs. pred" if eval_model else "gt"
cv2.namedWindow(window_name, 0)
cv2.resizeWindow(window_name, 1000, 500)
# -*- coding: utf-8 -*-
"""
Created on Wed Feb 3 16:24:19 2021
@author: vinic_000
"""
from models import Yolov4
model = Yolov4(weight_path='yolov4.weights',
class_name_path='class_names/coco_classes.txt')
model.predict('img/pessoas.png')
model.predict('img/vaca.png')
model.predict('img/cavalo.png')
model.predict('img/escola.jpg')
model.predict('img/kite.jpg')
def _get_date_str():
now = datetime.datetime.now()
return now.strftime('%Y-%m-%d_%H-%M')
if __name__ == "__main__":
logging = init_logger(log_dir='log')
cfg = get_args(**Cfg)
os.environ["CUDA_VISIBLE_DEVICES"] = cfg.gpu
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
logging.info(f'Using device {device}')
if cfg.use_darknet_cfg:
model = Darknet(cfg.cfgfile)
else:
model = Yolov4(cfg.pretrained, n_classes=cfg.classes)
if torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model)
model.to(device=device)
try:
train(
model=model,
config=cfg,
epochs=cfg.TRAIN_EPOCHS,
device=device,
)
except KeyboardInterrupt:
torch.save(model.state_dict(), 'INTERRUPTED.pth')
logging.info('Saved interrupt')
def train(model,
device,
config,
epochs=5,
batch_size=1,
save_cp=True,
log_step=20,
img_scale=0.5,
freeze_backbone=False):
""" Train the YOLOv4 network with given configurations """
train_dataset = Yolo_dataset(config.train_label, config, train=True)
val_dataset = Yolo_dataset(config.val_label, config, train=False)
n_train = len(train_dataset)
n_val = len(val_dataset)
train_loader = DataLoader(train_dataset,
batch_size=config.batch // config.subdivisions,
shuffle=True,
num_workers=8,
pin_memory=True,
drop_last=True,
collate_fn=collate)
val_loader = DataLoader(val_dataset,
batch_size=config.batch // config.subdivisions,
shuffle=True,
num_workers=8,
pin_memory=True,
drop_last=True,
collate_fn=val_collate)
global_step = 0
logging.info(f'''Starting training:
Epochs: {epochs}
Batch size: {config.batch}
Subdivisions: {config.subdivisions}
Learning rate: {config.learning_rate}
Training size: {n_train}
Validation size: {n_val}
Checkpoints: {save_cp}
Device: {device.type}
Images size: {config.width}
Optimizer: {config.TRAIN_OPTIMIZER}
Dataset classes: {config.classes}
Train label path:{config.train_label}
Pretrained: {config.pretrained}
''')
# Learning rate setup
def burnin_schedule(i):
if i < config.burn_in:
factor = pow(i / config.burn_in, 4)
elif i < config.steps[0]:
factor = 1.0
elif i < config.steps[1]:
factor = 0.1
else:
factor = 0.01
return factor
# Optimiser alternatives
if config.TRAIN_OPTIMIZER.lower() == 'adam':
optimizer = optim.Adam(
model.parameters(),
lr=config.learning_rate / config.batch,
betas=(0.9, 0.999),
eps=1e-08,
)
elif config.TRAIN_OPTIMIZER.lower() == 'sgd':
optimizer = optim.SGD(
params=model.parameters(),
lr=config.learning_rate / config.batch,
momentum=config.momentum,
weight_decay=config.decay,
)
scheduler = optim.lr_scheduler.LambdaLR(optimizer, burnin_schedule)
criterion = Yolo_loss(device=device,
batch=config.batch // config.subdivisions,
n_classes=config.classes)
save_prefix = 'Yolov4_epoch'
saved_models = deque()
# Set model to training mode
model.train()
for epoch in range(epochs):
epoch_loss = 0
epoch_step = 0
# Freeze backbone first epoch
if freeze_backbone and epoch < 2:
for name, p in model.named_parameters():
# Freeze everything but the head
if not 'head' in name.split('.')[0]:
p.requires_grad = False if (epoch == 0) else True
with tqdm(total=n_train,
desc=f'Epoch {epoch + 1}/{epochs}',
unit='img',
ncols=200) as progress_bar:
for batch in train_loader:
global_step += 1
epoch_step += 1
images = batch[0]
bboxes = batch[1]
images = images.to(device=device, dtype=torch.float32)
bboxes = bboxes.to(device=device)
bboxes_pred = model(images)
loss, loss_xy, loss_wh, loss_obj, loss_cls, loss_l2 = criterion(
bboxes_pred, bboxes)
loss.backward()
epoch_loss += loss.item()
if global_step % config.subdivisions == 0:
optimizer.step()
scheduler.step()
model.zero_grad()
if global_step % (log_step * config.subdivisions) == 0:
progress_bar.set_postfix(
**{
'loss (batch)': loss.item(),
'loss_xy': loss_xy.item(),
'loss_wh': loss_wh.item(),
'loss_obj': loss_obj.item(),
'loss_cls': loss_cls.item(),
'loss_l2': loss_l2.item(),
'lr': scheduler.get_lr()[0] * config.batch
})
logging.debug(
'Train step_{}: loss : {}, loss xy : {}, loss wh : {},'
' loss obj : {}, loss cls : {}, loss l2 : {}, lr : {}'.
format(global_step, loss.item(), loss_xy.item(),
loss_wh.item(), loss_obj.item(),
loss_cls.item(), loss_l2.item(),
scheduler.get_lr()[0] * config.batch))
progress_bar.update(images.shape[0])
eval_model = Yolov4(cfg.pretrained,
n_classes=cfg.classes,
inference=True)
if torch.cuda.device_count() > 1:
eval_model.load_state_dict(model.module.state_dict())
else:
eval_model.load_state_dict(model.state_dict())
eval_model.to(device)
"""
evaluator = evaluate(eval_model, val_loader, config, device)
del eval_model
stats = evaluator.coco_eval['bbox'].stats
writer.add_scalar('train/AP', stats[0], global_step)
writer.add_scalar('train/AP50', stats[1], global_step)
writer.add_scalar('train/AP75', stats[2], global_step)
writer.add_scalar('train/AP_small', stats[3], global_step)
writer.add_scalar('train/AP_medium', stats[4], global_step)
writer.add_scalar('train/AP_large', stats[5], global_step)
writer.add_scalar('train/AR1', stats[6], global_step)
writer.add_scalar('train/AR10', stats[7], global_step)
writer.add_scalar('train/AR100', stats[8], global_step)
writer.add_scalar('train/AR_small', stats[9], global_step)
writer.add_scalar('train/AR_medium', stats[10], global_step)
writer.add_scalar('train/AR_large', stats[11], global_step)
"""
# Save model to file
if save_cp:
try:
os.makedirs(config.checkpoints, exist_ok=True)
logging.info('Created checkpoint directory')
except OSError:
pass
save_path = os.path.join(config.checkpoints,
f'{save_prefix}{epoch + 1}.pth')
torch.save(model.state_dict(), save_path)
logging.info(f'Checkpoint {epoch + 1} saved !')
saved_models.append(save_path)
if len(saved_models) > config.keep_checkpoint_max > 0:
model_to_remove = saved_models.popleft()
try:
os.remove(model_to_remove)
except:
logging.info(f'failed to remove {model_to_remove}')
import tvm
from tvm import relay
import numpy as np
from tvm.contrib.download import download_testdata
import torch
from tvm.contrib import graph_runtime
from PIL import Image
from torchvision import transforms
from models import Yolov4
model = Yolov4(inference=True)
model.eval()
input_shape = [1, 3, 512, 512]
input_data = torch.randn(input_shape)
scripted_model = torch.jit.trace(model, input_data).eval()
img_url = "https://github.com/dmlc/mxnet.js/blob/master/data/cat.png?raw=true"
img_path = download_testdata(img_url, "cat.png", module="data")
img = Image.open(img_path).resize((512, 512))
my_preprocess = transforms.Compose(
[
transforms.Resize(512),
transforms.CenterCrop(512),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
]
)
img = my_preprocess(img)
img = np.expand_dims(img, 0)
input_name = "input0"
def evaluate_all(device=None):
"""
"""
device = device or torch.device("cuda")
print(f"device = {device}")
val_dataset = ACNE04(label_path=Cfg.val_label, cfg=Cfg, train=False)
val_loader = DataLoader(
dataset=val_dataset,
batch_size=Cfg.batch // Cfg.subdivisions,
shuffle=True,
num_workers=8,
pin_memory=True,
drop_last=True,
collate_fn=val_collate,
)
coco = convert_to_coco_api(val_loader.dataset, bbox_fmt='coco')
all_models = glob.glob(os.path.join(Cfg.checkpoints, "*.pth"))
all_models.sort(key=lambda fp: int(
os.path.splitext(os.path.basename(fp))[0].replace("Yolov4_epoch", "")))
state_dict = torch.load(model_path, map_location=device)
for model_path in all_models:
print(f"eval on {os.path.splitext(os.path.basename(model_path))[0]}")
model = Yolov4(None, 1, True)
model.load_state_dict(state_dict)
model.to(device)
model.eval()
coco_evaluator = CocoEvaluator(coco, iou_types=["bbox"])
for images, targets in val_loader:
model_input = [[cv2.resize(img, (Cfg.w, Cfg.h))] for img in images]
model_input = np.concatenate(model_input, axis=0)
model_input = model_input.transpose(0, 3, 1, 2)
model_input = torch.from_numpy(model_input).div(255.0)
model_input = model_input.to(device)
targets = [{k: v.to(device)
for k, v in t.items()} for t in targets]
if torch.cuda.is_available():
torch.cuda.synchronize()
model_time = time.time()
outputs = model(model_input)
# outputs = [{k: v.to(cpu_device) for k, v in t.items()} for t in outputs]
model_time = time.time() - model_time
outputs = outputs.cpu().detach().numpy()
res = {}
for img, target, output in zip(images, targets, outputs):
img_height, img_width = img.shape[:2]
boxes = output[..., :4].copy() # output boxes in yolo format
boxes[..., :2] = boxes[
..., :2] - boxes[..., 2:] / 2 # to coco format
boxes[..., 0] = boxes[..., 0] * img_width
boxes[..., 1] = boxes[..., 1] * img_height
boxes[..., 2] = boxes[..., 2] * img_width
boxes[..., 3] = boxes[..., 3] * img_height
boxes = torch.as_tensor(boxes, dtype=torch.float32)
labels = torch.as_tensor(np.zeros((len(output), )),
dtype=torch.int64)
scores = torch.as_tensor(output[..., -1], dtype=torch.float32)
res[target["image_id"].item()] = {
"boxes": boxes,
"scores": scores,
"labels": labels,
}
evaluator_time = time.time()
coco_evaluator.update(res)
evaluator_time = time.time() - evaluator_time
# gather the stats from all processes
coco_evaluator.synchronize_between_processes()
# accumulate predictions from all images
coco_evaluator.accumulate()
coco_evaluator.summarize()
del model
del coco_evaluator
console.setLevel(log_level)
formatter = logging.Formatter(fmt)
console.setFormatter(formatter)
logging.getLogger('').addHandler(console)
return logging
if __name__ == "__main__":
logging = init_logger(log_dir='log')
cfg = get_args(**Cfg)
os.environ["CUDA_VISIBLE_DEVICES"] = cfg.gpu
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
logging.info(f'Using device {device}')
model = Yolov4(cfg.pretrained)
if torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model)
model.to(device=device)
try:
train(
model=model,
config=cfg,
epochs=cfg.TRAIN_EPOCHS,
device=device,
)
except KeyboardInterrupt:
torch.save(model.state_dict(), 'INTERRUPTED.pth')
logging.info('Saved interrupt')
help='path of your image file. or video or 0 to webcam',
dest='source')
args = parser.parse_args()
return args
if __name__ == '__main__':
args = get_args()
# model selection, darknet or pytorch
if args.weightfile.split('.')[-1] == 'weights': # darknet
m = Darknet(args.cfgfile)
m.load_weights(args.weightfile)
elif args.weightfile.split('.')[-1] in ['pth', 'pt', 'ph']: # pytorch
m = Yolov4(n_classes=args.nclasses, inference=True)
state_dict = torch.load(args.weightfile)
m.load_state_dict(state_dict)
m.head.inference = True
m.eval()
if use_cuda:
torch.cuda.set_device(torch.device('cuda:1'))
m.cuda()
if args.source == '0': # webcam
pass
else:
form = args.source.split('.')[-1]
if form == 'jpg': # image
pass
# import config file and save it to log
update_config(cfg, args)
init_seed(cfg.SEED)
anchors = get_anchors(cfg)
log_dir = os.path.join("log", os.path.basename(args.config_file)[:-5])
logging = init_logger(log_dir=log_dir)
logging.info(pprint.pformat(args))
logging.info(cfg)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
logging.info(f'Using device {device}')
if cfg.use_darknet_cfg:
model = Darknet(cfg.cfgfile)
else:
model = Yolov4(anchors,
yolov4conv137weight=cfg.pretrained,
n_classes=cfg.classes)
if torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model)
model.to(device=device)
try:
train(model=model,
config=cfg,
tb_log_dir=log_dir,
anchors=anchors,
epochs=cfg.TRAIN_EPOCHS,
device=device)
except KeyboardInterrupt:
torch.save(model.state_dict(), 'INTERRUPTED.pth')
def train(model,
device,
config,
epochs=5,
batch_size=1,
save_cp=True,
log_step=20,
img_scale=0.5):
train_dataset = Yolo_dataset(config.train_label, config, train=True)
val_dataset = Yolo_dataset(config.val_label, config, train=False)
n_train = len(train_dataset)
n_val = len(val_dataset)
train_loader = DataLoader(train_dataset,
batch_size=config.batch // config.subdivisions,
shuffle=True,
num_workers=8,
pin_memory=True,
drop_last=True,
collate_fn=collate)
val_loader = DataLoader(val_dataset,
batch_size=config.batch // config.subdivisions,
shuffle=True,
num_workers=8,
pin_memory=True,
drop_last=True,
collate_fn=val_collate)
writer = SummaryWriter(
log_dir=config.TRAIN_TENSORBOARD_DIR,
filename_suffix=
f'OPT_{config.TRAIN_OPTIMIZER}_LR_{config.learning_rate}_BS_{config.batch}_Sub_{config.subdivisions}_Size_{config.width}',
comment=
f'OPT_{config.TRAIN_OPTIMIZER}_LR_{config.learning_rate}_BS_{config.batch}_Sub_{config.subdivisions}_Size_{config.width}'
)
# writer.add_images('legend',
# torch.from_numpy(train_dataset.label2colorlegend2(cfg.DATA_CLASSES).transpose([2, 0, 1])).to(
# device).unsqueeze(0))
max_itr = config.TRAIN_EPOCHS * n_train
# global_step = cfg.TRAIN_MINEPOCH * n_train
global_step = 0
logging.info(f'''Starting training:
Epochs: {epochs}
Batch size: {config.batch}
Subdivisions: {config.subdivisions}
Learning rate: {config.learning_rate}
Training size: {n_train}
Validation size: {n_val}
Checkpoints: {save_cp}
Device: {device.type}
Images size: {config.width}
Optimizer: {config.TRAIN_OPTIMIZER}
Dataset classes: {config.classes}
Train label path:{config.train_label}
Pretrained:
''')
# learning rate setup
def burnin_schedule(i):
if i < config.burn_in:
factor = pow(i / config.burn_in, 4)
elif i < config.steps[0]:
factor = 1.0
elif i < config.steps[1]:
factor = 0.1
else:
factor = 0.01
return factor
if config.TRAIN_OPTIMIZER.lower() == 'adam':
optimizer = optim.Adam(
model.parameters(),
lr=config.learning_rate / config.batch,
betas=(0.9, 0.999),
eps=1e-08,
)
elif config.TRAIN_OPTIMIZER.lower() == 'sgd':
optimizer = optim.SGD(
params=model.parameters(),
lr=config.learning_rate / config.batch,
momentum=config.momentum,
weight_decay=config.decay,
)
scheduler = optim.lr_scheduler.LambdaLR(optimizer, burnin_schedule)
criterion = Yolo_loss(device=device,
batch=config.batch // config.subdivisions,
n_classes=config.classes)
# scheduler = ReduceLROnPlateau(optimizer, mode='max', verbose=True, patience=6, min_lr=1e-7)
# scheduler = CosineAnnealingWarmRestarts(optimizer, 0.001, 1e-6, 20)
save_prefix = 'Yolov4_epoch'
saved_models = deque()
model.train()
for epoch in range(epochs):
# model.train()
epoch_loss = 0
epoch_step = 0
with tqdm(total=n_train,
desc=f'Epoch {epoch + 1}/{epochs}',
unit='img',
ncols=50) as pbar:
for i, batch in enumerate(train_loader):
global_step += 1
epoch_step += 1
images = batch[0]
bboxes = batch[1]
images = images.to(device=device, dtype=torch.float32)
bboxes = bboxes.to(device=device)
bboxes_pred = model(images)
loss, loss_xy, loss_wh, loss_obj, loss_cls, loss_l2 = criterion(
bboxes_pred, bboxes)
# loss = loss / config.subdivisions
loss.backward()
epoch_loss += loss.item()
if global_step % config.subdivisions == 0:
optimizer.step()
scheduler.step()
model.zero_grad()
if global_step % (log_step * config.subdivisions) == 0:
writer.add_scalar('train/Loss', loss.item(), global_step)
writer.add_scalar('train/loss_xy', loss_xy.item(),
global_step)
writer.add_scalar('train/loss_wh', loss_wh.item(),
global_step)
writer.add_scalar('train/loss_obj', loss_obj.item(),
global_step)
writer.add_scalar('train/loss_cls', loss_cls.item(),
global_step)
writer.add_scalar('train/loss_l2', loss_l2.item(),
global_step)
writer.add_scalar('lr',
scheduler.get_lr()[0] * config.batch,
global_step)
pbar.set_postfix(
**{
'loss (batch)': loss.item(),
'loss_xy': loss_xy.item(),
'loss_wh': loss_wh.item(),
'loss_obj': loss_obj.item(),
'loss_cls': loss_cls.item(),
'loss_l2': loss_l2.item(),
'lr': scheduler.get_lr()[0] * config.batch
})
logging.debug(
'Train step_{}: loss : {},loss xy : {},loss wh : {},'
'loss obj : {}，loss cls : {},loss l2 : {},lr : {}'.
format(global_step, loss.item(), loss_xy.item(),
loss_wh.item(), loss_obj.item(),
loss_cls.item(), loss_l2.item(),
scheduler.get_lr()[0] * config.batch))
pbar.update(images.shape[0])
if cfg.use_darknet_cfg:
eval_model = Darknet(cfg.cfgfile, inference=True)
else:
eval_model = Yolov4(cfg.pretrained,
n_classes=cfg.classes,
inference=True)
# eval_model = Yolov4(yolov4conv137weight=None, n_classes=config.classes, inference=True)
if torch.cuda.device_count() > 1:
eval_model.load_state_dict(model.module.state_dict())
else:
eval_model.load_state_dict(model.state_dict())
eval_model.to(device)
evaluator = evaluate(eval_model, val_loader, config, device)
del eval_model
stats = evaluator.coco_eval['bbox'].stats
writer.add_scalar('train/AP', stats[0], global_step)
writer.add_scalar('train/AP50', stats[1], global_step)
writer.add_scalar('train/AP75', stats[2], global_step)
writer.add_scalar('train/AP_small', stats[3], global_step)
writer.add_scalar('train/AP_medium', stats[4], global_step)
writer.add_scalar('train/AP_large', stats[5], global_step)
writer.add_scalar('train/AR1', stats[6], global_step)
writer.add_scalar('train/AR10', stats[7], global_step)
writer.add_scalar('train/AR100', stats[8], global_step)
writer.add_scalar('train/AR_small', stats[9], global_step)
writer.add_scalar('train/AR_medium', stats[10], global_step)
writer.add_scalar('train/AR_large', stats[11], global_step)
if save_cp:
try:
# os.mkdir(config.checkpoints)
os.makedirs(config.checkpoints, exist_ok=True)
logging.info('Created checkpoint directory')
except OSError:
pass
save_path = os.path.join(config.checkpoints,
f'{save_prefix}{epoch + 1}.pth')
torch.save(model.state_dict(), save_path)
logging.info(f'Checkpoint {epoch + 1} saved !')
saved_models.append(save_path)
if len(saved_models) > config.keep_checkpoint_max > 0:
model_to_remove = saved_models.popleft()
try:
os.remove(model_to_remove)
except:
logging.info(f'failed to remove {model_to_remove}')
writer.close()
console.setLevel(log_level)
formatter = logging.Formatter(fmt)
console.setFormatter(formatter)
logging.getLogger('').addHandler(console)
return logging
if __name__ == "__main__":
logging = init_logger(log_dir='log')
cfg = get_args(**Cfg)
os.environ["CUDA_VISIBLE_DEVICES"] = cfg.gpu
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
logging.info(f'Using device {device}')
model = Yolov4()
if torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model)
# model = model.cuda()
model.to(device=device)
try:
train(
model=model,
config=cfg,
epochs=cfg.TRAIN_EPOCHS,
device=device,
)
except KeyboardInterrupt:
torch.save(model.state_dict(), 'INTERRUPTED.pth')