[置顶] HEVC学习（十四） —— SAO函数解析之二

/** rate distortion optimization of all SAO units
 * \param saoParam SAO parameters
 * \param lambda 
 * \param lambdaChroma
 */
#if SAO_ENCODING_CHOICE
Void TEncSampleAdaptiveOffset::rdoSaoUnitAll(SAOParam *saoParam, Double lambda, Double lambdaChroma, Int depth)
#else
Void TEncSampleAdaptiveOffset::rdoSaoUnitAll(SAOParam *saoParam, Double lambda, Double lambdaChroma)
#endif
{

  Int idxY;
  Int idxX;
  Int frameHeightInCU = saoParam->numCuInHeight;
  Int frameWidthInCU  = saoParam->numCuInWidth;
  Int j, k;
  Int addr = 0;
  Int addrUp = -1;
  Int addrLeft = -1;
  Int compIdx = 0;
  SaoLcuParam mergeSaoParam[3][2];
  Double compDistortion[3];

  saoParam->bSaoFlag[0] = true;
  saoParam->bSaoFlag[1] = true;
  saoParam->oneUnitFlag[0] = false;
  saoParam->oneUnitFlag[1] = false;
  saoParam->oneUnitFlag[2] = false;

#if SAO_ENCODING_CHOICE
#if SAO_ENCODING_CHOICE_CHROMA
  Int numNoSao[2];
  numNoSao[0] = 0;// Luma 
  numNoSao[1] = 0;// Chroma 
  if( depth > 0 && m_depthSaoRate[0][depth-1] > SAO_ENCODING_RATE )
  {
    saoParam->bSaoFlag[0] = false;
  }
  if( depth > 0 && m_depthSaoRate[1][depth-1] > SAO_ENCODING_RATE_CHROMA )
  {
    saoParam->bSaoFlag[1] = false;
  }
#else
  Int numNoSao = 0;

  if( depth > 0 && m_depth0SaoRate > SAO_ENCODING_RATE )
  {
    saoParam->bSaoFlag[0] = false;
    saoParam->bSaoFlag[1] = false;
  }
#endif
#endif
  //!< 以LCU为单位对图像中的每个LCU进行遍历
  for (idxY = 0; idxY< frameHeightInCU; idxY++)
  {
    for (idxX = 0; idxX< frameWidthInCU; idxX++)
    {
      addr     = idxX  + frameWidthInCU*idxY; //!< 当前LCU地址
      addrUp   = addr < frameWidthInCU ? -1:idxX   + frameWidthInCU*(idxY-1); //!< 当前LCU上邻块地址
      addrLeft = idxX == 0               ? -1:idxX-1 + frameWidthInCU*idxY; //!< 当前LCU左邻块地址
      Int allowMergeLeft = 1;
      Int allowMergeUp   = 1;
      UInt rate;
      Double bestCost, mergeCost;
      if (idxX!=0) //!< 非第1列
      { 
        // check tile id and slice id //! 检查当前LCU与其左邻块是否属于同一个tile以及是否属于同一个slice，不同的话该邻块不可用
        if ( (m_pcPic->getPicSym()->getTileIdxMap(addr-1) != m_pcPic->getPicSym()->getTileIdxMap(addr)) || (m_pcPic->getCU(addr-1)->getSlice()->getSliceIdx() != m_pcPic->getCU(addr)->getSlice()->getSliceIdx()))
        {
          allowMergeLeft = 0; //!< 左邻块不可用
        }
      }
      else
      {
        allowMergeLeft = 0; //!< 第1列的左邻块均不可用
      }
      if (idxY!=0) //!< 非第1行
      {//! 检查当前LCU与其上邻块是否属于同一个tile以及是否属于同一个slice，不同的话该邻块不可用
        if ( (m_pcPic->getPicSym()->getTileIdxMap(addr-m_iNumCuInWidth) != m_pcPic->getPicSym()->getTileIdxMap(addr)) || (m_pcPic->getCU(addr-m_iNumCuInWidth)->getSlice()->getSliceIdx() != m_pcPic->getCU(addr)->getSlice()->getSliceIdx()))
        {
          allowMergeUp = 0; //!< 上邻块不可用
        }
      }
      else
      {
        allowMergeUp = 0; //!< 第1行的上邻块均不可用
      }

      compDistortion[0] = 0; //!< Y distortion
      compDistortion[1] = 0; //!< Cb distortion
      compDistortion[2] = 0; //!< Cr distortion
      m_pcRDGoOnSbacCoder->load(m_pppcRDSbacCoder[0][CI_CURR_BEST]);
      if (allowMergeLeft)
      {
        m_pcEntropyCoder->m_pcEntropyCoderIf->codeSaoMerge(0); //!< 编码句法元素sao_merge_left_flag
      }
      if (allowMergeUp)
      {
        m_pcEntropyCoder->m_pcEntropyCoderIf->codeSaoMerge(0); //!< 编码句法元素sao_merge_up_flag
      }
      m_pcRDGoOnSbacCoder->store( m_pppcRDSbacCoder[0][CI_TEMP_BEST] );
      // reset stats Y, Cb, Cr
      for ( compIdx=0;compIdx<3;compIdx++)
      {
        for ( j=0;j<MAX_NUM_SAO_TYPE;j++)
        {
          for ( k=0;k< MAX_NUM_SAO_CLASS;k++)
          {
            m_iOffset   [compIdx][j][k] = 0;
            if( m_saoLcuBasedOptimization && m_saoLcuBoundary ){ //!< true && false
              m_iCount    [compIdx][j][k] = m_count_PreDblk    [addr][compIdx][j][k];
              m_iOffsetOrg[compIdx][j][k] = m_offsetOrg_PreDblk[addr][compIdx][j][k];
            }
            else
            {
              m_iCount    [compIdx][j][k] = 0;
              m_iOffsetOrg[compIdx][j][k] = 0;
            }
          }  
        }
        saoParam->saoLcuParam[compIdx][addr].typeIdx       =  -1;
        saoParam->saoLcuParam[compIdx][addr].mergeUpFlag   = 0;
        saoParam->saoLcuParam[compIdx][addr].mergeLeftFlag = 0;
        saoParam->saoLcuParam[compIdx][addr].subTypeIdx    = 0;
#if SAO_ENCODING_CHOICE
  if( (compIdx ==0 && saoParam->bSaoFlag[0])|| (compIdx >0 && saoParam->bSaoFlag[1]) )
#endif
        {//! 统计BO和EO各个模式下，对应classIdx下滤波前的重建像素值与原始像素值的差值的总和，以及对classIdx的计数
          calcSaoStatsCu(addr, compIdx,  compIdx);
		}
      }
	  //!< Y分量最佳滤波模式的选择
      saoComponentParamDist(allowMergeLeft, allowMergeUp, saoParam, addr, addrUp, addrLeft, 0,  lambda, &mergeSaoParam[0][0], &compDistortion[0]);
	  //!< CbCr分量最佳滤波模式的选择
      sao2ChromaParamDist(allowMergeLeft, allowMergeUp, saoParam, addr, addrUp, addrLeft, lambdaChroma, &mergeSaoParam[1][0], &mergeSaoParam[2][0], &compDistortion[0]);
     if( saoParam->bSaoFlag[0] || saoParam->bSaoFlag[1] )
      {
        // Cost of new SAO_params
        m_pcRDGoOnSbacCoder->load(m_pppcRDSbacCoder[0][CI_CURR_BEST]);
        m_pcRDGoOnSbacCoder->resetBits();
        if (allowMergeLeft)
        {
          m_pcEntropyCoder->m_pcEntropyCoderIf->codeSaoMerge(0); 
        }
        if (allowMergeUp)
        {
          m_pcEntropyCoder->m_pcEntropyCoderIf->codeSaoMerge(0);
        }
        for ( compIdx=0;compIdx<3;compIdx++)
        {
        if( (compIdx ==0 && saoParam->bSaoFlag[0]) || (compIdx >0 && saoParam->bSaoFlag[1]))
          {
           m_pcEntropyCoder->encodeSaoOffset(&saoParam->saoLcuParam[compIdx][addr], compIdx);
          }
        }

        rate = m_pcEntropyCoder->getNumberOfWrittenBits();
        bestCost = compDistortion[0] + (Double)rate;
        m_pcRDGoOnSbacCoder->store(m_pppcRDSbacCoder[0][CI_TEMP_BEST]);

        // Cost of Merge
        for(Int mergeUp=0; mergeUp<2; ++mergeUp)
        {
          if ( (allowMergeLeft && (mergeUp==0)) || (allowMergeUp && (mergeUp==1)) )
          {
            m_pcRDGoOnSbacCoder->load(m_pppcRDSbacCoder[0][CI_CURR_BEST]);
            m_pcRDGoOnSbacCoder->resetBits();
            if (allowMergeLeft)
            {
              m_pcEntropyCoder->m_pcEntropyCoderIf->codeSaoMerge(1-mergeUp); 
            }
            if ( allowMergeUp && (mergeUp==1) )
            {
              m_pcEntropyCoder->m_pcEntropyCoderIf->codeSaoMerge(1); 
            }

            rate = m_pcEntropyCoder->getNumberOfWrittenBits();
            mergeCost = compDistortion[mergeUp+1] + (Double)rate;
            if (mergeCost < bestCost)
            {
              bestCost = mergeCost;
              m_pcRDGoOnSbacCoder->store(m_pppcRDSbacCoder[0][CI_TEMP_BEST]);              
              for ( compIdx=0;compIdx<3;compIdx++)
              {
                mergeSaoParam[compIdx][mergeUp].mergeLeftFlag = 1-mergeUp;
                mergeSaoParam[compIdx][mergeUp].mergeUpFlag = mergeUp;
                if( (compIdx==0 && saoParam->bSaoFlag[0]) || (compIdx>0 && saoParam->bSaoFlag[1]))
                {
                  copySaoUnit(&saoParam->saoLcuParam[compIdx][addr], &mergeSaoParam[compIdx][mergeUp] );             
                }
              }
            }
          }
        }
#if SAO_ENCODING_CHOICE
#if SAO_ENCODING_CHOICE_CHROMA
if( saoParam->saoLcuParam[0][addr].typeIdx == -1) //!< Y分量不存在SAO参数
{
  numNoSao[0]++;
}
if( saoParam->saoLcuParam[1][addr].typeIdx == -1) //!< CbCr分量不存在SAO参数
{
  numNoSao[1]+=2;
}
#else
        for ( compIdx=0;compIdx<3;compIdx++)
        {
          if( depth == 0 && saoParam->saoLcuParam[compIdx][addr].typeIdx == -1)
          {
            numNoSao++;
          }
        }
#endif
#endif
        m_pcRDGoOnSbacCoder->load(m_pppcRDSbacCoder[0][CI_TEMP_BEST]);
        m_pcRDGoOnSbacCoder->store(m_pppcRDSbacCoder[0][CI_CURR_BEST]);
      } //!< if( saoParam->bSaoFlag[0] || saoParam->bSaoFlag[1] )       
    } //!< for (idxX = 0; idxX< frameWidthInCU; idxX++)     
  } //!< for (idxY = 0; idxY< frameHeightInCU; idxY++)   
#if SAO_ENCODING_CHOICE
#if SAO_ENCODING_CHOICE_CHROMA
#if SAO_ENCODING_CHOICE_CHROMA_BF
  if( !saoParam->bSaoFlag[0]) 
  {
    m_depthSaoRate[0][depth] = 1.0;
  }
  else
  {
    m_depthSaoRate[0][depth] = numNoSao[0]/((Double) frameHeightInCU*frameWidthInCU);
  }
  if( !saoParam->bSaoFlag[1]) 
  {
    m_depthSaoRate[1][depth] = 1.0;
  }
  else 
  {
    m_depthSaoRate[1][depth] = numNoSao[1]/((Double) frameHeightInCU*frameWidthInCU*2);
  }
#else
m_depthSaoRate[0][depth] = numNoSao[0]/((Double) frameHeightInCU*frameWidthInCU);
m_depthSaoRate[1][depth] = numNoSao[1]/((Double) frameHeightInCU*frameWidthInCU*2);
#endif
#else
  if( depth == 0)
  {
    // update SAO Rate
    m_depth0SaoRate = numNoSao/((Double) frameHeightInCU*frameWidthInCU*3);
  }
#endif
#endif

}