Adds type hints

Author: robmarkcole

deepstack/core.py

@@ -3,54 +3,74 @@
 """
 import requests
 from PIL import Image
+from typing import Union, List, Set, Dict
 
 ## Const
 HTTP_OK = 200
 HTTP_BAD_REQUEST = 400
 HTTP_UNAUTHORIZED = 401
-TIMEOUT = 30  # seconds
+TIMEOUT = 20  # seconds
 
 
-def format_confidence(confidence):
+def format_confidence(confidence: Union[str, float]) -> float:
     """Takes a confidence from the API like 
        0.55623 and returne 55.6 (%).
     """
-    return round(float(confidence)*100, 1)
+    return round(float(confidence) * 100, 1)
 
 
-def get_matched_faces(predictions: dict):
+def get_confidences_above_threshold(
+    confidences: List[float], confidence_threshold: float
+) -> List[float]:
+    """Takes a list of confidences and returns those above a confidence_threshold."""
+    return [val for val in confidences if val >= confidence_threshold]
+
+
+def get_object_labels(predictions: List[Dict]) -> Set[str]:
+    """
+    Get a list of the unique object labels predicted.
+    """
+    labels = [pred["label"] for pred in predictions]
+    return set(labels)
+
+
+def get_label_confidences(predictions: List[Dict], target_label: str):
+    """
+    Return the list of confidences of instances of target label.
     """
-    Get the predicted faces and their confidence.
+    confidences = [
+        pred["confidence"] for pred in predictions if pred["label"] == target_label
+    ]
+    return confidences
+
+
+def get_objects_summary(predictions: List[Dict]):
     """
-    matched_faces = {}
-    matched_faces = {
-        face["userid"]: format_confidence(face["confidence"])
-        for face in predictions
-        if not face["userid"] == "unknown"
+    Get a summary of the objects detected.
+    """
+    labels = get_object_labels(predictions)
+    return {
+        label: len(get_label_confidences(predictions, target_label=label))
+        for label in labels
     }
-    return matched_faces
 
 
 def post_image(url: str, image: bytes):
-    """Post an image to the classifier."""
+    """Post an image to Deepstack."""
     response = requests.post(url, files={"image": image}, timeout=TIMEOUT)
     return response
 
 
-class DeepstackFace:
-    """Work with faces."""
+class DeepstackObject:
+    """The object detection API locates and classifies 80 
+    different kinds of objects in a single image.."""
 
     def __init__(self, ip_address: str, port: str):
 
-        self._url_check = "http://{}:{}/v1/vision/face/recognize".format(
+        self._url_object_detection = "http://{}:{}/v1/vision/detection".format(
             ip_address, port
         )
-
-        self._faces = None
-        self._matched = {}
-
-    def register_face(self, file_path: str, userid: str):
-        """Register a face with Deepstack."""
+        self._predictions = []
 
     def process_file(self, file_path: str):
         """Process an image file."""
@@ -59,22 +79,14 @@ def process_file(self, file_path: str):
 
     def process_image_bytes(self, image_bytes: bytes):
         """Process an image."""
-        response = post_image(self._url_check, image_bytes)
+        self._predictions = []
+
+        response = post_image(self._url_object_detection, image_bytes)
         if response:
             if response.status_code == HTTP_OK:
-                predictions_json = response.json()["predictions"]
-                self._faces = len(predictions_json)
-                self._matched = get_matched_faces(predictions_json)
-
-        else:
-            self._faces = None
-            self._matched = {}
+                self._predictions = response.json()["predictions"]
 
     @property
-    def attributes(self):
+    def predictions(self):
         """Return the classifier attributes."""
-        return {
-            "faces": self._faces,
-            "matched_faces": self._matched,
-            "total_matched_faces": len(self._matched),
-        }
+        return self._predictions