-
Notifications
You must be signed in to change notification settings - Fork 2
/
convert.py
155 lines (122 loc) · 6.16 KB
/
convert.py
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
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
import logging
from r2vstk.config_manager import ConfigManager
from r2vstk.house import House
import argparse
import os.path
import json
def run(conf: ConfigManager, source, output_path, scale_factor, save_previews,
save_room_json, skip_objects=False, adjust_short_walls=False,
classify_doors_and_windows=False, skip_rdr=False, r2v_annot=False):
"""
Converts Raster-to-Vector output/annotation to scene.json format.
"""
house = House()
# Pass in the scale factor
house.multiplication_factor = scale_factor
if r2v_annot:
# User has provided a r2v annotation file
house.load_r2v_annot_file(source)
else:
# User has provided a r2v output file
house.load_r2v_output_file(conf, source)
# Split walls that intersect with other walls into separate wall segments.
if conf.parser_config.split_walls.enabled:
house.split_source_walls(conf)
# Generate wall linkage graph and detect rooms.
house.generate_wall_graph(conf)
# Axis align nearly axis aligned walls.
if conf.parser_config.straighten_walls.enabled:
house.straighten_walls(conf)
# Assign annotation AABBs to rooms
house.populate_room_annotations(conf)
# Parse object annotation AABBs
if not skip_objects:
house.populate_object_annotations(conf)
# Separate doors from windows.
if classify_doors_and_windows:
house.classify_doors_and_windows()
# Populate additional details related to rooms such as room_id and room_type labels.
if r2v_annot:
house.populate_room_descriptions_from_r2v_annot(conf)
else:
house.populate_room_descriptions_from_r2v_output()
# Compute room-door-room connectivity graph edges
if not skip_rdr:
house.compute_rdr()
# Save preview sketches
if save_previews:
house.sketch_raw_annotations(conf, os.path.join(output_path, "raw_annot.png"))
logging.info("Saved {file}".format(file=os.path.join(output_path, "raw_annot.png")))
# Save wall mask
wall_mask = house.get_wall_mask()
with open(os.path.join(output_path, "wall_mask.svg"), 'w') as f:
f.write(wall_mask._repr_svg_())
# Save room level results
for i, room_key in enumerate(house.ordered_rooms):
if save_previews:
house.sketch_room_annotations(conf, room_key, os.path.join(output_path, "room_" + str(i) + ".png"))
logging.info("Saved {file}".format(file=os.path.join(output_path, "room_" + str(i) + ".png")))
if save_room_json:
room_json = house.get_room_json(conf, room_key, True, [], adjust_short_walls=adjust_short_walls)[0]
with open(os.path.join(output_path, room_json["id"] + ".arch.json"), "w") as f:
f.write(json.dumps(room_json, indent=3))
logging.info("Saved {file}".format(file=os.path.join(output_path, "room_" + str(i) + ".json")))
# Save scene.json
scene_json = house.get_scene_json(conf, adjust_short_walls=adjust_short_walls)
save_path = os.path.join(output_path, scene_json["scene"]["arch"]["id"] + ".scene.json")
with open(save_path, "w") as f:
f.write(json.dumps(scene_json, indent=3))
logging.info("Saved {file}".format(file=save_path))
# Save objectaabb.json
object_aabb_json = house.get_objectaabb_json(conf)
save_path = os.path.join(output_path, scene_json["scene"]["arch"]["id"] + ".objectaabb.json")
with open(save_path, "w") as f:
f.write(json.dumps(object_aabb_json, indent=3))
logging.info("Saved {file}".format(file=save_path))
return house
def run_args(conf: ConfigManager, args):
"""
Process command line args.
:param conf: ConfigManager
:param args: Command line args
"""
if not args.r2v_annot:
assert conf.room_types is not None
house = run(conf=conf, source=args.source[0], output_path=args.output_path[0], save_previews=not args.no_previews, save_room_json=args.room_json,
scale_factor=args.scale_factor[0], skip_objects=args.skip_objects,
adjust_short_walls=not args.do_not_adjust_short_walls,
classify_doors_and_windows=not args.do_not_classify_doors_and_windows,
skip_rdr=args.skip_rdr, r2v_annot=args.r2v_annot)
return house
def add_args(parser: argparse.ArgumentParser) -> None:
"""
Add command line args to the parser.
:param parser: Argument Parser
"""
parser.add_argument('output_path', metavar='output', type=str, nargs=1,
help='directory to save outputs from dataset')
parser.add_argument('source', metavar='source', type=str, nargs=1,
help='path to annotation file from dataset')
parser.add_argument("--r2v-annot", default=False, action="store_true",
help="Process a Raster-to-vector annotation file instead of R2V output file.")
parser.add_argument("--scale-factor", metavar="mf", type=float, nargs=1, default=[-1],
help="Scale factor to multiply all coordinates in the XY plane.")
parser.add_argument("--no-previews", default=False, action="store_true",
help="Don't generate PNG previews")
parser.add_argument("--room-json", default=False, action="store_true",
help="Generate individual room's json")
parser.add_argument("--skip-objects", default=False, action="store_true",
help="Don't place objects")
parser.add_argument("--do-not-adjust-short-walls", default=False, action="store_true",
help="Don't adjust height of short walls such as balconies")
parser.add_argument("--do-not-classify-doors-and-windows", default=False, action="store_true",
help="Don't classify holes as doors or windows")
parser.add_argument("--skip-rdr", default=False, action="store_true", help="Avoid computing RDR edges.")
if __name__ == "__main__":
conf = ConfigManager()
parser = argparse.ArgumentParser(description='Generate scene-toolkit compatible json files.')
conf.add_args(parser)
add_args(parser)
args = parser.parse_args()
conf.process_args(args, output_is_dir=True)
run_args(conf, args)