'r-')
plt.show()
for k in output.keys():
if isinstance(output[k], torch.Tensor):
output[k] = output[k].tolist()
results.append(output)
outpath_dataset = osp.join(output_dir, '{}.json'.format(name))
logger.info('Writing the results of the {} dataset into {}'.format(
name, outpath_dataset))
with open(outpath_dataset, 'w') as _out:
json.dump(results, _out)
if __name__ == "__main__":
cfg.merge_from_file(args.config_file)
cfg.merge_from_list(args.opts)
cfg.freeze()
output_dir = cfg.OUTPUT_DIR
logger = setup_logger('hawp', output_dir, out_file='test.log')
logger.info(args)
logger.info("Loaded configuration file {}".format(args.config_file))
test(cfg)
### Training
# import pdb; pdb.set_trace()
def test_building():
parser = argparse.ArgumentParser()
parser.add_argument("--config-file",
metavar="FILE",
help="path to config file",
type=str,
required=True,
)
parser.add_argument("--display",
default=False,
action='store_true')
parser.add_argument("opts",
help="Modify config options using the command-line",
default=None,
nargs=argparse.REMAINDER
)
parser.add_argument("--checkpoint", help="pytorch checkpoint")
parser.add_argument("--hawp_checkpoint", help="pytorch checkpoint", default="")
parser.add_argument("--fp16", action="store_true", help="training in fp16 mode")
parser.add_argument("--out_fname", help="output file name", type=str, required=True)
args = parser.parse_args()
cfg.merge_from_file(args.config_file)
cfg.merge_from_list(args.opts)
cfg.freeze()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = RFJunctionDetector(cfg).to(device)
test_dataset = build_building_test_dataset(cfg)
if args.checkpoint:
checkpoint = torch.load(args.checkpoint)
model.load_state_dict(checkpoint['model_state_dict'])
model = model.eval()
# hawp_model = WireframeDetector(cfg, test_info='junc')
# hawp_model = hawp_model.to(device)
# logger = logging.getLogger("hawp.testing")
# checkpointer = DetectronCheckpointer(cfg,
# hawp_model,
# save_dir=args.hawp_checkpoint,
# save_to_disk=True,
# logger=logger)
# _ = checkpointer.load()
# hawp_model = hawp_model.eval()
softmax_fn = torch.nn.Softmax(dim=0)
image_json = []
for i, (images, target) in enumerate(tqdm(test_dataset)):
# if test_dataset.dataset.filename(i) != 'building-180721-10001.21.27_pos_p655':
# continue
with torch.no_grad():
output, _ = model(images.to(device))
output = to_device(output,'cpu').squeeze(dim=0)
target = target.squeeze(dim=0)
if args.display:
output = softmax_fn(output)
fig = plt.figure(figsize=(24, 8))
# 01 origin
fig.add_subplot(1, 3, 1)
plt.imshow(((images - images.min()) * (255 / (images.max() - images.min()))).squeeze().permute(1, 2, 0).int())
plt.title('original image')
# 02 prediction
fig.add_subplot(1, 3, 2)
plt.imshow((output[1, :, :] * 255 / output[1, :, :].max()).int(), cmap='gray', vmin=0, vmax=255)
plt.title('prediction heatmap')
# 03 gt
fig.add_subplot(1, 3, 3)
plt.imshow((target * 255 / target.max()).int().permute(1, 2, 0), cmap='gray', vmin=0, vmax=255)
plt.title('ground truth')
# plt.show()
if not os.path.exists("./deform_attn_test"):
os.makedirs("./deform_attn_test")
save_dir = "./deform_attn_test/{}/".format(args.out_fname)
if not os.path.exists(save_dir):
os.makedirs(save_dir)
plt.savefig(save_dir + "{}.png".format(test_dataset.dataset.filename(i)))
plt.cla()
# pr curve
target = target.round()
binarize = lambda x, th: (x > th).int()
tp = lambda pred, gt: (pred * gt).sum() # [1, 1]
fp = lambda pred, gt: (pred * (1 - gt)).sum() # [1, 0]
# tn = lambda pred, gt: ((1 - pred) * (1 - gt)).sum() # [0, 0]
fn = lambda pred, gt: ((1 - pred) * gt).sum() # [0, 1]
precision = lambda pred, gt: (tp(pred, gt) / (tp(pred, gt) + fp(pred, gt))).numpy().item()
recall = lambda pred, gt: (tp(pred, gt) / (tp(pred, gt) + fn(pred, gt))).numpy().item()
rasters = {i/100: binarize(output[1, :, :], i/100) for i in range(1, 100)}
ps = {k: precision(v, target) for k, v in rasters.items() if not math.isnan(precision(v, target)) }
rs = {k: recall(v, target) for k, v in rasters.items() if not math.isnan(recall(v, target)) }
image_json.append({
'filename': test_dataset.dataset.filename(i),
'precision': ps,
'recall': rs
})
if args.display:
with open('deform_attn_test/{}.json'.format(args.out_fname), 'w') as f:
json.dump(image_json, f)
def predict():
parser = argparse.ArgumentParser()
parser.add_argument("--config-file",
metavar="FILE",
help="path to config file",
type=str,
required=True,
)
parser.add_argument("--display",
default=False,
action='store_true')
parser.add_argument("opts",
help="Modify config options using the command-line",
default=None,
nargs=argparse.REMAINDER
)
parser.add_argument("--checkpoint", help="pytorch checkpoint")
parser.add_argument("--ann_path", help="annotation path")
parser.add_argument("--im_path", help="image path")
parser.add_argument("--fp16", action="store_true", help="training in fp16 mode")
args = parser.parse_args()
cfg.merge_from_file(args.config_file)
cfg.merge_from_list(args.opts)
cfg.freeze()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = RFJunctionDetector(cfg).to(device)
if args.checkpoint:
checkpoint = torch.load(args.checkpoint)
model.load_state_dict(checkpoint['model_state_dict'], strict=False)
transform = build_transform(cfg)
image = io.imread(args.im_path)[:,:,:3]
image_tensor = transform(image.astype(float))[None].to(device)
meta = {
'filename': args.im_path,
'height': image.shape[0],
'width': image.shape[1],
}
model = model.eval()
with torch.no_grad():
output = model(image_tensor)
output = to_device(output,'cpu').squeeze(dim=0)
# target = target.squeeze(dim=0)
if args.display:
threshold = 0.1
fig = plt.figure(figsize=(8, 24))
fig.add_subplot(1, 4, 1)
plt.imshow(((image_tensor - image_tensor.min()) * (255 / (image_tensor.max() - image_tensor.min()))).squeeze().permute(1, 2, 0).int())
fig.add_subplot(1, 4, 2)
image_copy = ((image_tensor - image_tensor.min()) * (255 / (image_tensor.max() - image_tensor.min()))).squeeze()
# image_copy = image_copy.numpy()
# heat_idx = (output > threshold).nonzero().numpy()
bool_map = (output < threshold).int()
image_copy[0, ...] = image_copy[0, ...] * bool_map
image_copy[1, ...] = image_copy[1, ...] * bool_map
output[output < threshold] = 0
image_copy[0, ...] = image_copy[0, ...] + output * (255 / output.max())
image_copy[1, ...] = image_copy[1, ...] + output * (255 / output.max())
# image_copy[heat_idx] = output[heat_idx] * 255
plt.imshow(image_copy.permute(1, 2, 0).int())
fig.add_subplot(1, 4, 3)
plt.imshow((output * 255).int().permute(1, 2, 0))
fig.add_subplot(1, 4, 4)
plt.imshow((target * 255).int().permute(1, 2, 0))
plt.show()
plt.savefig("./test_fig.png")
def train_rf():
parser = argparse.ArgumentParser()
parser.add_argument(
"--config-file",
metavar="FILE",
help="path to config file",
type=str,
required=True,
)
parser.add_argument("opts",
help="Modify config options using the command-line",
default=None,
nargs=argparse.REMAINDER)
parser.add_argument("--checkpoint", help="pytorch checkpoint")
parser.add_argument("--resume_train",
action="store_true",
help="load checkpoint and resume training")
parser.add_argument("--fp16",
action="store_true",
help="training in fp16 mode")
args = parser.parse_args()
cfg.merge_from_file(args.config_file)
cfg.merge_from_list(args.opts)
cfg.freeze()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = RFJunctionDetector(cfg).to(device)
train_dataset = build_train_dataset(cfg)
if args.checkpoint and not args.resume_train:
checkpoint = torch.load(args.checkpoint)
filtered_weights = {
k: v for k, v \
in checkpoint['model_state_dict'].items() \
if "mflow_conv_g6_b3_joint_drop_conv_new_2" not in k
}
model.backbone.load_state_dict(filtered_weights, strict=False)
elif args.checkpoint:
checkpoint = torch.load(args.checkpoint)
model.load_state_dict(checkpoint['model_state_dict'])
params_res = []
params_refine = []
for k, v in model.named_parameters():
if k.split('.')[1].startswith(("conv1", "bn_conv1", "res", "bn")):
params_res.append(v)
else:
params_refine.append(v)
optimizer = torch.optim.Adam([{
"params": params_res,
"lr": 0.00000003
}, {
"params": params_refine,
"lr": 0.0000003
}],
betas=(0.9, 0.999),
weight_decay=0.0005)
if args.checkpoint and args.resume_train:
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
model.train()
if args.fp16:
scaler = torch.cuda.amp.GradScaler()
if args.checkpoint:
scaler.load_state_dict(checkpoint["scaler_state_dict"])
# loss_fn = torch.nn.CrossEntropyLoss()
loss_fn = torch.nn.MSELoss(reduce=True, size_average=True)
total_loss = 0.0
t = 0 if not args.resume_train else checkpoint["iter"]
for epoch in range(1000):
for it, (images, target) in enumerate(train_dataset):
images = images.to(device)
target = target.to(device)
# import matplotlib.pyplot as plt
# fig = plt.figure(figsize=(8, 16))
# fig.add_subplot(1, 2, 1)
# plt.imshow(images[0, ...].squeeze().permute(1, 2, 0))
# fig.add_subplot(1, 2, 2)
# plt.imshow(target[0, ...].squeeze())
# plt.show()
if args.fp16:
with torch.cuda.amp.autocast():
y = model(images)
loss = loss_fn(y, target.half())
else:
y = model(images)
loss = loss_fn(y, target)
total_loss += loss.item()
if (t + 1) % 100 == 0:
print("iter: {}, loss: {}".format(t + 1, total_loss / 100))
total_loss = 0.0
optimizer.zero_grad()
if args.fp16:
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
else:
loss.backward()
optimizer.step()
if t % 10000 == 0 and t != 0:
torch.save(
{
'iter': t,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
# 'scaler_state_dict': scaler.state_dict(),
'loss': loss,
},
"checkpoint/checkpoint_finetune_iter_{}".format(t))
t += 1
def train_building():
parser = argparse.ArgumentParser()
parser.add_argument(
"--config-file",
metavar="FILE",
help="path to config file",
type=str,
required=True,
)
parser.add_argument("opts",
help="Modify config options using the command-line",
default=None,
nargs=argparse.REMAINDER)
parser.add_argument("--checkpoint", help="checkpoint to be loaded")
parser.add_argument("--checkpoint_dir", help="checkpoint dir")
parser.add_argument("--checkpoint_gap",
default=10000,
type=int,
help="how many iters to save checkpoint")
parser.add_argument("--resume_train",
action="store_true",
help="load checkpoint and resume training")
parser.add_argument("--fp16",
action="store_true",
help="training in fp16 mode")
parser.add_argument("--lr_base",
default=1e-5,
type=float,
help="refine net base learning rate")
parser.add_argument("--lr_gap",
default=100,
type=int,
help="how many iters to print loss information")
parser.add_argument("--w_ctl",
default=1,
type=float,
help="contrastive loss weight")
args = parser.parse_args()
cfg.merge_from_file(args.config_file)
cfg.merge_from_list(args.opts)
cfg.freeze()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = RFJunctionDetector(cfg).to(device)
train_dataset = build_building_train_dataset(cfg)
if args.checkpoint and not args.resume_train:
checkpoint = torch.load(args.checkpoint)
filtered_weights = {
k: v for k, v \
in checkpoint['model_state_dict'].items() \
if "mflow_conv_g6_b3_joint_drop_conv_new_2" not in k
}
model.backbone.load_state_dict(filtered_weights, strict=False)
elif args.checkpoint:
checkpoint = torch.load(args.checkpoint)
model.load_state_dict(checkpoint['model_state_dict'])
params_res = []
params_refine = []
for k, v in model.named_parameters():
if k.split('.')[1].startswith(("conv1", "bn_conv1", "res", "bn")):
params_res.append(v)
else:
params_refine.append(v)
optimizer = torch.optim.Adam([{
"params": params_res,
"lr": 3 * args.lr_base / 10
}, {
"params": params_refine,
"lr": 3 * args.lr_base
}],
betas=(0.9, 0.999),
weight_decay=0.0005)
if args.checkpoint and args.resume_train:
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
model.train()
if args.fp16:
scaler = torch.cuda.amp.GradScaler()
if args.checkpoint:
scaler.load_state_dict(checkpoint["scaler_state_dict"])
loss_fn = torch.nn.CrossEntropyLoss()
# loss_fn = torch.nn.MSELoss(reduce=True, size_average=True)
total_loss, total_loss_main, total_loss_aux = 0.0, 0.0, 0.0
t = 0 if not args.resume_train else checkpoint["iter"]
if not os.path.exists(args.checkpoint_dir):
os.makedirs(args.checkpoint_dir)
with torch.autograd.set_detect_anomaly(True):
for epoch in range(1000):
for it, (images, target) in enumerate(train_dataset):
images = images.to(device)
target = target.to(device).long()
# import matplotlib.pyplot as plt
# fig = plt.figure(figsize=(8, 16))
# fig.add_subplot(1, 2, 1)
# plt.imshow(images[0, ...].squeeze().permute(1, 2, 0))
# fig.add_subplot(1, 2, 2)
# plt.imshow(target[0, ...].squeeze())
# plt.show()
if args.fp16:
with torch.cuda.amp.autocast():
y, contrastive_loss = model(images)
loss = loss_fn(y, target.half())
# contrastive_loss = sum(contrastive_loss)
else:
y, contrastive_loss = model(images)
contrastive_loss = sum([
torch.mean(c)
for c in contrastive_loss if not isinstance(c, int)
]) * 1e-3
loss_main = loss_fn(y, target.squeeze(1))
loss = loss_main + contrastive_loss * args.w_ctl
total_loss += loss.item()
total_loss_main += loss_main.item()
total_loss_aux += contrastive_loss.item(
) if cfg.MODEL.ATTN_USE_CTL else 0
if (t + 1) % args.lr_gap == 0:
print(
"iter: {}, loss_total: {}, loss_main: {}, loss_aux: {}"
.format(t + 1, total_loss / args.lr_gap,
total_loss_main / args.lr_gap,
total_loss_aux / args.lr_gap))
total_loss = 0.0
total_loss_main = 0.0
total_loss_aux = 0.0
optimizer.zero_grad()
if args.fp16:
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
else:
loss.backward()
optimizer.step()
if t % args.checkpoint_gap == 0 and t != 0:
torch.save(
{
'iter': t,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
# 'scaler_state_dict': scaler.state_dict(),
'loss': loss,
},
"{}/checkpoint_finetune_iter_{}".format(
args.checkpoint_dir, t))
t += 1