huffyuv.cpp

mpeg2的颜色转换程序,由BEN完成

/********************************************************************
********************************************************************/


// 0=unknown, -1=compressed YUY2, -2=compressed RGB, -3=compressed RGBA, 1=YUY2, 2=UYVY, 3=RGB 24-bit, 4=RGB 32-bit
static int GetBitmapType(LPBITMAPINFOHEADER lpbi) {
  if (!lpbi)
    return 0;
  const int fourcc = lpbi->biCompression;
  if (fourcc == FOURCC_VYUY || fourcc == FOURCC_YUY2)
    return 1;
  if (fourcc == FOURCC_UYVY)
    return 2;
  const int bitcount = GetBitCount(lpbi);
  if (fourcc == 0 || fourcc == ' BID')
    return (bitcount == 24) ? 3 : (bitcount == 32) ? 4 : 0;
  if (fourcc == FOURCC_HFYU)
    return ((bitcount&~7) == 16) ? -1 : ((bitcount&~7) == 24) ? -2 : ((bitcount&~7) == 32) ? -3 : 0;
  return 0;
}


static bool CanCompress(LPBITMAPINFOHEADER lpbiIn) {
  int intype = GetBitmapType(lpbiIn);
  return (intype == 1 || intype == 2 || intype == 3 || (intype == 4 && AllowRGBA()));
}


static bool CanCompress(LPBITMAPINFOHEADER lpbiIn, LPBITMAPINFOHEADER lpbiOut) {
  if (!lpbiOut) return CanCompress(lpbiIn);

  int intype = GetBitmapType(lpbiIn);
  int outtype = GetBitmapType(lpbiOut);

  if (outtype < 0) {
    if (!IsLegalMethod(GetMethod(lpbiOut), outtype != -1))
      return false;
  }

  switch (intype) {
    case 1: case 2:
      return (outtype == -1);
    case 3:
      return (outtype == 1 || outtype == -1 || outtype == -2);
    case 4:
      return (outtype == -3 && AllowRGBA());
    default:
      return false;
  }
}


/********************************************************************
********************************************************************/


const unsigned char* GetHuffTable(int method, bool rgb) {
  if (rgb) {
    if (method == methodLeft)
      return left_rgb;
    else if (method == methodLeft+flagDecorrelate)
      return left_decorrelate_rgb;
    else
      return grad_decorrelate_rgb;
  } else {
    if (method == methodLeft)
      return left_yuv;
    else if (method == methodGrad)
      return grad_yuv;
    else
      return med_yuv;
  }
}


const unsigned char* DecompressHuffmanTable(const unsigned char* hufftable, unsigned char* dst) {
  int i=0;
  do {
    int val = *hufftable & 31;
    int repeat = *hufftable++ >> 5;
    if (!repeat)
      repeat = *hufftable++;
    while (repeat--)
      dst[i++] = val;
  } while (i<256);
  return hufftable;
}


#define HUFFTABLE_CLASSIC_YUV ((const unsigned char*)-1)
#define HUFFTABLE_CLASSIC_RGB ((const unsigned char*)-2)
#define HUFFTABLE_CLASSIC_YUV_CHROMA ((const unsigned char*)-3)

const unsigned char* InitializeShiftAddTables(const unsigned char* hufftable, unsigned char* shift, unsigned* add_shifted) {
  // special-case the old tables, since they don't fit the new rules
  if (hufftable == HUFFTABLE_CLASSIC_YUV || hufftable == HUFFTABLE_CLASSIC_RGB) {
    DecompressHuffmanTable(classic_shift_luma, shift);
    for (int i=0; i<256; ++i) add_shifted[i] = classic_add_luma[i]<<(32-shift[i]);
    return (hufftable == HUFFTABLE_CLASSIC_YUV) ? HUFFTABLE_CLASSIC_YUV_CHROMA : hufftable;
  }
  else if (hufftable == HUFFTABLE_CLASSIC_YUV_CHROMA) {
    DecompressHuffmanTable(classic_shift_chroma, shift);
    for (int i=0; i<256; ++i) add_shifted[i] = classic_add_chroma[i]<<(32-shift[i]);
    return hufftable;
  }

  hufftable = DecompressHuffmanTable(hufftable, shift);
  // derive the actual bit patterns from the code lengths
  int min_already_processed = 32;
  unsigned __int32 bits = 0;
  do {
    int max_not_processed=0;
    for (int i=0; i<256; ++i) {
      if (shift[i] < min_already_processed && shift[i] > max_not_processed)
        max_not_processed = shift[i];
    }
    int bit = 1<<(32-max_not_processed);
    _ASSERTE(!(bits & (bit-1)));
    for (int j=0; j<256; ++j) {
      if (shift[j] == max_not_processed) {
        add_shifted[j] = bits;
        bits += bit;
      }
    }
    min_already_processed = max_not_processed;
  } while (bits&0xFFFFFFFF);
  return hufftable;
}


