refactor(YOLOX): refactor datasets and add demo_utils

datasets/README.md

@@ -0,0 +1,24 @@
+# Prepare datasets
+
+If you have a dataset directory, you could use os environment variable named `YOLOX_DATADIR`. Under this directory, YOLOX will look for datasets in the structure described below, if needed.
+```
+$YOLOX_DATADIR/
+  COCO/
+```
+You can set the location for builtin datasets by
+```shell
+export YOLOX_DATADIR=/path/to/your/datasets
+```
+If `YOLOX_DATADIR` is not set, the default value of dataset directory is `./datasets` relative to your current working directory.
+
+## Expected dataset structure for [COCO detection](https://cocodataset.org/#download):
+
+```
+COCO/
+  annotations/
+    instances_{train,val}2017.json
+  {train,val}2017/
+    # image files that are mentioned in the corresponding json
+```
+
+You can use the 2014 version of the dataset as well.

demo/ONNXRuntime/onnx_inference.py

@@ -1,14 +1,18 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import argparse
+import os
+
 import cv2
 import numpy as np
 
+import onnxruntime
+
 from yolox.data.data_augment import preproc as preprocess
 from yolox.data.datasets import COCO_CLASSES
-from yolox.utils.visualize import vis
-
-import argparse
-import onnxruntime
-import os
-from demo_utils import mkdir, multiclass_nms, postprocess
+from yolox.utils import mkdir, multiclass_nms, postprocess, vis
 
 
 def make_parser():

demo/OpenVINO/python/demo_utils.py

@@ -1,86 +0,0 @@
-import numpy as np
-
-import os
-
-
-def mkdir(path):
-    if not os.path.exists(path):
-        os.makedirs(path)
-
-
-def nms(boxes, scores, nms_thr):
-    """Single class NMS implemented in Numpy."""
-    x1 = boxes[:, 0]
-    y1 = boxes[:, 1]
-    x2 = boxes[:, 2]
-    y2 = boxes[:, 3]
-
-    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
-    order = scores.argsort()[::-1]
-
-    keep = []
-    while order.size > 0:
-        i = order[0]
-        keep.append(i)
-        xx1 = np.maximum(x1[i], x1[order[1:]])
-        yy1 = np.maximum(y1[i], y1[order[1:]])
-        xx2 = np.minimum(x2[i], x2[order[1:]])
-        yy2 = np.minimum(y2[i], y2[order[1:]])
-
-        w = np.maximum(0.0, xx2 - xx1 + 1)
-        h = np.maximum(0.0, yy2 - yy1 + 1)
-        inter = w * h
-        ovr = inter / (areas[i] + areas[order[1:]] - inter)
-
-        inds = np.where(ovr <= nms_thr)[0]
-        order = order[inds + 1]
-
-    return keep
-
-
-def multiclass_nms(boxes, scores, nms_thr, score_thr):
-    """Multiclass NMS implemented in Numpy"""
-    final_dets = []
-    num_classes = scores.shape[1]
-    for cls_ind in range(num_classes):
-        cls_scores = scores[:, cls_ind]
-        valid_score_mask = cls_scores > score_thr
-        if valid_score_mask.sum() == 0:
-            continue
-        else:
-            valid_scores = cls_scores[valid_score_mask]
-            valid_boxes = boxes[valid_score_mask]
-            keep = nms(valid_boxes, valid_scores, nms_thr)
-            if len(keep) > 0:
-                cls_inds = np.ones((len(keep), 1)) * cls_ind
-                dets = np.concatenate([valid_boxes[keep], valid_scores[keep, None], cls_inds], 1)
-                final_dets.append(dets)
-    return np.concatenate(final_dets, 0)
-
-
-def postprocess(outputs, img_size, p6=False):
-
-    grids = []
-    expanded_strides = []
-
-    if not p6:
-        strides = [8, 16, 32]
-    else:
-        strides = [8, 16, 32, 64]
-
-    hsizes = [img_size[0]//stride for stride in strides]
-    wsizes = [img_size[1]//stride for stride in strides]
-
-    for hsize, wsize, stride in zip(hsizes, wsizes, 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
-
-    return outputs

demo/OpenVINO/python/openvino_inference.py

@@ -2,6 +2,8 @@
 # -*- coding: utf-8 -*-
 # Copyright (C) 2018-2021 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Megvii, Inc. and its affiliates.
+
 import argparse
 import logging as log
 import os
@@ -10,11 +12,11 @@ import sys
 import cv2
 import numpy as np
 
-from demo_utils import mkdir, multiclass_nms, postprocess
 from openvino.inference_engine import IECore
+
 from yolox.data.data_augment import preproc as preprocess
 from yolox.data.datasets import COCO_CLASSES
-from yolox.utils.visualize import vis
+from yolox.utils import mkdir, multiclass_nms, postprocess, vis
 
 
 def parse_args() -> argparse.Namespace:

setup.cfg

@@ -6,7 +6,7 @@ known_standard_library = setuptools
 known_third_party = tqdm,loguru
 known_data_processing = cv2,numpy,scipy,PIL,matplotlib,scikit_image
 known_datasets = pycocotools
-known_deeplearning = torch,torchvision,caffe2,onnx,apex,timm,thop,torch2trt,tensorrt
+known_deeplearning = torch,torchvision,caffe2,onnx,apex,timm,thop,torch2trt,tensorrt,openvino,onnxruntime
 known_myself = yolox
 sections = FUTURE,STDLIB,THIRDPARTY,data_processing,datasets,deeplearning,myself,FIRSTPARTY,LOCALFOLDER
 no_lines_before=STDLIB,THIRDPARTY,datasets

yolox/data/dataloading.py

@@ -20,7 +20,8 @@ def get_yolox_datadir():
     yolox_datadir = os.getenv("YOLOX_DATADIR", None)
     if yolox_datadir is None:
         import yolox
-        yolox_datadir = os.path.join(os.path.dirname(yolox.__file__), "data")
+        yolox_path = os.path.dirname(os.path.dirname(yolox.__file__))
+        yolox_datadir = os.path.join(yolox_path, "datasets")
     return yolox_datadir
 
 

yolox/utils/__init__.py

@@ -5,6 +5,7 @@
 from .allreduce_norm import *
 from .boxes import *
 from .checkpoint import load_ckpt, save_checkpoint
+from .demo_utils import *
 from .dist import *
 from .ema import ModelEMA
 from .logger import setup_logger

{demo/ONNXRuntime → yolox/utils}/demo_utils.py

@@ -1,7 +1,13 @@
-import numpy as np
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
 
 import os
 
+import numpy as np
+
+__all__ = ["mkdir", "nms", "multiclass_nms", "postprocess"]
+
 
 def mkdir(path):
     if not os.path.exists(path):