Skip to content

Commit 2f4069e

Browse files
AlbertoValenzuelaRedondoAlberto ValenzuelaEdiolot
authored
GT 16 bits image support + lossless 16bits cache for HDR training [VOLINGA] (#1335)
* Templating load_image_t internally * FP32 image load support * Cuda kernel: HWC to CHW for fp32 * u16 loading path * Encoding jpeg2k * Decode wip (missing lanczos resize, maybe some decode path) * Finish u16 decoding path, adding support for float tensor on lanczos resize * Adding --use_8bit_color flag, disable with 0 to use HDR 16bits image input * Cahnges based on review * Restore rad and exclude export flag. Fixed 16uint pipeline * Added back decode stream --------- Co-authored-by: Alberto Valenzuela <avalenzuela@arquimea.com> Co-authored-by: Jorge Sierra <jsierra.siete@gmail.com>
1 parent 126f9d7 commit 2f4069e

14 files changed

Lines changed: 856 additions & 262 deletions

‎src/core/argument_parser.cpp‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -260,6 +260,7 @@ namespace {
260260
::args::ValueFlag<int> min_track_length(dataset_group, "min_track_length", "Minimum point track length for COLMAP sparse point import; 0 disables filtering", {"min-track-length"});
261261
::args::Flag no_cpu_cache(dataset_group, "no_cpu_cache", "Disable CPU memory caching (default: enabled)", {"no-cpu-cache"});
262262
::args::Flag no_fs_cache(dataset_group, "no_fs_cache", "Disable filesystem caching (default: enabled)", {"no-fs-cache"});
263+
::args::ValueFlag<bool> use_8bit_color(parser, "use_8bit_color", "if true - train with 8bit depth color images, else 16bits for HDR (default: true)", {"use_8bit_color"});
263264
::args::Flag undistort(dataset_group, "undistort", "Undistort images on-the-fly before training", {"undistort"});
264265
::args::MapFlag<std::string, std::string> centralize(dataset_group, "mode",
265266
"Centralize dataset origin: off, by_pointcloud, by_cameras (default: off)",
@@ -693,6 +694,7 @@ namespace {
693694
min_track_length_val = cli_option_present({"--min-track-length"}) ? std::optional<int>(::args::get(min_track_length)) : std::optional<int>(),
694695
no_cpu_cache_flag = static_cast<bool>(no_cpu_cache),
695696
no_fs_cache_flag = static_cast<bool>(no_fs_cache),
697+
use_8bit_color_val = use_8bit_color ? std::optional<bool>(::args::get(use_8bit_color)) : std::optional<bool>(),
696698
tcp_server_connection_port_val = tcp_server_connection_port ? std::optional<int>(::args::get(tcp_server_connection_port)) : std::optional<int>(),
697699
tcp_broadcast_connection_port_val = tcp_broadcast_connection_port ? std::optional<int>(::args::get(tcp_broadcast_connection_port)) : std::optional<int>(),
698700
tcp_connection_flag = bool(tcp_connection),
@@ -776,6 +778,7 @@ namespace {
776778
ds.loading_params.use_cpu_memory = false;
777779
if (no_fs_cache_flag)
778780
ds.loading_params.use_fs_cache = false;
781+
setVal(use_8bit_color_val, ds.loading_params.use_8bit_color);
779782
setVal(max_cap_val, opt.max_cap);
780783
setVal(tcp_server_connection_port_val, svs.tcp_server_connection_port);
781784
setVal(tcp_broadcast_connection_port_val, svs.tcp_broadcast_connection_port);

‎src/core/cuda/lanczos_resize/lanczos_resize.cu‎

Lines changed: 86 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,26 @@
2020

2121
namespace cg = cooperative_groups;
2222

23+
namespace {
24+
template <typename T>
25+
struct PixelTraits;
26+
27+
template <>
28+
struct PixelTraits<uint8_t> {
29+
__device__ __forceinline__ static float load_and_normalize(const uint8_t* ptr, int idx) {
30+
return static_cast<float>(ptr[idx]) / 255.0f;
31+
}
32+
};
33+
34+
template <>
35+
struct PixelTraits<float> {
36+
__device__ __forceinline__ static float load_and_normalize(const float* ptr, int idx) {
37+
return ptr[idx];
38+
}
39+
};
40+
41+
} // namespace
42+
2343
namespace lfs::core {
2444
namespace detail {
2545

@@ -68,15 +88,15 @@ namespace lfs::core {
6888
}
6989
}
7090

71-
template <uint32_t CHANNELS>
91+
template <uint32_t CHANNELS, typename T>
7292
__global__ void __launch_bounds__(BLOCK_X* BLOCK_Y)
7393
LanczosResampleCUDA(
7494
const int input_h, const int input_w,
7595
const int output_h, const int output_w,
7696
const int kernel_size,
7797
const float* __restrict__ pre_coef_x,
7898
const float* __restrict__ pre_coef_y,
79-
const uint8_t* __restrict__ input, // [H, W, C] uint8
99+
const T* __restrict__ input, // [H, W, C] T
80100
float* __restrict__ output // [C, H, W] float32
81101
) {
82102
const auto block = cg::this_thread_block();
@@ -115,8 +135,7 @@ namespace lfs::core {
115135
const float kernel_value = kernel_value_y * kernel_value_x;
116136

117137
for (int ch = 0; ch < CHANNELS; ch++) {
118-
const float pixel_value = (float)input[input_pix_id * 3 + ch] / 255.0f;
119-
accumulator[ch] += pixel_value * kernel_value;
138+
accumulator[ch] += PixelTraits<T>::load_and_normalize(input, input_pix_id * 3 + ch) * kernel_value;
120139
}
121140
}
122141
}
@@ -180,6 +199,60 @@ namespace lfs::core {
180199
output[pix.y * output_w + pix.x] = accumulator;
181200
}
182201

202+
template <typename T>
203+
static Tensor lanczos_resize_impl(
204+
const Tensor& input,
205+
int output_h,
206+
int output_w,
207+
int kernel_size,
208+
cudaStream_t cuda_stream) {
209+
const int input_h = static_cast<int>(input.size(0));
210+
const int input_w = static_cast<int>(input.size(1));
211+
const int channels = static_cast<int>(input.size(2));
212+
213+
if (channels != 3) {
214+
LOG_ERROR("lanczos_resize: Only 3-channel (RGB) images supported, got {}", channels);
215+
return Tensor();
216+
}
217+
218+
auto output = Tensor::empty(
219+
TensorShape({static_cast<size_t>(channels),
220+
static_cast<size_t>(output_h),
221+
static_cast<size_t>(output_w)}),
222+
Device::CUDA, DataType::Float32);
223+
224+
cudaMemsetAsync(output.data_ptr(), 0, output.bytes(), cuda_stream);
225+
226+
const uint32_t offset_step_x =
227+
(uint32_t)(kernel_size * (1.0 * input_w / output_w) * 2 + 1 + 0.5f);
228+
const uint32_t offset_step_y =
229+
(uint32_t)(kernel_size * (1.0 * input_h / output_h) * 2 + 1 + 0.5f);
230+
231+
float *coef_x, *coef_y;
232+
cudaMalloc(&coef_x, sizeof(float) * output_w * offset_step_x);
233+
cudaMalloc(&coef_y, sizeof(float) * output_h * offset_step_y);
234+
235+
const int threads = BLOCK_X * BLOCK_Y;
236+
detail::PreComputeCoef<<<(output_w + threads - 1) / threads, threads, 0, cuda_stream>>>(
237+
input_w, output_w, kernel_size, coef_x);
238+
detail::PreComputeCoef<<<(output_h + threads - 1) / threads, threads, 0, cuda_stream>>>(
239+
input_h, output_h, kernel_size, coef_y);
240+
241+
const dim3 tile_grid((output_w + BLOCK_X - 1) / BLOCK_X,
242+
(output_h + BLOCK_Y - 1) / BLOCK_Y);
243+
const dim3 block(BLOCK_X, BLOCK_Y, 1);
244+
245+
detail::LanczosResampleCUDA<NUM_CHANNELS, T>
246+
<<<tile_grid, block, 0, cuda_stream>>>(
247+
input_h, input_w, output_h, output_w, kernel_size,
248+
coef_x, coef_y,
249+
input.ptr<T>(), output.ptr<float>());
250+
251+
cudaFree(coef_x);
252+
cudaFree(coef_y);
253+
return output;
254+
}
255+
183256
} // namespace detail
184257

185258
Tensor lanczos_resize(
@@ -193,63 +266,22 @@ namespace lfs::core {
193266
LOG_ERROR("lanczos_resize: Input must be a valid CUDA tensor");
194267
return Tensor();
195268
}
196-
197-
if (input.dtype() != DataType::UInt8) {
198-
LOG_ERROR("lanczos_resize: Input must be UInt8 dtype");
199-
return Tensor();
200-
}
201-
202269
if (input.ndim() != 3) {
203270
LOG_ERROR("lanczos_resize: Input must be 3D tensor [H, W, C]");
204271
return Tensor();
205272
}
206273

207-
const int input_h = static_cast<int>(input.size(0));
208-
const int input_w = static_cast<int>(input.size(1));
209-
const int channels = static_cast<int>(input.size(2));
210-
211-
if (channels != 3) {
212-
LOG_ERROR("lanczos_resize: Only 3-channel (RGB) images supported, got {}", channels);
274+
if (input.dtype() == DataType::UInt8) {
275+
return detail::lanczos_resize_impl<uint8_t>(
276+
input, output_h, output_w, kernel_size, cuda_stream);
277+
} else if (input.dtype() == DataType::Float32) {
278+
return detail::lanczos_resize_impl<float>(
279+
input, output_h, output_w, kernel_size, cuda_stream);
280+
} else {
281+
LOG_ERROR("lanczos_resize: Unsupported dtype: {}",
282+
dtype_name(input.dtype()));
213283
return Tensor();
214284
}
215-
216-
auto output = Tensor::empty(
217-
TensorShape({static_cast<size_t>(channels), static_cast<size_t>(output_h), static_cast<size_t>(output_w)}),
218-
Device::CUDA,
219-
DataType::Float32);
220-
221-
cudaMemsetAsync(output.data_ptr(), 0, output.bytes(), cuda_stream);
222-
223-
const uint32_t offset_step_x = (uint32_t)(kernel_size * (1.0 * input_w / output_w) * 2 + 1 + 0.5f);
224-
const uint32_t offset_step_y = (uint32_t)(kernel_size * (1.0 * input_h / output_h) * 2 + 1 + 0.5f);
225-
226-
float* coef_x;
227-
float* coef_y;
228-
cudaMalloc(&coef_x, sizeof(float) * output_w * offset_step_x);
229-
cudaMalloc(&coef_y, sizeof(float) * output_h * offset_step_y);
230-
231-
detail::PreComputeCoef<<<(output_w + BLOCK_X * BLOCK_Y - 1) / (BLOCK_X * BLOCK_Y), BLOCK_X * BLOCK_Y, 0, cuda_stream>>>(
232-
input_w, output_w, kernel_size, coef_x);
233-
234-
detail::PreComputeCoef<<<(output_h + BLOCK_X * BLOCK_Y - 1) / (BLOCK_X * BLOCK_Y), BLOCK_X * BLOCK_Y, 0, cuda_stream>>>(
235-
input_h, output_h, kernel_size, coef_y);
236-
237-
const dim3 tile_grid((output_w + BLOCK_X - 1) / BLOCK_X, (output_h + BLOCK_Y - 1) / BLOCK_Y);
238-
const dim3 block(BLOCK_X, BLOCK_Y, 1);
239-
240-
detail::LanczosResampleCUDA<NUM_CHANNELS><<<tile_grid, block, 0, cuda_stream>>>(
241-
input_h, input_w,
242-
output_h, output_w,
243-
kernel_size,
244-
coef_x,
245-
coef_y,
246-
input.ptr<uint8_t>(),
247-
output.ptr<float>());
248-
249-
cudaFree(coef_x);
250-
cudaFree(coef_y);
251-
252-
return output;
253285
}
254286

255287
Tensor lanczos_resize_grayscale(

0 commit comments

Comments
 (0)