-
Notifications
You must be signed in to change notification settings - Fork 22
Expand file tree
/
Copy pathmodule.cpp
More file actions
319 lines (266 loc) · 15.2 KB
/
Copy pathmodule.cpp
File metadata and controls
319 lines (266 loc) · 15.2 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
/*
* 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 <cstdlib>
#include <iostream>
#include <string>
#include <optional>
#include <pybind11/stl_bind.h>
#include <ilogger.h>
#include <log.h>
#include <nvimgcodec.h>
#include "image.h"
#include "module.h"
namespace nvimgcodec {
uint32_t verbosity2severity(int verbose)
{
uint32_t result = 0;
if (verbose >= 1)
result |= NVIMGCODEC_DEBUG_MESSAGE_SEVERITY_FATAL | NVIMGCODEC_DEBUG_MESSAGE_SEVERITY_ERROR;
if (verbose >= 2)
result |= NVIMGCODEC_DEBUG_MESSAGE_SEVERITY_WARNING;
if (verbose >= 3)
result |= NVIMGCODEC_DEBUG_MESSAGE_SEVERITY_INFO;
if (verbose >= 4)
result |= NVIMGCODEC_DEBUG_MESSAGE_SEVERITY_DEBUG;
if (verbose >= 5)
result |= NVIMGCODEC_DEBUG_MESSAGE_SEVERITY_TRACE;
return result;
}
Module::Module()
: dbg_messenger_handle_(nullptr)
{
int verbosity = 2;
std::string verbosity_warning;
char* v = std::getenv("PYNVIMGCODEC_VERBOSITY");
try {
if (v) {
verbosity = std::stoi(v);
}
} catch (std::invalid_argument const& ex) {
verbosity_warning = "PYNVIMGCODEC_VERBOSITY has wrong value";
} catch (std::out_of_range const& ex) {
verbosity_warning = "PYNVIMGCODEC_VERBOSITY has out of range value";
}
if (verbosity > 0) {
dbg_messenger_ = std::make_unique<DefaultDebugMessenger>(verbosity2severity(verbosity), NVIMGCODEC_DEBUG_MESSAGE_CATEGORY_ALL);
logger_ = std::make_unique<Logger>("pynvimgcodec", dbg_messenger_.get());
if (!verbosity_warning.empty()) {
NVIMGCODEC_LOG_WARNING(logger_.get(), verbosity_warning);
}
} else {
logger_ = std::make_unique<Logger>("pynvimgcodec");
}
nvimgcodecInstanceCreateInfo_t instance_create_info{NVIMGCODEC_STRUCTURE_TYPE_INSTANCE_CREATE_INFO, sizeof(nvimgcodecInstanceCreateInfo_t), 0};
instance_create_info.load_builtin_modules = 1;
instance_create_info.load_extension_modules = 1;
instance_create_info.create_debug_messenger = verbosity > 0 ? 1 : 0;
instance_create_info.debug_messenger_desc = verbosity > 0 ? dbg_messenger_->getDesc() : nullptr;
nvimgcodecInstanceCreate(&instance_, &instance_create_info);
}
Module ::~Module()
{
nvimgcodecInstanceDestroy(instance_);
}
void Module::exportToPython(py::module& m, nvimgcodecInstance_t instance, ILogger* logger)
{
m.def(
"as_image",
[instance, logger](py::handle source, intptr_t cuda_stream,
std::optional<nvimgcodecSampleFormat_t> sample_format,
std::optional<nvimgcodecColorSpec_t> color_spec,
std::optional<int> precision) -> Image {
return Image(instance, logger, source.ptr(), cuda_stream, sample_format, color_spec, precision);
},
R"pbdoc(
Wraps an external buffer as an Image and ties the buffer lifetime to the Image.
The buffer must be C-style contiguous, but rows may have additional padding.
Args:
source: Input DLPack tensor encapsulated in a PyCapsule, or any object
exposing ``__cuda_array_interface__``, ``__array_interface__``
or ``__dlpack__``/``__dlpack_device__``.
cuda_stream: Optional ``cudaStream_t`` as a Python integer, used for any
synchronization the created Image needs.
sample_format: (keyword-only) ``nvimgcodec.SampleFormat`` selecting the
output layout:
* ``I_*`` (concrete interleaved) -> HWC; the last axis
must have at least the format's channel arity (2-D
``(H, W)`` only for ``I_Y``).
* ``P_*`` (concrete planar) -> CHW; the leading
axis must have at least the format's channel arity
(2-D ``(H, W)`` only for ``P_Y``).
* ``I_UNCHANGED`` -> HWC, no arity check.
* ``P_UNCHANGED`` -> CHW, no arity check.
* ``UNKNOWN`` -> raises ``ValueError``.
Extra channels beyond the format's arity are kept; too
few raise ``ValueError``.
When omitted, the value is inferred from the channel
count of the array using the HWC defaults:
* 1 channel -> ``I_Y``
* 2 channels -> ``I_YA``
* 3 channels -> ``I_RGB``
* 4 channels -> ``I_RGBA``
* 5+ channels -> ``UNKNOWN``
color_spec: (keyword-only) ``nvimgcodec.ColorSpec`` override. When omitted
the value is inferred from the channel count:
* 1 / 2 channels -> ``GRAY``
* 3 / 4 channels -> ``SRGB``
* 5+ channels -> ``UNKNOWN``
precision: (keyword-only) Optional integer giving the number of significant
bits per sample. Use this to describe lower-precision data
stored in a wider container, e.g. a 12-bit image held in a
uint16 buffer (``precision=12``). Accepted values are 0 (or
``None``, both meaning "use the full bitdepth of the sample
data type") and any integer in ``1..bitdepth(dtype)``. For
floating-point dtypes only 0/None or the dtype's full bitdepth
(e.g. 32 for float32) is accepted; sub-bitdepth values are an
integer-image concept and have no meaningful interpretation
for floats.
Returns:
nvimgcodec.Image
)pbdoc",
"source"_a, "cuda_stream"_a = 0, py::kw_only(), "sample_format"_a = py::none(), "color_spec"_a = py::none(),
"precision"_a = py::none(), py::keep_alive<0, 1>())
.def(
"as_images",
[instance, logger](const std::vector<py::handle>& sources, intptr_t cuda_stream,
std::optional<nvimgcodecSampleFormat_t> sample_format,
std::optional<nvimgcodecColorSpec_t> color_spec,
std::optional<int> precision) -> std::vector<py::object> {
std::vector<py::object> py_images;
py_images.reserve(sources.size());
for (auto& source : sources) {
Image img(instance, logger, source.ptr(), cuda_stream, sample_format, color_spec, precision);
py::object py_img = py::cast(img);
py_images.push_back(py_img);
py::detail::keep_alive_impl(py_img, source);
}
return py_images;
},
R"pbdoc(
Wrap a list of external buffers as Images and tie each buffer's lifetime
to the corresponding Image. Equivalent to calling :func:`as_image` on
each element of ``sources`` with the same ``sample_format``, ``color_spec``
and ``precision`` overrides.
Layout rules and sample_format / color_spec / precision inference are
identical to :func:`as_image`. The same overrides are applied uniformly
to every input - the elements of ``sources`` may otherwise differ in
shape, dtype and channel count.
Args:
sources: List of input DLPack tensors (PyCapsules) or objects exposing
``__cuda_array_interface__``, ``__array_interface__`` or
``__dlpack__``/``__dlpack_device__``.
cuda_stream: Optional ``cudaStream_t`` as a Python integer, used for
any synchronization the created Images need.
sample_format: (keyword-only) ``nvimgcodec.SampleFormat`` override
applied to every Image, selecting the output layout:
* ``I_*`` (concrete interleaved) -> HWC; the last axis
must have at least the format's channel arity (2-D
``(H, W)`` only for ``I_Y``).
* ``P_*`` (concrete planar) -> CHW; the leading
axis must have at least the format's channel arity
(2-D ``(H, W)`` only for ``P_Y``).
* ``I_UNCHANGED`` -> HWC, no arity check.
* ``P_UNCHANGED`` -> CHW, no arity check.
* ``UNKNOWN`` -> raises ``ValueError``.
Extra channels beyond the format's arity are kept; too
few raise ``ValueError``.
When omitted, each Image's value is inferred from its
own channel count using the HWC defaults:
* 1 channel -> ``I_Y``
* 2 channels -> ``I_YA``
* 3 channels -> ``I_RGB``
* 4 channels -> ``I_RGBA``
* 5+ channels -> ``UNKNOWN``
color_spec: (keyword-only) ``nvimgcodec.ColorSpec`` override applied
to every Image. When omitted, each Image's value is
inferred from its own channel count:
* 1 / 2 channels -> ``GRAY``
* 3 / 4 channels -> ``SRGB``
* 5+ channels -> ``UNKNOWN``
precision: (keyword-only) Optional integer giving the number of
significant bits per sample, applied uniformly to every
Image. Use this to describe lower-precision data stored in
a wider container, e.g. a 12-bit image held in a uint16
buffer (``precision=12``). Accepted values are 0 (or
``None``, both meaning "use the full bitdepth of the sample
data type") and any integer in ``1..bitdepth(dtype)``. For
floating-point dtypes only 0/None or the dtype's full
bitdepth (e.g. 32 for float32) is accepted; sub-bitdepth
values are an integer-image concept and have no meaningful
interpretation for floats. The value is validated against
each source's dtype.
Returns:
List of nvimgcodec.Image objects.
)pbdoc",
"sources"_a, "cuda_stream"_a = 0, py::kw_only(), "sample_format"_a = py::none(), "color_spec"_a = py::none(),
"precision"_a = py::none())
.def(
"from_dlpack",
[instance, logger](py::handle source, intptr_t cuda_stream,
std::optional<nvimgcodecSampleFormat_t> sample_format,
std::optional<nvimgcodecColorSpec_t> color_spec,
std::optional<int> precision) -> Image {
return Image(instance, logger, source.ptr(), cuda_stream, sample_format, color_spec, precision);
},
R"pbdoc(
Zero-copy conversion from a DLPack tensor to an Image.
The DLPack source must be a 3-dimensional, CUDA-accessible tensor
(device or CUDA-host memory). Layout selection and sample_format /
color_spec inference are identical to :func:`as_image`; see that
function for the full rules.
Args:
source: Input DLPack tensor encapsulated in a PyCapsule, or any object
exposing ``__dlpack__`` and ``__dlpack_device__`` methods.
cuda_stream: Optional ``cudaStream_t`` as a Python integer, used for any
synchronization the created Image needs.
sample_format: (keyword-only) ``nvimgcodec.SampleFormat`` selecting the
output layout:
* ``I_*`` (concrete interleaved) -> HWC; the last axis
must have at least the format's channel arity.
* ``P_*`` (concrete planar) -> CHW; the leading
axis must have at least the format's channel arity.
* ``I_UNCHANGED`` -> HWC, no arity check.
* ``P_UNCHANGED`` -> CHW, no arity check.
* ``UNKNOWN`` -> raises ``ValueError``.
Extra channels beyond the format's arity are kept; too
few raise ``ValueError``.
When omitted, the value is inferred from the channel
count of the tensor using the HWC defaults:
* 1 channel -> ``I_Y``
* 2 channels -> ``I_YA``
* 3 channels -> ``I_RGB``
* 4 channels -> ``I_RGBA``
* 5+ channels -> ``UNKNOWN``
color_spec: (keyword-only) ``nvimgcodec.ColorSpec`` override. When omitted
the value is inferred from the channel count:
* 1 / 2 channels -> ``GRAY``
* 3 / 4 channels -> ``SRGB``
* 5+ channels -> ``UNKNOWN``
precision: (keyword-only) Optional integer giving the number of significant
bits per sample. Use this to describe lower-precision data
stored in a wider container, e.g. a 12-bit image held in a
uint16 buffer (``precision=12``). Accepted values are 0 (or
``None``, both meaning "use the full bitdepth of the sample
data type") and any integer in ``1..bitdepth(dtype)``. See
:func:`as_image` for details.
Returns:
nvimgcodec.Image
)pbdoc",
"source"_a, "cuda_stream"_a = 0, py::kw_only(), "sample_format"_a = py::none(), "color_spec"_a = py::none(),
"precision"_a = py::none(), py::keep_alive<0, 1>());
}
} // namespace nvimgcodec