Add mmdetection and related tools (openvinotoolkit#96)

* Add mmdetection submodule

* Add ssd-related tools to object detection folder

Author: sovrasov

.gitmodules

@@ -5,3 +5,6 @@
 [submodule "models"]
 	path = external/models
 	url = https://github.com/tensorflow/models.git
+[submodule "external/mmdetection"]
+	path = external/mmdetection
+	url = https://github.com/sovrasov/mmdetection/

external/mmdetection

@@ -0,0 +1,30 @@
+#!/usr/bin/env bash
+
+work_dir=$(realpath "$(dirname $0)")
+
+cd ${work_dir}
+if [[ -e venv ]]; then
+  echo "Please remove a previously virtual environment folder '${work_dir}/venv'."
+  exit
+fi
+
+# Create virtual environment
+virtualenv venv -p python3 --prompt="(pytorch-toolbox) "
+echo "export PYTHONPATH=\$PYTHONPATH:${work_dir}" >> venv/bin/activate
+. venv/bin/activate
+pip install -r ${work_dir}/requirements.txt
+
+# Install OpenVino Model Optimizer (optional)
+mo_requirements_file="${INTEL_CVSDK_DIR}/deployment_tools/model_optimizer/requirements_tf.txt"
+if [[ -e "${mo_requirements_file}" ]]; then
+  pip install -qr ${mo_requirements_file}
+else
+  echo "Model optimizer requirements were not installed. Please install the OpenVino toolkit to use one."
+fi
+
+
+echo
+echo "===================================================="
+echo "To start to work, you need to activate a virtualenv:"
+echo "$ . venv/bin/activate"
+echo "===================================================="

pytorch_toolkit/object_detection/init_venv.sh

@@ -0,0 +1,3 @@
+cd ../../external/mmdetection/
+bash compile.sh
+python setup.py develop

pytorch_toolkit/object_detection/prepare_mmdet.sh

@@ -0,0 +1,4 @@
+torch==1.1
+cython
+matplotlib
+ptflops

pytorch_toolkit/object_detection/requirements.txt

@@ -0,0 +1,38 @@
+import argparse
+
+import torch
+from mmdet.apis import init_detector
+from ptflops import get_model_complexity_info
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description='MMDet count flops')
+    parser.add_argument('config', help='test config file path')
+    args = parser.parse_args()
+    return args
+
+
+def inp_fun(input_res):
+    batch = torch.FloatTensor(1, 3, *input_res).cuda()
+    return dict(img=[batch], img_meta=[[{'img_shape': (*input_res, 3),
+                                         'ori_shape': (*input_res, 3),
+                                         'scale_factor': 1.0}]],
+                rescale=True, return_loss=False)
+
+
+def main():
+    args = parse_args()
+    with torch.no_grad():
+        model = init_detector(args.config)
+        model.eval()
+        input_res = model.cfg.data['test']['img_scale']
+        flops, params = get_model_complexity_info(model, input_res,
+                                                  as_strings=True,
+                                                  print_per_layer_stat=True,
+                                                  input_constructor=inp_fun)
+        print('Computational complexity: ' + flops)
+        print('Number of parameters: ', params)
+
+
+if __name__ == '__main__':
+    main()

pytorch_toolkit/object_detection/tools/count_flops.py

