PNG encoder (24bitフルカラー -> 256インデックスカラーに変換)

pngtest.cpp

/*
 *  デスクトップのスクリーンショットを撮影して
 *  256カラーインデックスPNGに減色後、圧縮します
 */
#include <Windows.h>
#include "Screenshot.h"
#include <stdio.h>

#include "png.h"
#pragma comment(lib, "libpng15.lib")

#include "pngconf.h"
#include "pngdebug.h"
#include "pnginfo.h"
#include "pnglibconf.h"
#include "pngstruct.h"

#include <map>
#include <math.h>

typedef unsigned int uint;

/*=========================
 * 色操作
 ==========================*/
// 色の距離を算出
double between_color(BYTE r1, BYTE g1, BYTE b1, BYTE r2, BYTE g2, BYTE b2){
  return sqrt( (double)(( r1 - r2 ) * ( r1 - r2 ) + ( g1 - g2 ) * ( g1 - g2 ) + ( b1 - b2 ) * ( b1 - b2 )) );
}

// 指定された色と完全一致するインデックスを返します
int find_color_by_clt(png_colorp table, size_t size, BYTE r, BYTE g, BYTE b){
  for(int i=0; i<(int)size; i++){
    if(table[i].red == r && table[i].green == g && table[i].blue == b){
      return i;
    }
  }
  return -1;
}

// 指定された色に近い、パレットのインデックスを返します
int find_nearcolor_by_clt(png_colorp table, size_t size, BYTE r, BYTE g, BYTE b){
  double min_between = 999;
  int min_index = -1;

  for(int i=0; i<(int)size; i++){
    double between = between_color(r, g, b, table[i].red, table[i].green, table[i].blue);
    if( between < min_between ){
      min_between = between;
      min_index = i;
    }
  }
  return min_index;
}

int find_bestcolor_by_clt(png_colorp table, size_t size, BYTE r, BYTE g, BYTE b)
{
  int index = -1;
  if( (index = ::find_color_by_clt(table, size, r, g, b)) == -1 ){
    return ::find_nearcolor_by_clt(table, size, r, g, b);
  }else{
    return index;
  }
}

png_color *find_by_palette(png_colorp palette, int num_palette, int color_index){
  int i;
  for(i=0; i<num_palette; i++){
    if( i == color_index )
      return (png_color *)&palette[i];
  }
  return NULL;
}

void PNG_write(png_structp png_ptr, png_bytep buf, png_size_t size)
{
  FILE *fp = (FILE *)png_get_io_ptr(png_ptr);
  if(!fp) return;
  fwrite(buf, sizeof(BYTE), size, fp);
}

void PNG_flush(png_structp png_ptr){
  FILE *fp = (FILE *)png_get_io_ptr(png_ptr);
  if(!fp) return;
  fflush(fp);
}

