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>

typedef unsigned int uint;

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);
  
  //int color_type = PNG_COLOR_TYPE_RGB;
  int color_type = PNG_COLOR_TYPE_PALETTE;

  // IHDRチャンク
  // 24bitフルカラーで出力
  ::png_set_IHDR(png_ptr, info_ptr, width, height,
    depth,
    //PNG_COLOR_TYPE_RGB_ALPHA,
    color_type, // 透過色いらないのではずしてファイルサイズ下げる
    //PNG_COLOR_TYPE_PALETTE,
    PNG_INTERLACE_NONE,
    PNG_COMPRESSION_TYPE_DEFAULT,
    PNG_FILTER_TYPE_DEFAULT
    );

  // 色の統計を取得する
  // 最も使われてる上位256色のカラーパレットを作りたいのだ
  /*
  map<int, int> color_sets;
  color_sets.insert( map<int,int>::value_type(1, 7) );
  color_sets.insert( map<int,int>::value_type(2, 3) );

  map<int, int>::iterator it = color_sets.find( 2 );
  printf("[%d] = %d\n", it->first, it->second);

  it->second = 10;

  it = color_sets.find( 2 );
  printf("[%d] = %d\n", it->first, it->second);
  exit(0);
  */
  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の順番に格納する、んでそれをキーにしてsetに格納
      // mapの右側は個数(カウントされた)
      uint packed = (r << 0) | (g << 8) | (b << 16);
      //color_sets.insert( map<int, int>::value_type(packed, 1) );
      //printf("merge: %d, %d, %d -> %d\n", r, g, b, packed);

      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;

    //printf("%d, %d, %d, %d\n", rit->first, r, g, b);
    rit++; count++;
  }
  
  // 256色のパレットを作成する
  /*
  colors[0].blue = 255; //青色
  colors[0].red = 0; //青色
  colors[0].green = 0; //青色
  colors[1].blue = 0;
  colors[1].red = 255;
  colors[1].green = 0;
  */
  ::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);
    //memset(row_pointers[y], 0, width * pixel_bytes); // 完全に透過な黒に初期化

    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] = x % 2; // 青
        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){
        // カラーインデックスの色を参照する
        // 見つからなかった場合は、完全に透過な黒(0)にする
        bool found = false;
        int orig_r = p[y * bmp.bmWidthBytes + x * 4 + 2];
        int orig_g = p[y * bmp.bmWidthBytes + x * 4 + 1];
        int orig_b = p[y * bmp.bmWidthBytes + x * 4 + 0];

        int i = 0;
        for(; i<256; i++){
          // 色がカラーパレットと完全に一致した場合はその色を採用する
          if(orig_r == colors[i].red &&
            orig_g == colors[i].green &&
            orig_b == colors[i].blue){
              found = true;
              break;
          }
        }

        if(found)
          row_pointers[y][x * pixel_bytes + 0] = i; // 青
        else{
          // 見つからん買ったときはもっとも近い色を見つけてそれを採用します
          row_pointers[y][x * pixel_bytes + 0] = 0;
        }
      }else{
        ;
      }
    }
  }
  ::png_set_rows(png_ptr, info_ptr, (png_bytepp)row_pointers);
  //::png_write_image(png_ptr, 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);
  //::Screenshot::SaveToFileAutoDetectFormat(h, L"hoge.png");

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

  ::DeleteObject(h);
}