8000 feat: Add ProbabilisticKeyPoint3D by kunlin596 · Pull Request #171 · TRI-ML/dgp · GitHub
[go: up one dir, main page]
More Web Proxy on the site http://driver.im/
Skip to content

feat: Add ProbabilisticKeyPoint3D #171

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open

feat: Add ProbabilisticKeyPoint3D #171

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 7 additions & 3 deletions dgp/annotations/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
from dgp.annotations.semantic_segmentation_2d_annotation import SemanticSegmentation2DAnnotation # isort:skip
from dgp.annotations.key_line_2d_annotation import KeyLine2DAnnotationList # isort:skip
from dgp.annotations.key_line_3d_annotation import KeyLine3DAnnotationList # isort:skip
from dgp.annotations.key_line_3d_annotation import ProbabilisticKeyLine3DAnnotationList # isort:skip
from dgp.annotations.key_point_2d_annotation import KeyPoint2DAnnotationList # isort:skip
from dgp.annotations.key_point_3d_annotation import KeyPoint3DAnnotationList # isort:skip
from dgp.annotations.depth_annotation import DenseDepthAnnotation # isort:skip
Expand All @@ -36,12 +37,13 @@
"key_point_3d": KeyPointOntology,
"key_line_2d": KeyLineOntology,
"key_line_3d": KeyLineOntology,
"probabilistic_key_line_3d": KeyLineOntology,
"agent_behavior": AgentBehaviorOntology,
"depth": None,
"surface_normals_2d": None,
"surface_normals_3d": None,
"motion_vectors_2d": None,
"motion_vectors_3d": None
"motion_vectors_3d": None,
}

# Annotation objects for each annotation type
Expand All @@ -54,7 +56,8 @@
"key_point_3d": KeyPoint3DAnnotationList,
"key_line_2d": KeyLine2DAnnotationList,
"key_line_3d": KeyLine3DAnnotationList,
"depth": DenseDepthAnnotation
"probabilistic_key_line_3d": ProbabilisticKeyLine3DAnnotationList,
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I believe only if you want support for this type in the legacy dataloaders? I'm not sure if this is sufficient.

"depth": DenseDepthAnnotation,
}

# Annotation groups for each annotation type: 2d/3d
Expand All @@ -73,5 +76,6 @@
"key_line_2d": "2d",
"key_point_3d": "3d",
"key_line_3d": "3d",
"depth": "2d"
"probabilistic_key_line_3d": "3d",
"depth": "2d",
}
175 changes: 171 additions & 4 deletions dgp/annotations/key_line_3d_annotation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,16 +1,27 @@
# Copyright 2022 Woven Planet. All rights reserved.
from typing import Any, Dict, List, Union

import numpy as np

from dgp.annotations.base_annotation import Annotation
from dgp.annotations.ontology import KeyLineOntology
from dgp.proto.annotations_pb2 import KeyLine3DAnnotation, KeyLine3DAnnotations
from dgp.proto.annotations_pb2 import \
ProbabilisticKeyLine3DAnnotation as ProtoProbabilisticKeyLine3DAnnotation
from dgp.proto.annotations_pb2 import \
ProbabilisticKeyLine3DAnnotations as ProtoProbabilisticKeyLine3DAnnotations
from dgp.proto.annotations_pb2 import \
ProbabilisticKeyPoint3D as ProtoProbabilisticKeyPoint3D
from dgp.utils.protobuf import (
generate_uid_from_pbobject,
parse_pbobject,
save_pbobject_as_json,
)
from dgp.utils.structures.key_line_3d import KeyLine3D
from dgp.utils.structures.key_point_3d import KeyPoint3D
from dgp.utils.structures.key_line_3d import KeyLine3D, ProbabilisticKeyLine3D
from dgp.utils.structures.key_point_3d import (
KeyPoint3D,
ProbabilisticKeyPoint3D,
)


