Merge pull request #109 fengyuentau:fix_yolox_issues

Resolves #108:

- Renamed `YoloX.py` to `yolox.py` for import.

- Reimplemented batched-nms.

- Put anchor generation in the initialization stage to avoid re-generating in inference.

YoloX.py

@@ -17,6 +17,8 @@ class YoloX:
         self.net.setPreferableBackend(self.backendId)
         self.net.setPreferableTarget(self.targetId)
 
+        self.generateAnchors()
+
     @property
     def name(self):
         return self.__class__.__name__
@@ -43,51 +45,45 @@ class YoloX:
         return predictions
 
     def postprocess(self, outputs):
-        grids = []
-        expanded_strides = []
-        hsizes = [self.input_size[0] // stride for stride in self.strides]
-        wsizes = [self.input_size[1] // stride for stride in self.strides]
+        dets = outputs[0]
 
-        for hsize, wsize, stride in zip(hsizes, wsizes, self.strides):
-            xv, yv = np.meshgrid(np.arange(hsize), np.arange(wsize))
-            grid = np.stack((xv, yv), 2).reshape(1, -1, 2)
-            grids.append(grid)
-            shape = grid.shape[:2]
-            expanded_strides.append(np.full((*shape, 1), stride))
-
-        grids = np.concatenate(grids, 1)
-        expanded_strides = np.concatenate(expanded_strides, 1)
-        outputs[..., :2] = (outputs[..., :2] + grids) * expanded_strides
-        outputs[..., 2:4] = np.exp(outputs[..., 2:4]) * expanded_strides
-
-        predictions = outputs[0]
-
-        boxes = predictions[:, :4]
-        scores = predictions[:, 4:5] * predictions[:, 5:]
+        dets[:, :2] = (dets[:, :2] + self.grids) * self.expanded_strides
+        dets[:, 2:4] = np.exp(dets[:, 2:4]) * self.expanded_strides
 
+        # get boxes
+        boxes = dets[:, :4]
         boxes_xyxy = np.ones_like(boxes)
         boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2] / 2.
         boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3] / 2.
         boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2] / 2.
         boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3] / 2.
 
