Merge pull request #513 from jflatt-gia/scale2x

Add Optix 2X scaler to imgtool
This commit is contained in:
Matt Pharr 2025-10-22 12:35:02 -07:00 • committed by GitHub
commit 35bb9b88c0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 311 additions and 1 deletions

View file

@ -707,6 +707,7 @@ if (PBRT_CUDA_ENABLED)
set (PBRT_GPU_SOURCE
src/pbrt/gpu/optix/aggregate.cpp
src/pbrt/gpu/optix/denoiser.cpp
src/pbrt/gpu/optix/scaler.cpp
src/pbrt/gpu/memory.cpp
src/pbrt/gpu/util.cpp
)
@ -714,6 +715,7 @@ if (PBRT_CUDA_ENABLED)
src/pbrt/gpu/optix/aggregate.h
src/pbrt/gpu/cudagl.h
src/pbrt/gpu/optix/denoiser.h
src/pbrt/gpu/optix/scaler.h
src/pbrt/gpu/memory.h
src/pbrt/gpu/optix/optix.h
src/pbrt/gpu/util.h

View file

@ -9,6 +9,7 @@
#ifdef PBRT_BUILD_GPU_RENDERER
#ifndef __HIP_PLATFORM_AMD__
#include <pbrt/gpu/optix/denoiser.h>
#include <pbrt/gpu/optix/scaler.h>
#endif // __HIP_PLATFORM_AMD__
#include <pbrt/gpu/util.h>
#endif // PBRT_BUILD_GPU_RENDERER
@ -149,6 +150,14 @@ static std::map<std::string, CommandUsage> commandUsage = {
" be a multi-channel EXR as generated by pbrt's \"gbuffer\" film.",
std::string(R"( options:
--outfile <name> Filename to use for the denoised image.
)")}},
{"scale-optix",
{"scale-optix [options] <filename>",
"2X Scales the image using NVIDIA's OptiX denoiser which is\n"
" based on a deep neural network. The provided image should\n"
" be a multi-channel EXR as generated by pbrt's \"gbuffer\" film.",
std::string(R"( options:
--outfile <name> Filename to use for the scaled image.
)")}},
#endif // PBRT_BUILD_GPU_RENDERER
{"error",
@ -222,7 +231,7 @@ static std::map<std::string, CommandUsage> commandUsage = {
{"scalenormalmap",
{"scalenormalmap [options] <filename>",
"Scale the provided normal map by applying the given factor for x and y\n"
" and output the resulting normal map.\n",
" and output the resulting normal map.",
std::string(R"(
--scale <s> Scale factor. Default: 1
--outfile <name> Filename to store final image in.
@ -2335,6 +2344,125 @@ int denoise_optix(std::vector<std::string> args) {
return 0;
}
int scale_optix(std::vector<std::string> args) {
std::string inFilename, outFilename;
auto onError = [](const std::string &err) {
usage("scale-optix", "%s", err.c_str());
exit(1);
};
for (auto iter = args.begin(); iter != args.end(); ++iter) {
if (ParseArg(&iter, args.end(), "outfile", &outFilename, onError)) {
// success
} else if ((*iter)[0] == '-')
usage("scale-optix", "%s: unknown command flag", iter->c_str());
else if (inFilename.empty()) {
inFilename = *iter;
} else
usage("scale-optix", "multiple input filenames provided.");
}
if (inFilename.empty())
usage("scale-optix", "input image filename must be provided.");
if (outFilename.empty())
usage("scale-optix", "output image filename must be provided.");
ImageAndMetadata im = Image::Read(inFilename);
Image &image = im.image;
CUDA_CHECK(cudaFree(nullptr));
int nLayers = 3;
bool oldNormalNaming = false;
ImageChannelDesc desc[3] = {
image.GetChannelDesc({"R", "G", "B"}),
image.GetChannelDesc({"Albedo.R", "Albedo.G", "Albedo.B"}),
image.GetChannelDesc({"Ns.X", "Ns.Y", "Ns.Z"})};
if (!desc[0]) {
Error("%s: image doesn't have R, G, B channels.", inFilename);
return 1;
}
if (!desc[1]) {
Warning("%s: image doesn't have Albedo.{R,G,B} channels. "
"Denoising quality may suffer.",
inFilename);
nLayers = 1;
}
if (!desc[2]) {
// Try the old naming scheme
desc[2] = image.GetChannelDesc({"Nsx", "Nsy", "Nsz"});
if (desc[2])
oldNormalNaming = true;
else {
Warning("%s: image doesn't have Ns.X, Ns.Y, Ns.Z channels. "
"Denoising quality may suffer.",
inFilename);
nLayers = 1;
}
}
Scaler scaler((Vector2i)image.Resolution(), nLayers == 3);
size_t imageBytes = 3 * image.Resolution().x * image.Resolution().y * sizeof(float);
auto copyChannelsToGPU = [&](std::array<std::string, 3> ch, bool flipZ = false) {
void *bufGPU;
CUDA_CHECK(cudaMalloc(&bufGPU, imageBytes));
std::vector<float> hostStaging(imageBytes / sizeof(float));
ImageChannelDesc desc = image.GetChannelDesc(ch);
CHECK(desc);
int offset = 0;
for (int y = 0; y < image.Resolution().y; ++y)
for (int x = 0; x < image.Resolution().x; ++x) {
ImageChannelValues v = image.GetChannels({x, y}, desc);
if (flipZ)
v[2] *= -1; // flip normal's z--right handed...
for (int c = 0; c < 3; ++c)
hostStaging[offset++] = v[c];
}
CUDA_CHECK(
cudaMemcpy(bufGPU, hostStaging.data(), imageBytes, cudaMemcpyHostToDevice));
return bufGPU;
};
RGB *rgbGPU = (RGB *)copyChannelsToGPU({"R", "G", "B"});
RGB *albedoGPU = nullptr;
Normal3f *normalGPU = nullptr;
if (nLayers == 3) {
albedoGPU = (RGB *)copyChannelsToGPU({"Albedo.R", "Albedo.G", "Albedo.B"});
if (oldNormalNaming)
normalGPU = (Normal3f *)copyChannelsToGPU({"Nsx", "Nsy", "Nsz"}, true);
else
normalGPU = (Normal3f *)copyChannelsToGPU({"Ns.X", "Ns.Y", "Ns.Z"}, true);
}
RGB *rgbResultGPU;
Point2i destRes(image.Resolution().x * 2, image.Resolution().y * 2);
size_t destBytes = 3 * destRes.x * destRes.y * sizeof(float);
CUDA_CHECK(cudaMalloc(&rgbResultGPU, destBytes));
scaler.Scale(rgbGPU, normalGPU, albedoGPU, rgbResultGPU);
CUDA_CHECK(cudaDeviceSynchronize());
Image result(PixelFormat::Float, destRes, {"R", "G", "B"});
CUDA_CHECK(cudaMemcpy(result.RawPointer({0, 0}), (const void *)rgbResultGPU,
destBytes, cudaMemcpyDeviceToHost));
ImageMetadata outMetadata;
outMetadata.cameraFromWorld = im.metadata.cameraFromWorld;
outMetadata.NDCFromWorld = im.metadata.NDCFromWorld;
outMetadata.pixelBounds = Bounds2i(
{im.metadata.pixelBounds->pMin.x * 2, im.metadata.pixelBounds->pMin.y * 2},
{im.metadata.pixelBounds->pMax.x * 2, im.metadata.pixelBounds->pMax.y * 2});
outMetadata.fullResolution = destRes;
outMetadata.colorSpace = im.metadata.colorSpace;
CHECK(result.Write(outFilename, outMetadata));
return 0;
}
#endif // PBRT_BUILD_GPU_RENDERER
int main(int argc, char *argv[]) {
@ -2365,6 +2493,8 @@ int main(int argc, char *argv[]) {
#ifdef PBRT_BUILD_GPU_RENDERER
else if (cmd == "denoise-optix")
return denoise_optix(args);
else if (cmd == "scale-optix")
return scale_optix(args);
#endif // PBRT_BUILD_GPU_RENDERER
else if (cmd == "error")
return error(args);

View file

@ -0,0 +1,143 @@
// pbrt is Copyright(c) 1998-2020 Matt Pharr, Wenzel Jakob, and Greg Humphreys.
// The pbrt source code is licensed under the Apache License, Version 2.0.
// SPDX: Apache-2.0
#include <pbrt/gpu/optix/scaler.h>
#include <pbrt/gpu/memory.h>
#include <pbrt/gpu/util.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <array>
#include <optix.h>
#include <optix_stubs.h>
#define OPTIX_CHECK(EXPR) \
do { \
OptixResult res = EXPR; \
if (res != OPTIX_SUCCESS) \
LOG_FATAL("OptiX call " #EXPR " failed with code %d: \"%s\"", int(res), \
optixGetErrorString(res)); \
} while (false) /* eat semicolon */
// Stop that, Windows.
#ifdef RGB
#undef RGB
#endif
namespace pbrt {
Scaler::Scaler(Vector2i resolution, bool haveAlbedoAndNormal)
: resolution(resolution), haveAlbedoAndNormal(haveAlbedoAndNormal) {
CUcontext cudaContext;
CU_CHECK(cuCtxGetCurrent(&cudaContext));
CHECK(cudaContext != nullptr);
OPTIX_CHECK(optixInit());
OptixDeviceContext optixContext;
OPTIX_CHECK(optixDeviceContextCreate(cudaContext, 0, &optixContext));
OptixDenoiserOptions options = {};
#if (OPTIX_VERSION >= 80000)
options.denoiseAlpha = OPTIX_DENOISER_ALPHA_MODE_COPY;
#endif
#if (OPTIX_VERSION >= 70300)
if (haveAlbedoAndNormal)
options.guideAlbedo = options.guideNormal = 1;
OPTIX_CHECK(optixDenoiserCreate(optixContext, OPTIX_DENOISER_MODEL_KIND_UPSCALE2X, &options,
&denoiserHandle));
#else
options.inputKind = haveAlbedoAndNormal ? OPTIX_DENOISER_INPUT_RGB_ALBEDO_NORMAL
: OPTIX_DENOISER_INPUT_RGB;
OPTIX_CHECK(optixDenoiserCreate(optixContext, &options, &denoiserHandle));
OPTIX_CHECK(
optixDenoiserSetModel(denoiserHandle, OPTIX_DENOISER_MODEL_KIND_UPSCALE2X, nullptr, 0));
#endif
OPTIX_CHECK(optixDenoiserComputeMemoryResources(denoiserHandle, resolution.x,
resolution.y, &memorySizes));
CUDA_CHECK(cudaMalloc(&denoiserState, memorySizes.stateSizeInBytes));
CUDA_CHECK(cudaMalloc(&scratchBuffer, memorySizes.withoutOverlapScratchSizeInBytes));
OPTIX_CHECK(optixDenoiserSetup(
denoiserHandle, 0 /* stream */, resolution.x, resolution.y,
CUdeviceptr(denoiserState), memorySizes.stateSizeInBytes,
CUdeviceptr(scratchBuffer), memorySizes.withoutOverlapScratchSizeInBytes));
CUDA_CHECK(cudaMalloc(&intensity, sizeof(float)));
}
void Scaler::Scale(RGB *rgb, Normal3f *n, RGB *albedo, RGB *result) {
std::array<OptixImage2D, 3> inputLayers;
int nLayers = haveAlbedoAndNormal ? 3 : 1;
for (int i = 0; i < nLayers; ++i) {
inputLayers[i].width = resolution.x;
inputLayers[i].height = resolution.y;
inputLayers[i].rowStrideInBytes = resolution.x * 3 * sizeof(float);
inputLayers[i].pixelStrideInBytes = 0;
inputLayers[i].format = OPTIX_PIXEL_FORMAT_FLOAT3;
}
inputLayers[0].data = CUdeviceptr(rgb);
if (haveAlbedoAndNormal) {
CHECK(n != nullptr && albedo != nullptr);
inputLayers[1].data = CUdeviceptr(albedo);
inputLayers[2].data = CUdeviceptr(n);
} else
CHECK(n == nullptr && albedo == nullptr);
OptixImage2D outputImage;
outputImage.width = resolution.x * 2;
outputImage.height = resolution.y * 2;
outputImage.rowStrideInBytes = resolution.x * 2 * 3 * sizeof(float);
outputImage.pixelStrideInBytes = 0;
outputImage.format = OPTIX_PIXEL_FORMAT_FLOAT3;
outputImage.data = CUdeviceptr(result);
OPTIX_CHECK(optixDenoiserComputeIntensity(
denoiserHandle, 0 /* stream */, &inputLayers[0], CUdeviceptr(intensity),
CUdeviceptr(scratchBuffer), memorySizes.withoutOverlapScratchSizeInBytes));
OptixDenoiserParams params = {};
#if (OPTIX_VERSION >= 80000)
// denoiseAlpha is moved to OptixDenoiserOptions in OptiX 8.0
#elif (OPTIX_VERSION >= 70500)
params.denoiseAlpha = OPTIX_DENOISER_ALPHA_MODE_COPY;
#else
params.denoiseAlpha = 0;
#endif
params.hdrIntensity = CUdeviceptr(intensity);
params.blendFactor = 0; // TODO what should this be??
#if (OPTIX_VERSION >= 70300)
OptixDenoiserGuideLayer guideLayer;
if (haveAlbedoAndNormal) {
guideLayer.albedo = inputLayers[1];
guideLayer.normal = inputLayers[2];
}
OptixDenoiserLayer layers;
layers.input = inputLayers[0];
layers.output = outputImage;
OPTIX_CHECK(optixDenoiserInvoke(
denoiserHandle, 0 /* stream */, &params, CUdeviceptr(denoiserState),
memorySizes.stateSizeInBytes, &guideLayer, &layers, 1 /* # layers to denoise */,
0 /* offset x */, 0 /* offset y */, CUdeviceptr(scratchBuffer),
memorySizes.withoutOverlapScratchSizeInBytes));
#else
OPTIX_CHECK(optixDenoiserInvoke(
denoiserHandle, 0 /* stream */, &params, CUdeviceptr(denoiserState),
memorySizes.stateSizeInBytes, inputLayers.data(), nLayers, 0 /* offset x */,
0 /* offset y */, &outputImage, CUdeviceptr(scratchBuffer),
memorySizes.withoutOverlapScratchSizeInBytes));
#endif
}
} // namespace pbrt

View file

@ -0,0 +1,35 @@
// pbrt is Copyright(c) 1998-2020 Matt Pharr, Wenzel Jakob, and Greg Humphreys.
// The pbrt source code is licensed under the Apache License, Version 2.0.
// SPDX: Apache-2.0
#ifndef PBRT_GPU_SCALER_H
#define PBRT_GPU_SCALER_H
#include <pbrt/pbrt.h>
#include <pbrt/util/color.h>
#include <pbrt/util/vecmath.h>
#include <optix.h>
namespace pbrt {
class Scaler {
public:
Scaler(Vector2i resolution, bool haveAlbedoAndNormal);
// All pointers should be to GPU memory.
// |n| and |albedo| should be nullptr iff \haveAlbedoAndNormal| is false.
void Scale(RGB *rgb, Normal3f *n, RGB *albedo, RGB *result);
private:
Vector2i resolution;
bool haveAlbedoAndNormal;
OptixDenoiser denoiserHandle;
OptixDenoiserSizes memorySizes;
void *denoiserState, *scratchBuffer, *intensity;
};
} // namespace pbrt
#endif // PBRT_GPU_SCALER_H