Add Image.save_exr()

3 files changed, 322 insertions, 0 deletions

diff --git a/modules/tinyexr/image_saver_tinyexr.cpp b/modules/tinyexr/image_saver_tinyexr.cpp
new file mode 100644
index 0000000000..e1d42d3217
--- /dev/null
+++ b/modules/tinyexr/image_saver_tinyexr.cpp
@@ -0,0 +1,279 @@
+/*************************************************************************/
+/*  image_saver_tinyexr.cpp                                              */
+/*************************************************************************/
+/*                       This file is part of:                           */
+/*                           GODOT ENGINE                                */
+/*                      https://godotengine.org                          */
+/*************************************************************************/
+/* Copyright (c) 2007-2019 Juan Linietsky, Ariel Manzur.                 */
+/* Copyright (c) 2014-2019 Godot Engine contributors (cf. AUTHORS.md)    */
+/*                                                                       */
+/* Permission is hereby granted, free of charge, to any person obtaining */
+/* a copy of this software and associated documentation files (the       */
+/* "Software"), to deal in the Software without restriction, including   */
+/* without limitation the rights to use, copy, modify, merge, publish,   */
+/* distribute, sublicense, and/or sell copies of the Software, and to    */
+/* permit persons to whom the Software is furnished to do so, subject to */
+/* the following conditions:                                             */
+/*                                                                       */
+/* The above copyright notice and this permission notice shall be        */
+/* included in all copies or substantial portions of the Software.       */
+/*                                                                       */
+/* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,       */
+/* EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF    */
+/* MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.*/
+/* IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY  */
+/* CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,  */
+/* TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE     */
+/* SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.                */
+/*************************************************************************/
+
+#include "image_saver_tinyexr.h"
+#include "core/math/math_funcs.h"
+
+#include "thirdparty/tinyexr/tinyexr.h"
+
+static bool is_supported_format(Image::Format p_format) {
+	// This is checked before anything else.
+	// Mostly uncompressed formats are considered.
+	switch (p_format) {
+		case Image::FORMAT_RF:
+		case Image::FORMAT_RGF:
+		case Image::FORMAT_RGBF:
+		case Image::FORMAT_RGBAF:
+		case Image::FORMAT_RH:
+		case Image::FORMAT_RGH:
+		case Image::FORMAT_RGBH:
+		case Image::FORMAT_RGBAH:
+		case Image::FORMAT_R8:
+		case Image::FORMAT_RG8:
+		case Image::FORMAT_RGB8:
+		case Image::FORMAT_RGBA8:
+			return true;
+		default:
+			return false;
+	}
+}
+
+enum SrcPixelType {
+	SRC_FLOAT,
+	SRC_HALF,
+	SRC_BYTE
+};
+
+static SrcPixelType get_source_pixel_type(Image::Format p_format) {
+	switch (p_format) {
+		case Image::FORMAT_RF:
+		case Image::FORMAT_RGF:
+		case Image::FORMAT_RGBF:
+		case Image::FORMAT_RGBAF:
+			return SRC_FLOAT;
+		case Image::FORMAT_RH:
+		case Image::FORMAT_RGH:
+		case Image::FORMAT_RGBH:
+		case Image::FORMAT_RGBAH:
+			return SRC_HALF;
+		case Image::FORMAT_R8:
+		case Image::FORMAT_RG8:
+		case Image::FORMAT_RGB8:
+		case Image::FORMAT_RGBA8:
+			return SRC_BYTE;
+		default:
+			CRASH_NOW();
+	}
+}
+
+static int get_target_pixel_type(Image::Format p_format) {
+	switch (p_format) {
+		case Image::FORMAT_RF:
+		case Image::FORMAT_RGF:
+		case Image::FORMAT_RGBF:
+		case Image::FORMAT_RGBAF:
+			return TINYEXR_PIXELTYPE_FLOAT;
+		case Image::FORMAT_RH:
+		case Image::FORMAT_RGH:
+		case Image::FORMAT_RGBH:
+		case Image::FORMAT_RGBAH:
+		// EXR doesn't support 8-bit channels so in that case we'll convert
+		case Image::FORMAT_R8:
+		case Image::FORMAT_RG8:
+		case Image::FORMAT_RGB8:
+		case Image::FORMAT_RGBA8:
+			return TINYEXR_PIXELTYPE_HALF;
+		default:
+			CRASH_NOW();
+	}
+}
+
+static int get_pixel_type_size(int p_pixel_type) {
+	switch (p_pixel_type) {
+		case TINYEXR_PIXELTYPE_HALF:
+			return 2;
+		case TINYEXR_PIXELTYPE_FLOAT:
+			return 4;
+	}
+	CRASH_NOW();
+}
+
+static int get_channel_count(Image::Format p_format) {
+	switch (p_format) {
+		case Image::FORMAT_RF:
+		case Image::FORMAT_RH:
+		case Image::FORMAT_R8:
+			return 1;
+		case Image::FORMAT_RGF:
+		case Image::FORMAT_RGH:
+		case Image::FORMAT_RG8:
+			return 2;
+		case Image::FORMAT_RGBF:
+		case Image::FORMAT_RGBH:
+		case Image::FORMAT_RGB8:
+			return 3;
+		case Image::FORMAT_RGBAF:
+		case Image::FORMAT_RGBAH:
+		case Image::FORMAT_RGBA8:
+			return 4;
+		default:
+			CRASH_NOW();
+	}
+}
+
+Error save_exr(const String &p_path, const Ref<Image> &p_img, bool p_grayscale) {
+
+	Image::Format format = p_img->get_format();
+
+	if (!is_supported_format(format)) {
+		// Format not supported
+		print_error("Image format not supported for saving as EXR. Consider saving as PNG.");
+		return ERR_UNAVAILABLE;
+	}
+
+	EXRHeader header;
+	InitEXRHeader(&header);
+
+	EXRImage image;
+	InitEXRImage(&image);
+
+	const int max_channels = 4;
+
+	// Godot does not support more than 4 channels,
+	// so we can preallocate header infos on the stack and use only the subset we need
+	PoolByteArray channels[max_channels];
+	unsigned char *channels_ptrs[max_channels];
+	EXRChannelInfo channel_infos[max_channels];
+	int pixel_types[max_channels];
+	int requested_pixel_types[max_channels] = { -1 };
+
+	// Gimp and Blender are a bit annoying so order of channels isn't straightforward.
+	const int channel_mappings[4][4] = {
+		{ 0 }, // R
+		{ 1, 0 }, // GR
+		{ 2, 1, 0 }, // BGR
+		{ 2, 1, 0, 3 } // BGRA
+	};
+
+	int channel_count = get_channel_count(format);
+	ERR_FAIL_COND_V(p_grayscale && channel_count != 1, ERR_INVALID_PARAMETER);
+
+	int target_pixel_type = get_target_pixel_type(format);
+	int target_pixel_type_size = get_pixel_type_size(target_pixel_type);
+	SrcPixelType src_pixel_type = get_source_pixel_type(format);
+	const int pixel_count = p_img->get_width() * p_img->get_height();
+
+	const int *channel_mapping = channel_mappings[channel_count - 1];
+
+	{
+		PoolByteArray src_data = p_img->get_data();
+		PoolByteArray::Read src_r = src_data.read();
+
+		for (int channel_index = 0; channel_index < channel_count; ++channel_index) {
+
+			// De-interleave channels
+
+			PoolByteArray &dst = channels[channel_index];
+			dst.resize(pixel_count * target_pixel_type_size);
+
+			PoolByteArray::Write dst_w = dst.write();
+
+			if (src_pixel_type == SRC_FLOAT && target_pixel_type == TINYEXR_PIXELTYPE_FLOAT) {
+
+				// Note: we don't save mipmaps
+				CRASH_COND(src_data.size() < pixel_count * channel_count * target_pixel_type_size);
+
+				const float *src_rp = (float *)src_r.ptr();
+				float *dst_wp = (float *)dst_w.ptr();
+
+				for (int i = 0; i < pixel_count; ++i) {
+					dst_wp[i] = src_rp[channel_index + i * channel_count];
+				}
+
+			} else if (src_pixel_type == SRC_HALF && target_pixel_type == TINYEXR_PIXELTYPE_HALF) {
+
+				CRASH_COND(src_data.size() < pixel_count * channel_count * target_pixel_type_size);
+
+				const uint16_t *src_rp = (uint16_t *)src_r.ptr();
+				uint16_t *dst_wp = (uint16_t *)dst_w.ptr();
+
+				for (int i = 0; i < pixel_count; ++i) {
+					dst_wp[i] = src_rp[channel_index + i * channel_count];
+				}
+
+			} else if (src_pixel_type == SRC_BYTE && target_pixel_type == TINYEXR_PIXELTYPE_HALF) {
+
+				CRASH_COND(src_data.size() < pixel_count * channel_count);
+
+				const uint8_t *src_rp = (uint8_t *)src_r.ptr();
+				uint16_t *dst_wp = (uint16_t *)dst_w.ptr();
+
+				for (int i = 0; i < pixel_count; ++i) {
+					dst_wp[i] = Math::make_half_float(src_rp[channel_index + i * channel_count] / 255.f);
+				}
+
+			} else {
+				CRASH_NOW();
+			}
+
+			int remapped_index = channel_mapping[channel_index];
+
+			channels_ptrs[remapped_index] = dst_w.ptr();
+
+			// No conversion
+			pixel_types[remapped_index] = target_pixel_type;
+			requested_pixel_types[remapped_index] = target_pixel_type;
+
+			// Write channel name
+			if (p_grayscale) {
+				channel_infos[remapped_index].name[0] = 'Y';
+				channel_infos[remapped_index].name[1] = '\0';
+			} else {
+				const char *rgba = "RGBA";
+				channel_infos[remapped_index].name[0] = rgba[channel_index];
+				channel_infos[remapped_index].name[1] = '\0';
+			}
+		}
+	}
+
+	image.images = channels_ptrs;
+	image.num_channels = channel_count;
+	image.width = p_img->get_width();
+	image.height = p_img->get_height();
+
+	header.num_channels = image.num_channels;
+	header.channels = channel_infos;
+	header.pixel_types = pixel_types;
+	header.requested_pixel_types = requested_pixel_types;
+	// TODO DEBUG REMOVE
+	for (int i = 0; i < 4; ++i) {
+		print_line(String("requested_pixel_types{0}: {1}").format(varray(i, requested_pixel_types[i])));
+	}
+
+	CharString utf8_filename = p_path.utf8();
+	const char *err;
+	int ret = SaveEXRImageToFile(&image, &header, utf8_filename.ptr(), &err);
+	if (ret != TINYEXR_SUCCESS) {
+		print_error(String("Saving EXR failed. Error: {0}").format(varray(err)));
+		return ERR_FILE_CANT_WRITE;
+	}
+
+	return OK;
+}
diff --git a/modules/tinyexr/image_saver_tinyexr.h b/modules/tinyexr/image_saver_tinyexr.h
new file mode 100644
index 0000000000..298bd1d21c
--- /dev/null
+++ b/modules/tinyexr/image_saver_tinyexr.h
@@ -0,0 +1,38 @@
+/*************************************************************************/
+/*  image_saver_tinyexr.h                                                */
+/*************************************************************************/
+/*                       This file is part of:                           */
+/*                           GODOT ENGINE                                */
+/*                      https://godotengine.org                          */
+/*************************************************************************/
+/* Copyright (c) 2007-2019 Juan Linietsky, Ariel Manzur.                 */
+/* Copyright (c) 2014-2019 Godot Engine contributors (cf. AUTHORS.md)    */
+/*                                                                       */
+/* Permission is hereby granted, free of charge, to any person obtaining */
+/* a copy of this software and associated documentation files (the       */
+/* "Software"), to deal in the Software without restriction, including   */
+/* without limitation the rights to use, copy, modify, merge, publish,   */
+/* distribute, sublicense, and/or sell copies of the Software, and to    */
+/* permit persons to whom the Software is furnished to do so, subject to */
+/* the following conditions:                                             */
+/*                                                                       */
+/* The above copyright notice and this permission notice shall be        */
+/* included in all copies or substantial portions of the Software.       */
+/*                                                                       */
+/* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,       */
+/* EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF    */
+/* MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.*/
+/* IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY  */
+/* CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,  */
+/* TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE     */
+/* SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.                */
+/*************************************************************************/
+
+#ifndef IMAGE_SAVER_TINYEXR_H
+#define IMAGE_SAVER_TINYEXR_H
+
+#include "core/os/os.h"
+
+Error save_exr(const String &p_path, const Ref<Image> &p_img, bool p_grayscale);
+
+#endif // IMAGE_SAVER_TINYEXR_H
diff --git a/modules/tinyexr/register_types.cpp b/modules/tinyexr/register_types.cpp
index 5473a55687..233b3afa08 100644
--- a/modules/tinyexr/register_types.cpp
+++ b/modules/tinyexr/register_types.cpp
@@ -31,6 +31,7 @@
 #include "register_types.h"
 
 #include "image_loader_tinyexr.h"
+#include "image_saver_tinyexr.h"
 
 static ImageLoaderTinyEXR *image_loader_tinyexr = NULL;
 
@@ -38,9 +39,13 @@ void register_tinyexr_types() {
 
 	image_loader_tinyexr = memnew(ImageLoaderTinyEXR);
 	ImageLoader::add_image_format_loader(image_loader_tinyexr);
+
+	Image::save_exr_func = save_exr;
 }
 
 void unregister_tinyexr_types() {
 
 	memdelete(image_loader_tinyexr);
+
+	Image::save_exr_func = NULL;
 }