kopia lustrzana https://github.com/OpenDroneMap/ODM
92 wiersze
3.2 KiB
Python
92 wiersze
3.2 KiB
Python
import os
|
|
import sys
|
|
import shutil
|
|
import ctypes
|
|
from opendm import log
|
|
from repoze.lru import lru_cache
|
|
|
|
def gpu_disabled_by_user_env():
|
|
return bool(os.environ.get('ODM_NO_GPU'))
|
|
|
|
@lru_cache(maxsize=None)
|
|
def has_popsift_and_can_handle_texsize(width, height):
|
|
# We first check that we have the required compute capabilities
|
|
# As we do not support compute capabilities less than 3.5
|
|
try:
|
|
compute_major, compute_minor = get_cuda_compute_version(0)
|
|
if compute_major < 3 or (compute_major == 3 and compute_minor < 5):
|
|
# Not supported
|
|
log.ODM_INFO("CUDA compute platform is not supported (detected: %s.%s but we need at least 3.5)" % (compute_major, compute_minor))
|
|
return False
|
|
except Exception as e:
|
|
log.ODM_INFO("Using CPU for feature extraction: %s" % str(e))
|
|
return False
|
|
|
|
try:
|
|
from opensfm import pypopsift
|
|
return pypopsift.fits_texture(int(width * 1.02), int(height * 1.02))
|
|
except (ModuleNotFoundError, ImportError):
|
|
return False
|
|
except Exception as e:
|
|
log.ODM_WARNING(str(e))
|
|
return False
|
|
|
|
@lru_cache(maxsize=None)
|
|
def get_cuda_compute_version(device_id = 0):
|
|
cuda_lib = "libcuda.so"
|
|
if sys.platform == 'win32':
|
|
cuda_lib = os.path.join(os.environ.get('SYSTEMROOT'), 'system32', 'nvcuda.dll')
|
|
if not os.path.isfile(cuda_lib):
|
|
cuda_lib = "nvcuda.dll"
|
|
|
|
nvcuda = ctypes.cdll.LoadLibrary(cuda_lib)
|
|
|
|
nvcuda.cuInit.argtypes = (ctypes.c_uint32, )
|
|
nvcuda.cuInit.restypes = (ctypes.c_int32)
|
|
|
|
if nvcuda.cuInit(0) != 0:
|
|
raise Exception("Cannot initialize CUDA")
|
|
|
|
nvcuda.cuDeviceGetCount.argtypes = (ctypes.POINTER(ctypes.c_int32), )
|
|
nvcuda.cuDeviceGetCount.restypes = (ctypes.c_int32)
|
|
|
|
device_count = ctypes.c_int32()
|
|
if nvcuda.cuDeviceGetCount(ctypes.byref(device_count)) != 0:
|
|
raise Exception("Cannot get device count")
|
|
|
|
if device_count.value == 0:
|
|
raise Exception("No devices")
|
|
|
|
nvcuda.cuDeviceComputeCapability.argtypes = (ctypes.POINTER(ctypes.c_int32), ctypes.POINTER(ctypes.c_int32), ctypes.c_int32)
|
|
nvcuda.cuDeviceComputeCapability.restypes = (ctypes.c_int32)
|
|
compute_major = ctypes.c_int32()
|
|
compute_minor = ctypes.c_int32()
|
|
|
|
if nvcuda.cuDeviceComputeCapability(ctypes.byref(compute_major), ctypes.byref(compute_minor), device_id) != 0:
|
|
raise Exception("Cannot get CUDA compute version")
|
|
|
|
return (compute_major.value, compute_minor.value)
|
|
|
|
def has_gpu(args):
|
|
if gpu_disabled_by_user_env():
|
|
log.ODM_INFO("Disabling GPU features (ODM_NO_GPU is set)")
|
|
return False
|
|
if args.no_gpu:
|
|
log.ODM_INFO("Disabling GPU features (--no-gpu is set)")
|
|
return False
|
|
|
|
if sys.platform == 'win32':
|
|
nvcuda_path = os.path.join(os.environ.get('SYSTEMROOT'), 'system32', 'nvcuda.dll')
|
|
if os.path.isfile(nvcuda_path):
|
|
log.ODM_INFO("CUDA drivers detected")
|
|
return True
|
|
else:
|
|
log.ODM_INFO("No CUDA drivers detected, using CPU")
|
|
return False
|
|
else:
|
|
if shutil.which('nvidia-smi') is not None:
|
|
log.ODM_INFO("nvidia-smi detected")
|
|
return True
|
|
else:
|
|
return False
|