Spaces:
Sleeping
Sleeping
| from xml.dom import minidom | |
| import cv2 | |
| import numpy as np | |
| import math | |
| from annotation_builder import AnnotationBuilder as AnnBuild | |
| from pycocotools.coco import COCO | |
| from shapely import geometry | |
| from datetime import datetime | |
| import os | |
| import sys | |
| import random | |
| from from_root import from_root | |
| from rcnn_model.preprocessing.cleaning_single_image import preprocess_image | |
| from rcnn_model.utils.floorplan_vectorizer_utils import get_image_size, draw_from_coco | |
| ### After running, its split with https://github.com/akarazniewicz/cocosplit | |
| ### This may or may not be temporary | |
| ### Main functionality ### | |
| scale_factor = .5 | |
| dataset_root = str(from_root("dataset"))+"/" | |
| def main(): | |
| extract_all_cubicasa_anns(True) | |
| def extract_all_cubicasa_anns(export_image=False): | |
| #initialize annotation builder | |
| ann_builder = AnnBuild() | |
| ann_builder.set_info("converted from cubicasa 5k SVG file","cubicasa 5k","https://github.com/cubicasa/cubicasa5k",datetime(2019,5,24)) | |
| ann_builder.add_license("Creative Commons Attribution-NonCommercial 4.0 International License", "http://creativecommons.org/licenses/by-nc/4.0/") | |
| #iterate through cubicasa files | |
| for name in os.listdir(str(from_root(dataset_root+"cubicasa_data/"))): | |
| process_cubicasa_image(ann_builder, name) | |
| #save data | |
| print("SAVING TO annotations/cubicasa_coco.json") | |
| ann_builder.save_file(str(from_root(dataset_root+"annotations/cubicasa_coco.json"))) | |
| if(export_image): | |
| save_validation_images(str(from_root(dataset_root+"annotations/cubicasa_coco.json"))) | |
| def process_cubicasa_image(ann_builder, name): | |
| #load and preprocess image | |
| print("\nprocessing "+name) | |
| source_img_path = str(from_root(dataset_root+"cubicasa_data/"+name+"/F1_scaled.png")) | |
| processed_img_path = str(from_root(dataset_root+"preprocessed/casa"+name+".png")) | |
| apply_preprocessing(source_img_path, processed_img_path) | |
| #load svg | |
| source_svg_path = str(from_root(dataset_root+"cubicasa_data/"+name+"/model.svg")) | |
| print("from "+source_svg_path) | |
| print("image in "+processed_img_path) | |
| #extract data from svg | |
| try: | |
| width, height = get_image_size(processed_img_path) | |
| ann_builder = process_cubicasa(ann_builder, source_svg_path, processed_img_path, width, height) | |
| except: | |
| print("ERROR while extracting "+name) | |
| print(sys.exc_info()) | |
| def find_svg(path, name): | |
| for file in os.listdir(path): | |
| found_name = file.startswith(name+"_gt_") | |
| if(found_name): | |
| found_svg = file.endswith(".svg") | |
| if(found_svg): | |
| return path+file | |
| def process_cubicasa(ann_builder, sourve_svg_path, source_img_path, width, height): | |
| #Get points | |
| doc = minidom.parse(sourve_svg_path) | |
| walls = extract_casa_elements_with_id("Wall",doc) | |
| windows = extract_casa_elements_with_id("Window",doc) | |
| doors = extract_casa_elements_with_id("Door",doc) | |
| doc.unlink() | |
| #export to JSON and potentially imges for visual confirmation that the process works | |
| ann_builder = export_to_builder_casa(ann_builder,source_img_path,width,height,walls,doors,windows) | |
| return ann_builder | |
| ### Coco Formatting/Export ### | |
| def export_to_builder_casa(ann_builder,source_img,width,height,walls,doors,windows): | |
| #initialization | |
| id = ann_builder.add_image(source_img, width, height) | |
| #walls | |
| wall_polygons = get_features_from_ann_set(walls) | |
| door_polygons = get_features_from_ann_set(doors) | |
| window_polygons = get_features_from_ann_set(windows) | |
| features = wall_polygons + door_polygons + window_polygons | |
| rooms = create_rooms_from_features(features, width, height) | |
| for poly in rooms.geoms: | |
| ann_builder.add_annotation(id, 2, poly) | |
| return ann_builder | |
| def get_features_from_ann_set(set, coco = None, image_id = 0, category_id = 0): | |
| polygons = [] | |
| for points in set: | |
| poly = geometry.Polygon([[p[0], p[1]] for p in points]) | |
| if(coco is not None): | |
| coco.add_annotation(image_id, category_id, poly) | |
| polygons.append(poly) | |
| return polygons | |
| def create_rooms_from_features(features, width, height): | |
| room_polygons = geometry.Polygon([(0,0), | |
| (width,0), | |
| (width,height), | |
| (0,height) | |
| ]) | |
| for poly in features: | |
| room_polygons = room_polygons.difference(poly,3) | |
| return geometry.MultiPolygon(room_polygons.geoms[1:]) #this eliminates the exterior from the rooms | |
| def apply_preprocessing(source_path, processed_path): | |
| img = cv2.imread(source_path) | |
| small_img = cv2.resize(img, (0,0), fx=scale_factor, fy=scale_factor) | |
| cv2.imwrite(processed_path,small_img) | |
| processed_img = preprocess_image(processed_path) | |
| #small = cv2.resize(processed, (0,0), fx=scale_factor, fy=scale_factor) | |
| cv2.imwrite(processed_path,processed_img) | |
| print(get_image_size(source_path)) | |
| print(get_image_size(processed_path)) | |
| ### SVG element extraction ### | |
| def get_casa_size(doc): | |
| path = doc.getElementsByTagName('svg')[0] | |
| return int(float(path.getAttribute('width'))), int(float(path.getAttribute('height'))) | |
| def extract_casa_elements_with_id(id, doc): | |
| elements = [] | |
| for path in doc.getElementsByTagName('g'): | |
| #iterates through everything and finds items labelled as walls | |
| if(id in path.getAttribute('id')): | |
| #luckily, the first attribute after all of these is a polygon containing a list of coordinate points | |
| string = path.firstChild.getAttribute('points') | |
| points = points_string_to_int_points(string) | |
| elements.append(points) | |
| return elements | |
| ### Helper Functions ### | |
| def quadrilateral_to_line(points): | |
| base_point = [0,0] | |
| points.sort(key=lambda p: check_distance(base_point,p)) | |
| base_point = points[0] | |
| points.sort(key=lambda p: check_distance(base_point,p)) | |
| return np.array([get_midpoint(points[0], points[1]), get_midpoint(points[2], points[3])]) | |
| def check_distance(point_A, point_B): | |
| return math.sqrt(((point_A[0]-point_B[0]) ** 2) + ((point_A[1]-point_B[1]) ** 2)) | |
| def get_midpoint(point_A, point_B): | |
| return np.array([round((point_A[0]+point_B[0])/2), round((point_A[1]+point_B[1])/2)]) | |
| def points_string_to_int_points(string): | |
| return [[int(round(float(pi)*scale_factor)) for pi in p.split(",")] for p in string.split()] | |
| ### Validation Images ### | |
| def save_validation_images(filepath): | |
| count = 0 | |
| result = COCO(filepath) | |
| for id in random.sample(result.getImgIds(), 15): | |
| print("IMAGE "+str(result.imgs[id])) | |
| validation_path = str(from_root(dataset_root+"validation_images/casa_"+str(count))) | |
| draw_from_coco(id, result, validation_path) | |
| count+=1 | |
| main() |