def snap_on_wall(obj, wall, obj_poly, wall_poly): if wall_poly.intersects(obj_poly): intersection = wall_poly.intersection(obj_poly) if type(intersection) == shape_string: return elif type(intersection) == shape_poly: xy = np.array(intersection.exterior.coords.xy) ptp = xy.ptp(axis=1) for move in [(ptp[0], 0), (-ptp[0], 0), (0, ptp[1]), (0, -ptp[1])]: temp_obj = shape_poly(obj.get_coords() + move) if type(temp_obj.intersection(wall_poly)) == shape_string: obj.center += move return else: distance = wall_poly.distance(obj_poly) overlaps = [] dirs = [(distance, 0), (-distance, 0), (0, distance), (0, -distance)] for move in dirs: temp_obj = shape_poly(obj.get_coords() + move) if not temp_obj.intersects(wall_poly): overlaps.append(np.finfo(np.float).max) else: overlaps.append(obj_poly.intersection(wall_poly).area) snap_dir = np.argmin(np.array(overlaps)) obj.center += dirs[snap_dir] return print('No wall is found to fit the criteria...')
def overlaps(bbox_i, bbox_j): coord_i_xy = shape_poly(bbox_i[1].get_coords()) i_z = bbox_i[1].z coord_j_xy = shape_poly(bbox_j[1].get_coords()) j_z = bbox_j[1].z if coord_i_xy.intersects(coord_j_xy): z_overlap = get_z_overlaps(i_z, j_z) return coord_i_xy.intersection(coord_j_xy).area * z_overlap return 0
def process_hole_elements(model_id, cat): hole_cats = ['Door', 'Window', 'Hole'] if cat not in hole_cats: raise ValueError('{} is not among the hole categories') with open(model_id, 'r') as fp: data = json.load(fp) components = [f for f in data['mesh'] if f['type'] == cat and get_bbox(f)[1][1] > 0.05] polys = [mesh_to_xy_convex_hull(np.array(w['xyz']).reshape([-1,3])) for w in components] zs = [] for w in components: mmin, mmax = get_bbox(w) zs.append((mmin[1], mmax[1])) filtered_poly = [] filtered_indexes = [] for i,p in enumerate(polys): ith_poly = shape_poly(p) if intersects(ith_poly, filtered_poly): continue filtered_poly.append(ith_poly) filtered_indexes.append(i) zs = [p for i,p in enumerate(zs) if i in filtered_indexes] polys = [p for i,p in enumerate(polys) if i in filtered_indexes] rects = [minimum_bounding_rectangle(p) for p in polys] bboxes = [polygon_to_bbox(p[:,0], p[:,1], z, scale_factor=1.) for p,z in zip(rects, zs) if z[0] != z[1]] return polys,bboxes
def snap_out_of_wall(obj, wall, obj_poly, wall_poly): intersection = wall_poly.intersection(obj_poly) if type(intersection) == shape_string: overlap_margin = 0.001 for move in [(overlap_margin, 0), (-overlap_margin, 0), (0, overlap_margin), (0, -overlap_margin)]: temp_obj = shape_poly(obj.get_coords() + move) if not temp_obj.intersects(wall_poly): obj.center += move return elif type(intersection) == shape_poly: xy = np.array(intersection.exterior.coords.xy) ptp = xy.ptp(axis=1) + 0.001 for move in [(ptp[0], 0), (-ptp[0], 0), (0, ptp[1]), (0, -ptp[1])]: temp_obj = shape_poly(obj.get_coords() + move) if not temp_obj.intersects(wall_poly): obj.center += move return
def gen_structure(wall_bbox, hole_bbox, floor_bbox, model_dir, skirt_pts, save_obj=False, overwrite=False): # finding correspondence between wall and windows # print('Examining walls and holes...') walls_poly = [] for bbox in wall_bbox: polygon = shape_poly(bbox.get_coords()) walls_poly.append(polygon) holes_poly = [] for bbox, _ in hole_bbox: polygon = shape_poly(bbox.get_coords()) holes_poly.append(polygon) walls_to_holes = defaultdict(lambda: []) for hi, hole in enumerate(holes_poly): has_intersect = False distances = [] overlaps = [hole.intersects(wall) for wall in walls_poly] overlaps_area = [ hole.intersection(wall).area if overlap else 0 for overlap, wall in zip(overlaps, walls_poly) ] if sum(overlaps) > 1: max_idx = np.argmax(overlaps_area) for wi, o in enumerate(overlaps): if wi == max_idx: walls_to_holes[wi].append(hole_bbox[hi]) elif o: snap_out_of_wall(hole_bbox[hi][0], wall_bbox[wi], hole, walls_poly[wi]) elif sum(overlaps) == 1: wi = overlaps.index(True) walls_to_holes[wi].append(hole_bbox[hi]) else: continue if save_obj: structure_dir = os.path.join(model_dir, 'shape', 'visual') os.makedirs(structure_dir, exist_ok=True) vhacd_dir = os.path.join(model_dir, 'shape', 'collision') os.makedirs(vhacd_dir, exist_ok=True) verts = [] faces = [] structure_obj = [] num_verts = 0 for idx, w in enumerate(wall_bbox): wall = Wall_bbox(w, ceiling_height=w.z[1]) if idx in walls_to_holes: for bbox, oname in walls_to_holes[idx]: new_bb = wall.add_hole(bbox) if oname != '': structure_obj.append((new_bb, oname)) v, f = wall.build() if save_obj: cm_path = os.path.join(vhacd_dir, 'wall_{}_cm.obj'.format(idx)) if overwrite or not os.path.isfile(cm_path): wall.build_collision_mesh(cm_path) verts.append(v) f[:, :3] += num_verts faces.append(f) num_verts += v.shape[0] v, f = wall.add_skirt(skirt_pts) verts.append(v) f[:, :3] += num_verts faces.append(f) num_verts += v.shape[0] v = np.concatenate(verts, axis=0) f = np.concatenate(faces, axis=0) if save_obj: vm_path = os.path.join(structure_dir, 'wall_vm.obj'.format(idx)) if overwrite or not os.path.isfile(vm_path): with open(vm_path, 'w') as fp: for vs in v.reshape([-1, 3]): fp.write('v {} {} {}\n'.format(*vs)) for f1, f2, f3, fn in f: fp.write('f {f1} {f2} {f3}\n'.format(f1=f1, f2=f2, f3=f3)) ### Process floors floor_counter = 0 if len(floor_bbox) > 0 and type(floor_bbox[0]) is not list: floor_ids = set([i for _, i in floor_bbox]) if 0 in floor_ids: floor_id_remap = {} for i in floor_ids: floor_id_remap[i] = i + 1 else: floor_id_remap = {} for i in floor_ids: floor_id_remap[i] = i floor_id_remap[-1] = 0 floor_polygons = defaultdict(lambda: []) for idx, (bbox, id) in enumerate(floor_bbox): if floor_id_remap[id] == 0: # print(floor_counter) if save_obj: floor_path = os.path.join( structure_dir, 'floor_{}_vm.obj'.format(floor_counter)) v, f = gen_cube_obj( bbox, floor_path, is_color=False, should_save=(overwrite or not os.path.isfile(floor_path))) v = np.asarray(v) f = np.asarray(f) verts.append(v) f[:, :3] += num_verts faces.append(f) num_verts += v.shape[0] floor_counter += 1 else: floor_polygons[floor_id_remap[id]].append( shape_poly(bbox.get_coords())) else: floor_polygons = {} for i, f in enumerate(floor_bbox): floor_polygons[i] = [shape_poly(f)] thickness = 0.2 for k, polys in floor_polygons.items(): if len(polys) > 1: x, y = cascaded_union(polys).exterior.coords.xy else: x, y = polys[0].exterior.coords.xy v = np.asarray([x, y]).T v = v[:-1] v_flat = np.reshape(v, -1) f_t = np.array(earcut.earcut(v_flat)).reshape((-1, 3)) + 1 v_z = np.asarray([ 0., ] * int(v.shape[0]))[:, np.newaxis] top_v = np.concatenate([v, v_z], axis=-1) back_v = top_v + 0.0 back_v[:, -1] = -thickness f_b = f_t + len(top_v) temp = f_b[:, 1] + 0.0 f_b[:, 1] = f_b[:, 0] f_b[:, 0] = temp cross_faces = [] curr_point = 1 num_front_plane = len(top_v) for idx in range(num_front_plane): if idx == num_front_plane - 1: next_point = 1 else: next_point = curr_point + 1 cross_faces.append( (curr_point, next_point, curr_point + num_front_plane)) cross_faces.append((curr_point + num_front_plane, next_point, next_point + num_front_plane)) curr_point = next_point v = np.concatenate([top_v, back_v]) f = np.concatenate([f_t, f_b, cross_faces], axis=0) f_n = np.asarray([ 0, ] * int(f.shape[0]))[:, np.newaxis] f = np.concatenate([f, f_n], axis=-1) if save_obj: floor_path = os.path.join(structure_dir, 'floor_{}_vm.obj'.format(floor_counter)) if overwrite or not os.path.isfile(floor_path): with open(floor_path, 'w') as fp: for vs in v.reshape([-1, 3]): fp.write('v {} {} {}\n'.format(*vs)) for f1, f2, f3, fn in f.reshape([-1, 4]): fp.write('f {f1} {f2} {f3}\n'.format(f1=f1, f2=f2, f3=f3)) verts.append(v) f[:, :3] += num_verts faces.append(f) num_verts += v.shape[0] floor_counter += 1 verts = np.concatenate(verts, axis=0) faces = np.concatenate(faces, axis=0) if save_obj: with open(os.path.join(model_dir, 'misc', 'structure.obj'), 'w') as fp: for vs in verts.reshape([-1, 3]): fp.write('v {} {} {}\n'.format(*vs)) for f1, f2, f3, fn in faces.reshape([-1, 4]): fp.write('f {f1} {f2} {f3}\n'.format(f1=f1, f2=f2, f3=f3)) xa, ya, _ = verts.min(axis=0) x, y, _ = verts.max(axis=0) floor_coords = np.asarray([(x, y), (xa, y), (xa, ya), (x, ya)]) total_floor = BBox(floor_coords, 0, 0 - thickness) if save_obj: floor_cm_path = os.path.join(vhacd_dir, 'floor_cm.obj') if overwrite or not os.path.isfile(floor_cm_path): gen_cube_obj(total_floor, floor_cm_path, is_color=False) total_ceiling = BBox(floor_coords, ceiling_height + thickness, ceiling_height) ceiling_cm_path = os.path.join(vhacd_dir, 'ceiling_cm.obj') if overwrite or not os.path.isfile(ceiling_cm_path): gen_cube_obj(total_ceiling, ceiling_cm_path, is_color=False) ceiling_vm_path = os.path.join(structure_dir, 'ceiling_vm.obj') if overwrite or not os.path.isfile(ceiling_vm_path): gen_cube_obj(total_ceiling, ceiling_vm_path, is_color=False) bbox_dir = os.path.join(model_dir, 'misc', 'bbox') os.makedirs(bbox_dir, exist_ok=True) all_objects = [] for idx, (bbox, oclass) in enumerate(structure_obj): bbox_path = os.path.join(bbox_dir, '{}_{}_bbox.obj'.format(oclass, idx)) x, y, z = bbox.get_scale() if x > y: continue gen_cube_obj(bbox, bbox_path, is_color=True) bbox_dict = bbox.as_dict() bbox_dict['category'] = oclass bbox_dict['instance'] = None bbox_dict['is_fixed'] = True bbox_dict['theta'] = bbox.get_rotation_angle() all_objects.append(bbox_dict) return all_objects
def export_scene(fixed_obj_bbox, free_obj_bbox, wall_bbox, model_dir): walls_poly = [] for bbox in wall_bbox: polygon = shape_poly(bbox.get_coords()) walls_poly.append(polygon) free_poly = [] for _, bbox in free_obj_bbox: polygon = shape_poly(bbox.get_coords()) free_poly.append(polygon) for i, obj in enumerate(free_poly): for wi, wall in enumerate(walls_poly): if obj.intersects(wall): snap_out_of_wall(free_obj_bbox[i][1], wall_bbox[wi], obj, wall) free_poly_snapped = [ shape_poly(bbox.get_coords()) for _, bbox in free_obj_bbox ] to_delete = [] for i, obj in enumerate(free_poly_snapped): for wi, wall in enumerate(walls_poly): if obj.intersects(wall): to_delete.append(i) free_obj_bbox = [ f for i, f in enumerate(free_obj_bbox) if i not in to_delete ] fixed_poly = [] for _, bbox in fixed_obj_bbox: polygon = shape_poly(bbox.get_coords()) fixed_poly.append(polygon) vs = [] fs = [] all_objects = [] bbox_dir = os.path.join(model_dir, 'misc', 'bbox') os.makedirs(bbox_dir, exist_ok=True) for i, (oclass, bbox) in enumerate(free_obj_bbox): if oclass == 'shower_head': continue bbox_dict = bbox.as_dict() instance = None if '=' in oclass: oclass, instance = oclass.split('=') instance = int(instance) bbox_path = os.path.join(bbox_dir, 'free_{}_{}_bbox.obj'.format(oclass, i)) gen_cube_obj(bbox, bbox_path, is_color=True) bbox_dict['category'] = oclass bbox_dict['instance'] = instance bbox_dict['is_fixed'] = False bbox_dict['theta'] = bbox.get_rotation_angle() all_objects.append(bbox_dict) for i, (oclass, bbox) in enumerate(fixed_obj_bbox): if oclass == 'shower_head': continue bbox_dict = bbox.as_dict() instance = None if '=' in oclass: oclass, instance = oclass.split('=') instance = int(instance) bbox_path = os.path.join(bbox_dir, 'fixed_{}_{}_bbox.obj'.format(oclass, i)) gen_cube_obj(bbox, bbox_path, is_color=True) bbox_dict['category'] = oclass bbox_dict['instance'] = instance bbox_dict['is_fixed'] = True bbox_dict['theta'] = bbox.get_rotation_angle() all_objects.append(bbox_dict) return all_objects
def main(): args = parser.parse_args() model_dir = os.path.normpath(args.model_dir) print(model_dir) # model_id = "_".join(model_dir.split('/')[-2:]) model_id = os.path.basename(os.path.normpath(model_dir)) svg_file = os.path.join(model_dir, 'model.svg') img_path = os.path.join(model_dir, 'F1_scaled.png') svg = minidom.parse(svg_file) # parseString also exists fplan = cv2.imread(img_path) fplan = cv2.cvtColor(fplan, cv2.COLOR_BGR2RGB) # correct color channels height, width, nchannel = fplan.shape shape = height, width wall_bboxes, wall_image = get_wall(svg, shape) window_bboxes, door_bboxes = get_window_door(svg) furniture_bboxes = get_furniture(svg, wall_image) floor_polygons = get_floor(svg) overlaps = defaultdict(lambda: []) for i in range(len(furniture_bboxes)): for j in range(i + 1, len(furniture_bboxes)): if has_overlap(furniture_bboxes[i][1], furniture_bboxes[j][1]): overlaps[(i, furniture_bboxes[i][1])].append( (j, furniture_bboxes[j][1])) overlaps[(j, furniture_bboxes[j][1])].append( (i, furniture_bboxes[i][1])) to_delete = [] for _ in range(len(overlaps)): overlaps_list = list(overlaps.items()) overlaps_list.sort(key=lambda x: -get_volume(x[0][1])) overlaps_list.sort(key=lambda x: -len(x[1])) o = overlaps_list[0] # try shrinking shrink_success = False bbox = o[0][1] edge_x_og = bbox.edge_x[:] edge_y_og = bbox.edge_y[:] for scale_factor in range(20): scale = 1. - 0.01 * scale_factor # try scale edge_x bbox.edge_x = edge_x_og * scale overlap = False for i in o[1]: if has_overlap(bbox, i[1]): overlap = True break if not overlap: shrink_success = True break bbox.edge_x = edge_x_og # try scale edge_y bbox.edge_y = edge_y_og * scale overlap = False for i in o[1]: if has_overlap(bbox, i[1]): overlap = True break if not overlap: shrink_success = True break bbox.edge_y = edge_y_og # try scale both bbox.edge_y = edge_y_og * scale bbox.edge_x = edge_x_og * scale overlap = False for i in o[1]: if has_overlap(bbox, i[1]): overlap = True break if not overlap: shrink_success = True break if shrink_success: furniture_bboxes[o[0][0]] = (furniture_bboxes[o[0][0]][0], bbox) else: # add to delete to_delete.append(o[0][0]) # update graph for j in o[1]: overlaps[j].remove(o[0]) del overlaps[o[0]] for i in sorted(to_delete, reverse=True): del furniture_bboxes[i] ################################## # Splitting into separate floors # ################################## total_image = Image.new('1', (height, width), 0) d = ImageDraw.Draw(total_image) for group in [wall_bboxes, door_bboxes, window_bboxes]: for bbox in group: coords = bbox.get_coords() * 100. y = coords[:, 0] x = coords[:, 1] d.polygon(list(zip(y, x)), fill=1) for _, bbox in furniture_bboxes: coords = bbox.get_coords() * 100. y = coords[:, 0] x = coords[:, 1] d.polygon(list(zip(y, x)), fill=1) for _, coords in floor_polygons: y = coords[:, 0] * 100. x = coords[:, 1] * 100. d.polygon(list(zip(y, x)), fill=1) int_image = np.array(total_image).astype(np.uint8) binary = cv2.threshold(int_image, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)[1] ret, labels = cv2.connectedComponents(binary) masks = [] for label in range(1, ret): mask = np.zeros_like(int_image) mask[labels == label] = 1 masks.append(mask) # iterate through list of masks, each of which is a floor for floor_i, mask_i in enumerate(masks): wall_bboxes_i = [] window_bboxes_i = [] door_bboxes_i = [] furniture_bboxes_i = [] floor_polygons_i = [] for bbox in wall_bboxes: y, x = (bbox.center * 100.).astype(int) if y >= height: y = height - 1 if x >= width: x = width - 1 if mask_i[x, y] == 1: wall_bboxes_i.append(bbox) for bbox in door_bboxes: y, x = (bbox.center * 100.).astype(int) if y >= height: y = height - 1 if x >= width: x = width - 1 if mask_i[x, y] == 1: door_bboxes_i.append(bbox) for bbox in window_bboxes: y, x = (bbox.center * 100.).astype(int) if y >= height: y = height - 1 if x >= width: x = width - 1 if mask_i[x, y] == 1: window_bboxes_i.append(bbox) for c, bbox in furniture_bboxes: y, x = (bbox.center * 100.).astype(int) if y >= height: y = height - 1 if x >= width: x = width - 1 if mask_i[x, y] == 1: furniture_bboxes_i.append((c, bbox)) for poly in floor_polygons: y, x = (np.array(shape_poly(poly[1]).representative_point()) * 100.).astype(int) if y >= height: y = height - 1 if x >= width: x = width - 1 if mask_i[x, y] == 1: floor_polygons_i.append(poly) if len(wall_bboxes_i) < 4 or len(floor_polygons_i) < 1: # This suggests that the mask doesn't represent a floor continue model_dir = os.path.join(args.save_root, '{}_floor_{}'.format(model_id, floor_i)) os.makedirs(model_dir, exist_ok=True) save_dir = os.path.join(model_dir, 'misc') os.makedirs(save_dir, exist_ok=True) with open(os.path.join(save_dir, 'wall.json'), 'w') as fp: json.dump([bbox.as_dict() for bbox in wall_bboxes_i], fp) with open(os.path.join(save_dir, 'window.json'), 'w') as fp: json.dump([bbox.as_dict() for bbox in window_bboxes_i], fp) with open(os.path.join(save_dir, 'door.json'), 'w') as fp: json.dump([bbox.as_dict() for bbox in door_bboxes_i], fp) with open(os.path.join(save_dir, 'furniture.json'), 'w') as fp: json.dump([(cat, bbox.as_dict()) for cat, bbox in furniture_bboxes_i], fp) with open(os.path.join(save_dir, 'floor.json'), 'w') as fp: json.dump([(cat, poly.tolist()) for cat, poly in floor_polygons_i], fp) layout_dir = os.path.join(model_dir, 'layout') os.makedirs(layout_dir, exist_ok=True) coords = [bbox.get_coords() for bbox in wall_bboxes_i] stacked = np.vstack(coords) xmin, ymin = stacked.min(axis=0) xmax, ymax = stacked.max(axis=0) max_length = np.max( [np.abs(xmin), np.abs(ymin), np.abs(xmax), np.abs(ymax)]) max_length = np.ceil(max_length).astype(np.int) ins_image = Image.new('L', (2 * max_length * 100, 2 * max_length * 100), 0) d1 = ImageDraw.Draw(ins_image) sem_image = Image.new('L', (2 * max_length * 100, 2 * max_length * 100), 0) d2 = ImageDraw.Draw(sem_image) for i, (cat, poly) in enumerate(floor_polygons_i): room_id = rooms.index(cat) pts = [((x + max_length) * 100., (y + max_length) * 100.) for x, y in poly] d1.polygon(pts, fill=i + 1) d2.polygon(pts, fill=room_id + 1) ins_image.save(os.path.join(layout_dir, 'floor_insseg_0.png')) sem_image.save(os.path.join(layout_dir, 'floor_semseg_0.png')) padded_image = Image.new('L', (3000, 3000), 0) og_size = sem_image.size padded_image.paste(sem_image, ((3000 - og_size[0]) // 2, (3000 - og_size[1]) // 2)) light_image = semmap_to_lightmap(np.array(padded_image)) light_image.save(os.path.join(layout_dir, 'floor_lighttype_0.png')) if args.viz: from mpl_toolkits.axes_grid1 import make_axes_locatable fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(13, 5)) divider = make_axes_locatable(ax[0]) cax = divider.append_axes('right', size='5%', pad=0.05) im = ax[0].imshow(ins_image) fig.colorbar(im, cax=cax, orientation='vertical') divider = make_axes_locatable(ax[1]) cax = divider.append_axes('right', size='5%', pad=0.05) im = ax[1].imshow(sem_image) fig.colorbar(im, cax=cax, orientation='vertical') png_path = os.path.join(save_dir, 'room_sem_viz.png') fig.savefig(png_path) wall_patches = [ Polygon(bbox.get_coords(), True, color='b') for bbox in wall_bboxes_i ] window_patches = [ Polygon(bbox.get_coords(), True, color='b') for bbox in window_bboxes_i ] door_patches = [ Polygon(bbox.get_coords(), True, color='b') for bbox in door_bboxes_i ] furn_patches = [ Polygon(bbox.get_coords(), True, color='b') for c, bbox in furniture_bboxes_i ] floor_patches = [ Polygon(poly, True, color='b') for cat, poly in floor_polygons_i ] quivers = [ np.array([*bbox.center, *bbox.nrm]) for c, bbox in furniture_bboxes_i ] fig, ax = plt.subplots(figsize=(11, 11)) p0 = PatchCollection(floor_patches, cmap=matplotlib.cm.jet, alpha=0.3) colors = np.random.random(10) p0.set_array(colors) ax.add_collection(p0) p = PatchCollection(wall_patches, alpha=0.2) ax.add_collection(p) p1 = PatchCollection(window_patches, alpha=1) ax.add_collection(p1) p2 = PatchCollection(door_patches, alpha=1) p2.set_array(np.random.random(10)) ax.add_collection(p2) p3 = PatchCollection(furn_patches, cmap=matplotlib.cm.jet, alpha=1) colors = 100 * np.random.random(10) p3.set_array(colors) ax.add_collection(p3) ax.set_ylim([0, width / 100.]) ax.set_xlim([0, height / 100.]) if len(quivers) > 0: x, y, u, v = np.vstack(quivers).transpose() ax.quiver(x, y, u, v, scale=20) png_path = os.path.join(save_dir, 'viz.png') fig.savefig(png_path)
def get_bbox_vol(bbox): coord_xy = shape_poly(bbox[1].get_coords()) z = bbox[1].z return coord_xy.area * (z[1] - z[0])