Skip to content

Commit cea992a

Browse files
author
Felix Igelbrink
committed
started to implement the denoiser wrapper
1 parent 30533f9 commit cea992a

3 files changed

Lines changed: 239 additions & 1 deletion

File tree

‎optix/build.pyx‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,8 @@ __all__ = ['GeometryFlags',
1919
'BuildInputCurveArray',
2020
'BuildInputInstanceArray',
2121
'Instance',
22-
'AccelerationStructure'
22+
'AccelerationStructure',
23+
'CurveEndcapFlags'
2324
]
2425

2526
class GeometryFlags(IntEnum):

‎optix/denoiser.pxd‎

Lines changed: 135 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,135 @@
1+
from .common cimport OptixResult, CUstream, CUdeviceptr
2+
from .context cimport OptixDeviceContext, OptixContextObject
3+
from libcpp.vector cimport vector
4+
from .base cimport OptixObject
5+
from libc.stdint cimport uintptr_t
6+
7+
cdef extern from "optix_includes.h" nogil:
8+
cdef enum OptixDenoiserModelKind:
9+
OPTIX_DENOISER_MODEL_KIND_LDR
10+
OPTIX_DENOISER_MODEL_KIND_HDR
11+
OPTIX_DENOISER_MODEL_KIND_AOV
12+
OPTIX_DENOISER_MODEL_KIND_TEMPORAL
13+
OPTIX_DENOISER_MODEL_KIND_TEMPORAL_AOV
14+
15+
16+
cdef struct OptixDenoiserOptions:
17+
unsigned int guideAlbedo
18+
unsigned int guideNormal
19+
20+
cdef struct OptixDenoiserSizes:
21+
size_t stateSizeInBytes
22+
size_t withOverlapScratchSizeInBytes
23+
size_t withoutOverlapScratchSizeInBytes
24+
unsigned int overlapWindowSizeInPixels
25+
26+
cdef struct OptixDenoiserParams:
27+
unsigned int denoiseAlpha
28+
CUdeviceptr hdrIntensity
29+
float blendFactor
30+
CUdeviceptr hdrAverageColor
31+
32+
cdef enum OptixPixelFormat:
33+
OPTIX_PIXEL_FORMAT_HALF2
34+
OPTIX_PIXEL_FORMAT_HALF3
35+
OPTIX_PIXEL_FORMAT_HALF4
36+
OPTIX_PIXEL_FORMAT_FLOAT2
37+
OPTIX_PIXEL_FORMAT_FLOAT3
38+
OPTIX_PIXEL_FORMAT_FLOAT4
39+
OPTIX_PIXEL_FORMAT_UCHAR3
40+
OPTIX_PIXEL_FORMAT_UCHAR4
41+
42+
cdef struct OptixImage2D:
43+
CUdeviceptr data
44+
unsigned int width
45+
unsigned int height
46+
unsigned int rowStrideInBytes
47+
unsigned int pixelStrideInBytes
48+
OptixPixelFormat format
49+
50+
cdef struct OptixDenoiserLayer:
51+
OptixImage2D input
52+
OptixImage2D previousOutput
53+
OptixImage2D output
54+
55+
cdef struct OptixDenoiserGuideLayer:
56+
OptixImage2D albedo
57+
OptixImage2D normal
58+
OptixImage2D flow
59+
60+
ctypedef uintptr_t OptixDenoiser
61+
62+
OptixResult optixDenoiserCreate(OptixDeviceContext context,
63+
OptixDenoiserModelKind modelKind,
64+
const OptixDenoiserOptions* options,
65+
OptixDenoiser* denoiser
66+
)
67+
68+
OptixResult optixDenoiserCreateWithUserModel(OptixDeviceContext context,
69+
const void* userData,
70+
size_t userDataSizeInBytes,
71+
OptixDenoiser* denoiser
72+
)
73+
74+
OptixResult optixDenoiserDestroy(OptixDenoiser denoiser)
75+
76+
OptixResult optixDenoiserComputeMemoryResources(const OptixDenoiser denoiser,
77+
unsigned int inputWidth,
78+
unsigned int inputHeight,
79+
OptixDenoiserSizes* returnSizes)
80+
81+
OptixResult optixDenoiserSetup(
82+
OptixDenoiser denoiser,
83+
CUstream stream,
84+
unsigned int inputWidth,
85+
unsigned int inputHeight,
86+
CUdeviceptr denoiserState,
87+
size_t denoiserStateSizeInBytes,
88+
CUdeviceptr scratch,
89+
size_t scratchSizeInBytes)
90+
91+
OptixResult optixDenoiserInvoke(
92+
OptixDenoiser denoiser,
93+
CUstream stream,
94+
const OptixDenoiserParams * params,
95+
CUdeviceptr denoiserState,
96+
size_t denoiserStateSizeInBytes,
97+
const OptixDenoiserGuideLayer * guideLayer,
98+
const OptixDenoiserLayer * layers,
99+
unsigned int numLayers,
100+
unsigned int inputOffsetX,
101+
unsigned int inputOffsetY,
102+
CUdeviceptr scratch,
103+
size_t scratchSizeInBytes)
104+
105+
OptixResult optixDenoiserComputeAverageColor(
106+
OptixDenoiser denoiser,
107+
CUstream stream,
108+
const OptixImage2D * inputImage,
109+
CUdeviceptr outputAverageColor,
110+
CUdeviceptr scratch,
111+
size_t scratchSizeInBytes)
112+
113+
OptixResult optixDenoiserComputeIntensity(
114+
OptixDenoiser denoiser,
115+
CUstream stream,
116+
const OptixImage2D * inputImage,
117+
CUdeviceptr outputIntensity,
118+
CUdeviceptr scratch,
119+
size_t scratchSizeInBytes)
120+
121+
122+
class Image2D(OptixObject):
123+
cdef OptixImage2D image
124+
cdef object _d_data
125+
126+
class DenoiserLayer(OptixObject):
127+
cdef OptixDenoiserLayer layer
128+
cdef Image2D input
129+
cdef Image2D previous_output
130+
cdef Image2D output
131+
132+
class Denoiser(OptixContextObject):
133+
cdef OptixDenoiser denoiser
134+
cdef bool guide_albedo
135+
cdef bool guide_normals

