from data import Obj, ProjectionGenerator, Projection import numpy as np import math import scipy.misc as m from numba import jit import torch class TopDownView(): """ Take a room, pre-render top-down views Of floor, walls and individual objects That can be used to generate the multi-channel views used in our pipeline """ def __init__(self, height_cap=4.05, length_cap=6.05, size=512): #Padding by 0.05m to avoid problems with boundary cases """ Parameters ---------- height_cap (int): the maximum height (in meters) of rooms allowed, which will be rendered with a value of 1 in the depth channel. To separate the floor from empty spaces, floors will have a height of 0.5m. See zpad below length_cap (int): the maximum length/width of rooms allowed. size (int): size of the rendered top-down image Returns ------- visualization (Image): a visualization of the rendered room, which is simply the superimposition of all the rendered parts (floor, wall, nodes) (Triple[torch.Tensor, torch.Tensor, list[torch.Tensor]): rendered invidiual parts of the room, as 2D torch tensors this is the part used by the pipeline """ self.size = size self.pgen = ProjectionGenerator(room_size_cap=(length_cap, height_cap, length_cap), \ zpad=0.5, img_size=size) def render(self, room): projection = self.pgen.get_projection(room) visualization = np.zeros((self.size,self.size)) nodes = [] for node in room.nodes: modelId = node.modelId #Camelcase due to original json t = np.asarray(node.transform).reshape(4,4) o = Obj(modelId) t = projection.to_2d(t) o.transform(t) t = projection.to_2d() bbox_min = np.dot(np.asarray([node.xmin, node.zmin, node.ymin, 1]), t) bbox_max = np.dot(np.asarray([node.xmax, node.zmax, node.ymax, 1]), t) xmin = math.floor(bbox_min[0]) ymin = math.floor(bbox_min[2]) xsize = math.ceil(bbox_max[0]) - xmin + 1 ysize = math.ceil(bbox_max[2]) - ymin + 1 description = {} description["modelId"] = modelId description["transform"] = node.transform description["bbox_min"] = bbox_min description["bbox_max"] = bbox_max #Since it is possible that the bounding box information of a room #Was calculated without some doors/windows #We need to handle these cases if ymin < 0: ymin = 0 if xmin < 0: xmin = 0 rendered = self.render_object(o, xmin, ymin, xsize, ysize, self.size) description["height_map"] = torch.from_numpy(rendered).float() tmp = np.zeros((self.size, self.size)) tmp[xmin:xmin+rendered.shape[0],ymin:ymin+rendered.shape[1]] = rendered visualization += tmp nodes.append(description) #Render the floor o = Obj(room.modelId+"f", room.house_id, is_room=True) t = projection.to_2d() o.transform(t) floor = self.render_object(o, 0, 0, self.size, self.size, self.size) visualization += floor floor = torch.from_numpy(floor).float() #Render the walls o = Obj(room.modelId+"w", room.house_id, is_room=True) t = projection.to_2d() o.transform(t) wall = self.render_object(o, 0, 0, self.size, self.size, self.size) visualization += wall wall = torch.from_numpy(wall).float() return (visualization, (floor, wall, nodes)) @staticmethod @jit(nopython=True) def render_object_helper(triangles, xmin, ymin, xsize, ysize, img_size): result = np.zeros((img_size, img_size), dtype=np.float32) N, _, _ = triangles.shape for triangle in range(N): x0,z0,y0 = triangles[triangle][0] x1,z1,y1 = triangles[triangle][1] x2,z2,y2 = triangles[triangle][2] a = -y1*x2 + y0*(-x1+x2) + x0*(y1-y2) + x1*y2 if a != 0: for i in range(max(0,math.floor(min(x0,x1,x2))), \ min(img_size,math.ceil(max(x0,x1,x2)))): for j in range(max(0,math.floor(min(y0,y1,y2))), \ min(img_size,math.ceil(max(y0,y1,y2)))): x = i+0.5 y = j+0.5 s = (y0*x2 - x0*y2 + (y2-y0)*x + (x0-x2)*y)/a t = (x0*y1 - y0*x1 + (y0-y1)*x + (x1-x0)*y)/a if s < 0 and t < 0: s = -s t = -t if 0 < s < 1 and 0 < t < 1 and s + t <= 1: height = z0 *(1-s-t) + z1*s + z2*t result[i][j] = max(result[i][j], height) return result[xmin:xmin+xsize, ymin:ymin+ysize] @staticmethod def render_object(o, xmin, ymin, xsize, ysize, img_size): """ Render a cropped top-down view of object Parameters ---------- o (list[triple]): object to be rendered, represented as a triangle mesh xmin, ymin (int): min coordinates of the bounding box containing the object, with respect to the full image xsize, ysze (int); size of the bounding box containing the object img_size (int): size of the full image """ triangles = np.asarray(list(o.get_triangles()), dtype=np.float32) return TopDownView.render_object_helper(triangles, xmin, ymin, xsize, ysize, img_size) @staticmethod @jit(nopython=True) def render_object_full_size_helper(triangles, size): result = np.zeros((size, size), dtype=np.float32) N, _, _ = triangles.shape for triangle in range(N): x0,z0,y0 = triangles[triangle][0] x1,z1,y1 = triangles[triangle][1] x2,z2,y2 = triangles[triangle][2] a = -y1*x2 + y0*(-x1+x2) + x0*(y1-y2) + x1*y2 if a != 0: for i in range(max(0,math.floor(min(x0,x1,x2))), \ min(size,math.ceil(max(x0,x1,x2)))): for j in range(max(0,math.floor(min(y0,y1,y2))), \ min(size,math.ceil(max(y0,y1,y2)))): x = i+0.5 y = j+0.5 s = (y0*x2 - x0*y2 + (y2-y0)*x + (x0-x2)*y)/a t = (x0*y1 - y0*x1 + (y0-y1)*x + (x1-x0)*y)/a if s < 0 and t < 0: s = -s t = -t if 0 < s < 1 and 0 < t < 1 and s + t <= 1: height = z0 *(1-s-t) + z1*s + z2*t result[i][j] = max(result[i][j], height) return result @staticmethod def render_object_full_size(o, size): """ Render a full-sized top-down view of the object, see render_object """ triangles = np.asarray(list(o.get_triangles()), dtype=np.float32) return TopDownView.render_object_full_size_helper(triangles, size)