WongKinYiu · Dec 14, 2022 · Jun 15, 2023 · Jun 15, 2023
Showing with 236 additions and 9 deletions.

+13 −0 Dockerfile

+5 −0 build-docker.sh

+16 −8 detect.py

+15 −0 kserve-detect.py

+6 −0 kserve_wrapper/modelpath.py

+154 −0 kserve_wrapper/predictor.py

+26 −0 model_singleton.py

+1 −1 requirements.txt
diff --git a/Dockerfile b/Dockerfile
@@ -0,0 +1,13 @@
+FROM nvcr.io/nvidia/pytorch:21.08-py3
+
+RUN apt-get update && \
+    apt-get install -y zip htop screen libgl1-mesa-glx
+
+COPY requirements.txt requirements.txt
+
+RUN pip install --no-cache-dir --upgrade pip && pip install --no-cache-dir kserve==v0.10.1
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY . .
+
+ENTRYPOINT ["python"]
diff --git a/build-docker.sh b/build-docker.sh
@@ -0,0 +1,5 @@
+#!/usr/bin/env bash
+
+TAG=v0.0.3
+docker build -t footprintai/yolov7-objectdetector:${TAG} -f Dockerfile .
+docker push footprintai/yolov7-objectdetector:${TAG}
diff --git a/detect.py b/detect.py
@@ -14,8 +14,9 @@
 from utils.plots import plot_one_box
 from utils.torch_utils import select_device, load_classifier, time_synchronized, TracedModel
 
+from model_singleton import ModelSingleton
 
