robmarkcole · robmarkcole · Sep 14, 2019 · Sep 14, 2019 · Sep 14, 2019 · Sep 14, 2019
diff --git a/README.md b/README.md
@@ -4,14 +4,16 @@
 
 # deepstack-python
 Unofficial python API for [DeepStack](https://python.deepstack.cc/). Provides class for making requests to the object detection endpoint, and functions for processing the result. See the Jupyter notebooks for usage.
-
-## Services
-Face and object detection endpoints return bounding boxes of faces and objects respectively. 
-
-TODO: add face registration and recognition.
+Run deepstack (CPU, noAVX mode):
+```
+docker run -e VISION-DETECTION=True -e VISION-FACE=True -e MODE=High -d \
+      -v localstorage:/datastore -p 5000:5000 \
+      -e API-KEY="Mysecretkey" \
+       --name deepstack deepquestai/deepstack:noavx
+```
 
 ## Development
 * Use `venv` -> `source venv/bin/activate`
 * `pip install -r requirements-dev.txt`
 * Run tests with `venv/bin/pytest tests/*`
-* Black format with `venv/bin/black deepstack/core.py`
+* Black format with `venv/bin/black deepstack/core.py` and `venv/bin/black tests/test_deepstack.py`
diff --git a/deepstack/core.py b/deepstack/core.py
@@ -12,6 +12,8 @@
 ## API urls
 URL_OBJECT_DETECTION = "http://{}:{}/v1/vision/detection"
 URL_FACE_DETECTION = "http://{}:{}/v1/vision/face"
+URL_FACE_REGISTRATION = "http://{}:{}/v1/vision/face/register"
+URL_FACE_RECOGNITION = "http://{}:{}/v1/vision/face/recognize"
 
 
 def format_confidence(confidence: Union[str, float]) -> float:
@@ -28,20 +30,35 @@ def get_confidences_above_threshold(
     return [val for val in confidences if val >= confidence_threshold]
 
 
-def get_object_labels(predictions: List[Dict]) -> List[str]:
+def get_recognised_faces(predictions: List[Dict]) -> List[Dict]:
     """
-    Get a list of the unique object labels predicted.
+    Get the recognised faces.
+    """
+    try:
+        matched_faces = {
+            face["userid"]: round(face["confidence"] * 100, 1)
+            for face in predictions
+            if not face["userid"] == "unknown"
+        }
+        return matched_faces
+    except:
+        return {}
+
+
+def get_objects(predictions: List[Dict]) -> List[str]:
+    """
+    Get a list of the unique objects predicted.
     """
     labels = [pred["label"] for pred in predictions]
     return list(set(labels))
 
 
-def get_label_confidences(predictions: List[Dict], target_label: str):
+def get_object_confidences(predictions: List[Dict], target_object: str):
     """
     Return the list of confidences of instances of target label.
     """
     confidences = [
-        pred["confidence"] for pred in predictions if pred["label"] == target_label
+        pred["confidence"] for pred in predictions if pred["label"] == target_object
     ]
     return confidences
 
@@ -50,24 +67,29 @@ def get_objects_summary(predictions: List[Dict]):
     """
     Get a summary of the objects detected.
     """
- labels = get_object_labels(predictions)
+ objects = get_objects(predictions)
     return {
- label: len(get_label_confidences(predictions, target_label=label))
-        for label in labels
+ target_object: len(get_object_confidences(predictions, target_object))
+        for target_object in objects
     }
 
 
-def post_image(url: str, image: bytes, api_key: str, timeout: int):
+def post_image(
+    url: str, image_bytes: bytes, api_key: str, timeout: int, data: dict = {}
+):
     """Post an image to Deepstack."""
     try:
+        data["api_key"] = api_key
         response = requests.post(
-            url, files={"image": image}, data={"api_key": api_key}, timeout=timeout
+            url, files={"image": image_bytes}, data=data, timeout=timeout
         )
         return response
     except requests.exceptions.Timeout:
         raise DeepstackException(
             f"Timeout connecting to Deepstack, current timeout is {timeout} seconds"
         )
+    except requests.exceptions.ConnectionError as exc:
+        raise DeepstackException(f"Connection error: {exc}")
 
 
 class DeepstackException(Exception):
@@ -93,13 +115,8 @@ def __init__(
         self._timeout = timeout
         self._predictions = []
 
-    def process_file(self, file_path: str):
-        """Process an image file."""
-        with open(file_path, "rb") as image_bytes:
-            self.process_image_bytes(image_bytes)
-
-    def process_image_bytes(self, image_bytes: bytes):
-        """Process an image."""
+    def detect(self, image_bytes: bytes):
+        """Process image_bytes, performing detection."""
         self._predictions = []
         url = self._url_detection.format(self._ip_address, self._port)
 
@@ -132,6 +149,7 @@ def __init__(
             ip_address, port, api_key, timeout, url_detection=URL_OBJECT_DETECTION
         )
 
+
 class DeepstackFace(Deepstack):
     """Work with objects"""
 
@@ -145,3 +163,36 @@ def __init__(
         super().__init__(
             ip_address, port, api_key, timeout, url_detection=URL_FACE_DETECTION
         )
+
+    def register_face(self, name: str, image_bytes: bytes):
+        """
+        Register a face name to a file.
+        """
+
+        response = post_image(
+            url=URL_FACE_REGISTRATION.format(self._ip_address, self._port),
+            image_bytes=image_bytes,
+            api_key=self._api_key,
+            timeout=self._timeout,
+            data={"userid": name},
+        )
+
+        if response.status_code == 200 and response.json()["success"] == True:
+            return
+        elif response.status_code == 200 and response.json()["success"] == False:
+            error = response.json()["error"]
+            raise DeepstackException(f"Error from Deepstack: {error}")
+
+    def recognise(self, image_bytes: bytes):
+        """Process image_bytes, performing recognition."""
+        self._predictions = []
+        url = URL_FACE_RECOGNITION.format(self._ip_address, self._port)
+
+        response = post_image(url, image_bytes, self._api_key, self._timeout)
+
+        if response.status_code == HTTP_OK:
+            if response.json()["success"]:
+                self._predictions = response.json()["predictions"]
+            else:
+                error = response.json()["error"]
+                raise DeepstackException(f"Error from Deepstack: {error}")
diff --git a/setup.py b/setup.py
@@ -1,6 +1,6 @@
 from setuptools import setup, find_packages
 
-VERSION = "0.3"
+VERSION = "0.4"
 
 REQUIRES = ["requests"]
 

diff --git a/tests/test_deepstack.py b/tests/test_deepstack.py
@@ -1,5 +1,3 @@
-# Placeholder
-
 import deepstack.core as ds
 import requests
 import requests_mock
@@ -13,7 +11,7 @@
 MOCK_API_KEY = "mock_api_key"
 MOCK_TIMEOUT = 8
 
-MOCK_RESPONSE = {
+MOCK_OBJECT_DETECTION_RESPONSE = {
     "success": True,
     "predictions": [
         {
@@ -43,53 +41,87 @@
     ],
 }
 
-MOCK_PREDICTIONS = MOCK_RESPONSE["predictions"]
-MOCK_CONFIDENCES = [0.6998661, 0.7996547]
+MOCK_FACE_RECOGNITION_RESPONSE = {
+    "success": True,
+    "predictions": [
+        {
+            "confidence": 0.74999994,
+            "userid": "Idris Elba",
+            "y_min": 176,
+            "x_min": 209,
+            "y_max": 825,
+            "x_max": 677,
+        },
+        {
+            "confidence": 0,
+            "userid": "unknown",
+            "y_min": 230,
+            "x_min": 867,
+            "y_max": 729,
+            "x_max": 1199,
+        },
+    ],
+}
+
+MOCK_OBJECT_PREDICTIONS = MOCK_OBJECT_DETECTION_RESPONSE["predictions"]
+MOCK_OBJECT_CONFIDENCES = [0.6998661, 0.7996547]
 CONFIDENCE_THRESHOLD = 0.7
+MOCK_RECOGNISED_FACES = {"Idris Elba": 75.0}
 
 
-def test_DeepstackObject_process_image_bytes():
+def test_DeepstackObject_detect():
     """Test a good response from server."""
     with requests_mock.Mocker() as mock_req:
-        mock_req.post(MOCK_URL, status_code=ds.HTTP_OK, json=MOCK_RESPONSE)
+        mock_req.post(
+            MOCK_URL, status_code=ds.HTTP_OK, json=MOCK_OBJECT_DETECTION_RESPONSE
+        )
 
         dsobject = ds.DeepstackObject(MOCK_IP_ADDRESS, MOCK_PORT)
-        dsobject.process_image_bytes(MOCK_BYTES)
-        assert dsobject.predictions == MOCK_PREDICTIONS
+        dsobject.detect(MOCK_BYTES)
+        assert dsobject.predictions == MOCK_OBJECT_PREDICTIONS
 
 
-def test_DeepstackObject_process_image_bytes_timeout():
+def test_DeepstackObject_detect_timeout():
     """Test a timeout. THIS SHOULD FAIL"""
     with pytest.raises(ds.DeepstackException) as excinfo:
         with requests_mock.Mocker() as mock_req:
             mock_req.post(MOCK_URL, exc=requests.exceptions.ConnectTimeout)
             dsobject = ds.DeepstackObject(MOCK_IP_ADDRESS, MOCK_PORT)
-            dsobject.process_image_bytes(MOCK_BYTES)
+            dsobject.detect(MOCK_BYTES)
             assert False
             assert "SHOULD FAIL" in str(excinfo.value)
 
 
-def test_get_object_labels():
+def test_get_objects():
     """Cant always be sure order of returned list items."""
- object_labels = ds.get_object_labels(MOCK_PREDICTIONS)
-    assert type(object_labels) is list
-    assert "dog" in object_labels
-    assert "person" in object_labels
-    assert len(object_labels) == 2
+ objects = ds.get_objects(MOCK_OBJECT_PREDICTIONS)
+    assert type(objects) is list
+    assert "dog" in objects
+    assert "person" in objects
+    assert len(objects) == 2
 
 
 def test_get_objects_summary():
-    objects_summary = ds.get_objects_summary(MOCK_PREDICTIONS)
+    objects_summary = ds.get_objects_summary(MOCK_OBJECT_PREDICTIONS)
     assert objects_summary == {"dog": 1, "person": 2}
 
 
-def test_get_label_confidences():
- label_confidences = ds.get_label_confidences(MOCK_PREDICTIONS, "person")
-    assert label_confidences == MOCK_CONFIDENCES
+def test_get_object_confidences():
+ object_confidences = ds.get_object_confidences(MOCK_OBJECT_PREDICTIONS, "person")
+    assert object_confidences == MOCK_OBJECT_CONFIDENCES
 
 
 def test_get_confidences_above_threshold():
     assert (
-        len(ds.get_confidences_above_threshold(MOCK_CONFIDENCES, CONFIDENCE_THRESHOLD))
+        len(
+            ds.get_confidences_above_threshold(
+                MOCK_OBJECT_CONFIDENCES, CONFIDENCE_THRESHOLD
+            )
+        )
         == 1
     )
+
+
+def test_get_recognised_faces():
+    predictions = MOCK_FACE_RECOGNITION_RESPONSE["predictions"]
+    assert ds.get_recognised_faces(predictions) == MOCK_RECOGNISED_FACES