Merge pull request #20 from 5sControl/dev

Dev

connection/ModelPredictionsReceiver.py

@@ -0,0 +1,63 @@
+import requests
+import numpy as np
+from logging import Logger
+from PIL import Image
+import io
+
+
+PORT = 5000
+
+
+class ModelPredictionsReceiver:
+    def __init__(self, server_url: str, logger: Logger) -> None:
+        self._server_url = server_url
+        self._logger = logger
+
+    @staticmethod
+    def _convert_image2bytes(image: np.array, format='PNG') -> io.BytesIO:
+        pil_image = Image.fromarray(image)
+        img_byte_arr = io.BytesIO()
+        pil_image.save(img_byte_arr, format=format)
+        img_byte_arr.seek(0)
+        return img_byte_arr
+
+    def predict_human(self, img: np.array):
+        try:
+            response = requests.post(
+                f"{self._server_url}:{PORT}/predict_human",
+                files={
+                    "image": ("image", self._convert_image2bytes(img), "image/png")
+                }
+            )
+            response.raise_for_status()
+            return np.array(response.json().get("coordinates"))
+        except Exception as exc:
+            self._logger.critical("Cannot send request to model server. Error - {}".format(exc))
+
+    def predict_bottles(self, img: np.array):
+        try:
+            response = requests.post(
+                f"{self._server_url}:{PORT}/predict_bottles",
+                files={
+                    "image": ("image", self._convert_image2bytes(img), "image/png")
+                }
+            )
+            response.raise_for_status()
+            return np.array(response.json().get("coordinates"))
+        except Exception as exc:
+            self._logger.critical("Cannot send request to model server. Error - {}".format(exc))
+
+    def predict_boxes(self, img: np.array):
+        try:
+            response = requests.post(
+                f"{self._server_url}:{PORT}/predict_boxes",
+                files={
+                    "image": ("image", self._convert_image2bytes(img), "image/png")
+                }
+            )
+            response.raise_for_status()
+            return np.array(response.json().get("coordinates"))
+        except Exception as exc:
+            self._logger.critical("Cannot send request to model server. Error - {}".format(exc))
+
+

connection/__init__.py

@@ -0,0 +1,2 @@
+from .HTTPLIB2Capture import HTTPLIB2Capture
+from .ModelPredictionsReceiver import ModelPredictionsReceiver

main.py

@@ -1,4 +1,4 @@
-from min_max_utils.HTTPLIB2Capture import HTTPLIB2Capture
+from connection import HTTPLIB2Capture
 from min_max_utils.min_max_utils import create_logger
 import warnings
 import os

min_max_utils/MinMaxReporter.py

@@ -39,7 +39,7 @@ def add_empty_zone(zones: list):
         return zones
 
     def create_report(self, n_boxes_history: list, img: np.array, areas: list, boxes_coords: list, zones: list) -> dict:
-        red_lines = find_red_line(img)
+        red_lines = find_red_line(img) 
         report = []
         if not zones:
             for item in areas:
