Merge pull request #13 from 5sControl/dev

Dev

Author: eugenos-programos

IdleAlgorithm.py

@@ -0,0 +1,45 @@
+from logging import Logger
+import asyncio
+from connection import ImageHTTPExtractor, ModelPredictionsReceiver, IdleReporter
+from confs import configs
+import utils
+import numpy as np
+
+
+class IdleAlgorithm:
+    def __init__(
+            self, 
+            logger: Logger, 
+            image_extractor: ImageHTTPExtractor, 
+            model_predictor: ModelPredictionsReceiver,
+            reporter: IdleReporter
+        ) -> None:
+        self._logger = logger
+        self._image_extractor = image_extractor
+        self._model_predictor = model_predictor
+        self._reporter = reporter
+
+    async def run(self):
+        iter_idx = 0
+        while True:
+            iter_idx += 1
+            if iter_idx % 60 == 0:
+                self._logger.info("60 iterations passed")
+            img, start_tracking = self._image_extractor.get_snapshot()
+            if img is None:
+                asyncio.sleep(1)
+                continue
+            preds = self._model_predictor.predict(img)
+            if preds is None:
+                asyncio.sleep(1)
+                continue
+            if preds.size != 0 and not np.any(preds == 1.):
+                self._logger.info("Telephone is detected")
+                if utils.bboxes_not_equal(prev_preds, preds, configs["threshold"]):
+                    utils.save_cropped_bbox(img, preds[:, :4])
+                    img = utils.put_rectangle(img, preds[:, :4], preds[:, 4])
+                    self._reporter.send_report(self._reporter.create_report(img, str(start_tracking)))
+                else:
+                    self._logger.debug("Equal bboxes")
+                prev_preds = preds
+            asyncio.sleep(2)

connection/ImageHTTPExtractor.py

@@ -4,26 +4,42 @@
 import numpy as np
 import cv2
 from typing import Tuple, Union
+import socketio
+
+
+sio = socketio.AsyncClient()
+images = {}
 
 
 class ImageHTTPExtractor:
-    def __init__(self, server_url: str, logger: logging.Logger, **credentials) -> None:
+    def __init__(self, camera_ip: str, logger: logging.Logger, **credentials) -> None:
         self.http_connection = httplib2.Http(".cache")
         self.http_connection.add_credentials(credentials.get("username"), credentials.get("password"))
         self.logger = logger
-        self.server_url = server_url
+        self.camera_ip = camera_ip
 
     def get_snapshot(self) -> Tuple[Union[cv2.Mat, None], Union[datetime.time, None]]:
         try:
+            global images
+            image = images[self.camera_ip]            
             curr_time = datetime.datetime.now()
-            response, content = self.http_connection.request(
-                self.server_url,
-                "GET", 
-                body="foobar"
-            )
-            nparr = np.frombuffer(content, np.uint8)
+            nparr = np.frombuffer(image, np.uint8)
             img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
             return img, curr_time
         except Exception as exc:
             self.logger.error(f"Cannot retrieve image. Following error raised - {exc}")
             return None, None
+
+@sio.event
+async def connect():
+    print("Connection")
+
+@sio.event
+async def snapshot_updated(data):
+    camera_url, screen = data.get("camera_ip"), data.get("screenshot")
+    global images
+    images[camera_url] = screen
+
+async def run_sio(url):
+    await sio.connect(url)
+    await sio.wait()

connection/__init__.py

@@ -1,3 +1,3 @@
 from .IdleReporter import IdleReporter
-from .ImageHTTPExtractor import ImageHTTPExtractor
+from .ImageHTTPExtractor import ImageHTTPExtractor, run_sio
 from .ModelPredictionsReceiver import ModelPredictionsReceiver

idle_model/IdleObjectDetectionModel.py

@@ -6,7 +6,7 @@
 
 
 class IdleObjectDetectionModel:
-    def __init__(self, model_path: str, config_path: str, conf_thresh, iou_thresh, classes) -> None:
+    def __init__(self, model_path: str, config_path: str, conf_thresh: float, iou_thresh: float, classes: list) -> None:
         self.model, self.device = get_model(model_path, config_path)
         self.conf_thresh = conf_thresh
         self.iou_thresh = iou_thresh
