Merge pull request #184 from tenghehan-feature/mmdetection_support

support MMDetection

Signed-off-by: ZQPei <dfzspzq@163.com>

Author: ZQPei

.gitmodules

@@ -1,3 +1,6 @@
 [submodule "thirdparty/fast-reid"]
 	path = thirdparty/fast-reid
 	url = https://github.com/JDAI-CV/fast-reid.git
+[submodule "thirdparty/mmdetection"]
+	path = thirdparty/mmdetection
+	url = https://github.com/open-mmlab/mmdetection.git

README.md

@@ -79,26 +79,38 @@ cd ../../..
 Notice:
 If compiling failed, the simplist way is to **Upgrade your pytorch >= 1.1 and torchvision >= 0.3" and you can avoid the troublesome compiling problems which are most likely caused by either `gcc version too low` or `libraries missing`.
 
-5. (Optional) Prepare [fast-reid](https://github.com/JDAI-CV/fast-reid)
+5. (Optional) Prepare third party submodules
+
+[fast-reid](https://github.com/JDAI-CV/fast-reid)
 
 This library supports bagtricks, AGW and other mainstream ReID methods through providing an fast-reid adapter.
 
+to prepare our bundled fast-reid, then follow instructions in its README to install it.
+
+Please refer to `configs/fastreid.yaml` for a sample of using fast-reid. See [Model Zoo](https://github.com/JDAI-CV/fast-reid/blob/master/docs/MODEL_ZOO.md) for available methods and trained models.
+
+[MMDetection](https://github.com/open-mmlab/mmdetection)
+
+This library supports Faster R-CNN and other mainstream detection methods through providing an MMDetection adapter.
+
+to prepare our bundled MMDetection, then follow instructions in its README to install it.
+
+Please refer to `configs/mmdet.yaml` for a sample of using MMDetection. See [Model Zoo](https://github.com/open-mmlab/mmdetection/blob/master/docs/model_zoo.md) for available methods and trained models.
+
 Run
 
 ```
 git submodule update --init --recursive
 ```
 
-to prepare our bundled fast-reid, then follow instructions in its README to install it.
-
-Please refer to `configs/fastreid.yaml` for a sample of using fast-reid. See [Model Zoo](https://github.com/JDAI-CV/fast-reid/blob/master/docs/MODEL_ZOO.md) for available methods and trained models.
-
 
 6. Run demo
 ```
 usage: deepsort.py [-h]
                    [--fastreid]
                    [--config_fastreid CONFIG_FASTREID]
+                   [--mmdet]
+                   [--config_mmdetection CONFIG_MMDETECTION]
                    [--config_detection CONFIG_DETECTION]
                    [--config_deepsort CONFIG_DEEPSORT] [--display]
                    [--frame_interval FRAME_INTERVAL]
@@ -121,6 +133,9 @@ python3 deepsort.py /dev/video0 --config_detection ./configs/yolov3_tiny.yaml --
 
 # fast-reid + deepsort
 python deepsort.py [VIDEO_PATH] --fastreid [--config_fastreid ./configs/fastreid.yaml]
+
+# MMDetection + deepsort
+python deepsort.py [VIDEO_PATH] --mmdet [--config_mmdetection ./configs/mmdet.yaml]
 ```
 Use `--display` to enable display.  
 Results will be saved to `./output/results.avi` and `./output/results.txt`.

configs/mmdet.yaml

@@ -0,0 +1,5 @@
+MMDET:
+  CFG: "thirdparty/mmdetection/configs/faster_rcnn/faster_rcnn_r50_fpn_1x_coco.py"
+  CHECKPOINT: "detector/MMDet/weight/faster_rcnn_r50_fpn_1x_coco_20200130-047c8118.pth"
+
+  SCORE_THRESH: 0.5

deepsort.py

@@ -135,10 +135,12 @@ def run(self):
 def parse_args():
     parser = argparse.ArgumentParser()
     parser.add_argument("VIDEO_PATH", type=str)
+    parser.add_argument("--config_mmdetection", type=str, default="./configs/mmdet.yaml")
     parser.add_argument("--config_detection", type=str, default="./configs/yolov3.yaml")
     parser.add_argument("--config_deepsort", type=str, default="./configs/deep_sort.yaml")
     parser.add_argument("--config_fastreid", type=str, default="./configs/fastreid.yaml")
     parser.add_argument("--fastreid", action="store_true")
+    parser.add_argument("--mmdet", action="store_true")
     # parser.add_argument("--ignore_display", dest="display", action="store_false", default=True)
     parser.add_argument("--display", action="store_true")
     parser.add_argument("--frame_interval", type=int, default=1)
@@ -153,7 +155,12 @@ def parse_args():
 if __name__ == "__main__":
     args = parse_args()
     cfg = get_config()
-    cfg.merge_from_file(args.config_detection)
+    if args.mmdet:
+        cfg.merge_from_file(args.config_mmdetection)
+        cfg.USE_MMDET = True
+    else:
+        cfg.merge_from_file(args.config_detection)
+        cfg.USE_MMDET = False
     cfg.merge_from_file(args.config_deepsort)
     if args.fastreid:
         cfg.merge_from_file(args.config_fastreid)

detector/MMDet/__init__.py

@@ -0,0 +1,2 @@
+from .detector import MMDet
+__all__ = ['MMDet']

detector/MMDet/detector.py

@@ -0,0 +1,55 @@
+import logging
+import numpy as np
+import torch
+
+from mmdet.apis import init_detector, inference_detector
+from .mmdet_utils import xyxy_to_xywh
+
+class MMDet(object):
+    def __init__(self, cfg_file, checkpoint_file, score_thresh=0.7,
+                is_xywh=False, use_cuda=True):
+        # net definition
+        self.device = "cuda" if use_cuda else "cpu"
+        self.net = init_detector(cfg_file, checkpoint_file, device=self.device)
+        logger = logging.getLogger("root.detector")
+        logger.info('Loading weights from %s... Done!' % (checkpoint_file))
+
+        #constants
+        self.score_thresh = score_thresh
+        self.use_cuda = use_cuda
+        self.is_xywh = is_xywh
+        self.class_names = self.net.CLASSES
+        self.num_classes = len(self.class_names)
+
+    def __call__(self, ori_img):
+        # forward
+        bbox_result = inference_detector(self.net, ori_img)
+        bboxes = np.vstack(bbox_result)
+
+        if len(bboxes) == 0:
+            bbox = np.array([]).reshape([0, 4])
+            cls_conf = np.array([])
+            cls_ids = np.array([])
+            return bbox, cls_conf, cls_ids
+
+        bbox = bboxes[:, :4]
+        cls_conf = bboxes[:, 4]
+        cls_ids = [
+            np.full(bbox.shape[0], i, dtype=np.int32)
+            for i, bbox in enumerate(bbox_result)
+        ]
+        cls_ids = np.concatenate(cls_ids)
+
+        selected_idx = cls_conf > self.score_thresh
+        bbox = bbox[selected_idx, :]
+        cls_conf = cls_conf[selected_idx]
+        cls_ids = cls_ids[selected_idx]
+
+        if self.is_xywh:
+            bbox = xyxy_to_xywh(bbox)
+
+        return bbox, cls_conf, cls_ids
+
+
+            
+

detector/MMDet/mmdet_utils.py

@@ -0,0 +1,15 @@
+import torch
+import numpy as np
+
+def xyxy_to_xywh(boxes_xyxy):
+    if isinstance(boxes_xyxy, torch.Tensor):
+        boxes_xywh = boxes_xyxy.clone()
+    elif isinstance(boxes_xyxy, np.ndarray):
+        boxes_xywh = boxes_xyxy.copy()
+
+    boxes_xywh[:, 0] = (boxes_xyxy[:, 0] + boxes_xyxy[:, 2]) / 2.
+    boxes_xywh[:, 1] = (boxes_xyxy[:, 1] + boxes_xyxy[:, 3]) / 2.
+    boxes_xywh[:, 2] = boxes_xyxy[:, 2] - boxes_xyxy[:, 0]
+    boxes_xywh[:, 3] = boxes_xyxy[:, 3] - boxes_xyxy[:, 1]
+
+    return boxes_xywh

detector/__init__.py

@@ -1,9 +1,15 @@
 from .YOLOv3 import YOLOv3
+from .MMDet import MMDet
 
 
 __all__ = ['build_detector']
 
 def build_detector(cfg, use_cuda):
-    return YOLOv3(cfg.YOLOV3.CFG, cfg.YOLOV3.WEIGHT, cfg.YOLOV3.CLASS_NAMES, 
+    if cfg.USE_MMDET:
+        return MMDet(cfg.MMDET.CFG, cfg.MMDET.CHECKPOINT,
+                    score_thresh=cfg.MMDET.SCORE_THRESH,
+                    is_xywh=True, use_cuda=use_cuda)
+    else:
+        return YOLOv3(cfg.YOLOV3.CFG, cfg.YOLOV3.WEIGHT, cfg.YOLOV3.CLASS_NAMES, 
                     score_thresh=cfg.YOLOV3.SCORE_THRESH, nms_thresh=cfg.YOLOV3.NMS_THRESH, 
                     is_xywh=True, use_cuda=use_cuda)