tidalove/yolox

1#!/usr/bin/env python2# Copyright (c) Megvii, Inc. and its affiliates. All Rights Reserved3 4import re5import setuptools6import sys7 8TORCH_AVAILABLE = True9try:10    import torch11    from torch.utils import cpp_extension12except ImportError:13    TORCH_AVAILABLE = False14    print("[WARNING] Unable to import torch, pre-compiling ops will be disabled.")15 16 17def get_package_dir():18    pkg_dir = {19        "yolox.tools": "tools",20        "yolox.exp.default": "exps/default",21    }22    return pkg_dir23 24 25def get_install_requirements():26    with open("requirements.txt", "r", encoding="utf-8") as f:27        reqs = [x.strip() for x in f.read().splitlines()]28    reqs = [x for x in reqs if not x.startswith("#")]29    return reqs30 31 32def get_yolox_version():33    with open("yolox/__init__.py", "r") as f:34        version = re.search(35            r'^__version__\s*=\s*[\'"]([^\'"]*)[\'"]',36            f.read(), re.MULTILINE37        ).group(1)38    return version39 40 41def get_long_description():42    with open("README.md", "r", encoding="utf-8") as f:43        long_description = f.read()44    return long_description45 46 47def get_ext_modules():48    ext_module = []49    if sys.platform != "win32":  # pre-compile ops on linux50        assert TORCH_AVAILABLE, "torch is required for pre-compiling ops, please install it first."51        # if any other op is added, please also add it here52        from yolox.layers import FastCOCOEvalOp53        ext_module.append(FastCOCOEvalOp().build_op())54    return ext_module55 56 57def get_cmd_class():58    cmdclass = {}59    if TORCH_AVAILABLE:60        cmdclass["build_ext"] = cpp_extension.BuildExtension61    return cmdclass62 63 64setuptools.setup(65    name="yolox",66    version=get_yolox_version(),67    author="megvii basedet team",68    url="https://github.com/Megvii-BaseDetection/YOLOX",69    package_dir=get_package_dir(),70    packages=setuptools.find_packages(exclude=("tests", "tools")) + list(get_package_dir().keys()),71    python_requires=">=3.6",72    install_requires=get_install_requirements(),73    setup_requires=["wheel"],  # avoid building error when pip is not updated74    long_description=get_long_description(),75    long_description_content_type="text/markdown",76    include_package_data=True,  # include files in MANIFEST.in77    ext_modules=get_ext_modules(),78    cmdclass=get_cmd_class(),79    classifiers=[80        "Programming Language :: Python :: 3", "Operating System :: OS Independent",81        "License :: OSI Approved :: Apache Software License",82    ],83    project_urls={84        "Documentation": "https://yolox.readthedocs.io",85        "Source": "https://github.com/Megvii-BaseDetection/YOLOX",86        "Tracker": "https://github.com/Megvii-BaseDetection/YOLOX/issues",87    },88)89