更换文档检测模型

paddle_detection/ppdet/ext_op/setup.py

@@ -0,0 +1,33 @@
+import os
+import glob
+import paddle
+from paddle.utils.cpp_extension import CppExtension, CUDAExtension, setup
+
+
+def get_extensions():
+    root_dir = os.path.dirname(os.path.abspath(__file__))
+    ext_root_dir = os.path.join(root_dir, 'csrc')
+    sources = []
+    for ext_name in os.listdir(ext_root_dir):
+        ext_dir = os.path.join(ext_root_dir, ext_name)
+        source = glob.glob(os.path.join(ext_dir, '*.cc'))
+        kwargs = dict()
+        if paddle.device.is_compiled_with_cuda():
+            source += glob.glob(os.path.join(ext_dir, '*.cu'))
+
+        if not source:
+            continue
+
+        sources += source
+
+    if paddle.device.is_compiled_with_cuda():
+        extension = CUDAExtension(
+            sources, extra_compile_args={'cxx': ['-DPADDLE_WITH_CUDA']})
+    else:
+        extension = CppExtension(sources)
+
+    return extension
+
+
+if __name__ == "__main__":
+    setup(name='ext_op', ext_modules=get_extensions())