const unsigned char* InitializeEncodeTables(const unsigned char* hufftable) {
  hufftable = InitializeShiftAddTables(hufftable, encode1_shift, encode1_add_shifted);
  hufftable = InitializeShiftAddTables(hufftable, encode2_shift, encode2_add_shifted);
  hufftable = InitializeShiftAddTables(hufftable, encode3_shift, encode3_add_shifted);
//  _ASSERTE(!*hufftable);
  return hufftable;
}


DWORD CodecInst::CompressQuery(LPBITMAPINFOHEADER lpbiIn, LPBITMAPINFOHEADER lpbiOut) {
  Msg("CompressQuery: input = %s, output = %s\n", &PrintBitmapType(lpbiIn), &PrintBitmapType(lpbiOut));
  return CanCompress(lpbiIn, lpbiOut) ? ICERR_OK : ICERR_BADFORMAT;
}


DWORD CodecInst::CompressGetFormat(LPBITMAPINFOHEADER lpbiIn, LPBITMAPINFOHEADER lpbiOut) {
  if (!CanCompress(lpbiIn))
    return ICERR_BADFORMAT;

  int intype = GetBitmapType(lpbiIn);

  bool compress_as_yuy2 = (intype<3);
  int method;
  if (!compress_as_yuy2) {
    method = GetPrivateProfileInt("general", "rgbmethod", methodGrad+flagDecorrelate, "huffyuv.ini");
    if (method==methodConvertToYUY2)
      compress_as_yuy2 = true;
  }
  if (compress_as_yuy2)
    method = GetPrivateProfileInt("general", "yuy2method", methodMedian, "huffyuv.ini");

  const unsigned char* hufftable = GetHuffTable(method, !compress_as_yuy2);

  int hufftable_length = (lstrlen((const char*)hufftable)+4) & -4;

  if (!lpbiOut)
    return sizeof(BITMAPINFOHEADER)+4+hufftable_length;

  int real_bit_count = (intype==4) ? 32 : (compress_as_yuy2) ? 16 : 24;

  *lpbiOut = *lpbiIn;
  lpbiOut->biSize = sizeof(BITMAPINFOHEADER)+4+hufftable_length;
  lpbiOut->biCompression = FOURCC_HFYU;
  lpbiOut->biClrImportant = lpbiOut->biClrUsed = 0;
  lpbiOut->biBitCount = max(24, real_bit_count);
  unsigned char* extra_data = (unsigned char*)lpbiOut + sizeof(BITMAPINFOHEADER);
  extra_data[0] = method;
  extra_data[1] = real_bit_count;
  extra_data[2] = 0;
  extra_data[3] = 0;
  lstrcpy((char*)extra_data+4, (const char*)hufftable);
  return ICERR_OK;
}


const unsigned char* GetEmbeddedHuffmanTable(LPBITMAPINFOHEADER lpbi) {
  if (lpbi->biSize == sizeof(BITMAPINFOHEADER) && !(lpbi->biBitCount&7))
    return lpbi->biBitCount>=24 ? HUFFTABLE_CLASSIC_RGB : HUFFTABLE_CLASSIC_YUV;
  else
    return (const unsigned char*)lpbi + sizeof(BITMAPINFOHEADER) + ((lpbi->biBitCount&7) ? 0 : 4);
}