-        # multi-class nms
-        final_dets = []
-        for cls_ind in range(scores.shape[1]):
-            cls_scores = scores[:, cls_ind]
-            valid_score_mask = cls_scores > self.confThreshold
-            if valid_score_mask.sum() == 0:
-                continue
-            else:
-                # call nms
-                indices = cv2.dnn.NMSBoxes(boxes_xyxy.tolist(), cls_scores.tolist(), self.confThreshold, self.nmsThreshold)
-
-                classids_ = np.ones((len(indices), 1)) * cls_ind
-                final_dets.append(
-                    np.concatenate([boxes_xyxy[indices], cls_scores[indices, None], classids_], axis=1)
-                )
-
-        if len(final_dets) == 0:
-            return np.array([])
-
-        return np.concatenate(final_dets, 0)
+        # get scores and class indices
+        scores = dets[:, 4:5] * dets[:, 5:]
+        max_scores = np.amax(scores, axis=1)
+        max_scores_idx = np.argmax(scores, axis=1)
+
+        # batched-nms, TODO: replace with cv2.dnn.NMSBoxesBatched when OpenCV 4.7.0 is released
+        max_coord = boxes_xyxy.max()
+        offsets = max_scores_idx * (max_coord + 1)
+        boxes_for_nms = boxes_xyxy + offsets[:, None]
+        keep = cv2.dnn.NMSBoxes(boxes_for_nms.tolist(), max_scores.tolist(), self.confThreshold, self.nmsThreshold)
+
+        candidates = np.concatenate([boxes_xyxy, max_scores[:, None], max_scores_idx[:, None]], axis=1)
+        return candidates[keep]
+
+    def generateAnchors(self):
+        self.grids = []
+        self.expanded_strides = []
+        hsizes = [self.input_size[0] // stride for stride in self.strides]
+        wsizes = [self.input_size[1] // stride for stride in self.strides]
+
+        for hsize, wsize, stride in zip(hsizes, wsizes, self.strides):
+            xv, yv = np.meshgrid(np.arange(hsize), np.arange(wsize))
+            grid = np.stack((xv, yv), 2).reshape(1, -1, 2)
+            self.grids.append(grid)
+            shape = grid.shape[:2]
+            self.expanded_strides.append(np.full((*shape, 1), stride))
+
+        self.grids = np.concatenate(self.grids, 1)
+        self.expanded_strides = np.concatenate(self.expanded_strides, 1)

yolox.py

@@ -0,0 +1,89 @@
+import numpy as np
+import cv2
+
+class YoloX:
+    def __init__(self, modelPath, confThreshold=0.35, nmsThreshold=0.5, objThreshold=0.5, backendId=0, targetId=0):
+        self.num_classes = 80
+        self.net = cv2.dnn.readNet(modelPath)
+        self.input_size = (640, 640)
+        self.mean = np.array([0.485, 0.456, 0.406], dtype=np.float32).reshape(1, 1, 3)
+        self.std = np.array([0.229, 0.224, 0.225], dtype=np.float32).reshape(1, 1, 3)
+        self.strides = [8, 16, 32]
+        self.confThreshold = confThreshold
+        self.nmsThreshold = nmsThreshold
+        self.objThreshold = objThreshold
+        self.backendId = backendId
+        self.targetId = targetId
+        self.net.setPreferableBackend(self.backendId)
+        self.net.setPreferableTarget(self.targetId)
+
+        self.generateAnchors()
+
+    @property
+    def name(self):
+        return self.__class__.__name__
+
+    def setBackend(self, backenId):
+        self.backendId = backendId
+        self.net.setPreferableBackend(self.backendId)
+
+    def setTarget(self, targetId):
+        self.targetId = targetId
+        self.net.setPreferableTarget(self.targetId)
+
+    def preprocess(self, img):
+        blob = np.transpose(img, (2, 0, 1))
+        return blob[np.newaxis, :, :, :]
+
+    def infer(self, srcimg):
+        input_blob = self.preprocess(srcimg)
+
+        self.net.setInput(input_blob)
+        outs = self.net.forward(self.net.getUnconnectedOutLayersNames())
+
+        predictions = self.postprocess(outs[0])
+        return predictions
+
+    def postprocess(self, outputs):
+        dets = outputs[0]
+
+        dets[:, :2] = (dets[:, :2] + self.grids) * self.expanded_strides
+        dets[:, 2:4] = np.exp(dets[:, 2:4]) * self.expanded_strides
+
+        # get boxes
+        boxes = dets[:, :4]
+        boxes_xyxy = np.ones_like(boxes)
+        boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2] / 2.
+        boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3] / 2.
+        boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2] / 2.
+        boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3] / 2.
+
+        # get scores and class indices
+        scores = dets[:, 4:5] * dets[:, 5:]
+        max_scores = np.amax(scores, axis=1)
+        max_scores_idx = np.argmax(scores, axis=1)
+
+        # batched-nms, TODO: replace with cv2.dnn.NMSBoxesBatched when OpenCV 4.7.0 is released
+        max_coord = boxes_xyxy.max()
+        offsets = max_scores_idx * (max_coord + 1)
+        boxes_for_nms = boxes_xyxy + offsets[:, None]
+        keep = cv2.dnn.NMSBoxes(boxes_for_nms.tolist(), max_scores.tolist(), self.confThreshold, self.nmsThreshold)
+
+        candidates = np.concatenate([boxes_xyxy, max_scores[:, None], max_scores_idx[:, None]], axis=1)
+        return candidates[keep]
+
+    def generateAnchors(self):
+        self.grids = []
+        self.expanded_strides = []
+        hsizes = [self.input_size[0] // stride for stride in self.strides]
+        wsizes = [self.input_size[1] // stride for stride in self.strides]
+
+        for hsize, wsize, stride in zip(hsizes, wsizes, self.strides):
+            xv, yv = np.meshgrid(np.arange(hsize), np.arange(wsize))
+            grid = np.stack((xv, yv), 2).reshape(1, -1, 2)
+            self.grids.append(grid)
+            shape = grid.shape[:2]
+            self.expanded_strides.append(np.full((*shape, 1), stride))
+
+        self.grids = np.concatenate(self.grids, 1)
+        self.expanded_strides = np.concatenate(self.expanded_strides, 1)