OpenCV build with CUDA

opencv-build.sh

#!/usr/bin/env bash
#
# ========================================================================
#   Script for compilation and installation of OpenCV with CUDA hardware
#       acceleration for a Python environment on an Ubuntu-based machine.
# ========================================================================
#       Originally written for OpenCV 4.6.0 and Ubuntu 20.04
#
# The script attempts to automatically infer the CUDA version and location
#   (thank you NVIDIA for making this more complicated than necessary!).
#
# === Troubleshooting
#
# A. For build issues with CUDA / NVCC and GCC version errors:
#   1. find CUDA_BIN_DIR with `locate gcc | grep cuda` (usually /usr/local/cuda/bin)
#       MAX_GCC_VERSION=8
#   2. Install correct gcc version
#       sudo apt-get install -y gcc-$MAX_GCC_VERSION g++-$MAX_GCC_VERSION
#   3. Create links
#       sudo ln -s /usr/bin/gcc-$MAX_GCC_VERSION $CUDA_BIN_DIR/gcc
#       sudo ln -s /usr/bin/g++-$MAX_GCC_VERSION $CUDA_BIN_DIR/g++
#       It may be necessary to change default gcc linked to /usr/bin/gcc:
#           sudo mv /usr/bin/gcc /usr/bin/gcc.bak
#           sudo ln -s /usr/bin/gcc-$MAX_GCC_VERSION /usr/bin/gcc
#
#  B. Undefined reference to 'culibosInit'
#
#   Solution that worked:
#       Edit ./build/opencv/cmake/FindCUDNN.cmake with contents in:
#       https://github.com/opencv/opencv/issues/18082#issuecomment-673513426
#

set -e

# if script is sourced, don't run it
if [[ "${BASH_SOURCE[0]}" != "${0}" ]]; then
    return
fi

# main function
function main() {
    echo "Starting script..."
    requirement_check
    install_prerequisites
    set_directory_variables
    compile_opencv
    test_commands
    cleanup --dry-run
    cd "$DIR_BACK" || exit 0
}

# Checks if a command or program exists
function exists() {
    if [ -z "$1" ]; then
        echo "Usage: exists <program>"
        return 1
    fi
    command -v "$1" >/dev/null 2>&1
}

# Makes sure the system requirements are met
function requirement_check() {
    # The main pre-requisite is to have the NVIDIA CUDA Toolkit and the NVIDIA CUDA Deep Neural Network library (cuDNN) library installed.
    if ! exists "nvidia-smi"; then
        echo "nvidia-smi is not installed. Please install it first."
        return 1
    else
        echo "nvidia-smi is available."
    fi

    CUDA_CAPABILITY=$(nvidia-smi --query-gpu=compute_cap --format=csv | tail -n 1)
    # if empty, set to 7.5
    if [ -z "$CUDA_CAPABILITY" ]; then
        echo -e "\n\tWARNING :: COULD NOT AUTOMATICALLY GET THE CUDA CAPABILITY. SETTING TO 7.5. THIS NEEDS A MANUAL CHECK.\n"
        # press enter to continue
        read -rp "Press enter to continue..."
        CUDA_CAPABILITY=7.5
    fi
    GPU_MODEL=$(nvidia-smi -L)
    echo "Assuming compute capability of $CUDA_CAPABILITY ($GPU_MODEL)"
    echo -e "\n\t>>> Check https://developer.nvidia.com/cuda-gpus to see if\n\t\tyour hardware matches a capability of $CUDA_CAPABILITY <<<\n\n"

    # Tries to automatically grab cuDNN version and location.
    CUDNN_LOCATION=$(whereis cudnn.h | cut -f2 -d' ')
    echo -e "\tCUDNN Location:  $CUDNN_LOCATION"
    # make sure file exists
    if [ ! -f "$CUDNN_LOCATION" ]; then
        echo "Command 'whereis cudnn.h' did not return a location. Please install cuDNN first."
        return 1
    fi
    CUDNN_VERSION=$(less "$CUDNN_LOCATION" | grep -e "CUDNN_MAJOR\|CUDNN_MINOR\|CUDNN_PATCH" | head -n 3 | cut -d' ' -f3 | tr '\n' '.' | head -c -1)
    # if CUDNN_VERSION is empty, try cudnn_version.h in CUDNN_LOCATION directory
    if [ -z "$CUDNN_VERSION" ]; then
        CUDNN_VERSION=$(less "$(dirname "$CUDNN_LOCATION")/cudnn_version.h" | grep -e "CUDNN_MAJOR\|CUDNN_MINOR\|CUDNN_PATCH" | head -n 3 | cut -d' ' -f3 | tr '\n' '.' | head -c -1)
    fi
    echo -e "\tCUDNN Version is $CUDNN_VERSION"
    # CUDNN_LIBRARY=$(whereis libcudnn_static_v | cut -f2 -d' ')
    # # if blank or has a colon, try something else
    # if [ -z "$CUDNN_LIBRARY" ] || [[ "$CUDNN_LIBRARY" == *":"* ]]; then
    #     # CUDNN_LIBRARY="$(dirname "$(find / -name "*libcudnn*" -maxdepth 4 2>/dev/null | grep static | head -n 1)")"
    #     CUDNN_LIBRARY="$(find / -name "*libcudnn*" -maxdepth 4 2>/dev/null | grep static | head -n 1)"
    # fi
    # # make sure file exists
    # echo -e "\tCUDNN Library is $CUDNN_LIBRARY"
    # if [ ! -f "$CUDNN_LIBRARY" ]; then
    #     echo "$CUDNN_LIBRARY is an invalid location. Please install cuDNN first."
    #     return 1
    # fi
    CUDNN_INCLUDE_DIR=$(dirname "$CUDNN_LOCATION")
    echo -e "\tInclude directory is $CUDNN_INCLUDE_DIR"

    echo "System requirements are met."

    export CUDA_CAPABILITY
    export CUDNN_VERSION
    export CUDNN_LOCATION
    # export CUDNN_LIBRARY

}

