def detect(cfgfile, weightfile, imgfile):
m = Darknet(cfgfile)
checkpoint = torch.load(weightfile)
model_dict = m.state_dict()
pretrained_dict = checkpoint
keys = []
for k, v in pretrained_dict.items():
keys.append(k)
i = 0
for k, v in model_dict.items():
if v.size() == pretrained_dict[keys[i]].size():
model_dict[k] = pretrained_dict[keys[i]]
i = i + 1
m.load_state_dict(model_dict)
# m.load_state_dict(torch.load(weightfile))
# m.print_network()
# m.load_weights(weightfile)
print('Loading weights from %s... Done!' % (weightfile))
namesfile = 'data/mydata.names'
use_cuda = 1
if use_cuda:
m.cuda()
input_img = cv2.imread(imgfile)
# orig_img = Image.open(imgfile).convert('RGB')
start = time.time()
boxes, scale = do_detect(m, input_img, 0.5, 0.4, use_cuda)
finish = time.time()
print('%s: Predicted in %f seconds.' % (imgfile, (finish - start)))
class_names = load_class_names(namesfile)
# draw_boxes(input_img,boxes,scale=scale)
plot_boxes_cv2(input_img,
boxes,
'predictions1.jpg',
class_names=class_names,
scale=scale)
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)
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')
try:
sys.exit(0)
except SystemExit:
os._exit(0)
model = Darknet(cfg.cfgfile)
start_epoch = cfg.start_epoch
_, extension = os.path.splitext(cfg.pretrained)
if extension == '.weights':
model.load_weights(cfg.pretrained)
elif extension == '.pth':
ckpt = torch.load(cfg.pretrained)
model.load_state_dict(ckpt['state_dict'])
if 'epoch' in ckpt:
start_epoch = ckpt['epoch']
# if torch.cuda.device_count() > 1:
# model = torch.nn.DataParallel(model)
model.to(device=device)
try:
train(model=model,
config=cfg,
epochs=cfg.num_epochs,
device=device,
start_epoch=start_epoch,
batch_size=cfg.batch_size)
except KeyboardInterrupt:
save_dict = {'state_dict': model.state_dict()}
torch.save(save_dict, 'INTERRUPTED.pth')
logging.info('Saved interrupt')
try:
sys.exit(0)
except SystemExit:
os._exit(0)
# get logger and args
logging = init_logger(log_dir="log")
cfg = get_args(**Cfg)
# look for available devices
os.environ["CUDA_VISIBLE_DEVICES"] = cfg.gpu
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logging.info(f"Using device {device}")
# load model and push to device
model = Darknet(cfg.cfgfile, model_type="BEV_grid")
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(), "checkpoints/INTERRUPTED.pth")
logging.info("Saved interrupt to checkpoints/INTERRUPTED.pth")
try:
sys.exit(0)
except SystemExit:
os._exit(0)
