Skip to content

Commit da800f9

Browse files
feat: add face annotations support in OCRResult (#333)
* feat: add face annotations support in OCRResult * fix: use new enum Likelihood instead of SafeSearchLikelihood --------- Co-authored-by: Areeb Ahmed <[email protected]>
1 parent e383de1 commit da800f9

File tree

1 file changed

+44
-0
lines changed

1 file changed

+44
-0
lines changed

openfoodfacts/ocr.py

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,7 @@ class OCRResult:
109109
"logo_annotations",
110110
"safe_search_annotation",
111111
"label_annotations",
112+
"face_annotations",
112113
)
113114

114115
def __init__(self, data: JSONType):
@@ -117,6 +118,7 @@ def __init__(self, data: JSONType):
117118
self.logo_annotations: List[LogoAnnotation] = []
118119
self.label_annotations: List[LabelAnnotation] = []
119120
self.safe_search_annotation: Optional[SafeSearchAnnotation] = None
121+
self.face_annotations: List[FaceAnnotation] = []
120122

121123
for text_annotation_data in data.get("textAnnotations", []):
122124
text_annotation = OCRTextAnnotation(text_annotation_data)
@@ -145,6 +147,10 @@ def __init__(self, data: JSONType):
145147
data["safeSearchAnnotation"]
146148
)
147149

150+
for face_annotation_data in data.get("faceAnnotations", []):
151+
face_annotation = FaceAnnotation(face_annotation_data)
152+
self.face_annotations.append(face_annotation)
153+
148154
def get_full_text(self) -> str:
149155
return (
150156
self.full_text_annotation.text
@@ -186,6 +192,9 @@ def get_logo_annotations(self) -> List["LogoAnnotation"]:
186192
def get_label_annotations(self) -> List["LabelAnnotation"]:
187193
return self.label_annotations
188194

195+
def get_face_annotations(self) -> List["FaceAnnotation"]:
196+
return self.face_annotations
197+
189198
def get_safe_search_annotation(self):
190199
return self.safe_search_annotation
191200

@@ -1130,6 +1139,31 @@ def __init__(self, data: JSONType):
11301139
self.description = data["description"]
11311140

11321141

1142+
class FaceAnnotation:
1143+
__slots__ = (
1144+
"detection_confidence",
1145+
"joy_likelihood",
1146+
"sorrow_likelihood",
1147+
"anger_likelihood",
1148+
"surprise_likelihood",
1149+
"under_exposed_likelihood",
1150+
"blurred_likelihood",
1151+
"headwear_likelihood",
1152+
)
1153+
1154+
def __init__(self, data: JSONType):
1155+
self.detection_confidence = data.get("detectionConfidence", 0.0)
1156+
self.joy_likelihood = Likelihood[data.get("joyLikelihood", "UNKNOWN")]
1157+
self.sorrow_likelihood = Likelihood[data.get("sorrowLikelihood", "UNKNOWN")]
1158+
self.anger_likelihood = Likelihood[data.get("angerLikelihood", "UNKNOWN")]
1159+
self.surprise_likelihood = Likelihood[data.get("surpriseLikelihood", "UNKNOWN")]
1160+
self.under_exposed_likelihood = Likelihood[
1161+
data.get("underExposedLikelihood", "UNKNOWN")
1162+
]
1163+
self.blurred_likelihood = Likelihood[data.get("blurredLikelihood", "UNKNOWN")]
1164+
self.headwear_likelihood = Likelihood[data.get("headwearLikelihood", "UNKNOWN")]
1165+
1166+
11331167
class SafeSearchAnnotation:
11341168
__slots__ = ("adult", "spoof", "medical", "violence", "racy")
11351169

@@ -1151,6 +1185,16 @@ def __init__(self, data: JSONType):
11511185
]
11521186

11531187

1188+
class Likelihood(enum.IntEnum):
1189+
UNKNOWN = 1
1190+
VERY_UNLIKELY = 2
1191+
UNLIKELY = 3
1192+
POSSIBLE = 4
1193+
LIKELY = 5
1194+
VERY_LIKELY = 6
1195+
1196+
1197+
# This class should be deleted and replaced by `Likelihood` in the future.
11541198
class SafeSearchAnnotationLikelihood(enum.IntEnum):
11551199
UNKNOWN = 1
11561200
VERY_UNLIKELY = 2

0 commit comments

Comments
 (0)