def testBatch(options, model, dataset):
model.eval()
dataloader = DataLoader(dataset,
batch_size=1,
shuffle=False,
num_workers=16)
data_iterator = tqdm(dataloader, total=len(dataset))
for sampleIndex, sample in enumerate(data_iterator):
images, img_path = sample[0].cuda(), sample[1][0]
img_name = os.path.splitext(img_path)[0].split('/')[-1]
#if int(img_name) < 61: continue
corner_pred, icon_pred, room_pred = model(images)
visualizeBatch(
options,
images.detach().cpu().numpy(),
[('pred', {
'corner': corner_pred.max(-1)[1].detach().cpu().numpy(),
'icon': icon_pred.max(-1)[1].detach().cpu().numpy(),
'room': room_pred.max(-1)[1].detach().cpu().numpy()
})],
batch_img_name=img_name)
for batchIndex in range(len(images)):
corner_heatmaps = torch.sigmoid(
corner_pred[batchIndex]).detach().cpu().numpy()
icon_heatmaps = torch.sigmoid(
icon_pred[batchIndex]).detach().cpu().numpy()
room_heatmaps = torch.sigmoid(
room_pred[batchIndex]).detach().cpu().numpy()
reconstructFloorplan(
corner_heatmaps[:, :, :NUM_WALL_CORNERS],
corner_heatmaps[:, :, NUM_WALL_CORNERS:NUM_WALL_CORNERS + 8],
corner_heatmaps[:, :, -4:],
icon_heatmaps,
room_heatmaps,
output_prefix=options.test_dir + '/' + img_name + '_',
densityImage=None,
gt_dict=None,
gt=False,
gap=-1,
distanceThreshold=-1,
lengthThreshold=-1,
debug_prefix='test',
heatmapValueThresholdWall=None,
heatmapValueThresholdDoor=None,
heatmapValueThresholdIcon=None,
enableAugmentation=True)
continue
continue
model.train()
return
def main(img_path):
options = parse_args()
model = Model(options)
model.load_state_dict(
torch.load('checkpoint.pth', map_location=torch.device('cpu')))
corner_pred, icon_pred, room_pred = model(
torch.tensor(load_image(img_path)))
corner_heatmaps = corner_pred[0].detach().cpu().numpy()
icon_heatmaps = torch.nn.functional.softmax(icon_pred[0],
dim=-1).detach().cpu().numpy()
room_heatmaps = torch.nn.functional.softmax(room_pred[0],
dim=-1).detach().cpu().numpy()
wallCornerHeatmaps = corner_heatmaps[:, :, :NUM_WALL_CORNERS]
doorCornerHeatmaps = corner_heatmaps[:, :,
NUM_WALL_CORNERS:NUM_WALL_CORNERS + 4]
iconCornerHeatmaps = corner_heatmaps[:, :, -4:]
maps = {
'original': cv2.imread(img_path),
'corner_heatmaps': corner_heatmaps.max(-1),
'icon_heatmaps': icon_heatmaps.max(-1),
'room_heatmaps': room_heatmaps.max(-1),
'corner_pred':
np.squeeze(corner_pred.max(-1)[1].detach().cpu().numpy()),
'icon_pred': np.squeeze(icon_pred.max(-1)[1].detach().cpu().numpy()),
'room_pred': np.squeeze(room_pred.max(-1)[1].detach().cpu().numpy()),
'wallCornerHeatmaps': wallCornerHeatmaps.max(-1),
'doorCornerHeatmaps': doorCornerHeatmaps.max(-1),
'iconCornerHeatmaps': iconCornerHeatmaps.max(-1),
}
np_print(maps)
plot_images(maps)
reconstructFloorplan(wallCornerHeatmaps,
doorCornerHeatmaps,
iconCornerHeatmaps,
icon_heatmaps,
room_heatmaps,
output_prefix='output-',
densityImage=None,
gt_dict=None,
gt=False,
gap=-1,
distanceThreshold=-1,
lengthThreshold=-1,
debug_prefix='test',
heatmapValueThresholdWall=None,
heatmapValueThresholdDoor=None,
heatmapValueThresholdIcon=None,
enableAugmentation=True)
def testOneEpoch(options, model, dataset):
model.eval()
dataloader = DataLoader(dataset, batch_size=options.batchSize, shuffle=False, num_workers=1)
epoch_losses = []
data_iterator = tqdm(dataloader, total=len(dataset) // options.batchSize + 1)
for sampleIndex, sample in enumerate(data_iterator):
images, corner_gt, icon_gt, room_gt = sample[0].cuda(), sample[1].cuda(), sample[2].cuda(), sample[3].cuda()
corner_pred, icon_pred, room_pred = model(images)
corner_loss = torch.nn.functional.binary_cross_entropy(corner_pred, corner_gt)
icon_loss = torch.nn.functional.cross_entropy(icon_pred.view(-1, NUM_ICONS + 2), icon_gt.view(-1))
room_loss = torch.nn.functional.cross_entropy(room_pred.view(-1, NUM_ROOMS + 2), room_gt.view(-1))
losses = [corner_loss, icon_loss, room_loss]
loss = sum(losses)
loss_values = [l.data.item() for l in losses]
epoch_losses.append(loss_values)
status = 'val loss: '
for l in loss_values:
status += '%0.5f '%l
continue
data_iterator.set_description(status)
if sampleIndex % 500 == 0:
visualizeBatch(options, images.detach().cpu().numpy(), [('gt', {'corner': corner_gt.detach().cpu().numpy(), 'icon': icon_gt.detach().cpu().numpy(), 'room': room_gt.detach().cpu().numpy()}), ('pred', {'corner': corner_pred.max(-1)[1].detach().cpu().numpy(), 'icon': icon_pred.max(-1)[1].detach().cpu().numpy(), 'room': room_pred.max(-1)[1].detach().cpu().numpy()})])
for batchIndex in range(len(images)):
corner_heatmaps = corner_pred[batchIndex].detach().cpu().numpy()
icon_heatmaps = torch.nn.functional.softmax(icon_pred[batchIndex], dim=-1).detach().cpu().numpy()
room_heatmaps = torch.nn.functional.softmax(room_pred[batchIndex], dim=-1).detach().cpu().numpy()
reconstructFloorplan(corner_heatmaps[:, :, :NUM_WALL_CORNERS], corner_heatmaps[:, :, NUM_WALL_CORNERS:NUM_WALL_CORNERS + 4], corner_heatmaps[:, :, -4:], icon_heatmaps, room_heatmaps, output_prefix=options.test_dir + '/' + str(batchIndex) + '_', densityImage=None, gt_dict=None, gt=False, gap=-1, distanceThreshold=-1, lengthThreshold=-1, debug_prefix='test', heatmapValueThresholdWall=None, heatmapValueThresholdDoor=None, heatmapValueThresholdIcon=None, enableAugmentation=True)
continue
if options.visualizeMode == 'debug':
exit(1)
pass
continue
print('validation loss', np.array(epoch_losses).mean(0))
model.train()
return
def evaluateBatch(options, gt_dict=None, pred_dict=None):
datasetFlag = 1
if options.useCache != -1:
if options.loss != '5':
gt_dict = np.load(options.test_dir + '/dummy/gt_dict.npy')[()]
pred_dict = np.load(options.test_dir + '/dummy/pred_dict.npy')[()]
else:
gt_dict = np.load(options.test_dir.replace('loss5', 'loss0') + '/dummy/gt_dict.npy')[()]
pred_wc = np.load(options.test_dir.replace('loss5', 'loss0') + '/dummy/pred_dict.npy')[()]['corner'][:, :, :, :NUM_WALL_CORNERS]
pred_oc = np.load(options.test_dir.replace('loss5', 'loss1') + '/dummy/pred_dict.npy')[()]['corner'][:, :, :, NUM_WALL_CORNERS:NUM_WALL_CORNERS + 4]
pred_ic = np.load(options.test_dir.replace('loss5', 'loss2').replace('hybrid14', 'hybrid1') + '/dummy/pred_dict.npy')[()]['corner'][:, :, :, NUM_WALL_CORNERS + 4:NUM_WALL_CORNERS + 8]
pred_icon = np.load(options.test_dir.replace('loss5', 'loss3').replace('hybrid14', 'hybrid1') + '/dummy/pred_dict.npy')[()]['icon']
pred_room = np.load(options.test_dir.replace('loss5', 'loss4').replace('hybrid14', 'hybrid1') + '/dummy/pred_dict.npy')[()]['room']
pred_dict = {'corner': np.concatenate([pred_wc, pred_oc, pred_ic], axis=-1), 'icon': pred_icon, 'room': pred_room}
pass
if options.separateIconLoss:
pred_icon_separate = softmax(np.load(options.test_dir.replace('wsf', 'wsf_loss3') + '/dummy/pred_dict.npy')[()]['icon'])
pass
#pred_dict['icon'] = np.load(options.test_dir.replace('wsf', 'wsf_loss3').replace('hybrid1', 'hybrid14').replace('dataset_1', '') + '/dummy/pred_dict.npy')[()]['icon']
#pred_dict['corner'][:, :, :, NUM_WALL_CORNERS + 4:NUM_WALL_CORNERS + 8] = np.load(options.test_dir.replace('wsf', 'wsf_loss2') + '/dummy/pred_dict.npy')[()]['corner'][:, :, :, NUM_WALL_CORNERS + 4:NUM_WALL_CORNERS + 8]
#pass
if options.cornerLossType != 'mse':
threshold = np.ones((HEIGHT, WIDTH, 1)) * 0.5
else:
threshold = np.ones((HEIGHT, WIDTH, 1)) * 0.5# HEATMAP_SCALE / 2
pass
statisticsSum = {k: [0.0, 0.0, 0.0] for k in ['wall', 'door', 'icon', 'room', 'neighbor', 'neighbor_all']}
#print(pred_dict['corner'].max())
pred_wc = pred_dict['corner'][:, :, :, :NUM_WALL_CORNERS]
pred_oc = pred_dict['corner'][:, :, :, NUM_WALL_CORNERS:NUM_WALL_CORNERS + 4]
pred_ic = pred_dict['corner'][:, :, :, NUM_WALL_CORNERS + 4:NUM_WALL_CORNERS + 8]
if options.branches != '5':
pred_wc = sigmoid(pred_wc)
pred_oc = sigmoid(pred_oc)
pred_ic = sigmoid(pred_ic)
else:
threshold = np.ones((HEIGHT, WIDTH, 1)) * 0.3
pass
gt_wc = gt_dict['corner'][:, :, :, :NUM_WALL_CORNERS]
gt_oc = gt_dict['corner'][:, :, :, NUM_WALL_CORNERS:NUM_WALL_CORNERS + 4]
gt_ic = gt_dict['corner'][:, :, :, NUM_WALL_CORNERS + 4:NUM_WALL_CORNERS + 8]
names = []
for batchIndex in xrange(gt_dict['corner'].shape[0]):
#if batchIndex == 0:
#continue
#if options.branches == '4' and gt_dict['image_flags'][batchIndex] == 0:
if options.evaluateImage and gt_dict['image_flags'][batchIndex] == 0:
continue
density = np.minimum(gt_dict['density'][batchIndex] * 255, 255).astype(np.uint8)
density = np.stack([density, density, density], axis=2)
pred_icon = softmax(pred_dict['icon'][batchIndex])
pred_room = softmax(pred_dict['room'][batchIndex])
if options.separateIconLoss:
pred_icon[:, :, :-2] = pred_icon_separate[batchIndex][:, :, :-2]
#pred_icon = pred_icon_separate[batchIndex]
pass
if False:
#print('batch index', batchIndex)
cv2.imwrite(options.test_dir + '/' + str(batchIndex) + '_density.png', density)
#print('heatmap max value', pred_wc[batchIndex].max())
if datasetFlag in [0, 1, 4]:
cornerImage = drawSegmentationImage(np.concatenate([threshold, pred_wc[batchIndex]], axis=2), blackIndex=0)
cornerImage[cornerImage == 0] = density[cornerImage == 0]
cv2.imwrite(options.test_dir + '/' + str(batchIndex) + '_corner_pred.png', cornerImage)
cornerImage = drawSegmentationImage(np.concatenate([threshold, gt_wc[batchIndex]], axis=2), blackIndex=0)
cornerImage[cornerImage == 0] = density[cornerImage == 0]
cv2.imwrite(options.test_dir + '/' + str(batchIndex) + '_corner_gt.png', cornerImage)
pass
if False:
corner_heat = np.max(pred_wc[batchIndex], axis=-1)
#print('corner_shape', corner_heat.shape)
cmap = plt.get_cmap('jet')
corner_rgba_img = cmap(corner_heat)
corner_rgb_img = np.delete(corner_rgba_img, 3, 2)
#print('rgb_out', corner_rgb_img.shape, corner_rgb_img.max(), corner_rgb_img.min())
corner_rgb_img = (corner_rgb_img * 255).round().astype('uint8')
#print('rgb_out', corner_rgb_img.shape, corner_rgb_img.max())
cv2.imwrite(options.test_dir + '/' + str(batchIndex) + '_corner_heatmap.png', corner_rgb_img)
pass
if datasetFlag in [1, 4]:
cornerImage = drawSegmentationImage(np.concatenate([threshold, pred_oc[batchIndex]], axis=2), blackIndex=0)
cornerImage[cornerImage == 0] = density[cornerImage == 0]
cv2.imwrite(options.test_dir + '/' + str(batchIndex) + '_opening_corner_pred.png', cornerImage)
cornerImage = drawSegmentationImage(np.concatenate([threshold, pred_ic[batchIndex]], axis=2), blackIndex=0)
cornerImage[cornerImage == 0] = density[cornerImage == 0]
cv2.imwrite(options.test_dir + '/' + str(batchIndex) + '_icon_corner_pred.png', cornerImage)
cornerImage = drawSegmentationImage(np.concatenate([threshold, gt_oc[batchIndex]], axis=2), blackIndex=0)
cornerImage[cornerImage == 0] = density[cornerImage == 0]
cv2.imwrite(options.test_dir + '/' + str(batchIndex) + '_opening_corner_gt.png', cornerImage)
cornerImage = drawSegmentationImage(np.concatenate([threshold, gt_ic[batchIndex]], axis=2), blackIndex=0)
cornerImage[cornerImage == 0] = density[cornerImage == 0]
cv2.imwrite(options.test_dir + '/' + str(batchIndex) + '_icon_corner_gt.png', cornerImage)
pass
if datasetFlag in [1, 2, 3, 4]:
icon_density = drawSegmentationImage(gt_dict['icon'][batchIndex], blackIndex=0)
icon_density[icon_density == 0] = density[icon_density == 0]
cv2.imwrite(options.test_dir + '/' + str(batchIndex) + '_icon_gt.png', icon_density)
icon_density = drawSegmentationImage(pred_dict['icon'][batchIndex], blackIndex=0)
icon_density[icon_density == 0] = density[icon_density == 0]
cv2.imwrite(options.test_dir + '/' + str(batchIndex) + '_icon_pred.png', icon_density)
pass
if datasetFlag in [1, 3, 4]:
room_density = drawSegmentationImage(gt_dict['room'][batchIndex], blackIndex=0)
room_density[room_density == 0] = density[room_density == 0]
cv2.imwrite(options.test_dir + '/' + str(batchIndex) + '_room_gt.png', room_density)
room_density = drawSegmentationImage(pred_dict['room'][batchIndex], blackIndex=0)
room_density[room_density == 0] = density[room_density == 0]
cv2.imwrite(options.test_dir + '/' + str(batchIndex) + '_room_pred.png', room_density)
pass
if batchIndex == 0 and False:
for c in xrange(22):
cv2.imwrite(options.test_dir + '/mask_' + str(c) + '.png', cv2.dilate(drawMaskImage(corner_segmentation[batchIndex] == c), np.ones((3, 3)), 3))
continue
continue
if batchIndex < options.visualizeReconstruction or True:
if options.debug >= 0 and batchIndex != options.debug:
continue
names.append((batchIndex, gt_dict['image_path'][batchIndex]))
print(batchIndex, 'start reconstruction', gt_dict['image_path'][batchIndex])
if True:
if options.debug == -1:
blockPrint()
pass
# gtHeatmaps = gt_dict['corner'][batchIndex]
#result_gt = reconstructFloorplan(gtHeatmaps[:, :, :NUM_WALL_CORNERS], gtHeatmaps[:, :, NUM_WALL_CORNERS:NUM_WALL_CORNERS + 4], gtHeatmaps[:, :, NUM_WALL_CORNERS + 4:NUM_WALL_CORNERS + 8], segmentation2Heatmaps(gt_dict['icon'][batchIndex], NUM_ICONS), segmentation2Heatmaps(gt_dict['room'][batchIndex], NUM_ROOMS), density[:, :, 0], gt=True)
orientationCorners = getOrientationCorners(gt_dict['corner_values'][batchIndex][:gt_dict['num_corners'][batchIndex]])
result_gt = reconstructFloorplan(orientationCorners[:NUM_WALL_CORNERS], orientationCorners[NUM_WALL_CORNERS:NUM_WALL_CORNERS + 4], orientationCorners[NUM_WALL_CORNERS + 4:NUM_WALL_CORNERS + 8], segmentation2Heatmaps(gt_dict['icon'][batchIndex], NUM_ICONS), segmentation2Heatmaps(gt_dict['room'][batchIndex], NUM_ROOMS), density[:, :, 0], gt=True)
#if batchIndex == 1:
#exit(1)
#pred_debug_dir = options.test_dir + '/' + str(batchIndex) + '_debug'
pred_debug_dir = options.test_dir
try:
os.mkdir(pred_debug_dir)
pass
except OSError as e:
pass
result_pred = reconstructFloorplan(pred_wc[batchIndex], pred_oc[batchIndex], pred_ic[batchIndex], pred_icon, pred_room, density[:, :, 0], gt_dict=result_gt, gt=False, debug_prefix=pred_debug_dir)
if True:
try:
newWidth = newHeight = 1000
resizeResult(result_gt, newWidth, newHeight, WIDTH, HEIGHT)
resultImageGT = drawResultImageFinal(newWidth, newHeight, result_gt)
#cv2.imwrite(options.test_dir + '/' + str(batchIndex) + '_result_gt.png', resultImageGT)
cv2.imwrite(options.test_dir + '/' + gt_dict['image_path'][batchIndex] + '_gt.png', resultImageGT)
resizeResult(result_pred, newWidth, newHeight, WIDTH, HEIGHT)
resultImagePred = drawResultImageFinal(newWidth, newHeight, result_pred)
cv2.imwrite(options.test_dir + '/' + gt_dict['image_path'][batchIndex] + '_pred.png', resultImagePred)
except:
continue
continue
if 'wall' not in result_pred or 'wall' not in result_gt:
print('invalid result')
continue
statistics = findMatches(result_pred, result_gt, distanceThreshold=10)
if options.drawFinal:
newWidth = newHeight = 1000
resizeResult(result_gt, newWidth, newHeight, WIDTH, HEIGHT)
resultImageGT = drawResultImageFinal(newWidth, newHeight, result_gt)
cv2.imwrite(options.test_dir + '/' + str(batchIndex) + '_result_gt.png', resultImageGT)
resizeResult(result_pred, newWidth, newHeight, WIDTH, HEIGHT)
resultImagePred = drawResultImageFinal(newWidth, newHeight, result_pred)
cv2.imwrite(options.test_dir + '/' + str(batchIndex) + '_result_pred.png', resultImagePred)
writeRepresentation('popup/data/floorplan_' + str(batchIndex) + '_gt.txt', newWidth, newHeight, result_gt)
writeRepresentation('popup/data/floorplan_' + str(batchIndex) + '_pred.txt', newWidth, newHeight, result_pred)
cv2.imwrite('popup/data/floorplan_' + str(batchIndex) + '_gt.png', resultImageGT)
cv2.imwrite('popup/data/floorplan_' + str(batchIndex) + '_pred.png', resultImagePred)
exit(1)
else:
resultImage, iconImage = drawResultImage(WIDTH, HEIGHT, result_gt)
cv2.imwrite(options.test_dir + '/' + str(batchIndex) + '_reconstruction_wall_gt.png', resultImage)
iconImage[iconImage == 0] = density[iconImage == 0]
cv2.imwrite(options.test_dir + '/' + str(batchIndex) + '_reconstruction_icon_gt.png', iconImage)
resultImage, iconImage = drawResultImage(WIDTH, HEIGHT, result_pred)
cv2.imwrite(options.test_dir + '/' + str(batchIndex) + '_reconstruction_wall_pred.png', resultImage)
iconImage[iconImage == 0] = density[iconImage == 0]
cv2.imwrite(options.test_dir + '/' + str(batchIndex) + '_reconstruction_icon_pred.png', iconImage)
pass
if options.debug == -1:
enablePrint()
pass
if len(result_pred) == 0:
continue
# print(result_pred)
# print(result_pred['door'])
# print('gt')
# print(result_gt)
# print(result_gt['door'])
# exit(1)
print('find predictions among ground-truths')
#print(result_pred['wall'][2])
#statistics = findMatches(result_pred, result_gt, distanceThreshold=10)
#statistics = findMatches(result_gt, result_pred, distanceThreshold=10)
#print('find ground-truths among predictions')
#statistics = findMatches(result_gt, result_pred, distanceThreshold=10)
#print(statistics)
print('statistics', [(k, float(v[0]) / max(v[1], 1), float(v[0]) / max(v[2], 1)) for k, v in statistics.iteritems()])
#print('topology statistics', [(k, float(v[0]) / max(v[1], 1), float(v[0]) / max(v[2], 1)) for k, v in topologyStatistics.iteritems()])
print('finish reconstruction', gt_dict['image_path'][batchIndex])
for k, v in statistics.iteritems():
if k in statisticsSum:
for c in xrange(3):
statisticsSum[k][c] += v[c]
continue
else:
print(k, 'not in', statisticsSum)
continue
if options.debug >= 0:
exit(1)
pass
pass
# except Exception as e:
# #traceback.print_tb(e)
# print('exception-----------: ', e)
# #raise e
continue
print(names)
print('final statistics', [(k, float(v[0]) / max(v[1], 1), float(v[0]) / max(v[2], 1)) for k, v in statisticsSum.iteritems()])
np.save(options.test_dir + '/numbers.npy', statisticsSum)
#print(statisticsSum)
return
def testBatch_unet(options, model, dataset):
model.eval()
dataloader = DataLoader(dataset,
batch_size=1,
shuffle=False,
num_workers=16)
data_iterator = tqdm(dataloader, total=len(dataset))
for sampleIndex, sample in enumerate(data_iterator):
images, img_path = sample[0].cuda(), sample[1][0]
img_name = os.path.splitext(img_path)[0].split('/')[-1]
corner_pred, icon_pred, room_pred = model(images)
room_unet = np.load('../../UNets/sb_preds_1000/%s.npy' %
img_name).transpose(1, 2, 0)
#corner_hg = np.load('../../pytorch-pose/val_res_0725_mse/%s.npy' % img_name)[0].transpose(1,2,0)
visualizeBatch(
options,
images.detach().cpu().numpy(),
[('pred', {
'corner': corner_pred.max(-1)[1].detach().cpu().numpy(),
'icon': icon_pred.max(-1)[1].detach().cpu().numpy(),
'room': room_pred.max(-1)[1].detach().cpu().numpy()
})],
batch_img_name=img_name)
for batchIndex in range(len(images)):
corner_heatmaps = torch.sigmoid(
corner_pred[batchIndex]).detach().cpu().numpy()
#corner_heatmaps = corner_hg
#print(corner_heatmaps.shape); exit(1)
icon_heatmaps = torch.sigmoid(
icon_pred[batchIndex]).detach().cpu().numpy()
#room_heatmaps = torch.sigmoid(room_pred[batchIndex]).detach().cpu().numpy()
room_heatmaps = torch.nn.functional.softmax(
torch.from_numpy(room_unet), dim=-1).detach().cpu().numpy()
reconstructFloorplan(
corner_heatmaps[:, :, :NUM_WALL_CORNERS],
corner_heatmaps[:, :, NUM_WALL_CORNERS:NUM_WALL_CORNERS + 8],
corner_heatmaps[:, :, -4:],
icon_heatmaps,
room_heatmaps,
output_prefix=options.test_dir + '/' + img_name + '_',
densityImage=None,
gt_dict=None,
gt=False,
gap=-1,
distanceThreshold=-1,
lengthThreshold=-1,
debug_prefix='test',
heatmapValueThresholdWall=None,
heatmapValueThresholdDoor=None,
heatmapValueThresholdIcon=None,
enableAugmentation=True)
continue
continue
return
def evaluate_batch(options, gt_dict, pred_dict):
if options.separateIconLoss:
# pred_icon_separate = softmax(np.load(options.test_dir.replace('wsf', 'wsf_loss3') + '/dummy/pred_dict.npy')[()]['icon'])
raise NotImplementedError
stat = {key: [0, 0, 0] for key in ('corner', 'edge', 'room', 'room++')}
pred_wc = pred_dict['corner'][:, :, :, :NUM_WALL_CORNERS]
pred_oc = pred_dict['corner'][:, :, :, NUM_WALL_CORNERS:NUM_WALL_CORNERS + 4]
pred_ic = pred_dict['corner'][:, :, :, NUM_WALL_CORNERS + 4:NUM_WALL_CORNERS + 8]
if options.branches != '5':
pred_wc = sigmoid(pred_wc)
pred_oc = sigmoid(pred_oc)
pred_ic = sigmoid(pred_ic)
else:
raise NotImplementedError
for batch_idx in xrange(gt_dict['corner'].shape[0]):
if options.evaluateImage and gt_dict['image_flags'][batch_idx] == 0:
raise NotImplementedError
# Build density map for visualization
density = np.minimum(gt_dict['density'][batch_idx] * 255, 255).astype(np.uint8)
density = np.stack([density, density, density], axis=-1)
pred_icon = softmax(pred_dict['icon'][batch_idx])
pred_room = softmax(pred_dict['room'][batch_idx])
if options.separateIconLoss:
raise NotImplementedError
print(batch_idx, 'start reconstruction', gt_dict['image_path'][batch_idx])
debug_dir = os.path.join(options.test_dir, 'gt_debug_')
output_dir = os.path.join(options.test_dir, 'gt_output_')
num_corners = gt_dict['num_corners'][batch_idx]
orientationCorners = getOrientationCorners(gt_dict['corner_values'][batch_idx][:num_corners])
result_gt = reconstructFloorplan(
orientationCorners[:NUM_WALL_CORNERS],
orientationCorners[NUM_WALL_CORNERS:NUM_WALL_CORNERS + 4],
orientationCorners[NUM_WALL_CORNERS + 4:NUM_WALL_CORNERS + 8],
segmentation2Heatmaps(gt_dict['icon'][batch_idx], NUM_ICONS),
segmentation2Heatmaps(gt_dict['room'][batch_idx], NUM_ROOMS),
output_prefix=output_dir,
densityImage=density[:, :, 0],
gt=True,
debug_prefix=debug_dir)
debug_dir = os.path.join(options.test_dir, 'pred_debug_')
output_dir = os.path.join(options.test_dir, 'pred_output_')
result_pred = reconstructFloorplan(
pred_wc[batch_idx],
pred_oc[batch_idx],
pred_ic[batch_idx],
pred_icon,
pred_room,
output_prefix=output_dir,
densityImage=density[:, :, 0],
gt_dict=result_gt,
gt=False, debug_prefix=debug_dir)
# TODO: Visualize result and fix visualize in reconstructFloorplan
# Evaluate corner metric
result = eval_corner_wrapper(
result_pred, result_gt,
dist_threshold=CORNER_DIST_THRESHOLD
)
stat['corner'][0] += result[0]
stat['corner'][1] += result[1]
stat['corner'][2] += result[2]
result = eval_edges_wrapper(
result_pred, result_gt,
dist_threshold=CORNER_DIST_THRESHOLD
)
stat['edge'][0] += result[0]
stat['edge'][1] += result[1]
stat['edge'][2] += result[2]
result = eval_room_wrapper(result_pred, result_gt, ROOM_IOU_THRESHOLD)
stat['room'][0] += result[0]
stat['room'][1] += result[1]
stat['room'][2] += result[2]
result = eval_room_pp_wrapper(result_pred, result_gt, ROOM_IOU_THRESHOLD)
stat['room++'][0] += result[0]
stat['room++'][1] += result[1]
stat['room++'][2] += result[2]
# visualize pred
background = density[:, :, :3].copy()
result_path = os.path.join(options.test_dir, '%s_pred_result.png' % batch_idx)
result_img = draw_result(result_pred, background)
cv2.imwrite(result_path, result_img)
logging.info('Write visualization result to %s' % result_path)
# visualize GT
background = density[:, :, :3].copy()
result_path = os.path.join(options.test_dir, '%s_gt_result.png' % batch_idx)
result_img = draw_result(result_gt, background)
cv2.imwrite(result_path, result_img)
logging.info('Write visualization result to %s' % result_path)
logging.info(stat)
