Spaces:
Runtime error
Runtime error
File size: 4,507 Bytes
fadb92b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | import os
import time
import cv2
import numpy as np
import torch
from s3d_floorplan_eval.DataRW.DataRW import DataRW
from s3d_floorplan_eval.S3DLoader.S3DLoader import S3DLoader
class S3DRW(DataRW):
def __init__(self, options, mode):
"""
Class for accessing FloorNet dataset related data
:param options:
"""
# initialize the base class variables
super(DataRW, self).__init__()
self.options = options
self.dataset_path = options.dataset_path
self.scene_id = options.scene_id
self.mcts_path = options.mcts_path
self.creation_time = int(time.time())
self.device = torch.device("cpu")
# mode = "train"
# mode = "online_eval"
# For validation only
# self.loader = S3DLoader(options, 'online_eval').dataset
self.loader = S3DLoader(options, mode).dataset
# gt_sample = iter(floornet_loader.dataset[int(self.scene_id)])
# self.gt_sample = floornet_loader.load_sample(list(iter(floornet_loader.dataset))[int(self.scene_id)])
if mode == "online_eval":
scene_ind = int(self.scene_id[6:]) - 3000
elif mode == "test":
scene_ind = int(self.scene_id[6:]) - 3250
elif mode == "train":
scene_ind = int(self.scene_id[6:])
else:
assert False
self.gt_sample = self.loader[scene_ind]
self.gt_sample["density_map"] = torch.tensor(self.gt_sample["density_map"][None], device=self.device)
self.gt_sample["room_map"] = torch.tensor(self.gt_sample["room_map"][None, :, :, None], device=self.device)
self.gt_sample["wall_map"] = torch.tensor(self.gt_sample["wall_map"][None, :, :, None], device=self.device)
self.density_map = self.gt_sample["density_map"][:, :, :, None]
self.h, self.w = self.density_map.shape[1], self.density_map.shape[2]
self.generate_input_map_from_props = self.generate_input_dict_from_room_props
def get_gt_solution(self):
"""
Read top-view density map of the scene
:return:
"""
img_path = os.path.join(self.dataset_path, str(self.scene_id) + "_density.png")
density_map = cv2.imread(img_path, cv2.IMREAD_ANYDEPTH | cv2.IMREAD_ANYCOLOR)[:, :, 0][None, :, :, None]
density_map = torch.from_numpy(density_map).to(self.device)
dm_min = torch.min(density_map)
dm_max = torch.max(density_map)
density_map = (density_map - dm_min) / (dm_max - dm_min)
return density_map.type(torch.cuda.FloatTensor)
def polygonize_mask(self, pm, return_mask=True):
pm_np = pm.cpu().numpy()
room_mask = 255 * (pm_np == 1)
room_mask = room_mask.astype(np.uint8)
room_mask_inv = 255 - room_mask
ret, thresh = cv2.threshold(room_mask_inv, 250, 255, cv2.THRESH_BINARY_INV)
contours, hierarchy = cv2.findContours(thresh, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_NONE)
cnt = contours[0]
max_area = cv2.contourArea(cnt)
for cont in contours:
if cv2.contourArea(cont) > max_area:
cnt = cont
max_area = cv2.contourArea(cont)
# define main island contour approx. and hull
epsilon = 0.01 * cv2.arcLength(cnt, True)
approx = cv2.approxPolyDP(cnt, epsilon, True)
approx = approx.astype(np.int32).reshape((1, -1, 2))
if return_mask:
room_filled_map = np.zeros((self.h, self.w))
cv2.fillPoly(room_filled_map, approx, color=1.0)
room_filled_map = torch.tensor(room_filled_map[:, :], dtype=torch.float32, device=self.device)
return room_filled_map
else:
approx_tensor = torch.tensor(approx, device=self.device)
return approx_tensor
def generate_input_dict_from_room_props(self, room_prop_list, score_function, use_thresh=False):
"""
:param room_prop_list:
:type room_prop_list: list of FloorPlanRoomProp
:param score_function:
:return:
"""
if score_function == "room_maskrcnn_iou":
inputs = self.generate_input_dict_for_room_maskrcnn_iou(room_prop_list)
elif score_function == "room_iou":
inputs = self.generate_input_dict_for_room_iou(room_prop_list, use_thresh=use_thresh)
else:
assert "generate_input_dict_from_room_props for %s not implemented" % score_function
return inputs
|