@@ -67,13 +67,12 @@ def create_report(self, n_boxes_history: list, img: np.array, areas: list, boxes
             for subarr_idx, coord in enumerate(item['coords']):
                 area_coords = convert_coords_from_dict_to_list(coord)
 
-                is_red_line_in_subarea = False
-                for idx, line in enumerate(red_lines):
-                    if is_line_in_area(area_coords, line):
-                        debug_user_image = draw_line(debug_user_image, line, area_coords, thickness=4)
-                        is_red_line_in_subarea = is_red_line_in_item = True
-
                 if multi_row:
+                    is_red_line_in_subarea = False
+                    for idx, line in enumerate(red_lines):
+                        if is_line_in_area(area_coords, line):
+                            debug_user_image = draw_line(debug_user_image, line, area_coords, thickness=4)
+                            is_red_line_in_subarea = is_red_line_in_item = True
                     text_item = f"{item_name}: {n_boxes_history[item_index][subarr_idx] if not is_red_line_in_subarea else 'low stock level'}"
                 else:
                     text_item = f"{item_name}: "

min_max_utils/min_max_utils.py

@@ -109,7 +109,7 @@ def check_box_in_area(box_coord, area_coord):
     return False
 
 
-def filter_boxes(main_item_coords, _, boxes_coords, area_coords=None, check=True):
+def filter_boxes(main_item_coords, boxes_coords, area_coords=None, check=True):
     result = []
     for box_coord in boxes_coords:
         box_coord = transfer_coords(box_coord, main_item_coords)

model_image/ObjectDetectionModel.py

@@ -1,16 +1,17 @@
 from ultralytics import YOLO
 import torch
+import numpy as np
 
 
-class ObjDetectionModel:
+class YOLOv8ObjDetectionModel:
     def __init__(self, path: str, conf_thresh: float, iou_thresh: float, classes: list) -> None:
         self.model = YOLO(path)
         self.conf_thresh = conf_thresh
         self.iou_thresh = iou_thresh
         self.classes = classes
 
     @torch.no_grad()
-    def __call__(self, img, classes: list = None) -> list:
+    def __call__(self, img: np.array, classes: list = None) -> list:
         results = self.model(
             source=img,
             conf=self.conf_thresh,
@@ -19,7 +20,6 @@ def __call__(self, img, classes: list = None) -> list:
             classes=self.classes if classes is None else classes,
             verbose=False
         )[0].boxes
-        n_boxes = len(results)
         coords_with_confs = torch.hstack((results.xyxy, results.conf.unsqueeze(-1)))
-        return [n_boxes, coords_with_confs]
+        return coords_with_confs
 

model_image/YOLORObjectDetectionModel.py

@@ -0,0 +1,34 @@
+import torch
+from yolor.model import get_model
+import numpy as np
+from yolor.utils.datasets import letterbox
+from yolor.utils.general import non_max_suppression, scale_coords
+
+
+class YOLORObjectDetectionModel:
+    def __init__(self, model_path: str, config_path: str, conf_thresh, iou_thresh, classes) -> None:
+        self.model, self.device = get_model(model_path, config_path)
+        self.conf_thresh = conf_thresh
+        self.iou_thresh = iou_thresh
+        self.classes = classes
+
+    def __preprocess_image__(self, img: np.array) -> np.array:
+        self.img_shape = img.shape
+        img = letterbox(img.copy(), new_shape=1280, auto_size=64)[0]
+        img = img[:, :, ::-1].transpose(2, 0, 1)
+        img = np.ascontiguousarray(img)
+        img = torch.from_numpy(img).to(self.device)
+        img = img.float()
+        img /= 255.0
+        img = img.unsqueeze(0)
+        return img
+
+    @torch.no_grad()
+    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=self.classes, agnostic=False)[0]
+        pred[:, :4] = scale_coords(
+            img.shape[2:], pred[:, :4], self.img_shape).round()
+        return pred[:, :5]

model_image/app.py

@@ -1,16 +1,18 @@
 from PIL import Image
 from flask import Flask, jsonify, request
 from flask_configs.load_configs import *
-from ObjectDetectionModel import ObjDetectionModel
+from ObjectDetectionModel import YOLOv8ObjDetectionModel
+from YOLORObjectDetectionModel import YOLORObjectDetectionModel
 import numpy as np
 import colorlog
 import logging
+import io
 
 
 app = Flask(__name__)
-human_model = ObjDetectionModel(HUMAN_MODEL_PATH, CONF_THRES, IOU_THRES, CLASSES)
-box_model = ObjDetectionModel(BOX_MODEL_PATH, CONF_THRES, IOU_THRES, CLASSES)
-bottle_model = ObjDetectionModel(BOTTLE_MODEL_PATH, CONF_THRES, IOU_THRES, CLASSES)
+human_model = YOLORObjectDetectionModel(HUMAN_MODEL_PATH, CONF_PATH, CONF_THRES, IOU_THRES, CLASSES)
+box_model = YOLOv8ObjDetectionModel(BOX_MODEL_PATH, CONF_THRES, IOU_THRES, CLASSES)
+bottle_model = YOLORObjectDetectionModel(HUMAN_MODEL_PATH, CONF_PATH, 0.25, IOU_THRES, [39])
 
 
 logger = logging.getLogger('min_max_logger')
