def psoft(cls_dets, thresh):
cls_dets = soft_nms(cls_dets, method=2)
indexes = np.where(cls_dets[:, -1] > thresh)[0]
cls_dets = cls_dets[indexes]
return cls_dets
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 pred_eval(predictor, test_data, imdb, cfg, vis = False, thresh = 1e-3, logger = None, ignore_cache = True):
det_file = os.path.join(imdb.result_path, imdb.name + '_detections.pkl')
if os.path.exists(det_file) and not ignore_cache:
with open(det_file, 'rb') as fid:
all_boxes = pickle.load(fid)
info_str = imdb.evaluate_detections(all_boxes)
if logger:
logger.info('evaluate detections: \n{}'.format(info_str))
return
assert vis or not test_data.shuffle
data_names = [k[0] for k in test_data.provide_data]
#if not isinstance(test_data,mx.io.PrefetchingIter):
# test_data = mx.io.PrefetchingIter(test_data)
max_per_image = cfg.TEST.max_per_image
num_images = imdb.num_images
for test_scale_index, test_scale in enumerate(cfg.TEST_SCALES):
det_file_single_scale = os.path.join(imdb.result_path, imdb.name + "_detections_" + str(test_scale_index) + '.pkl')
cfg.SCALES = [test_scale]
test_data.reset()
all_boxes_single_scale = [[[] for _ in range(num_images)]
for _ in range(imdb.num_classes)]
detect_at_single_scale(predictor, data_names, imdb, test_data, cfg, thresh, vis, all_boxes_single_scale, logger)
with open(det_file_single_scale,'wb') as f:
pickle.dump(all_boxes_single_scale, f, protocol = pickle.HIGHEST_PROTOCOL)
all_boxes = [[[] for _ in range(num_images)] for _ in range(imdb.num_classes)]
for test_scale_index, test_scale in enumerate(cfg.TEST_SCALES):
det_file_single_scale = os.path.join(imdb.result_path, imdb.name + '_detections_' + str(test_scale_index) + '.pkl')
if os.path.exists(det_file_single_scale):
with open(det_file_single_scale, 'rb') as fid:
all_boxes_single_scale = pickle.load(fid)
for idx_class in range(1, imdb.num_classes):
for idx_im in range(0,num_images):
if len(all_boxes[idx_class][idx_im]) == 0:
all_boxes[idx_class][idx_im] = all_boxes_single_scale[idx_class][idx_im]
else:
all_boxes[idx_class][idx_im] = np.vstack((all_boxes[idx_class][idx_im], all_boxes_single_scale[idx_class][idx_im]))
for idx_class in range(1, imdb.num_classes):
for idx_im in range(0, num_images):
# keep = py_nms(all_boxes[idx_class][idx_im],cfg.TEST.NMS)
keep = soft_nms(all_boxes[idx_class][idx_im],cfg.TEST.NMS)
all_boxes[idx_class][idx_im] = all_boxes[idx_class][idx_im][keep, :]
if max_per_image > 0:
for idx_im in range(0,num_images):
image_scores = np.hstack([all_boxes[j][idx_im][:,-1]
for j in range(1, imdb.num_classes)])
if len(image_scores) > max_per_image:
image_thresh = np.sort(image_scores)[-max_per_image]
for j in range(1, imdb.num_classes):
keep = np.where(all_boxes[j][idx_im][:,-1] >= image_thresh)[0]
all_boxes[j][idx_im] = all_boxes[j][idx_im][keep,:]
if vis:
test_data.reset()
for i in range(num_images):
im_info, data = test_data.next()
img = data.data[0].asnumpy().transpose((0,2,3,1))[0]
img = (img * np.array([[[0.229, 0.224, 0.225]]]) +np.array([[[0.485, 0.456, 0.406]]])) * 255
img = np.clip(img,0,255)
img = img.astype(np.uint8)
print(type(img))
image = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
scale = im_info[0,2]
for j in range(1,imdb.num_classes):
for box in all_boxes[j][i]:
if box[-1] < 0.3: continue
box[:4] = (box[:4] * scale)
box = box.astype(np.int64)
print(box)
cv2.rectangle(image,tuple(box[:2]),tuple(box[2:4]),(255,0,0),1)
cv2.putText(image,names[j], tuple(box[:2]),cv2.FONT_HERSHEY_COMPLEX,1,(0,0,255),1)
cv2.imwrite("./det_images/det_img_{}.jpg".format(i),image)
print("detected image saved in ./det_images/det_img_{}.jpg".format(i))
import pdb
pdb.set_trace()
with open(det_file, 'wb') as f:
pickle.dump(all_boxes, f, protocol = pickle.HIGHEST_PROTOCOL)
info_str = imdb.evaluate_detections(all_boxes)
if logger:
logger.info("evaluate detections: \n{}".format(info_str))
def psoft(cls_dets):
keep = soft_nms(cls_dets, method=2)
return cls_dets[keep]
def main():
cfg.device = torch.device('cuda')
torch.backends.cudnn.benchmark = False
max_per_image = 100
image = cv2.imread(cfg.img_dir)
# orig_image = image
height, width = image.shape[0:2]
padding = 127 if 'hourglass' in cfg.arch else 31
imgs = {}
for scale in cfg.test_scales:
new_height = int(height * scale)
new_width = int(width * scale)
if cfg.img_size > 0:
img_height, img_width = cfg.img_size, cfg.img_size
center = np.array([new_width / 2., new_height / 2.],
dtype=np.float32)
scaled_size = max(height, width) * 1.0
scaled_size = np.array([scaled_size, scaled_size],
dtype=np.float32)
else:
img_height = (new_height | padding) + 1
img_width = (new_width | padding) + 1
center = np.array([new_width // 2, new_height // 2],
dtype=np.float32)
scaled_size = np.array([img_width, img_height], dtype=np.float32)
img = cv2.resize(image, (new_width, new_height))
trans_img = get_affine_transform(center, scaled_size, 0,
[img_width, img_height])
img = cv2.warpAffine(img, trans_img, (img_width, img_height))
img = img.astype(np.float32) / 255.
img -= np.array(COCO_MEAN if cfg.dataset == 'coco' else VOC_MEAN,
dtype=np.float32)[None, None, :]
img /= np.array(COCO_STD if cfg.dataset == 'coco' else VOC_STD,
dtype=np.float32)[None, None, :]
img = img.transpose(2, 0,
1)[None, :, :, :] # from [H, W, C] to [1, C, H, W]
if cfg.test_flip:
img = np.concatenate((img, img[:, :, :, ::-1].copy()), axis=0)
imgs[scale] = {
'image': torch.from_numpy(img).float(),
'center': np.array(center),
'scale': np.array(scaled_size),
'fmap_h': np.array(img_height // 4),
'fmap_w': np.array(img_width // 4)
}
print('Creating model...')
if 'hourglass' in cfg.arch:
model = get_hourglass[cfg.arch]
elif 'resdcn' in cfg.arch:
model = get_pose_net(num_layers=int(cfg.arch.split('_')[-1]),
num_classes=80 if cfg.dataset == 'coco' else 20)
else:
raise NotImplementedError
model = load_model(model, cfg.ckpt_dir)
model = model.to(cfg.device)
model.eval()
with torch.no_grad():
detections = []
for scale in imgs:
imgs[scale]['image'] = imgs[scale]['image'].to(cfg.device)
output = model(imgs[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], imgs[scale]['center'], imgs[scale]['scale'],
(imgs[scale]['fmap_w'], imgs[scale]['fmap_h']))
dets[:, 2:4] = transform_preds(
dets[:, 2:4], imgs[scale]['center'], imgs[scale]['scale'],
(imgs[scale]['fmap_w'], imgs[scale]['fmap_h']))
cls = dets[:, -1]
for j in range(80):
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, 81 if cfg.dataset == 'coco' else 21):
bbox_and_scores[j] = np.concatenate([d[j] for d in detections],
axis=0)
if len(cfg.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, 81 if cfg.dataset == 'coco' else 21)
])
if len(scores) > max_per_image:
kth = len(scores) - max_per_image
thresh = np.partition(scores, kth)[kth]
for j in range(1, 81 if cfg.dataset == 'coco' else 21):
keep_inds = (bbox_and_scores[j][:, 4] >= thresh)
bbox_and_scores[j] = bbox_and_scores[j][keep_inds]
# plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
# plt.show()
fig = plt.figure(0)
colors = COCO_COLORS if cfg.dataset == 'coco' else VOC_COLORS
names = COCO_NAMES if cfg.dataset == 'coco' else VOC_NAMES
plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
for lab in bbox_and_scores:
for boxes in bbox_and_scores[lab]:
x1, y1, x2, y2, score = boxes
if score > 0.3:
plt.gca().add_patch(
Rectangle((x1, y1),
x2 - x1,
y2 - y1,
linewidth=2,
edgecolor=colors[lab],
facecolor='none'))
plt.text(x1 + 3,
y1 + 3,
names[lab] + '%.2f' % score,
bbox=dict(facecolor=colors[lab], alpha=0.5),
fontsize=7,
color='k')
fig.patch.set_visible(False)
plt.axis('off')
plt.savefig('data/demo_results.png', dpi=300, transparent=True)
plt.show()
def psoft(cls_dets):
cls_dets = soft_nms(cls_dets, method=1)
return cls_dets
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 = COCO_eval if cfg.dataset == 'coco' else PascalVOC_eval
dataset = Dataset_eval(cfg.data_dir,
split='val',
img_size=cfg.img_size,
test_scales=cfg.test_scales,
test_flip=cfg.test_flip)
data_loader = torch.utils.data.DataLoader(dataset,
batch_size=1,
shuffle=False,
num_workers=1,
pin_memory=False,
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(num_layers=int(cfg.arch.split('_')[-1]),
num_classes=dataset.num_classes)
else:
raise NotImplementedError
model = load_model(model, cfg.pretrain_dir)
model = model.to(cfg.device)
model.eval()
results = {}
with torch.no_grad():
for inputs in data_loader:
img_id, inputs = inputs[0]
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]
results[img_id] = bbox_and_scores
eval_results = dataset.run_eval(results, cfg.ckpt_dir)
print(eval_results)