@@ -0,0 +1,67 @@
+import argparse
+import cv2 as cv
+from mmdet.apis import inference_detector, init_detector
+
+
+def decode_detections(detections, conf_t=0.5):
+    results = []
+    for detection in detections:
+        confidence = detection[4]
+
+        if confidence > conf_t:
+            left, top, right, bottom = detection[:4]
+            results.append(((int(left), int(top), int(right), int(bottom)),
+                            confidence))
+
+    return results
+
+
+def draw_detections(frame, detections, class_name):
+    """Draws detections and labels"""
+    for i, rect in enumerate(detections):
+        left, top, right, bottom = rect[0]
+        cv.rectangle(frame, (left, top), (right, bottom),
+                     (0, 255, 0), thickness=2)
+        label = class_name + '(' + str(round(rect[1], 2)) + ')'
+        label_size, base_line = cv.getTextSize(label,
+                                               cv.FONT_HERSHEY_SIMPLEX, 1, 1)
+        top = max(top, label_size[1])
+        cv.rectangle(frame, (left, top - label_size[1]),
+                     (left + label_size[0], top + base_line),
+                     (255, 255, 255), cv.FILLED)
+        cv.putText(frame, label, (left, top),
+                   cv.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0))
+
+    return frame
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Face detection live \
+                                                  demo script')
+    parser.add_argument('--cam_id', type=int, default=0, help='Input cam')
+    parser.add_argument('config', help='test config file path')
+    parser.add_argument('checkpoint', help='checkpoint file')
+    parser.add_argument('--d_thresh', type=float, default=0.5,
+                        help='Threshold for FD')
+    args = parser.parse_args()
+
+    model = init_detector(args.config, args.checkpoint)
+
+    cap = cv.VideoCapture(args.cam_id)
+    cap.set(cv.CAP_PROP_FRAME_WIDTH, 1280)
+    cap.set(cv.CAP_PROP_FRAME_HEIGHT, 720)
+    cap.set(cv.CAP_PROP_FOURCC, cv.VideoWriter_fourcc(*'MJPG'))
+
+    while cv.waitKey(1) != 27:
+        has_frame, frame = cap.read()
+        if not has_frame:
+            return
+        results = inference_detector(model, frame)
+        for i, class_result in enumerate(results):
+            class_boxes = decode_detections(class_result, args.d_thresh)
+            frame = draw_detections(frame, class_boxes, model.CLASSES[i])
+        cv.imshow('Detection Demo', frame)
+
+
+if __name__ == '__main__':
+    main()

pytorch_toolkit/object_detection/tools/detection_live_demo.py

