Add app

.gitignore

@@ -0,0 +1,169 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+# Usually these files are written by a python script from a template
+# before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+# For a library or package, you might want to ignore these files since the code is
+# intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+# However, in case of collaboration, if having platform-specific dependencies or dependencies
+# having no cross-platform support, pipenv may install dependencies that don't work, or not
+# install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+# This is especially recommended for binary packages to ensure reproducibility, and is more
+# commonly ignored for libraries.
+# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+# in version control.
+# https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+# and can be added to the global gitignore or merged into this file. For a more nuclear
+# option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+*.png
+*.jpg
+*.mp4
+
+YOLOX_outputs/
+artifacts/
+*.engine
+*.pth

README.md

@@ -1,6 +1,6 @@
 ---
 title: OpenLenda
-emoji: 🏃
+emoji: 🚥
 colorFrom: blue
 colorTo: purple
 sdk: gradio

app.py

@@ -0,0 +1,108 @@
+
+from yolox.exp import get_exp
+from yolox.data.datasets import COCO_CLASSES
+from predictor import Predictor
+
+import cv2
+import gradio as gr
+import torch
+
+import subprocess
+import tempfile
+import time
+from pathlib import Path
+
+exp = get_exp("exps/openlenda_s.py", None)
+model = exp.get_model()
+model.eval()
+ckpt_file = "models/openlenda_s.pth"
+model.load_state_dict(torch.load(ckpt_file, map_location="cpu")["model"])
+predictor = Predictor(
+ model, COCO_CLASSES, "cpu", False, False
+)
+
+
+def image_inference(image, confthre, nmsthre):
+ cv2.cvtColor(image, cv2.COLOR_RGB2BGR, image)
+ outputs, img_info = predictor.inference(image, confthre, nmsthre)
+ result_image = predictor.visual(outputs[0], img_info)
+ cv2.cvtColor(result_image, cv2.COLOR_BGR2RGB, result_image)
+ return result_image
+
+
+image_interface = gr.Interface(
+ fn=image_inference,
+ inputs=[
+ "image",
+ gr.Slider(0, 1, value=0.5, step=0.01, label="Confidence Threshold", ),
+ gr.Slider(0, 1, value=0.01, step=0.01, label="NMS Threshold")
+ ],
+ examples=[["assets/sample.png", 0.5, 0.01]],
+ outputs=gr.Image(type="pil"),
+ title="OpenLenda image demo"
+)
+
+
+def video_inference(video_file, confthre, nmsthre, start_sec, duration):
+ start_timestamp = time.strftime("%H:%M:%S", time.gmtime(start_sec))
+ end_timestamp = time.strftime("%H:%M:%S", time.gmtime(start_sec + duration))
+
+ suffix = Path(video_file).suffix
+
+ clip_temp_file = tempfile.NamedTemporaryFile(suffix=suffix)
+ subprocess.call(
+ f"ffmpeg -y -ss {start_timestamp} -i {video_file} -to {end_timestamp} -c copy {clip_temp_file.name}".split()
+ )
+
+ cap = cv2.VideoCapture(clip_temp_file.name)
+ width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+ height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+ fps = cap.get(cv2.CAP_PROP_FPS)
+
+ with tempfile.NamedTemporaryFile(suffix=".mp4") as temp_file:
+ out = cv2.VideoWriter(temp_file.name, cv2.VideoWriter_fourcc(*"MP4V"), fps, (width, height))
+
+ num_frames = 0
+ max_frames = duration * fps
+ while cap.isOpened():
+ try:
+ ret, frame = cap.read()
+ if not ret:
+ break
+ except Exception as e:
+ print(e)
+ continue
+ outputs, img_info = predictor.inference(frame, confthre, nmsthre)
+ result_frame = predictor.visual(outputs[0], img_info)
+ out.write(result_frame)
+ num_frames += 1
+ if num_frames == max_frames:
+ break
+
+ out.release()
+
+ out_file = tempfile.NamedTemporaryFile(suffix="out.mp4", delete=False)
+ subprocess.run(f"ffmpeg -y -loglevel quiet -stats -i {temp_file.name} -c:v libx264 {out_file.name}".split())
+
+ return out_file.name
+
+
+video_interface = gr.Interface(
+ fn=video_inference,
+ inputs=[
+ gr.Video(),
+ gr.Slider(0, 1, value=0.5, step=0.01, label="Confidence Threshold", ),
+ gr.Slider(0, 1, value=0.01, step=0.01, label="NMS Threshold"),
+ gr.Slider(0, 60, value=0, step=1, label="Start Second"),
+ gr.Slider(0, 10, value=3, step=1, label="Duration"),
+ ],
+ outputs=gr.Video(),
+ title="OpenLenda video demo"
+)
+
+if __name__ == "__main__":
+ gr.TabbedInterface(
+ [image_interface, video_interface],
+ ["Image", "Video"],
+ title="OpenLenda demo!",
+ ).launch()

exps/default/__init__.py

@@ -0,0 +1,3 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.

exps/default/yolov3.py

@@ -0,0 +1,33 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import os
+
+import torch.nn as nn
+
+from yolox.exp import Exp as MyExp
+
+
+class Exp(MyExp):
+ def __init__(self):
+ super(Exp, self).__init__()
+ self.depth = 1.0
+ self.width = 1.0
+ self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
+
+ def get_model(self, sublinear=False):
+ def init_yolo(M):
+ for m in M.modules():
+ if isinstance(m, nn.BatchNorm2d):
+ m.eps = 1e-3
+ m.momentum = 0.03
+ if "model" not in self.__dict__:
+ from yolox.models import YOLOX, YOLOFPN, YOLOXHead
+ backbone = YOLOFPN()
+ head = YOLOXHead(self.num_classes, self.width, in_channels=[128, 256, 512], act="lrelu")
+ self.model = YOLOX(backbone, head)
+ self.model.apply(init_yolo)
+ self.model.head.initialize_biases(1e-2)
+
+ return self.model

exps/default/yolox_l.py

@@ -0,0 +1,15 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import os
+
+from yolox.exp import Exp as MyExp
+
+
+class Exp(MyExp):
+ def __init__(self):
+ super(Exp, self).__init__()
+ self.depth = 1.0
+ self.width = 1.0
+ self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]

exps/default/yolox_m.py

@@ -0,0 +1,15 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import os
+
+from yolox.exp import Exp as MyExp
+
+
+class Exp(MyExp):
+ def __init__(self):
+ super(Exp, self).__init__()
+ self.depth = 0.67
+ self.width = 0.75
+ self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]

exps/default/yolox_nano.py

@@ -0,0 +1,48 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import os
+
+import torch.nn as nn
+
+from yolox.exp import Exp as MyExp
+
+
+class Exp(MyExp):
+ def __init__(self):
+ super(Exp, self).__init__()
+ self.depth = 0.33
+ self.width = 0.25
+ self.input_size = (416, 416)
+ self.random_size = (10, 20)
+ self.mosaic_scale = (0.5, 1.5)
+ self.test_size = (416, 416)
+ self.mosaic_prob = 0.5
+ self.enable_mixup = False
+ self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
+
+ def get_model(self, sublinear=False):
+
+ def init_yolo(M):
+ for m in M.modules():
+ if isinstance(m, nn.BatchNorm2d):
+ m.eps = 1e-3
+ m.momentum = 0.03
+ if "model" not in self.__dict__:
+ from yolox.models import YOLOX, YOLOPAFPN, YOLOXHead
+ in_channels = [256, 512, 1024]
+ # NANO model use depthwise = True, which is main difference.
+ backbone = YOLOPAFPN(
+ self.depth, self.width, in_channels=in_channels,
+ act=self.act, depthwise=True,
+ )
+ head = YOLOXHead(
+ self.num_classes, self.width, in_channels=in_channels,
+ act=self.act, depthwise=True
+ )
+ self.model = YOLOX(backbone, head)
+
+ self.model.apply(init_yolo)
+ self.model.head.initialize_biases(1e-2)
+ return self.model

exps/default/yolox_s.py

@@ -0,0 +1,15 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import os
+
+from yolox.exp import Exp as MyExp
+
+
+class Exp(MyExp):
+ def __init__(self):
+ super(Exp, self).__init__()
+ self.depth = 0.33
+ self.width = 0.50
+ self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]

exps/default/yolox_tiny.py

@@ -0,0 +1,20 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import os
+
+from yolox.exp import Exp as MyExp
+
+
+class Exp(MyExp):
+ def __init__(self):
+ super(Exp, self).__init__()
+ self.depth = 0.33
+ self.width = 0.375
+ self.input_size = (416, 416)
+ self.mosaic_scale = (0.5, 1.5)
+ self.random_size = (10, 20)
+ self.test_size = (416, 416)
+ self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
+ self.enable_mixup = False

exps/default/yolox_x.py

@@ -0,0 +1,15 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import os
+
+from yolox.exp import Exp as MyExp
+
+
+class Exp(MyExp):
+ def __init__(self):
+ super(Exp, self).__init__()
+ self.depth = 1.33
+ self.width = 1.25
+ self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]

exps/openlenda_nano.py

@@ -0,0 +1,53 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import os
+
+import torch.nn as nn
+
+from yolox.exp import Exp as MyExp
+
+
+class Exp(MyExp):
+ def __init__(self):
+ super(Exp, self).__init__()
+ self.depth = 0.33
+ self.width = 0.25
+ self.input_size = (416, 416)
+ self.random_size = (10, 20)
+ self.mosaic_scale = (0.5, 1.5)
+ self.test_size = (416, 416)
+ self.mosaic_prob = 0.5
+ self.enable_mixup = False
+ self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
+ # max training epoch
+ self.max_epoch = 30
+ self.num_classes = 8
+ # --------------- transform config ----------------- #
+ self.flip_prob = 0
+
+ def get_model(self, sublinear=False):
+
+ def init_yolo(M):
+ for m in M.modules():
+ if isinstance(m, nn.BatchNorm2d):
+ m.eps = 1e-3
+ m.momentum = 0.03
+ if "model" not in self.__dict__:
+ from yolox.models import YOLOX, YOLOPAFPN, YOLOXHead
+ in_channels = [256, 512, 1024]
+ # NANO model use depthwise = True, which is main difference.
+ backbone = YOLOPAFPN(
+ self.depth, self.width, in_channels=in_channels,
+ act=self.act, depthwise=True,
+ )
+ head = YOLOXHead(
+ self.num_classes, self.width, in_channels=in_channels,
+ act=self.act, depthwise=True
+ )
+ self.model = YOLOX(backbone, head)
+
+ self.model.apply(init_yolo)
+ self.model.head.initialize_biases(1e-2)
+ return self.model

exps/openlenda_s.py

@@ -0,0 +1,21 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import os
+
+from yolox.exp import Exp as MyExp
+
+
+class Exp(MyExp):
+ def __init__(self):
+ super(Exp, self).__init__()
+ self.depth = 0.33
+ self.width = 0.50
+ self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
+ # max training epoch
+ self.max_epoch = 30
+ self.num_classes = 8
+ # --------------- transform config ----------------- #
+ self.flip_prob = 0
+ self.input_size = (1280, 1280) # (height, width)

exps/openlenda_tiny.py

@@ -0,0 +1,25 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import os
+
+from yolox.exp import Exp as MyExp
+
+
+class Exp(MyExp):
+ def __init__(self):
+ super(Exp, self).__init__()
+ self.depth = 0.33
+ self.width = 0.375
+ self.input_size = (416, 416)
+ self.mosaic_scale = (0.5, 1.5)
+ self.random_size = (10, 20)
+ self.test_size = (416, 416)
+ self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
+ self.enable_mixup = False
+ # max training epoch
+ self.max_epoch = 30
+ self.num_classes = 8
+ # --------------- transform config ----------------- #
+ self.flip_prob = 0

exps/openlenda_x.py

@@ -0,0 +1,20 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import os
+
+from yolox.exp import Exp as MyExp
+
+
+class Exp(MyExp):
+ def __init__(self):
+ super(Exp, self).__init__()
+ self.depth = 1.33
+ self.width = 1.25
+ self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
+ # max training epoch
+ self.max_epoch = 30
+ self.num_classes = 8
+ # --------------- transform config ----------------- #
+ self.input_size = (640, 800) # (height, width)

predictor.py

@@ -0,0 +1,87 @@
+import os
+import time
+from loguru import logger
+
+import cv2
+
+import torch
+
+from yolox.data.data_augment import ValTransform
+from yolox.data.datasets import COCO_CLASSES
+from yolox.utils import postprocess, vis
+
+
+class Predictor(object):
+ def __init__(
+ self,
+ model,
+ cls_names=COCO_CLASSES,
+ device="cpu",
+ fp16=False,
+ legacy=False,
+ ):
+ self.model = model
+ self.cls_names = cls_names
+ self.num_classes = len(COCO_CLASSES)
+ self.confthre = 0.01
+ self.nmsthre = 0.01
+ self.test_size = (640, 640)
+ self.device = device
+ self.fp16 = fp16
+ self.preproc = ValTransform(legacy=legacy)
+
+ def inference(self, img, confthre=None, nmsthre=None, test_size=None):
+ if confthre is not None:
+ self.confthre = confthre
+ if nmsthre is not None:
+ self.nmsthre = nmsthre
+ if test_size is not None:
+ self.test_size = test_size
+ img_info = {"id": 0}
+ if isinstance(img, str):
+ img_info["file_name"] = os.path.basename(img)
+ img = cv2.imread(img)
+ else:
+ img_info["file_name"] = None
+ cv2.imwrite("test.png", img)
+ height, width = img.shape[:2]
+ img_info["height"] = height
+ img_info["width"] = width
+ img_info["raw_img"] = img
+
+ ratio = min(self.test_size[0] / img.shape[0], self.test_size[1] / img.shape[1])
+ img_info["ratio"] = ratio
+
+ img, _ = self.preproc(img, None, self.test_size)
+ img = torch.from_numpy(img).unsqueeze(0)
+ img = img.float()
+ if self.device == "gpu":
+ img = img.cuda()
+ if self.fp16:
+ img = img.half() # to FP16
+
+ with torch.no_grad():
+ outputs = self.model(img)
+ outputs = postprocess(
+ outputs, self.num_classes, self.confthre,
+ self.nmsthre
+ )
+ return outputs, img_info
+
+ def visual(self, output, img_info):
+ ratio = img_info["ratio"]
+ img = img_info["raw_img"]
+ if output is None:
+ return img
+ output = output.cpu()
+
+ bboxes = output[:, 0:4]
+
+ # preprocessing: resize
+ bboxes /= ratio
+
+ cls = output[:, 6]
+ scores = output[:, 4] * output[:, 5]
+
+ vis_res = vis(img, bboxes, scores, cls, self.confthre, self.cls_names)
+ return vis_res

requirements.txt

@@ -0,0 +1,7 @@
+loguru
+tabulate
+psutil
+pycocotools
+torch >= 2.0.1
+torchvision >= 0.15.2
+opencv-python

yolox/__init__.py

@@ -0,0 +1,4 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+
+__version__ = "0.3.0"

yolox/core/__init__.py

@@ -0,0 +1,6 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+from .launch import launch
+from .trainer import Trainer

yolox/core/launch.py

@@ -0,0 +1,147 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Code are based on
+# https://github.com/facebookresearch/detectron2/blob/master/detectron2/engine/launch.py
+# Copyright (c) Facebook, Inc. and its affiliates.
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import sys
+from datetime import timedelta
+from loguru import logger
+
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+
+import yolox.utils.dist as comm
+
+__all__ = ["launch"]
+
+
+DEFAULT_TIMEOUT = timedelta(minutes=30)
+
+
+def _find_free_port():
+ """
+ Find an available port of current machine / node.
+ """
+ import socket
+
+ sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+ # Binding to port 0 will cause the OS to find an available port for us
+ sock.bind(("", 0))
+ port = sock.getsockname()[1]
+ sock.close()
+ # NOTE: there is still a chance the port could be taken by other processes.
+ return port
+
+
+def launch(
+ main_func,
+ num_gpus_per_machine,
+ num_machines=1,
+ machine_rank=0,
+ backend="nccl",
+ dist_url=None,
+ args=(),
+ timeout=DEFAULT_TIMEOUT,
+):
+ """
+ Args:
+ main_func: a function that will be called by `main_func(*args)`
+ num_machines (int): the total number of machines
+ machine_rank (int): the rank of this machine (one per machine)
+ dist_url (str): url to connect to for distributed training, including protocol
+ e.g. "tcp://127.0.0.1:8686".
+ Can be set to auto to automatically select a free port on localhost
+ args (tuple): arguments passed to main_func
+ """
+ world_size = num_machines * num_gpus_per_machine
+ if world_size > 1:
+ # https://github.com/pytorch/pytorch/pull/14391
+ # TODO prctl in spawned processes
+
+ if dist_url == "auto":
+ assert (
+ num_machines == 1
+ ), "dist_url=auto cannot work with distributed training."
+ port = _find_free_port()
+ dist_url = f"tcp://127.0.0.1:{port}"
+
+ start_method = "spawn"
+ cache = vars(args[1]).get("cache", False)
+
+ # To use numpy memmap for caching image into RAM, we have to use fork method
+ if cache:
+ assert sys.platform != "win32", (
+ "As Windows platform doesn't support fork method, "
+ "do not add --cache in your training command."
+ )
+ start_method = "fork"
+
+ mp.start_processes(
+ _distributed_worker,
+ nprocs=num_gpus_per_machine,
+ args=(
+ main_func,
+ world_size,
+ num_gpus_per_machine,
+ machine_rank,
+ backend,
+ dist_url,
+ args,
+ ),
+ daemon=False,
+ start_method=start_method,
+ )
+ else:
+ main_func(*args)
+
+
+def _distributed_worker(
+ local_rank,
+ main_func,
+ world_size,
+ num_gpus_per_machine,
+ machine_rank,
+ backend,
+ dist_url,
+ args,
+ timeout=DEFAULT_TIMEOUT,
+):
+ assert (
+ torch.cuda.is_available()
+ ), "cuda is not available. Please check your installation."
+ global_rank = machine_rank * num_gpus_per_machine + local_rank
+ logger.info("Rank {} initialization finished.".format(global_rank))
+ try:
+ dist.init_process_group(
+ backend=backend,
+ init_method=dist_url,
+ world_size=world_size,
+ rank=global_rank,
+ timeout=timeout,
+ )
+ except Exception:
+ logger.error("Process group URL: {}".format(dist_url))
+ raise
+
+ # Setup the local process group (which contains ranks within the same machine)
+ assert comm._LOCAL_PROCESS_GROUP is None
+ num_machines = world_size // num_gpus_per_machine
+ for i in range(num_machines):
+ ranks_on_i = list(
+ range(i * num_gpus_per_machine, (i + 1) * num_gpus_per_machine)
+ )
+ pg = dist.new_group(ranks_on_i)
+ if i == machine_rank:
+ comm._LOCAL_PROCESS_GROUP = pg
+
+ # synchronize is needed here to prevent a possible timeout after calling init_process_group
+ # See: https://github.com/facebookresearch/maskrcnn-benchmark/issues/172
+ comm.synchronize()
+
+ assert num_gpus_per_machine <= torch.cuda.device_count()
+ torch.cuda.set_device(local_rank)
+
+ main_func(*args)

yolox/core/trainer.py