DWORD CodecInst::CompressBegin(LPBITMAPINFOHEADER lpbiIn, LPBITMAPINFOHEADER lpbiOut) {
  Msg("CompressBegin: input = %s, output = %s\n", &PrintBitmapType(lpbiIn), &PrintBitmapType(lpbiOut));

  if (!CanCompress(lpbiIn, lpbiOut))
    return ICERR_BADFORMAT;

  CompressEnd();  // free resources if necessary

  int intype = GetBitmapType(lpbiIn);
  int outtype = GetBitmapType(lpbiOut);
  int method = GetMethod(lpbiOut);

  if (outtype < 0) {
    InitializeEncodeTables(GetEmbeddedHuffmanTable(lpbiOut));
    encode_table_owner = this;
  }

  // allocate buffer if compressing RGB->YUY2->HFYU
  if (outtype == -1 && (method == methodGrad || intype == 3))
    yuy2_buffer = new unsigned char[lpbiIn->biWidth * lpbiIn->biHeight * 2 + 4];
  if (method == methodMedian)
    median_buffer = new unsigned char[lpbiIn->biWidth * lpbiIn->biHeight * 2 + 8];
  if (outtype == -2 && method == methodGrad+flagDecorrelate)
    rgb_buffer = new unsigned char[lpbiIn->biWidth * lpbiIn->biHeight * 3 + 4];
  if (outtype == -3 && method == methodGrad+flagDecorrelate)
    rgb_buffer = new unsigned char[lpbiIn->biWidth * lpbiIn->biHeight * 4 + 4];

  InitClip();

  return ICERR_OK;
}

DWORD CodecInst::CompressGetSize(LPBITMAPINFOHEADER lpbiIn, LPBITMAPINFOHEADER lpbiOut) {
  // Assume 24 bpp worst-case for YUY2 input, 40 bpp for RGB.
  // The actual worst case is 43/51 bpp currently, but this is exceedingly improbable
  // (probably impossible with real captured video)
  return lpbiIn->biWidth * lpbiIn->biHeight * ((GetBitmapType(lpbiIn) <= 2) ? 3 : 5);
}


template<typename T> static inline T* Align8(T* ptr) { return (T*)((unsigned(ptr)+7)&~7); }


