Spaces:
Runtime error
Runtime error
File size: 1,830 Bytes
fadb92b ef36c4f 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 | def build_dataset(image_set, args):
from .poly_data import build as build_poly
if args.dataset_name in ["stru3d", "cubicasa", "waffle", "r2g"]:
print(f"Build {args.dataset_name} {image_set} dataset")
return build_poly(image_set, args)
raise ValueError(f"dataset {args.dataset_name} not supported")
def get_dataset_class_labels(dataset_name):
semantics_label = None
if dataset_name == "stru3d":
semantics_label = {
0: "Living Room",
1: "Kitchen",
2: "Bedroom",
3: "Bathroom",
4: "Balcony",
5: "Corridor",
6: "Dining room",
7: "Study",
8: "Studio",
9: "Store room",
10: "Garden",
11: "Laundry room",
12: "Office",
13: "Basement",
14: "Garage",
15: "Misc.",
16: "Door",
17: "Window",
}
elif dataset_name == "cubicasa":
semantics_label = {
"Outdoor": 0,
"Kitchen": 1,
"Living Room": 2,
"Bed Room": 3,
"Bath": 4,
"Entry": 5,
"Storage": 6,
"Garage": 7,
"Undefined": 8,
"Window": 9,
"Door": 10,
}
elif dataset_name == "r2g":
semantics_label = {
"unknown": 0,
"living_room": 1,
"kitchen": 2,
"bedroom": 3,
"bathroom": 4,
"restroom": 5,
"balcony": 6,
"closet": 7,
"corridor": 8,
"washing_room": 9,
"PS": 10,
"outside": 11,
}
id2class = {v: k for k, v in semantics_label.items()} if semantics_label else None
return semantics_label, id2class
|