@@ -0,0 +1,390 @@
+#!/usr/bin/env python3
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import datetime
+import os
+import time
+from loguru import logger
+
+import torch
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.utils.tensorboard import SummaryWriter
+
+from yolox.data import DataPrefetcher
+from yolox.exp import Exp
+from yolox.utils import (
+ MeterBuffer,
+ ModelEMA,
+ WandbLogger,
+ adjust_status,
+ all_reduce_norm,
+ get_local_rank,
+ get_model_info,
+ get_rank,
+ get_world_size,
+ gpu_mem_usage,
+ is_parallel,
+ load_ckpt,
+ mem_usage,
+ occupy_mem,
+ save_checkpoint,
+ setup_logger,
+ synchronize
+)
+
+
+class Trainer:
+ def __init__(self, exp: Exp, args):
+ # init function only defines some basic attr, other attrs like model, optimizer are built in
+ # before_train methods.
+ self.exp = exp
+ self.args = args
+
+ # training related attr
+ self.max_epoch = exp.max_epoch
+ self.amp_training = args.fp16
+ self.scaler = torch.cuda.amp.GradScaler(enabled=args.fp16)
+ self.is_distributed = get_world_size() > 1
+ self.rank = get_rank()
+ self.local_rank = get_local_rank()
+ self.device = "cuda:{}".format(self.local_rank)
+ self.use_model_ema = exp.ema
+ self.save_history_ckpt = exp.save_history_ckpt
+
+ # data/dataloader related attr
+ self.data_type = torch.float16 if args.fp16 else torch.float32
+ self.input_size = exp.input_size
+ self.best_ap = 0
+
+ # metric record
+ self.meter = MeterBuffer(window_size=exp.print_interval)
+ self.file_name = os.path.join(exp.output_dir, args.experiment_name)
+
+ if self.rank == 0:
+ os.makedirs(self.file_name, exist_ok=True)
+
+ setup_logger(
+ self.file_name,
+ distributed_rank=self.rank,
+ filename="train_log.txt",
+ mode="a",
+ )
+
+ def train(self):
+ self.before_train()
+ try:
+ self.train_in_epoch()
+ except Exception:
+ raise
+ finally:
+ self.after_train()
+
+ def train_in_epoch(self):
+ for self.epoch in range(self.start_epoch, self.max_epoch):
+ self.before_epoch()
+ self.train_in_iter()
+ self.after_epoch()
+
+ def train_in_iter(self):
+ for self.iter in range(self.max_iter):
+ self.before_iter()
+ self.train_one_iter()
+ self.after_iter()
+
+ def train_one_iter(self):
+ iter_start_time = time.time()
+
+ inps, targets = self.prefetcher.next()
+ inps = inps.to(self.data_type)
+ targets = targets.to(self.data_type)
+ targets.requires_grad = False
+ inps, targets = self.exp.preprocess(inps, targets, self.input_size)
+ data_end_time = time.time()
+
+ with torch.cuda.amp.autocast(enabled=self.amp_training):
+ outputs = self.model(inps, targets)
+
+ loss = outputs["total_loss"]
+
+ self.optimizer.zero_grad()
+ self.scaler.scale(loss).backward()
+ self.scaler.step(self.optimizer)
+ self.scaler.update()
+
+ if self.use_model_ema:
+ self.ema_model.update(self.model)
+
+ lr = self.lr_scheduler.update_lr(self.progress_in_iter + 1)
+ for param_group in self.optimizer.param_groups:
+ param_group["lr"] = lr
+
+ iter_end_time = time.time()
+ self.meter.update(
+ iter_time=iter_end_time - iter_start_time,
+ data_time=data_end_time - iter_start_time,
+ lr=lr,
+ **outputs,
+ )
+
+ def before_train(self):
+ logger.info("args: {}".format(self.args))
+ logger.info("exp value:\n{}".format(self.exp))
+
+ # model related init
+ torch.cuda.set_device(self.local_rank)
+ model = self.exp.get_model()
+ logger.info(
+ "Model Summary: {}".format(get_model_info(model, self.exp.test_size))
+ )
+ model.to(self.device)
+
+ # solver related init
+ self.optimizer = self.exp.get_optimizer(self.args.batch_size)
+
+ # value of epoch will be set in `resume_train`
+ model = self.resume_train(model)
+
+ # data related init
+ self.no_aug = self.start_epoch >= self.max_epoch - self.exp.no_aug_epochs
+ self.train_loader = self.exp.get_data_loader(
+ batch_size=self.args.batch_size,
+ is_distributed=self.is_distributed,
+ no_aug=self.no_aug,
+ cache_img=self.args.cache,
+ )
+ logger.info("init prefetcher, this might take one minute or less...")
+ self.prefetcher = DataPrefetcher(self.train_loader)
+ # max_iter means iters per epoch
+ self.max_iter = len(self.train_loader)
+
+ self.lr_scheduler = self.exp.get_lr_scheduler(
+ self.exp.basic_lr_per_img * self.args.batch_size, self.max_iter
+ )
+ if self.args.occupy:
+ occupy_mem(self.local_rank)
+
+ if self.is_distributed:
+ model = DDP(model, device_ids=[self.local_rank], broadcast_buffers=False)
+
+ if self.use_model_ema:
+ self.ema_model = ModelEMA(model, 0.9998)
+ self.ema_model.updates = self.max_iter * self.start_epoch
+
+ self.model = model
+
+ self.evaluator = self.exp.get_evaluator(
+ batch_size=self.args.batch_size, is_distributed=self.is_distributed
+ )
+ # Tensorboard and Wandb loggers
+ if self.rank == 0:
+ if self.args.logger == "tensorboard":
+ self.tblogger = SummaryWriter(os.path.join(self.file_name, "tensorboard"))
+ elif self.args.logger == "wandb":
+ self.wandb_logger = WandbLogger.initialize_wandb_logger(
+ self.args,
+ self.exp,
+ self.evaluator.dataloader.dataset
+ )
+ else:
+ raise ValueError("logger must be either 'tensorboard' or 'wandb'")
+
+ logger.info("Training start...")
+ logger.info("\n{}".format(model))
+
+ def after_train(self):
+ logger.info(
+ "Training of experiment is done and the best AP is {:.2f}".format(self.best_ap * 100)
+ )
+ if self.rank == 0:
+ if self.args.logger == "wandb":
+ self.wandb_logger.finish()
+
+ def before_epoch(self):
+ logger.info("---> start train epoch{}".format(self.epoch + 1))
+
+ if self.epoch + 1 == self.max_epoch - self.exp.no_aug_epochs or self.no_aug:
+ logger.info("--->No mosaic aug now!")
+ self.train_loader.close_mosaic()
+ logger.info("--->Add additional L1 loss now!")
+ if self.is_distributed:
+ self.model.module.head.use_l1 = True
+ else:
+ self.model.head.use_l1 = True
+ self.exp.eval_interval = 1
+ if not self.no_aug:
+ self.save_ckpt(ckpt_name="last_mosaic_epoch")
+
+ def after_epoch(self):
+ self.save_ckpt(ckpt_name="latest")
+
+ if (self.epoch + 1) % self.exp.eval_interval == 0:
+ all_reduce_norm(self.model)
+ self.evaluate_and_save_model()
+
+ def before_iter(self):
+ pass
+
+ def after_iter(self):
+ """
+ `after_iter` contains two parts of logic:
+ * log information
+ * reset setting of resize
+ """
+ # log needed information
+ if (self.iter + 1) % self.exp.print_interval == 0:
+ # TODO check ETA logic
+ left_iters = self.max_iter * self.max_epoch - (self.progress_in_iter + 1)
+ eta_seconds = self.meter["iter_time"].global_avg * left_iters
+ eta_str = "ETA: {}".format(datetime.timedelta(seconds=int(eta_seconds)))
+
+ progress_str = "epoch: {}/{}, iter: {}/{}".format(
+ self.epoch + 1, self.max_epoch, self.iter + 1, self.max_iter
+ )
+ loss_meter = self.meter.get_filtered_meter("loss")
+ loss_str = ", ".join(
+ ["{}: {:.1f}".format(k, v.latest) for k, v in loss_meter.items()]
+ )
+
+ time_meter = self.meter.get_filtered_meter("time")
+ time_str = ", ".join(
+ ["{}: {:.3f}s".format(k, v.avg) for k, v in time_meter.items()]
+ )
+
+ mem_str = "gpu mem: {:.0f}Mb, mem: {:.1f}Gb".format(gpu_mem_usage(), mem_usage())
+
+ logger.info(
+ "{}, {}, {}, {}, lr: {:.3e}".format(
+ progress_str,
+ mem_str,
+ time_str,
+ loss_str,
+ self.meter["lr"].latest,
+ )
+ + (", size: {:d}, {}".format(self.input_size[0], eta_str))
+ )
+
+ if self.rank == 0:
+ if self.args.logger == "tensorboard":
+ self.tblogger.add_scalar(
+ "train/lr", self.meter["lr"].latest, self.progress_in_iter)
+ for k, v in loss_meter.items():
+ self.tblogger.add_scalar(
+ f"train/{k}", v.latest, self.progress_in_iter)
+ if self.args.logger == "wandb":
+ metrics = {"train/" + k: v.latest for k, v in loss_meter.items()}
+ metrics.update({
+ "train/lr": self.meter["lr"].latest
+ })
+ self.wandb_logger.log_metrics(metrics, step=self.progress_in_iter)
+
+ self.meter.clear_meters()
+
+ # random resizing
+ if (self.progress_in_iter + 1) % 10 == 0:
+ self.input_size = self.exp.random_resize(
+ self.train_loader, self.epoch, self.rank, self.is_distributed
+ )
+
+ @property
+ def progress_in_iter(self):
+ return self.epoch * self.max_iter + self.iter
+
+ def resume_train(self, model):
+ if self.args.resume:
+ logger.info("resume training")
+ if self.args.ckpt is None:
+ ckpt_file = os.path.join(self.file_name, "latest" + "_ckpt.pth")
+ else:
+ ckpt_file = self.args.ckpt
+
+ ckpt = torch.load(ckpt_file, map_location=self.device)
+ # resume the model/optimizer state dict
+ model.load_state_dict(ckpt["model"])
+ self.optimizer.load_state_dict(ckpt["optimizer"])
+ self.best_ap = ckpt.pop("best_ap", 0)
+ # resume the training states variables
+ start_epoch = (
+ self.args.start_epoch - 1
+ if self.args.start_epoch is not None
+ else ckpt["start_epoch"]
+ )
+ self.start_epoch = start_epoch
+ logger.info(
+ "loaded checkpoint '{}' (epoch {})".format(
+ self.args.resume, self.start_epoch
+ )
+ ) # noqa
+ else:
+ if self.args.ckpt is not None:
+ logger.info("loading checkpoint for fine tuning")
+ ckpt_file = self.args.ckpt
+ ckpt = torch.load(ckpt_file, map_location=self.device)["model"]
+ model = load_ckpt(model, ckpt)
+ self.start_epoch = 0
+
+ return model
+
+ def evaluate_and_save_model(self):
+ if self.use_model_ema:
+ evalmodel = self.ema_model.ema
+ else:
+ evalmodel = self.model
+ if is_parallel(evalmodel):
+ evalmodel = evalmodel.module
+
+ with adjust_status(evalmodel, training=False):
+ (ap50_95, ap50, summary), predictions = self.exp.eval(
+ evalmodel, self.evaluator, self.is_distributed, return_outputs=True
+ )
+
+ update_best_ckpt = ap50_95 > self.best_ap
+ self.best_ap = max(self.best_ap, ap50_95)
+
+ if self.rank == 0:
+ if self.args.logger == "tensorboard":
+ self.tblogger.add_scalar("val/COCOAP50", ap50, self.epoch + 1)
+ self.tblogger.add_scalar("val/COCOAP50_95", ap50_95, self.epoch + 1)
+ if self.args.logger == "wandb":
+ self.wandb_logger.log_metrics({
+ "val/COCOAP50": ap50,
+ "val/COCOAP50_95": ap50_95,
+ "train/epoch": self.epoch + 1,
+ })
+ self.wandb_logger.log_images(predictions)
+ logger.info("\n" + summary)
+ synchronize()
+
+ self.save_ckpt("last_epoch", update_best_ckpt, ap=ap50_95)
+ if self.save_history_ckpt:
+ self.save_ckpt(f"epoch_{self.epoch + 1}", ap=ap50_95)
+
+ def save_ckpt(self, ckpt_name, update_best_ckpt=False, ap=None):
+ if self.rank == 0:
+ save_model = self.ema_model.ema if self.use_model_ema else self.model
+ logger.info("Save weights to {}".format(self.file_name))
+ ckpt_state = {
+ "start_epoch": self.epoch + 1,
+ "model": save_model.state_dict(),
+ "optimizer": self.optimizer.state_dict(),
+ "best_ap": self.best_ap,
+ "curr_ap": ap,
+ }
+ save_checkpoint(
+ ckpt_state,
+ update_best_ckpt,
+ self.file_name,
+ ckpt_name,
+ )
+
+ if self.args.logger == "wandb":
+ self.wandb_logger.save_checkpoint(
+ self.file_name,
+ ckpt_name,
+ update_best_ckpt,
+ metadata={
+ "epoch": self.epoch + 1,
+ "optimizer": self.optimizer.state_dict(),
+ "best_ap": self.best_ap,
+ "curr_ap": ap
+ }
+ )

yolox/data/__init__.py

@@ -0,0 +1,9 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+from .data_augment import TrainTransform, ValTransform
+from .data_prefetcher import DataPrefetcher
+from .dataloading import DataLoader, get_yolox_datadir, worker_init_reset_seed
+from .datasets import *
+from .samplers import InfiniteSampler, YoloBatchSampler

yolox/data/data_augment.py

@@ -0,0 +1,243 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+"""
+Data augmentation functionality. Passed as callable transformations to
+Dataset classes.
+
+The data augmentation procedures were interpreted from @weiliu89's SSD paper
+http://arxiv.org/abs/1512.02325
+"""
+
+import math
+import random
+
+import cv2
+import numpy as np
+
+from yolox.utils import xyxy2cxcywh
+
+
+def augment_hsv(img, hgain=5, sgain=30, vgain=30):
+ hsv_augs = np.random.uniform(-1, 1, 3) * [hgain, sgain, vgain] # random gains
+ hsv_augs *= np.random.randint(0, 2, 3) # random selection of h, s, v
+ hsv_augs = hsv_augs.astype(np.int16)
+ img_hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV).astype(np.int16)
+
+ img_hsv[..., 0] = (img_hsv[..., 0] + hsv_augs[0]) % 180
+ img_hsv[..., 1] = np.clip(img_hsv[..., 1] + hsv_augs[1], 0, 255)
+ img_hsv[..., 2] = np.clip(img_hsv[..., 2] + hsv_augs[2], 0, 255)
+
+ cv2.cvtColor(img_hsv.astype(img.dtype), cv2.COLOR_HSV2BGR, dst=img) # no return needed
+
+
+def get_aug_params(value, center=0):
+ if isinstance(value, float):
+ return random.uniform(center - value, center + value)
+ elif len(value) == 2:
+ return random.uniform(value[0], value[1])
+ else:
+ raise ValueError(
+ "Affine params should be either a sequence containing two values\
+ or single float values. Got {}".format(value)
+ )
+
+
+def get_affine_matrix(
+ target_size,
+ degrees=10,
+ translate=0.1,
+ scales=0.1,
+ shear=10,
+):
+ twidth, theight = target_size
+
+ # Rotation and Scale
+ angle = get_aug_params(degrees)
+ scale = get_aug_params(scales, center=1.0)
+
+ if scale <= 0.0:
+ raise ValueError("Argument scale should be positive")
+
+ R = cv2.getRotationMatrix2D(angle=angle, center=(0, 0), scale=scale)
+
+ M = np.ones([2, 3])
+ # Shear
+ shear_x = math.tan(get_aug_params(shear) * math.pi / 180)
+ shear_y = math.tan(get_aug_params(shear) * math.pi / 180)
+
+ M[0] = R[0] + shear_y * R[1]
+ M[1] = R[1] + shear_x * R[0]
+
+ # Translation
+ translation_x = get_aug_params(translate) * twidth # x translation (pixels)
+ translation_y = get_aug_params(translate) * theight # y translation (pixels)
+
+ M[0, 2] = translation_x
+ M[1, 2] = translation_y
+
+ return M, scale
+
+
+def apply_affine_to_bboxes(targets, target_size, M, scale):
+ num_gts = len(targets)
+
+ # warp corner points
+ twidth, theight = target_size
+ corner_points = np.ones((4 * num_gts, 3))
+ corner_points[:, :2] = targets[:, [0, 1, 2, 3, 0, 3, 2, 1]].reshape(
+ 4 * num_gts, 2
+ ) # x1y1, x2y2, x1y2, x2y1
+ corner_points = corner_points @ M.T # apply affine transform
+ corner_points = corner_points.reshape(num_gts, 8)
+
+ # create new boxes
+ corner_xs = corner_points[:, 0::2]
+ corner_ys = corner_points[:, 1::2]
+ new_bboxes = (
+ np.concatenate(
+ (corner_xs.min(1), corner_ys.min(1), corner_xs.max(1), corner_ys.max(1))
+ )
+ .reshape(4, num_gts)
+ .T
+ )
+
+ # clip boxes
+ new_bboxes[:, 0::2] = new_bboxes[:, 0::2].clip(0, twidth)
+ new_bboxes[:, 1::2] = new_bboxes[:, 1::2].clip(0, theight)
+
+ targets[:, :4] = new_bboxes
+
+ return targets
+
+
+def random_affine(
+ img,
+ targets=(),
+ target_size=(640, 640),
+ degrees=10,
+ translate=0.1,
+ scales=0.1,
+ shear=10,
+):
+ M, scale = get_affine_matrix(target_size, degrees, translate, scales, shear)
+
+ img = cv2.warpAffine(img, M, dsize=target_size, borderValue=(114, 114, 114))
+
+ # Transform label coordinates
+ if len(targets) > 0:
+ targets = apply_affine_to_bboxes(targets, target_size, M, scale)
+
+ return img, targets
+
+
+def _mirror(image, boxes, prob=0.5):
+ _, width, _ = image.shape
+ if random.random() < prob:
+ image = image[:, ::-1]
+ boxes[:, 0::2] = width - boxes[:, 2::-2]
+ return image, boxes
+
+
+def preproc(img, input_size, swap=(2, 0, 1)):
+ if len(img.shape) == 3:
+ padded_img = np.ones((input_size[0], input_size[1], 3), dtype=np.uint8) * 114
+ else:
+ padded_img = np.ones(input_size, dtype=np.uint8) * 114
+
+ r = min(input_size[0] / img.shape[0], input_size[1] / img.shape[1])
+ resized_img = cv2.resize(
+ img,
+ (int(img.shape[1] * r), int(img.shape[0] * r)),
+ interpolation=cv2.INTER_LINEAR,
+ ).astype(np.uint8)
+ padded_img[: int(img.shape[0] * r), : int(img.shape[1] * r)] = resized_img
+
+ padded_img = padded_img.transpose(swap)
+ padded_img = np.ascontiguousarray(padded_img, dtype=np.float32)
+ return padded_img, r
+
+
+class TrainTransform:
+ def __init__(self, max_labels=50, flip_prob=0.5, hsv_prob=1.0):
+ self.max_labels = max_labels
+ self.flip_prob = flip_prob
+ self.hsv_prob = hsv_prob
+
+ def __call__(self, image, targets, input_dim):
+ boxes = targets[:, :4].copy()
+ labels = targets[:, 4].copy()
+ if len(boxes) == 0:
+ targets = np.zeros((self.max_labels, 5), dtype=np.float32)
+ image, r_o = preproc(image, input_dim)
+ return image, targets
+
+ image_o = image.copy()
+ targets_o = targets.copy()
+ height_o, width_o, _ = image_o.shape
+ boxes_o = targets_o[:, :4]
+ labels_o = targets_o[:, 4]
+ # bbox_o: [xyxy] to [c_x,c_y,w,h]
+ boxes_o = xyxy2cxcywh(boxes_o)
+
+ if random.random() < self.hsv_prob:
+ augment_hsv(image)
+ image_t, boxes = _mirror(image, boxes, self.flip_prob)
+ height, width, _ = image_t.shape
+ image_t, r_ = preproc(image_t, input_dim)
+ # boxes [xyxy] 2 [cx,cy,w,h]
+ boxes = xyxy2cxcywh(boxes)
+ boxes *= r_
+
+ mask_b = np.minimum(boxes[:, 2], boxes[:, 3]) > 1
+ boxes_t = boxes[mask_b]
+ labels_t = labels[mask_b]
+
+ if len(boxes_t) == 0:
+ image_t, r_o = preproc(image_o, input_dim)
+ boxes_o *= r_o
+ boxes_t = boxes_o
+ labels_t = labels_o
+
+ labels_t = np.expand_dims(labels_t, 1)
+
+ targets_t = np.hstack((labels_t, boxes_t))
+ padded_labels = np.zeros((self.max_labels, 5))
+ padded_labels[range(len(targets_t))[: self.max_labels]] = targets_t[
+ : self.max_labels
+ ]
+ padded_labels = np.ascontiguousarray(padded_labels, dtype=np.float32)
+ return image_t, padded_labels
+
+
+class ValTransform:
+ """
+ Defines the transformations that should be applied to test PIL image
+ for input into the network
+
+ dimension -> tensorize -> color adj
+
+ Arguments:
+ resize (int): input dimension to SSD
+ rgb_means ((int,int,int)): average RGB of the dataset
+ (104,117,123)
+ swap ((int,int,int)): final order of channels
+
+ Returns:
+ transform (transform) : callable transform to be applied to test/val
+ data
+ """
+
+ def __init__(self, swap=(2, 0, 1), legacy=False):
+ self.swap = swap
+ self.legacy = legacy
+
+ # assume input is cv2 img for now
+ def __call__(self, img, res, input_size):
+ img, _ = preproc(img, input_size, self.swap)
+ if self.legacy:
+ img = img[::-1, :, :].copy()
+ img /= 255.0
+ img -= np.array([0.485, 0.456, 0.406]).reshape(3, 1, 1)
+ img /= np.array([0.229, 0.224, 0.225]).reshape(3, 1, 1)
+ return img, np.zeros((1, 5))

