mirror of
https://github.com/mmp/pbrt-v4
synced 2026-09-26 16:20:07 +03:00
Merge pull request #513 from jflatt-gia/scale2x
Add Optix 2X scaler to imgtool
This commit is contained in:
commit
35bb9b88c0
4 changed files with 311 additions and 1 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
143
src/pbrt/gpu/optix/scaler.cpp
Normal file
143
src/pbrt/gpu/optix/scaler.cpp
Normal 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 */, ¶ms, 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 */, ¶ms, CUdeviceptr(denoiserState),
|
||||
memorySizes.stateSizeInBytes, inputLayers.data(), nLayers, 0 /* offset x */,
|
||||
0 /* offset y */, &outputImage, CUdeviceptr(scratchBuffer),
|
||||
memorySizes.withoutOverlapScratchSizeInBytes));
|
||||
#endif
|
||||
}
|
||||
|
||||
} // namespace pbrt
|
||||
35
src/pbrt/gpu/optix/scaler.h
Normal file
35
src/pbrt/gpu/optix/scaler.h
Normal 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
|
||||
Loading…
Add table
Add a link
Reference in a new issue