BOOL SaveToPngFile(HBITMAP h, LPCTSTR fileName)
{
  png_structp png_ptr = png_create_write_struct(
    PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
  if (!png_ptr) return (ERROR);

  if (setjmp(png_jmpbuf(png_ptr))){
    png_destroy_write_struct(&png_ptr, NULL);
    return (ERROR);
  }

  png_infop info_ptr = png_create_info_struct(png_ptr);
  if (!info_ptr) {
    png_destroy_write_struct(&png_ptr, 0);
    return (ERROR);
  }

  BITMAP bmp;
  ::GetObject(h, sizeof(BITMAP), &bmp);

  int len = bmp.bmWidth * bmp.bmHeight * (bmp.bmBitsPixel / 8);

  BYTE *p = (BYTE *)malloc(len * sizeof(BYTE));
  ::GetBitmapBits(h, len, p);

  int width = bmp.bmWidth;
  int height = bmp.bmHeight;
  int depth = 8;

  png_set_compression_level(png_ptr, Z_BEST_COMPRESSION);

  // IHDRチャンク
  int color_type = PNG_COLOR_TYPE_PALETTE;
  ::png_set_IHDR(png_ptr, info_ptr, width, height,
    depth,
    color_type,
    PNG_INTERLACE_NONE,
    PNG_COMPRESSION_TYPE_DEFAULT,
    PNG_FILTER_TYPE_DEFAULT
    );

  // 色の統計を取得する
  // 最も使われてる上位256色のカラーパレットを作りたい
  map<uint, uint>color_sets;
  for(int y=0; y<height; y++){
    for(int x=0; x<width; x++){
      int r = p[y * bmp.bmWidthBytes + x * 4 + 2];   // R
      int g = p[y * bmp.bmWidthBytes + x * 4 + 1];   // G
      int b = p[y * bmp.bmWidthBytes + x * 4 + 0];   // B

      // intを32bitとして、R,G,Bの順番に格納する、んでそれをキーにしてmapに格納
      // mapの右側は個数(カウントされた)
      uint packed = (r << 0) | (g << 8) | (b << 16);
    
      map<uint, uint>::iterator it = color_sets.find(packed);
      if( it == color_sets.end() ){
        color_sets.insert( map<uint,uint>::value_type(packed, 1) );
      }else{
        it->second++;
      }
    }
  }
  
  // 同じカウントの色が消えてしまう問題
  // value -> keyのmapに変換して(自動的にソートする)
  map<uint,uint>::iterator it = color_sets.begin();
  map<uint,uint> sorted_map;
  while(it != color_sets.end()){
    int count = it->second; // 色の出現回数を取得

    // その出現回数がすでに出力先マップに提供されていたら
    // かぶらないように+1する
    // まだかぶっていたらまた+1する、繰り返す
    map<uint,uint>::iterator tt = sorted_map.find(count);
    if(tt != sorted_map.end()){ // かぶってたら回避する
      count++;
    }
    sorted_map.insert(map<uint,uint>::value_type(count, it->first));
    // 同じ回数出現する色だとこのときに消えてしまう可能性あり
    it++;
  }

  // valueの高い順に256個取得してパレットを設定する
  int count = 0;
  png_color colors[256] = {0}; // 完全に透過な黒が初期化値
  map<uint,uint>::reverse_iterator rit = sorted_map.rbegin();
  while(rit != sorted_map.rend()){
    if(count >= 256){
      perror("256以上のカラーインデックスがあります");
      break;
    }

    int color = rit->second;
    int r = (color & 0xFF);
    int g = (color & 0xFF00) >> 8;
    int b = (color & 0xFF0000) >> 16;

    colors[count].red = r;
    colors[count].green = g;
    colors[count].blue = b;

    rit++; count++;
  }
  printf("%d個のカラーテーブルがあります\n", count);

  int color_table_count = count;

  ::png_set_PLTE(png_ptr, info_ptr, colors, color_table_count);

  // カキコミ関数の設定
  FILE *fp;
  ::_wfopen_s(&fp, fileName, L"wb");
  png_init_io (png_ptr, fp);
  ::png_set_write_fn(png_ptr, fp, PNG_write, PNG_flush);

  int pixel_bytes = 0;
  if( color_type == PNG_COLOR_TYPE_RGB )
    pixel_bytes = 3;
  else if( color_type == PNG_COLOR_TYPE_RGB_ALPHA )
    pixel_bytes = 4;
  else if( color_type == PNG_COLOR_TYPE_PALETTE )
    pixel_bytes = 1;
  else
    pixel_bytes = 4; // default = PNG_COLOR_TYPE_RGB_ALPHA

  png_byte **row_pointers = (png_byte **)png_malloc(png_ptr, sizeof(png_byte *) * height);
  for(int y=0; y<height; y++){
    row_pointers[y] = (png_byte*)png_malloc(png_ptr, pixel_bytes * width);
    for(int x=0; x<width; x++){
      if(color_type == PNG_COLOR_TYPE_RGB){
        row_pointers[y][x * pixel_bytes + 0] = (png_byte)p[y * bmp.bmWidthBytes + x * 4 + 2];   // R
        row_pointers[y][x * pixel_bytes + 1] = (png_byte)p[y * bmp.bmWidthBytes + x * 4 + 1];   // G
        row_pointers[y][x * pixel_bytes + 2] = (png_byte)p[y * bmp.bmWidthBytes + x * 4 + 0];   // B
      }else if(color_type == PNG_COLOR_TYPE_RGB_ALPHA){
        row_pointers[y][x * pixel_bytes + 0] = (png_byte)p[y * bmp.bmWidthBytes + x * 4 + 2];   // R
        row_pointers[y][x * pixel_bytes + 1] = (png_byte)p[y * bmp.bmWidthBytes + x * 4 + 1];   // G
        row_pointers[y][x * pixel_bytes + 2] = (png_byte)p[y * bmp.bmWidthBytes + x * 4 + 0];   // B
        row_pointers[y][x * pixel_bytes + 3] = (png_byte)255; // A(255 = 不透明, 0 = 透明)
      }else if(color_type == PNG_COLOR_TYPE_PALETTE){
        // カラーインデックスの色を参照する
        BYTE orig_r = p[y * bmp.bmWidthBytes + x * 4 + 2];
        BYTE orig_g = p[y * bmp.bmWidthBytes + x * 4 + 1];
        BYTE orig_b = p[y * bmp.bmWidthBytes + x * 4 + 0];
        
        int index = ::find_bestcolor_by_clt(colors, color_table_count, orig_r, orig_g, orig_b);
        if(index == -1){
          perror("近い色が見つかりませんでした");
          exit(1);
        }
        row_pointers[y][x * pixel_bytes + 0] = index;
      }else{
        perror("未対応のフォーマットです"); // 未対応
      }
    }
  }
  ::png_set_rows(png_ptr, info_ptr, (png_bytepp)row_pointers);
  ::png_write_png(png_ptr, info_ptr, PNG_TRANSFORM_IDENTITY, NULL);
  ::png_destroy_write_struct(&png_ptr, NULL);
  ::png_write_end(png_ptr, info_ptr);

  ::fclose(fp);
  return TRUE;
}

int main()
{
  // スクリーンショット撮影して、libpngで保存してみる
  HWND hWnd = ::GetDesktopWindow();
  RECT rect;
  ::GetWindowRect(hWnd, &rect);
  rect.right = ::GetSystemMetrics(SM_CXVIRTUALSCREEN);
  rect.bottom = ::GetSystemMetrics(SM_CYVIRTUALSCREEN);

  HBITMAP h = ::Screenshot::ScreenshotInMemory(NULL, &rect);

  // hbitmapをpngで保存してみる
  SaveToPngFile(h, L"output.png");

  ::DeleteObject(h);
}