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

pngtest.cpp

#include <Windows.h>
#include "Screenshot.h"
#include <stdio.h>
#include <iostream>
#include "png.h"
#pragma comment(lib, "libpng15.lib")

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

#include <set>
#include <map>
#include <vector>

#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<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);
  //fwrite("test", 1, strlen("test") * sizeof(BYTE), 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チャンク
  // 24bitフルカラーで出力
  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()){
    sorted_map.insert(map<uint,uint>::value_type(it->second, 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){
      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++;
  }
  ::png_set_PLTE(png_ptr, info_ptr, colors, 256);

  // カキコミ関数の設定
  FILE *fp = ::_wfopen(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, 256, 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;
  rect.right = ::GetSystemMetrics(SM_CXVIRTUALSCREEN);
  rect.bottom = ::GetSystemMetrics(SM_CYVIRTUALSCREEN);

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

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

  ::DeleteObject(h);
}