| 13 | |
| 14 | |
| 15 | def get_extensions(): |
| 16 | this_dir = path.dirname(path.abspath(__file__)) |
| 17 | extensions_dir = path.join(this_dir, "yolox", "layers", "csrc") |
| 18 | |
| 19 | main_source = path.join(extensions_dir, "vision.cpp") |
| 20 | sources = glob.glob(path.join(extensions_dir, "**", "*.cpp")) |
| 21 | |
| 22 | sources = [main_source] + sources |
| 23 | extension = CppExtension |
| 24 | |
| 25 | extra_compile_args = {"cxx": ["-O3"]} |
| 26 | define_macros = [] |
| 27 | |
| 28 | include_dirs = [extensions_dir] |
| 29 | |
| 30 | ext_modules = [ |
| 31 | extension( |
| 32 | "yolox._C", |
| 33 | sources, |
| 34 | include_dirs=include_dirs, |
| 35 | define_macros=define_macros, |
| 36 | extra_compile_args=extra_compile_args, |
| 37 | ) |
| 38 | ] |
| 39 | |
| 40 | return ext_modules |
| 41 | |
| 42 | |
| 43 | with open("yolox/__init__.py", "r") as f: |