yolox/data/data_prefetcher.py

@@ -0,0 +1,51 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import torch
+
+
+class DataPrefetcher:
+ """
+ DataPrefetcher is inspired by code of following file:
+ https://github.com/NVIDIA/apex/blob/master/examples/imagenet/main_amp.py
+ It could speedup your pytorch dataloader. For more information, please check
+ https://github.com/NVIDIA/apex/issues/304#issuecomment-493562789.
+ """
+
+ def __init__(self, loader):
+ self.loader = iter(loader)
+ self.stream = torch.cuda.Stream()
+ self.input_cuda = self._input_cuda_for_image
+ self.record_stream = DataPrefetcher._record_stream_for_image
+ self.preload()
+
+ def preload(self):
+ try:
+ self.next_input, self.next_target, _, _ = next(self.loader)
+ except StopIteration:
+ self.next_input = None
+ self.next_target = None
+ return
+
+ with torch.cuda.stream(self.stream):
+ self.input_cuda()
+ self.next_target = self.next_target.cuda(non_blocking=True)
+
+ def next(self):
+ torch.cuda.current_stream().wait_stream(self.stream)
+ input = self.next_input
+ target = self.next_target
+ if input is not None:
+ self.record_stream(input)
+ if target is not None:
+ target.record_stream(torch.cuda.current_stream())
+ self.preload()
+ return input, target
+
+ def _input_cuda_for_image(self):
+ self.next_input = self.next_input.cuda(non_blocking=True)
+
+ @staticmethod
+ def _record_stream_for_image(input):
+ input.record_stream(torch.cuda.current_stream())

yolox/data/dataloading.py

@@ -0,0 +1,113 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import os
+import random
+import uuid
+
+import numpy as np
+
+import torch
+from torch.utils.data.dataloader import DataLoader as torchDataLoader
+from torch.utils.data.dataloader import default_collate
+
+from .samplers import YoloBatchSampler
+
+
+def get_yolox_datadir():
+ """
+ get dataset dir of YOLOX. If environment variable named `YOLOX_DATADIR` is set,
+ this function will return value of the environment variable. Otherwise, use data
+ """
+ yolox_datadir = os.getenv("YOLOX_DATADIR", None)
+ if yolox_datadir is None:
+ import yolox
+
+ yolox_path = os.path.dirname(os.path.dirname(yolox.__file__))
+ yolox_datadir = os.path.join(yolox_path, "datasets")
+ return yolox_datadir
+
+
+class DataLoader(torchDataLoader):
+ """
+ Lightnet dataloader that enables on the fly resizing of the images.
+ See :class:`torch.utils.data.DataLoader` for more information on the arguments.
+ Check more on the following website:
+ https://gitlab.com/EAVISE/lightnet/-/blob/master/lightnet/data/_dataloading.py
+ """
+
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+ self.__initialized = False
+ shuffle = False
+ batch_sampler = None
+ if len(args) > 5:
+ shuffle = args[2]
+ sampler = args[3]
+ batch_sampler = args[4]
+ elif len(args) > 4:
+ shuffle = args[2]
+ sampler = args[3]
+ if "batch_sampler" in kwargs:
+ batch_sampler = kwargs["batch_sampler"]
+ elif len(args) > 3:
+ shuffle = args[2]
+ if "sampler" in kwargs:
+ sampler = kwargs["sampler"]
+ if "batch_sampler" in kwargs:
+ batch_sampler = kwargs["batch_sampler"]
+ else:
+ if "shuffle" in kwargs:
+ shuffle = kwargs["shuffle"]
+ if "sampler" in kwargs:
+ sampler = kwargs["sampler"]
+ if "batch_sampler" in kwargs:
+ batch_sampler = kwargs["batch_sampler"]
+
+ # Use custom BatchSampler
+ if batch_sampler is None:
+ if sampler is None:
+ if shuffle:
+ sampler = torch.utils.data.sampler.RandomSampler(self.dataset)
+ # sampler = torch.utils.data.DistributedSampler(self.dataset)
+ else:
+ sampler = torch.utils.data.sampler.SequentialSampler(self.dataset)
+ batch_sampler = YoloBatchSampler(
+ sampler,
+ self.batch_size,
+ self.drop_last,
+ input_dimension=self.dataset.input_dim,
+ )
+ # batch_sampler = IterationBasedBatchSampler(batch_sampler, num_iterations =
+
+ self.batch_sampler = batch_sampler
+
+ self.__initialized = True
+
+ def close_mosaic(self):
+ self.batch_sampler.mosaic = False
+
+
+def list_collate(batch):
+ """
+ Function that collates lists or tuples together into one list (of lists/tuples).
+ Use this as the collate function in a Dataloader, if you want to have a list of
+ items as an output, as opposed to tensors (eg. Brambox.boxes).
+ """
+ items = list(zip(*batch))
+
+ for i in range(len(items)):
+ if isinstance(items[i][0], (list, tuple)):
+ items[i] = list(items[i])
+ else:
+ items[i] = default_collate(items[i])
+
+ return items
+
+
+def worker_init_reset_seed(worker_id):
+ seed = uuid.uuid4().int % 2**32
+ random.seed(seed)
+ torch.set_rng_state(torch.manual_seed(seed).get_state())
+ np.random.seed(seed)

yolox/data/datasets/__init__.py

@@ -0,0 +1,9 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+from .coco import COCODataset
+from .coco_classes import COCO_CLASSES
+from .datasets_wrapper import CacheDataset, ConcatDataset, Dataset, MixConcatDataset
+from .mosaicdetection import MosaicDetection
+from .voc import VOCDetection

yolox/data/datasets/coco.py

@@ -0,0 +1,188 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+import copy
+import os
+
+import cv2
+import numpy as np
+from pycocotools.coco import COCO
+
+from ..dataloading import get_yolox_datadir
+from .datasets_wrapper import CacheDataset, cache_read_img
+
+
+def remove_useless_info(coco):
+ """
+ Remove useless info in coco dataset. COCO object is modified inplace.
+ This function is mainly used for saving memory (save about 30% mem).
+ """
+ if isinstance(coco, COCO):
+ dataset = coco.dataset
+ dataset.pop("info", None)
+ dataset.pop("licenses", None)
+ for img in dataset["images"]:
+ img.pop("license", None)
+ img.pop("coco_url", None)
+ img.pop("date_captured", None)
+ img.pop("flickr_url", None)
+ if "annotations" in coco.dataset:
+ for anno in coco.dataset["annotations"]:
+ anno.pop("segmentation", None)
+
+
+class COCODataset(CacheDataset):
+ """
+ COCO dataset class.
+ """
+
+ def __init__(
+ self,
+ data_dir=None,
+ json_file="instances_train2017.json",
+ name="train2017",
+ img_size=(416, 416),
+ preproc=None,
+ cache=False,
+ cache_type="ram",
+ ):
+ """
+ COCO dataset initialization. Annotation data are read into memory by COCO API.
+ Args:
+ data_dir (str): dataset root directory
+ json_file (str): COCO json file name
+ name (str): COCO data name (e.g. 'train2017' or 'val2017')
+ img_size (int): target image size after pre-processing
+ preproc: data augmentation strategy
+ """
+ if data_dir is None:
+ data_dir = os.path.join(get_yolox_datadir(), "COCO")
+ self.data_dir = data_dir
+ self.json_file = json_file
+
+ self.coco = COCO(os.path.join(self.data_dir, "annotations", self.json_file))
+ remove_useless_info(self.coco)
+ self.ids = self.coco.getImgIds()
+ self.num_imgs = len(self.ids)
+ self.class_ids = sorted(self.coco.getCatIds())
+ self.cats = self.coco.loadCats(self.coco.getCatIds())
+ self._classes = tuple([c["name"] for c in self.cats])
+ self.name = name
+ self.img_size = img_size
+ self.preproc = preproc
+ self.annotations = self._load_coco_annotations()
+
+ path_filename = [os.path.join(name, anno[3]) for anno in self.annotations]
+ super().__init__(
+ input_dimension=img_size,
+ num_imgs=self.num_imgs,
+ data_dir=data_dir,
+ cache_dir_name=f"cache_{name}",
+ path_filename=path_filename,
+ cache=cache,
+ cache_type=cache_type
+ )
+
+ def __len__(self):
+ return self.num_imgs
+
+ def _load_coco_annotations(self):
+ return [self.load_anno_from_ids(_ids) for _ids in self.ids]
+
+ def load_anno_from_ids(self, id_):
+ im_ann = self.coco.loadImgs(id_)[0]
+ width = im_ann["width"]
+ height = im_ann["height"]
+ anno_ids = self.coco.getAnnIds(imgIds=[int(id_)], iscrowd=False)
+ annotations = self.coco.loadAnns(anno_ids)
+ objs = []
+ for obj in annotations:
+ x1 = np.max((0, obj["bbox"][0]))
+ y1 = np.max((0, obj["bbox"][1]))
+ x2 = np.min((width, x1 + np.max((0, obj["bbox"][2]))))
+ y2 = np.min((height, y1 + np.max((0, obj["bbox"][3]))))
+ if obj["area"] > 0 and x2 >= x1 and y2 >= y1:
+ obj["clean_bbox"] = [x1, y1, x2, y2]
+ objs.append(obj)
+
+ num_objs = len(objs)
+
+ res = np.zeros((num_objs, 5))
+ for ix, obj in enumerate(objs):
+ cls = self.class_ids.index(obj["category_id"])
+ res[ix, 0:4] = obj["clean_bbox"]
+ res[ix, 4] = cls
+
+ r = min(self.img_size[0] / height, self.img_size[1] / width)
+ res[:, :4] *= r
+
+ img_info = (height, width)
+ resized_info = (int(height * r), int(width * r))
+
+ file_name = (
+ im_ann["file_name"]
+ if "file_name" in im_ann
+ else "{:012}".format(id_) + ".jpg"
+ )
+
+ return (res, img_info, resized_info, file_name)
+
+ def load_anno(self, index):
+ return self.annotations[index][0]
+
+ def load_resized_img(self, index):
+ img = self.load_image(index)
+ r = min(self.img_size[0] / img.shape[0], self.img_size[1] / img.shape[1])
+ resized_img = cv2.resize(
+ img,
+ (int(img.shape[1] * r), int(img.shape[0] * r)),
+ interpolation=cv2.INTER_LINEAR,
+ ).astype(np.uint8)
+ return resized_img
+
+ def load_image(self, index):
+ file_name = self.annotations[index][3]
+
+ img_file = os.path.join(self.data_dir, self.name, file_name)
+
+ img = cv2.imread(img_file)
+ assert img is not None, f"file named {img_file} not found"
+
+ return img
+
+ @cache_read_img(use_cache=True)
+ def read_img(self, index):
+ return self.load_resized_img(index)
+
+ def pull_item(self, index):
+ id_ = self.ids[index]
+ label, origin_image_size, _, _ = self.annotations[index]
+ img = self.read_img(index)
+
+ return img, copy.deepcopy(label), origin_image_size, np.array([id_])
+
+ @CacheDataset.mosaic_getitem
+ def __getitem__(self, index):
+ """
+ One image / label pair for the given index is picked up and pre-processed.
+
+ Args:
+ index (int): data index
+
+ Returns:
+ img (numpy.ndarray): pre-processed image
+ padded_labels (torch.Tensor): pre-processed label data.
+ The shape is :math:`[max_labels, 5]`.
+ each label consists of [class, xc, yc, w, h]:
+ class (float): class index.
+ xc, yc (float) : center of bbox whose values range from 0 to 1.
+ w, h (float) : size of bbox whose values range from 0 to 1.
+ info_img : tuple of h, w.
+ h, w (int): original shape of the image
+ img_id (int): same as the input index. Used for evaluation.
+ """
+ img, target, img_info, img_id = self.pull_item(index)
+
+ if self.preproc is not None:
+ img, target = self.preproc(img, target, self.input_dim)
+ return img, target, img_info, img_id

yolox/data/datasets/coco_classes.py

@@ -0,0 +1,5 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+COCO_CLASSES = ("red", "green", "yellow", "empty", "straight", "left", "right", "other")

yolox/data/datasets/datasets_wrapper.py

@@ -0,0 +1,300 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import bisect
+import copy
+import os
+import random
+from abc import ABCMeta, abstractmethod
+from functools import partial, wraps
+from multiprocessing.pool import ThreadPool
+import psutil
+from loguru import logger
+from tqdm import tqdm
+
+import numpy as np
+
+from torch.utils.data.dataset import ConcatDataset as torchConcatDataset
+from torch.utils.data.dataset import Dataset as torchDataset
+
+
+class ConcatDataset(torchConcatDataset):
+ def __init__(self, datasets):
+ super(ConcatDataset, self).__init__(datasets)
+ if hasattr(self.datasets[0], "input_dim"):
+ self._input_dim = self.datasets[0].input_dim
+ self.input_dim = self.datasets[0].input_dim
+
+ def pull_item(self, idx):
+ if idx < 0:
+ if -idx > len(self):
+ raise ValueError(
+ "absolute value of index should not exceed dataset length"
+ )
+ idx = len(self) + idx
+ dataset_idx = bisect.bisect_right(self.cumulative_sizes, idx)
+ if dataset_idx == 0:
+ sample_idx = idx
+ else:
+ sample_idx = idx - self.cumulative_sizes[dataset_idx - 1]
+ return self.datasets[dataset_idx].pull_item(sample_idx)
+
+
+class MixConcatDataset(torchConcatDataset):
+ def __init__(self, datasets):
+ super(MixConcatDataset, self).__init__(datasets)
+ if hasattr(self.datasets[0], "input_dim"):
+ self._input_dim = self.datasets[0].input_dim
+ self.input_dim = self.datasets[0].input_dim
+
+ def __getitem__(self, index):
+
+ if not isinstance(index, int):
+ idx = index[1]
+ if idx < 0:
+ if -idx > len(self):
+ raise ValueError(
+ "absolute value of index should not exceed dataset length"
+ )
+ idx = len(self) + idx
+ dataset_idx = bisect.bisect_right(self.cumulative_sizes, idx)
+ if dataset_idx == 0:
+ sample_idx = idx
+ else:
+ sample_idx = idx - self.cumulative_sizes[dataset_idx - 1]
+ if not isinstance(index, int):
+ index = (index[0], sample_idx, index[2])
+
+ return self.datasets[dataset_idx][index]
+
+
+class Dataset(torchDataset):
+ """ This class is a subclass of the base :class:`torch.utils.data.Dataset`,
+ that enables on the fly resizing of the ``input_dim``.
+
+ Args:
+ input_dimension (tuple): (width,height) tuple with default dimensions of the network
+ """
+
+ def __init__(self, input_dimension, mosaic=True):
+ super().__init__()
+ self.__input_dim = input_dimension[:2]
+ self.enable_mosaic = mosaic
+
+ @property
+ def input_dim(self):
+ """
+ Dimension that can be used by transforms to set the correct image size, etc.
+ This allows transforms to have a single source of truth
+ for the input dimension of the network.
+
+ Return:
+ list: Tuple containing the current width,height
+ """
+ if hasattr(self, "_input_dim"):
+ return self._input_dim
+ return self.__input_dim
+
+ @staticmethod
+ def mosaic_getitem(getitem_fn):
+ """
+ Decorator method that needs to be used around the ``__getitem__`` method. |br|
+ This decorator enables the closing mosaic
+
+ Example:
+ >>> class CustomSet(ln.data.Dataset):
+ ... def __len__(self):
+ ... return 10
+ ... @ln.data.Dataset.mosaic_getitem
+ ... def __getitem__(self, index):
+ ... return self.enable_mosaic
+ """
+
+ @wraps(getitem_fn)
+ def wrapper(self, index):
+ if not isinstance(index, int):
+ self.enable_mosaic = index[0]
+ index = index[1]
+
+ ret_val = getitem_fn(self, index)
+
+ return ret_val
+
+ return wrapper
+
+
+class CacheDataset(Dataset, metaclass=ABCMeta):
+ """ This class is a subclass of the base :class:`yolox.data.datasets.Dataset`,
+ that enables cache images to ram or disk.
+
+ Args:
+ input_dimension (tuple): (width,height) tuple with default dimensions of the network
+ num_imgs (int): datset size
+ data_dir (str): the root directory of the dataset, e.g. `/path/to/COCO`.
+ cache_dir_name (str): the name of the directory to cache to disk,
+ e.g. `"custom_cache"`. The files cached to disk will be saved
+ under `/path/to/COCO/custom_cache`.
+ path_filename (str): a list of paths to the data relative to the `data_dir`,
+ e.g. if you have data `/path/to/COCO/train/1.jpg`, `/path/to/COCO/train/2.jpg`,
+ then `path_filename = ['train/1.jpg', ' train/2.jpg']`.
+ cache (bool): whether to cache the images to ram or disk.
+ cache_type (str): the type of cache,
+ "ram" : Caching imgs to ram for fast training.
+ "disk": Caching imgs to disk for fast training.
+ """
+
+ def __init__(
+ self,
+ input_dimension,
+ num_imgs=None,
+ data_dir=None,
+ cache_dir_name=None,
+ path_filename=None,
+ cache=False,
+ cache_type="ram",
+ ):
+ super().__init__(input_dimension)
+ self.cache = cache
+ self.cache_type = cache_type
+
+ if self.cache and self.cache_type == "disk":
+ self.cache_dir = os.path.join(data_dir, cache_dir_name)
+ self.path_filename = path_filename
+
+ if self.cache and self.cache_type == "ram":
+ self.imgs = None
+
+ if self.cache:
+ self.cache_images(
+ num_imgs=num_imgs,
+ data_dir=data_dir,
+ cache_dir_name=cache_dir_name,
+ path_filename=path_filename,
+ )
+
+ def __del__(self):
+ if self.cache and self.cache_type == "ram":
+ del self.imgs
+
+ @abstractmethod
+ def read_img(self, index):
+ """
+ Given index, return the corresponding image
+
+ Args:
+ index (int): image index
+ """
+ raise NotImplementedError
+
+ def cache_images(
+ self,
+ num_imgs=None,
+ data_dir=None,
+ cache_dir_name=None,
+ path_filename=None,
+ ):
+ assert num_imgs is not None, "num_imgs must be specified as the size of the dataset"
+ if self.cache_type == "disk":
+ assert (data_dir and cache_dir_name and path_filename) is not None, \
+ "data_dir, cache_name and path_filename must be specified if cache_type is disk"
+ self.path_filename = path_filename
+
+ mem = psutil.virtual_memory()
+ mem_required = self.cal_cache_occupy(num_imgs)
+ gb = 1 << 30
+
+ if self.cache_type == "ram":
+ if mem_required > mem.available:
+ self.cache = False
+ else:
+ logger.info(
+ f"{mem_required / gb:.1f}GB RAM required, "
+ f"{mem.available / gb:.1f}/{mem.total / gb:.1f}GB RAM available, "
+ f"Since the first thing we do is cache, "
+ f"there is no guarantee that the remaining memory space is sufficient"
+ )
+
+ if self.cache and self.imgs is None:
+ if self.cache_type == 'ram':
+ self.imgs = [None] * num_imgs
+ logger.info("You are using cached images in RAM to accelerate training!")
+ else: # 'disk'
+ if not os.path.exists(self.cache_dir):
+ os.mkdir(self.cache_dir)
+ logger.warning(
+ f"\n*******************************************************************\n"
+ f"You are using cached images in DISK to accelerate training.\n"
+ f"This requires large DISK space.\n"
+ f"Make sure you have {mem_required / gb:.1f} "
+ f"available DISK space for training your dataset.\n"
+ f"*******************************************************************\\n"
+ )
+ else:
+ logger.info(f"Found disk cache at {self.cache_dir}")
+ return
+
+ logger.info(
+ "Caching images...\n"
+ "This might take some time for your dataset"
+ )
+
+ num_threads = min(8, max(1, os.cpu_count() - 1))
+ b = 0
+ load_imgs = ThreadPool(num_threads).imap(
+ partial(self.read_img, use_cache=False),
+ range(num_imgs)
+ )
+ pbar = tqdm(enumerate(load_imgs), total=num_imgs)
+ for i, x in pbar: # x = self.read_img(self, i, use_cache=False)
+ if self.cache_type == 'ram':
+ self.imgs[i] = x
+ else: # 'disk'
+ cache_filename = f'{self.path_filename[i].split(".")[0]}.npy'
+ cache_path_filename = os.path.join(self.cache_dir, cache_filename)
+ os.makedirs(os.path.dirname(cache_path_filename), exist_ok=True)
+ np.save(cache_path_filename, x)
+ b += x.nbytes
+ pbar.desc = \
+ f'Caching images ({b / gb:.1f}/{mem_required / gb:.1f}GB {self.cache_type})'
+ pbar.close()
+
+ def cal_cache_occupy(self, num_imgs):
+ cache_bytes = 0
+ num_samples = min(num_imgs, 32)
+ for _ in range(num_samples):
+ img = self.read_img(index=random.randint(0, num_imgs - 1), use_cache=False)
+ cache_bytes += img.nbytes
+ mem_required = cache_bytes * num_imgs / num_samples
+ return mem_required
+
+
+def cache_read_img(use_cache=True):
+ def decorator(read_img_fn):
+ """
+ Decorate the read_img function to cache the image
+
+ Args:
+ read_img_fn: read_img function
+ use_cache (bool, optional): For the decorated read_img function,
+ whether to read the image from cache.
+ Defaults to True.
+ """
+ @wraps(read_img_fn)
+ def wrapper(self, index, use_cache=use_cache):
+ cache = self.cache and use_cache
+ if cache:
+ if self.cache_type == "ram":
+ img = self.imgs[index]
+ img = copy.deepcopy(img)
+ elif self.cache_type == "disk":
+ img = np.load(
+ os.path.join(
+ self.cache_dir, f"{self.path_filename[index].split('.')[0]}.npy"))
+ else:
+ raise ValueError(f"Unknown cache type: {self.cache_type}")
+ else:
+ img = read_img_fn(self, index)
+ return img
+ return wrapper
+ return decorator

