#!/usr/bin/env python3
#
# Copyright 2021 Graviti. Licensed under MIT License.
#
# pylint: disable=invalid-name
"""Dataloader of NeolixOD dataset."""
import os
from quaternion import from_rotation_vector
from tensorbay.dataset import Data, Dataset
from tensorbay.label import LabeledBox3D
from tensorbay.opendataset._utility import glob
DATASET_NAME = "NeolixOD"
[docs]def NeolixOD(path: str) -> Dataset:
"""`Neolix OD <https://gas.graviti.cn/dataset\
/graviti-open-dataset/NeolixOD>`_ dataset.
The file structure should be like::
<path>
bins/
<id>.bin
labels/
<id>.txt
...
Arguments:
path: The root directory of the dataset.
Returns:
Loaded :class:`~tensorbay.dataset.dataset.Dataset` instance.
"""
root_path = os.path.abspath(os.path.expanduser(path))
dataset = Dataset(DATASET_NAME)
dataset.load_catalog(os.path.join(os.path.dirname(__file__), "catalog.json"))
segment = dataset.create_segment()
point_cloud_paths = glob(os.path.join(root_path, "bins", "*.bin"))
for point_cloud_path in point_cloud_paths:
data = Data(point_cloud_path)
data.label.box3d = []
point_cloud_id = os.path.basename(point_cloud_path)[:6]
label_path = os.path.join(root_path, "labels", f"{point_cloud_id}.txt")
with open(label_path, encoding="utf-8") as fp:
for label_value_raw in fp:
label_value = label_value_raw.rstrip().split()
label = LabeledBox3D(
size=[float(label_value[10]), float(label_value[9]), float(label_value[8])],
translation=[
float(label_value[11]),
float(label_value[12]),
float(label_value[13]) + 0.5 * float(label_value[8]),
],
rotation=from_rotation_vector((0, 0, float(label_value[14]))),
category=label_value[0],
attributes={
"Occlusion": int(label_value[1]),
"Truncation": bool(int(label_value[2])),
"Alpha": float(label_value[3]),
},
)
data.label.box3d.append(label)
segment.append(data)
return dataset