1#include "taichi/common/logging.h"
2#include "taichi/util/image_io.h"
3
4#include "stb_image.h"
5#include "stb_image_write.h"
6
7namespace taichi {
8
9void imwrite(const std::string &filename,
10 size_t ptr,
11 int resx,
12 int resy,
13 int comp) {
14 void *data = (void *)ptr;
15 TI_ASSERT_INFO(filename.size() >= 5, "Bad image file name");
16 int result = 0;
17 std::string suffix = filename.substr(filename.size() - 4);
18 if (suffix == ".png") {
19 result =
20 stbi_write_png(filename.c_str(), resx, resy, comp, data, comp * resx);
21 } else if (suffix == ".bmp") {
22 result = stbi_write_bmp(filename.c_str(), resx, resy, comp, data);
23 } else if (suffix == ".jpg") {
24 result = stbi_write_jpg(filename.c_str(), resx, resy, comp, data, 95);
25 } else {
26 TI_ERROR("Unknown image file suffix {}", suffix);
27 }
28
29 if (!result) {
30 TI_ERROR("Cannot write image file [{}]", filename);
31 }
32 TI_TRACE("saved image {}: {}x{}x{}", filename, resx, resy, comp);
33}
34
35std::vector<size_t> imread(const std::string &filename, int comp) {
36 int resx = 0, resy = 0;
37 void *data = stbi_load(filename.c_str(), &resx, &resy, &comp, comp);
38 if (!data) {
39 TI_ERROR("Cannot read image file [{}]", filename);
40 }
41 TI_TRACE("loaded image {}: {}x{}x{}", filename, resx, resy, comp);
42
43 std::vector<size_t> ret = {(size_t)data, (size_t)resx, (size_t)resy,
44 (size_t)comp};
45 return ret;
46}
47
48} // namespace taichi
49