OpenDroneMap-ODM/opendm/gpu.py

import os
import sys
import shutil
from opendm import log
from repoze.lru import lru_cache

def gpu_disabled_by_user():
    return bool(os.environ.get('ODM_NO_GPU'))

@lru_cache(maxsize=None)
def has_popsift_and_can_handle_texsize(width, height):
    try:
        from opensfm import pypopsift
        fits = pypopsift.fits_texture(int(width * 1.025), int(height * 1.025))
        if not fits:
            log.ODM_WARNING("Image size (%sx%spx) would not fit in GPU memory, falling back to CPU" % (width, height))
        return fits
    except (ModuleNotFoundError, ImportError):
        return False
    except Exception as e:
        log.ODM_WARNING(str(e))
        return False

@lru_cache(maxsize=None)
def has_gpu():
    if gpu_disabled_by_user():
        log.ODM_INFO("Disabling GPU features (ODM_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:
            log.ODM_INFO("nvidia-smi not found in PATH, using CPU")
            return False