def load_remote_scene_dataset_json(scene_dataset_json_url, save_dir=None):
"""Fetch scene dataset JSON from remote given its remote uri.
Parameters
----------
scene_dataset_json_uri: str
Remote scene dataset JSON URI.
save_dir: str, default: None
Directory to save dataset JSON.
Returns
-------
dataset: dgp.proto.dataset_pb2.SceneDataset
SceneDataset proto object
"""
logging.info('Fetching scene dataset {}'.format(scene_dataset_json_url))
bucket_name, s3_base_path = convert_uri_to_bucket_path(
scene_dataset_json_url)
dataset_blob = get_string_from_s3_file(s3_bucket(bucket_name),
s3_base_path)
dataset = Parse(dataset_blob, SceneDataset())
if save_dir is not None:
dataset_dir = os.path.join(
save_dir,
'{name}_{version}'.format(name=dataset.metadata.name,
version=dataset.metadata.version))
save_path = os.path.join(dataset_dir, os.path.basename(s3_base_path))
if not os.path.exists(save_path):
os.makedirs(dataset_dir, exist_ok=True)
save_pbobject_as_json(dataset, save_path)
return dataset
def save(self, save_dir):
"""Serialize Annotation object and saved to specified directory. Annotations are saved in format <save_dir>/<sha>.<ext>
Returns
-------
output_annotation_file: str
Full path to saved annotation
"""
return save_pbobject_as_json(self.to_proto(), save_path=save_dir)
def write_dataset(self, upload=False):
"""Write the final scene dataset JSON.
Parameters
----------
upload: bool, default: False
If true, upload the dataset to the scene pool in s3.
Returns
-------
scene_dataset_json_path: str
Path of the scene dataset JSON file created.
"""
scene_dataset_json_path = os.path.join(
self.local_output_path,
'{}_v{}.json'.format(self.scene_dataset_pb2.metadata.name, self.scene_dataset_pb2.metadata.version)
)
save_pbobject_as_json(self.scene_dataset_pb2, scene_dataset_json_path)
# Printing SceneDataset scene counts per split (post-merging)
logging.info('-' * 80)
logging.info(
'Output SceneDataset {} has: {} train, {} val, {} test'.format(
scene_dataset_json_path, len(self.scene_dataset_pb2.scene_splits[dataset_pb2.TRAIN].filenames),
len(self.scene_dataset_pb2.scene_splits[dataset_pb2.VAL].filenames),
len(self.scene_dataset_pb2.scene_splits[dataset_pb2.TEST].filenames)
)
)
s3_path = os.path.join(
self.scene_dataset_pb2.metadata.bucket_path.value, os.path.basename(scene_dataset_json_path)
)
if upload:
s3_copy(scene_dataset_json_path, s3_path)
else:
logging.info(
'Upload the DGP-compliant scene dataset JSON to s3 via `aws s3 cp --acl bucket-owner-full-control {} {}`'
.format(scene_dataset_json_path, s3_path)
)
return scene_dataset_json_path
def save(self, save_dir):
"""Write out ontology items to `<sha>.json`. SHA generated from Ontology proto object.
Parameters
----------
save_dir: str
Directory in which to save serialized ontology.
Returns
-------
output_ontology_file: str
Path to serialized ontology file.
"""
os.makedirs(save_dir, exist_ok=True)
return save_pbobject_as_json(self.to_proto(), save_path=save_dir)
def _write_annotation(filename, annotation):
"""Utility function for writing 3d annotations."""
_mkdir(os.path.dirname(filename))
save_pbobject_as_json(annotation, filename)
def clone_scene_as_autolabel(dataset_root, autolabel_root, autolabel_model,
autolabel_type):
"""Helper function to copy a scene directory for use as autolabels
Parameters
----------
dataset_root: str
Path to dataset root folder containing scene folders
autolabel_root: str
Path to where autolabels should be stored
autolabel_model: str
Name of autolabel model
autolabel_type: str
Annotation type i.e., 'bounding_box_3d', 'depth' etc
"""
# For each scene dir, copy the requested annotation into the new autolabel folder
autolabel_dirs = []
for scene_dir in os.listdir(dataset_root):
if not os.path.isdir(os.path.join(dataset_root, scene_dir)):
continue
# Clear any existing folders
autolabel_scene_dir = os.path.join(autolabel_root, scene_dir,
AUTOLABEL_FOLDER, autolabel_model)
if os.path.exists(autolabel_scene_dir):
rmtree(autolabel_scene_dir)
os.makedirs(autolabel_scene_dir, exist_ok=True)
full_scene_dir = os.path.join(dataset_root, scene_dir)
autolabel_dirs.append(autolabel_scene_dir)
for scene_json in os.listdir(full_scene_dir):
if 'scene' in scene_json and scene_json.endswith('json'):
base_scene = open_pbobject(
os.path.join(full_scene_dir, scene_json), Scene)
for i in range(len(base_scene.data)):
name = base_scene.data[i].id.name
datum = base_scene.data[i].datum
datum_type = datum.WhichOneof('datum_oneof')
datum_value = getattr(
datum, datum_type
) # This is datum.image or datum.point_cloud etc
annotation_type_id = ANNOTATION_KEY_TO_TYPE_ID[
autolabel_type]
current_annotation = datum_value.annotations[
annotation_type_id]
# NOTE: this should not actually change the path but is included for clarity
datum_value.annotations[annotation_type_id] = os.path.join(
ANNOTATION_TYPE_ID_TO_FOLDER[autolabel_type], name,
os.path.basename(current_annotation))
save_pbobject_as_json(
base_scene,
os.path.join(autolabel_scene_dir,
AUTOLABEL_SCENE_JSON_NAME))
# Only modify one scene.json, test scene should not contain multiple scene.jsons
break
ontology_dir = os.path.join(autolabel_scene_dir, ONTOLOGY_FOLDER)
if os.path.exists(ontology_dir):
rmtree(ontology_dir)
copytree(os.path.join(full_scene_dir, ONTOLOGY_FOLDER), ontology_dir)
annotation_dir = os.path.join(
autolabel_scene_dir, ANNOTATION_TYPE_ID_TO_FOLDER[autolabel_type])
if os.path.exists(annotation_dir):
rmtree(annotation_dir)
copytree(
os.path.join(full_scene_dir,
ANNOTATION_TYPE_ID_TO_FOLDER[autolabel_type]),
annotation_dir)
return autolabel_dirs