‎optix/denoiser.pyx‎

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
from .common cimport optix_check_return, optix_init
2+
from .context cimport DeviceContext
3+
import cupy as cp
4+
import numpy as np
5+
from enum import IntEnum, IntFlag
6+
from libc.string cimport memcpy, memset
7+
from libcpp.vector cimport vector
8+
from .common import ensure_iterable
9+
10+
optix_init()
11+
12+
class DenoiserModelKind(IntEnum):
13+
"""
14+
Wraps the OptixDenoiserModelKind enum.
15+
"""
16+
LHR = OPTIX_DENOISER_MODEL_KIND_LDR
17+
HDR = OPTIX_DENOISER_MODEL_KIND_HDR
18+
AOV = OPTIX_DENOISER_MODEL_KIND_AOV
19+
TEMPORAL = OPTIX_DENOISER_MODEL_KIND_TEMPORAL
20+
TEMPORAL_AOV = OPTIX_DENOISER_MODEL_KIND_TEMPORAL_AOV
21+
22+
class Image2D(OptixObject):
23+
_dtype_to_pixeltype = {
24+
(np.float16, 2): OPTIX_PIXEL_FORMAT_HALF2,
25+
(np.float16, 3): OPTIX_PIXEL_FORMAT_HALF3,
26+
(np.float16, 4): OPTIX_PIXEL_FORMAT_HALF4,
27+
(np.float32, 2): OPTIX_PIXEL_FORMAT_FLOAT2,
28+
(np.float32, 3): OPTIX_PIXEL_FORMAT_FLOAT3,
29+
(np.float32, 4): OPTIX_PIXEL_FORMAT_FLOAT4,
30+
}
31+
def __init__(self, data):
32+
self._d_data = cp.asarray(data) # push data to the gpu
33+
34+
if len(self._d_data.shape) != 3:
35+
raise ValueError("Invalid shape of input array. Expected a 3D (height, width, channels) array.")
36+
try:
37+
pixel_type = self._dtype_to_pixeltype[(self._d_data.dtype, self._d_data.shape[2])]
38+
except KeyError:
39+
raise ValueError(f"Invalid dtype {self._d_data.dtype} of data. Only float32 and float16 are supported for Image2D")
40+
41+
# fill the underlying struct values
42+
self.image.data = self._d_data.data.ptr
43+
self.image.height = self._d_data.shape[0]
44+
self.image.width = self._d_data.shape[1]
45+
self.image.rowStrideInBytes = self._d_data.strides[0]
46+
self.image.pixelStrideInBytes = self._d_data.strides[1]
47+
self.image.format = pixel_type
48+
49+
@property
50+
def dtype(self):
51+
return self._d_data.dtype
52+
53+
@property
54+
def shape(self):
55+
return self._d_data.shape
56+
57+
58+
class DenoiserLayer(OptixObject):
59+
def __init__(self, input, previous_output=None, output=None):
60+
self.input = Image2D(input) if not isinstance(input, Image2D) else input
61+
62+
if previous_output is not None:
63+
self.previous_output = Image2D(previous_output) if not isinstance(previous_output, Image2D) else previous_output
64+
65+
if output is not None:
66+
self.output = Image2D(output) if not isinstance(output, Image2D) else output
67+
68+
self.layer.input = (<Image2D>self.input).image
69+
self.layer.previousOutput = (<Image2D>self.previous_output).image
70+
self.layer.output = (<Image2D>self.output).image
71+
72+
73+
class Denoiser(OptixContextObject):
74+
def __init__(self, DeviceContext context, DenoiserModelKind model_kind=DenoiserModelKind.LHR, enable_albedo=False, enable_normals=False):
75+
super().__init__(context)
76+
cdef OptixDenoiserOptions options
77+
78+
if model_kind is not None:
79+
options.guideAlbedo = 1 if enable_albedo else 0
80+
options.guideNormal = 1 if enable_normals else 0
81+
self.guide_albedo = enable_albedo
82+
self.guide_normals = enable_normals
83+
84+
optix_check_return(optixDenoiserCreate(self.context.c_context,
85+
model_kind.value,
86+
&options,
87+
&self.denoiser))
88+
89+
@classmethod
90+
def create_with_user_model(cls, DeviceContext context, user_model):
91+
raise NotImplementedError()
92+
#obj = cls(context, model_kind=None)
93+
#optix_check_return(optixDenoiserCreateWithUserModel(obj.context.c_context,
94+
# user_model, #TODO
95+
# len(user_model), #TODO
96+
# &obj.denoiser))
97+
#return obj
98+
99+
def __dealloc__(self):
100+
optix_check_return(optixDenoiserDestroy(self.denoiser))
101+
102+

0 commit comments

Comments
 (0)