diff --git a/applications/orthorectification_with_optix/Dockerfile b/applications/orthorectification_with_optix/Dockerfile new file mode 100755 index 000000000..b54a35d60 --- /dev/null +++ b/applications/orthorectification_with_optix/Dockerfile @@ -0,0 +1,59 @@ +# syntax = docker/dockerfile:1.2 + +ARG FROM_IMAGE="nvcr.io/nvidia/clara-holoscan/holoscan:v0.6.0-dgpu" + +# ============ Stage: base ============ +FROM ${FROM_IMAGE} AS base + +# Configure the base conda environment +ARG CONDA_ENV_NAME=ortho_holoscan +ARG PYTHON_VER=3.8 + +# Install miniconda +ENV CONDA_DIR /opt/conda +RUN wget --quiet https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda.sh && \ + /bin/bash ~/miniconda.sh -b -p /opt/conda + +# Install libarchive +RUN apt update && DEBIAN_FRONTEND=noninteractive apt install -y libarchive13 + +# Put conda in path so we can use conda activate +ENV PATH=$CONDA_DIR/bin:$PATH +RUN conda init bash +RUN ln -sf /bin/bash /bin/sh + +# Install mamba to speed the solve up +RUN conda config --set ssl_verify false &&\ + conda config --add pkgs_dirs /opt/conda/pkgs &&\ + conda config --env --add channels conda-forge &&\ + /opt/conda/bin/conda install -y -n base -c conda-forge "mamba >=0.22" "boa >=0.10" python=${PYTHON_VER} + +# Create a base environment +RUN --mount=type=cache,id=conda_pkgs,target=/opt/conda/pkgs,sharing=locked \ + # Create the environment and install as little dependencies as possible + CONDA_ALWAYS_YES=true /opt/conda/bin/mamba create -n ${CONDA_ENV_NAME} -c conda-forge python=${PYTHON_VER} + +RUN source activate ${CONDA_ENV_NAME} && mamba install -c conda-forge -y numpy cupy pillow pytest gdal shapely && \ + pip install pynvrtc tqdm opencv-python odm_report_shot_coverage + +# MUST RUN BUILD COMMAND FROM ABOVE HOLOHUB-INTERNAL REPO +WORKDIR /work +COPY ./NVIDIA-OptiX-SDK-7.4.0-linux64-x86_64 /work/NVIDIA-OptiX-SDK-7.4.0-linux64-x86_64 + +RUN git clone https://github.com/NVIDIA/optix-toolkit.git +WORKDIR /work/optix-toolkit +RUN git checkout v0.8.1 +RUN git submodule update --init --recursive PyOptiX +WORKDIR /work/optix-toolkit/PyOptiX/optix + +RUN source activate ${CONDA_ENV_NAME} &&\ + export PYOPTIX_CMAKE_ARGS="-DOptiX_INSTALL_DIR=/work/NVIDIA-OptiX-SDK-7.4.0-linux64-x86_64" &&\ + export PYOPTIX_STDDEF_DIR="/usr/include/linux" &&\ + python setup.py install +WORKDIR /work + +# now install holoscan python bindings +RUN source activate ${CONDA_ENV_NAME} && pip install holoscan +WORKDIR /work + +RUN echo "conda activate ${CONDA_ENV_NAME}" >> ~/.bashrc \ No newline at end of file diff --git a/applications/orthorectification_with_optix/README.md b/applications/orthorectification_with_optix/README.md new file mode 100644 index 000000000..a2a7374b7 --- /dev/null +++ b/applications/orthorectification_with_optix/README.md @@ -0,0 +1,56 @@ +# HoloHub Orthorectification Application + +This application is an example of utilizing the nvidia OptiX SDK via the PyOptix bindings to create per-frame orthorectified imagery. In this example, one can create a visualization of mapping frames from a drone mapping mission processed with [Open Drone Map](https://opendronemap.org/). A typical output of a mapping mission is a single merged mosaic. While this product is useful for GIS applications, it is difficult to apply algorithms on a such a large single image without incurring additional steps like image chipping. Additionally, the mosaic process introduces image artifacts which can negativley impact algorithm performance. + +Since this holoscan pipeline processes each frame individually, it opens the door for one to apply an algorithm to the original un-modififed imagery and then map the result. If custom image processing is desired, it is recommended to insert custom operators before the Ray Trace Ortho operator in the application flow. + + +![](docs/odm_ortho_pipeline.png)
+Fig. 1 Orthorectification sample application workflow + +Steps for running the application: + +a) Download and Prep the ODM Dataset
+1. Download the [Lafayette Square Dataset](https://www.opendronemap.org/odm/datasets/) and place into ~/Data. + +2. Process the dataset with ODM via docker command:
+```docker run -ti --rm -v ~/Data/lafayette_square:/datasets/code opendronemap/odm --project-path /datasets --camera-lens perspective --dsm``` + +If you run out of memory add the following argument to preserve some memory: ```--feature-quality medium``` + +b) Clone holohub and navigate to this application directory + +c) Download [OptiX SDK 7.4.0](https://developer.nvidia.com/optix/downloads/7.4.0/linux64-x86_64) and extract the package in the same directory as the source code +(i.e. applications/orthorectification_with_optix). + +d) Build development container
+1. ```DOCKER_BUILDKIT=1 docker build -t holohub-ortho-optix:latest .``` + +You can now run the docker container by:
+1. ```xhost +local:docker``` +2. ```nvidia_icd_json=$(find /usr/share /etc -path '*/vulkan/icd.d/nvidia_icd.json' -type f 2>/dev/null | grep .) || (echo "nvidia_icd.json not found" >&2 && false)``` +3. ```docker run -it --rm --net host --runtime=nvidia -v ~/Data:/root/Data -v .:/work/ -v /tmp/.X11-unix:/tmp/.X11-unix -v $nvidia_icd_json:$nvidia_icd_json:ro -e NVIDIA_DRIVER_CAPABILITIES=graphics,video,compute,utility,display -e DISPLAY=$DISPLAY holohub-ortho-optix``` + +Finish prepping the input data:
+1. ```gdal_translate -tr 0.25 0.25 -r cubic ~/Data/lafayette_square/odm_dem/dsm.tif ~/Data/lafayette_square/odm_dem/dsm_small.tif``` +2. ```gdal_fillnodata.py -md 0 ~/Data/lafayette_square/odm_dem/dsm_small.tif ~/Data/lafayette_square/odm_dem/dsm_small_filled.tif``` + +Finally run the application:
+1. ```python ./python/ortho_with_pyoptix.py``` + +You can modify the applications settings in the file "ortho_with_pyoptix.py" + +``` +sensor_resize = 0.25 # resizes the raw sensor pixels +ncpu = 8 # how many cores to use to load sensor simulation +gsd = 0.25 # controls how many pixels are in the rendering +iterations = 425 # how many frames to render from the source images (in this case 425 is max) +use_mosaic_bbox = True # render to a static bounds on the ground as defined by the DEM +write_geotiff = False +nb=3 # how many bands to write to the GeoTiff +render_scale = 0.5 # scale the holoview window up or down +fps = 8.0 # rate limit the simulated sensor feed to this many frames per second +``` + +![](docs/holohub_ortho_app.gif)
+Fig. 2 Running the orthorectification sample application \ No newline at end of file diff --git a/applications/orthorectification_with_optix/cpp/src/optix/optixOrtho.cu b/applications/orthorectification_with_optix/cpp/src/optix/optixOrtho.cu new file mode 100644 index 000000000..4613af412 --- /dev/null +++ b/applications/orthorectification_with_optix/cpp/src/optix/optixOrtho.cu @@ -0,0 +1,152 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include + +#include "optixOrtho.h" +#include + +#include + +extern "C" { +__constant__ Params params; +} + +static __forceinline__ __device__ void setPayload(float3 p) { + optixSetPayload_0(float_as_int(p.x)); + optixSetPayload_1(float_as_int(p.y)); + optixSetPayload_2(float_as_int(p.z)); +} + +static __forceinline__ __device__ float2 +computeSensorUV(const float3 hit_point) { + const float3 origin_to_hit_point = + hit_point - params.sensor_focal_plane_origin; + const float dist_x = length(cross(params.sensor_up, origin_to_hit_point)); + const float dist_y = length(cross(params.sensor_right, origin_to_hit_point)); + return make_float2(dist_x, dist_y) / params.sensor_focal_plane_size; +} + +static __device__ __inline__ float3 traceSensorRay(float3 ray_origin, + float3 ray_direction) { + unsigned int p0 = 0, p1 = 0, p2 = 0; + optixTrace(params.handle, ray_origin, ray_direction, + 0.0f, // Min intersection distance + 1e16f, // Max intersection distance + 0.0f, // rayTime -- used for motion blur + OptixVisibilityMask(255), // Specify always visible + OPTIX_RAY_FLAG_NONE, + 0, // SBT offset -- See SBT discussion + 1, // SBT stride -- See SBT discussion + 0, // missSBTIndex -- See SBT discussion + p0, p1, p2); + + const float3 sensor_col = + make_float3(int_as_float(p0), int_as_float(p1), int_as_float(p2)); + return sensor_col; +} + +extern "C" __global__ void __raygen__rg() { + // Lookup our location within the launch grid + const uint3 idx = optixGetLaunchIndex(); + const uint3 dim = optixGetLaunchDimensions(); + + // for nearest + const float3 ray_origin = make_float3( + idx.x * params.image_gsd + params.image_corner_coords.x, + idx.y * params.image_gsd + params.image_corner_coords.y, -10.f); + + const float3 ray_direction = make_float3(0.f, 0.f, 1.f); + + // Trace the ray against our scene hierarchy + unsigned int p0, p1, p2; + optixTrace(params.handle, ray_origin, ray_direction, + 0.0f, // Min intersection distance + 1e16f, // Max intersection distance + 0.0f, // rayTime -- used for motion blur + OptixVisibilityMask(255), // Specify always visible + OPTIX_RAY_FLAG_NONE, + 0, // SBT offset -- See SBT discussion + 1, // SBT stride -- See SBT discussion + 0, // missSBTIndex -- See SBT discussion + p0, p1, p2); + + if ((p0 > 0) && (p1 > 0) && (p2 > 0)) { + const float3 result = + make_float3(int_as_float(p0), int_as_float(p1), int_as_float(p2)); + + const uchar4 clr = make_color(result); + // Record results in the output raster + params.image[idx.y * params.image_width + idx.x] = clr; + } +} + +extern "C" __global__ void __miss__ms() { + const float3 sensor_col = make_float3(0., 0., 0.); + setPayload(sensor_col); +} + +extern "C" __global__ void __closesthit__terrain_ch() { + // When built-in triangle intersection is used, a number of fundamental + // attributes are provided by the OptiX API, indlucing barycentric + // coordinates. + const float3 ray_orig = optixGetWorldRayOrigin(); + const float3 ray_dir = optixGetWorldRayDirection(); // incident direction + const float ray_t = optixGetRayTmax(); + + // Lookup our location within the launch grid + const uint3 idx = optixGetLaunchIndex(); + const uint3 dim = optixGetLaunchDimensions(); + + const float3 hit_point = ray_orig + ray_t * ray_dir; + const int index = idx.y * params.image_width + idx.x; + + if (hit_point.z < params.terrain_zmax) { // We hit terrain, cast ray to sensor + if (ray_orig.z > 0.f) { + const float3 sensor_col = make_float3(0., 0., 0.); + setPayload(sensor_col); + } else { + // // if you want to cast from sensor back to terrain + // const float3 to_terrain = normalize(params.sensor_pos - hit_point); + // const float3 sensor_col = traceSensorRay(params.sensor_pos, + // to_terrain); + + // from terrain to sensor + const float3 to_sensor = normalize(params.sensor_pos - hit_point); + const float3 sensor_col = + traceSensorRay(hit_point + 0.01f * to_sensor, to_sensor); + + setPayload(sensor_col); + } + } else { // We hit the sensor plane + const float2 sensor_uv = computeSensorUV(hit_point); + // for nearest lookup + const uchar4 sensor_rgba = + tex2D(params.sensor_tex, sensor_uv.x, sensor_uv.y); + const float3 sensor_col = + make_float3(sensor_rgba.x, sensor_rgba.y, sensor_rgba.z) / 255.; + + // // for linear lookup + // const float4 sensor_rgba = tex2D(params.sensor_tex, sensor_uv.x, + // sensor_uv.y); const float3 sensor_col = make_float3(sensor_rgba.x, + // sensor_rgba.y, sensor_rgba.z); + + setPayload(sensor_col); + } +} diff --git a/applications/orthorectification_with_optix/cpp/src/optix/optixOrtho.h b/applications/orthorectification_with_optix/cpp/src/optix/optixOrtho.h new file mode 100644 index 000000000..3afb7e2d1 --- /dev/null +++ b/applications/orthorectification_with_optix/cpp/src/optix/optixOrtho.h @@ -0,0 +1,45 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +struct Params { + OptixTraversableHandle handle; + cudaTextureObject_t sensor_tex; + uchar4 *image; + unsigned int image_width; + unsigned int image_height; + float2 image_corner_coords; + float image_gsd; + float sensor_focal_length; + float terrain_zmax; + float3 sensor_pos; + float3 sensor_up; + float3 sensor_right; + float3 sensor_focal_plane_origin; + float2 sensor_focal_plane_size; +}; + +struct RayGenData { + // No data needed +}; + +struct MissData { + float3 bg_color; +}; + +struct HitGroupData { + // No data needed +}; diff --git a/applications/orthorectification_with_optix/docs/holohub_ortho_app.gif b/applications/orthorectification_with_optix/docs/holohub_ortho_app.gif new file mode 100644 index 000000000..104197b68 Binary files /dev/null and b/applications/orthorectification_with_optix/docs/holohub_ortho_app.gif differ diff --git a/applications/orthorectification_with_optix/docs/odm_ortho_pipeline.png b/applications/orthorectification_with_optix/docs/odm_ortho_pipeline.png new file mode 100755 index 000000000..6e5c82aa0 Binary files /dev/null and b/applications/orthorectification_with_optix/docs/odm_ortho_pipeline.png differ diff --git a/applications/orthorectification_with_optix/python/metadata.json b/applications/orthorectification_with_optix/python/metadata.json new file mode 100644 index 000000000..0702280b8 --- /dev/null +++ b/applications/orthorectification_with_optix/python/metadata.json @@ -0,0 +1,39 @@ +{ + "application": { + "name": "Endoscopy Out of Body Detection Pipeline in C++", + "authors": [ + { + "name": "Brent Bartlett", + "affiliation": "NVIDIA" + } + ], + "language": "Python", + "version": "1.0", + "changelog": { + "1.0": "Initial Release" + }, + "holoscan_sdk": { + "minimum_required_version": "0.6.0", + "tested_versions": [ + "0.6.0" + ] + }, + "platforms": [ + "amd64" + ], + "tags": [ + "Orthorectification", + "Drone", + "OptiX" + ], + "ranking": 4, + "dependencies": { + "OptiX-SDK": { + "version": "4.7.0" + }, + "OptiX-Toolit": { + "version": "0.8.1" + } + } + } +} diff --git a/applications/orthorectification_with_optix/python/odm_utils.py b/applications/orthorectification_with_optix/python/odm_utils.py new file mode 100644 index 000000000..740d5d945 --- /dev/null +++ b/applications/orthorectification_with_optix/python/odm_utils.py @@ -0,0 +1,183 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import os +from functools import partial +from multiprocessing import Pool + +import cv2 +import numpy as np +import tqdm +from numpy.linalg import norm +from odm_report_shot_coverage.models.reconstruction import parse_reconstruction +from ortho_utils import extract_extent_nativeCRS, ray_plane_intersections_cpu +from osgeo import gdal +from shapely.geometry import MultiPoint, Point, box + + +def norm2uv_odm(norm_coords, shot): + x_n, y_n = norm_coords + r_2 = x_n * x_n + y_n * y_n + d = 1 + r_2 * shot.camera.k1 + r_2 * r_2 * shot.camera.k2 + return shot.camera.focal * d * x_n, shot.camera.focal * d * y_n + + +def load_frame_odm( + reconstruction, + offsets, + mosaic_bbox, + use_mosaic_bbox, + data_dir, + sensor_resize, + sensor_dims, + min_el, + fnum, +): + shot = reconstruction.shots[fnum] + utm_sensor_pos = [ + shot.translation[0] + offsets[0], + shot.translation[1] + offsets[1], + shot.translation[2], + ] + + utm_sensor_pos_str = ",".join([str(item) for item in utm_sensor_pos]) + + uvx = norm2uv_odm((0.5, 0), shot)[0] + scale_ratio = shot.camera.height / shot.camera.width + uvy = norm2uv_odm((0, scale_ratio * 0.5), shot)[1] + + hfov = 2 * np.rad2deg(np.arctan((uvx / shot.camera.focal) / shot.camera.focal)) + vfov = 2 * np.rad2deg(np.arctan((uvy / shot.camera.focal) / shot.camera.focal)) + + # to get direction vector - start with camera coords at x,y,z = (0,0,1), then rotate + rc = [0, 0, 1] + sensor_dir = shot._transfo_rotation.apply(rc, inverse=True) + + # to get sensor up vector - start with camera coords at x,y,z = (0, -1, 0), then rotate back + rc = [0, -1, 0] + sensor_up = shot._transfo_rotation.apply(rc, inverse=True) + + sensor_dir_str = ",".join([str(item) for item in sensor_dir]) + sensor_up_str = ",".join([str(item) for item in sensor_up]) + + # use camera model to get bounding box estimates in world coordinates assuming + # flat plane elevation with min elevation + if not use_mosaic_bbox: + # vector in direction of corners + x_edge = shot.camera.focal * np.tan(np.deg2rad(hfov / 2.0)) + y_edge = shot.camera.focal * np.tan(np.deg2rad(vfov / 2.0)) + + bot_right = np.array( + shot._transfo_rotation.apply([x_edge, y_edge, shot.camera.focal], inverse=True) + ) + bot_right = bot_right / norm(bot_right) + + bot_left = np.array( + shot._transfo_rotation.apply([-x_edge, y_edge, shot.camera.focal], inverse=True) + ) + bot_left = bot_left / norm(bot_left) + + top_left = np.array( + shot._transfo_rotation.apply([-x_edge, -y_edge, shot.camera.focal], inverse=True) + ) + top_left = top_left / norm(top_left) + + top_right = np.array( + shot._transfo_rotation.apply([x_edge, -y_edge, shot.camera.focal], inverse=True) + ) + top_right = top_right / norm(top_right) + + corners = np.vstack([bot_right, bot_left, top_left, top_right, bot_right]) + + utm_corners = ray_plane_intersections_cpu( + np.array([0, 0, 1]), np.array([0, 0, min_el]), corners, np.array(utm_sensor_pos) + ) + + bpts = MultiPoint(list(utm_corners)) + bbox = bpts.convex_hull + else: + bbox = mosaic_bbox + + camera_loc = Point(utm_sensor_pos) + + meta = { + "wkt": bbox.wkt, + "wkt_camera_loc": camera_loc.wkt, + "image_name": shot.image_name, + "sensor_dir": sensor_dir_str, + "sensor_up": sensor_up_str, + "utm_pos": utm_sensor_pos_str, + "hfov": str(hfov), + "vfov": str(vfov), + "frame_number": fnum, + } + + sensor_image_fname = os.path.join(data_dir, "images", meta["image_name"]) + sensor_pix = cv2.imread(sensor_image_fname) + if sensor_resize is not None: + sensor_pix = cv2.resize(sensor_pix, dsize=sensor_dims, interpolation=cv2.INTER_AREA) + sensor_pix = cv2.cvtColor(sensor_pix, cv2.COLOR_BGR2RGB) + + return {"sensor_pix": sensor_pix, "meta": meta} + + +def load_all_frames_odm( + use_mosaic_bbox, + dem_image_fname, + geo_path, + data_dir, + iterations, + ncpu, + sensor_resize, + sensor_dims, +): + dem_info = gdal.Info(dem_image_fname, options=gdal.InfoOptions(format="json", stats=True)) + min_el = float(dem_info["bands"][0]["metadata"][""]["STATISTICS_MINIMUM"]) + + with open(geo_path) as fp: + _ = fp.readline() + offsets = [int(item) for item in fp.readline().strip().split(" ")] + + reconstruction = parse_reconstruction(data_dir) + + bbox = None + if use_mosaic_bbox: + bbox_bounds = extract_extent_nativeCRS(dem_image_fname) + bbox = box(bbox_bounds[0], bbox_bounds[1], bbox_bounds[2], bbox_bounds[3]) + + staged_dataset = {} + frame_nums = list(range(iterations)) + func = partial( + load_frame_odm, + reconstruction, + offsets, + bbox, + use_mosaic_bbox, + data_dir, + sensor_resize, + sensor_dims, + min_el, + ) + + pool = Pool(ncpu) + for item in tqdm.tqdm( + pool.imap(func, frame_nums), desc="simulating sensor feed", total=iterations + ): + staged_dataset[item["meta"]["frame_number"]] = item + pool.close() + pool.join() + + return staged_dataset diff --git a/applications/orthorectification_with_optix/python/optix_utils_for_ortho.py b/applications/orthorectification_with_optix/python/optix_utils_for_ortho.py new file mode 100755 index 000000000..e262a6005 --- /dev/null +++ b/applications/orthorectification_with_optix/python/optix_utils_for_ortho.py @@ -0,0 +1,706 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import ctypes # C interop helpers +import os + +import cupy as cp +import numpy as np +import optix +from osgeo import gdal +from pynvrtc.compiler import Program + + +class State: + def __init__(self): + self.bounding_box_estimate = None + self.utm_sensor_pos = None + self.frame_number = None + self.image = None + self.image_width = None + self.image_height = None + self.image_corner_coords = None + self.image_gsd = None + self.image_geot = None + self.image_wkt = None + + self.terrain_pix = None + self.terrain_width = None + self.terrain_height = None + self.terrain_gsd = None + self.terrain_geot = None + self.terrain_wkt = None + self.terrain_xmin = None + self.terrain_ymin = None + self.terrain_zmax = None + self.terrain_vert_grid = None + self.terrain_tri_grid = None + + self.sensor_isRGB = None + self.sensor_hor_fov = None + self.sensor_ver_fov = None + self.sensor_focal_length = None + self.sensor_pos = None + self.sensor_lookat = None + self.sensor_dir = None + self.sensor_up = None + self.sensor_right = None + self.sensor_focal_plane_origin = None + self.sensor_focal_plane_size = None + self.sensor_focal_plane_center = None + self.half_focal_plane_x = None + self.half_focal_plane_y = None + self.sensor_tex_pixels = None + self.sensor_tex_width = None + self.sensor_tex_height = None + self.sensor_vert_grid = None + self.sensor_tri_grid = None + + self.sensor_alpha = None + + +################### +## Terrain Stuff +################### + + +def build_terrain_geom(state: State): + # build verts + x = ( + np.arange(0, state.terrain_width) * state.terrain_geot[1] + + state.terrain_geot[0] + - state.terrain_xmin + ) + y = ( + np.arange(0, state.terrain_height) * state.terrain_geot[5] + + state.terrain_geot[3] + - state.terrain_ymin + ) + xx, yy = np.meshgrid(x, y) + zz = state.terrain_pix + vert_grid = np.vstack((xx, yy, zz)).reshape([3, -1]).transpose() + + # build triangles + ai = np.arange(0, state.terrain_width - 1) + aj = np.arange(0, state.terrain_height - 1) + aii, ajj = np.meshgrid(ai, aj) + a = aii + ajj * state.terrain_width + a = a.flatten() + + tria = np.vstack( + ( + a, + a + state.terrain_width, + a + state.terrain_width + 1, + a, + a + state.terrain_width + 1, + a + 1, + ) + ) + tri_grid = np.transpose(tria).reshape([-1, 3]) + + state.terrain_vert_grid = vert_grid.reshape(vert_grid.size) + state.terrain_tri_grid = tri_grid.reshape(tri_grid.size) + + +def load_terrain_host(state: State, terrain_fname): + dset = gdal.Open(terrain_fname, gdal.GA_ReadOnly) + + nx = dset.RasterXSize + ny = dset.RasterYSize + state.terrain_width = nx + state.terrain_height = ny + + geot = dset.GetGeoTransform() + state.terrain_wkt = dset.GetProjection() + + state.terrain_geot = geot + state.terrain_xmin = geot[0] + geot[1] + state.terrain_ymin = geot[3] + (ny - 1) * geot[5] + state.terrain_gsd = geot[1] + + state.terrain_pix = dset.GetRasterBand(1).ReadAsArray() + state.terrain_zmax = np.max(state.terrain_pix) + + dset = None + + +################### +## Sensor Stuff +################### + + +def build_sensor_geom(state: State): + vert_grid = [] + tri_grid = [] + + # top left corner + for i in range(3): + vert_grid.append( + state.sensor_focal_plane_center[i] + - state.half_focal_plane_x[i] + + state.half_focal_plane_y[i] + ) + + # top right corner + for i in range(3): + vert_grid.append( + state.sensor_focal_plane_center[i] + + state.half_focal_plane_x[i] + + state.half_focal_plane_y[i] + ) + + # bottom left corner + for i in range(3): + vert_grid.append( + state.sensor_focal_plane_center[i] + - state.half_focal_plane_x[i] + - state.half_focal_plane_y[i] + ) + + # bottom right corner + for i in range(3): + vert_grid.append( + state.sensor_focal_plane_center[i] + + state.half_focal_plane_x[i] + - state.half_focal_plane_y[i] + ) + + # lower left triangle + tri_grid.extend([0, 3, 2]) + tri_grid.extend([0, 1, 3]) + + vert_grid = np.array(vert_grid, dtype=np.float32) + tri_grid = np.array(tri_grid, dtype=np.uint32) + + state.sensor_vert_grid = vert_grid + state.sensor_tri_grid = tri_grid + + +def crop_center_image(img, cropx, cropy): + y, x, nb = img.shape + startx = x // 2 - (cropx // 2) + starty = y // 2 - (cropy // 2) + cropped = [] + for bn in range(nb): + cropped.append(img[starty : starty + cropy, startx : startx + cropx, bn]) + + return np.dstack(cropped) + + +def load_sensor_texture(state: State, impix): + if not state.sensor_isRGB: + impix = crop_center_image(impix, state.sensor_tex_width, state.sensor_tex_height) + state.sensor_tex_pixels = impix + + +################### +## Rendering Params and math helpers +################### + + +def init_render_params(state: State): + xmin, ymin, xmax, ymax = state.bounding_box_estimate + + box_size_x = xmax - xmin + box_size_y = ymax - ymin + + state.image_width = int(box_size_x / state.image_gsd) + state.image_height = int(box_size_y / state.image_gsd) + + state.image_corner_coords = np.zeros(2) + state.image_corner_coords[0] = xmin - state.terrain_xmin + state.image_corner_coords[1] = ymin - state.terrain_ymin + + state.image_geot = np.zeros(6) + state.image_geot[0] = xmin + state.terrain_gsd / 2.0 + state.image_geot[3] = ymax - state.terrain_gsd / 2.0 + + state.image_geot[1] = state.image_gsd + state.image_geot[5] = -1.0 * state.image_gsd + + state.image_geot[2] = 0.0 + state.image_geot[4] = 0.0 + + # TODO fix in case these are different + state.image_wkt = state.terrain_wkt + + state.sensor_pos = np.array( + [ + state.utm_sensor_pos[0] - state.terrain_xmin, + state.utm_sensor_pos[1] - state.terrain_ymin, + state.utm_sensor_pos[2], + ] + ) + + cross_p = np.cross(state.sensor_up, state.sensor_dir) + state.sensor_right = cross_p / np.linalg.norm(cross_p) + + # need to define focal length in our virtual camera such that we maximize the + # focal plane size - this gives us better accuracy in our hit locations + # so subtract z location of sensor from max terrain height as an estimate + # TODO more elegant way to slide this down towards terrain + state.sensor_focal_length = (state.utm_sensor_pos[2] - state.terrain_zmax) * 0.5 + state.sensor_focal_plane_size = np.array( + [ + state.sensor_focal_length * 2.0 * np.tan(state.sensor_hor_fov / 2.0), + state.sensor_focal_length * 2.0 * np.tan(state.sensor_ver_fov / 2.0), + ] + ) + + state.sensor_focal_plane_center = ( + state.sensor_pos - state.sensor_dir * state.sensor_focal_length + ) + state.half_focal_plane_x = state.sensor_right * state.sensor_focal_plane_size[0] * 0.5 + state.half_focal_plane_y = state.sensor_up * state.sensor_focal_plane_size[1] * 0.5 + state.sensor_focal_plane_origin = ( + state.sensor_focal_plane_center - state.half_focal_plane_x - state.half_focal_plane_y + ) + + +################### +## Optix Setup +################### +class Logger: + def __init__(self): + self.num_mssgs = 0 + + def __call__(self, level, tag, mssg): + print("[{:>2}][{:>12}]: {}".format(level, tag, mssg)) + self.num_mssgs += 1 + + +def log_callback(level, tag, mssg): + print("[{:>2}][{:>12}]: {}".format(level, tag, mssg)) + + +def init_optix(): + print("Initializing cuda ...") + cp.cuda.runtime.free(0) + + print("Initializing optix ...") + optix.init() + + +def round_up(val, mult_of): + return val if val % mult_of == 0 else val + mult_of - val % mult_of + + +def get_aligned_itemsize(formats, alignment): + names = [] + for i in range(len(formats)): + names.append("x" + str(i)) + + temp_dtype = np.dtype({"names": names, "formats": formats, "align": True}) + return round_up(temp_dtype.itemsize, alignment) + + +def optix_version_gte(version): + if optix.version()[0] > version[0]: + return True + if optix.version()[0] == version[0] and optix.version()[1] >= version[1]: + return True + return False + + +def array_to_device_memory(numpy_array, stream=cp.cuda.Stream()): + byte_size = numpy_array.size * numpy_array.dtype.itemsize + + h_ptr = ctypes.c_void_p(numpy_array.ctypes.data) + d_mem = cp.cuda.memory.alloc(byte_size) + d_mem.copy_from_async(h_ptr, byte_size, stream) + return d_mem + + +def create_ctx(): + print("Creating optix device context ...") + + # Note that log callback data is no longer needed. We can + # instead send a callable class instance as the log-function + # which stores any data needed + global logger + logger = Logger() + + # OptiX param struct fields can be set with optional + # keyword constructor arguments. + ctx_options = optix.DeviceContextOptions(logCallbackFunction=logger, logCallbackLevel=4) + + # They can also be set and queried as properties on the struct + if optix.version()[1] >= 2: + ctx_options.validationMode = optix.DEVICE_CONTEXT_VALIDATION_MODE_ALL + + cu_ctx = 0 + return optix.deviceContextCreate(cu_ctx, ctx_options) + + +def compile_cuda( + cuda_file, cuda_tk_path, include_path, project_include_path, stddef_path="/usr/include/linux" +): + with open(cuda_file, "rb") as f: + src = f.read() + nvrtc_dll = os.environ.get("NVRTC_DLL") + if nvrtc_dll is None: + nvrtc_dll = "" + # print("NVRTC_DLL = {}".format(nvrtc_dll)) + prog = Program(src.decode(), cuda_file, lib_name=nvrtc_dll) + compile_options = [ + "-use_fast_math", + "-lineinfo", + "-default-device", + "-std=c++11", + "-rdc", + "true", + #'-IC:\\Program Files\\NVIDIA GPU Computing Toolkit\CUDA\\v11.1\include' + f"-I{cuda_tk_path}", + f"-I{include_path}", + f"-I{project_include_path}", + ] + # Optix 7.0 compiles need path to system stddef.h + # the value of optix.stddef_path is compiled in constant. When building + # the module, the value can be specified via an environment variable, e.g. + # export PYOPTIX_STDDEF_DIR="/usr/include/linux" + if optix.version()[1] == 0: + compile_options.append(f"-I{stddef_path}") + + ptx = prog.compile(compile_options) + return ptx + + +def set_pipeline_options(): + if optix.version()[1] >= 2: + return optix.PipelineCompileOptions( + usesMotionBlur=False, + traversableGraphFlags=int(optix.TRAVERSABLE_GRAPH_FLAG_ALLOW_SINGLE_GAS), + numPayloadValues=3, + numAttributeValues=3, + exceptionFlags=int(optix.EXCEPTION_FLAG_NONE), + pipelineLaunchParamsVariableName="params", + usesPrimitiveTypeFlags=optix.PRIMITIVE_TYPE_FLAGS_TRIANGLE, + ) + else: + return optix.PipelineCompileOptions( + usesMotionBlur=False, + traversableGraphFlags=int(optix.TRAVERSABLE_GRAPH_FLAG_ALLOW_SINGLE_GAS), + numPayloadValues=3, + numAttributeValues=3, + exceptionFlags=int(optix.EXCEPTION_FLAG_NONE), + pipelineLaunchParamsVariableName="params", + ) + + +def create_module(ctx, pipeline_options, triangle_ptx): + print("Creating optix module ...") + + module_options = optix.ModuleCompileOptions( + maxRegisterCount=optix.COMPILE_DEFAULT_MAX_REGISTER_COUNT, + optLevel=optix.COMPILE_OPTIMIZATION_DEFAULT, + debugLevel=optix.COMPILE_DEBUG_LEVEL_DEFAULT, + ) + + module, log = ctx.moduleCreateFromPTX(module_options, pipeline_options, triangle_ptx) + print("\tModule create log: <<<{}>>>".format(log)) + return module + + +def create_program_groups(ctx, module): + print("Creating program groups ... ") + + raygen_prog_group_desc = optix.ProgramGroupDesc() + raygen_prog_group_desc.raygenModule = module + raygen_prog_group_desc.raygenEntryFunctionName = "__raygen__rg" + raygen_prog_group, log = ctx.programGroupCreate([raygen_prog_group_desc]) + print("\tProgramGroup raygen create log: <<<{}>>>".format(log)) + + miss_prog_group_desc = optix.ProgramGroupDesc() + miss_prog_group_desc.missModule = module + miss_prog_group_desc.missEntryFunctionName = "__miss__ms" + optix.ProgramGroupOptions() + miss_prog_group, log = ctx.programGroupCreate([miss_prog_group_desc]) + print("\tProgramGroup miss create log: <<<{}>>>".format(log)) + + hitgroup_prog_group_desc = optix.ProgramGroupDesc() + hitgroup_prog_group_desc.hitgroupModuleCH = module + hitgroup_prog_group_desc.hitgroupEntryFunctionNameCH = "__closesthit__terrain_ch" + hitgroup_prog_group, log = ctx.programGroupCreate([hitgroup_prog_group_desc]) + print("\tProgramGroup hitgroup create log: <<<{}>>>".format(log)) + + return [raygen_prog_group, miss_prog_group, hitgroup_prog_group] + + +def create_pipeline(ctx, program_groups, pipeline_compile_options): + print("Creating pipeline ... ") + + max_trace_depth = 2 + pipeline_link_options = optix.PipelineLinkOptions() + pipeline_link_options.maxTraceDepth = max_trace_depth + pipeline_link_options.debugLevel = optix.COMPILE_DEBUG_LEVEL_MODERATE + + log = "" + pipeline = ctx.pipelineCreate( + pipeline_compile_options, pipeline_link_options, program_groups, log + ) + print("pipeline create log: " + log) + + stack_sizes = optix.StackSizes() + for prog_group in program_groups: + optix.util.accumulateStackSizes(prog_group, stack_sizes) + + ( + dc_stack_size_from_trav, + dc_stack_size_from_state, + cc_stack_size, + ) = optix.util.computeStackSizes( + stack_sizes, max_trace_depth, 0, 0 # maxCCDepth # maxDCDepth + ) + + pipeline.setStackSize( + dc_stack_size_from_trav, dc_stack_size_from_state, cc_stack_size, 1 # maxTraversableDepth + ) + + return pipeline + + +def create_sbt(prog_groups): + # print( "Creating sbt ... " ) + + (raygen_prog_group, miss_prog_group, hitgroup_prog_group) = prog_groups + + global d_raygen_sbt + global d_miss_sbt + + header_format = "{}B".format(optix.SBT_RECORD_HEADER_SIZE) + + # + # raygen record + # + formats = [header_format] + itemsize = get_aligned_itemsize(formats, optix.SBT_RECORD_ALIGNMENT) + dtype = np.dtype({"names": ["header"], "formats": formats, "itemsize": itemsize, "align": True}) + h_raygen_sbt = np.array([0], dtype=dtype) + optix.sbtRecordPackHeader(raygen_prog_group, h_raygen_sbt) + global d_raygen_sbt + d_raygen_sbt = array_to_device_memory(h_raygen_sbt) + + # + # miss record + # + formats = [header_format, "f4", "f4", "f4"] + itemsize = get_aligned_itemsize(formats, optix.SBT_RECORD_ALIGNMENT) + dtype = np.dtype( + { + "names": ["header", "r", "g", "b"], + "formats": formats, + "itemsize": itemsize, + "align": True, + } + ) + h_miss_sbt = np.array([(0, 0.3, 0.1, 0.2)], dtype=dtype) + # h_miss_sbt = np.array( [ (0, 0.0, 0.0, 0.0) ], dtype=dtype ) + optix.sbtRecordPackHeader(miss_prog_group, h_miss_sbt) + d_miss_sbt = array_to_device_memory(h_miss_sbt) + + # + # hitgroup record + # + formats = [header_format] + itemsize = get_aligned_itemsize(formats, optix.SBT_RECORD_ALIGNMENT) + dtype = np.dtype({"names": ["header"], "formats": formats, "itemsize": itemsize, "align": True}) + h_hitgroup_sbt = np.array([(0)], dtype=dtype) + optix.sbtRecordPackHeader(hitgroup_prog_group, h_hitgroup_sbt) + global d_hitgroup_sbt + d_hitgroup_sbt = array_to_device_memory(h_hitgroup_sbt) + + return optix.ShaderBindingTable( + raygenRecord=d_raygen_sbt.ptr, + missRecordBase=d_miss_sbt.ptr, + missRecordStrideInBytes=h_miss_sbt.dtype.itemsize, + missRecordCount=1, + hitgroupRecordBase=d_hitgroup_sbt.ptr, + hitgroupRecordStrideInBytes=h_hitgroup_sbt.dtype.itemsize, + hitgroupRecordCount=1, + ) + + +def create_accel(state: State, ctx, cuda_stream): + vert_grid_ter = state.terrain_vert_grid + vert_grid_sen = state.sensor_vert_grid + + vert_grid = np.concatenate([vert_grid_ter, vert_grid_sen]) + + # update inds in sensor + tri_grid_sen = state.sensor_tri_grid + tri_grid_sen = tri_grid_sen + int(vert_grid_ter.size / 3) + + tri_grid_ter = state.terrain_tri_grid + tri_grid = np.concatenate([tri_grid_ter, tri_grid_sen]) + + accel_options = optix.AccelBuildOptions( + buildFlags=int(optix.BUILD_FLAG_ALLOW_RANDOM_VERTEX_ACCESS), + operation=optix.BUILD_OPERATION_BUILD, + ) + + d_vert_grid = cp.array(vert_grid, dtype="f4") + nverts = int(len(d_vert_grid) / 3) + + d_tri_grid = cp.array(tri_grid, dtype=np.uint32) + n_ind_triplets = int(len(d_tri_grid) / 3) + + ortho_input_flags = [optix.GEOMETRY_FLAG_NONE] + ortho_input = optix.BuildInputTriangleArray() + ortho_input.vertexFormat = optix.VERTEX_FORMAT_FLOAT3 + ortho_input.numVertices = nverts + ortho_input.vertexBuffers = [d_vert_grid.data.ptr] + + ortho_input.indexFormat = optix.INDICES_FORMAT_UNSIGNED_INT3 + ortho_input.numIndexTriplets = n_ind_triplets + ortho_input.indexBuffer = d_tri_grid.data.ptr + ortho_input.indexStrideInBytes = 0 + + ortho_input.flags = ortho_input_flags + ortho_input.numSbtRecords = 1 + + gas_buffer_sizes = ctx.accelComputeMemoryUsage([accel_options], [ortho_input]) + + d_temp_buffer_gas = cp.cuda.alloc(gas_buffer_sizes.tempSizeInBytes) + d_gas_output_buffer = cp.cuda.alloc(gas_buffer_sizes.outputSizeInBytes) + + gas_handle = ctx.accelBuild( + cuda_stream.ptr, # CUDA stream + # 0, + [accel_options], + [ortho_input], + d_temp_buffer_gas.ptr, + gas_buffer_sizes.tempSizeInBytes, + d_gas_output_buffer.ptr, + gas_buffer_sizes.outputSizeInBytes, + [], # emitted properties + ) + # cuda_stream.synchronize() + + return (gas_handle, d_gas_output_buffer) + + +def launch(state: State, pipeline, sbt, trav_handle, cuda_stream): + d_tex_pix = cp.array(state.sensor_tex_pixels) + alpha = state.sensor_alpha + d_tex_pix = cp.dstack([d_tex_pix, alpha]) + + tex_height, tex_width, tex_chan = d_tex_pix.shape + d_tex_pix = cp.reshape(d_tex_pix, (tex_height, tex_chan * tex_width)) + + chan_desc = cp.cuda.texture.ChannelFormatDescriptor( + 8, 8, 8, 8, cp.cuda.runtime.cudaChannelFormatKindUnsigned + ) + + d_cuda_arr = cp.cuda.texture.CUDAarray(chan_desc, tex_width, tex_height) + + d_cuda_arr.copy_from(d_tex_pix, stream=cuda_stream) + # cuda_stream.synchronize() + + res_desc = cp.cuda.texture.ResourceDescriptor(cp.cuda.runtime.cudaResourceTypeArray, d_cuda_arr) + + tex_desc = cp.cuda.texture.TextureDescriptor( + addressModes=[cp.cuda.runtime.cudaAddressModeBorder, cp.cuda.runtime.cudaAddressModeBorder], + filterMode=cp.cuda.runtime.cudaFilterModePoint, + readMode=cp.cuda.runtime.cudaReadModeElementType, + normalizedCoords=1, + ) + + pix_width = state.image_width + pix_height = state.image_height + d_pix = state.image + + float2_dtype = np.dtype([("x", "f4"), ("y", "f4")], align=True) + + tex_obj = cp.cuda.texture.TextureObject(res_desc, tex_desc) + + # TODO remove the padding variable after data alignment is addressed + params = [ + ("u8", "trav_handle", trav_handle), + ("u8", "sensor_tex", tex_obj.ptr), + ("u8", "image", d_pix.data.ptr), + ("u4", "image_width", pix_width), + ("u4", "image_height", pix_height), + ("f4", "image_corner_coords_x", state.image_corner_coords[0].astype(np.float32)), + ("f4", "image_corner_coords_y", state.image_corner_coords[1].astype(np.float32)), + ("f4", "image_gsd", state.image_gsd.astype(np.float32)), + ("f4", "sensor_focal_length", state.sensor_focal_length.astype(np.float32)), + ("f4", "terrain_zmax", state.terrain_zmax.astype(np.float32)), + ("f4", "sensor_pos_x", state.sensor_pos[0].astype(np.float32)), + ("f4", "sensor_pos_y", state.sensor_pos[1].astype(np.float32)), + ("f4", "sensor_pos_z", state.sensor_pos[2].astype(np.float32)), + ("f4", "sensor_up_x", state.sensor_up[0].astype(np.float32)), + ("f4", "sensor_up_y", state.sensor_up[1].astype(np.float32)), + ("f4", "sensor_up_z", state.sensor_up[2].astype(np.float32)), + ("f4", "sensor_right_x", state.sensor_right[0].astype(np.float32)), + ("f4", "sensor_right_y", state.sensor_right[1].astype(np.float32)), + ("f4", "sensor_right_z", state.sensor_right[2].astype(np.float32)), + ( + "f4", + "sensor_focal_plane_origin_x", + state.sensor_focal_plane_origin[0].astype(np.float32), + ), + ( + "f4", + "sensor_focal_plane_origin_y", + state.sensor_focal_plane_origin[1].astype(np.float32), + ), + ( + "f4", + "sensor_focal_plane_origin_z", + state.sensor_focal_plane_origin[2].astype(np.float32), + ), + ("f4", "DUMMY", 1.0), + ( + float2_dtype, + "sensor_focal_plane_size", + ( + state.sensor_focal_plane_size[0].astype(np.float32), + state.sensor_focal_plane_size[1].astype(np.float32), + ), + ), + ] + + formats = [x[0] for x in params] + names = [x[1] for x in params] + values = [x[2] for x in params] + + itemsize = get_aligned_itemsize(formats, 8) + params_dtype = np.dtype( + {"names": names, "formats": formats, "itemsize": itemsize, "align": True} + ) + + h_params = np.array([tuple(values)], dtype=params_dtype) + d_params = array_to_device_memory(h_params, stream=cuda_stream) + # cuda_stream.synchronize() + + optix.launch( + pipeline, + cuda_stream.ptr, + d_params.ptr, + h_params.dtype.itemsize, + sbt, + pix_width, + pix_height, + 10, # depth + ) + cuda_stream.synchronize() + + d_pix = cp.reshape(d_pix, (state.image_height, state.image_width, 4)) + d_pix = cp.flipud(d_pix) + h_pix = cp.asnumpy(d_pix) + + return h_pix diff --git a/applications/orthorectification_with_optix/python/ortho_utils.py b/applications/orthorectification_with_optix/python/ortho_utils.py new file mode 100644 index 000000000..dfc3ac260 --- /dev/null +++ b/applications/orthorectification_with_optix/python/ortho_utils.py @@ -0,0 +1,73 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import os + +import numpy as np +from osgeo import gdal +from shapely.geometry import MultiPoint + + +def ray_plane_intersections_cpu(planeNormal, planePoint, rayDirections, rayPoint, epsilon=1e-6): + # ndotu = planeNormal.dot(rayDirections[0,:]) + npoints = rayDirections.shape[0] + planeNormals = np.repeat(planeNormal[np.newaxis, :], npoints, axis=0) + ndotus = (planeNormals * rayDirections).sum(1) + if (np.abs(ndotus) < epsilon).any(): + raise RuntimeError("no intersection or line is within plane") + + w = rayPoint - planePoint + # si = -planeNormal.dot(w) / ndotu + + ws = np.repeat(w[np.newaxis, :], npoints, axis=0) + sis = -(planeNormals * ws).sum(1) / ndotus + + # Psi = w + si * rayDirections[0,:] + planePoint + sis = np.repeat(sis[:, np.newaxis], 3, axis=1) + inters = ws + sis * rayDirections + planePoint + + return inters + + +def extract_extent_nativeCRS(image_fname): + info = gdal.Info(image_fname, options=gdal.InfoOptions(format="json", stats=True)) + top_left = [float(item) for item in info["cornerCoordinates"]["upperLeft"]] + bot_left = [float(item) for item in info["cornerCoordinates"]["lowerLeft"]] + top_right = [float(item) for item in info["cornerCoordinates"]["upperRight"]] + bot_right = [float(item) for item in info["cornerCoordinates"]["lowerRight"]] + corners = np.vstack([bot_right, bot_left, top_left, top_right, bot_right]) + bpts = MultiPoint(list(corners)) + return bpts.convex_hull.bounds + + +def load_basemap(basemap_fname, nx, ny, bands=[1, 2, 3]): + warp_ops = gdal.WarpOptions( + width=nx, height=ny, resampleAlg="cubic", creationOptions=["COMPRESS=LZW"] + ) + warped_fname = os.path.splitext(basemap_fname)[0] + "_resized.tif" + dset = gdal.Warp(warped_fname, basemap_fname, options=warp_ops) + img = [] + for bn in bands: + band = dset.GetRasterBand(bn) + frame = band.ReadAsArray() + img.append(np.flipud(frame)) + dset = None + + return np.dstack(img) + + +def create_blank_basemap(nx, ny, bands=3): + return np.zeros((ny, nx, bands), dtype=np.uint8) diff --git a/applications/orthorectification_with_optix/python/ortho_with_pyoptix.py b/applications/orthorectification_with_optix/python/ortho_with_pyoptix.py new file mode 100755 index 000000000..c1b86a59a --- /dev/null +++ b/applications/orthorectification_with_optix/python/ortho_with_pyoptix.py @@ -0,0 +1,373 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import os +import time + +import cupy as cp +import holoscan as hs +import numpy as np +import tqdm +from holoscan.conditions import CountCondition +from holoscan.core import Application, Operator, OperatorSpec +from holoscan.gxf import Entity +from holoscan.operators import HolovizOp +from odm_utils import load_all_frames_odm +from optix_utils_for_ortho import ( + State, + build_sensor_geom, + build_terrain_geom, + compile_cuda, + create_accel, + create_ctx, + create_module, + create_pipeline, + create_program_groups, + create_sbt, + init_render_params, + launch, + load_sensor_texture, + load_terrain_host, + set_pipeline_options, +) +from ortho_utils import create_blank_basemap, extract_extent_nativeCRS, load_basemap +from osgeo import gdal +from shapely.wkt import loads as wkt_load + +# dataset paths within the container +data_dir = os.path.expanduser("~/Data/lafayette_square") +image_dir = os.path.join(data_dir, "images") +geo_path = os.path.join(data_dir, "odm_georeferencing", "odm_georeferencing_model_geo.txt") +dem_image_fname = os.path.join(data_dir, "odm_dem", "dsm_small_filled.tif") +basemap_fname = None + +# config and include paths within container +cuda_tk_path = "/usr/local/cuda-11.6/include" +include_path = os.path.normpath("/work/NVIDIA-OptiX-SDK-7.4.0-linux64-x86_64/include") +proot = os.path.normpath("/work/") +project_include_path = os.path.normpath("/work/NVIDIA-OptiX-SDK-7.4.0-linux64-x86_64/SDK") +project_cu = os.path.join(proot, "cpp", "src", "optix", "optixOrtho.cu") + +# general configuration parameters +sensor_dims = (5472, 3648) +sensor_resize = 0.25 # resizes the raw sensor pixels +ncpu = 8 # how many cores to use to load sensor simulation +gsd = 0.25 # controls how many pixels are in the rendering +iterations = 18 # how many frames to render from the source images (in this case 425 is max) +use_mosaic_bbox = True # render to a static bounds on the ground as defined by the DEM +write_geotiff = False +nb = 3 # how many bands to write to the GeoTiff +render_scale = 0.5 # scale the holoview window up or down +fps = 8.0 # rate limit the simulated sensor feed to this many frames per second + +# ---------------Helper Functionality -------------------------# +# TODO Extend to camera models other than perspective + +if sensor_resize is not None: + sensor_dims = [ + int(item) for item in (sensor_dims[0] * sensor_resize, sensor_dims[1] * sensor_resize) + ] + +pbar = tqdm.tqdm(total=iterations, desc="ortho frames") + +frame_rate = 1 / fps +dem_bounds = extract_extent_nativeCRS(dem_image_fname) +xmin, ymin, xmax, ymax = dem_bounds + +box_size_x = xmax - xmin +box_size_y = ymax - ymin + +render_width = int(box_size_x / gsd) +render_height = int(box_size_y / gsd) + + +if use_mosaic_bbox: + print("Creating mosaic basemap") + if basemap_fname is not None: + basemap = load_basemap(basemap_fname, render_width, render_height) + else: + basemap = create_blank_basemap(render_width, render_height) + + mosaic_d = cp.array(basemap) + mosaic_alpha_d = cp.full((render_height, render_width, 1), 255, cp.uint8) + + mosaic_image_d = cp.dstack([mosaic_d, mosaic_alpha_d]) + + +# ---------------BEGIN HOLOSCAN -------------------------# +class FrameSimulatorODMOp(Operator): + def __init__(self, *args, **kwargs): + self.frame_count = 0 + self.prior_time = time.time() + # Need to call the base class constructor last + super().__init__(*args, **kwargs) + + def setup(self, spec: OperatorSpec): + print("pre-fetch frames from disk to simulate streaming sensor") + + self.staged_dataset = load_all_frames_odm( + use_mosaic_bbox, + dem_image_fname, + geo_path, + data_dir, + iterations, + ncpu, + sensor_resize, + sensor_dims, + ) + + spec.output("sensor_meta") + spec.output("sensor_pix") + + def compute(self, op_input, op_output, context): + data = self.staged_dataset[self.frame_count] + + sensor_pix = data["sensor_pix"] + meta = data["meta"] + + self.frame_count = self.frame_count + 1 + pbar.update(1) + + curr_time = time.time() + delta_time = curr_time - self.prior_time + if delta_time < frame_rate: + time.sleep(frame_rate - delta_time) + self.prior_time = time.time() + + op_output.emit(meta, "sensor_meta") + op_output.emit(sensor_pix, "sensor_pix") + + +class LoadTerrain(Operator): + def __init__(self, *args, **kwargs): + # Need to call the base class constructor last + super().__init__(*args, **kwargs) + + def setup(self, spec: OperatorSpec): + self.state = State() + self.state.image_gsd = np.float64(gsd) + self.state.sensor_alpha = cp.full((sensor_dims[1], sensor_dims[0], 1), 255, cp.uint8) + if use_mosaic_bbox: + self.state.image = mosaic_image_d + + load_terrain_host(self.state, dem_image_fname) + build_terrain_geom(self.state) + + spec.input("sensor_meta") + spec.input("sensor_pix") + + spec.output("optix_state") + spec.output("sensor_meta") + spec.output("sensor_pix") + + def compute(self, op_input, op_output, context): + sensor_meta = op_input.receive("sensor_meta") + sensor_pix = op_input.receive("sensor_pix") + + # TODO dynamically update / load terrain based on image footprint location + + op_output.emit(self.state, "optix_state") + op_output.emit(sensor_meta, "sensor_meta") + op_output.emit(sensor_pix, "sensor_pix") + + +class PrepOdmMeta4Ortho(Operator): + def __init__(self, *args, **kwargs): + # Need to call the base class constructor last + super().__init__(*args, **kwargs) + + def setup(self, spec: OperatorSpec): + self.state = State() + + spec.input("sensor_meta") + spec.input("sensor_pix") + spec.input("optix_state") + + spec.output("optix_state") + spec.output("sensor_pix") + + def compute(self, op_input, op_output, context): + sensor_meta = op_input.receive("sensor_meta") + sensor_pix = op_input.receive("sensor_pix") + optix_state = op_input.receive("optix_state") + + optix_state.frame_number = sensor_meta["frame_number"] + optix_state.bounding_box_estimate = wkt_load(sensor_meta["wkt"]).bounds + optix_state.sensor_dir = np.array( + [float(item) for item in sensor_meta["sensor_dir"].split(",")] + ) + optix_state.sensor_up = np.array( + [float(item) for item in sensor_meta["sensor_up"].split(",")] + ) + + optix_state.sensor_tex_width = sensor_dims[0] + optix_state.sensor_tex_height = sensor_dims[1] + + optix_state.sensor_hor_fov = np.deg2rad(float(sensor_meta["hfov"])) + optix_state.sensor_ver_fov = np.deg2rad(float(sensor_meta["vfov"])) + + optix_state.utm_sensor_pos = np.array( + [float(item) for item in sensor_meta["utm_pos"].split(",")] + ) + + init_render_params(optix_state) + + if not use_mosaic_bbox: + optix_state.image = cp.zeros( + (optix_state.image_width, optix_state.image_height, 4), cp.uint8 + ) + + op_output.emit(sensor_pix, "sensor_pix") + op_output.emit(optix_state, "optix_state") + + +class OptixOrthoOp(Operator): + def __init__(self, *args, **kwargs): + self.project_ptx = compile_cuda( + project_cu, cuda_tk_path, include_path, project_include_path + ) + self.ctx = create_ctx() + # Need to call the base class constructor last + super().__init__(*args, **kwargs) + + def setup(self, spec: OperatorSpec): + pipeline_options = set_pipeline_options() + module = create_module(self.ctx, pipeline_options, self.project_ptx) + self.prog_groups = create_program_groups(self.ctx, module) + + self.pipeline = create_pipeline(self.ctx, self.prog_groups, pipeline_options) + self.stream = cp.cuda.Stream() + + spec.input("optix_state") + spec.input("sensor_pix") + + spec.output("gtiff_meta") + spec.output("ortho_pix") + + def compute(self, op_input, op_output, context): + optix_state = op_input.receive("optix_state") + sensor_pix = op_input.receive("sensor_pix") + + load_sensor_texture(optix_state, sensor_pix) + build_sensor_geom(optix_state) + + gas_handle, d_gas_output_buffer = create_accel(optix_state, self.ctx, self.stream) + + sbt = create_sbt(self.prog_groups) + ortho_pix = launch(optix_state, self.pipeline, sbt, gas_handle, self.stream) + + gtiff_meta = {} + gtiff_meta["xsize"] = optix_state.image_width + gtiff_meta["ysize"] = optix_state.image_height + gtiff_meta["geot"] = optix_state.image_geot + gtiff_meta["wkt"] = optix_state.image_wkt + gtiff_meta["frame_number"] = optix_state.frame_number + + op_output.emit(gtiff_meta, "gtiff_meta") + op_output.emit(ortho_pix, "ortho_pix") + + +class ProcessOrthoOp(Operator): + def __init__(self, *args, **kwargs): + # Need to call the base class constructor last + super().__init__(*args, **kwargs) + + def setup(self, spec: OperatorSpec): + spec.input("gtiff_meta") + spec.input("ortho_pix") + + spec.output("outputs") + + def compute(self, op_input, op_output, context): + ortho_meta = op_input.receive("gtiff_meta") + ortho_pix = op_input.receive("ortho_pix") + + if write_geotiff: + img = ortho_pix[:, :, 0:nb] + + frame_str = str(ortho_meta["frame_number"]).zfill(8) + ortho_fname = os.path.join(data_dir, frame_str + "_frame_pyoptix_ortho.tif") + + driver = gdal.GetDriverByName("GTiff") + dst_ds = driver.Create( + ortho_fname, + xsize=ortho_meta["xsize"], + ysize=ortho_meta["ysize"], + bands=nb, + eType=gdal.GDT_Byte, + options=["TILED=YES", "COMPRESS=LZW"], + ) + + dst_ds.SetGeoTransform(ortho_meta["geot"]) + dst_ds.SetProjection(ortho_meta["wkt"]) + for bn in range(nb): + band = dst_ds.GetRasterBand(bn + 1) + band.SetNoDataValue(0) + + band.WriteArray(np.squeeze(img[:, :, bn])) + + dst_ds = None + + out_message = Entity(context) + out_message.add(hs.as_tensor(ortho_pix), "pixels") + + op_output.emit(out_message, "outputs") + + def on_error(e): + print("There was an error: {}".format(e)) + + +# ---------------END HOLOSCAN-------------------------# +class BasicOrthoFlow(Application): + def __init__(self): + super().__init__() + + def compose(self): + src = FrameSimulatorODMOp(self, CountCondition(self, iterations), name="src") + meta = PrepOdmMeta4Ortho(self, name="prep_odm_meta") + + terrain = LoadTerrain(self, name="terrain_loader") + ortho = OptixOrthoOp(self, name="optix_ortho") + ortho_proc = ProcessOrthoOp(self, name="proc_ortho") + + rh = int(render_height * render_scale) + rw = int(render_width * render_scale) + sink = HolovizOp( + self, name="holoviz ortho viewer", width=rw, height=rh, **self.kwargs("holoviz") + ) + + self.add_flow(src, terrain, {("sensor_meta", "sensor_meta"), ("sensor_pix", "sensor_pix")}) + self.add_flow( + terrain, + meta, + { + ("sensor_meta", "sensor_meta"), + ("sensor_pix", "sensor_pix"), + ("optix_state", "optix_state"), + }, + ) + self.add_flow(meta, ortho, {("optix_state", "optix_state"), ("sensor_pix", "sensor_pix")}) + self.add_flow(ortho, ortho_proc, {("gtiff_meta", "gtiff_meta"), ("ortho_pix", "ortho_pix")}) + self.add_flow(ortho_proc, sink, {("outputs", "receivers")}) + + +if __name__ == "__main__": + app = BasicOrthoFlow() + app.config("") + app.run() + pbar.close() + + print("done with app")