Merge pull request #41925 from reduz/improve-exr-loading

Make EXR import format support all depths and channel configurations

1 files changed, 105 insertions, 53 deletions

diff --git a/modules/tinyexr/image_loader_tinyexr.cpp b/modules/tinyexr/image_loader_tinyexr.cpp
index 9e7266b95a..5bdcb84244 100644
--- a/modules/tinyexr/image_loader_tinyexr.cpp
+++ b/modules/tinyexr/image_loader_tinyexr.cpp
@@ -73,8 +73,10 @@ Error ImageLoaderTinyEXR::load_image(Ref<Image> p_image, FileAccess *f, bool p_f
 	}
 
 	// Read HALF channel as FLOAT. (GH-13490)
+	bool use_float16 = false;
 	for (int i = 0; i < exr_header.num_channels; i++) {
 		if (exr_header.pixel_types[i] == TINYEXR_PIXELTYPE_HALF) {
+			use_float16 = true;
 			exr_header.requested_pixel_types[i] = TINYEXR_PIXELTYPE_FLOAT;
 		}
 	}
@@ -102,33 +104,10 @@ Error ImageLoaderTinyEXR::load_image(Ref<Image> p_image, FileAccess *f, bool p_f
 			idxB = c;
 		} else if (strcmp(exr_header.channels[c].name, "A") == 0) {
 			idxA = c;
-		}
-	}
-
-	if (exr_header.num_channels == 1) {
-		// Grayscale channel only.
-		idxR = 0;
-		idxG = 0;
-		idxB = 0;
-		idxA = 0;
-	} else {
-		// Assume RGB(A)
-		if (idxR == -1) {
-			ERR_PRINT("TinyEXR: R channel not found.");
-			// @todo { free exr_image }
-			return ERR_FILE_CORRUPT;
-		}
-
-		if (idxG == -1) {
-			ERR_PRINT("TinyEXR: G channel not found.");
-			// @todo { free exr_image }
-			return ERR_FILE_CORRUPT;
-		}
-
-		if (idxB == -1) {
-			ERR_PRINT("TinyEXR: B channel not found.");
-			// @todo { free exr_image }
-			return ERR_FILE_CORRUPT;
+		} else if (strcmp(exr_header.channels[c].name, "Y") == 0) {
+			idxR = c;
+			idxG = c;
+			idxB = c;
 		}
 	}
 