class KeyLine3DAnnotationList(Annotation):
Expand Down Expand Up @@ -79,10 +90,10 @@ def to_proto(self):
class_id=line.class_id,
instance_id=line.instance_id,
color=line.color,
attributes=line.attributes
attributes=line.attributes,
).to_proto() for x, y, z in zip(line.x, line.y, line.z)
],
attributes=line.attributes
attributes=line.attributes,
) for line in self._linelist
]
)
Expand Down Expand Up @@ -138,3 +149,159 @@ def instance_ids(self):
def hexdigest(self):
"""Reproducible hash of annotation."""
return generate_uid_from_pbobject(self.to_proto())


class ProbabilisticKeyLine3DAnnotationList(Annotation):
"""Container for probabilistic 3D key line annotations.

Parameters
----------
ontology: KeyLineOntology
Ontology for 3D key line tasks.

linelist: list[ProbabilisticKeyLine3D]
List of ProbabilisticKeyLine3D objects. See `dgp/utils/structures/key_line_3d` for more details.
"""
def __init__(self, ontology: KeyLineOntology, linelist: List[ProbabilisticKeyLine3D]) -> None:
Annotation.__init__(self, ontology)
assert isinstance(self._ontology, KeyLineOntology), "Trying to load annotation with wrong type of ontology!"
for line in linelist:
assert isinstance(
line, ProbabilisticKeyLine3D
), f"Can only instantiate an annotation from a list of ProbabilisticKeyLine3D, not {type(line)}!"
self._linelist = linelist

@classmethod
def load(
cls, annotation_file: Union[str, bytes], ontology: KeyLineOntology
) -> "ProbabilisticKeyLine3DAnnotationList":
"""Load annotation from annotation file and ontology.

Parameters
----------
annotation_file: str or bytes
Full path to annotation or bytestring

ontology: KeyLineOntology
Ontology for 3D key line tasks.

Returns
-------
KeyLine3DAnnotationList
Annotation object instantiated from file.
"""
_annotation_pb2 = parse_pbobject(
annotation_file,
ProtoProbabilisticKeyLine3DAnnotations,
)
annotation_list = [
ProbabilisticKeyLine3D(
points=[
ProbabilisticKeyPoint3D(
point=np.asarray([
vert.x,
vert.y,
vert.z,
]),
covariance=np.asarray([
vert.var_x,
vert.cov_xy,
vert.cov_xz,
vert.var_y,
vert.cov_yz,
vert.var_z,
]),
) for vert in annotation.vertices
],
class_id=ontology.class_id_to_contiguous_id[annotation.class_id],
instance_id=annotation.key,
color=ontology.colormap[annotation.class_id],
attributes=getattr(annotation, "attributes", {}),
) for annotation in _annotation_pb2.annotations
]
return cls( linelist=annotation_list)

def to_proto(self):
"""Return annotation as pb object.

Returns
-------
KeyLine3DAnnotations
Annotation as defined in `proto/annotations.proto`
"""
return ProtoProbabilisticKeyLine3DAnnotations(
annotations=[
ProtoProbabilisticKeyLine3DAnnotation(
class_id=self._ontology.contiguous_id_to_class_id[line.class_id],
vertices=[
ProtoProbabilisticKeyPoint3D(
x=point.x,
y=point.y,
z=point.z,
var_x=point.cov3.var_x,
cov_xy=point.cov3.cov_xy,
cov_xz=point.cov3.cov_xz,
var_y=point.cov3.var_y,
cov_yz=point.cov3.cov_yz,
var_z=point.cov3.var_z,
) for point in line
],
key=line.instance_id,
attributes=line.attributes,
) for line in self._linelist
]
)

def save(self, save_dir: str) -> None:
"""Serialize Annotation object and saved to specified directory.

Annotations are saved in format <save_dir>/<sha>.<ext>

Parameters
----------
save_dir: str
Directory in which annotation is saved.

Returns
-------
output_annotation_file: str
Full path to saved annotation.
"""
return save_pbobject_as_json(self.to_proto(), save_path=save_dir)

def __len__(self) -> int:
return len(self._linelist)

def __getitem__(self, index: int) -> ProbabilisticKeyLine3D:
"""Return a single 3D keyline"""
return self._linelist[index]

def render(self):
"""Batch rendering function for keylines."""
raise NotImplementedError

@property
def xyz(self) -> np.ndarray:
"""Return lines as (N, 3) np.ndarray in format ([x, y, z])"""
return np.array([line.xyz.tolist() for line in self._linelist], dtype=np.float32)

