streamtools/png-out.cpp

   1 #include "png-out.hpp"
   2 #include <zlib.h>
   3 #include <cstdio>
   4 #include <cstring>
   5 #include <stdexcept>
   6 #include <sstream>
   7
   8 namespace
   9 {
  10         void encode32(unsigned char* buf, uint32_t val)
  11         {
  12                 buf[0] = ((val >> 24) & 0xFF);
  13                 buf[1] = ((val >> 16) & 0xFF);
  14                 buf[2] = ((val >> 8) & 0xFF);
  15                 buf[3] = (val & 0xFF);
  16         }
  17
  18         uint32_t decode32(const unsigned char* buf)
  19         {
  20                 uint32_t v = 0;
  21                 v |= ((uint32_t)buf[0] << 24);
  22                 v |= ((uint32_t)buf[1] << 16);
  23                 v |= ((uint32_t)buf[2] << 8);
  24                 v |= ((uint32_t)buf[3]);
  25                 return v;
  26         }
  27
  28         void write_hunk(FILE* filp, uint32_t typecode, const unsigned char* data, uint32_t datasize)
  29         {
  30                 uLong crc = crc32(0, NULL, 0);
  31                 unsigned char fixed[12];
  32                 encode32(fixed, datasize);
  33                 encode32(fixed + 4, typecode);
  34                 crc = crc32(crc, fixed + 4, 4);
  35                 if(datasize > 0)
  36                         crc = crc32(crc, data, datasize);
  37                 encode32(fixed + 8, crc);
  38
  39                 if(fwrite(fixed, 1, 8, filp) < 8)
  40                         throw std::runtime_error("Can't write PNG file hunk header");
  41                 if(datasize > 0 && fwrite(data, 1, datasize, filp) < datasize)
  42                         throw std::runtime_error("Can't write PNG file hunk contents");
  43                 if(fwrite(fixed + 8, 1, 4, filp) < 4)
  44                         throw std::runtime_error("Can't write PNG file hunk checksum");
  45         }
  46
  47         void write_magic(FILE* filp)
  48         {
  49                 unsigned char magic[] = {137, 80, 78, 71, 13, 10, 26, 10};
  50                 if(fwrite(magic, 1, sizeof(magic), filp) < sizeof(magic))
  51                         throw std::runtime_error("Can't write PNG file magic");
  52         }
  53
  54         void write_ihdr(FILE* filp, struct image_frame& image)
  55         {
  56                 unsigned char ihdr[] = {25, 25, 25, 25, 25, 25, 25, 25, 8, 2, 0, 0, 0};
  57                 encode32(ihdr, image.get_width());
  58                 encode32(ihdr + 4, image.get_height());
  59                 write_hunk(filp, 0x49484452, ihdr, sizeof(ihdr));
  60         }
  61
  62         void write_iend(FILE* filp)
  63         {
  64                 write_hunk(filp, 0x49454E44, NULL, 0);
  65         }
  66
  67         void write_idat(FILE* filp, const unsigned char* data, uint32_t datasize)
  68         {
  69                 write_hunk(filp, 0x49444154, data, datasize);
  70         }
  71
  72         void write_idat_zlib_flush(FILE* filp, z_stream* s, uint32_t buffer)
  73         {
  74                 if(s->avail_out >= buffer)
  75                         return;
  76                 write_idat(filp, s->next_out + s->avail_out - buffer, buffer - s->avail_out);
  77         }
  78
  79         #define INBUF_SIZE 16384
  80         #define OUTBUF_SIZE 16384
  81
  82         int output_png2(FILE* filp, struct image_frame& image)
  83         {
  84                 write_magic(filp);
  85                 write_ihdr(filp, image);
  86
  87                 z_stream s;
  88                 memset(&s, 0, sizeof(s));
  89                 int r = deflateInit(&s, 9);
  90                 if(r < 0) {
  91                         std::stringstream str;
  92                         str << "deflateInit: zlib error: " << s.msg;
  93                         throw std::runtime_error(str.str());
  94                 }
  95
  96                 unsigned char in[INBUF_SIZE];
  97                 unsigned char out[OUTBUF_SIZE];
  98                 s.avail_out = OUTBUF_SIZE;
  99                 s.next_out = out;
 100                 int finish_flag = 0;
 101                 uint32_t total_data = (3 * image.get_width() * image.get_height() + image.get_height());
 102                 uint32_t data_emitted = 0;
 103                 uint32_t filter_divisior = 3 * image.get_width() + 1;
 104                 uint32_t dataptr = 0;
 105                 while(1) {
 106                         if(s.avail_in == 0) {
 107                                 s.next_in = in;
 108                                 //Input buffer empty. Fill it.
 109                                 while(data_emitted < total_data && s.avail_in < INBUF_SIZE) {
 110                                         if(data_emitted % filter_divisior == 0)
 111                                                 s.next_in[s.avail_in++] = 0;    //Filter none.
 112                                         else
 113                                                 s.next_in[s.avail_in++] = (image.get_pixels())[dataptr++];
 114                                         data_emitted++;
 115                                 }
 116                                 if(data_emitted == total_data)
 117                                         finish_flag = 1;
 118                         }
 119                         r = deflate(&s, finish_flag ? Z_FINISH : Z_NO_FLUSH);
 120                         if(r < 0) {
 121                                 std::stringstream str;
 122                                 str << "deflate: zlib error: " << s.msg;
 123                                 deflateEnd(&s);
 124                                 throw std::runtime_error(str.str());
 125                         }
 126                         if(s.avail_out == 0 || finish_flag) {
 127                                 try {
 128                                         write_idat_zlib_flush(filp, &s, OUTBUF_SIZE);
 129                                 } catch(...) {
 130                                         deflateEnd(&s);
 131                                         throw;
 132                                 }
 133                                 s.avail_out = OUTBUF_SIZE;
 134                                 s.next_out = out;
 135                         }
 136                         if(r == Z_STREAM_END)
 137                                 break;
 138                 }
 139
 140                 deflateEnd(&s);
 141                 write_iend(filp);
 142                 return 0;
 143         }
 144
 145
 146         bool output_png(const char* name, struct image_frame& image)
 147         {
 148                 if(!image.get_width() || !image.get_height())
 149                         return false;
 150
 151                 FILE* filp = fopen(name, "wb");
 152                 if(!filp) {
 153                         std::stringstream str;
 154                         str << "Can't open PNG output file '" << name << "'";
 155                         throw std::runtime_error(str.str());
 156                 }
 157
 158                 try {
 159                         output_png2(filp, image);
 160                 } catch(...) {
 161                         fclose(filp);
 162                         throw;
 163                 }
 164
 165                 if(fclose(filp) < 0) {
 166                         std::stringstream str;
 167                         str << "Can't close PNG output file '" << name << "'";
 168                         throw std::runtime_error(str.str());
 169                 }
 170                 return true;
 171         }
 172
 173         void copy4to3(unsigned char* target, const unsigned char* src, uint32_t pixels)
 174         {
 175                 for(uint32_t i = 0; i < pixels; i++) {
 176                         target[3 * i + 0] = src[4 * i + 0];
 177                         target[3 * i + 1] = src[4 * i + 1];
 178                         target[3 * i + 2] = src[4 * i + 2];
 179                 }
 180         }
 181
 182         void decode_zlib(unsigned char* target, const unsigned char* src, uint32_t insize, uint32_t pixels)
 183         {
 184                 unsigned char out[INBUF_SIZE];
 185                 unsigned char in[OUTBUF_SIZE];
 186                 uint32_t dptr = 0;
 187
 188                 z_stream s;
 189                 memset(&s, 0, sizeof(s));
 190                 int r = inflateInit(&s);
 191                 if(r < 0) {
 192                         std::stringstream str;
 193                         str << "inflateInit: zlib error: " << s.msg;
 194                         throw std::runtime_error(str.str());
 195                 }
 196
 197                 s.next_in = in;
 198                 s.avail_in = 0;
 199                 s.next_out = out;
 200                 s.avail_out = INBUF_SIZE;
 201
 202                 while(1) {
 203                         uint32_t old_avail_out = s.avail_out;
 204
 205                         if(s.avail_in == 0) {
 206                                 if(insize > OUTBUF_SIZE) {
 207                                         s.next_in = in;
 208                                         s.avail_in = OUTBUF_SIZE;
 209                                         memcpy(s.next_in, src, OUTBUF_SIZE);
 210                                         src += OUTBUF_SIZE;
 211                                         insize -= OUTBUF_SIZE;
 212                                 } else {
 213                                         s.next_in = in;
 214                                         s.avail_in = insize;
 215                                         memcpy(s.next_in, src, insize);
 216                                         src += insize;
 217                                         insize = 0;
 218                                 }
 219                         }
 220
 221                         r = inflate(&s, Z_NO_FLUSH);
 222                         if(r < 0) {
 223                                 inflateEnd(&s);
 224                                 std::stringstream str;
 225                                 str << "inflate: zlib error: " << s.msg;
 226                                 throw std::runtime_error(str.str());
 227                         }
 228                         dptr += (old_avail_out - s.avail_out);
 229                         if(s.avail_out == 0) {
 230                                 copy4to3(target, out, INBUF_SIZE / 4);
 231                                 target += 3 * (INBUF_SIZE / 4);
 232                                 s.avail_out = INBUF_SIZE;
 233                                 s.next_out = out;
 234                         }
 235                         if(dptr == 4 * pixels) {
 236                                 copy4to3(target, out, (INBUF_SIZE - s.avail_out) / 4);
 237                                 target += 3 * ((INBUF_SIZE - s.avail_out) / 4);
 238                                 s.avail_out = INBUF_SIZE;
 239                                 s.next_out = out;
 240                                 break;
 241                         }
 242                         if(r == Z_STREAM_END) {
 243                                 inflateEnd(&s);
 244                                 std::stringstream str;
 245                                 str << "Uncompressed stream truncated";
 246                                 throw std::runtime_error(str.str());
 247                         }
 248                 }
 249                 inflateEnd(&s);
 250         }
 251 }
 252
 253 const unsigned char* image_frame::get_pixels() const
 254 {
 255         return imagedata;
 256 }
 257
 258 unsigned char* image_frame::get_pixels()
 259 {
 260         return imagedata;
 261 }
 262
 263 uint32_t image_frame::get_height() const
 264 {
 265         return height;
 266 }
 267
 268 uint32_t image_frame::get_width() const
 269 {
 270         return width;
 271 }
 272
 273 image_frame::~image_frame()
 274 {
 275         delete[] imagedata;
 276 }
 277
 278 image_frame::image_frame(uint32_t width, uint32_t height)
 279 {
 280         this->width = width;
 281         this->height = height;
 282         this->imagedata = new unsigned char[3 * width * height];
 283 }
 284
 285 image_frame::image_frame(const image_frame& x)
 286 {
 287         this->width = x.width;
 288         this->height = x.height;
 289         this->imagedata = new unsigned char[3 * width * height];
 290         memcpy(imagedata, x.imagedata, 3 * width * height);
 291 }
 292
 293 image_frame& image_frame::operator=(const image_frame& x)
 294 {
 295         if(this == &x)
 296                 return *this;
 297         unsigned char* old_imagedata = imagedata;
 298         this->imagedata = new unsigned char[3 * width * height];
 299         this->width = x.width;
 300         this->height = x.height;
 301         memcpy(imagedata, x.imagedata, 3 * width * height);
 302         delete[] old_imagedata;
 303         return *this;
 304 }
 305
 306 image_frame::image_frame(struct packet& p)
 307 {
 308         if(p.rp_major != 0) {
 309                 std::stringstream str;
 310                 str << "frame_from_packet: Incorrect major type (" << p.rp_major << ", should be 0)";
 311                 throw std::runtime_error(str.str());
 312         }
 313         if(p.rp_minor != 0 && p.rp_minor != 1) {
 314                 std::stringstream str;
 315                 str << "frame_from_packet: Unknown minor type (" << p.rp_minor << ", should be 0 or 1)";
 316                 throw std::runtime_error(str.str());
 317         }
 318         if(p.rp_payload.size() < 4)
 319                 throw std::runtime_error("frame_from_packet: Malformed payload (image parameters missing)");
 320
 321         uint32_t ihdr = decode32(&p.rp_payload[0]);
 322         width = ihdr / 65536;
 323         height = ihdr % 65536;
 324         imagedata = new unsigned char[3 * width * height];
 325
 326         if(p.rp_minor == 0)
 327                 copy4to3(imagedata, &p.rp_payload[4], width * height);
 328         else if(p.rp_minor == 1)
 329                 try {
 330                         decode_zlib(imagedata, &p.rp_payload[4], p.rp_payload.size() - 4, width * height);
 331                 } catch(...) {
 332                         delete[] imagedata;
 333                         imagedata = NULL;
 334                 }
 335 }
 336
 337 bool image_frame::save_png(const std::string& name)
 338 {
 339         return output_png(name.c_str(), *this);
 340 }