DWORD CodecInst::Compress(ICCOMPRESS* icinfo, DWORD dwSize) {

  if (icinfo->lpckid)
    *icinfo->lpckid = FOURCC_HFYU;
  *icinfo->lpdwFlags = AVIIF_KEYFRAME;

  int intype = GetBitmapType(icinfo->lpbiInput);
  int outtype = GetBitmapType(icinfo->lpbiOutput);

  if (outtype < 0) {    // compressing (as opposed to converting RGB->YUY2)
    const int method = GetMethod(icinfo->lpbiOutput);

    if (encode_table_owner != this) {
      InitializeEncodeTables(GetEmbeddedHuffmanTable(icinfo->lpbiOutput));
      encode_table_owner = this;
    }

    const unsigned char* const in = (unsigned char*)icinfo->lpInput;
    unsigned long* const out = (unsigned long*)icinfo->lpOutput;

    unsigned long* out_end;

    if (outtype == -1) {    // compressing to HFYU16 (4:2:2 YUV)
      int stride = icinfo->lpbiInput->biWidth * 2;
      int size = stride * icinfo->lpbiInput->biHeight;
      if (icinfo->lpbiInput->biHeight > 288) stride *= 2;    // if image is interlaced, double stride so fields are treated separately
      const unsigned char* yuy2in = in;

      if (intype == 3) {    // RGB24->YUY2->HFYU16
        ConvertRGB24toYUY2(in, yuy2_buffer, icinfo->lpbiInput->biWidth, icinfo->lpbiInput->biHeight);
        yuy2in = yuy2_buffer;
        size = icinfo->lpbiInput->biWidth * icinfo->lpbiInput->biHeight * 2;
      }
      if (method == methodGrad) {
        // Attempt to match yuy2_buffer alignment to input alignment.
        // Why, oh why, can't standard libraries just quadword-align
        // big memory blocks and be done with it?
        unsigned char* const aligned_yuy2_buffer = yuy2_buffer + 4*((int(yuy2in)&7) != (int(yuy2_buffer)&7));
        mmx_RowDiff(yuy2in, aligned_yuy2_buffer, yuy2in+size, stride);
        yuy2in = aligned_yuy2_buffer;
      }
      if (intype==2) {    // UYVY->HFYU16
        if (method == methodMedian) {
          unsigned char* aligned_median_buffer = Align8(median_buffer);
          mmx_MedianPredictUYVY(yuy2in, aligned_median_buffer, yuy2in+size, stride);
          out_end = asm_CompressUYVY(aligned_median_buffer, out, aligned_median_buffer+size);
        }
        else {
          out_end = asm_CompressUYVYDelta(yuy2in,out,yuy2in+size);
        }
      } else {    // YUY2->HFYU16
        if (method == methodMedian) {
          unsigned char* aligned_median_buffer = Align8(median_buffer);
          mmx_MedianPredictYUY2(yuy2in, aligned_median_buffer, yuy2in+size, stride);
          out_end = asm_CompressYUY2(aligned_median_buffer, out, aligned_median_buffer+size);
        }
        else {
          out_end = asm_CompressYUY2Delta(yuy2in,out,yuy2in+size);
        }
      }
    }
    else {    // compressing to HFYU24 (RGB) or HFYU32 (RGBA)
      int stride = (icinfo->lpbiInput->biWidth * GetBitCount(icinfo->lpbiInput)) >> 3;
      int size = stride * icinfo->lpbiInput->biHeight;
      if (icinfo->lpbiInput->biHeight > 288) stride *= 2;    // if image is interlaced, double stride so fields are treated separately
      const unsigned char* rgbin = in;

      if ((method&~flagDecorrelate) == methodGrad) {
        unsigned char* const aligned_rgb_buffer = rgb_buffer + 4*((int(in)&7) != (int(rgb_buffer)&7));
        mmx_RowDiff(in, aligned_rgb_buffer, in+size, stride);
        rgbin = aligned_rgb_buffer;
      }
      if (outtype == -2) {    // RGB24->HFYU24
        if (method == methodLeft || method == methodOld)
          out_end = asm_CompressRGBDelta(rgbin,out,rgbin+size);
        else
          out_end = asm_CompressRGBDeltaDecorrelate(rgbin,out,rgbin+size);
      }
      else if (outtype == -3) {     // RGBA->HFYU32
        if (method == methodLeft || method == methodOld)
          out_end = asm_CompressRGBADelta(rgbin,out,rgbin+size);
        else
          out_end = asm_CompressRGBADeltaDecorrelate(rgbin,out,rgbin+size);
      }
      else
        return ICERR_BADFORMAT;
    }
    icinfo->lpbiOutput->biSizeImage = DWORD(out_end) - DWORD(out);
    return ICERR_OK;
  }
  else if (outtype == 1) {    // RGB24->YUY2
    ConvertRGB24toYUY2((unsigned char*)icinfo->lpInput, (unsigned char*)icinfo->lpOutput,
      icinfo->lpbiInput->biWidth, icinfo->lpbiInput->biHeight);
    icinfo->lpbiOutput->biSizeImage = (icinfo->lpbiOutput->biWidth * icinfo->lpbiOutput->biHeight * GetBitCount(icinfo->lpbiOutput)) >> 3;
    return ICERR_OK;
  }
  else
    return ICERR_BADFORMAT;
}


void CodecInst::ConvertRGB24toYUY2(const unsigned char* src, unsigned char* dst, int width, int height) {
  const int cyb = int(0.114*219/255*65536+0.5);
  const int cyg = int(0.587*219/255*65536+0.5);
  const int cyr = int(0.299*219/255*65536+0.5);

  for (int row = 0; row < height; ++row) {
    const unsigned char* rgb = src + width * 3 * (height-1-row);
    unsigned char* yuv = dst + width * 2 * row;
    for (int col = 0; col < width; col += 2) {
      // y1 and y2 can't overflow
      int y1 = (cyb*rgb[0] + cyg*rgb[1] + cyr*rgb[2] + 0x108000) >> 16;
      yuv[0] = y1;
      int y2 = (cyb*rgb[3] + cyg*rgb[4] + cyr*rgb[5] + 0x108000) >> 16;
      yuv[2] = y2;