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)