@@ -0,0 +1,135 @@
+import argparse
+
+import numpy as np
+import torch
+from mmcv.parallel import MMDataParallel
+
+from mmdet.apis import init_detector
+from mmdet.models import detectors
+
+from tools.ssd_export_helpers import (get_proposals, PriorBox,
+                                      PriorBoxClustered, DetectionOutput)
+
+
+def onnx_export(self, img, img_meta, export_name='', **kwargs):
+    self._export_mode = True
+    self.img_metas = img_meta
+    torch.onnx.export(self, img, export_name, verbose=False)
+
+
+def forward(self, img, img_meta=[None], return_loss=True, **kwargs): #passing None here is a hack to fool the jit engine
+    if self._export_mode:
+        return self.forward_export(img)
+    if return_loss:
+        return self.forward_train(img, img_meta, **kwargs)
+    else:
+        return self.forward_test(img, img_meta, **kwargs)
+
+
+def forward_export_detector(self, img):
+    x = self.extract_feat(img)
+    outs = self.bbox_head(x)
+    bbox_result = self.bbox_head.export_forward(*outs, self.test_cfg, True,
+                                                self.img_metas, x, img)
+    return bbox_result
+
+
+def export_forward_ssd_head(self, cls_scores, bbox_preds, cfg, rescale,
+                            img_metas, feats, img_tensor):
+    num_levels = len(cls_scores)
+
+    anchors = []
+    for i in range(num_levels):
+        if self.anchor_generators[i].manual_anchors:
+            anchors.append(PriorBoxClustered.apply(
+                self.anchor_generators[i], self.anchor_strides[i],
+                feats[i], img_tensor, self.target_stds))
+        else:
+            anchors.append(PriorBox.apply(self.anchor_generators[i],
+                                          self.anchor_strides[i],
+                                          feats[i],
+                                          img_tensor, self.target_stds))
+    anchors = torch.cat(anchors, 2)
+    cls_scores, bbox_preds = self._prepare_cls_scores_bbox_preds(
+                                                cls_scores, bbox_preds)
+
+    return DetectionOutput.apply(cls_scores, bbox_preds, img_metas, cfg,
+                                 rescale, anchors, self.cls_out_channels,
+                                 self.use_sigmoid_cls, self.target_means,
+                                 self.target_stds)
+
+
+def prepare_cls_scores_bbox_preds_ssd_head(self, cls_scores, bbox_preds):
+    scores_list = []
+    for o in cls_scores:
+        score = o.permute(0, 2, 3, 1).contiguous().view(o.size(0), -1)
+        scores_list.append(score)
+    cls_scores = torch.cat(scores_list, 1)
+    cls_scores = cls_scores.view(cls_scores.size(0), -1, self.num_classes)
+    if self.use_sigmoid_cls:
+        cls_scores = cls_scores.sigmoid()
+    else:
+        cls_scores = cls_scores.softmax(-1)
+    cls_scores = cls_scores.view(cls_scores.size(0), -1)
+    bbox_list = []
+    for o in bbox_preds:
+        boxes = o.permute(0, 2, 3, 1).contiguous().view(o.size(0), -1)
+        bbox_list.append(boxes)
+    bbox_preds = torch.cat(bbox_list, 1)
+    return cls_scores, bbox_preds
+
+
+def get_bboxes_ssd_head(self, cls_scores, bbox_preds, img_metas, cfg,
+               rescale=False):
+    assert len(cls_scores) == len(bbox_preds)
+    num_levels = len(cls_scores)
+    mlvl_anchors = [
+        self.anchor_generators[i].grid_anchors(cls_scores[i].size()[-2:],
+                                               self.anchor_strides[i])
+        for i in range(num_levels)
+    ]
+    mlvl_anchors = torch.cat(mlvl_anchors, 0)
+    cls_scores, bbox_preds = self._prepare_cls_scores_bbox_preds(
+                                                    cls_scores, bbox_preds)
+    bboxes_list = get_proposals(img_metas, cls_scores, bbox_preds,
+                                mlvl_anchors, cfg, rescale,
+                                self.cls_out_channels,
+                                self.use_sigmoid_cls, self.target_means,
+                                self.target_stds)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description='MMDet onnx exporter for \
+                                                  SSD detector')
+    parser.add_argument('config', help='config file path')
+    parser.add_argument('checkpoint', help='checkpoint file')
+    parser.add_argument('output', help='onnx file')
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = parse_args()
+
+    model = init_detector(args.config, args.checkpoint)
+    cfg = model.cfg
+    assert getattr(detectors, cfg.model['type']) is \
+        detectors.SingleStageDetector
+    model = MMDataParallel(model, device_ids=[0])
+
+    batch = torch.FloatTensor(1, 3, cfg.input_size, cfg.input_size).cuda()
+    input_shape = (cfg.input_size, cfg.input_size, 3)
+    scale = np.array([1, 1, 1, 1], dtype=np.float32)
+    data = dict(img=batch, img_meta=[{'img_shape': input_shape,
+                                      'scale_factor': scale}])
+    model.eval()
+    model.module.onnx_export = onnx_export.__get__(model.module)
+    model.module.forward = forward.__get__(model.module)
+    model.module.forward_export = forward_export_detector.__get__(model.module)
+    model.module.bbox_head.export_forward = export_forward_ssd_head.__get__(model.module.bbox_head)
+    model.module.bbox_head._prepare_cls_scores_bbox_preds = prepare_cls_scores_bbox_preds_ssd_head.__get__(model.module.bbox_head)
+    model.module.bbox_head.get_bboxes = get_bboxes_ssd_head.__get__(model.module.bbox_head)
+    model.module.onnx_export(export_name=args.output, **data)
+
+if __name__ == '__main__':
+    main()