yolox/data/datasets/mosaicdetection.py

@@ -0,0 +1,234 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import random
+
+import cv2
+import numpy as np
+
+from yolox.utils import adjust_box_anns, get_local_rank
+
+from ..data_augment import random_affine
+from .datasets_wrapper import Dataset
+
+
+def get_mosaic_coordinate(mosaic_image, mosaic_index, xc, yc, w, h, input_h, input_w):
+ # TODO update doc
+ # index0 to top left part of image
+ if mosaic_index == 0:
+ x1, y1, x2, y2 = max(xc - w, 0), max(yc - h, 0), xc, yc
+ small_coord = w - (x2 - x1), h - (y2 - y1), w, h
+ # index1 to top right part of image
+ elif mosaic_index == 1:
+ x1, y1, x2, y2 = xc, max(yc - h, 0), min(xc + w, input_w * 2), yc
+ small_coord = 0, h - (y2 - y1), min(w, x2 - x1), h
+ # index2 to bottom left part of image
+ elif mosaic_index == 2:
+ x1, y1, x2, y2 = max(xc - w, 0), yc, xc, min(input_h * 2, yc + h)
+ small_coord = w - (x2 - x1), 0, w, min(y2 - y1, h)
+ # index2 to bottom right part of image
+ elif mosaic_index == 3:
+ x1, y1, x2, y2 = xc, yc, min(xc + w, input_w * 2), min(input_h * 2, yc + h) # noqa
+ small_coord = 0, 0, min(w, x2 - x1), min(y2 - y1, h)
+ return (x1, y1, x2, y2), small_coord
+
+
+class MosaicDetection(Dataset):
+ """Detection dataset wrapper that performs mixup for normal dataset."""
+
+ def __init__(
+ self, dataset, img_size, mosaic=True, preproc=None,
+ degrees=10.0, translate=0.1, mosaic_scale=(0.5, 1.5),
+ mixup_scale=(0.5, 1.5), shear=2.0, enable_mixup=True,
+ mosaic_prob=1.0, mixup_prob=1.0, *args
+ ):
+ """
+
+ Args:
+ dataset(Dataset) : Pytorch dataset object.
+ img_size (tuple):
+ mosaic (bool): enable mosaic augmentation or not.
+ preproc (func):
+ degrees (float):
+ translate (float):
+ mosaic_scale (tuple):
+ mixup_scale (tuple):
+ shear (float):
+ enable_mixup (bool):
+ *args(tuple) : Additional arguments for mixup random sampler.
+ """
+ super().__init__(img_size, mosaic=mosaic)
+ self._dataset = dataset
+ self.preproc = preproc
+ self.degrees = degrees
+ self.translate = translate
+ self.scale = mosaic_scale
+ self.shear = shear
+ self.mixup_scale = mixup_scale
+ self.enable_mosaic = mosaic
+ self.enable_mixup = enable_mixup
+ self.mosaic_prob = mosaic_prob
+ self.mixup_prob = mixup_prob
+ self.local_rank = get_local_rank()
+
+ def __len__(self):
+ return len(self._dataset)
+
+ @Dataset.mosaic_getitem
+ def __getitem__(self, idx):
+ if self.enable_mosaic and random.random() < self.mosaic_prob:
+ mosaic_labels = []
+ input_dim = self._dataset.input_dim
+ input_h, input_w = input_dim[0], input_dim[1]
+
+ # yc, xc = s, s # mosaic center x, y
+ yc = int(random.uniform(0.5 * input_h, 1.5 * input_h))
+ xc = int(random.uniform(0.5 * input_w, 1.5 * input_w))
+
+ # 3 additional image indices
+ indices = [idx] + [random.randint(0, len(self._dataset) - 1) for _ in range(3)]
+
+ for i_mosaic, index in enumerate(indices):
+ img, _labels, _, img_id = self._dataset.pull_item(index)
+ h0, w0 = img.shape[:2] # orig hw
+ scale = min(1. * input_h / h0, 1. * input_w / w0)
+ img = cv2.resize(
+ img, (int(w0 * scale), int(h0 * scale)), interpolation=cv2.INTER_LINEAR
+ )
+ # generate output mosaic image
+ (h, w, c) = img.shape[:3]
+ if i_mosaic == 0:
+ mosaic_img = np.full((input_h * 2, input_w * 2, c), 114, dtype=np.uint8)
+
+ # suffix l means large image, while s means small image in mosaic aug.
+ (l_x1, l_y1, l_x2, l_y2), (s_x1, s_y1, s_x2, s_y2) = get_mosaic_coordinate(
+ mosaic_img, i_mosaic, xc, yc, w, h, input_h, input_w
+ )
+
+ mosaic_img[l_y1:l_y2, l_x1:l_x2] = img[s_y1:s_y2, s_x1:s_x2]
+ padw, padh = l_x1 - s_x1, l_y1 - s_y1
+
+ labels = _labels.copy()
+ # Normalized xywh to pixel xyxy format
+ if _labels.size > 0:
+ labels[:, 0] = scale * _labels[:, 0] + padw
+ labels[:, 1] = scale * _labels[:, 1] + padh
+ labels[:, 2] = scale * _labels[:, 2] + padw
+ labels[:, 3] = scale * _labels[:, 3] + padh
+ mosaic_labels.append(labels)
+
+ if len(mosaic_labels):
+ mosaic_labels = np.concatenate(mosaic_labels, 0)
+ np.clip(mosaic_labels[:, 0], 0, 2 * input_w, out=mosaic_labels[:, 0])
+ np.clip(mosaic_labels[:, 1], 0, 2 * input_h, out=mosaic_labels[:, 1])
+ np.clip(mosaic_labels[:, 2], 0, 2 * input_w, out=mosaic_labels[:, 2])
+ np.clip(mosaic_labels[:, 3], 0, 2 * input_h, out=mosaic_labels[:, 3])
+
+ mosaic_img, mosaic_labels = random_affine(
+ mosaic_img,
+ mosaic_labels,
+ target_size=(input_w, input_h),
+ degrees=self.degrees,
+ translate=self.translate,
+ scales=self.scale,
+ shear=self.shear,
+ )
+
+ # -----------------------------------------------------------------
+ # CopyPaste: https://arxiv.org/abs/2012.07177
+ # -----------------------------------------------------------------
+ if (
+ self.enable_mixup
+ and not len(mosaic_labels) == 0
+ and random.random() < self.mixup_prob
+ ):
+ mosaic_img, mosaic_labels = self.mixup(mosaic_img, mosaic_labels, self.input_dim)
+ mix_img, padded_labels = self.preproc(mosaic_img, mosaic_labels, self.input_dim)
+ img_info = (mix_img.shape[1], mix_img.shape[0])
+
+ # -----------------------------------------------------------------
+ # img_info and img_id are not used for training.
+ # They are also hard to be specified on a mosaic image.
+ # -----------------------------------------------------------------
+ return mix_img, padded_labels, img_info, img_id
+
+ else:
+ self._dataset._input_dim = self.input_dim
+ img, label, img_info, img_id = self._dataset.pull_item(idx)
+ img, label = self.preproc(img, label, self.input_dim)
+ return img, label, img_info, img_id
+
+ def mixup(self, origin_img, origin_labels, input_dim):
+ jit_factor = random.uniform(*self.mixup_scale)
+ FLIP = random.uniform(0, 1) > 0.5
+ cp_labels = []
+ while len(cp_labels) == 0:
+ cp_index = random.randint(0, self.__len__() - 1)
+ cp_labels = self._dataset.load_anno(cp_index)
+ img, cp_labels, _, _ = self._dataset.pull_item(cp_index)
+
+ if len(img.shape) == 3:
+ cp_img = np.ones((input_dim[0], input_dim[1], 3), dtype=np.uint8) * 114
+ else:
+ cp_img = np.ones(input_dim, dtype=np.uint8) * 114
+
+ cp_scale_ratio = min(input_dim[0] / img.shape[0], input_dim[1] / img.shape[1])
+ resized_img = cv2.resize(
+ img,
+ (int(img.shape[1] * cp_scale_ratio), int(img.shape[0] * cp_scale_ratio)),
+ interpolation=cv2.INTER_LINEAR,
+ )
+
+ cp_img[
+ : int(img.shape[0] * cp_scale_ratio), : int(img.shape[1] * cp_scale_ratio)
+ ] = resized_img
+
+ cp_img = cv2.resize(
+ cp_img,
+ (int(cp_img.shape[1] * jit_factor), int(cp_img.shape[0] * jit_factor)),
+ )
+ cp_scale_ratio *= jit_factor
+
+ if FLIP:
+ cp_img = cp_img[:, ::-1, :]
+
+ origin_h, origin_w = cp_img.shape[:2]
+ target_h, target_w = origin_img.shape[:2]
+ padded_img = np.zeros(
+ (max(origin_h, target_h), max(origin_w, target_w), 3), dtype=np.uint8
+ )
+ padded_img[:origin_h, :origin_w] = cp_img
+
+ x_offset, y_offset = 0, 0
+ if padded_img.shape[0] > target_h:
+ y_offset = random.randint(0, padded_img.shape[0] - target_h - 1)
+ if padded_img.shape[1] > target_w:
+ x_offset = random.randint(0, padded_img.shape[1] - target_w - 1)
+ padded_cropped_img = padded_img[
+ y_offset: y_offset + target_h, x_offset: x_offset + target_w
+ ]
+
+ cp_bboxes_origin_np = adjust_box_anns(
+ cp_labels[:, :4].copy(), cp_scale_ratio, 0, 0, origin_w, origin_h
+ )
+ if FLIP:
+ cp_bboxes_origin_np[:, 0::2] = (
+ origin_w - cp_bboxes_origin_np[:, 0::2][:, ::-1]
+ )
+ cp_bboxes_transformed_np = cp_bboxes_origin_np.copy()
+ cp_bboxes_transformed_np[:, 0::2] = np.clip(
+ cp_bboxes_transformed_np[:, 0::2] - x_offset, 0, target_w
+ )
+ cp_bboxes_transformed_np[:, 1::2] = np.clip(
+ cp_bboxes_transformed_np[:, 1::2] - y_offset, 0, target_h
+ )
+
+ cls_labels = cp_labels[:, 4:5].copy()
+ box_labels = cp_bboxes_transformed_np
+ labels = np.hstack((box_labels, cls_labels))
+ origin_labels = np.vstack((origin_labels, labels))
+ origin_img = origin_img.astype(np.float32)
+ origin_img = 0.5 * origin_img + 0.5 * padded_cropped_img.astype(np.float32)
+
+ return origin_img.astype(np.uint8), origin_labels

yolox/data/datasets/voc.py

@@ -0,0 +1,331 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Code are based on
+# https://github.com/fmassa/vision/blob/voc_dataset/torchvision/datasets/voc.py
+# Copyright (c) Francisco Massa.
+# Copyright (c) Ellis Brown, Max deGroot.
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import os
+import os.path
+import pickle
+import xml.etree.ElementTree as ET
+
+import cv2
+import numpy as np
+
+from yolox.evaluators.voc_eval import voc_eval
+
+from .datasets_wrapper import CacheDataset, cache_read_img
+from .voc_classes import VOC_CLASSES
+
+
+class AnnotationTransform(object):
+
+ """Transforms a VOC annotation into a Tensor of bbox coords and label index
+ Initilized with a dictionary lookup of classnames to indexes
+
+ Arguments:
+ class_to_ind (dict, optional): dictionary lookup of classnames -> indexes
+ (default: alphabetic indexing of VOC's 20 classes)
+ keep_difficult (bool, optional): keep difficult instances or not
+ (default: False)
+ height (int): height
+ width (int): width
+ """
+
+ def __init__(self, class_to_ind=None, keep_difficult=True):
+ self.class_to_ind = class_to_ind or dict(
+ zip(VOC_CLASSES, range(len(VOC_CLASSES)))
+ )
+ self.keep_difficult = keep_difficult
+
+ def __call__(self, target):
+ """
+ Arguments:
+ target (annotation) : the target annotation to be made usable
+ will be an ET.Element
+ Returns:
+ a list containing lists of bounding boxes [bbox coords, class name]
+ """
+ res = np.empty((0, 5))
+ for obj in target.iter("object"):
+ difficult = obj.find("difficult")
+ if difficult is not None:
+ difficult = int(difficult.text) == 1
+ else:
+ difficult = False
+ if not self.keep_difficult and difficult:
+ continue
+ name = obj.find("name").text.strip()
+ bbox = obj.find("bndbox")
+
+ pts = ["xmin", "ymin", "xmax", "ymax"]
+ bndbox = []
+ for i, pt in enumerate(pts):
+ cur_pt = int(float(bbox.find(pt).text)) - 1
+ # scale height or width
+ # cur_pt = cur_pt / width if i % 2 == 0 else cur_pt / height
+ bndbox.append(cur_pt)
+ label_idx = self.class_to_ind[name]
+ bndbox.append(label_idx)
+ res = np.vstack((res, bndbox)) # [xmin, ymin, xmax, ymax, label_ind]
+ # img_id = target.find('filename').text[:-4]
+
+ width = int(target.find("size").find("width").text)
+ height = int(target.find("size").find("height").text)
+ img_info = (height, width)
+
+ return res, img_info
+
+
+class VOCDetection(CacheDataset):
+
+ """
+ VOC Detection Dataset Object
+
+ input is image, target is annotation
+
+ Args:
+ root (string): filepath to VOCdevkit folder.
+ image_set (string): imageset to use (eg. 'train', 'val', 'test')
+ transform (callable, optional): transformation to perform on the
+ input image
+ target_transform (callable, optional): transformation to perform on the
+ target `annotation`
+ (eg: take in caption string, return tensor of word indices)
+ dataset_name (string, optional): which dataset to load
+ (default: 'VOC2007')
+ """
+
+ def __init__(
+ self,
+ data_dir,
+ image_sets=[("2007", "trainval"), ("2012", "trainval")],
+ img_size=(416, 416),
+ preproc=None,
+ target_transform=AnnotationTransform(),
+ dataset_name="VOC0712",
+ cache=False,
+ cache_type="ram",
+ ):
+ self.root = data_dir
+ self.image_set = image_sets
+ self.img_size = img_size
+ self.preproc = preproc
+ self.target_transform = target_transform
+ self.name = dataset_name
+ self._annopath = os.path.join("%s", "Annotations", "%s.xml")
+ self._imgpath = os.path.join("%s", "JPEGImages", "%s.jpg")
+ self._classes = VOC_CLASSES
+ self.cats = [
+ {"id": idx, "name": val} for idx, val in enumerate(VOC_CLASSES)
+ ]
+ self.class_ids = list(range(len(VOC_CLASSES)))
+ self.ids = list()
+ for (year, name) in image_sets:
+ self._year = year
+ rootpath = os.path.join(self.root, "VOC" + year)
+ for line in open(
+ os.path.join(rootpath, "ImageSets", "Main", name + ".txt")
+ ):
+ self.ids.append((rootpath, line.strip()))
+ self.num_imgs = len(self.ids)
+
+ self.annotations = self._load_coco_annotations()
+
+ path_filename = [
+ (self._imgpath % self.ids[i]).split(self.root + "/")[1]
+ for i in range(self.num_imgs)
+ ]
+ super().__init__(
+ input_dimension=img_size,
+ num_imgs=self.num_imgs,
+ data_dir=self.root,
+ cache_dir_name=f"cache_{self.name}",
+ path_filename=path_filename,
+ cache=cache,
+ cache_type=cache_type
+ )
+
+ def __len__(self):
+ return self.num_imgs
+
+ def _load_coco_annotations(self):
+ return [self.load_anno_from_ids(_ids) for _ids in range(self.num_imgs)]
+
+ def load_anno_from_ids(self, index):
+ img_id = self.ids[index]
+ target = ET.parse(self._annopath % img_id).getroot()
+
+ assert self.target_transform is not None
+ res, img_info = self.target_transform(target)
+ height, width = img_info
+
+ r = min(self.img_size[0] / height, self.img_size[1] / width)
+ res[:, :4] *= r
+ resized_info = (int(height * r), int(width * r))
+
+ return (res, img_info, resized_info)
+
+ def load_anno(self, index):
+ return self.annotations[index][0]
+
+ def load_resized_img(self, index):
+ img = self.load_image(index)
+ r = min(self.img_size[0] / img.shape[0], self.img_size[1] / img.shape[1])
+ resized_img = cv2.resize(
+ img,
+ (int(img.shape[1] * r), int(img.shape[0] * r)),
+ interpolation=cv2.INTER_LINEAR,
+ ).astype(np.uint8)
+
+ return resized_img
+
+ def load_image(self, index):
+ img_id = self.ids[index]
+ img = cv2.imread(self._imgpath % img_id, cv2.IMREAD_COLOR)
+ assert img is not None, f"file named {self._imgpath % img_id} not found"
+
+ return img
+
+ @cache_read_img(use_cache=True)
+ def read_img(self, index):
+ return self.load_resized_img(index)
+
+ def pull_item(self, index):
+ """Returns the original image and target at an index for mixup
+
+ Note: not using self.__getitem__(), as any transformations passed in
+ could mess up this functionality.
+
+ Argument:
+ index (int): index of img to show
+ Return:
+ img, target
+ """
+ target, img_info, _ = self.annotations[index]
+ img = self.read_img(index)
+
+ return img, target, img_info, index
+
+ @CacheDataset.mosaic_getitem
+ def __getitem__(self, index):
+ img, target, img_info, img_id = self.pull_item(index)
+
+ if self.preproc is not None:
+ img, target = self.preproc(img, target, self.input_dim)
+
+ return img, target, img_info, img_id
+
+ def evaluate_detections(self, all_boxes, output_dir=None):
+ """
+ all_boxes is a list of length number-of-classes.
+ Each list element is a list of length number-of-images.
+ Each of those list elements is either an empty list []
+ or a numpy array of detection.
+
+ all_boxes[class][image] = [] or np.array of shape #dets x 5
+ """
+ self._write_voc_results_file(all_boxes)
+ IouTh = np.linspace(
+ 0.5, 0.95, int(np.round((0.95 - 0.5) / 0.05)) + 1, endpoint=True
+ )
+ mAPs = []
+ for iou in IouTh:
+ mAP = self._do_python_eval(output_dir, iou)
+ mAPs.append(mAP)
+
+ print("--------------------------------------------------------------")
+ print("map_5095:", np.mean(mAPs))
+ print("map_50:", mAPs[0])
+ print("--------------------------------------------------------------")
+ return np.mean(mAPs), mAPs[0]
+
+ def _get_voc_results_file_template(self):
+ filename = "comp4_det_test" + "_{:s}.txt"
+ filedir = os.path.join(self.root, "results", "VOC" + self._year, "Main")
+ if not os.path.exists(filedir):
+ os.makedirs(filedir)
+ path = os.path.join(filedir, filename)
+ return path
+
+ def _write_voc_results_file(self, all_boxes):
+ for cls_ind, cls in enumerate(VOC_CLASSES):
+ cls_ind = cls_ind
+ if cls == "__background__":
+ continue
+ print("Writing {} VOC results file".format(cls))
+ filename = self._get_voc_results_file_template().format(cls)
+ with open(filename, "wt") as f:
+ for im_ind, index in enumerate(self.ids):
+ index = index[1]
+ dets = all_boxes[cls_ind][im_ind]
+ if dets == []:
+ continue
+ for k in range(dets.shape[0]):
+ f.write(
+ "{:s} {:.3f} {:.1f} {:.1f} {:.1f} {:.1f}\n".format(
+ index,
+ dets[k, -1],
+ dets[k, 0] + 1,
+ dets[k, 1] + 1,
+ dets[k, 2] + 1,
+ dets[k, 3] + 1,
+ )
+ )
+
+ def _do_python_eval(self, output_dir="output", iou=0.5):
+ rootpath = os.path.join(self.root, "VOC" + self._year)
+ name = self.image_set[0][1]
+ annopath = os.path.join(rootpath, "Annotations", "{:s}.xml")
+ imagesetfile = os.path.join(rootpath, "ImageSets", "Main", name + ".txt")
+ cachedir = os.path.join(
+ self.root, "annotations_cache", "VOC" + self._year, name
+ )
+ if not os.path.exists(cachedir):
+ os.makedirs(cachedir)
+ aps = []
+ # The PASCAL VOC metric changed in 2010
+ use_07_metric = True if int(self._year) < 2010 else False
+ print("Eval IoU : {:.2f}".format(iou))
+ if output_dir is not None and not os.path.isdir(output_dir):
+ os.mkdir(output_dir)
+ for i, cls in enumerate(VOC_CLASSES):
+
+ if cls == "__background__":
+ continue
+
+ filename = self._get_voc_results_file_template().format(cls)
+ rec, prec, ap = voc_eval(
+ filename,
+ annopath,
+ imagesetfile,
+ cls,
+ cachedir,
+ ovthresh=iou,
+ use_07_metric=use_07_metric,
+ )
+ aps += [ap]
+ if iou == 0.5:
+ print("AP for {} = {:.4f}".format(cls, ap))
+ if output_dir is not None:
+ with open(os.path.join(output_dir, cls + "_pr.pkl"), "wb") as f:
+ pickle.dump({"rec": rec, "prec": prec, "ap": ap}, f)
+ if iou == 0.5:
+ print("Mean AP = {:.4f}".format(np.mean(aps)))
+ print("~~~~~~~~")
+ print("Results:")
+ for ap in aps:
+ print("{:.3f}".format(ap))
+ print("{:.3f}".format(np.mean(aps)))
+ print("~~~~~~~~")
+ print("")
+ print("--------------------------------------------------------------")
+ print("Results computed with the **unofficial** Python eval code.")
+ print("Results should be very close to the official MATLAB eval code.")
+ print("Recompute with `./tools/reval.py --matlab ...` for your paper.")
+ print("-- Thanks, The Management")
+ print("--------------------------------------------------------------")
+
+ return np.mean(aps)

