def test_class_count_panoptic(self):
# test scene_01
scene_json = 'tests/data/dgp/test_scene/scene_01/scene_a8dc5ed1da0923563f85ea129f0e0a83e7fe1867.json'
scene_dir = os.path.dirname(scene_json)
scene = open_pbobject(scene_json, Scene)
annotation_enum = annotations_pb2.BOUNDING_BOX_2D
ontology = None
class_stats = get_scene_class_statistics(scene, scene_dir,
annotation_enum, ontology)
self.assertTrue(class_stats['Bicycle'] == 0)
self.assertTrue(class_stats['Bus/RV/Caravan'] == 3)
self.assertTrue(class_stats['Car'] == 152)
self.assertTrue(class_stats['Motorcycle'] == 0)
self.assertTrue(class_stats['Person'] == 4)
self.assertTrue(class_stats['Towed Object'] == 0)
self.assertTrue(class_stats['Trailer'] == 0)
self.assertTrue(class_stats['Train'] == 0)
self.assertTrue(class_stats['Truck'] == 0)
self.assertTrue(class_stats['Wheeled Slow'] == 0)
# test scene_02
scene_json = 'tests/data/dgp/test_scene/scene_02/scene_fe9f29d3bde25d182dcf88caf1011acd8cc13624.json'
scene_dir = os.path.dirname(scene_json)
scene = open_pbobject(scene_json, Scene)
annotation_enum = annotations_pb2.BOUNDING_BOX_2D
ontology = None
class_stats = get_scene_class_statistics(scene, scene_dir,
annotation_enum, ontology)
self.assertTrue(class_stats['Bicycle'] == 0)
self.assertTrue(class_stats['Bus/RV/Caravan'] == 0)
self.assertTrue(class_stats['Car'] == 33)
self.assertTrue(class_stats['Motorcycle'] == 0)
self.assertTrue(class_stats['Person'] == 0)
self.assertTrue(class_stats['Towed Object'] == 0)
self.assertTrue(class_stats['Trailer'] == 0)
self.assertTrue(class_stats['Train'] == 0)
self.assertTrue(class_stats['Truck'] == 6)
self.assertTrue(class_stats['Wheeled Slow'] == 0)
# test assertion
wrong_ontology = Ontology(items=[
OntologyItem(name='dummy',
id=0,
color=OntologyItem.Color(r=220, b=60, g=20),
isthing=True),
])
self.assertRaises(AssertionError, get_scene_class_statistics, scene,
scene_dir, annotation_enum, wrong_ontology)
def load(cls, annotation_file, ontology):
"""Load annotation from annotation file and ontology.
Parameters
----------
annotation_file: str
Full path to annotation
ontology: KeyPointOntology
Ontology for 2D keypoint tasks.
Returns
-------
KeyPoint2DAnnotationList
Annotation object instantiated from file.
"""
_annotation_pb2 = open_pbobject(annotation_file, KeyPoint2DAnnotations)
pointlist = [
KeyPoint2D(
point=np.float32([ann.point.x, ann.point.y]),
class_id=ontology.class_id_to_contiguous_id[ann.class_id],
color=ontology.colormap[ann.class_id],
attributes=getattr(ann, "attributes", {}),
) for ann in _annotation_pb2.annotations
]
return cls(ontology, pointlist)
def load(cls, annotation_file, ontology):
"""Load annotation from annotation file and ontology.
Parameters
----------
annotation_file: str
Full path to annotation
ontology: BoundingBoxOntology
Ontology for 3D bounding box tasks.
Returns
-------
BoundingBox3DAnnotationList
Annotation object instantiated from file.
"""
_annotation_pb2 = open_pbobject(annotation_file,
BoundingBox3DAnnotations)
boxlist = [
BoundingBox3D(
pose=Pose.load(ann.box.pose),
sizes=np.float32(
[ann.box.width, ann.box.length, ann.box.height]),
class_id=ontology.class_id_to_contiguous_id[ann.class_id],
instance_id=ann.instance_id,
color=ontology.colormap[ann.class_id],
attributes=getattr(ann, "attributes", {}),
num_points=ann.num_points,
occlusion=ann.box.occlusion,
truncation=ann.box.truncation)
for ann in _annotation_pb2.annotations
]
return cls(ontology, boxlist)
def load(cls, annotation_file, ontology):
"""Load annotation from annotation file and ontology.
Parameters
----------
annotation_file: str
Full path to annotation
ontology: KeyLineOntology
Ontology for 2D keyline tasks.
Returns
-------
KeyLine2DAnnotationList
Annotation object instantiated from file.
"""
_annotation_pb2 = open_pbobject(annotation_file, KeyLine2DAnnotations)
linelist = [
KeyLine2D(
line=np.float32([[vertex.x, vertex.y]
for vertex in ann.vertices]),
class_id=ontology.class_id_to_contiguous_id[ann.class_id],
color=ontology.colormap[ann.class_id],
attributes=getattr(ann, "attributes", {}),
) for ann in _annotation_pb2.annotations
]
return cls(ontology, linelist)
def upload_scenes(scene_dataset_json, s3_dst_dir):
"""Parallelized upload for scenes from a scene dataset JSON.
NOTE: This tool only verifies the presence of a scene, not the validity any of its contents.
"""
bucket_name, s3_base_path = convert_uri_to_bucket_path(s3_dst_dir)
dataset = open_pbobject(scene_dataset_json, SceneDataset)
local_dataset_root = os.path.dirname(os.path.abspath(scene_dataset_json))
if not dataset:
print('Failed to parse dataset artifacts {}'.format(scene_dataset_json))
sys.exit(0)
scene_dirs = []
for split in dataset.scene_splits.keys():
scene_dirs.extend([
os.path.join(local_dataset_root, os.path.dirname(filename))
for filename in dataset.scene_splits[split].filenames
])
# Make sure the scenes exist
with Pool(cpu_count()) as proc:
file_list = list(itertools.chain.from_iterable(proc.map(_get_scene_files, scene_dirs)))
# Upload the scene JSON, too.
file_list += [scene_dataset_json]
print("Creating file manifest for S3 for {} files".format(len(file_list)))
s3_file_list = [os.path.join(s3_base_path, os.path.relpath(_f, local_dataset_root)) for _f in file_list]
print("Done. Uploading to S3.")
parallel_upload_s3_objects(file_list, s3_file_list, bucket_name)
def load(cls, annotation_file, ontology):
"""Load annotation from annotation file and ontology.
Parameters
----------
annotation_file: str
Full path to annotation
ontology: BoundingBoxOntology
Ontology for 2D bounding box tasks.
Returns
-------
BoundingBox2DAnnotationList
Annotation object instantiated from file.
"""
_annotation_pb2 = open_pbobject(annotation_file, BoundingBox2DAnnotations)
boxlist = [
BoundingBox2D(
box=np.float32([ann.box.x, ann.box.y, ann.box.w, ann.box.h]),
class_id=ontology.class_id_to_contiguous_id[ann.class_id],
instance_id=ann.instance_id,
color=ontology.colormap[ann.class_id],
attributes=getattr(ann, "attributes", {}),
) for ann in _annotation_pb2.annotations
]
return cls(ontology, boxlist)
def __init__(
self,
scene_dataset_json_paths,
description,
version,
local_output_path=None,
bucket_path=None,
name=None,
email=None,
):
assert len(scene_dataset_json_paths) >= 2, \
'At least 2 scene datasets required in order to merge'
self.scene_dataset_json_paths = scene_dataset_json_paths
assert all([f.endswith('.json') for f in scene_dataset_json_paths]), 'Invalid scene dataset JSON.'
self.scene_datasets = [
open_remote_pb_object(f, SceneDataset) if f.startswith('s3://') else open_pbobject(f, SceneDataset)
for f in scene_dataset_json_paths
]
assert all([item is not None for item in self.scene_datasets]), \
'Some of the scene dataset failed to load'
self.local_output_path = local_output_path
self.scene_dataset_pb2 = SceneDataset()
# Copy Metadata.
self.scene_dataset_pb2.metadata.description = description
self.scene_dataset_pb2.metadata.version = version
self.scene_dataset_pb2.metadata.origin = self.scene_datasets[0].metadata.origin
self.scene_dataset_pb2.metadata.creation_date = get_date()
# raw_path for merged dataset is combination of *bucket_paths* for input datasets
self.scene_dataset_pb2.metadata.raw_path.value = ",".join(
set([
scene_dataset.metadata.bucket_path.value # NOTE: *bucket_path*
for scene_dataset in self.scene_datasets
if scene_dataset.metadata.bucket_path.value
])
)
# bucket_path is specified by the user (and defaults to the bucket_path of the first dataset)
assert bucket_path is None or bucket_path.startswith('s3://')
self.scene_dataset_pb2.metadata.bucket_path.value = bucket_path or self.scene_datasets[
0].metadata.bucket_path.value
# Write in available annotation types
ann_types = []
for scene_dataset in self.scene_datasets:
ann_types.extend(scene_dataset.metadata.available_annotation_types)
self.scene_dataset_pb2.metadata.available_annotation_types.extend(set(ann_types))
# Dataset name and creator email
self.scene_dataset_pb2.metadata.name = name or self.scene_datasets[0].metadata.name
self.scene_dataset_pb2.metadata.creator = email or self.scene_datasets[0].metadata.creator
def is_empty_annotation(annotation_file, annotation_type):
"""Check if JSON style annotation files are empty
Parameters
----------
annotations: str
Path to JSON file containing annotations for 2D/3D bounding boxes
Returns
-------
bool:
True if empty annotation, otherwise False
"""
with open(annotation_file) as _f:
annotations = open_pbobject(annotation_file, annotation_type)
return len(list(annotations.annotations)) == 0
def parse_annotations_3d_proto(annotation_file,
json_category_id_to_contiguous_id):
"""Parse annotations from BoundingBox2DAnnotations structure.
Parameters
----------
annotations: str
Path to JSON file containing annotations for 2D bounding boxes
json_category_id_to_contiguous_id: dict
Lookup from COCO style JSON id's to contiguous id's
transformation: Pose
Pose object that can be used to convert annotations to a new reference frame.
Returns
-------
tuple holding:
boxes: list of BoundingBox3D
Tensor containing bounding boxes for this sample
(pose.quat.qw, pose.quat.qx, pose.quat.qy, pose.quat.qz,
pose.tvec.x, pose.tvec.y, pose.tvec.z, width, length, height)
in absolute scale
class_ids: np.int64 array
Numpy array containing class ids (aligned with ``boxes``)
instance_ids: dict
Map from instance_id to tuple of (box, class_id)
"""
# *CAVEAT*: `attributes` field is defined in proto, but not being read here.
# TODO: read attributes (see above); change outputs of all function calls.
with open(annotation_file) as _f:
annotations = open_pbobject(annotation_file, BoundingBox3DAnnotations)
boxes, class_ids, instance_ids = [], [], {}
for i, ann in enumerate(list(annotations.annotations)):
boxes.append(
BoundingBox3D(
Pose.from_pose_proto(ann.box.pose),
np.float32([ann.box.width, ann.box.length,
ann.box.height]), ann.num_points,
ann.box.occlusion, ann.box.truncation))
class_ids.append(json_category_id_to_contiguous_id[ann.class_id])
instance_ids[ann.instance_id] = (boxes[i], class_ids[i])
return boxes, class_ids, instance_ids
def parse_annotations_2d_proto(annotation_file,
json_category_id_to_contiguous_id):
"""Parse annotations from BoundingBox2DAnnotations structure.
Parameters
----------
annotations: str
Path to JSON file containing annotations for 2D bounding boxes
json_category_id_to_contiguous_id: dict
Lookup from COCO style JSON id's to contiguous id's
Returns
-------
tuple holding:
boxes: torch.FloatTensor
Tensor containing bounding boxes for this sample
(x, y, w, h) in absolute pixel coordinates
class_ids: np.int64 array
Numpy array containing class ids (aligned with ``boxes``)
instance_ids: dict
Map from instance_id to tuple of (box, class_id)
attributes: list
list of dict mapping attribute names to values.
"""
with open(annotation_file) as _f:
annotations = open_pbobject(annotation_file, BoundingBox2DAnnotations)
boxes, class_ids, instance_ids, attributes = [], [], {}, []
for i, ann in enumerate(list(annotations.annotations)):
boxes.append(
np.float32([ann.box.x, ann.box.y, ann.box.w, ann.box.h]))
class_ids.append(json_category_id_to_contiguous_id[ann.class_id])
instance_ids[ann.instance_id] = (boxes[i], class_ids[i])
attributes.append(getattr(ann, 'attributes', {}))
return np.float32(boxes), np.int64(class_ids), instance_ids, attributes
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
def get_scene_class_statistics(scene,
scene_dir,
annotation_enum,
ontology=None):
"""Given a DGP scene, return class counts of the annotations in the scene.
Parameters
----------
scene: dgp.proto.scene_pb2.Scene
Scene Object.
scene_dir: string
s3 URL or local path to scene.
annotation_enum: dgp.proto.ontology_pb2.AnnotationType
Annotation type enum
ontology: dgp.proto.ontology_pb2.Ontology or None
Stats will be computed for this ontology. If None, the ontology will be read from the scene.
Returns
--------
scene_stats: OrderedDict
class_name: int
Counts of annotations for each class.
"""
datum_type, annotation_pb = _get_bounding_box_annotation_info(
annotation_enum)
if ontology is not None:
# Scene ontology must match input ontology
class_counts = OrderedDict({item.name: 0 for item in ontology.items})
id_name_map = {item.id: item.name for item in ontology.items}
ontology_sha = generate_uid_from_pbobject(ontology)
assert annotation_enum in scene.ontologies, 'Given annotation_enum not in scene.ontologies!'
assert scene.ontologies[
annotation_enum] == ontology_sha, 'Input ontology does not match Scene ontology!'
else:
# Just grab the ontology in the scene
ontology_path = os.path.join(
scene_dir, ONTOLOGY_FOLDER,
scene.ontologies[annotation_enum] + '.json')
ontology = Ontology(open_ontology_pbobject(ontology_path))
id_name_map = ontology.id_to_name
class_counts = OrderedDict({name: 0 for name in ontology.name_to_id})
for datum in scene.data:
# Get the annotation file for each object
if not datum.datum.WhichOneof('datum_oneof') == datum_type:
continue
datum_value = getattr(datum.datum, datum_type)
annotations = datum_value.annotations
if annotation_enum not in annotations:
continue
annotation_path = os.path.join(scene_dir, annotations[annotation_enum])
if annotation_path.startswith('s3://'):
annotation_object = open_remote_pb_object(annotation_path,
annotation_pb)
else:
annotation_object = open_pbobject(annotation_path, annotation_pb)
# Update class counts
for annotation in annotation_object.annotations:
class_counts[id_name_map[annotation.class_id]] += 1
return class_counts
