cassdlcm/mediapipe
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fixed a to_pb2() bug in the landmark_detection_result container

Browse Source
This commit is contained in:
kinaryml 2022-11-08 02:55:19 -08:00
parent 0402ee383f
commit 9c9519eeb8
1 changed files with 19 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
mediapipe/tasks/python/components/containers/landmark_detection_result.py
View File

@ -53,19 +53,29 @@ class LandmarksDetectionResult:
def to_pb2(self) -> _LandmarksDetectionResultProto: def to_pb2(self) -> _LandmarksDetectionResultProto:
"""Generates a LandmarksDetectionResult protobuf object.""" """Generates a LandmarksDetectionResult protobuf object."""
landmarks = _NormalizedLandmarkListProto()
landmarks.landmark.extend([
landmark.to_pb2() for landmark in self.landmarks
])
classifications = _ClassificationListProto() classifications = _ClassificationListProto()
for category in self.categories: classifications.classification.extend([
classifications.classification.append(
_ClassificationProto( _ClassificationProto(
index=category.index, index=category.index,
score=category.score, score=category.score,
label=category.category_name, label=category.category_name,
display_name=category.display_name)) display_name=category.display_name) for category in self.categories
])
world_landmarks = _LandmarkListProto()
world_landmarks.landmark.extend([
world_landmark.to_pb2() for world_landmark in self.world_landmarks
])
return _LandmarksDetectionResultProto( return _LandmarksDetectionResultProto(
landmarks=_NormalizedLandmarkListProto(self.landmarks), landmarks=landmarks,
classifications=classifications, classifications=classifications,
world_landmarks=_LandmarkListProto(self.world_landmarks), world_landmarks=world_landmarks,
rect=self.rect.to_pb2()) rect=self.rect.to_pb2())
@classmethod @classmethod