yolox/data/datasets/voc_classes.py

@@ -0,0 +1,27 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+# VOC_CLASSES = ( '__background__', # always index 0
+VOC_CLASSES = (
+ "aeroplane",
+ "bicycle",
+ "bird",
+ "boat",
+ "bottle",
+ "bus",
+ "car",
+ "cat",
+ "chair",
+ "cow",
+ "diningtable",
+ "dog",
+ "horse",
+ "motorbike",
+ "person",
+ "pottedplant",
+ "sheep",
+ "sofa",
+ "train",
+ "tvmonitor",
+)

yolox/data/samplers.py

@@ -0,0 +1,85 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import itertools
+from typing import Optional
+
+import torch
+import torch.distributed as dist
+from torch.utils.data.sampler import BatchSampler as torchBatchSampler
+from torch.utils.data.sampler import Sampler
+
+
+class YoloBatchSampler(torchBatchSampler):
+ """
+ This batch sampler will generate mini-batches of (mosaic, index) tuples from another sampler.
+ It works just like the :class:`torch.utils.data.sampler.BatchSampler`,
+ but it will turn on/off the mosaic aug.
+ """
+
+ def __init__(self, *args, mosaic=True, **kwargs):
+ super().__init__(*args, **kwargs)
+ self.mosaic = mosaic
+
+ def __iter__(self):
+ for batch in super().__iter__():
+ yield [(self.mosaic, idx) for idx in batch]
+
+
+class InfiniteSampler(Sampler):
+ """
+ In training, we only care about the "infinite stream" of training data.
+ So this sampler produces an infinite stream of indices and
+ all workers cooperate to correctly shuffle the indices and sample different indices.
+ The samplers in each worker effectively produces `indices[worker_id::num_workers]`
+ where `indices` is an infinite stream of indices consisting of
+ `shuffle(range(size)) + shuffle(range(size)) + ...` (if shuffle is True)
+ or `range(size) + range(size) + ...` (if shuffle is False)
+ """
+
+ def __init__(
+ self,
+ size: int,
+ shuffle: bool = True,
+ seed: Optional[int] = 0,
+ rank=0,
+ world_size=1,
+ ):
+ """
+ Args:
+ size (int): the total number of data of the underlying dataset to sample from
+ shuffle (bool): whether to shuffle the indices or not
+ seed (int): the initial seed of the shuffle. Must be the same
+ across all workers. If None, will use a random seed shared
+ among workers (require synchronization among all workers).
+ """
+ self._size = size
+ assert size > 0
+ self._shuffle = shuffle
+ self._seed = int(seed)
+
+ if dist.is_available() and dist.is_initialized():
+ self._rank = dist.get_rank()
+ self._world_size = dist.get_world_size()
+ else:
+ self._rank = rank
+ self._world_size = world_size
+
+ def __iter__(self):
+ start = self._rank
+ yield from itertools.islice(
+ self._infinite_indices(), start, None, self._world_size
+ )
+
+ def _infinite_indices(self):
+ g = torch.Generator()
+ g.manual_seed(self._seed)
+ while True:
+ if self._shuffle:
+ yield from torch.randperm(self._size, generator=g)
+ else:
+ yield from torch.arange(self._size)
+
+ def __len__(self):
+ return self._size // self._world_size

yolox/evaluators/__init__.py

@@ -0,0 +1,6 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+from .coco_evaluator import COCOEvaluator
+from .voc_evaluator import VOCEvaluator

yolox/evaluators/coco_evaluator.py

@@ -0,0 +1,317 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import contextlib
+import io
+import itertools
+import json
+import tempfile
+import time
+from collections import ChainMap, defaultdict
+from loguru import logger
+from tabulate import tabulate
+from tqdm import tqdm
+
+import numpy as np
+
+import torch
+
+from yolox.data.datasets import COCO_CLASSES
+from yolox.utils import (
+ gather,
+ is_main_process,
+ postprocess,
+ synchronize,
+ time_synchronized,
+ xyxy2xywh
+)
+
+
+def per_class_AR_table(coco_eval, class_names=COCO_CLASSES, headers=["class", "AR"], colums=6):
+ per_class_AR = {}
+ recalls = coco_eval.eval["recall"]
+ # dimension of recalls: [TxKxAxM]
+ # recall has dims (iou, cls, area range, max dets)
+ assert len(class_names) == recalls.shape[1]
+
+ for idx, name in enumerate(class_names):
+ recall = recalls[:, idx, 0, -1]
+ recall = recall[recall > -1]
+ ar = np.mean(recall) if recall.size else float("nan")
+ per_class_AR[name] = float(ar * 100)
+
+ num_cols = min(colums, len(per_class_AR) * len(headers))
+ result_pair = [x for pair in per_class_AR.items() for x in pair]
+ row_pair = itertools.zip_longest(*[result_pair[i::num_cols] for i in range(num_cols)])
+ table_headers = headers * (num_cols // len(headers))
+ table = tabulate(
+ row_pair, tablefmt="pipe", floatfmt=".3f", headers=table_headers, numalign="left",
+ )
+ return table
+
+
+def per_class_AP_table(coco_eval, class_names=COCO_CLASSES, headers=["class", "AP"], colums=6):
+ per_class_AP = {}
+ precisions = coco_eval.eval["precision"]
+ # dimension of precisions: [TxRxKxAxM]
+ # precision has dims (iou, recall, cls, area range, max dets)
+ assert len(class_names) == precisions.shape[2]
+
+ for idx, name in enumerate(class_names):
+ # area range index 0: all area ranges
+ # max dets index -1: typically 100 per image
+ precision = precisions[:, :, idx, 0, -1]
+ precision = precision[precision > -1]
+ ap = np.mean(precision) if precision.size else float("nan")
+ per_class_AP[name] = float(ap * 100)
+
+ num_cols = min(colums, len(per_class_AP) * len(headers))
+ result_pair = [x for pair in per_class_AP.items() for x in pair]
+ row_pair = itertools.zip_longest(*[result_pair[i::num_cols] for i in range(num_cols)])
+ table_headers = headers * (num_cols // len(headers))
+ table = tabulate(
+ row_pair, tablefmt="pipe", floatfmt=".3f", headers=table_headers, numalign="left",
+ )
+ return table
+
+
+class COCOEvaluator:
+ """
+ COCO AP Evaluation class. All the data in the val2017 dataset are processed
+ and evaluated by COCO API.
+ """
+
+ def __init__(
+ self,
+ dataloader,
+ img_size: int,
+ confthre: float,
+ nmsthre: float,
+ num_classes: int,
+ testdev: bool = False,
+ per_class_AP: bool = True,
+ per_class_AR: bool = True,
+ ):
+ """
+ Args:
+ dataloader (Dataloader): evaluate dataloader.
+ img_size: image size after preprocess. images are resized
+ to squares whose shape is (img_size, img_size).
+ confthre: confidence threshold ranging from 0 to 1, which
+ is defined in the config file.
+ nmsthre: IoU threshold of non-max supression ranging from 0 to 1.
+ per_class_AP: Show per class AP during evalution or not. Default to True.
+ per_class_AR: Show per class AR during evalution or not. Default to True.
+ """
+ self.dataloader = dataloader
+ self.img_size = img_size
+ self.confthre = confthre
+ self.nmsthre = nmsthre
+ self.num_classes = num_classes
+ self.testdev = testdev
+ self.per_class_AP = per_class_AP
+ self.per_class_AR = per_class_AR
+
+ def evaluate(
+ self, model, distributed=False, half=False, trt_file=None,
+ decoder=None, test_size=None, return_outputs=False
+ ):
+ """
+ COCO average precision (AP) Evaluation. Iterate inference on the test dataset
+ and the results are evaluated by COCO API.
+
+ NOTE: This function will change training mode to False, please save states if needed.
+
+ Args:
+ model : model to evaluate.
+
+ Returns:
+ ap50_95 (float) : COCO AP of IoU=50:95
+ ap50 (float) : COCO AP of IoU=50
+ summary (sr): summary info of evaluation.
+ """
+ # TODO half to amp_test
+ tensor_type = torch.cuda.HalfTensor if half else torch.cuda.FloatTensor
+ model = model.eval()
+ if half:
+ model = model.half()
+ ids = []
+ data_list = []
+ output_data = defaultdict()
+ progress_bar = tqdm if is_main_process() else iter
+
+ inference_time = 0
+ nms_time = 0
+ n_samples = max(len(self.dataloader) - 1, 1)
+
+ if trt_file is not None:
+ from torch2trt import TRTModule
+
+ model_trt = TRTModule()
+ model_trt.load_state_dict(torch.load(trt_file))
+
+ x = torch.ones(1, 3, test_size[0], test_size[1]).cuda()
+ model(x)
+ model = model_trt
+
+ for cur_iter, (imgs, _, info_imgs, ids) in enumerate(
+ progress_bar(self.dataloader)
+ ):
+ with torch.no_grad():
+ imgs = imgs.type(tensor_type)
+
+ # skip the last iters since batchsize might be not enough for batch inference
+ is_time_record = cur_iter < len(self.dataloader) - 1
+ if is_time_record:
+ start = time.time()
+
+ outputs = model(imgs)
+ if decoder is not None:
+ outputs = decoder(outputs, dtype=outputs.type())
+
+ if is_time_record:
+ infer_end = time_synchronized()
+ inference_time += infer_end - start
+
+ outputs = postprocess(
+ outputs, self.num_classes, self.confthre, self.nmsthre
+ )
+ if is_time_record:
+ nms_end = time_synchronized()
+ nms_time += nms_end - infer_end
+
+ data_list_elem, image_wise_data = self.convert_to_coco_format(
+ outputs, info_imgs, ids, return_outputs=True)
+ data_list.extend(data_list_elem)
+ output_data.update(image_wise_data)
+
+ statistics = torch.cuda.FloatTensor([inference_time, nms_time, n_samples])
+ if distributed:
+ # different process/device might have different speed,
+ # to make sure the process will not be stucked, sync func is used here.
+ synchronize()
+ data_list = gather(data_list, dst=0)
+ output_data = gather(output_data, dst=0)
+ data_list = list(itertools.chain(*data_list))
+ output_data = dict(ChainMap(*output_data))
+ torch.distributed.reduce(statistics, dst=0)
+
+ eval_results = self.evaluate_prediction(data_list, statistics)
+ synchronize()
+
+ if return_outputs:
+ return eval_results, output_data
+ return eval_results
+
+ def convert_to_coco_format(self, outputs, info_imgs, ids, return_outputs=False):
+ data_list = []
+ image_wise_data = defaultdict(dict)
+ for (output, img_h, img_w, img_id) in zip(
+ outputs, info_imgs[0], info_imgs[1], ids
+ ):
+ if output is None:
+ continue
+ output = output.cpu()
+
+ bboxes = output[:, 0:4]
+
+ # preprocessing: resize
+ scale = min(
+ self.img_size[0] / float(img_h), self.img_size[1] / float(img_w)
+ )
+ bboxes /= scale
+ cls = output[:, 6]
+ scores = output[:, 4] * output[:, 5]
+
+ image_wise_data.update({
+ int(img_id): {
+ "bboxes": [box.numpy().tolist() for box in bboxes],
+ "scores": [score.numpy().item() for score in scores],
+ "categories": [
+ self.dataloader.dataset.class_ids[int(cls[ind])]
+ for ind in range(bboxes.shape[0])
+ ],
+ }
+ })
+
+ bboxes = xyxy2xywh(bboxes)
+
+ for ind in range(bboxes.shape[0]):
+ label = self.dataloader.dataset.class_ids[int(cls[ind])]
+ pred_data = {
+ "image_id": int(img_id),
+ "category_id": label,
+ "bbox": bboxes[ind].numpy().tolist(),
+ "score": scores[ind].numpy().item(),
+ "segmentation": [],
+ } # COCO json format
+ data_list.append(pred_data)
+
+ if return_outputs:
+ return data_list, image_wise_data
+ return data_list
+
+ def evaluate_prediction(self, data_dict, statistics):
+ if not is_main_process():
+ return 0, 0, None
+
+ logger.info("Evaluate in main process...")
+
+ annType = ["segm", "bbox", "keypoints"]
+
+ inference_time = statistics[0].item()
+ nms_time = statistics[1].item()
+ n_samples = statistics[2].item()
+
+ a_infer_time = 1000 * inference_time / (n_samples * self.dataloader.batch_size)
+ a_nms_time = 1000 * nms_time / (n_samples * self.dataloader.batch_size)
+
+ time_info = ", ".join(
+ [
+ "Average {} time: {:.2f} ms".format(k, v)
+ for k, v in zip(
+ ["forward", "NMS", "inference"],
+ [a_infer_time, a_nms_time, (a_infer_time + a_nms_time)],
+ )
+ ]
+ )
+
+ info = time_info + "\n"
+
+ # Evaluate the Dt (detection) json comparing with the ground truth
+ if len(data_dict) > 0:
+ cocoGt = self.dataloader.dataset.coco
+ # TODO: since pycocotools can't process dict in py36, write data to json file.
+ if self.testdev:
+ json.dump(data_dict, open("./yolox_testdev_2017.json", "w"))
+ cocoDt = cocoGt.loadRes("./yolox_testdev_2017.json")
+ else:
+ _, tmp = tempfile.mkstemp()
+ json.dump(data_dict, open(tmp, "w"))
+ cocoDt = cocoGt.loadRes(tmp)
+ try:
+ from yolox.layers import COCOeval_opt as COCOeval
+ except ImportError:
+ from pycocotools.cocoeval import COCOeval
+
+ logger.warning("Use standard COCOeval.")
+
+ cocoEval = COCOeval(cocoGt, cocoDt, annType[1])
+ cocoEval.evaluate()
+ cocoEval.accumulate()
+ redirect_string = io.StringIO()
+ with contextlib.redirect_stdout(redirect_string):
+ cocoEval.summarize()
+ info += redirect_string.getvalue()
+ cat_ids = list(cocoGt.cats.keys())
+ cat_names = [cocoGt.cats[catId]['name'] for catId in sorted(cat_ids)]
+ if self.per_class_AP:
+ AP_table = per_class_AP_table(cocoEval, class_names=cat_names)
+ info += "per class AP:\n" + AP_table + "\n"
+ if self.per_class_AR:
+ AR_table = per_class_AR_table(cocoEval, class_names=cat_names)
+ info += "per class AR:\n" + AR_table + "\n"
+ return cocoEval.stats[0], cocoEval.stats[1], info
+ else:
+ return 0, 0, info

yolox/evaluators/voc_eval.py

@@ -0,0 +1,183 @@
+#!/usr/bin/env python3
+# Code are based on
+# https://github.com/rbgirshick/py-faster-rcnn/blob/master/lib/datasets/voc_eval.py
+# Copyright (c) Bharath Hariharan.
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import os
+import pickle
+import xml.etree.ElementTree as ET
+
+import numpy as np
+
+
+def parse_rec(filename):
+ """Parse a PASCAL VOC xml file"""
+ tree = ET.parse(filename)
+ objects = []
+ for obj in tree.findall("object"):
+ obj_struct = {}
+ obj_struct["name"] = obj.find("name").text
+ obj_struct["pose"] = obj.find("pose").text
+ obj_struct["truncated"] = int(obj.find("truncated").text)
+ obj_struct["difficult"] = int(obj.find("difficult").text)
+ bbox = obj.find("bndbox")
+ obj_struct["bbox"] = [
+ int(bbox.find("xmin").text),
+ int(bbox.find("ymin").text),
+ int(bbox.find("xmax").text),
+ int(bbox.find("ymax").text),
+ ]
+ objects.append(obj_struct)
+
+ return objects
+
+
+def voc_ap(rec, prec, use_07_metric=False):
+ """
+ Compute VOC AP given precision and recall.
+ If use_07_metric is true, uses the
+ VOC 07 11 point method (default:False).
+ """
+ if use_07_metric:
+ # 11 point metric
+ ap = 0.0
+ for t in np.arange(0.0, 1.1, 0.1):
+ if np.sum(rec >= t) == 0:
+ p = 0
+ else:
+ p = np.max(prec[rec >= t])
+ ap = ap + p / 11.0
+ else:
+ # correct AP calculation
+ # first append sentinel values at the end
+ mrec = np.concatenate(([0.0], rec, [1.0]))
+ mpre = np.concatenate(([0.0], prec, [0.0]))
+
+ # compute the precision envelope
+ for i in range(mpre.size - 1, 0, -1):
+ mpre[i - 1] = np.maximum(mpre[i - 1], mpre[i])
+
+ # to calculate area under PR curve, look for points
+ # where X axis (recall) changes value
+ i = np.where(mrec[1:] != mrec[:-1])[0]
+
+ # and sum (\Delta recall) * prec
+ ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1])
+ return ap
+
+
+def voc_eval(
+ detpath,
+ annopath,
+ imagesetfile,
+ classname,
+ cachedir,
+ ovthresh=0.5,
+ use_07_metric=False,
+):
+ # first load gt
+ if not os.path.isdir(cachedir):
+ os.mkdir(cachedir)
+ cachefile = os.path.join(cachedir, "annots.pkl")
+ # read list of images
+ with open(imagesetfile, "r") as f:
+ lines = f.readlines()
+ imagenames = [x.strip() for x in lines]
+
+ if not os.path.isfile(cachefile):
+ # load annots
+ recs = {}
+ for i, imagename in enumerate(imagenames):
+ recs[imagename] = parse_rec(annopath.format(imagename))
+ if i % 100 == 0:
+ print(f"Reading annotation for {i + 1}/{len(imagenames)}")
+ # save
+ print(f"Saving cached annotations to {cachefile}")
+ with open(cachefile, "wb") as f:
+ pickle.dump(recs, f)
+ else:
+ # load
+ with open(cachefile, "rb") as f:
+ recs = pickle.load(f)
+
+ # extract gt objects for this class
+ class_recs = {}
+ npos = 0
+ for imagename in imagenames:
+ R = [obj for obj in recs[imagename] if obj["name"] == classname]
+ bbox = np.array([x["bbox"] for x in R])
+ difficult = np.array([x["difficult"] for x in R]).astype(bool)
+ det = [False] * len(R)
+ npos = npos + sum(~difficult)
+ class_recs[imagename] = {"bbox": bbox, "difficult": difficult, "det": det}
+
+ # read dets
+ detfile = detpath.format(classname)
+ with open(detfile, "r") as f:
+ lines = f.readlines()
+
+ if len(lines) == 0:
+ return 0, 0, 0
+
+ splitlines = [x.strip().split(" ") for x in lines]
+ image_ids = [x[0] for x in splitlines]
+ confidence = np.array([float(x[1]) for x in splitlines])
+ BB = np.array([[float(z) for z in x[2:]] for x in splitlines])
+
+ # sort by confidence
+ sorted_ind = np.argsort(-confidence)
+ BB = BB[sorted_ind, :]
+ image_ids = [image_ids[x] for x in sorted_ind]
+
+ # go down dets and mark TPs and FPs
+ nd = len(image_ids)
+ tp = np.zeros(nd)
+ fp = np.zeros(nd)
+ for d in range(nd):
+ R = class_recs[image_ids[d]]
+ bb = BB[d, :].astype(float)
+ ovmax = -np.inf
+ BBGT = R["bbox"].astype(float)
+
+ if BBGT.size > 0:
+ # compute overlaps
+ # intersection
+ ixmin = np.maximum(BBGT[:, 0], bb[0])
+ iymin = np.maximum(BBGT[:, 1], bb[1])
+ ixmax = np.minimum(BBGT[:, 2], bb[2])
+ iymax = np.minimum(BBGT[:, 3], bb[3])
+ iw = np.maximum(ixmax - ixmin + 1.0, 0.0)
+ ih = np.maximum(iymax - iymin + 1.0, 0.0)
+ inters = iw * ih
+
+ # union
+ uni = (
+ (bb[2] - bb[0] + 1.0) * (bb[3] - bb[1] + 1.0)
+ + (BBGT[:, 2] - BBGT[:, 0] + 1.0) * (BBGT[:, 3] - BBGT[:, 1] + 1.0) - inters
+ )
+
+ overlaps = inters / uni
+ ovmax = np.max(overlaps)
+ jmax = np.argmax(overlaps)
+
+ if ovmax > ovthresh:
+ if not R["difficult"][jmax]:
+ if not R["det"][jmax]:
+ tp[d] = 1.0
+ R["det"][jmax] = 1
+ else:
+ fp[d] = 1.0
+ else:
+ fp[d] = 1.0
+
+ # compute precision recall
+ fp = np.cumsum(fp)
+ tp = np.cumsum(tp)
+ rec = tp / float(npos)
+ # avoid divide by zero in case the first detection matches a difficult
+ # ground truth
+ prec = tp / np.maximum(tp + fp, np.finfo(np.float64).eps)
+ ap = voc_ap(rec, prec, use_07_metric)
+
+ return rec, prec, ap

