Skip to content

Commit dfe59cb

Browse files
authored
refactor: define the method to load tables as function (#180)
Signed-off-by: ktro2828 <kotaro.uetake@tier4.jp>
1 parent bc64ee1 commit dfe59cb

1 file changed

Lines changed: 45 additions & 41 deletions

File tree

t4_devkit/tier4.py

Lines changed: 45 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@
4141
Visibility,
4242
)
4343

44-
__all__ = ["DBMetadata", "load_metadata", "Tier4"]
44+
__all__ = ["DBMetadata", "load_metadata", "load_table", "Tier4"]
4545

4646

4747
@define
@@ -92,6 +92,29 @@ def load_metadata(db_root: str, revision: str | None = None) -> DBMetadata:
9292
return DBMetadata(data_root=data_root, dataset_id=dataset_id, version=version)
9393

9494

95+
def load_table(annotation_dir: str, schema: SchemaName) -> list[SchemaTable]:
96+
"""Load schema table from a JSON file.
97+
98+
If the schema is optional and there is no corresponding JSON file in dataset,
99+
returns empty list.
100+
101+
Args:
102+
annotation_dir (str): Path to the directory of JSON annotation schema files.
103+
schema (SchemaName): An enum member of `SchemaName`.
104+
105+
Returns:
106+
Loaded table data saved in `.json`.
107+
"""
108+
filepath = osp.join(annotation_dir, schema.filename)
109+
if not osp.exists(filepath) and schema.is_optional():
110+
return []
111+
112+
if not osp.exists(filepath):
113+
raise FileNotFoundError(f"{schema.value} is mandatory.")
114+
115+
return build_schema(schema, filepath)
116+
117+
95118
class Tier4:
96119
"""Database class for T4 dataset to help query and retrieve information from the database."""
97120

@@ -152,27 +175,29 @@ def __init__(
152175
print("======\nLoading T4 tables...")
153176

154177
# assign tables explicitly
155-
self.attribute: list[Attribute] = self.__load_table__(SchemaName.ATTRIBUTE)
156-
self.calibrated_sensor: list[CalibratedSensor] = self.__load_table__(
157-
SchemaName.CALIBRATED_SENSOR
178+
self.attribute: list[Attribute] = load_table(self.annotation_dir, SchemaName.ATTRIBUTE)
179+
self.calibrated_sensor: list[CalibratedSensor] = load_table(
180+
self.annotation_dir, SchemaName.CALIBRATED_SENSOR
158181
)
159-
self.category: list[Category] = self.__load_table__(SchemaName.CATEGORY)
160-
self.ego_pose: list[EgoPose] = self.__load_table__(SchemaName.EGO_POSE)
161-
self.instance: list[Instance] = self.__load_table__(SchemaName.INSTANCE)
162-
self.keypoint: list[Keypoint] = self.__load_table__(SchemaName.KEYPOINT)
163-
self.log: list[Log] = self.__load_table__(SchemaName.LOG)
164-
self.map: list[Map] = self.__load_table__(SchemaName.MAP)
165-
self.object_ann: list[ObjectAnn] = self.__load_table__(SchemaName.OBJECT_ANN)
166-
self.sample_annotation: list[SampleAnnotation] = self.__load_table__(
167-
SchemaName.SAMPLE_ANNOTATION
182+
self.category: list[Category] = load_table(self.annotation_dir, SchemaName.CATEGORY)
183+
self.ego_pose: list[EgoPose] = load_table(self.annotation_dir, SchemaName.EGO_POSE)
184+
self.instance: list[Instance] = load_table(self.annotation_dir, SchemaName.INSTANCE)
185+
self.keypoint: list[Keypoint] = load_table(self.annotation_dir, SchemaName.KEYPOINT)
186+
self.log: list[Log] = load_table(self.annotation_dir, SchemaName.LOG)
187+
self.map: list[Map] = load_table(self.annotation_dir, SchemaName.MAP)
188+
self.object_ann: list[ObjectAnn] = load_table(self.annotation_dir, SchemaName.OBJECT_ANN)
189+
self.sample_annotation: list[SampleAnnotation] = load_table(
190+
self.annotation_dir, SchemaName.SAMPLE_ANNOTATION
168191
)
169-
self.sample_data: list[SampleData] = self.__load_table__(SchemaName.SAMPLE_DATA)
170-
self.sample: list[Sample] = self.__load_table__(SchemaName.SAMPLE)
171-
self.scene: list[Scene] = self.__load_table__(SchemaName.SCENE)
172-
self.sensor: list[Sensor] = self.__load_table__(SchemaName.SENSOR)
173-
self.surface_ann: list[SurfaceAnn] = self.__load_table__(SchemaName.SURFACE_ANN)
174-
self.vehicle_state: list[VehicleState] = self.__load_table__(SchemaName.VEHICLE_STATE)
175-
self.visibility: list[Visibility] = self.__load_table__(SchemaName.VISIBILITY)
192+
self.sample_data: list[SampleData] = load_table(self.annotation_dir, SchemaName.SAMPLE_DATA)
193+
self.sample: list[Sample] = load_table(self.annotation_dir, SchemaName.SAMPLE)
194+
self.scene: list[Scene] = load_table(self.annotation_dir, SchemaName.SCENE)
195+
self.sensor: list[Sensor] = load_table(self.annotation_dir, SchemaName.SENSOR)
196+
self.surface_ann: list[SurfaceAnn] = load_table(self.annotation_dir, SchemaName.SURFACE_ANN)
197+
self.vehicle_state: list[VehicleState] = load_table(
198+
self.annotation_dir, SchemaName.VEHICLE_STATE
199+
)
200+
self.visibility: list[Visibility] = load_table(self.annotation_dir, SchemaName.VISIBILITY)
176201

177202
# make reverse indexes for common lookups
178203
self.__make_reverse_index__(verbose)
@@ -217,27 +242,6 @@ def bag_dir(self) -> str:
217242
"""Return the path to ROS bag directory."""
218243
return osp.join(self.data_root, "input_bag")
219244

220-
def __load_table__(self, schema: SchemaName) -> list[SchemaTable]:
221-
"""Load schema table from a json file.
222-
223-
If the schema is optional and there is no corresponding json file in dataset,
224-
returns empty list.
225-
226-
Args:
227-
schema (SchemaName): An enum member of `SchemaName`.
228-
229-
Returns:
230-
Loaded table data saved in `.json`.
231-
"""
232-
filepath = osp.join(self.annotation_dir, schema.filename)
233-
if not osp.exists(filepath) and schema.is_optional():
234-
return []
235-
236-
if not osp.exists(filepath):
237-
raise FileNotFoundError(f"{schema.value} is mandatory.")
238-
239-
return build_schema(schema, filepath)
240-
241245
def __make_reverse_index__(self, verbose: bool) -> None:
242246
"""De-normalize database to create reverse indices for common cases.
243247

0 commit comments

Comments
 (0)