-def detect(save_img=False):
+def detect(opt: argparse.Namespace, save_img=False):
     source, weights, view_img, save_txt, imgsz, trace = opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size, not opt.no_trace
     save_img = not opt.nosave and not source.endswith('.txt')  # save inference images
     webcam = source.isnumeric() or source.endswith('.txt') or source.lower().startswith(
@@ -27,12 +28,17 @@ def detect(save_img=False):
 
     # Initialize
     set_logging()
+
+    ms = ModelSingleton(weights, opt.device)
+    model = ms.get_model()
+    stride = ms.get_stride()
+
     device = select_device(opt.device)
     half = device.type != 'cpu'  # half precision only supported on CUDA
 
     # Load model
-    model = attempt_load(weights, map_location=device)  # load FP32 model
-    stride = int(model.stride.max())  # model stride
+    #model = attempt_load(weights, map_location=device)  # load FP32 model
+    #stride = int(model.stride.max())  # model stride
     imgsz = check_img_size(imgsz, s=stride)  # check img_size
 
     if trace:
@@ -162,8 +168,7 @@ def detect(save_img=False):
 
     print(f'Done. ({time.time() - t0:.3f}s)')
 
-
-if __name__ == '__main__':
+def make_parser():
     parser = argparse.ArgumentParser()
     parser.add_argument('--weights', nargs='+', type=str, default='yolov7.pt', help='model.pt path(s)')
     parser.add_argument('--source', type=str, default='inference/images', help='source')  # file/folder, 0 for webcam
@@ -183,14 +188,17 @@ def detect(save_img=False):
     parser.add_argument('--name', default='exp', help='save results to project/name')
     parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
     parser.add_argument('--no-trace', action='store_true', help='don`t trace model')
-    opt = parser.parse_args()
+    return parser
+
+if __name__ == '__main__':
+    opt = make_parser().parse_args()
     print(opt)
     #check_requirements(exclude=('pycocotools', 'thop'))
 
     with torch.no_grad():
         if opt.update:  # update all models (to fix SourceChangeWarning)
             for opt.weights in ['yolov7.pt']:
-                detect()
+                detect(otp)
                 strip_optimizer(opt.weights)
         else:
-            detect()
+            detect(opt)
diff --git a/kserve-detect.py b/kserve-detect.py
@@ -0,0 +1,15 @@
+import argparse
+import os
+import kserve
+
+from kserve_wrapper.predictor import YoloV7Model
+
+DEFAULT_MODEL_NAME = os.getenv('DEFAULT_MODEL_NAME')
+parser = argparse.ArgumentParser(parents=[kserve.model_server.parser])
+parser.add_argument('--model_name', default=DEFAULT_MODEL_NAME)
+args, _ = parser.parse_known_args()
+
+if __name__ == "__main__":
+    model = YoloV7Model(args.model_name)
+    model.load()
+    kserve.ModelServer().start([model])
diff --git a/kserve_wrapper/modelpath.py b/kserve_wrapper/modelpath.py
@@ -0,0 +1,6 @@
+import os
+
+MODEL_PATH="/mnt/models"
+
+def modelpath_join(suffix_path):
+    return os.path.join(MODEL_PATH, suffix_path)
diff --git a/kserve_wrapper/predictor.py b/kserve_wrapper/predictor.py
@@ -0,0 +1,154 @@
+import base64
+import io
+import kserve
+import os
+
+from detect import make_parser, detect
+from kserve_wrapper.modelpath import modelpath_join
+from model_singleton import ModelSingleton
+from PIL import Image
+from typing import Dict
+
+class YoloV7Model(kserve.Model):
+    """ YoloV7 implements paper - YOLOv7: Trainable bag-of-freebies sets new state-of-the-art for real-time object detectors
+    see the orignal link: https://github.com/WongKinYiu/yolov7
+    """
+    def __init__(self, name: str):
+        super().__init__(name)
+        self.name = name
+        self.cfg = None
+        self.device = self._detect_device()
+        self.ms = None
+
+    def _detect_device(self) -> str:
+        if "YOLOV7_DEVICE" not in os.environ:
+            return "" # empty string for auto detect
+        return os.environ["YOLOV7_DEVICE"]
+
+    def load(self):
+        weights = modelpath_join("yolov7.pt")
+        self.ms = ModelSingleton(weights, self.device)
+
+        self.ready = True
+        return self.ready
+
+    def _query_ts(self) -> int:
+        import datetime
+
+        now = datetime.datetime.utcnow()
+        return int(now.timestamp()*1000.0)
+
+    def _ensure_folders_exists(self, strs: []):
+        for astr in strs:
+            basedir = os.path.dirname(astr)
+            os.makedirs(basedir, exist_ok=True)
+
+    def _load_label_txt(self, path: str) -> []:
+        if not os.path.exists(path):
+            return None
+
+        import csv
+
+        labels_list = []
+        with open(path, 'r') as file:
+            csvreader = csv.reader(file, delimiter=' ')
+            for row in csvreader:
+                if not ''.join(row).strip():
+                    continue
+                labels_list.append(self._make_label_dict(row))
+        return labels_list
+
+    def _make_label_dict(self, csvrow):
+        # all string
+        return {
+            "label_index": csvrow[0],
+            "x": csvrow[1],
+            "y": csvrow[2],
+            "w": csvrow[3],
+            "h": csvrow[4],
+        }
+
+    def predict(self, request: Dict, headers: Dict[str, str] = None) -> Dict:
+        inputs = request["instances"]
+        # request is wrapped the following format
+        # {
+        #   "instances": [
+        #     {
+        #       "image_bytes": {
+        #           "b64": "<b64-encoded>",
+        #       },
+        #       "key": "somekeys",
+        #     },
+        #   ],
+        # }
+        # and response is wrapped into the following
+        # {
+        #  "predictions: [
+        #    {
+        #      "image_bytes": {
+        #          "b64": "<b64-encoded>",
+        #      },
+        #      "label_and_positions: [
+        #        { "label_index": '17', "x": '0.534282', "y": '0.519531', "w":
+        #        '0.111255', "h": '0.21875'},
+        #      ],
+        #      "key": "somekeys",
+        #      "type": "yolov7-object-detector",
+        #    },
+        #  ]
+        # }
+
+        ts = self._query_ts()
+        ts_str = "{:010d}".format(ts)
+        query_file_path = "/tmp/query/{}/input.jpg".format(ts_str)
+        detect_label_txt_path = "/tmp/kserve/{}/labels/input.txt".format(ts_str)
+        detect_photo_path = "/tmp/kserve/{}/input.jpg".format(ts_str)
+
+        self._ensure_folders_exists([query_file_path, detect_label_txt_path,
+            detect_photo_path])
+
+        data = inputs[0]["image_bytes"]["b64"]
+        key = inputs[0]["key"]
+        raw_img_data = base64.b64decode(data)
+        input_image = Image.open(io.BytesIO(raw_img_data))
+        input_image.save(query_file_path)
+
+
+        # generate online argument for this simple query
+        runtime_argparse = make_parser().parse_args([
+            "--weights",
+            "/mnt/models/yolov7.pt",
+            "--conf",
+            "0.25",
+            "--img-size",
+            "640",
+            "--source",
+            query_file_path,
+            "--project",
+            "/tmp/kserve",
+            "--name",
+            ts_str,
+            "--save-txt",
+            "--exist-ok",
+         ])
+        detect(runtime_argparse, True)
+
+        # convert image to b64
+        grid_result_img = Image.open(detect_photo_path)
+        rgb_im = grid_result_img.convert('RGB')
+
+        buffered = io.BytesIO()
+        rgb_im.save(buffered, format="JPEG")
+        img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
+        return {
+                "predictions": [
+                {
+                    "image_bytes": {
+                        "b64": img_str,
+                    },
+                    "label_and_position": self._load_label_txt(detect_label_txt_path),
+                    "key": key,
+                    "type": "yolov7-object-detector",
+                },
+            ]
+        }
diff --git a/model_singleton.py b/model_singleton.py
@@ -0,0 +1,26 @@
+
+from models.experimental import attempt_load
+from utils.torch_utils import select_device
+
+class ModelSingleton:
+    _model = None
+    _stride = None
+
+    _instance = None
+    def __new__(cls, *args, **kwargs):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def __init__(self, weights_path: str, device_name: str):
+        device = select_device(device_name)
+
+        # Load model
+        self._model = attempt_load(weights_path, map_location=device)  # load FP32 model
+        self._stride = int(self._model.stride.max())  # model stride
+
+    def get_model(self):
+        return self._model
+
+    def get_stride(self):
+        return self._stride
diff --git a/requirements.txt b/requirements.txt
@@ -3,7 +3,7 @@
 # Base ----------------------------------------
 matplotlib>=3.2.2
 numpy>=1.18.5
-opencv-python>=4.1.1
+opencv-python==4.5.5.64
 Pillow>=7.1.2
 PyYAML>=5.3.1
 requests>=2.23.0