yolox/evaluators/voc_evaluator.py

@@ -0,0 +1,187 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii, Inc. and its affiliates.
+
+import sys
+import tempfile
+import time
+from collections import ChainMap
+from loguru import logger
+from tqdm import tqdm
+
+import numpy as np
+
+import torch
+
+from yolox.utils import gather, is_main_process, postprocess, synchronize, time_synchronized
+
+
+class VOCEvaluator:
+ """
+ VOC AP Evaluation class.
+ """
+
+ def __init__(self, dataloader, img_size, confthre, nmsthre, num_classes):
+ """
+ Args:
+ dataloader (Dataloader): evaluate dataloader.
+ img_size (int): image size after preprocess. images are resized
+ to squares whose shape is (img_size, img_size).
+ confthre (float): confidence threshold ranging from 0 to 1, which
+ is defined in the config file.
+ nmsthre (float): IoU threshold of non-max supression ranging from 0 to 1.
+ """
+ self.dataloader = dataloader
+ self.img_size = img_size
+ self.confthre = confthre
+ self.nmsthre = nmsthre
+ self.num_classes = num_classes
+ self.num_images = len(dataloader.dataset)
+
+ def evaluate(
+ self, model, distributed=False, half=False, trt_file=None,
+ decoder=None, test_size=None, return_outputs=False,
+ ):
+ """
+ VOC average precision (AP) Evaluation. Iterate inference on the test dataset
+ and the results are evaluated by COCO API.
+
+ NOTE: This function will change training mode to False, please save states if needed.
+
+ Args:
+ model : model to evaluate.
+
+ Returns:
+ ap50_95 (float) : COCO style AP of IoU=50:95
+ ap50 (float) : VOC 2007 metric AP of IoU=50
+ summary (sr): summary info of evaluation.
+ """
+ # TODO half to amp_test
+ tensor_type = torch.cuda.HalfTensor if half else torch.cuda.FloatTensor
+ model = model.eval()
+ if half:
+ model = model.half()
+ ids = []
+ data_list = {}
+ progress_bar = tqdm if is_main_process() else iter
+
+ inference_time = 0
+ nms_time = 0
+ n_samples = max(len(self.dataloader) - 1, 1)
+
+ if trt_file is not None:
+ from torch2trt import TRTModule
+
+ model_trt = TRTModule()
+ model_trt.load_state_dict(torch.load(trt_file))
+
+ x = torch.ones(1, 3, test_size[0], test_size[1]).cuda()
+ model(x)
+ model = model_trt
+
+ for cur_iter, (imgs, _, info_imgs, ids) in enumerate(progress_bar(self.dataloader)):
+ with torch.no_grad():
+ imgs = imgs.type(tensor_type)
+
+ # skip the last iters since batchsize might be not enough for batch inference
+ is_time_record = cur_iter < len(self.dataloader) - 1
+ if is_time_record:
+ start = time.time()
+
+ outputs = model(imgs)
+ if decoder is not None:
+ outputs = decoder(outputs, dtype=outputs.type())
+
+ if is_time_record:
+ infer_end = time_synchronized()
+ inference_time += infer_end - start
+
+ outputs = postprocess(
+ outputs, self.num_classes, self.confthre, self.nmsthre
+ )
+ if is_time_record:
+ nms_end = time_synchronized()
+ nms_time += nms_end - infer_end
+
+ data_list.update(self.convert_to_voc_format(outputs, info_imgs, ids))
+
+ statistics = torch.cuda.FloatTensor([inference_time, nms_time, n_samples])
+ if distributed:
+ data_list = gather(data_list, dst=0)
+ data_list = ChainMap(*data_list)
+ torch.distributed.reduce(statistics, dst=0)
+
+ eval_results = self.evaluate_prediction(data_list, statistics)
+ synchronize()
+ if return_outputs:
+ return eval_results, data_list
+ return eval_results
+
+ def convert_to_voc_format(self, outputs, info_imgs, ids):
+ predictions = {}
+ for output, img_h, img_w, img_id in zip(outputs, info_imgs[0], info_imgs[1], ids):
+ if output is None:
+ predictions[int(img_id)] = (None, None, None)
+ continue
+ output = output.cpu()
+
+ bboxes = output[:, 0:4]
+
+ # preprocessing: resize
+ scale = min(self.img_size[0] / float(img_h), self.img_size[1] / float(img_w))
+ bboxes /= scale
+
+ cls = output[:, 6]
+ scores = output[:, 4] * output[:, 5]
+
+ predictions[int(img_id)] = (bboxes, cls, scores)
+ return predictions
+
+ def evaluate_prediction(self, data_dict, statistics):
+ if not is_main_process():
+ return 0, 0, None
+
+ logger.info("Evaluate in main process...")
+
+ inference_time = statistics[0].item()
+ nms_time = statistics[1].item()
+ n_samples = statistics[2].item()
+
+ a_infer_time = 1000 * inference_time / (n_samples * self.dataloader.batch_size)
+ a_nms_time = 1000 * nms_time / (n_samples * self.dataloader.batch_size)
+
+ time_info = ", ".join(
+ [
+ "Average {} time: {:.2f} ms".format(k, v)
+ for k, v in zip(
+ ["forward", "NMS", "inference"],
+ [a_infer_time, a_nms_time, (a_infer_time + a_nms_time)],
+ )
+ ]
+ )
+ info = time_info + "\n"
+
+ all_boxes = [
+ [[] for _ in range(self.num_images)] for _ in range(self.num_classes)
+ ]
+ for img_num in range(self.num_images):
+ bboxes, cls, scores = data_dict[img_num]
+ if bboxes is None:
+ for j in range(self.num_classes):
+ all_boxes[j][img_num] = np.empty([0, 5], dtype=np.float32)
+ continue
+ for j in range(self.num_classes):
+ mask_c = cls == j
+ if sum(mask_c) == 0:
+ all_boxes[j][img_num] = np.empty([0, 5], dtype=np.float32)
+ continue
+
+ c_dets = torch.cat((bboxes, scores.unsqueeze(1)), dim=1)
+ all_boxes[j][img_num] = c_dets[mask_c].numpy()
+
+ sys.stdout.write(f"im_eval: {img_num + 1}/{self.num_images} \r")
+ sys.stdout.flush()
+
+ with tempfile.TemporaryDirectory() as tempdir:
+ mAP50, mAP70 = self.dataloader.dataset.evaluate_detections(all_boxes, tempdir)
+ return mAP50, mAP70, info

yolox/exp/__init__.py

@@ -0,0 +1,6 @@
+#!/usr/bin/env python3
+# Copyright (c) Megvii Inc. All rights reserved.
+
+from .base_exp import BaseExp
+from .build import get_exp
+from .yolox_base import Exp, check_exp_value

yolox/exp/base_exp.py

@@ -0,0 +1,90 @@
+#!/usr/bin/env python3
+# Copyright (c) Megvii Inc. All rights reserved.
+
+import ast
+import pprint
+from abc import ABCMeta, abstractmethod
+from typing import Dict, List, Tuple
+from tabulate import tabulate
+
+import torch
+from torch.nn import Module
+
+from yolox.utils import LRScheduler
+
+
+class BaseExp(metaclass=ABCMeta):
+ """Basic class for any experiment."""
+
+ def __init__(self):
+ self.seed = None
+ self.output_dir = "./YOLOX_outputs"
+ self.print_interval = 100
+ self.eval_interval = 10
+ self.dataset = None
+
+ @abstractmethod
+ def get_model(self) -> Module:
+ pass
+
+ @abstractmethod
+ def get_dataset(self, cache: bool = False, cache_type: str = "ram"):
+ pass
+
+ @abstractmethod
+ def get_data_loader(
+ self, batch_size: int, is_distributed: bool
+ ) -> Dict[str, torch.utils.data.DataLoader]:
+ pass
+
+ @abstractmethod
+ def get_optimizer(self, batch_size: int) -> torch.optim.Optimizer:
+ pass
+
+ @abstractmethod
+ def get_lr_scheduler(
+ self, lr: float, iters_per_epoch: int, **kwargs
+ ) -> LRScheduler:
+ pass
+
+ @abstractmethod
+ def get_evaluator(self):
+ pass
+
+ @abstractmethod
+ def eval(self, model, evaluator, weights):
+ pass
+
+ def __repr__(self):
+ table_header = ["keys", "values"]
+ exp_table = [
+ (str(k), pprint.pformat(v))
+ for k, v in vars(self).items()
+ if not k.startswith("_")
+ ]
+ return tabulate(exp_table, headers=table_header, tablefmt="fancy_grid")
+
+ def merge(self, cfg_list):
+ assert len(cfg_list) % 2 == 0, f"length must be even, check value here: {cfg_list}"
+ for k, v in zip(cfg_list[0::2], cfg_list[1::2]):
+ # only update value with same key
+ if hasattr(self, k):
+ src_value = getattr(self, k)
+ src_type = type(src_value)
+
+ # pre-process input if source type is list or tuple
+ if isinstance(src_value, (List, Tuple)):
+ v = v.strip("[]()")
+ v = [t.strip() for t in v.split(",")]
+
+ # find type of tuple
+ if len(src_value) > 0:
+ src_item_type = type(src_value[0])
+ v = [src_item_type(t) for t in v]
+
+ if src_value is not None and src_type != type(v):
+ try:
+ v = src_type(v)
+ except Exception:
+ v = ast.literal_eval(v)
+ setattr(self, k, v)

yolox/exp/build.py

@@ -0,0 +1,42 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii Inc. All rights reserved.
+
+import importlib
+import os
+import sys
+
+
+def get_exp_by_file(exp_file):
+ try:
+ sys.path.append(os.path.dirname(exp_file))
+ current_exp = importlib.import_module(os.path.basename(exp_file).split(".")[0])
+ exp = current_exp.Exp()
+ except Exception:
+ raise ImportError("{} doesn't contains class named 'Exp'".format(exp_file))
+ return exp
+
+
+def get_exp_by_name(exp_name):
+ exp = exp_name.replace("-", "_") # convert string like "yolox-s" to "yolox_s"
+ module_name = ".".join(["yolox", "exp", "default", exp])
+ exp_object = importlib.import_module(module_name).Exp()
+ return exp_object
+
+
+def get_exp(exp_file=None, exp_name=None):
+ """
+ get Exp object by file or name. If exp_file and exp_name
+ are both provided, get Exp by exp_file.
+
+ Args:
+ exp_file (str): file path of experiment.
+ exp_name (str): name of experiment. "yolo-s",
+ """
+ assert (
+ exp_file is not None or exp_name is not None
+ ), "plz provide exp file or exp name."
+ if exp_file is not None:
+ return get_exp_by_file(exp_file)
+ else:
+ return get_exp_by_name(exp_name)

yolox/exp/default/__init__.py

@@ -0,0 +1,28 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii Inc. All rights reserved.
+
+# This file is used for package installation and find default exp file
+
+import sys
+from importlib import abc, util
+from pathlib import Path
+
+_EXP_PATH = Path(__file__).resolve().parent.parent.parent.parent / "exps" / "default"
+
+if _EXP_PATH.is_dir():
+ # This is true only for in-place installation (pip install -e, setup.py develop),
+ # where setup(package_dir=) does not work: https://github.com/pypa/setuptools/issues/230
+
+ class _ExpFinder(abc.MetaPathFinder):
+ 
+ def find_spec(self, name, path, target=None):
+ if not name.startswith("yolox.exp.default"):
+ return
+ project_name = name.split(".")[-1] + ".py"
+ target_file = _EXP_PATH / project_name
+ if not target_file.is_file():
+ return
+ return util.spec_from_file_location(name, target_file)
+
+ sys.meta_path.append(_ExpFinder())

yolox/exp/yolox_base.py

@@ -0,0 +1,358 @@
+#!/usr/bin/env python3
+# Copyright (c) Megvii Inc. All rights reserved.
+
+import os
+import random
+
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+
+from .base_exp import BaseExp
+
+__all__ = ["Exp", "check_exp_value"]
+
+
+class Exp(BaseExp):
+ def __init__(self):
+ super().__init__()
+
+ # ---------------- model config ---------------- #
+ # detect classes number of model
+ self.num_classes = 80
+ # factor of model depth
+ self.depth = 1.00
+ # factor of model width
+ self.width = 1.00
+ # activation name. For example, if using "relu", then "silu" will be replaced to "relu".
+ self.act = "silu"
+
+ # ---------------- dataloader config ---------------- #
+ # set worker to 4 for shorter dataloader init time
+ # If your training process cost many memory, reduce this value.
+ self.data_num_workers = 4
+ self.input_size = (640, 640) # (height, width)
+ # Actual multiscale ranges: [640 - 5 * 32, 640 + 5 * 32].
+ # To disable multiscale training, set the value to 0.
+ self.multiscale_range = 5
+ # You can uncomment this line to specify a multiscale range
+ # self.random_size = (14, 26)
+ # dir of dataset images, if data_dir is None, this project will use `datasets` dir
+ self.data_dir = None
+ # name of annotation file for training
+ self.train_ann = "instances_train2017.json"
+ # name of annotation file for evaluation
+ self.val_ann = "instances_val2017.json"
+ # name of annotation file for testing
+ self.test_ann = "instances_test2017.json"
+
+ # --------------- transform config ----------------- #
+ # prob of applying mosaic aug
+ self.mosaic_prob = 1.0
+ # prob of applying mixup aug
+ self.mixup_prob = 1.0
+ # prob of applying hsv aug
+ self.hsv_prob = 1.0
+ # prob of applying flip aug
+ self.flip_prob = 0.5
+ # rotation angle range, for example, if set to 2, the true range is (-2, 2)
+ self.degrees = 10.0
+ # translate range, for example, if set to 0.1, the true range is (-0.1, 0.1)
+ self.translate = 0.1
+ self.mosaic_scale = (0.1, 2)
+ # apply mixup aug or not
+ self.enable_mixup = True
+ self.mixup_scale = (0.5, 1.5)
+ # shear angle range, for example, if set to 2, the true range is (-2, 2)
+ self.shear = 2.0
+
+ # -------------- training config --------------------- #
+ # epoch number used for warmup
+ self.warmup_epochs = 5
+ # max training epoch
+ self.max_epoch = 300
+ # minimum learning rate during warmup
+ self.warmup_lr = 0
+ self.min_lr_ratio = 0.05
+ # learning rate for one image. During training, lr will multiply batchsize.
+ self.basic_lr_per_img = 0.01 / 64.0
+ # name of LRScheduler
+ self.scheduler = "yoloxwarmcos"
+ # last #epoch to close augmention like mosaic
+ self.no_aug_epochs = 15
+ # apply EMA during training
+ self.ema = True
+
+ # weight decay of optimizer
+ self.weight_decay = 5e-4
+ # momentum of optimizer
+ self.momentum = 0.9
+ # log period in iter, for example,
+ # if set to 1, user could see log every iteration.
+ self.print_interval = 10
+ # eval period in epoch, for example,
+ # if set to 1, model will be evaluate after every epoch.
+ self.eval_interval = 10
+ # save history checkpoint or not.
+ # If set to False, yolox will only save latest and best ckpt.
+ self.save_history_ckpt = True
+ # name of experiment
+ self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
+
+ # ----------------- testing config ------------------ #
+ # output image size during evaluation/test
+ self.test_size = (640, 640)
+ # confidence threshold during evaluation/test,
+ # boxes whose scores are less than test_conf will be filtered
+ self.test_conf = 0.01
+ # nms threshold
+ self.nmsthre = 0.65
+
+ def get_model(self):
+ from yolox.models import YOLOX, YOLOPAFPN, YOLOXHead
+
+ def init_yolo(M):
+ for m in M.modules():
+ if isinstance(m, nn.BatchNorm2d):
+ m.eps = 1e-3
+ m.momentum = 0.03
+
+ if getattr(self, "model", None) is None:
+ in_channels = [256, 512, 1024]
+ backbone = YOLOPAFPN(self.depth, self.width, in_channels=in_channels, act=self.act)
+ head = YOLOXHead(self.num_classes, self.width, in_channels=in_channels, act=self.act)
+ self.model = YOLOX(backbone, head)
+
+ self.model.apply(init_yolo)
+ self.model.head.initialize_biases(1e-2)
+ self.model.train()
+ return self.model
+
+ def get_dataset(self, cache: bool = False, cache_type: str = "ram"):
+ """
+ Get dataset according to cache and cache_type parameters.
+ Args:
+ cache (bool): Whether to cache imgs to ram or disk.
+ cache_type (str, optional): Defaults to "ram".
+ "ram" : Caching imgs to ram for fast training.
+ "disk": Caching imgs to disk for fast training.
+ """
+ from yolox.data import COCODataset, TrainTransform
+
+ return COCODataset(
+ data_dir=self.data_dir,
+ json_file=self.train_ann,
+ img_size=self.input_size,
+ preproc=TrainTransform(
+ max_labels=50,
+ flip_prob=self.flip_prob,
+ hsv_prob=self.hsv_prob
+ ),
+ cache=cache,
+ cache_type=cache_type,
+ )
+
+ def get_data_loader(self, batch_size, is_distributed, no_aug=False, cache_img: str = None):
+ """
+ Get dataloader according to cache_img parameter.
+ Args:
+ no_aug (bool, optional): Whether to turn off mosaic data enhancement. Defaults to False.
+ cache_img (str, optional): cache_img is equivalent to cache_type. Defaults to None.
+ "ram" : Caching imgs to ram for fast training.
+ "disk": Caching imgs to disk for fast training.
+ None: Do not use cache, in this case cache_data is also None.
+ """
+ from yolox.data import (
+ TrainTransform,
+ YoloBatchSampler,
+ DataLoader,
+ InfiniteSampler,
+ MosaicDetection,
+ worker_init_reset_seed,
+ )
+ from yolox.utils import wait_for_the_master
+
+ # if cache is True, we will create self.dataset before launch
+ # else we will create self.dataset after launch
+ if self.dataset is None:
+ with wait_for_the_master():
+ assert cache_img is None, \
+ "cache_img must be None if you didn't create self.dataset before launch"
+ self.dataset = self.get_dataset(cache=False, cache_type=cache_img)
+
+ self.dataset = MosaicDetection(
+ dataset=self.dataset,
+ mosaic=not no_aug,
+ img_size=self.input_size,
+ preproc=TrainTransform(
+ max_labels=120,
+ flip_prob=self.flip_prob,
+ hsv_prob=self.hsv_prob),
+ degrees=self.degrees,
+ translate=self.translate,
+ mosaic_scale=self.mosaic_scale,
+ mixup_scale=self.mixup_scale,
+ shear=self.shear,
+ enable_mixup=self.enable_mixup,
+ mosaic_prob=self.mosaic_prob,
+ mixup_prob=self.mixup_prob,
+ )
+
+ if is_distributed:
+ batch_size = batch_size // dist.get_world_size()
+
+ sampler = InfiniteSampler(len(self.dataset), seed=self.seed if self.seed else 0)
+
+ batch_sampler = YoloBatchSampler(
+ sampler=sampler,
+ batch_size=batch_size,
+ drop_last=False,
+ mosaic=not no_aug,
+ )
+
+ dataloader_kwargs = {"num_workers": self.data_num_workers, "pin_memory": True}
+ dataloader_kwargs["batch_sampler"] = batch_sampler
+
+ # Make sure each process has different random seed, especially for 'fork' method.
+ # Check https://github.com/pytorch/pytorch/issues/63311 for more details.
+ dataloader_kwargs["worker_init_fn"] = worker_init_reset_seed
+
+ train_loader = DataLoader(self.dataset, **dataloader_kwargs)
+
+ return train_loader
+
+ def random_resize(self, data_loader, epoch, rank, is_distributed):
+ tensor = torch.LongTensor(2).cuda()
+
+ if rank == 0:
+ size_factor = self.input_size[1] * 1.0 / self.input_size[0]
+ if not hasattr(self, 'random_size'):
+ min_size = int(self.input_size[0] / 32) - self.multiscale_range
+ max_size = int(self.input_size[0] / 32) + self.multiscale_range
+ self.random_size = (min_size, max_size)
+ size = random.randint(*self.random_size)
+ size = (int(32 * size), 32 * int(size * size_factor))
+ tensor[0] = size[0]
+ tensor[1] = size[1]
+
+ if is_distributed:
+ dist.barrier()
+ dist.broadcast(tensor, 0)
+
+ input_size = (tensor[0].item(), tensor[1].item())
+ return input_size
+
+ def preprocess(self, inputs, targets, tsize):
+ scale_y = tsize[0] / self.input_size[0]
+ scale_x = tsize[1] / self.input_size[1]
+ if scale_x != 1 or scale_y != 1:
+ inputs = nn.functional.interpolate(
+ inputs, size=tsize, mode="bilinear", align_corners=False
+ )
+ targets[..., 1::2] = targets[..., 1::2] * scale_x
+ targets[..., 2::2] = targets[..., 2::2] * scale_y
+ return inputs, targets
+
+ def get_optimizer(self, batch_size):
+ if "optimizer" not in self.__dict__:
+ if self.warmup_epochs > 0:
+ lr = self.warmup_lr
+ else:
+ lr = self.basic_lr_per_img * batch_size
+
+ pg0, pg1, pg2 = [], [], [] # optimizer parameter groups
+
+ for k, v in self.model.named_modules():
+ if hasattr(v, "bias") and isinstance(v.bias, nn.Parameter):
+ pg2.append(v.bias) # biases
+ if isinstance(v, nn.BatchNorm2d) or "bn" in k:
+ pg0.append(v.weight) # no decay
+ elif hasattr(v, "weight") and isinstance(v.weight, nn.Parameter):
+ pg1.append(v.weight) # apply decay
+
+ optimizer = torch.optim.SGD(
+ pg0, lr=lr, momentum=self.momentum, nesterov=True
+ )
+ optimizer.add_param_group(
+ {"params": pg1, "weight_decay": self.weight_decay}
+ ) # add pg1 with weight_decay
+ optimizer.add_param_group({"params": pg2})
+ self.optimizer = optimizer
+
+ return self.optimizer
+
+ def get_lr_scheduler(self, lr, iters_per_epoch):
+ from yolox.utils import LRScheduler
+
+ scheduler = LRScheduler(
+ self.scheduler,
+ lr,
+ iters_per_epoch,
+ self.max_epoch,
+ warmup_epochs=self.warmup_epochs,
+ warmup_lr_start=self.warmup_lr,
+ no_aug_epochs=self.no_aug_epochs,
+ min_lr_ratio=self.min_lr_ratio,
+ )
+ return scheduler
+
+ def get_eval_dataset(self, **kwargs):
+ from yolox.data import COCODataset, ValTransform
+ testdev = kwargs.get("testdev", False)
+ legacy = kwargs.get("legacy", False)
+
+ return COCODataset(
+ data_dir=self.data_dir,
+ json_file=self.val_ann if not testdev else self.test_ann,
+ name="val2017" if not testdev else "test2017",
+ img_size=self.test_size,
+ preproc=ValTransform(legacy=legacy),
+ )
+
+ def get_eval_loader(self, batch_size, is_distributed, **kwargs):
+ valdataset = self.get_eval_dataset(**kwargs)
+
+ if is_distributed:
+ batch_size = batch_size // dist.get_world_size()
+ sampler = torch.utils.data.distributed.DistributedSampler(
+ valdataset, shuffle=False
+ )
+ else:
+ sampler = torch.utils.data.SequentialSampler(valdataset)
+
+ dataloader_kwargs = {
+ "num_workers": self.data_num_workers,
+ "pin_memory": True,
+ "sampler": sampler,
+ }
+ dataloader_kwargs["batch_size"] = batch_size
+ val_loader = torch.utils.data.DataLoader(valdataset, **dataloader_kwargs)
+
+ return val_loader
+
+ def get_evaluator(self, batch_size, is_distributed, testdev=False, legacy=False):
+ from yolox.evaluators import COCOEvaluator
+
+ return COCOEvaluator(
+ dataloader=self.get_eval_loader(batch_size, is_distributed,
+ testdev=testdev, legacy=legacy),
+ img_size=self.test_size,
+ confthre=self.test_conf,
+ nmsthre=self.nmsthre,
+ num_classes=self.num_classes,
+ testdev=testdev,
+ )
+
+ def get_trainer(self, args):
+ from yolox.core import Trainer
+ trainer = Trainer(self, args)
+ # NOTE: trainer shouldn't be an attribute of exp object
+ return trainer
+
+ def eval(self, model, evaluator, is_distributed, half=False, return_outputs=False):
+ return evaluator.evaluate(model, is_distributed, half, return_outputs=return_outputs)
+
+
+def check_exp_value(exp: Exp):
+ h, w = exp.input_size
+ assert h % 32 == 0 and w % 32 == 0, "input size must be multiples of 32"

