1
2 #include <fstream>
3 #include <memory>
4
5 #include <pangolin/image/typed_image.h>
6
7 #ifdef HAVE_ZSTD
8 # include <zstd.h>
9 #endif
10
11 namespace pangolin {
12
13 #pragma pack(push, 1)
14 struct zstd_image_header
15 {
16 char magic[4];
17 char fmt[16];
18 size_t w, h;
19 };
20 #pragma pack(pop)
21
SaveZstd(const Image<unsigned char> & image,const pangolin::PixelFormat & fmt,std::ostream & out,int compression_level)22 void SaveZstd(const Image<unsigned char>& image, const pangolin::PixelFormat& fmt, std::ostream& out, int compression_level)
23 {
24 #ifdef HAVE_ZSTD
25 // Write out header, uncompressed
26 zstd_image_header header;
27 strncpy(header.magic,"ZSTD",4);
28 strncpy(header.fmt, fmt.format.c_str(), sizeof(header.fmt));
29 header.w = image.w;
30 header.h = image.h;
31 out.write((char*)&header, sizeof(header));
32
33 // Write out image data
34 const size_t output_buffer_size = ZSTD_CStreamOutSize();
35 std::unique_ptr<char[]> output_buffer(new char[output_buffer_size]);
36
37 ZSTD_CStream* const cstream = ZSTD_createCStream();
38 if (cstream==nullptr) {
39 throw std::runtime_error("ZSTD_createCStream() error");
40 }
41
42 size_t const initResult = ZSTD_initCStream(cstream, compression_level);
43 if (ZSTD_isError(initResult)) {
44 throw std::runtime_error(FormatString("ZSTD_initCStream() error : %", ZSTD_getErrorName(initResult)));
45 }
46
47 const size_t row_size_bytes = (fmt.bpp * image.w)/8;
48
49 for(size_t y=0; y < image.h; ++y) {
50 ZSTD_inBuffer input = { image.RowPtr(y), row_size_bytes, 0 };
51
52 while (input.pos < input.size) {
53 ZSTD_outBuffer output = { output_buffer.get(), output_buffer_size, 0 };
54 size_t left_to_read = ZSTD_compressStream(cstream, &output , &input);
55 if (ZSTD_isError(left_to_read)) {
56 throw std::runtime_error(FormatString("ZSTD_compressStream() error : %", ZSTD_getErrorName(left_to_read)));
57 }
58 out.write(output_buffer.get(), output.pos);
59 }
60 }
61
62 ZSTD_outBuffer output = { output_buffer.get(), output_buffer_size, 0 };
63 size_t const remainingToFlush = ZSTD_endStream(cstream, &output); /* close frame */
64 if (remainingToFlush) {
65 throw std::runtime_error("not fully flushed");
66 }
67 out.write(output_buffer.get(), output.pos);
68
69 ZSTD_freeCStream(cstream);
70 #else
71 PANGOLIN_UNUSED(image);
72 PANGOLIN_UNUSED(fmt);
73 PANGOLIN_UNUSED(out);
74 PANGOLIN_UNUSED(compression_level);
75 throw std::runtime_error("Rebuild Pangolin for ZSTD support.");
76 #endif // HAVE_ZSTD
77 }
78
LoadZstd(std::istream & in)79 TypedImage LoadZstd(std::istream& in)
80 {
81 #ifdef HAVE_ZSTD
82 // Read in header, uncompressed
83 zstd_image_header header;
84 in.read( (char*)&header, sizeof(header));
85
86 TypedImage img(header.w, header.h, PixelFormatFromString(header.fmt));
87
88 const size_t input_buffer_size = ZSTD_DStreamInSize();
89 std::unique_ptr<char[]> input_buffer(new char[input_buffer_size]);
90
91 ZSTD_DStream* dstream = ZSTD_createDStream();
92 if(!dstream) {
93 throw std::runtime_error("ZSTD_createDStream() error");
94 }
95
96 size_t read_size_hint = ZSTD_initDStream(dstream);
97 if (ZSTD_isError(read_size_hint)) {
98 throw std::runtime_error(FormatString("ZSTD_initDStream() error : % \n", ZSTD_getErrorName(read_size_hint)));
99 }
100
101 // Image represents our fixed buffer.
102 ZSTD_outBuffer output = { img.ptr, img.SizeBytes(), 0 };
103
104 while(read_size_hint)
105 {
106 const size_t read = in.readsome(input_buffer.get(), read_size_hint);
107 ZSTD_inBuffer input = { input_buffer.get(), read, 0 };
108 while (input.pos < input.size) {
109 read_size_hint = ZSTD_decompressStream(dstream, &output , &input);
110 if (ZSTD_isError(read_size_hint)) {
111 throw std::runtime_error(FormatString("ZSTD_decompressStream() error : %", ZSTD_getErrorName(read_size_hint)));
112 }
113 }
114 }
115
116 ZSTD_freeDStream(dstream);
117
118 return img;
119 #else
120 PANGOLIN_UNUSED(in);
121 throw std::runtime_error("Rebuild Pangolin for ZSTD support.");
122 #endif // HAVE_ZSTD
123 }
124
125 }
126