# Installs the necessary dependencies
function install_prerequisites() {

    # asks for sudo permission
    sudo echo "Superuser permission granted to install requirements" || return 0

    # basic tooling:
    sudo apt-get install -y git cmake gcc g++ curl wget build-essential

    # required:
    sudo apt-get install -y libavcodec-dev libavformat-dev libswscale-dev \
        libgstreamer-plugins-base1.0-dev libgstreamer1.0-dev libgtk-3-dev

    # optional:
    sudo apt-get install -y libpng-dev libjpeg-dev libopenexr-dev libtiff-dev libwebp-dev

}

# Compiles the OpenCV library
function compile_opencv() {

    DIR_BUILD="$DIR_SCRIPT/build"
    OPENCV_VERSION=4.6.0

    export DIR_BUILD
    export OPENCV_VERSION

    # create build directory
    mkdir -p "$DIR_BUILD"
    cd "$DIR_BUILD" || exit 1

    # clone repositories
    if ! [ -d "$DIR_BUILD/opencv" ]; then
        git clone https://github.com/opencv/opencv.git "$DIR_BUILD/opencv"
    fi
    if ! [ -d "$DIR_BUILD/opencv-contrib" ]; then
        git clone https://github.com/opencv/opencv_contrib.git "$DIR_BUILD/opencv-contrib"
    fi

    # checkout correct versions
    for repo in opencv opencv-contrib; do
        cd "$DIR_BUILD/$repo" || exit 1
        git checkout "$OPENCV_VERSION"
    done

    # get virtual environment path information
    DIR_POETRY_ENV=$(poetry env info -p)
    if ! [ -d "$DIR_POETRY_ENV" ]; then
        echo "Poetry environment dir not found. Please run 'poetry install' first."
        exit 3
    fi
    echo " >> Using installation prefix $DIR_POETRY_ENV"

    # start opencv build
    DIR_CV2_BUILD="$DIR_BUILD/opencv/build"
    mkdir -p "$DIR_CV2_BUILD"
    cd "$DIR_CV2_BUILD" || exit 1
    cmake \
        -D CMAKE_BUILD_TYPE=RELEASE \
        -D CMAKE_INSTALL_PREFIX="$DIR_POETRY_ENV" \
        -D WITH_CUDA=ON \
        -D WITH_CUDNN=ON \
        -D WITH_CUBLAS=ON \
        -D WITH_FFMPEG=ON \
        -D WITH_TBB=ON \
        -D CUDA_FAST_MATH=ON \
        -D ENABLE_FAST_MATH=ON \
        -D OPENCV_DNN_CUDA=ON \
        -D OPENCV_ENABLE_NONFREE=ON \
        -D CUDA_ARCH_BIN=$CUDA_CAPABILITY \
        -D OPENCV_EXTRA_MODULES_PATH="$DIR_BUILD/opencv-contrib/modules" \
        -D BUILD_EXAMPLES=OFF \
        -D PYTHON_EXECUTABLE="$(which python)" \
        -D PYTHON_DEFAULT_EXECUTABLE="$(which python)" \
        -D PYTHON3_EXECUTABLE="$(which python3)" \
        -D PYTHON3_INCLUDE_DIR="$(python -c "from distutils.sysconfig import get_python_inc; print(get_python_inc())")" \
        -D PYTHON3_PACKAGES_PATH="$(python -c "from distutils.sysconfig import get_python_lib; print(get_python_lib())")" \
        -D PYTHON_INCLUDE_DIR="$(python -c "from distutils.sysconfig import get_python_inc; print(get_python_inc())")" \
        -D PYTHON_INCLUDE_DIR2="$(python -c "from os.path import dirname; from distutils.sysconfig import get_config_h_filename; print(dirname(get_config_h_filename()))")" \
        -D PYTHON_LIBRARY="$(python -c "from distutils.sysconfig import get_config_var;from os.path import dirname,join ; print(join(dirname(get_config_var('LIBPC')),get_config_var('LDLIBRARY')))")" \
        -D PYTHON3_NUMPY_INCLUDE_DIRS="$(python -c "import numpy; print(numpy.get_include())")" \
        -D PYTHON3_PACKAGES_PATH="$(python -c "from distutils.sysconfig import get_python_lib; print(get_python_lib())")" \
        -D BUILD_opencv_python3=ON \
        -D WITH_OPENGL=ON \
        -D WITH_OPENCL=ON \
        -D WITH_IPP=ON \
        -D WITH_TBB=ON \
        -D WITH_EIGEN=ON \
        -D WITH_V4L=ON \
        -D WITH_NVCUVID=ON \
        -D CUDNN_VERSION="$CUDNN_VERSION" \
        -D INSTALL_PYTHON_EXAMPLES=ON \
        ..

    # Options reference: https://docs.opencv.org/4.x/db/d05/tutorial_config_reference.html
    # Other (disabled) flags:
    # -D CMAKE_INSTALL_PREFIX="$(python -c "import sys; print(sys.prefix)")" \
    # -D CUDNN_INCLUDE_DIR="$CUDNN_INCLUDE_DIR" \

    make -j "${N_CORES}"
    make install

    # create symbolic links to the OpenCV bindings for Python 3 to be used by the environment.
    PYTHON_VERSION=$(python --version | xargs | cut -d' ' -f2 | cut -d '.' -f1,2)
    DIR_PYTHON_LIB="$DIR_POETRY_ENV/lib/python$PYTHON_VERSION"
    # make sure DIR_PYTHON_LIB exists
    if ! [ -d "$DIR_PYTHON_LIB" ]; then
        echo -e "\n\tError: $DIR_PYTHON_LIB does not exist -- check the script. Wrong Python version?"
        exit 2
    fi
    mkdir -p "$DIR_PYTHON_LIB/dist-packages"
    if ! [ -L "$DIR_PYTHON_LIB/dist-packages/cv2" ]; then
        ln -s "$DIR_PYTHON_LIB/site-packages/cv2" "$DIR_PYTHON_LIB/dist-packages/cv2"
    fi

}

