def main():
logger = create_logger(save_dir=cfg.log_dir)
print = logger.info
print(cfg)
cfg.device = torch.device('cuda')
torch.backends.cudnn.benchmark = False
max_per_image = 100
Dataset_eval = Damage_eval # your own data set
# Crack RE Spalling
dataset = Dataset_eval(cfg.data_dir, split='val', test_scales=cfg.test_scales, test_flip=cfg.test_flip) # split test
data_loader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False,
num_workers=1, pin_memory=True,
collate_fn=dataset.collate_fn)
print('Creating model...')
if 'hourglass' in cfg.arch:
model = get_hourglass[cfg.arch]
elif 'resdcn' in cfg.arch:
model = get_pose_net_resdcn(num_layers=18, head_conv=64, num_classes=3)
elif cfg.arch == 'resnet':
model = get_pose_net(num_layers=18, head_conv=64, num_classes=3)
elif cfg.arch == 'res_CBAM':
model = get_pose_net_resnet_CBAM(num_layers=18, head_conv=64, num_classes=3)
elif cfg.arch == 'resnet_PAM':
model = get_pose_net_resnet_PAM(num_layers=18, head_conv=64, num_classes=3)
elif cfg.arch == 'resnet_SE':
model = get_pose_net_resnet_SE(num_layers=18, head_conv=64, num_classes=3)
model = load_model(model, cfg.pretrain_dir)
model = model.to(cfg.device)
model.eval()
results = {}
with torch.no_grad():
for inputs in tqdm(data_loader):
img_id, inputs,img_path = inputs[0]
print('id%s ',img_id)
detections = []
for scale in inputs:
inputs[scale]['image'] = inputs[scale]['image'].to(cfg.device)
output = model(inputs[scale]['image'])[-1]
dets = ctdet_decode(*output, K=cfg.test_topk)
dets = dets.detach().cpu().numpy().reshape(1, -1, dets.shape[2])[0]
top_preds = {}
dets[:, :2] = transform_preds(dets[:, 0:2],
inputs[scale]['center'],
inputs[scale]['scale'],
(inputs[scale]['fmap_w'], inputs[scale]['fmap_h']))
dets[:, 2:4] = transform_preds(dets[:, 2:4],
inputs[scale]['center'],
inputs[scale]['scale'],
(inputs[scale]['fmap_w'], inputs[scale]['fmap_h']))
cls = dets[:, -1]
for j in range(dataset.num_classes):
inds = (cls == j)
top_preds[j + 1] = dets[inds, :5].astype(np.float32)
top_preds[j + 1][:, :4] /= scale
detections.append(top_preds)
bbox_and_scores = {}
for j in range(1, dataset.num_classes + 1):
bbox_and_scores[j] = np.concatenate([d[j] for d in detections], axis=0)
if len(dataset.test_scales) > 1:
soft_nms(bbox_and_scores[j], Nt=0.5, method=2)
scores = np.hstack([bbox_and_scores[j][:, 4] for j in range(1, dataset.num_classes + 1)])
if len(scores) > max_per_image:
kth = len(scores) - max_per_image
thresh = np.partition(scores, kth)[kth]
for j in range(1, dataset.num_classes + 1):
keep_inds = (bbox_and_scores[j][:, 4] >= thresh)
bbox_and_scores[j] = bbox_and_scores[j][keep_inds]
images_test = cv2.imread(img_path)
fig = plt.figure(0)
colors = COCO_COLORS
names = COCO_NAMES
#cv2.imwrite('E:/test1.png',images_test)
plt.imshow(cv2.cvtColor(images_test, cv2.COLOR_BGR2RGB))
for lab in bbox_and_scores:
for boxes in bbox_and_scores[lab]:
x1, y1, x2, y2, score = boxes
if (x1 < 0):
x1 = 0
if (y1 < 0):
y1 = 0
if (x2 > 511):
x2 = 511
if (y2 > 511):
y2 = 511
if score > 0.2:
plt.gca().add_patch(Rectangle((x1, y1), x2 - x1, y2 - y1, linewidth=2, edgecolor=colors[lab], facecolor='none'))
plt.text(x1 -12 , y1 - 12 , names[lab], bbox=dict(facecolor=colors[lab], alpha=0.5), fontsize=7, color='k')
fig.patch.set_visible(False)
Save_dir = 'data/damage/Predict_images' # save images
Image_name = img_path[-10:]
Save_dir = os.path.join(Save_dir, Image_name)
plt.axis('off')
plt.savefig(Save_dir, dpi=400, transparent=True, bbox_inches="tight", pad_inches=0.1) # 保存
plt.close(0)
results[img_id] = bbox_and_scores
eval_results = dataset.run_eval(results, cfg.ckpt_dir)
print(eval_results)
def main():
logger = create_logger(save_dir=cfg.log_dir)
print = logger.info
print(cfg)
cfg.device = torch.device('cuda')
torch.backends.cudnn.benchmark = False
max_per_image = 100
Dataset_eval = Damage_eval
dataset = Dataset_eval(cfg.data_dir,
split='train',
test_scales=cfg.test_scales,
test_flip=cfg.test_flip) # split test
data_loader = torch.utils.data.DataLoader(dataset,
batch_size=1,
shuffle=False,
num_workers=1,
pin_memory=True,
collate_fn=dataset.collate_fn)
print('Creating model...')
if 'hourglass' in cfg.arch:
model = get_hourglass[cfg.arch]
elif 'resdcn' in cfg.arch:
model = get_pose_net_resdcn(num_layers=18, head_conv=64, num_classes=3)
elif cfg.arch == 'resnet':
model = get_pose_net(num_layers=18, head_conv=64, num_classes=3)
elif cfg.arch == 'res_CBAM':
model = get_pose_net_resnet_CBAM(num_layers=18,
head_conv=64,
num_classes=3)
elif cfg.arch == 'resnet_PAM':
model = get_pose_net_resnet_PAM(num_layers=18,
head_conv=64,
num_classes=3)
elif cfg.arch == 'resnet_SE':
model = get_pose_net_resnet_SE(num_layers=18,
head_conv=64,
num_classes=3)
def Evaluate(epoch, model):
print('\n Evaluate@Epoch: %d' % epoch)
start_time = time.clock()
print('Start time %s Seconds' % start_time)
model.eval()
torch.cuda.empty_cache()
max_per_image = 100
results = {}
with torch.no_grad():
for inputs in data_loader:
img_id, inputs, img_path = inputs[0]
detections = []
for scale in inputs:
inputs[scale]['image'] = inputs[scale]['image'].to(
cfg.device) # (1,3)
output = model(
inputs[scale]['image'])[-1] # hmap, regs, pxpy
dets = ctdet_decode(
*output, K=cfg.test_topk
) # torch.cat([bboxes, scores, clses], dim=2)
dets = dets.detach().cpu().numpy().reshape(
1, -1, dets.shape[2])[0]
top_preds = {}
dets[:, :2] = transform_preds(
dets[:, 0:2], inputs[scale]['center'],
inputs[scale]['scale'],
(inputs[scale]['fmap_w'], inputs[scale]['fmap_h']))
dets[:, 2:4] = transform_preds(
dets[:, 2:4], inputs[scale]['center'],
inputs[scale]['scale'],
(inputs[scale]['fmap_w'], inputs[scale]['fmap_h']))
clses = dets[:, -1]
for j in range(dataset.num_classes):
inds = (clses == j)
top_preds[j + 1] = dets[inds, :5].astype(np.float32)
top_preds[j + 1][:, :4] /= scale
detections.append(top_preds)
bbox_and_scores = {
j: np.concatenate([d[j] for d in detections], axis=0)
for j in range(1, dataset.num_classes + 1)
}
scores = np.hstack([
bbox_and_scores[j][:, 4]
for j in range(1, dataset.num_classes + 1)
])
if len(scores) > max_per_image:
kth = len(scores) - max_per_image
thresh = np.partition(scores, kth)[kth]
for j in range(1, dataset.num_classes + 1):
keep_inds = (bbox_and_scores[j][:, 4] >= thresh)
bbox_and_scores[j] = bbox_and_scores[j][keep_inds]
results[img_id] = bbox_and_scores
end_time = time.clock()
eval_results = dataset.run_eval(results, save_dir=cfg.ckpt_dir)
print(eval_results)
print('End time %s Seconds' % end_time)
Run_time = end_time - start_time
FPS = 100 / Run_time # replace 100 with the number of images
print('FPS %s ' % FPS)
#summary_writer.add_scalar('Evaluate_mAP/mAP', eval_results[0], epoch)
return eval_results[0]
num_epochs = 60 # replace 60 with the number of epoch
Max_mAP = 0
for epoch in range(1, num_epochs + 1):
cfg.pretrain_dir = os.path.join(cfg.ckpt_dir, 'checkpoint_epoch' +
str(epoch) + '.t7') # the address
model = load_model(model, cfg.pretrain_dir)
model = model.to(cfg.device)
mAP = Evaluate(epoch, model)
if mAP > Max_mAP:
Max_mAP = mAP
print('Max_AP=%s' % Max_mAP)
def main():
saver = create_saver(cfg.local_rank, save_dir=cfg.ckpt_dir)
logger = create_logger(cfg.local_rank, save_dir=cfg.log_dir)
summary_writer = create_summary(cfg.local_rank, log_dir=cfg.log_dir)
print = logger.info
print(cfg)
torch.manual_seed(300)
torch.backends.cudnn.benchmark = True
'''
# you can also set like this. If you do like this, the random seed will be fixed.
torch.manual_seed(350)
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True # consistent results on the cpu and gpu
'''
num_gpus = torch.cuda.device_count()
if cfg.dist:
cfg.device = torch.device('cuda:%d' % cfg.local_rank)
torch.cuda.set_device(cfg.local_rank)
dist.init_process_group(backend='nccl',
init_method='env://',
world_size=num_gpus,
rank=cfg.local_rank)
else:
cfg.device = torch.device('cuda')
print('Setting up data...')
Dataset = Damage
train_dataset = Dataset(cfg.data_dir,
'train',
split_ratio=cfg.split_ratio,
img_size=cfg.img_size)
train_sampler = torch.utils.data.distributed.DistributedSampler(
train_dataset, num_replicas=num_gpus, rank=cfg.local_rank)
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=cfg.batch_size // num_gpus if cfg.dist else cfg.batch_size,
shuffle=not cfg.dist,
num_workers=cfg.num_workers,
pin_memory=True,
drop_last=True,
sampler=train_sampler if cfg.dist else None)
Dataset_eval = Damage_eval
test_dataset = Dataset_eval(cfg.data_dir,
'test',
test_scales=[1.],
test_flip=False)
test_loader = torch.utils.data.DataLoader(
test_dataset,
batch_size=1, # 测试集的batch_size
shuffle=False,
num_workers=1,
pin_memory=True, # 测试集的num_workers
collate_fn=test_dataset.collate_fn)
val_dataset = Dataset_eval(cfg.data_dir,
'val',
test_scales=[1.],
test_flip=False)
val_loader = torch.utils.data.DataLoader(
val_dataset,
batch_size=1, # 验证集的batch_size
shuffle=False,
num_workers=1,
pin_memory=True, # 验证集的num_workers
collate_fn=val_dataset.collate_fn)
print('Creating model...')
if 'hourglass' in cfg.arch:
model = get_hourglass[cfg.arch]
elif 'resdcn' in cfg.arch:
model = get_pose_net_resdcn(num_layers=18, head_conv=64, num_classes=3)
elif cfg.arch == 'resnet':
model = get_pose_net(num_layers=18, head_conv=64, num_classes=3)
elif cfg.arch == 'resnet_CBAM':
model = get_pose_net_resnet_CBAM(num_layers=18,
head_conv=64,
num_classes=3)
elif cfg.arch == 'resnet_PAM':
model = get_pose_net_resnet_PAM(num_layers=18,
head_conv=64,
num_classes=3)
elif cfg.arch == 'resnet_SE':
model = get_pose_net_resnet_SE(num_layers=18,
head_conv=64,
num_classes=3)
if cfg.dist:
# model = nn.SyncBatchNorm.convert_sync_batchnorm(model)
model = model.to(cfg.device)
model = nn.parallel.DistributedDataParallel(
model, device_ids=[
cfg.local_rank,
], output_device=cfg.local_rank)
else:
model = nn.DataParallel(model).to(cfg.device)
#if os.path.isfile(cfg.pretrain_dir):
# model = load_model(model, cfg.pretrain_dir) # 不加载预训练模型
optimizer = torch.optim.Adam(model.parameters(), cfg.lr)
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,
cfg.lr_step,
gamma=0.1) # adjust lr
def train(epoch):
print('\n Epoch: %d' % epoch)
model.train()
tic = time.perf_counter()
for batch_idx, batch in enumerate(train_loader):
for k in batch:
if k != 'meta':
batch[k] = batch[k].to(device=cfg.device,
non_blocking=True)
outputs = model(batch['image'])
hmap, regs, w_h_, pxpy = zip(*outputs)
# batch * C（channel） * W * H
regs = [
_tranpose_and_gather_feature(r, batch['inds']) for r in regs
]
pxpy = [
_tranpose_and_gather_feature(r, batch['inds']) for r in pxpy
]
w_h_ = [
_tranpose_and_gather_feature(r, batch['inds']) for r in w_h_
]
# batch * K * C= batch * 128 *2
hmap_loss = _neg_loss(hmap, batch['hmap'])
reg_loss = _SmoothL1Loss(regs, batch['regs'], batch['ind_masks'])
pxpy_loss = _reg_loss(pxpy, batch['pxpy'], batch['ind_masks'])
w_h_loss = _SmoothL1Loss(w_h_, batch['w_h_'], batch['ind_masks'])
loss = hmap_loss + 10 * reg_loss + 0.1 * w_h_loss + 0.1 * pxpy_loss
optimizer.zero_grad()
loss.backward()
optimizer.step()
if batch_idx % cfg.log_interval == 0:
duration = time.perf_counter() - tic
tic = time.perf_counter()
print(
'[%d/%d-%d/%d] ' %
(epoch, cfg.num_epochs, batch_idx, len(train_loader)) +
' hmap_loss= %.5f reg_loss= %.5f w_h_loss= %.5f pxpy_loss= %.5f'
% (hmap_loss.item(), reg_loss.item(), w_h_loss.item(),
pxpy_loss.item()) + ' (%d samples/sec)' %
(cfg.batch_size * cfg.log_interval / duration))
step = len(train_loader) * epoch + batch_idx
summary_writer.add_scalar('hmap_loss', hmap_loss.item(), step)
summary_writer.add_scalar('reg_loss', reg_loss.item(), step)
summary_writer.add_scalar('w_h_loss', w_h_loss.item(), step)
summary_writer.add_scalar('pxpy_loss', pxpy_loss.item(), step)
return
#--------------------test set--------------------#
def test_map(epoch):
print('\n Test@Epoch: %d' % epoch)
start_time = time.clock()
print('Start time %s Seconds' % start_time)
model.eval()
torch.cuda.empty_cache()
max_per_image = 100
results = {}
with torch.no_grad():
for inputs in test_loader:
img_id, inputs, img_path = inputs[0]
detections = []
for scale in inputs:
inputs[scale]['image'] = inputs[scale]['image'].to(
cfg.device) # (1,3)
output = model(inputs[scale]['image'])[-1]
dets = ctdet_decode(
*output, K=cfg.test_topk
) # torch.cat([bboxes, scores, clses], dim=2)
dets = dets.detach().cpu().numpy().reshape(
1, -1, dets.shape[2])[0]
top_preds = {}
dets[:, :2] = transform_preds(
dets[:, 0:2], inputs[scale]['center'],
inputs[scale]['scale'],
(inputs[scale]['fmap_w'], inputs[scale]['fmap_h']))
dets[:, 2:4] = transform_preds(
dets[:, 2:4], inputs[scale]['center'],
inputs[scale]['scale'],
(inputs[scale]['fmap_w'], inputs[scale]['fmap_h']))
clses = dets[:, -1]
for j in range(test_dataset.num_classes):
inds = (clses == j)
top_preds[j + 1] = dets[inds, :5].astype(np.float32)
top_preds[j + 1][:, :4] /= scale
detections.append(top_preds)
bbox_and_scores = {
j: np.concatenate([d[j] for d in detections], axis=0)
for j in range(1, test_dataset.num_classes + 1)
}
scores = np.hstack([
bbox_and_scores[j][:, 4]
for j in range(1, test_dataset.num_classes + 1)
])
if len(scores) > max_per_image:
kth = len(scores) - max_per_image
thresh = np.partition(scores, kth)[kth]
for j in range(1, test_dataset.num_classes + 1):
keep_inds = (bbox_and_scores[j][:, 4] >= thresh)
bbox_and_scores[j] = bbox_and_scores[j][keep_inds]
results[img_id] = bbox_and_scores
end_time = time.clock()
eval_results = test_dataset.run_eval(results, save_dir=cfg.ckpt_dir)
print(eval_results)
print('End time %s Seconds' % end_time)
Run_time = end_time - start_time
FPS = 100 / Run_time # replace 100 with the number of images
print('FPS %s ' % FPS)
summary_writer.add_scalar('test_mAP/mAP', eval_results[0], epoch)
return eval_results[0]
#--------------------end of test set--------------------#
#--------------------validation set--------------------#
def val_map(epoch):
print('\n Val@Epoch: %d' % epoch)
start_time = time.clock()
print('Start time %s Seconds' % start_time)
model.eval()
torch.cuda.empty_cache()
max_per_image = 100
results = {}
with torch.no_grad():
for inputs in val_loader:
img_id, inputs, img_path = inputs[0]
detections = []
for scale in inputs:
inputs[scale]['image'] = inputs[scale]['image'].to(
cfg.device) # (1,3)
output = model(
inputs[scale]['image'])[-1] # hmap, regs, pxpy
dets = ctdet_decode(
*output, K=cfg.test_topk
) # torch.cat([bboxes, scores, clses], dim=2)
dets = dets.detach().cpu().numpy().reshape(
1, -1, dets.shape[2])[0]
top_preds = {}
dets[:, :2] = transform_preds(
dets[:, 0:2], inputs[scale]['center'],
inputs[scale]['scale'],
(inputs[scale]['fmap_w'], inputs[scale]['fmap_h']))
dets[:, 2:4] = transform_preds(
dets[:, 2:4], inputs[scale]['center'],
inputs[scale]['scale'],
(inputs[scale]['fmap_w'], inputs[scale]['fmap_h']))
clses = dets[:, -1]
for j in range(val_dataset.num_classes):
inds = (clses == j)
top_preds[j + 1] = dets[inds, :5].astype(np.float32)
top_preds[j + 1][:, :4] /= scale
detections.append(top_preds)
bbox_and_scores = {
j: np.concatenate([d[j] for d in detections], axis=0)
for j in range(1, val_dataset.num_classes + 1)
}
scores = np.hstack([
bbox_and_scores[j][:, 4]
for j in range(1, val_dataset.num_classes + 1)
])
if len(scores) > max_per_image:
kth = len(scores) - max_per_image
thresh = np.partition(scores, kth)[kth]
for j in range(1, val_dataset.num_classes + 1):
keep_inds = (bbox_and_scores[j][:, 4] >= thresh)
bbox_and_scores[j] = bbox_and_scores[j][keep_inds]
results[img_id] = bbox_and_scores
end_time = time.clock()
eval_results = val_dataset.run_eval(results, save_dir=cfg.ckpt_dir)
print(eval_results)
print('End time %s Seconds' % end_time)
Run_time = end_time - start_time
FPS = 100 / Run_time # replace 100 with the number of images
print('FPS %s ' % FPS)
summary_writer.add_scalar('val_mAP/mAP', eval_results[0], epoch)
return eval_results[0]
#--------------------end of validation set--------------------#
print('Starting training...')
Max_test_AP = 0 # max test AP
Max_val_AP = 0 # max validation AP
flag_epoch = 1
for epoch in range(1, cfg.num_epochs + 1):
train_sampler.set_epoch(epoch)
train(epoch)
if epoch >= flag_epoch:
test_mAP = test_map(epoch)
val_mAP = val_map(epoch)
if (test_mAP > Max_test_AP):
Max_test_AP = test_mAP
if (val_mAP > Max_val_AP):
print(
saver.save(model.module.state_dict(),
'checkpoint_MaxAP_epoch' + str(epoch)))
Max_val_AP = val_mAP
print(saver.save(model.module.state_dict(),
'checkpoint')) # save current epoch
total = sum([param.nelement()
for param in model.parameters()]) # calculate parameters
print("Number of parameter: %.2fM" % (total / 1e6))
print('Max_test_AP=%s' % Max_test_AP)
print('Max_val_AP=%s' % Max_val_AP)
lr_scheduler.step(epoch) # move to here after pytorch1.1.0
summary_writer.close()