yolox/layers/__init__.py

@@ -0,0 +1,13 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii Inc. All rights reserved.
+
+# import torch first to make jit op work without `ImportError of libc10.so`
+import torch # noqa
+
+from .jit_ops import FastCOCOEvalOp, JitOp
+
+try:
+ from .fast_coco_eval_api import COCOeval_opt
+except ImportError: # exception will be raised when users build yolox from source
+ pass

yolox/layers/cocoeval/cocoeval.cpp

@@ -0,0 +1,502 @@
+// Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+#include "cocoeval.h"
+#include <time.h>
+#include <algorithm>
+#include <cstdint>
+#include <numeric>
+
+using namespace pybind11::literals;
+
+namespace COCOeval {
+
+// Sort detections from highest score to lowest, such that
+// detection_instances[detection_sorted_indices[t]] >=
+// detection_instances[detection_sorted_indices[t+1]]. Use stable_sort to match
+// original COCO API
+void SortInstancesByDetectionScore(
+ const std::vector<InstanceAnnotation>& detection_instances,
+ std::vector<uint64_t>* detection_sorted_indices) {
+ detection_sorted_indices->resize(detection_instances.size());
+ std::iota(
+ detection_sorted_indices->begin(), detection_sorted_indices->end(), 0);
+ std::stable_sort(
+ detection_sorted_indices->begin(),
+ detection_sorted_indices->end(),
+ [&detection_instances](size_t j1, size_t j2) {
+ return detection_instances[j1].score > detection_instances[j2].score;
+ });
+}
+
+// Partition the ground truth objects based on whether or not to ignore them
+// based on area
+void SortInstancesByIgnore(
+ const std::array<double, 2>& area_range,
+ const std::vector<InstanceAnnotation>& ground_truth_instances,
+ std::vector<uint64_t>* ground_truth_sorted_indices,
+ std::vector<bool>* ignores) {
+ ignores->clear();
+ ignores->reserve(ground_truth_instances.size());
+ for (auto o : ground_truth_instances) {
+ ignores->push_back(
+ o.ignore || o.area < area_range[0] || o.area > area_range[1]);
+ }
+
+ ground_truth_sorted_indices->resize(ground_truth_instances.size());
+ std::iota(
+ ground_truth_sorted_indices->begin(),
+ ground_truth_sorted_indices->end(),
+ 0);
+ std::stable_sort(
+ ground_truth_sorted_indices->begin(),
+ ground_truth_sorted_indices->end(),
+ [&ignores](size_t j1, size_t j2) {
+ return (int)(*ignores)[j1] < (int)(*ignores)[j2];
+ });
+}
+
+// For each IOU threshold, greedily match each detected instance to a ground
+// truth instance (if possible) and store the results
+void MatchDetectionsToGroundTruth(
+ const std::vector<InstanceAnnotation>& detection_instances,
+ const std::vector<uint64_t>& detection_sorted_indices,
+ const std::vector<InstanceAnnotation>& ground_truth_instances,
+ const std::vector<uint64_t>& ground_truth_sorted_indices,
+ const std::vector<bool>& ignores,
+ const std::vector<std::vector<double>>& ious,
+ const std::vector<double>& iou_thresholds,
+ const std::array<double, 2>& area_range,
+ ImageEvaluation* results) {
+ // Initialize memory to store return data matches and ignore
+ const int num_iou_thresholds = iou_thresholds.size();
+ const int num_ground_truth = ground_truth_sorted_indices.size();
+ const int num_detections = detection_sorted_indices.size();
+ std::vector<uint64_t> ground_truth_matches(
+ num_iou_thresholds * num_ground_truth, 0);
+ std::vector<uint64_t>& detection_matches = results->detection_matches;
+ std::vector<bool>& detection_ignores = results->detection_ignores;
+ std::vector<bool>& ground_truth_ignores = results->ground_truth_ignores;
+ detection_matches.resize(num_iou_thresholds * num_detections, 0);
+ detection_ignores.resize(num_iou_thresholds * num_detections, false);
+ ground_truth_ignores.resize(num_ground_truth);
+ for (auto g = 0; g < num_ground_truth; ++g) {
+ ground_truth_ignores[g] = ignores[ground_truth_sorted_indices[g]];
+ }
+
+ for (auto t = 0; t < num_iou_thresholds; ++t) {
+ for (auto d = 0; d < num_detections; ++d) {
+ // information about best match so far (match=-1 -> unmatched)
+ double best_iou = std::min(iou_thresholds[t], 1 - 1e-10);
+ int match = -1;
+ for (auto g = 0; g < num_ground_truth; ++g) {
+ // if this ground truth instance is already matched and not a
+ // crowd, it cannot be matched to another detection
+ if (ground_truth_matches[t * num_ground_truth + g] > 0 &&
+ !ground_truth_instances[ground_truth_sorted_indices[g]].is_crowd) {
+ continue;
+ }
+
+ // if detected instance matched to a regular ground truth
+ // instance, we can break on the first ground truth instance
+ // tagged as ignore (because they are sorted by the ignore tag)
+ if (match >= 0 && !ground_truth_ignores[match] &&
+ ground_truth_ignores[g]) {
+ break;
+ }
+
+ // if IOU overlap is the best so far, store the match appropriately
+ if (ious[d][ground_truth_sorted_indices[g]] >= best_iou) {
+ best_iou = ious[d][ground_truth_sorted_indices[g]];
+ match = g;
+ }
+ }
+ // if match was made, store id of match for both detection and
+ // ground truth
+ if (match >= 0) {
+ detection_ignores[t * num_detections + d] = ground_truth_ignores[match];
+ detection_matches[t * num_detections + d] =
+ ground_truth_instances[ground_truth_sorted_indices[match]].id;
+ ground_truth_matches[t * num_ground_truth + match] =
+ detection_instances[detection_sorted_indices[d]].id;
+ }
+
+ // set unmatched detections outside of area range to ignore
+ const InstanceAnnotation& detection =
+ detection_instances[detection_sorted_indices[d]];
+ detection_ignores[t * num_detections + d] =
+ detection_ignores[t * num_detections + d] ||
+ (detection_matches[t * num_detections + d] == 0 &&
+ (detection.area < area_range[0] || detection.area > area_range[1]));
+ }
+ }
+
+ // store detection score results
+ results->detection_scores.resize(detection_sorted_indices.size());
+ for (size_t d = 0; d < detection_sorted_indices.size(); ++d) {
+ results->detection_scores[d] =
+ detection_instances[detection_sorted_indices[d]].score;
+ }
+}
+
+std::vector<ImageEvaluation> EvaluateImages(
+ const std::vector<std::array<double, 2>>& area_ranges,
+ int max_detections,
+ const std::vector<double>& iou_thresholds,
+ const ImageCategoryInstances<std::vector<double>>& image_category_ious,
+ const ImageCategoryInstances<InstanceAnnotation>&
+ image_category_ground_truth_instances,
+ const ImageCategoryInstances<InstanceAnnotation>&
+ image_category_detection_instances) {
+ const int num_area_ranges = area_ranges.size();
+ const int num_images = image_category_ground_truth_instances.size();
+ const int num_categories =
+ image_category_ious.size() > 0 ? image_category_ious[0].size() : 0;
+ std::vector<uint64_t> detection_sorted_indices;
+ std::vector<uint64_t> ground_truth_sorted_indices;
+ std::vector<bool> ignores;
+ std::vector<ImageEvaluation> results_all(
+ num_images * num_area_ranges * num_categories);
+
+ // Store results for each image, category, and area range combination. Results
+ // for each IOU threshold are packed into the same ImageEvaluation object
+ for (auto i = 0; i < num_images; ++i) {
+ for (auto c = 0; c < num_categories; ++c) {
+ const std::vector<InstanceAnnotation>& ground_truth_instances =
+ image_category_ground_truth_instances[i][c];
+ const std::vector<InstanceAnnotation>& detection_instances =
+ image_category_detection_instances[i][c];
+
+ SortInstancesByDetectionScore(
+ detection_instances, &detection_sorted_indices);
+ if ((int)detection_sorted_indices.size() > max_detections) {
+ detection_sorted_indices.resize(max_detections);
+ }
+
+ for (size_t a = 0; a < area_ranges.size(); ++a) {
+ SortInstancesByIgnore(
+ area_ranges[a],
+ ground_truth_instances,
+ &ground_truth_sorted_indices,
+ &ignores);
+
+ MatchDetectionsToGroundTruth(
+ detection_instances,
+ detection_sorted_indices,
+ ground_truth_instances,
+ ground_truth_sorted_indices,
+ ignores,
+ image_category_ious[i][c],
+ iou_thresholds,
+ area_ranges[a],
+ &results_all
+ [c * num_area_ranges * num_images + a * num_images + i]);
+ }
+ }
+ }
+
+ return results_all;
+}
+
+// Convert a python list to a vector
+template <typename T>
+std::vector<T> list_to_vec(const py::list& l) {
+ std::vector<T> v(py::len(l));
+ for (int i = 0; i < (int)py::len(l); ++i) {
+ v[i] = l[i].cast<T>();
+ }
+ return v;
+}
+
+// Helper function to Accumulate()
+// Considers the evaluation results applicable to a particular category, area
+// range, and max_detections parameter setting, which begin at
+// evaluations[evaluation_index]. Extracts a sorted list of length n of all
+// applicable detection instances concatenated across all images in the dataset,
+// which are represented by the outputs evaluation_indices, detection_scores,
+// image_detection_indices, and detection_sorted_indices--all of which are
+// length n. evaluation_indices[i] stores the applicable index into
+// evaluations[] for instance i, which has detection score detection_score[i],
+// and is the image_detection_indices[i]'th of the list of detections
+// for the image containing i. detection_sorted_indices[] defines a sorted
+// permutation of the 3 other outputs
+int BuildSortedDetectionList(
+ const std::vector<ImageEvaluation>& evaluations,
+ const int64_t evaluation_index,
+ const int64_t num_images,
+ const int max_detections,
+ std::vector<uint64_t>* evaluation_indices,
+ std::vector<double>* detection_scores,
+ std::vector<uint64_t>* detection_sorted_indices,
+ std::vector<uint64_t>* image_detection_indices) {
+ assert(evaluations.size() >= evaluation_index + num_images);
+
+ // Extract a list of object instances of the applicable category, area
+ // range, and max detections requirements such that they can be sorted
+ image_detection_indices->clear();
+ evaluation_indices->clear();
+ detection_scores->clear();
+ image_detection_indices->reserve(num_images * max_detections);
+ evaluation_indices->reserve(num_images * max_detections);
+ detection_scores->reserve(num_images * max_detections);
+ int num_valid_ground_truth = 0;
+ for (auto i = 0; i < num_images; ++i) {
+ const ImageEvaluation& evaluation = evaluations[evaluation_index + i];
+
+ for (int d = 0;
+ d < (int)evaluation.detection_scores.size() && d < max_detections;
+ ++d) { // detected instances
+ evaluation_indices->push_back(evaluation_index + i);
+ image_detection_indices->push_back(d);
+ detection_scores->push_back(evaluation.detection_scores[d]);
+ }
+ for (auto ground_truth_ignore : evaluation.ground_truth_ignores) {
+ if (!ground_truth_ignore) {
+ ++num_valid_ground_truth;
+ }
+ }
+ }
+
+ // Sort detections by decreasing score, using stable sort to match
+ // python implementation
+ detection_sorted_indices->resize(detection_scores->size());
+ std::iota(
+ detection_sorted_indices->begin(), detection_sorted_indices->end(), 0);
+ std::stable_sort(
+ detection_sorted_indices->begin(),
+ detection_sorted_indices->end(),
+ [&detection_scores](size_t j1, size_t j2) {
+ return (*detection_scores)[j1] > (*detection_scores)[j2];
+ });
+
+ return num_valid_ground_truth;
+}
+
+// Helper function to Accumulate()
+// Compute a precision recall curve given a sorted list of detected instances
+// encoded in evaluations, evaluation_indices, detection_scores,
+// detection_sorted_indices, image_detection_indices (see
+// BuildSortedDetectionList()). Using vectors precisions and recalls
+// and temporary storage, output the results into precisions_out, recalls_out,
+// and scores_out, which are large buffers containing many precion/recall curves
+// for all possible parameter settings, with precisions_out_index and
+// recalls_out_index defining the applicable indices to store results.
+void ComputePrecisionRecallCurve(
+ const int64_t precisions_out_index,
+ const int64_t precisions_out_stride,
+ const int64_t recalls_out_index,
+ const std::vector<double>& recall_thresholds,
+ const int iou_threshold_index,
+ const int num_iou_thresholds,
+ const int num_valid_ground_truth,
+ const std::vector<ImageEvaluation>& evaluations,
+ const std::vector<uint64_t>& evaluation_indices,
+ const std::vector<double>& detection_scores,
+ const std::vector<uint64_t>& detection_sorted_indices,
+ const std::vector<uint64_t>& image_detection_indices,
+ std::vector<double>* precisions,
+ std::vector<double>* recalls,
+ std::vector<double>* precisions_out,
+ std::vector<double>* scores_out,
+ std::vector<double>* recalls_out) {
+ assert(recalls_out->size() > recalls_out_index);
+
+ // Compute precision/recall for each instance in the sorted list of detections
+ int64_t true_positives_sum = 0, false_positives_sum = 0;
+ precisions->clear();
+ recalls->clear();
+ precisions->reserve(detection_sorted_indices.size());
+ recalls->reserve(detection_sorted_indices.size());
+ assert(!evaluations.empty() || detection_sorted_indices.empty());
+ for (auto detection_sorted_index : detection_sorted_indices) {
+ const ImageEvaluation& evaluation =
+ evaluations[evaluation_indices[detection_sorted_index]];
+ const auto num_detections =
+ evaluation.detection_matches.size() / num_iou_thresholds;
+ const auto detection_index = iou_threshold_index * num_detections +
+ image_detection_indices[detection_sorted_index];
+ assert(evaluation.detection_matches.size() > detection_index);
+ assert(evaluation.detection_ignores.size() > detection_index);
+ const int64_t detection_match =
+ evaluation.detection_matches[detection_index];
+ const bool detection_ignores =
+ evaluation.detection_ignores[detection_index];
+ const auto true_positive = detection_match > 0 && !detection_ignores;
+ const auto false_positive = detection_match == 0 && !detection_ignores;
+ if (true_positive) {
+ ++true_positives_sum;
+ }
+ if (false_positive) {
+ ++false_positives_sum;
+ }
+
+ const double recall =
+ static_cast<double>(true_positives_sum) / num_valid_ground_truth;
+ recalls->push_back(recall);
+ const int64_t num_valid_detections =
+ true_positives_sum + false_positives_sum;
+ const double precision = num_valid_detections > 0
+ ? static_cast<double>(true_positives_sum) / num_valid_detections
+ : 0.0;
+ precisions->push_back(precision);
+ }
+
+ (*recalls_out)[recalls_out_index] = !recalls->empty() ? recalls->back() : 0;
+
+ for (int64_t i = static_cast<int64_t>(precisions->size()) - 1; i > 0; --i) {
+ if ((*precisions)[i] > (*precisions)[i - 1]) {
+ (*precisions)[i - 1] = (*precisions)[i];
+ }
+ }
+
+ // Sample the per instance precision/recall list at each recall threshold
+ for (size_t r = 0; r < recall_thresholds.size(); ++r) {
+ // first index in recalls >= recall_thresholds[r]
+ std::vector<double>::iterator low = std::lower_bound(
+ recalls->begin(), recalls->end(), recall_thresholds[r]);
+ size_t precisions_index = low - recalls->begin();
+
+ const auto results_ind = precisions_out_index + r * precisions_out_stride;
+ assert(results_ind < precisions_out->size());
+ assert(results_ind < scores_out->size());
+ if (precisions_index < precisions->size()) {
+ (*precisions_out)[results_ind] = (*precisions)[precisions_index];
+ (*scores_out)[results_ind] =
+ detection_scores[detection_sorted_indices[precisions_index]];
+ } else {
+ (*precisions_out)[results_ind] = 0;
+ (*scores_out)[results_ind] = 0;
+ }
+ }
+}
+py::dict Accumulate(
+ const py::object& params,
+ const std::vector<ImageEvaluation>& evaluations) {
+ const std::vector<double> recall_thresholds =
+ list_to_vec<double>(params.attr("recThrs"));
+ const std::vector<int> max_detections =
+ list_to_vec<int>(params.attr("maxDets"));
+ const int num_iou_thresholds = py::len(params.attr("iouThrs"));
+ const int num_recall_thresholds = py::len(params.attr("recThrs"));
+ const int num_categories = params.attr("useCats").cast<int>() == 1
+ ? py::len(params.attr("catIds"))
+ : 1;
+ const int num_area_ranges = py::len(params.attr("areaRng"));
+ const int num_max_detections = py::len(params.attr("maxDets"));
+ const int num_images = py::len(params.attr("imgIds"));
+
+ std::vector<double> precisions_out(
+ num_iou_thresholds * num_recall_thresholds * num_categories *
+ num_area_ranges * num_max_detections,
+ -1);
+ std::vector<double> recalls_out(
+ num_iou_thresholds * num_categories * num_area_ranges *
+ num_max_detections,
+ -1);
+ std::vector<double> scores_out(
+ num_iou_thresholds * num_recall_thresholds * num_categories *
+ num_area_ranges * num_max_detections,
+ -1);
+
+ // Consider the list of all detected instances in the entire dataset in one
+ // large list. evaluation_indices, detection_scores,
+ // image_detection_indices, and detection_sorted_indices all have the same
+ // length as this list, such that each entry corresponds to one detected
+ // instance
+ std::vector<uint64_t> evaluation_indices; // indices into evaluations[]
+ std::vector<double> detection_scores; // detection scores of each instance
+ std::vector<uint64_t> detection_sorted_indices; // sorted indices of all
+ // instances in the dataset
+ std::vector<uint64_t>
+ image_detection_indices; // indices into the list of detected instances in
+ // the same image as each instance
+ std::vector<double> precisions, recalls;
+
+ for (auto c = 0; c < num_categories; ++c) {
+ for (auto a = 0; a < num_area_ranges; ++a) {
+ for (auto m = 0; m < num_max_detections; ++m) {
+ // The COCO PythonAPI assumes evaluations[] (the return value of
+ // COCOeval::EvaluateImages() is one long list storing results for each
+ // combination of category, area range, and image id, with categories in
+ // the outermost loop and images in the innermost loop.
+ const int64_t evaluations_index =
+ c * num_area_ranges * num_images + a * num_images;
+ int num_valid_ground_truth = BuildSortedDetectionList(
+ evaluations,
+ evaluations_index,
+ num_images,
+ max_detections[m],
+ &evaluation_indices,
+ &detection_scores,
+ &detection_sorted_indices,
+ &image_detection_indices);
+
+ if (num_valid_ground_truth == 0) {
+ continue;
+ }
+
+ for (auto t = 0; t < num_iou_thresholds; ++t) {
+ // recalls_out is a flattened vectors representing a
+ // num_iou_thresholds X num_categories X num_area_ranges X
+ // num_max_detections matrix
+ const int64_t recalls_out_index =
+ t * num_categories * num_area_ranges * num_max_detections +
+ c * num_area_ranges * num_max_detections +
+ a * num_max_detections + m;
+
+ // precisions_out and scores_out are flattened vectors
+ // representing a num_iou_thresholds X num_recall_thresholds X
+ // num_categories X num_area_ranges X num_max_detections matrix
+ const int64_t precisions_out_stride =
+ num_categories * num_area_ranges * num_max_detections;
+ const int64_t precisions_out_index = t * num_recall_thresholds *
+ num_categories * num_area_ranges * num_max_detections +
+ c * num_area_ranges * num_max_detections +
+ a * num_max_detections + m;
+
+ ComputePrecisionRecallCurve(
+ precisions_out_index,
+ precisions_out_stride,
+ recalls_out_index,
+ recall_thresholds,
+ t,
+ num_iou_thresholds,
+ num_valid_ground_truth,
+ evaluations,
+ evaluation_indices,
+ detection_scores,
+ detection_sorted_indices,
+ image_detection_indices,
+ &precisions,
+ &recalls,
+ &precisions_out,
+ &scores_out,
+ &recalls_out);
+ }
+ }
+ }
+ }
+
+ time_t rawtime;
+ struct tm local_time;
+ std::array<char, 200> buffer;
+ time(&rawtime);
+#ifdef _WIN32
+ localtime_s(&local_time, &rawtime);
+#else
+ localtime_r(&rawtime, &local_time);
+#endif
+ strftime(
+ buffer.data(), 200, "%Y-%m-%d %H:%num_max_detections:%S", &local_time);
+ return py::dict(
+ "params"_a = params,
+ "counts"_a = std::vector<int64_t>({num_iou_thresholds,
+ num_recall_thresholds,
+ num_categories,
+ num_area_ranges,
+ num_max_detections}),
+ "date"_a = buffer,
+ "precision"_a = precisions_out,
+ "recall"_a = recalls_out,
+ "scores"_a = scores_out);
+}
+
+} // namespace COCOeval