function set_directory_variables() {
    DIR_BACK=$(pwd)
    DIR_SCRIPT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
    N_CORES=$(nproc)
    MIN_CORES=2
    BUFFER_CORES=6
    N_CORES=$(((N_CORES - BUFFER_CORES) > MIN_CORES ? (N_CORES - BUFFER_CORES) : MIN_CORES))
    echo "Using $N_CORES cores for building processes."
    cd "$DIR_SCRIPT" || exit 1
    export DIR_BACK
    export DIR_SCRIPT
}

function test_commands() {
    python -c "import cv2; print(cv2.__version__); [ print(cv2.videoio_registry.getBackendName(x)) for x in cv2.videoio_registry.getBackends()]"
}

# cleanup script
function cleanup() {
    dry_run="${1:---dry-run}"
    if [ "$dry_run" != "--delete" ]; then
        echo "Dry run. No files will be deleted."
    fi
    declare -a dir_list=(
        "${DIR_BUILD:-${DIR_SCRIPT}/build}/opencv"
        "${DIR_BUILD:-${DIR_SCRIPT}/build}/opencv-contrib"
    )
    for dir in "${dir_list[@]}"; do
        if [ -d "$dir" ]; then
            if [ "$dry_run" == "--delete" ]; then
                # make sure the directory is not home, root, or current
                if [ "$dir" != "$HOME" ] && [ "$dir" != "/" ] && [ "$dir" != "." ]; then
                    echo "Removing '$dir'"
                    # commented out for safety:
                    # rm -rf "$dir"
                fi
            else
                echo "Removing '$dir' (dry run)"
            fi
        fi
    done
}

main