@@ -27,8 +27,6 @@ def __preprocess_image__(self, img: np.array) -> np.array:
     def __call__(self, img: np.array) -> list:
         img = self.__preprocess_image__(img)
         pred = self.model(img, augment=False)[0]
-        pred = non_max_suppression(
-            pred, 0.45, 0.5, classes=[67], agnostic=False)[0]
-        pred[:, :4] = scale_coords(
-            img.shape[2:], pred[:, :4], self.img_shape).round()
+        pred = non_max_suppression(pred, self.conf_thresh, self.iou_thresh, classes=[67], agnostic=False)[0]
+        pred[:, :4] = scale_coords(img.shape[2:], pred[:, :4], self.img_shape).round()
         return pred[:, :5]

main.py

@@ -1,10 +1,11 @@
-import time
 import utils
 import connection
 from confs import configs
 from dotenv import load_dotenv
+from IdleAlgorithm import IdleAlgorithm
 import os
 import numpy as np
+import asyncio
 
 
 password = os.environ.get("password")
@@ -13,37 +14,24 @@
 password = os.environ.get("password")
 username = os.environ.get("username")
 server_url = os.environ.get("server_url")
-camera_url = os.environ.get("camera_url")
+camera_ip = os.environ.get("camera_ip")
 folder = os.environ.get("folder")
 
 logger = utils.create_logger()
 
 prev_preds = np.array([[]]).astype(np.float32)
 reporter = connection.IdleReporter(folder, server_url, configs["wait_time"], logger)
-image_extractor = connection.ImageHTTPExtractor(camera_url, logger, username=username, password=password)
+image_extractor = connection.ImageHTTPExtractor(camera_ip, logger, username=username, password=password)
 model_predictor = connection.ModelPredictionsReceiver(server_url, logger)
+algo = IdleAlgorithm(logger, image_extractor, model_predictor, reporter)
 
-iter_idx = 0
+async def main():
+    await asyncio.gather(connection.run_sio(server_url + ':3456'), algo.run())
 
-while True:
-    iter_idx += 1
-    if iter_idx % 60 == 0:
-        logger.info("60 iterations passed")
-    img, start_tracking = image_extractor.get_snapshot()
-    if img is None:
-        time.sleep(1)
-        continue
-    preds = model_predictor.predict(img)
-    if preds is None:
-        time.sleep(1)
-        continue
-    if preds.size != 0 and not np.any(preds == 1.):
-        logger.info("Telephone is detected")
-        if utils.bboxes_not_equal(prev_preds, preds, configs["threshold"]):
-            utils.save_cropped_bbox(img, preds[:, :4])
-            img = utils.put_rectangle(img, preds[:, :4], preds[:, 4])
-            reporter.send_report(reporter.create_report(img, str(start_tracking)))
-        else:
-            logger.debug("Equal bboxes")
-        prev_preds = preds
-    time.sleep(2)
+loop = asyncio.get_event_loop()
+try:
+    loop.run_until_complete(main())
+except Exception as exc:
+    print(exc)
+finally:
+    loop.close()

utils/functional.py

@@ -21,7 +21,10 @@ def bboxes_not_equal(coords_1: np.array, coords_2: np.array, threshold: float) -
 
 def save_cropped_bbox(img: np.array, bboxes: np.array):
     os.makedirs("images/debug/", exist_ok=True)
-    for bbox in bboxes:
-        x1, y1, x2, y2 = bbox.astype(int)
-        c = 15
-        cv2.imwrite(f"images/debug/{uuid.uuid4()}.png", img[x1 - c:x2 + c, y1 - c:y2 + c])
+    try:
+        for bbox in bboxes:
+            x1, y1, x2, y2 = bbox.astype(int)
+            c = 15
+            cv2.imwrite(f"images/debug/{uuid.uuid4()}.png", img[x1 - c:x2 + c, y1 - c:y2 + c])
+    except Exception:
+        pass