@@ -28,40 +30,40 @@
 logger.setLevel(logging.DEBUG)
 logger.propagate = False
 
+convert_bytes2image = lambda bytes: np.array(Image.open(io.BytesIO(bytes)), dtype=np.uint8)
+
 @app.route('/predict_human', methods=['POST'])
 def predict_human():
     if request.method == 'POST':
-        image = np.array(request.json['image']).astype(np.float32)
-        n_boxes, coords = human_model(image)
+        image = convert_bytes2image(request.files["image"].read()).astype(np.float32)
+        coords = human_model(image)
+        logger.info(f"request to predict_human: {len(coords)}")
         return jsonify(
             {
-                "n_boxes": n_boxes,
                 "coordinates": coords.tolist()
             }
         )
 
 @app.route('/predict_boxes', methods=['POST'])
 def predict_boxes():
     if request.method == 'POST':
-        image = np.array(request.json['image']).astype(np.float32)
-        n_boxes, coords = box_model(image)
-        logger.info("Request to predict_boxes: " + str(n_boxes))
+        image = convert_bytes2image(request.files["image"].read()).astype(np.float32)
+        coords = box_model(image)
+        logger.info(f"request to predict_boxes: {len(coords)}")
         return jsonify(
             {
-                "n_items": n_boxes,
                 "coordinates": coords.tolist()
             }
         )
 
 @app.route('/predict_bottles', methods=['POST'])
 def predict_bottles():
     if request.method == 'POST':
-        image = np.array(request.json['image']).astype(np.float32)
-        n_bottles, coords = bottle_model(image)
-        logger.info("Request to predict_bottles: " + str(n_bottles))
+        image = convert_bytes2image(request.files["image"].read()).astype(np.float32)
+        coords = bottle_model(image)
+        logger.info(f"request to predict_bottles: {len(coords)}")
         return jsonify(
             {
-                "n_items": n_bottles,
                 "coordinates": coords.tolist()
             }
         )

model_image/flask_configs/flask_confs.json

@@ -4,9 +4,9 @@
     ],
     "iou_thres": 0.5,
     "conf_thres": 0.46,
-    "box_detect_model": "min_max_v0.3.10.pt",
-    "human_detect_model": "min_max_v1.0h.pt",
-    "bottle_detect_model": "min_max_v0.4.3b.pt",
+    "config_path": "weights/yolor_csp_x.cfg",
+    "box_detect_model": "weights/min_max_v0.3.10.pt",
+    "human_detect_model": "weights/min_max_v0.4.4b.pt",
     "img_size": 640,
     "port": 5000
 }

model_image/flask_configs/load_configs.py

@@ -7,7 +7,7 @@
     IOU_THRES = configs.get("iou_thres")
     BOX_MODEL_PATH = configs.get("box_detect_model")
     HUMAN_MODEL_PATH = configs.get("human_detect_model")
-    BOTTLE_MODEL_PATH = configs.get("bottle_detect_model")
+    CONF_PATH = configs.get("config_path")
     CLASSES = configs.get("classes")
     IMG_SIZE = configs.get("img_size")
     PORT = configs.get("port")

model_image/yolor/__init__.py

@@ -0,0 +1 @@
+from .model import get_model

model_image/yolor/model.py

@@ -0,0 +1,12 @@
+import torch
+from .utils.torch_utils import select_device
+from .models.models import Darknet
+
+
+def get_model(weights, cfg):
+    imgsz = 1280
+    device = select_device('cpu')
+    model = Darknet(cfg, imgsz)
+    model.load_state_dict(torch.load(weights, map_location=device)['model'])
+    model.to(device).eval()
+    return model, device

model_image/yolor/models/__init__.py

@@ -0,0 +1 @@
+