Simple detector

core/image/__init__.py

@@ -1,2 +1,2 @@
-from core.image.common import read_image, read_video, preprocess_image, postprocess_image
-from core.image.draw import draw_bboxes, Shader
+from core.image.common import read_image, read_video, preprocess_image, postprocess_image, postprocess_bboxes, BoundingBox
+from core.image.draw import draw_bboxes, Shader

core/image/common.py

@@ -3,6 +3,17 @@
 import cv2
 
 
+class BoundingBox:
+    def __init__(self, xmin, ymin, xmax, ymax, score, cls):
+        self.xmin = xmin
+        self.ymin = ymin
+        self.xmax = xmax
+        self.ymax = ymax
+
+        self.score = score
+        self.cls = int(cls)
+
+
 def read_image(*args, **kwargs):
     return cv2.cvtColor(cv2.imread(*args, **kwargs), cv2.COLOR_BGR2RGB)
 
@@ -11,6 +22,26 @@ def read_video(*args, **kwargs):
     return cv2.VideoCapture(*args, **kwargs)
 
 
+def postprocess_bboxes(input_size, output_size, bboxes, scores, classes):
+    input_w, input_h = input_size
+    output_w, output_h = output_size
+
+    scale = min(input_w / output_w, input_h / output_h)
+    nw, nh = int(scale * output_w), int(scale * output_h)
+    dw, dh = (input_w - nw) // 2, (input_h - nh) // 2
+
+    bboxes = bboxes.astype(np.float32)
+    bboxes[:, [0, 2]] = np.clip((bboxes[:, [0, 2]] - dw) / scale, 0., output_w)
+    bboxes[:, [1, 3]] = np.clip((bboxes[:, [1, 3]] - dh) / scale, 0., output_h)
+
+    boxes = []
+    for i in range(len(bboxes)):
+        x1, y1, x2, y2 = bboxes[i][:4]
+        boxes.append(BoundingBox(x1, y1, x2, y2, scores[i], classes[i]))
+
+    return boxes
+
+
 def postprocess_image(image, size, bboxes=None):
     """
     :param image: RGB, uint8

core/image/draw.py

@@ -33,7 +33,8 @@ def h2rgb(h):
             elif i == 5:
                 return c + m, m, x + m
 
-        self._colors = {i: h2rgb(i / num_colors * 6) for i in range(num_colors)}
+        self._colors = {i: h2rgb(i / num_colors * 6)
+                        for i in range(num_colors)}
 
     def get_color(self, index):
         """
@@ -62,7 +63,8 @@ def _draw_bboxes_relative(img, bboxes, scores, classes, names, shader):
 
     def _draw_bbox(img, bbox, score, cls):
         msg = '{} {:.2%}'.format(names[int(cls)], score)
-        (x, y), base = cv2.getTextSize(msg, cv2.FONT_HERSHEY_COMPLEX_SMALL, 0.5, 1)
+        (x, y), base = cv2.getTextSize(
+            msg, cv2.FONT_HERSHEY_COMPLEX_SMALL, 0.5, 1)
         color = shader.get_color(int(cls))
 
         h, w = img.shape[:2]
@@ -95,7 +97,8 @@ def _draw_bboxes_absolute(img, bboxes, scores, classes, names, shader):
 
     def _draw_bbox(img, bbox, score, cls):
         msg = '{} {:.2%}'.format(names[int(cls)], score)
-        (x, y), base = cv2.getTextSize(msg, cv2.FONT_HERSHEY_COMPLEX_SMALL, 0.5, 1)
+        (x, y), base = cv2.getTextSize(
+            msg, cv2.FONT_HERSHEY_COMPLEX_SMALL, 0.5, 1)
         color = shader.get_color(int(cls))
 
         x1, y1, x2, y2 = bbox[:4]
@@ -115,3 +118,36 @@ def _draw_bbox(img, bbox, score, cls):
         img = _draw_bbox(img, bbox, score, cls)
 
     return img
+
+
+def draw_bboxes(img, boxes, names, shader):
+    """
+    BBoxes is absolute Format
+    :param img: BGR, uint8
+    :param bboxes: x1, y1, x2, y2, int
+    :return: img, BGR, uint8
+    """
+
+    def _draw_bbox(img, bbox):
+        msg = '{} {:.2%}'.format(names[bbox.cls], bbox.score)
+        (x, y), base = cv2.getTextSize(
+            msg, cv2.FONT_HERSHEY_COMPLEX_SMALL, 0.5, 1)
+        color = shader.get_color(bbox.cls)
+
+        x1, y1, x2, y2 = bbox.xmin, bbox.ymin, bbox.xmax, bbox.ymax
+        x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
+        img = cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
+        img = cv2.rectangle(img, (x1, y1 - y - base),
+                            (x1 + x, y1),
+                            color, -1)
+        img = cv2.putText(img,
+                          msg,
+                          (x1, y1),
+                          cv2.FONT_HERSHEY_COMPLEX_SMALL, 0.5, (255, 255, 255), 1)
+
+        return img
+
+    for bbox in boxes:
+        img = _draw_bbox(img, bbox)
+
+    return img

simple_detector.py

@@ -0,0 +1,68 @@
+import numpy as np
+import cv2
+import glob
+import os
+from absl import app, flags
+
+from tensorflow.keras.models import load_model
+from core.image import draw_bboxes, preprocess_image, postprocess_bboxes, Shader, read_image
+
+
+class Inference:
+    def __init__(self, model_path):
+        assert os.path.isdir(model_path), \
+            "Model {} does not exist.".format(model_path)
+
+        # Init tf model
+        self.model = load_model(model_path)
+
+    def predict(self, image, net_input_size):
+
+        input_h, input_w = image.shape[:2]
+        image = preprocess_image(
+            image, (net_input_size, net_input_size)).astype(np.float32)
+        output_h, output_w = image.shape[:2]
+
+        images = np.expand_dims(image, axis=0)
+        bboxes, scores, classes, valid_detections = self.model.predict(images)
+
+        bboxes = bboxes[0][:valid_detections[0]]
+        scores = scores[0][:valid_detections[0]]
+        classes = classes[0][:valid_detections[0]]
+
+        return postprocess_bboxes((output_w, output_h), (input_w, input_h), bboxes, scores, classes)
+
+
+flags.DEFINE_string(
+    'model', 'models/full_yolo4', 'Path to the keras full model (not only the weights)')
+flags.DEFINE_integer('image_size', '416', 'CNN input image size')
+flags.DEFINE_string('image_folder', '/data/Images/',
+                    'Path to the image folder')
+flags.DEFINE_multi_string('names', ['class1', 'class2'], 'class names')
+FLAGS = flags.FLAGS
+
+
+def main(_argv):
+    model = Inference(FLAGS.model)
+
+    shader = Shader(1)
+
+    for img in glob.glob(FLAGS.image_folder+"*"):
+
+        image = read_image(img)
+
+        bboxes = model.predict(image, FLAGS.image_size)
+
+        for bbox in bboxes:
+            if bbox.score > 0.45:
+                image = draw_bboxes(image, [bbox], FLAGS.names, shader)
+
+        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+        cv2.imshow('Image', image)
+
+        if cv2.waitKey(0) == ord('q'):
+            exit(0)
+
+
+if __name__ == "__main__":
+    app.run(main)