1/*******************************************************************************
2 Copyright (c) The Taichi Authors (2016- ). All Rights Reserved.
3 The use of this software is governed by the LICENSE file.
4*******************************************************************************/
5
6#include "taichi/math/math.h"
7#include "taichi/math/linalg.h"
8#include "taichi/util/base64.h"
9
10#define STBI_FAILURE_USERMSG
11#define STB_IMAGE_IMPLEMENTATION
12#include "stb_image.h"
13#define STB_IMAGE_WRITE_IMPLEMENTATION
14#include "stb_image_write.h"
15#define STB_TRUETYPE_IMPLEMENTATION
16#include "stb_truetype.h"
17
18namespace taichi {
19
20template <typename T>
21void Array2D<T>::load_image(const std::string &filename, bool linearize) {
22 int channels;
23 FILE *f = fopen(filename.c_str(), "rb");
24 TI_ASSERT_INFO(f != nullptr, "Image file not found: " + filename);
25 stbi_ldr_to_hdr_gamma(1.0_f);
26 float32 *data =
27 stbi_loadf(filename.c_str(), &this->res[0], &this->res[1], &channels, 0);
28 TI_ASSERT_INFO(data != nullptr,
29 "Image file load failed: " + filename +
30 " # Msg: " + std::string(stbi_failure_reason()));
31 TI_ASSERT_INFO(channels == 1 || channels == 3 || channels == 4,
32 "Image must have channel 1, 3 or 4: " + filename);
33 this->initialize(Vector2i(this->res[0], this->res[1]));
34
35 for (int i = 0; i < this->res[0]; i++) {
36 for (int j = 0; j < this->res[1]; j++) {
37 float32 *pixel_ =
38 data + ((this->res[1] - 1 - j) * this->res[0] + i) * channels;
39 Vector4 pixel;
40 if (channels == 1) {
41 pixel = Vector4(pixel_[0]);
42 } else {
43 pixel = Vector4(pixel_[0], pixel_[1], pixel_[2], pixel_[3]);
44 }
45 if (linearize) {
46 pixel = pixel.pow(2.2f);
47 }
48 (*this)[i][j][0] = pixel[0];
49 (*this)[i][j][1] = pixel[1];
50 (*this)[i][j][2] = pixel[2];
51 if (channels == 4 && std::is_same<T, Vector4>::value)
52 (*this)[i][j][3] = pixel[3];
53 }
54 }
55
56 stbi_image_free(data);
57}
58
59template <typename T>
60void Array2D<T>::write_as_image(const std::string &filename) {
61 int comp = 3;
62 std::vector<unsigned char> data(this->res[0] * this->res[1] * comp);
63 for (int i = 0; i < this->res[0]; i++) {
64 for (int j = 0; j < this->res[1]; j++) {
65 for (int k = 0; k < comp; k++) {
66 data[j * this->res[0] * comp + i * comp + k] =
67 (unsigned char)(255.0f *
68 clamp(VectorND<3, real>(
69 this->data[i * this->res[1] +
70 (this->res[1] - j - 1)])[k],
71 0.0_f, 1.0_f));
72 }
73 }
74 }
75 TI_ASSERT(filename.size() >= 5);
76 int write_result = 0;
77 std::string suffix = filename.substr(filename.size() - 4);
78 if (suffix == ".png") {
79 write_result = stbi_write_png(filename.c_str(), this->res[0], this->res[1],
80 comp, &data[0], comp * this->res[0]);
81 } else if (suffix == ".bmp") {
82 // TODO: test
83 write_result = stbi_write_bmp(filename.c_str(), this->res[0], this->res[1],
84 comp, &data[0]);
85 } else if (suffix == ".jpg") {
86 // TODO: test
87 write_result = stbi_write_jpg(filename.c_str(), this->res[0], this->res[1],
88 comp, &data[0], 95);
89 } else {
90 TI_ERROR("Unknown suffix {}", suffix);
91 }
92
93 TI_ASSERT_INFO((bool)write_result, "Cannot write image file");
94}
95
96std::map<std::string, stbtt_fontinfo> fonts;
97std::map<std::string, std::vector<uint8>> font_buffers;
98
99template <typename T>
100void Array2D<T>::write_text(const std::string &font_fn,
101 const std::string &content_,
102 real size,
103 int dx,
104 int dy,
105 T color) {
106 std::vector<unsigned char> screen_buffer(
107 (size_t)(this->res[0] * this->res[1]), (unsigned char)0);
108
109 int i, j, ascent, baseline, ch = 0;
110 float xpos = 2; // leave a little padding in case the character extends left
111
112 stbtt_fontinfo font;
113 if (fonts.find(font_fn) == fonts.end()) {
114 auto buffer_size = 24 << 20;
115 font_buffers[font_fn] =
116 std::vector<unsigned char>(buffer_size, (unsigned char)0);
117 if (font_fn != "") {
118 FILE *font_file = fopen(font_fn.c_str(), "rb");
119 TI_ASSERT_INFO(font_file != nullptr,
120 "Font file not found: " + std::string(font_fn));
121 trash(fread(&font_buffers[font_fn][0], 1, buffer_size, font_file));
122 fclose(font_file);
123 } else {
124 TI_NOT_IMPLEMENTED
125 }
126 stbtt_InitFont(&font, &font_buffers[font_fn][0], 0);
127 fonts[font_fn] = font;
128 } else {
129 font = fonts[font_fn];
130 }
131
132 real scale = stbtt_ScaleForPixelHeight(&font, size);
133
134 stbtt_GetFontVMetrics(&font, &ascent, nullptr, nullptr);
135 baseline = (int)(ascent * scale);
136 const std::string c_content = content_;
137 const char *content = c_content.c_str();
138 while (content[ch]) {
139 int advance, lsb, x0, y0, x1, y1;
140 float32 x_shift = xpos - (float32)floor(xpos);
141 stbtt_GetCodepointHMetrics(&font, content[ch], &advance, &lsb);
142 stbtt_GetCodepointBitmapBoxSubpixel(&font, content[ch], scale, scale,
143 x_shift, 0, &x0, &y0, &x1, &y1);
144 stbtt_MakeCodepointBitmapSubpixel(
145 &font,
146 &screen_buffer[0] + this->res[0] * (baseline + y0) + (int)xpos + x0,
147 x1 - x0, y1 - y0, this->res[0], scale, scale, x_shift, 0, content[ch]);
148 // note that this stomps the old data, so where character boxes overlap
149 // (e.g. 'lj') it's wrong
150 xpos += (advance * scale);
151 if (content[ch + 1])
152 xpos += scale * stbtt_GetCodepointKernAdvance(&font, content[ch],
153 content[ch + 1]);
154 ++ch;
155 }
156 for (j = 0; j < this->res[1]; ++j) {
157 for (i = 0; i < this->res[0]; ++i) {
158 int x = dx + i, y = dy + j - this->res[1];
159 auto index = ((this->res[1] - j - 1) * this->res[0] + i);
160 real alpha = screen_buffer[index] / 255.0f;
161 if (inside(x, y) && alpha != 0) {
162 (*this)[x][y] = lerp(alpha, this->get(x, y), color);
163 }
164 }
165 }
166}
167
168template void Array2D<Vector3>::write_text(const std::string &font_fn,
169 const std::string &content_,
170 real size,
171 int dx,
172 int dy,
173 Vector3);
174
175template void Array2D<Vector4>::write_text(const std::string &font_fn,
176 const std::string &content_,
177 real size,
178 int dx,
179 int dy,
180 Vector4);
181
182template void Array2D<Vector3f>::load_image(const std::string &filename, bool);
183
184template void Array2D<Vector4f>::load_image(const std::string &filename, bool);
185
186template void Array2D<Vector3d>::load_image(const std::string &filename, bool);
187
188template void Array2D<Vector4d>::load_image(const std::string &filename, bool);
189
190template void Array2D<float32>::write_as_image(const std::string &filename);
191
192template void Array2D<float64>::write_as_image(const std::string &filename);
193
194template void Array2D<Vector3f>::write_as_image(const std::string &filename);
195
196template void Array2D<Vector4f>::write_as_image(const std::string &filename);
197
198template void Array2D<Vector3d>::write_as_image(const std::string &filename);
199
200template void Array2D<Vector4d>::write_as_image(const std::string &filename);
201
202void write_pgm(Array2D<real> img, const std::string &fn) {
203 std::ofstream fs(fn, std::ios_base::binary);
204 Vector2i res = img.get_res();
205 fs << fmt::format("P5\n{} {}\n{}\n", res[0], res[1], 255);
206 for (int j = 0; j < res[1]; j++) {
207 std::string line;
208 for (int i = 0; i < res[0]; i++) {
209 uint8_t v = clamp((int)(img[i][res[1] - j - 1] * 255), 0, 255);
210 line.push_back(v);
211 }
212 fs.write(line.c_str(), line.size());
213 }
214}
215
216} // namespace taichi
217