@property
def class_ids(self) -> np.ndarray:
"""Return class ID for each line, with ontology applied:
class IDs mapped to a contiguous set.
"""
return np.array([line.class_id for line in self._linelist], dtype=np.int64)

@property
def attributes(self) -> List[Dict[str, Any]]:
"""Return a list of dictionaries of attribute name to value."""
return [line.attributes for line in self._linelist]

@property
def instance_ids(self) -> np.ndarray:
return np.array([line.instance_id for line in self._linelist], dtype=np.int64)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this dtype can cause problems. instance_id is optional and when None gets set to a string. This is a known issue in BoundingBox3DAnnotationList where printing throws an error


@property
def hexdigest(self) -> str:
"""Reproducible hash of annotation."""
return generate_uid_from_pbobject(self.to_proto())
65 changes: 65 additions & 0 deletions dgp/proto/annotations.proto
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -252,6 +252,22 @@ message KeyPoint3D {
float z = 3;
}

// 3D point with uncertainty information.
message ProbabilisticKeyPoint3D {
// (x, y, z) point (in 3D Cartesian coordinates).
float x = 1;
float y = 2;
float z = 3;

// The correlation terms in the 3x3 covariance matrix.
float var_x = 4;
float cov_xy = 5;
float cov_xz = 6;
float var_y = 7;
float cov_yz = 8;
float var_z = 9;
}

// 3D point annotation.
message KeyPoint3DAnnotation {
// Class identifier (should be in [0, num_classes - 1]),
Expand All @@ -271,6 +287,26 @@ message KeyPoint3DAnnotation {
string key = 4;
}


// 3D probabilistic point annotation.
message ProbabilisticKeyPoint3DAnnotation {
// Class identifier (should be in [0, num_classes - 1]),
// where num_classes is the total number of classes in your ontology.
uint32 class_id = 1;

// 3D point.
ProbabilisticKeyPoint3D point = 2;

// A map of attribute names to their values.
// Add only key/value pairs that are stored in a project document accessible
// to project contributors.
map<string, string> attributes = 3;

// An identifier key. Used to link with other annotations, which specify
// this key in their instance to link to corresponding ProbabilisticKeyPoint3DAnnotation.
string key = 4;
}

// 3D line annotation.
message KeyLine3DAnnotation{
// Class identifier (should be in [0, num_classes - 1]),
Expand All @@ -290,6 +326,25 @@ message KeyLine3DAnnotation{
string key = 4;
}

// 3D probabilistic line annotation.
message ProbabilisticKeyLine3DAnnotation {
// Class identifier (should be in [0, num_classes - 1]),
// where num_classes is the total number of classes in your ontology.
uint32 class_id = 1;

// 3D line.
repeated ProbabilisticKeyPoint3D vertices = 2;

// A map of attribute names to their values.
// Add only key/value pairs that are stored in a project document accessible
// to project contributors.
map<string, string> attributes = 3;

// An identifier key. Used to link with other annotations, which specify
// this key in their instance to link to corresponding KeyLine3DAnnotation.
string key = 4;
}

message PolygonPoint3D {
// (x, y, z) point (in 3D Cartesian coordinates).
float x = 1;
Expand Down Expand Up @@ -401,11 +456,21 @@ message KeyPoint3DAnnotations {
repeated KeyPoint3DAnnotation annotations = 1;
}

// List of ProbabilisticKeyPoint3DAnnotation.
message ProbabilisticKeyPoint3DAnnotations {
repeated ProbabilisticKeyPoint3DAnnotation annotations = 1;
}

// List of KeyLine3DAnnotation.
message KeyLine3DAnnotations {
repeated KeyLine3DAnnotation annotations = 1;
}

// List of ProbabilisticKeyLine3DAnnotation.
message ProbabilisticKeyLine3DAnnotations {
repeated ProbabilisticKeyLine3DAnnotation annotations = 1;
}

// List of Polygon3DAnnotation.
message Polygon3DAnnotations {
repeated Polygon3DAnnotation annotations = 1;
Expand Down
Loading
Loading
0