Refactor core

Author: robmarkcole

deepstack/core.py

@@ -19,13 +19,15 @@
 
 
 def format_confidence(confidence: Union[str, float]) -> float:
-    """Takes a confidence from the API like 
-       0.55623 and returns 55.6 (%).
+    """Takes a confidence from the API like
+    0.55623 and returns 55.6 (%).
     """
     return round(float(confidence) * 100, 1)
 
 
-def get_confidences_above_threshold(confidences: List[float], confidence_threshold: float) -> List[float]:
+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]
 
@@ -57,7 +59,9 @@ 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_object]
+    confidences = [
+        pred["confidence"] for pred in predictions if pred["label"] == target_object
+    ]
     return confidences
 
 
@@ -66,17 +70,26 @@ def get_objects_summary(predictions: List[Dict]):
     Get a summary of the objects detected.
     """
     objects = get_objects(predictions)
-    return {target_object: len(get_object_confidences(predictions, target_object)) for target_object in objects}
+    return {
+        target_object: len(get_object_confidences(predictions, target_object))
+        for target_object in objects
+    }
 
 
-def post_image(url: str, image_bytes: bytes, api_key: str, timeout: int, data: dict = {}):
+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_bytes}, data=data, timeout=timeout)
+        response = requests.post(
+            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")
+        raise DeepstackException(
+            f"Timeout connecting to Deepstack, current timeout is {timeout} seconds"
+        )
     except requests.exceptions.ConnectionError as exc:
         raise DeepstackException(f"Connection error: {exc}")
 
@@ -89,25 +102,32 @@ class Deepstack(object):
     """Base class for deepstack."""
 
     def __init__(
-        self, ip: str, port: str, api_key: str = "", timeout: int = DEFAULT_TIMEOUT, url: str = "",
+        self,
+        api_key: str = "",
+        timeout: int = DEFAULT_TIMEOUT,
+        url_detect: str = None,
+        url_recognise: str = None,
+        url_register: str = None,
     ):
 
-        self._ip = ip
-        self._port = port
-        self._url = url
+        self._url_detect = url_detect
+        self._url_recognise = url_recognise
+        self._url_register = url_register
         self._api_key = api_key
         self._timeout = timeout
         self._response = None
 
     def detect(self, image_bytes: bytes):
         """Process image_bytes, performing detection."""
         self._response = None
-        url = self._url.format(ip=self._ip, port=self._port)
-
-        response = post_image(url, image_bytes, self._api_key, self._timeout)
+        response = post_image(
+            self._url_detect, image_bytes, self._api_key, self._timeout
+        )
 
         if not response.status_code == HTTP_OK:
-            raise DeepstackException(f"Error from request, status code: {response.status_code}")
+            raise DeepstackException(
+                f"Error from request, status code: {response.status_code}"
+            )
             return
 
         self._response = response.json()
@@ -125,12 +145,27 @@ class DeepstackObject(Deepstack):
     """Work with objects"""
 
     def __init__(
-        self, ip: str, port: str, api_key: str = "", timeout: int = DEFAULT_TIMEOUT, custom_model: str = None,
+        self,
+        ip: str,
+        port: str,
+        api_key: str = "",
+        timeout: int = DEFAULT_TIMEOUT,
+        custom_model: str = None,
     ):
         if not custom_model:
-            super().__init__(ip, port, api_key, timeout, url=URL_OBJECT_DETECTION)
+            super().__init__(
+                api_key,
+                timeout,
+                url_detect=URL_OBJECT_DETECTION.format(ip=ip, port=port),
+            )
         else:
-            super().__init__(ip, port, api_key, timeout, url=URL_CUSTOM.format(custom_model=custom_model))
+            super().__init__(
+                api_key,
+                timeout,
+                url_detect=URL_CUSTOM.format(
+                    ip=ip, port=port, custom_model=custom_model
+                ),
+            )
 
     @property
     def predictions(self):
@@ -142,9 +177,17 @@ class DeepstackScene(Deepstack):
     """Work with scenes"""
 
     def __init__(
-        self, ip: str, port: str, api_key: str = "", timeout: int = DEFAULT_TIMEOUT,
+        self,
+        ip: str,
+        port: str,
+        api_key: str = "",
+        timeout: int = DEFAULT_TIMEOUT,
     ):
-        super().__init__(ip, port, api_key, timeout, url=URL_SCENE_DETECTION)
+        super().__init__(
+            api_key,
+            timeout,
+            url_detect=URL_SCENE_DETECTION.format(ip=self._ip, port=self._port),
+        )
 
     @property
     def predictions(self):
@@ -156,9 +199,19 @@ class DeepstackFace(Deepstack):
     """Work with objects"""
 
     def __init__(
-        self, ip: str, port: str, api_key: str = "", timeout: int = DEFAULT_TIMEOUT,
+        self,
+        ip: str,
+        port: str,
+        api_key: str = "",
+        timeout: int = DEFAULT_TIMEOUT,
     ):
-        super().__init__(ip, port, api_key, timeout, url=URL_FACE_DETECTION)
+        super().__init__(
+            api_key,
+            timeout,
+            url_detect=URL_FACE_DETECTION.format(ip=self._ip, port=self._port),
+            url_register=URL_FACE_REGISTRATION.format(ip=self._ip, port=self._port),
+            url_recognise=URL_FACE_RECOGNITION.format(ip=self._ip, port=self._port),
+        )
 
     @property
     def predictions(self):
@@ -171,7 +224,7 @@ def register_face(self, name: str, image_bytes: bytes):
         """
 
         response = post_image(
-            url=URL_FACE_REGISTRATION.format(ip=self._ip, port=self._port),
+            url=self._url_register,
             image_bytes=image_bytes,
             api_key=self._api_key,
             timeout=self._timeout,
@@ -180,18 +233,22 @@ def register_face(self, name: str, image_bytes: bytes):
 
         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."""
-        url = URL_FACE_RECOGNITION.format(ip=self._ip, port=self._port)
 
-        response = post_image(url, image_bytes, self._api_key, self._timeout)
+        response = post_image(
+            self._url_recognise, image_bytes, self._api_key, self._timeout
+        )
 
         if not response.status_code == HTTP_OK:
-            raise DeepstackException(f"Error from request, status code: {response.status_code}")
+            raise DeepstackException(
+                f"Error from request, status code: {response.status_code}"
+            )
             return
 
         self._response = response.json()