def main():
aug = aug_aocr(args)
datasets = data.fetch_data(args, args.datasets, batch_size=args.batch_size_per_gpu,
batch_size_val=args.batch_size_per_gpu_val, k_fold=1, split_val=0.1,
pre_process=None, aug=aug)
for idx, (train_set, val_set) in enumerate(datasets):
losses = []
lev_dises, str_accus = [], []
print("\n =============== Cross Validation: %s/%s ================ " %
(idx + 1, len(datasets)))
# Prepare Network
encoder = att_model.Attn_CNN(backbone_require_grad=True)
decoder = att_model.AttnDecoder(args)
encoder.apply(init.init_rnn).apply(init.init_others)
decoder.apply(init.init_rnn).apply(init.init_others)
criterion = nn.NLLLoss()
encoder = torch.nn.DataParallel(encoder).cuda()
decoder = torch.nn.DataParallel(decoder).cuda()
torch.backends.cudnn.benchmark = True
if args.finetune:
encoder, decoder = util.load_latest_model(args, [encoder, decoder],
prefix=["encoder", "decoder"], strict=False)
# Prepare loss function and optimizer
encoder_optimizer = AdaBound(encoder.parameters(), lr=args.learning_rate,
final_lr=args.learning_rate * 10, weight_decay=args.weight_decay)
decoder_optimizer = AdaBound(decoder.parameters(), lr=args.learning_rate,
final_lr=args.learning_rate * 10, weight_decay=args.weight_decay)
for epoch in range(args.epoch_num):
loss = fit(args, encoder, decoder, train_set, encoder_optimizer,
decoder_optimizer, criterion, is_train=True)
losses.append(loss)
train_losses = [np.asarray(losses)]
if val_set is not None:
lev_dis, str_accu = fit(args, encoder, decoder, val_set, encoder_optimizer,
decoder_optimizer, criterion, is_train=False)
lev_dises.append(lev_dis)
str_accus.append(str_accu)
val_scores = [np.asarray(lev_dises), np.asarray(str_accus)]
if epoch % 5 == 0:
util.save_model(args, args.curr_epoch, encoder.state_dict(), prefix="encoder",
keep_latest=20)
util.save_model(args, args.curr_epoch, decoder.state_dict(), prefix="decoder",
keep_latest=20)
if epoch > 4:
vb.plot_multi_loss_distribution(
multi_line_data= [train_losses, val_scores],
multi_line_labels= [["NLL Loss"], ["Levenstein", "String-Level"]],
save_path = args.loss_log, window=5, name = dt,
bound=[None, {"low": 0.0, "high": 100.0}],
titles=["Train Loss", "Validation Score"]
)
def test():
aug = aug_test(args)
dataset = data.fetch_detection_data(args,
sources=args.test_sources,
k_fold=1,
auxiliary_info=args.test_aux,
aug=aug,
batch_size=1 /
torch.cuda.device_count(),
shuffle=False)[0][0]
net = model.SSD(cfg,
connect_loc_to_conf=args.loc_to_conf,
fix_size=args.fix_size,
conf_incep=args.conf_incep,
loc_incep=args.loc_incep,
nms_thres=args.nms_threshold,
loc_preconv=args.loc_preconv,
conf_preconv=args.conf_preconv,
FPN=args.feature_pyramid_net,
SA=args.self_attention,
in_wid=args.inner_filters,
m_factor=args.inner_m_factor)
net = torch.nn.DataParallel(net, device_ids=[0], output_device=0).cuda()
net = util.load_latest_model(args,
net,
prefix=args.model_prefix_finetune,
strict=True)
detector = model.Detect(num_classes=2,
bkg_label=0,
top_k=args.detector_top_k,
conf_thresh=args.detector_conf_threshold,
nms_thresh=args.detector_nms_threshold)
with torch.no_grad():
for batch_idx, (images, targets) in enumerate(dataset):
images = images.cuda()
print(images.shape)
targets = [ann.cuda() for ann in targets]
ratios = images.size(3) / images.size(2)
if ratios != 1.0:
print(ratios)
out = net(images, is_train=False)
loc_data, conf_data, prior_data = out
#prior_data = prior_data.to("cuda:%d" % (device_id))
det_result = detector(loc_data, conf_data, prior_data)
eval_result = evaluate(images,
det_result.data,
targets,
batch_idx,
args.threshold,
visualize=True,
post_combine=True)
print(eval_result)
def test():
args.train = False
dataset = data.fetch_probaV_data(args,
sources=args.test_sources,
shuffle=False,
batch_size=1 / torch.cuda.device_count(),
auxiliary_info=[2, 2])[0][0]
if args.which_model.lower() == "carn":
net = model.CARN(args.n_selected_img,
args.filters,
3,
s_MSE=True,
trellis=args.trellis)
elif args.which_model.lower() == "rdn":
net = model.RDN(args.n_selected_img,
3,
3,
filters=args.filters,
s_MSE=True,
group=args.n_selected_img,
trellis=args.trellis)
elif args.which_model.lower() == "basic":
net = model.ProbaV_basic(inchannel=args.n_selected_img)
else:
print("args.which_model or -wm should be one of [carn, rdn, basic], "
"your -wm %s is illegal, and switched to 'basic' automatically" %
(args.which_model.lower()))
net = model.ProbaV_basic(inchannel=args.n_selected_img)
net = torch.nn.DataParallel(net, device_ids=[0], output_device=0).cuda()
torch.backends.cudnn.benchmark = True
net = util.load_latest_model(args,
net,
prefix=args.model_prefix_finetune,
strict=True)
with torch.no_grad():
for batch_idx, (images, blend_target, unblend_target,
norm) in enumerate(dataset):
print(batch_idx)
images, blend_target = images.cuda(), blend_target.cuda(),
prediction, mae, s_mse = net(images, blend_target, train=False)
pred = vb.plot_tensor(args, prediction, margin=0)
cv2.imwrite(
os.path.expanduser("~/Pictures/result/%s.jpg" %
str(batch_idx).zfill(4)),
pred / 65536 * 255)
def main():
test_folder = os.path.expanduser(
"~/Pictures/dataset/reid/eval_lp/HHR_Body")
diff_person_choose = 4
diff_view = 8
net = model.Encoder(args)
net = torch.nn.DataParallel(net).cuda()
net.eval()
net = util.load_latest_model(args, net, prefix=args.model_prefix, nth=1)
#test_data, label_mapping = data.fetch_dataset(args, verbose=False, for_test=True)
person_id = os.listdir(test_folder)
gt_label, views = [], []
for p_id in sorted(person_id):
views += random.sample(glob.glob(os.path.join(test_folder, p_id, "*")),
diff_view)
test_imgs = load_img(views, augment())
gt_label = torch.arange(len(person_id)).unsqueeze(-1).repeat(
1, diff_view).view(-1).int()
with torch.no_grad():
test(net, test_imgs, gt_label)
def test_rotation(opt):
result_dir = os.path.join(args.path, args.code_name,
"result+" + "-".join(opt.model_prefix_list))
if not os.path.exists(result_dir):
os.makedirs(result_dir)
# Load
assert len(opt.model_prefix_list) <= torch.cuda.device_count(), \
"number of models should not exceed the device numbers"
nets = []
for _, prefix in enumerate(opt.model_prefix_list):
net = model.SSD(cfg,
connect_loc_to_conf=True,
fix_size=False,
incep_conf=True,
incep_loc=True)
device_id = opt.device_id if len(opt.model_prefix_list) == 1 else _
net = net.to("cuda:%d" % (device_id))
net_dict = net.state_dict()
weight_dict = util.load_latest_model(args,
net,
prefix=prefix,
return_state_dict=True,
nth=opt.nth_best_model)
loading_fail_signal = False
for i, key in enumerate(net_dict.keys()):
if "module." + key not in weight_dict:
net_dict[key] = torch.zeros(net_dict[key].shape)
for key in weight_dict.keys():
if key[7:] in net_dict:
if net_dict[key[7:]].shape == weight_dict[key].shape:
net_dict[key[7:]] = weight_dict[key]
else:
print(
"Key: %s from disk has shape %s copy to the model with shape %s"
% (key[7:], str(weight_dict[key].shape),
str(net_dict[key[7:]].shape)))
loading_fail_signal = True
else:
print("Key: %s does not exist in net_dict" % (key[7:]))
if loading_fail_signal:
raise RuntimeError(
'Shape Error happens, remove "%s" from your -mpl settings.' %
(prefix))
net.load_state_dict(net_dict)
net.eval()
nets.append(net.half())
print("Above model loaded with out a problem")
detector = model.Detect(num_classes=2,
bkg_label=0,
top_k=opt.detector_top_k,
conf_thresh=opt.detector_conf_threshold,
nms_thresh=opt.detector_nms_threshold)
# Enumerate test folder
root_path = os.path.expanduser(opt.test_dataset_root)
if not os.path.exists(root_path):
raise FileNotFoundError(
"%s does not exists, please check your -tdr/--test_dataset_root settings"
% (root_path))
img_list = glob.glob(root_path + "/*.%s" % (opt.extension))
precisions, recalls = [], []
for i, img_file in enumerate(sorted(img_list)):
start = time.time()
name = img_file[img_file.rfind("/") + 1:-4]
img = cv2.imread(img_file)
height_ori, width_ori = img.shape[0], img.shape[1]
# detect rotation for returning the image back
transform_det = {"rotation": 0}
# Resize the longer side to a certain length
if height_ori >= width_ori:
resize_aug = augmenters.Sequential([
augmenters.Resize(size={
"height": opt.test_size,
"width": "keep-aspect-ratio"
})
])
else:
resize_aug = augmenters.Sequential([
augmenters.Resize(size={
"height": "keep-aspect-ratio",
"width": opt.test_size
})
])
resize_aug = resize_aug.to_deterministic()
image = resize_aug.augment_image(img)
h_re, w_re = image.shape[0], image.shape[1]
# Pad the image into a square image
pad_aug = augmenters.Sequential(
augmenters.PadToFixedSize(width=opt.test_size,
height=opt.test_size,
pad_cval=255,
position="center"))
pad_aug = pad_aug.to_deterministic()
image = pad_aug.augment_image(image)
h_final, w_final = image.shape[0], image.shape[1]
# Prepare image tensor and test
image_t = torch.Tensor(util.normalize_image(args, image)).unsqueeze(0)
image_t = image_t.permute(0, 3, 1, 2)
#visualize_bbox(args, cfg, image, [torch.Tensor(rot_coord).cuda()], net.prior, height_final/width_final)
text_boxes = []
for _, net in enumerate(nets):
device_id = opt.device_id if len(nets) == 1 else _
image_t = image_t.to("cuda:%d" % (device_id)).half()
out = net(image_t, is_train=False)
loc_data, conf_data, prior_data = out
prior_data = prior_data.to("cuda:%d" % (device_id))
det_result = detector(loc_data, conf_data, prior_data)
# Extract the predicted bboxes
idx = det_result.data[0, 1, :, 0] >= 0.1
text_boxes.append(det_result.data[0, 1, idx, 1:])
text_boxes = torch.cat(text_boxes, dim=0)
text_boxes = combine_boxes(text_boxes, img=image_t)
pred = [[float(coor) for coor in area] for area in text_boxes]
BBox = [
imgaug.augmentables.bbs.BoundingBox(box[0] * w_final,
box[1] * h_final,
box[2] * w_final,
box[3] * h_final)
for box in pred
]
BBoxes = imgaug.augmentables.bbs.BoundingBoxesOnImage(
BBox, shape=image.shape)
return_aug = augment_back(transform_det, height_ori, width_ori,
(h_final - h_re) / 2, (w_final - w_re) / 2)
return_aug = return_aug.to_deterministic()
img_ori = return_aug.augment_image(image)
bbox = return_aug.augment_bounding_boxes([BBoxes])[0]
f = open(os.path.join(result_dir, name + ".txt"), "w")
pred_final = []
for box in bbox.bounding_boxes:
x1, y1, x2, y2 = int(round(box.x1)), int(round(box.y1)), int(
round(box.x2)), int(round(box.y2))
pred_final.append([x1, y1, x2, y2])
#box_tensors.append(torch.tensor([x1, y1, x2, y2]))
# 4-point to 8-point: x1, y1, x2, y1, x2, y2, x1, y2
f.write("%d,%d,%d,%d,%d,%d,%d,%d\n" %
(x1, y1, x2, y1, x2, y2, x1, y2))
cv2.rectangle(img, (x1, y1), (x2, y2), (255, 105, 65), 2)
#accu, precision, recall = measure(torch.Tensor(pred_final).cuda(), torch.Tensor(gt_coords).cuda(),
#width=img.shape[1], height=img.shape[0])
img_save_directory = os.path.join(
args.path, args.code_name,
"val+" + "-".join(opt.model_prefix_list))
if not os.path.exists(img_save_directory):
os.mkdir(img_save_directory)
_imgh, _imgw, _imgc = img.shape
_imgh = _imgh * opt.test_size / _imgw
img = cv2.resize(img, (opt.test_size, int(_imgh)))
#cv2.imwrite(os.path.join(img_save_directory, name + ".jpg"), img)
cv2.imwrite(os.path.join(img_save_directory, "%04d.jpg" % i), img)
f.close()
print("%d th image cost %.2f seconds" % (i, time.time() - start))
}
def init_weight(m):
if type(m) == nn.Linear:
torch.nn.init.xavier_normal_(m.weight)
#torch.nn.init.xavier_uniform_(m.bias)
if type(m) == nn.Conv2d:
torch.nn.init.xavier_normal_(m.weight)
#torch.nn.init.kaiming_uniform_(m.bias)
if __name__ == "__main__":
args = util.get_args(presets.PRESET)
with torch.cuda.device(2):
net = model.CifarNet_Vanilla()
if args.finetune:
net = util.load_latest_model(args, net)
else:
#net.apply(init_weight)
keras_model = get_keras_model()
model_path = os.path.join(os.getcwd(), 'test', 'models', "cifar10_cnn.h5")
net = weight_transfer.initialize_with_keras_hdf5(keras_model, map_dict, net, model_path)
omth_util.save_model(args, args.curr_epoch, net.state_dict())
#net.to(args.device)
net.cuda()
#summary(net, input_size=(3, 32, 32), device=device)
#train_set = fetch_data(args, [("data_batch_1", "data_batch_2", "data_batch_3", "data_batch_4", "data_batch_5")])
#test_set = fetch_data(args, ["test_batch"])
transform = transforms.Compose([
transforms.ToTensor(),
def main():
aug = aug_temp(args)
dt = datetime.datetime.now().strftime("%Y-%m-%d_%H:%M")
datasets = data.fetch_detection_data(args,
sources=args.train_sources,
k_fold=1,
batch_size=args.batch_size_per_gpu,
batch_size_val=1,
auxiliary_info=args.train_aux,
split_val=0.1,
aug=aug)
for idx, (train_set, val_set) in enumerate(datasets):
loc_loss, conf_loss = [], []
accuracy, precision, recall, f1_score = [], [], [], []
print("\n =============== Cross Validation: %s/%s ================ " %
(idx + 1, len(datasets)))
net = model.SSD(cfg,
connect_loc_to_conf=args.loc_to_conf,
fix_size=args.fix_size,
conf_incep=args.conf_incep,
loc_incep=args.loc_incep,
nms_thres=args.nms_threshold,
loc_preconv=args.loc_preconv,
conf_preconv=args.conf_preconv,
FPN=args.feature_pyramid_net,
SA=args.self_attention,
in_wid=args.inner_filters,
m_factor=args.inner_m_factor)
net = torch.nn.DataParallel(net,
device_ids=args.gpu_id,
output_device=args.output_gpu_id).cuda()
detector = model.Detect(num_classes=2,
bkg_label=0,
top_k=800,
conf_thresh=0.05,
nms_thresh=0.3)
#detector = None
# Input dimension of bbox is different in each step
torch.backends.cudnn.benchmark = True
if args.fix_size:
net.module.prior = net.module.prior.cuda()
if args.finetune:
net = util.load_latest_model(args,
net,
prefix=args.model_prefix_finetune,
strict=True)
# Using the latest optimizer, better than Adam and SGD
optimizer = AdaBound(
net.parameters(),
lr=args.learning_rate,
final_lr=20 * args.learning_rate,
weight_decay=args.weight_decay,
)
for epoch in range(args.epoch_num):
loc_avg, conf_avg = fit(args,
cfg,
net,
detector,
train_set,
optimizer,
is_train=True)
loc_loss.append(loc_avg)
conf_loss.append(conf_avg)
train_losses = [np.asarray(loc_loss), np.asarray(conf_loss)]
if val_set is not None:
accu, pre, rec, f1 = fit(args,
cfg,
net,
detector,
val_set,
optimizer,
is_train=False)
accuracy.append(accu)
precision.append(pre)
recall.append(rec)
f1_score.append(f1)
val_losses = [
np.asarray(accuracy),
np.asarray(precision),
np.asarray(recall),
np.asarray(f1_score)
]
if epoch != 0 and epoch % 10 == 0:
util.save_model(args,
args.curr_epoch,
net.state_dict(),
prefix=args.model_prefix,
keep_latest=3)
if epoch > 5:
vb.plot_multi_loss_distribution(
multi_line_data=[train_losses, val_losses],
multi_line_labels=[["location", "confidence"],
[
"Accuracy", "Precision", "Recall",
"F1-Score"
]],
save_path=args.loss_log,
window=5,
name=dt,
bound=[{
"low": 0.0,
"high": 3.0
}, {
"low": 0.0,
"high": 1.0
}],
titles=["Train Loss", "Validation Score"])
# Clean the data for next cross validation
del net, optimizer
args.curr_epoch = 0
def main():
dt = datetime.datetime.now().strftime("%Y-%m-%d_%H:%M")
datasets = data.fetch_probaV_data(args,
sources=args.train_sources,
k_fold=args.cross_val,
split_val=0.1,
batch_size=args.batch_size_per_gpu,
auxiliary_info=[2, 2])
for idx, (train_set, val_set) in enumerate(datasets):
Loss, Measure = [], []
val_Loss, val_Measure = [], []
print("\n =============== Cross Validation: %s/%s ================ " %
(idx + 1, len(datasets)))
if args.which_model.lower() == "carn":
net = model.CARN(args.n_selected_img,
args.filters,
3,
s_MSE=args.s_MSE,
trellis=args.trellis)
elif args.which_model.lower() == "rdn":
net = model.RDN(args.n_selected_img,
3,
3,
filters=args.filters,
s_MSE=args.s_MSE,
group=args.n_selected_img,
trellis=args.trellis)
elif args.which_model.lower() == "meta_rdn":
net = RDN_Meta(args.n_selected_img,
filters=args.filters,
scale=3,
s_MSE=args.s_MSE,
group=args.n_selected_img,
trellis=args.trellis)
elif args.which_model.lower() == "basic":
net = model.ProbaV_basic(inchannel=args.n_selected_img)
else:
print(
"args.which_model or -wm should be one of [carn, rdn, basic], "
"your -wm %s is illegal, and switched to 'basic' automatically"
% (args.which_model.lower()))
net = model.ProbaV_basic(inchannel=args.n_selected_img)
net.apply(init_cnn)
if args.half_precision:
net.half()
net = torch.nn.DataParallel(net,
device_ids=args.gpu_id,
output_device=args.output_gpu_id).cuda()
torch.backends.cudnn.benchmark = True
if args.finetune:
net = util.load_latest_model(args,
net,
prefix=args.model_prefix_finetune,
strict=True)
optimizer = AdaBound(net.parameters(),
lr=args.learning_rate,
final_lr=10 * args.learning_rate,
weight_decay=args.weight_decay)
#criterion = ListedLoss(type="l1", reduction="mean")
#criterion = torch.nn.DataParallel(criterion, device_ids=args.gpu_id, output_device=args.output_gpu_id).cuda()
measure = MultiMeasure(type="l2",
reduction="mean",
half_precision=args.half_precision)
#measure = torch.nn.DataParallel(measure, device_ids=args.gpu_id, output_device=args.output_gpu_id).cuda()
for epoch in range(args.epoch_num):
_l, _m = fit(args,
net,
train_set,
optimizer,
measure,
is_train=True)
Loss.append(_l)
Measure.append(_m)
if val_set is not None:
_vl, _vm = val(args, net, val_set, optimizer, measure)
val_Loss.append(_vl)
val_Measure.append(_vm)
if (epoch + 1) % 10 == 0:
util.save_model(args,
args.curr_epoch,
net.state_dict(),
prefix=args.model_prefix,
keep_latest=10)
if (epoch + 1) > 5:
vb.plot_multi_loss_distribution(
multi_line_data=[
to_array(Loss) + to_array(val_Loss),
to_array(Measure) + to_array(val_Measure)
],
multi_line_labels=[[
"train_mae", "train_smse", "val_mae", "val_smse"
], [
"train_PSNR",
"train_L1",
"val_PSNR",
"val_L1",
]],
save_path=args.loss_log,
window=3,
name=dt + "cv_%d" % (idx + 1),
bound=[{
"low": 0.0,
"high": 15
}, {
"low": 10,
"high": 50
}],
titles=["Loss", "Measure"])
# Clean the data for next cross validation
del net, optimizer, measure
args.curr_epoch = 0
def test_rotation(opt):
result_dir = os.path.join(args.path, args.code_name,
"result+" + "-".join(opt.model_prefix_list))
if not os.path.exists(result_dir):
os.makedirs(result_dir)
# Load
if type(opt.model_prefix_list) is str:
opt.model_prefix_list = [opt.model_prefix_list]
assert len(opt.model_prefix_list) <= torch.cuda.device_count(), \
"number of models should not exceed the device numbers"
nets = []
for _, prefix in enumerate(opt.model_prefix_list):
#net = model.SSD(cfg, connect_loc_to_conf=True, fix_size=False,
#conf_incep=True, loc_incep=True)
net = model.SSD(cfg,
connect_loc_to_conf=opt.loc_to_conf,
fix_size=False,
conf_incep=opt.conf_incep,
loc_incep=opt.loc_incep,
loc_preconv=opt.loc_preconv,
conf_preconv=opt.conf_preconv,
FPN=opt.feature_pyramid_net,
SA=opt.self_attention,
in_wid=opt.inner_filters,
m_factor=opt.inner_m_factor)
device_id = opt.device_id if len(opt.model_prefix_list) == 1 else _
net = net.to("cuda:%d" % (device_id))
net_dict = net.state_dict()
weight_dict = util.load_latest_model(args,
net,
prefix=prefix,
return_state_dict=True,
nth=opt.nth_best_model)
loading_fail_signal = False
for i, key in enumerate(net_dict.keys()):
if "module." + key not in weight_dict:
net_dict[key] = torch.zeros(net_dict[key].shape)
for key in weight_dict.keys():
if key[7:] in net_dict:
if net_dict[key[7:]].shape == weight_dict[key].shape:
net_dict[key[7:]] = weight_dict[key]
else:
print(
"Key: %s from disk has shape %s copy to the model with shape %s"
% (key[7:], str(weight_dict[key].shape),
str(net_dict[key[7:]].shape)))
loading_fail_signal = True
else:
print("Key: %s does not exist in net_dict" % (key[7:]))
if loading_fail_signal:
raise RuntimeError(
'Shape Error happens, remove "%s" from your -mpl settings.' %
(prefix))
net.load_state_dict(net_dict)
net.eval()
nets.append(net)
print("Above model loaded with out a problem")
detector = model.Detect(num_classes=2,
bkg_label=0,
top_k=opt.detector_top_k,
conf_thresh=opt.detector_conf_threshold,
nms_thresh=opt.detector_nms_threshold)
# Enumerate test folder
root_path = os.path.expanduser(opt.test_dataset_root)
if not os.path.exists(root_path):
raise FileNotFoundError(
"%s does not exists, please check your -tdr/--test_dataset_root settings"
% (root_path))
img_list = glob.glob(root_path + "/*.%s" % (opt.extension))
precisions, recalls = [], []
for i, img_file in enumerate(sorted(img_list)):
start = time.time()
name = img_file[img_file.rfind("/") + 1:-4]
img = cv2.imread(img_file)
height_ori, width_ori = img.shape[0], img.shape[1]
do_rotation = False
if do_rotation:
# detect rotation for returning the image back
img, transform_det = estimate_angle(img, args, None, None, None)
transform_det["rotation"] = 0
if transform_det["rotation"] != 0:
rot_aug = augmenters.Affine(rotate=transform_det["rotation"],
cval=args.aug_bg_color)
else:
rot_aug = None
# Perform Augmentation
if rot_aug:
rot_aug = augmenters.Sequential(
augmenters.Affine(rotate=transform_det["rotation"],
cval=args.aug_bg_color))
image = rot_aug.augment_image(img)
else:
image = img
else:
image = img
# Resize the longer side to a certain length
if height_ori >= width_ori:
resize_aug = augmenters.Sequential([
augmenters.Resize(size={
"height": square,
"width": "keep-aspect-ratio"
})
])
else:
resize_aug = augmenters.Sequential([
augmenters.Resize(size={
"height": "keep-aspect-ratio",
"width": square
})
])
resize_aug = resize_aug.to_deterministic()
image = resize_aug.augment_image(image)
h_re, w_re = image.shape[0], image.shape[1]
# Pad the image into a square image
pad_aug = augmenters.Sequential(
augmenters.PadToFixedSize(width=square,
height=square,
pad_cval=255,
position="center"))
pad_aug = pad_aug.to_deterministic()
image = pad_aug.augment_image(image)
h_final, w_final = image.shape[0], image.shape[1]
# Prepare image tensor and test
image_t = torch.Tensor(util.normalize_image(args, image)).unsqueeze(0)
image_t = image_t.permute(0, 3, 1, 2)
# visualize_bbox(args, cfg, image, [torch.Tensor(rot_coord).cuda()], net.prior, height_final/width_final)
text_boxes = []
for _, net in enumerate(nets):
device_id = opt.device_id if len(nets) == 1 else _
image_t = image_t.to("cuda:%d" % (device_id))
out = net(image_t, is_train=False)
loc_data, conf_data, prior_data = out
prior_data = prior_data.to("cuda:%d" % (device_id))
det_result = detector(loc_data, conf_data, prior_data)
# Extract the predicted bboxes
idx = det_result.data[0, 1, :, 0] >= 0.1
text_boxes.append(det_result.data[0, 1, idx, 1:])
text_boxes = torch.cat(text_boxes, dim=0)
text_boxes = combine_boxes(text_boxes, img=image_t)
pred = [[float(coor) for coor in area] for area in text_boxes]
BBox = [
imgaug.augmentables.bbs.BoundingBox(box[0] * w_final,
box[1] * h_final,
box[2] * w_final,
box[3] * h_final)
for box in pred
]
BBoxes = imgaug.augmentables.bbs.BoundingBoxesOnImage(
BBox, shape=image.shape)
return_aug = augment_back(height_ori, width_ori, (h_final - h_re) / 2,
(w_final - w_re) / 2)
return_aug = return_aug.to_deterministic()
img_ori = return_aug.augment_image(image)
bbox = return_aug.augment_bounding_boxes([BBoxes])[0]
# print_box(blue_boxes=pred, idx=i, img=vb.plot_tensor(args, image_t, margin=0),
# save_dir=args.val_log)
f = open(os.path.join(result_dir, name + ".txt"), "w")
gt_box_file = os.path.join(opt.test_dataset_root,
name + "." + opt.ground_truth_extension)
coords = tb_data.parse_file(os.path.expanduser(gt_box_file))
gt_coords = []
for coord in coords:
x1, x2 = min(coord[::2]), max(coord[::2])
y1, y2 = min(coord[1::2]), max(coord[1::2])
gt_coords.append([x1, y1, x2, y2])
cv2.rectangle(img, (x1, y1), (x2, y2), (70, 67, 238), 2)
pred_final = []
for box in bbox.bounding_boxes:
x1, y1, x2, y2 = int(round(box.x1)), int(round(box.y1)), int(
round(box.x2)), int(round(box.y2))
pred_final.append([x1, y1, x2, y2])
# box_tensors.append(torch.tensor([x1, y1, x2, y2]))
# 4-point to 8-point: x1, y1, x2, y1, x2, y2, x1, y2
f.write("%d,%d,%d,%d,%d,%d,%d,%d\n" %
(x1, y1, x2, y1, x2, y2, x1, y2))
cv2.rectangle(img, (x1, y1), (x2, y2), (255, 105, 65), 2)
accu, precision, recall = measure(torch.Tensor(pred_final).cuda(),
torch.Tensor(gt_coords).cuda(),
width=img.shape[1],
height=img.shape[0])
precisions.append(precision)
recalls.append(recall)
img_save_directory = os.path.join(
args.path, args.code_name,
"val+" + "-".join(opt.model_prefix_list))
if not os.path.exists(img_save_directory):
os.mkdir(img_save_directory)
name = "%s_%.2f_%.2f" % (str(i).zfill(4), precision, recall)
cv2.imwrite(os.path.join(img_save_directory, name + ".jpg"), img)
f.close()
if opt.verbose:
print(
"%d th image cost %.2f seconds, precision: %.2f, recall: %.2f"
% (i, time.time() - start, precision, recall))
print("Precision: %.2f, Recall: %.2f" % (avg(precisions), avg(recalls)))