picture.cpp 2.37 KB
Newer Older
1 2
#include "picture.hpp"

3 4 5 6
#include <cmath>
#include <csetjmp>
#include <cstdio>

Lars Kuehne's avatar
Lars Kuehne committed
7
#include <iostream>
8
#include <stdexcept>
Lars Kuehne's avatar
Lars Kuehne committed
9 10 11 12 13 14
#include <vector>

#include <png.h>

#include "colormaps.hpp"

15 16
// all functions for picture transformation and output:

17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35
class FileWrapper {
  FILE *file_;

 public:
  FileWrapper(const std::string &filename, const char *open_mode)
      : file_(std::fopen(filename.c_str(), open_mode)) {
    if (!file_) throw std::runtime_error("error opening file " + filename);
  }

  ~FileWrapper() {
    if (file_) {
      fclose(file_);
      file_ = nullptr;
    }
  }

  operator FILE *() { return file_; }
};

36
// function: map result to color vector
Lars Kuehne's avatar
Lars Kuehne committed
37 38
bool write_png(const char *filename, const float *result, float threshold,
               int width, int height) {
39 40 41 42 43 44 45 46 47 48
  std::vector<unsigned char> colors_gray(width * height);
  std::vector<unsigned char> colors_rgb(3 * width * height);

  for (int i = 0; i < width * height; ++i) {
    if (result[i] > threshold) {
      colors_gray[i] = 255;
      colors_rgb[3 * i] = 255;
      colors_rgb[3 * i + 1] = 255;
      colors_rgb[3 * i + 2] = 255;
    } else {
49
      colors_gray[i] = std::floor(254 * result[i] / threshold);
50
      // RGB color gradient: viridis from matplotlib
51
      int idx = std::floor(255 * result[i] / threshold);
52

53 54 55
      unsigned char r = std::floor(255 * viridis[idx][0]);
      unsigned char g = std::floor(255 * viridis[idx][1]);
      unsigned char b = std::floor(255 * viridis[idx][2]);
56 57 58 59 60 61

      colors_rgb[3 * i] = r;      // red
      colors_rgb[3 * i + 1] = g;  // green
      colors_rgb[3 * i + 2] = b;  // blue
    }
  }
62 63 64 65 66 67 68 69 70

  FileWrapper file(filename, "wb");
  png_structp png =
      png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
  if (!png) return false;

  png_infop info = png_create_info_struct(png);
  if (!info) return false;

71
  if (setjmp(png_jmpbuf(png))) return false;
72 73 74 75 76 77 78 79 80 81

  png_init_io(png, file);

  // Output is 8bit depth, RGBA format.
  png_set_IHDR(png, info, width, height, 8, PNG_COLOR_TYPE_RGB,
               PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_DEFAULT,
               PNG_FILTER_TYPE_DEFAULT);
  png_write_info(png, info);

  std::vector<png_bytep> row_pointers(height);
82
  const int row_stride = width * 3;
83 84
  for (int i = 0; i < row_pointers.size(); ++i)
    row_pointers[i] = colors_rgb.data() + i * row_stride;
85

86 87 88
  png_write_image(png, row_pointers.data());
  png_write_end(png, NULL);
  return true;
89
}