@@ -138,14 +117,27 @@ Error ImageLoaderTinyEXR::load_image(Ref<Image> p_image, FileAccess *f, bool p_f
 	Image::Format format;
 	int output_channels = 0;
 
+	int channel_size = use_float16 ? 2 : 4;
 	if (idxA != -1) {
-		imgdata.resize(exr_image.width * exr_image.height * 8); //RGBA16
-		format = Image::FORMAT_RGBAH;
+		imgdata.resize(exr_image.width * exr_image.height * 4 * channel_size); //RGBA
+		format = use_float16 ? Image::FORMAT_RGBAH : Image::FORMAT_RGBAF;
 		output_channels = 4;
-	} else {
-		imgdata.resize(exr_image.width * exr_image.height * 6); //RGB16
-		format = Image::FORMAT_RGBH;
+	} else if (idxB != -1) {
+		ERR_FAIL_COND_V(idxG == -1, ERR_FILE_CORRUPT);
+		ERR_FAIL_COND_V(idxR == -1, ERR_FILE_CORRUPT);
+		imgdata.resize(exr_image.width * exr_image.height * 3 * channel_size); //RGB
+		format = use_float16 ? Image::FORMAT_RGBH : Image::FORMAT_RGBF;
 		output_channels = 3;
+	} else if (idxG != -1) {
+		ERR_FAIL_COND_V(idxR == -1, ERR_FILE_CORRUPT);
+		imgdata.resize(exr_image.width * exr_image.height * 2 * channel_size); //RG
+		format = use_float16 ? Image::FORMAT_RGH : Image::FORMAT_RGF;
+		output_channels = 2;
+	} else {
+		ERR_FAIL_COND_V(idxR == -1, ERR_FILE_CORRUPT);
+		imgdata.resize(exr_image.width * exr_image.height * 1 * channel_size); //R
+		format = use_float16 ? Image::FORMAT_RH : Image::FORMAT_RF;
+		output_channels = 1;
 	}
 
 	EXRTile single_image_tile;
@@ -175,9 +167,11 @@ Error ImageLoaderTinyEXR::load_image(Ref<Image> p_image, FileAccess *f, bool p_f
 		exr_tiles = exr_image.tiles;
 	}
 
+	//print_line("reading format: " + Image::get_format_name(format));
 	{
 		uint8_t *wd = imgdata.ptrw();
-		uint16_t *iw = (uint16_t *)wd;
+		uint16_t *iw16 = (uint16_t *)wd;
+		float *iw32 = (float *)wd;
 
 		// Assume `out_rgba` have enough memory allocated.
 		for (int tile_index = 0; tile_index < num_tiles; tile_index++) {
@@ -187,41 +181,99 @@ Error ImageLoaderTinyEXR::load_image(Ref<Image> p_image, FileAccess *f, bool p_f
 			int th = tile.height;
 
 			const float *r_channel_start = reinterpret_cast<const float *>(tile.images[idxR]);
-			const float *g_channel_start = reinterpret_cast<const float *>(tile.images[idxG]);
-			const float *b_channel_start = reinterpret_cast<const float *>(tile.images[idxB]);
+			const float *g_channel_start = nullptr;
+			const float *b_channel_start = nullptr;
 			const float *a_channel_start = nullptr;
 
+			if (idxG != -1) {
+				g_channel_start = reinterpret_cast<const float *>(tile.images[idxG]);
+			}
+			if (idxB != -1) {
+				b_channel_start = reinterpret_cast<const float *>(tile.images[idxB]);
+			}
 			if (idxA != -1) {
 				a_channel_start = reinterpret_cast<const float *>(tile.images[idxA]);
 			}
 
-			uint16_t *first_row_w = iw + (tile.offset_y * tile_height * exr_image.width + tile.offset_x * tile_width) * output_channels;
+			uint16_t *first_row_w16 = iw16 + (tile.offset_y * tile_height * exr_image.width + tile.offset_x * tile_width) * output_channels;
+			float *first_row_w32 = iw32 + (tile.offset_y * tile_height * exr_image.width + tile.offset_x * tile_width) * output_channels;
 
 			for (int y = 0; y < th; y++) {
 				const float *r_channel = r_channel_start + y * tile_width;
-				const float *g_channel = g_channel_start + y * tile_width;
-				const float *b_channel = b_channel_start + y * tile_width;
+				const float *g_channel = nullptr;
+				const float *b_channel = nullptr;
 				const float *a_channel = nullptr;
-
+				if (g_channel_start) {
+					g_channel = g_channel_start + y * tile_width;
+				}
+				if (b_channel_start) {
+					b_channel = b_channel_start + y * tile_width;
+				}
 				if (a_channel_start) {
 					a_channel = a_channel_start + y * tile_width;
 				}
 
-				uint16_t *row_w = first_row_w + (y * exr_image.width * output_channels);
-
-				for (int x = 0; x < tw; x++) {
-					Color color(*r_channel++, *g_channel++, *b_channel++);
-
-					if (p_force_linear) {
-						color = color.to_linear();
+				if (use_float16) {
+					uint16_t *row_w = first_row_w16 + (y * exr_image.width * output_channels);
+
+					for (int x = 0; x < tw; x++) {
+						Color color;
+						color.r = *r_channel++;
+						if (g_channel) {
+							color.g = *g_channel++;
+						}
+						if (b_channel) {
+							color.b = *b_channel++;
+						}
+						if (a_channel) {
+							color.a = *a_channel++;
+						}
+
+						if (p_force_linear) {
+							color = color.to_linear();
+						}
+
+						*row_w++ = Math::make_half_float(color.r);
+						if (g_channel) {
+							*row_w++ = Math::make_half_float(color.g);
+						}
+						if (b_channel) {
+							*row_w++ = Math::make_half_float(color.b);
+						}
+						if (a_channel) {
+							*row_w++ = Math::make_half_float(color.a);
+						}
 					}
-
-					*row_w++ = Math::make_half_float(color.r);
-					*row_w++ = Math::make_half_float(color.g);
-					*row_w++ = Math::make_half_float(color.b);
-
-					if (idxA != -1) {
-						*row_w++ = Math::make_half_float(*a_channel++);
+				} else {
+					float *row_w = first_row_w32 + (y * exr_image.width * output_channels);
+
+					for (int x = 0; x < tw; x++) {
+						Color color;
+						color.r = *r_channel++;
+						if (g_channel) {
+							color.g = *g_channel++;
+						}
+						if (b_channel) {
+							color.b = *b_channel++;
+						}
+						if (a_channel) {
+							color.a = *a_channel++;
+						}
+
+						if (p_force_linear) {
+							color = color.to_linear();
+						}
+
+						*row_w++ = color.r;
+						if (g_channel) {
+							*row_w++ = color.g;
+						}
+						if (b_channel) {
+							*row_w++ = color.b;
+						}
+						if (a_channel) {
+							*row_w++ = color.a;
+						}
 					}
 				}
 			}