yolox/layers/cocoeval/cocoeval.h

@@ -0,0 +1,98 @@
+// Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+#pragma once
+
+#include <pybind11/numpy.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+#include <pybind11/stl_bind.h>
+#include <vector>
+
+namespace py = pybind11;
+
+namespace COCOeval {
+
+// Annotation data for a single object instance in an image
+struct InstanceAnnotation {
+ InstanceAnnotation(
+ uint64_t id,
+ double score,
+ double area,
+ bool is_crowd,
+ bool ignore)
+ : id{id}, score{score}, area{area}, is_crowd{is_crowd}, ignore{ignore} {}
+ uint64_t id;
+ double score = 0.;
+ double area = 0.;
+ bool is_crowd = false;
+ bool ignore = false;
+};
+
+// Stores intermediate results for evaluating detection results for a single
+// image that has D detected instances and G ground truth instances. This stores
+// matches between detected and ground truth instances
+struct ImageEvaluation {
+ // For each of the D detected instances, the id of the matched ground truth
+ // instance, or 0 if unmatched
+ std::vector<uint64_t> detection_matches;
+
+ // The detection score of each of the D detected instances
+ std::vector<double> detection_scores;
+
+ // Marks whether or not each of G instances was ignored from evaluation (e.g.,
+ // because it's outside area_range)
+ std::vector<bool> ground_truth_ignores;
+
+ // Marks whether or not each of D instances was ignored from evaluation (e.g.,
+ // because it's outside aRng)
+ std::vector<bool> detection_ignores;
+};
+
+template <class T>
+using ImageCategoryInstances = std::vector<std::vector<std::vector<T>>>;
+
+// C++ implementation of COCO API cocoeval.py::COCOeval.evaluateImg(). For each
+// combination of image, category, area range settings, and IOU thresholds to
+// evaluate, it matches detected instances to ground truth instances and stores
+// the results into a vector of ImageEvaluation results, which will be
+// interpreted by the COCOeval::Accumulate() function to produce precion-recall
+// curves. The parameters of nested vectors have the following semantics:
+// image_category_ious[i][c][d][g] is the intersection over union of the d'th
+// detected instance and g'th ground truth instance of
+// category category_ids[c] in image image_ids[i]
+// image_category_ground_truth_instances[i][c] is a vector of ground truth
+// instances in image image_ids[i] of category category_ids[c]
+// image_category_detection_instances[i][c] is a vector of detected
+// instances in image image_ids[i] of category category_ids[c]
+std::vector<ImageEvaluation> EvaluateImages(
+ const std::vector<std::array<double, 2>>& area_ranges, // vector of 2-tuples
+ int max_detections,
+ const std::vector<double>& iou_thresholds,
+ const ImageCategoryInstances<std::vector<double>>& image_category_ious,
+ const ImageCategoryInstances<InstanceAnnotation>&
+ image_category_ground_truth_instances,
+ const ImageCategoryInstances<InstanceAnnotation>&
+ image_category_detection_instances);
+
+// C++ implementation of COCOeval.accumulate(), which generates precision
+// recall curves for each set of category, IOU threshold, detection area range,
+// and max number of detections parameters. It is assumed that the parameter
+// evaluations is the return value of the functon COCOeval::EvaluateImages(),
+// which was called with the same parameter settings params
+py::dict Accumulate(
+ const py::object& params,
+ const std::vector<ImageEvaluation>& evalutations);
+
+} // namespace COCOeval
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
+{
+ m.def("COCOevalAccumulate", &COCOeval::Accumulate, "COCOeval::Accumulate");
+ m.def(
+ "COCOevalEvaluateImages",
+ &COCOeval::EvaluateImages,
+ "COCOeval::EvaluateImages");
+ pybind11::class_<COCOeval::InstanceAnnotation>(m, "InstanceAnnotation")
+ .def(pybind11::init<uint64_t, double, double, bool, bool>());
+ pybind11::class_<COCOeval::ImageEvaluation>(m, "ImageEvaluation")
+ .def(pybind11::init<>());
+}

yolox/layers/fast_coco_eval_api.py

@@ -0,0 +1,151 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# This file comes from
+# https://github.com/facebookresearch/detectron2/blob/master/detectron2/evaluation/fast_eval_api.py
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+# Copyright (c) Megvii Inc. All rights reserved.
+
+import copy
+import time
+
+import numpy as np
+from pycocotools.cocoeval import COCOeval
+
+from .jit_ops import FastCOCOEvalOp
+
+
+class COCOeval_opt(COCOeval):
+ """
+ This is a slightly modified version of the original COCO API, where the functions evaluateImg()
+ and accumulate() are implemented in C++ to speedup evaluation
+ """
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+ self.module = FastCOCOEvalOp().load()
+
+ def evaluate(self):
+ """
+ Run per image evaluation on given images and store results in self.evalImgs_cpp, a
+ datastructure that isn't readable from Python but is used by a c++ implementation of
+ accumulate(). Unlike the original COCO PythonAPI, we don't populate the datastructure
+ self.evalImgs because this datastructure is a computational bottleneck.
+ :return: None
+ """
+ tic = time.time()
+
+ print("Running per image evaluation...")
+ p = self.params
+ # add backward compatibility if useSegm is specified in params
+ if p.useSegm is not None:
+ p.iouType = "segm" if p.useSegm == 1 else "bbox"
+ print(
+ "useSegm (deprecated) is not None. Running {} evaluation".format(
+ p.iouType
+ )
+ )
+ print("Evaluate annotation type *{}*".format(p.iouType))
+ p.imgIds = list(np.unique(p.imgIds))
+ if p.useCats:
+ p.catIds = list(np.unique(p.catIds))
+ p.maxDets = sorted(p.maxDets)
+ self.params = p
+
+ self._prepare()
+
+ # loop through images, area range, max detection number
+ catIds = p.catIds if p.useCats else [-1]
+
+ if p.iouType == "segm" or p.iouType == "bbox":
+ computeIoU = self.computeIoU
+ elif p.iouType == "keypoints":
+ computeIoU = self.computeOks
+ self.ious = {
+ (imgId, catId): computeIoU(imgId, catId)
+ for imgId in p.imgIds
+ for catId in catIds
+ }
+
+ maxDet = p.maxDets[-1]
+
+ # <<<< Beginning of code differences with original COCO API
+ def convert_instances_to_cpp(instances, is_det=False):
+ # Convert annotations for a list of instances in an image to a format that's fast
+ # to access in C++
+ instances_cpp = []
+ for instance in instances:
+ instance_cpp = self.module.InstanceAnnotation(
+ int(instance["id"]),
+ instance["score"] if is_det else instance.get("score", 0.0),
+ instance["area"],
+ bool(instance.get("iscrowd", 0)),
+ bool(instance.get("ignore", 0)),
+ )
+ instances_cpp.append(instance_cpp)
+ return instances_cpp
+
+ # Convert GT annotations, detections, and IOUs to a format that's fast to access in C++
+ ground_truth_instances = [
+ [convert_instances_to_cpp(self._gts[imgId, catId]) for catId in p.catIds]
+ for imgId in p.imgIds
+ ]
+ detected_instances = [
+ [
+ convert_instances_to_cpp(self._dts[imgId, catId], is_det=True)
+ for catId in p.catIds
+ ]
+ for imgId in p.imgIds
+ ]
+ ious = [[self.ious[imgId, catId] for catId in catIds] for imgId in p.imgIds]
+
+ if not p.useCats:
+ # For each image, flatten per-category lists into a single list
+ ground_truth_instances = [
+ [[o for c in i for o in c]] for i in ground_truth_instances
+ ]
+ detected_instances = [
+ [[o for c in i for o in c]] for i in detected_instances
+ ]
+
+ # Call C++ implementation of self.evaluateImgs()
+ self._evalImgs_cpp = self.module.COCOevalEvaluateImages(
+ p.areaRng,
+ maxDet,
+ p.iouThrs,
+ ious,
+ ground_truth_instances,
+ detected_instances,
+ )
+ self._evalImgs = None
+
+ self._paramsEval = copy.deepcopy(self.params)
+ toc = time.time()
+ print("COCOeval_opt.evaluate() finished in {:0.2f} seconds.".format(toc - tic))
+ # >>>> End of code differences with original COCO API
+
+ def accumulate(self):
+ """
+ Accumulate per image evaluation results and store the result in self.eval. Does not
+ support changing parameter settings from those used by self.evaluate()
+ """
+ print("Accumulating evaluation results...")
+ tic = time.time()
+ if not hasattr(self, "_evalImgs_cpp"):
+ print("Please run evaluate() first")
+
+ self.eval = self.module.COCOevalAccumulate(self._paramsEval, self._evalImgs_cpp)
+
+ # recall is num_iou_thresholds X num_categories X num_area_ranges X num_max_detections
+ self.eval["recall"] = np.array(self.eval["recall"]).reshape(
+ self.eval["counts"][:1] + self.eval["counts"][2:]
+ )
+
+ # precision and scores are num_iou_thresholds X num_recall_thresholds X num_categories X
+ # num_area_ranges X num_max_detections
+ self.eval["precision"] = np.array(self.eval["precision"]).reshape(
+ self.eval["counts"]
+ )
+ self.eval["scores"] = np.array(self.eval["scores"]).reshape(self.eval["counts"])
+ toc = time.time()
+ print(
+ "COCOeval_opt.accumulate() finished in {:0.2f} seconds.".format(toc - tic)
+ )

yolox/layers/jit_ops.py

@@ -0,0 +1,138 @@
+#!/usr/bin/env python3
+# Copyright (c) Megvii, Inc. and its affiliates. All Rights Reserved
+
+import glob
+import importlib
+import os
+import sys
+import time
+from typing import List
+
+__all__ = ["JitOp", "FastCOCOEvalOp"]
+
+
+class JitOp:
+ """
+ Just-in-time compilation of ops.
+
+ Some code of `JitOp` is inspired by `deepspeed.op_builder`,
+ check the following link for more details:
+ https://github.com/microsoft/DeepSpeed/blob/master/op_builder/builder.py
+ """
+
+ def __init__(self, name):
+ self.name = name
+
+ def absolute_name(self) -> str:
+ """Get absolute build path for cases where the op is pre-installed."""
+ pass
+
+ def sources(self) -> List:
+ """Get path list of source files of op.
+
+ NOTE: the path should be elative to root of package during building,
+ Otherwise, exception will be raised when building package.
+ However, for runtime building, path will be absolute.
+ """
+ pass
+
+ def include_dirs(self) -> List:
+ """
+ Get list of include paths, relative to root of package.
+
+ NOTE: the path should be elative to root of package.
+ Otherwise, exception will be raised when building package.
+ """
+ return []
+
+ def define_macros(self) -> List:
+ """Get list of macros to define for op"""
+ return []
+
+ def cxx_args(self) -> List:
+ """Get optional list of compiler flags to forward"""
+ args = ["-O2"] if sys.platform == "win32" else ["-O3", "-std=c++14", "-g", "-Wno-reorder"]
+ return args
+
+ def nvcc_args(self) -> List:
+ """Get optional list of compiler flags to forward to nvcc when building CUDA sources"""
+ args = [
+ "-O3", "--use_fast_math",
+ "-std=c++17" if sys.platform == "win32" else "-std=c++14",
+ "-U__CUDA_NO_HALF_OPERATORS__",
+ "-U__CUDA_NO_HALF_CONVERSIONS__",
+ "-U__CUDA_NO_HALF2_OPERATORS__",
+ ]
+ return args
+
+ def build_op(self):
+ from torch.utils.cpp_extension import CppExtension
+ return CppExtension(
+ name=self.absolute_name(),
+ sources=self.sources(),
+ include_dirs=self.include_dirs(),
+ define_macros=self.define_macros(),
+ extra_compile_args={
+ "cxx": self.cxx_args(),
+ },
+ )
+
+ def load(self, verbose=True):
+ try:
+ # try to import op from pre-installed package
+ return importlib.import_module(self.absolute_name())
+ except Exception: # op not compiled, jit load
+ from yolox.utils import wait_for_the_master
+ with wait_for_the_master(): # to avoid race condition
+ return self.jit_load(verbose)
+
+ def jit_load(self, verbose=True):
+ from torch.utils.cpp_extension import load
+ from loguru import logger
+ try:
+ import ninja # noqa
+ except ImportError:
+ if verbose:
+ logger.warning(
+ f"Ninja is not installed, fall back to normal installation for {self.name}."
+ )
+
+ build_tik = time.time()
+ # build op and load
+ op_module = load(
+ name=self.name,
+ sources=self.sources(),
+ extra_cflags=self.cxx_args(),
+ extra_cuda_cflags=self.nvcc_args(),
+ verbose=verbose,
+ )
+ build_duration = time.time() - build_tik
+ if verbose:
+ logger.info(f"Load {self.name} op in {build_duration:.3f}s.")
+ return op_module
+
+ def clear_dynamic_library(self):
+ """Remove dynamic libraray files generated by JIT compilation."""
+ module = self.load()
+ os.remove(module.__file__)
+
+
+class FastCOCOEvalOp(JitOp):
+
+ def __init__(self, name="fast_cocoeval"):
+ super().__init__(name=name)
+
+ def absolute_name(self):
+ return f'yolox.layers.{self.name}'
+
+ def sources(self):
+ sources = glob.glob(os.path.join("yolox", "layers", "cocoeval", "*.cpp"))
+ if not sources: # source will be empty list if the so file is removed after install
+ # use abosolute path to compile
+ import yolox
+ code_path = os.path.join(yolox.__path__[0], "layers", "cocoeval", "*.cpp")
+ sources = glob.glob(code_path)
+ return sources
+
+ def include_dirs(self):
+ return [os.path.join("yolox", "layers", "cocoeval")]

yolox/models/__init__.py

@@ -0,0 +1,11 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+# Copyright (c) Megvii Inc. All rights reserved.
+
+from .build import *
+from .darknet import CSPDarknet, Darknet
+from .losses import IOUloss
+from .yolo_fpn import YOLOFPN
+from .yolo_head import YOLOXHead
+from .yolo_pafpn import YOLOPAFPN
+from .yolox import YOLOX