1// SPDX-License-Identifier: Apache-2.0
2// ----------------------------------------------------------------------------
3// Copyright 2021 Arm Limited
4//
5// Licensed under the Apache License, Version 2.0 (the "License"); you may not
6// use this file except in compliance with the License. You may obtain a copy
7// of the License at:
8//
9//     http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing, software
12// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
13// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
14// License for the specific language governing permissions and limitations
15// under the License.
16// ----------------------------------------------------------------------------
17
18// This is a utility tool to encode HDR into RGBM, or decode RGBM into HDR.
19
20#include <stdint.h>
21#include <stdio.h>
22#include <stdlib.h>
23
24#include "astcenc_mathlib.h"
25
26#define STB_IMAGE_IMPLEMENTATION
27#include "stb_image.h"
28
29#define STB_IMAGE_WRITE_IMPLEMENTATION
30#include "stb_image_write.h"
31
32#define MODE_ENCODE 0
33#define MODE_DECODE 1
34
35int main(int argc, char **argv)
36{
37	// Parse command line
38	if (argc != 6)
39	{
40		printf("Usage: astc_rgbm_codec [-ch|-dh] <M> <low_clamp> <source> <dest>\n");
41		exit(1);
42	}
43
44	int opmode;
45	if (strcmp(argv[1], "-ch") == 0)
46	{
47		opmode = MODE_ENCODE;
48	}
49	else if (strcmp(argv[1], "-dh") == 0)
50	{
51		opmode = MODE_DECODE;
52	}
53	else
54	{
55		printf("ERROR: Bad operation mode\n");
56		exit(1);
57	}
58
59 	float rgbm_multiplier = atof(argv[2]);
60 	float low_clamp = atof(argv[3]);
61
62	const char* src_file = argv[4];
63	const char* dst_file = argv[5];
64
65	// Convert an HDR input file into an RGBM encoded LDR file
66	if (opmode == MODE_ENCODE)
67	{
68		// Load the input image
69		int dim_x;
70		int dim_y;
71		const float* data_in = stbi_loadf(src_file, &dim_x, &dim_y, nullptr, 4);
72		if (!data_in)
73		{
74			printf("ERROR: Failed to load input image.\n");
75			exit(1);
76		}
77
78		// Allocate the output image
79		uint8_t* data_out = (uint8_t*)malloc(4 * dim_y * dim_x);
80		if (!data_out)
81		{
82			printf("ERROR: Failed to allow output image.\n");
83			exit(1);
84		}
85
86		// For each pixel apply RGBM encoding
87		for (int y = 0; y < dim_y; y++)
88		{
89			const float* row_in = data_in + (4 * dim_x * y);
90			uint8_t* row_out = data_out + (4 * dim_x * y);
91
92			for (int x = 0; x < dim_x; x++)
93			{
94				const float* pixel_in = row_in + 4 * x;
95				uint8_t* pixel_out = row_out + 4 * x;
96
97				float r_in = pixel_in[0] / rgbm_multiplier;
98				float g_in = pixel_in[1] / rgbm_multiplier;
99				float b_in = pixel_in[2] / rgbm_multiplier;
100
101				float max_rgb = astc::max(r_in, g_in, b_in);
102
103				// Ensure we always round up to next largest M
104				float m_scale = astc::min(1.0f, ceil(max_rgb * 255.0f) / 255.0f);
105
106				// But keep well above zero to avoid clamps in the compressor
107				m_scale = astc::max(m_scale, low_clamp / 255.0f);
108
109				float r_scale = astc::min(1.0f, r_in / m_scale);
110				float g_scale = astc::min(1.0f, g_in / m_scale);
111				float b_scale = astc::min(1.0f, b_in / m_scale);
112
113				pixel_out[0] = (uint8_t)(r_scale * 255.0f);
114				pixel_out[1] = (uint8_t)(g_scale * 255.0f);
115				pixel_out[2] = (uint8_t)(b_scale * 255.0f);
116				pixel_out[3] = (uint8_t)(m_scale * 255.0f);
117			}
118		}
119
120		// Write out the result
121		stbi_write_png(dst_file, dim_x, dim_y, 4, data_out, 4 * dim_x);
122	}
123	// Convert an RGBM encoded LDR file into an HDR file
124	else
125	{
126		// Load the input image
127		int dim_x;
128		int dim_y;
129		const uint8_t* data_in = stbi_load(src_file, &dim_x, &dim_y, nullptr, 4);
130		if (!data_in)
131		{
132			printf("ERROR: Failed to load input image.\n");
133			exit(1);
134		}
135
136		// Allocate the output image
137		float* data_out = (float*)malloc(4 * dim_y * dim_x * sizeof(float));
138		if (!data_out)
139		{
140			printf("ERROR: Failed to allow output image.\n");
141			exit(1);
142		}
143
144		// For each pixel apply RGBM decoding
145		for (int y = 0; y < dim_y; y++)
146		{
147			const uint8_t* row_in = data_in + (4 * dim_x * y);
148			float* row_out = data_out + (4 * dim_x * y);
149
150			for (int x = 0; x < dim_x; x++)
151			{
152				const uint8_t* pixel_in = row_in + 4 * x;
153				float* pixel_out = row_out + 4 * x;
154
155				float r_scale = ((float)pixel_in[0]) / 255.0f;
156				float g_scale = ((float)pixel_in[1]) / 255.0f;
157				float b_scale = ((float)pixel_in[2]) / 255.0f;
158
159				float m_scale = ((float)pixel_in[3]) / 255.0f;
160
161				pixel_out[0] = r_scale * (m_scale * rgbm_multiplier);
162				pixel_out[1] = g_scale * (m_scale * rgbm_multiplier);
163				pixel_out[2] = b_scale * (m_scale * rgbm_multiplier);
164				pixel_out[3] = 1.0f;
165			}
166		}
167
168		// Write out the result
169		stbi_write_hdr(dst_file, dim_x, dim_y, 4, data_out);
170	}
171
172	return 0;
173}
174