denoiser wrapper working, not all features have been tested though

Author: Felix Igelbrink

examples/denolser.py

@@ -0,0 +1,65 @@
+import optix as ox
+import cupy as cp
+import argparse
+import logging
+import sys
+import imageio
+import matplotlib.pyplot as plt
+
+logging.basicConfig(stream=sys.stdout, level=logging.INFO)
+log = logging.getLogger()
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser("Denoise images using the builtin Optix denoiser")
+
+    parser.add_argument('color_file', help="The input color image to denoise (exr format)")
+    parser.add_argument('-n', '--normal', type=str, required=False,
+                        help="Additional normals image (exr format)")
+    parser.add_argument('-a', '--albedo', type=str, required=False,
+                        help="Additional albedo image (exr format)")
+    parser.add_argument('-f', '--flow', type=str, required=False,
+                        help=f"Additional flow image for temporal denoising")
+    parser.add_argument('-o', '--out', type=str, default="denoised.exr", required=False,
+                        help="Output file. Defaults to \"denoised.exr\"")
+    parser.add_argument('-F', '--Frames', type=int, nargs=2, required=False,
+                        help="First-Last frame number of a sequence")
+    parser.add_argument('-e', '--exposure', type=float, required=False,
+                        help="apply exposure on output images")
+    parser.add_argument('-t', '--tilesize', type=int, nargs=2, required=False,
+                        help="Use tiling to save GPU memory")
+    parser.add_argument('-z', action='store_true', required=False,
+                        help="Apply flow to input images (no denoising) and write output")
+    parser.add_argument('-k', action='store_true', required=False,
+                        help="Use kernel prediction model even if there are no AOVs")
+
+    args = parser.parse_args()
+
+
+    # a color image is always required
+    color_image = imageio.read(args.color_file).get_data(0)
+    plt.imshow(color_image / 10000)
+    plt.show()
+    normal_image = None
+    albedo_image = None
+    flow_image = None
+    if args.normal is not None:
+        normal_image = imageio.read(args.normal).get_data(0)
+
+    if args.albedo is not None:
+        albedo_image = imageio.read(args.albedo).get_data(0)
+
+    # setup optix context and denoiser
+
+    logger = ox.Logger(log)
+    ctx = ox.DeviceContext(validation_mode=True, log_callback_function=logger, log_callback_level=3)
+
+    model_kind = ox.DenoiserModelKind.HDR
+
+    denoiser = ox.Denoiser(ctx, model_kind=model_kind, guide_albedo=args.albedo is not None, guide_normals=args.normal is not None, kp_mode=args.k)
+
+    ret = denoiser.invoke(color_image, albedo=albedo_image if args.albedo else None, normals=normal_image if args.normal else None)
+    ret = cp.asnumpy(ret)
+
+    ret /= 10000
+    plt.imshow(ret)
+    plt.show()

optix/__init__.py

@@ -6,6 +6,7 @@
 from .shader_binding_table import ShaderBindingTable
 from .pipeline import CompileDebugLevel, ExceptionFlags, TraversableGraphFlags, \
     PrimitiveTypeFlags, PipelineCompileOptions, PipelineLinkOptions, Pipeline
+from .denoiser import *
 from .logging_utility import Logger
 from ._version import __version__
 

optix/denoiser.pxd

@@ -3,14 +3,22 @@ from .context cimport OptixDeviceContext, OptixContextObject
 from libcpp.vector cimport vector
 from .base cimport OptixObject
 from libc.stdint cimport uintptr_t
+from libcpp cimport bool
 
 cdef extern from "optix_includes.h" nogil:
-    cdef enum OptixDenoiserModelKind:
-        OPTIX_DENOISER_MODEL_KIND_LDR
-        OPTIX_DENOISER_MODEL_KIND_HDR
-        OPTIX_DENOISER_MODEL_KIND_AOV
-        OPTIX_DENOISER_MODEL_KIND_TEMPORAL
-        OPTIX_DENOISER_MODEL_KIND_TEMPORAL_AOV
+    IF _OPTIX_VERSION > 70300:
+        cdef enum OptixDenoiserModelKind:
+            OPTIX_DENOISER_MODEL_KIND_LDR
+            OPTIX_DENOISER_MODEL_KIND_HDR
+            OPTIX_DENOISER_MODEL_KIND_AOV
+            OPTIX_DENOISER_MODEL_KIND_TEMPORAL
+            OPTIX_DENOISER_MODEL_KIND_TEMPORAL_AOV
+    ELSE:
+        cdef enum OptixDenoiserModelKind:
+            OPTIX_DENOISER_MODEL_KIND_LDR
+            OPTIX_DENOISER_MODEL_KIND_HDR
+            OPTIX_DENOISER_MODEL_KIND_AOV
+            OPTIX_DENOISER_MODEL_KIND_TEMPORAL
 
 
     cdef struct OptixDenoiserOptions:
@@ -57,7 +65,8 @@ cdef extern from "optix_includes.h" nogil:
         OptixImage2D normal
         OptixImage2D flow
 
-    ctypedef uintptr_t OptixDenoiser
+    ctypedef struct OptixDenoiser:
+        pass
 
     OptixResult optixDenoiserCreate(OptixDeviceContext context,
                                     OptixDenoiserModelKind modelKind,
@@ -118,18 +127,66 @@ cdef extern from "optix_includes.h" nogil:
             CUdeviceptr         scratch,
             size_t              scratchSizeInBytes)
 
+cdef extern from "optix_denoiser_tiling.h" nogil:
+    OptixResult optixUtilDenoiserInvokeTiled(OptixDenoiser denoiser,
+                                             CUstream stream,
+                                             const OptixDenoiserParams *params,
+                                             CUdeviceptr denoiserState,
+                                             size_t denoiserStateSizeInBytes,
+                                             const OptixDenoiserGuideLayer *guideLayer,
+                                             const OptixDenoiserLayer *layers,
+                                             unsigned int numLayers,
+                                             CUdeviceptr scratch,
+                                             size_t scratchSizeInBytes,
+                                             unsigned int overlapWindowSizeInPixels,
+                                             unsigned int tileWidth,
+                                             unsigned int tileHeight)
+
+    cdef struct OptixUtilDenoiserImageTile:
+        OptixImage2D input
+        OptixImage2D output
+        unsigned int inputOffsetX
+        unsigned int inputOffsetY
+
+    OptixResult optixUtilDenoiserSplitImage(const OptixImage2D & input,
+                                            const OptixImage2D & output,
+                                            unsigned int overlapWindowSizeInPixels, unsigned int tileWidth,
+                                            unsigned int tileHeight,
+                                            vector[OptixUtilDenoiserImageTile]& tiles)
+
+    unsigned int optixUtilGetPixelStride(const OptixImage2D & image)
 
-class Image2D(OptixObject):
+
+
+cdef class Image2D(OptixObject):
     cdef OptixImage2D image
     cdef object _d_data
 
-class DenoiserLayer(OptixObject):
+cdef class DenoiserLayer(OptixObject):
     cdef OptixDenoiserLayer layer
     cdef Image2D input
     cdef Image2D previous_output
     cdef Image2D output
 
-class Denoiser(OptixContextObject):
+cdef class DenoiserGuideLayer(OptixObject):
+    cdef OptixDenoiserGuideLayer layer
+    cdef Image2D albedo
+    cdef Image2D normal
+    cdef Image2D flow
+
+cdef class Denoiser(OptixContextObject):
     cdef OptixDenoiser denoiser
+    cdef object model_kind
     cdef bool guide_albedo
     cdef bool guide_normals
+    cdef bool kp_mode
+    cdef tuple tile_size
+    cdef object _d_state
+    cdef size_t _state_size
+    cdef object _d_scratch
+    cdef size_t _scratch_size
+    cdef object _d_window
+    cdef size_t _window_size
+    cdef object _d_intensity
+    cdef object _d_avg_color
+    cdef unsigned int _overlap