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 | |
7 | namespace taichi { |
8 | |
9 | void 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 | |
35 | std::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 | |