H.266/VVC代码学习笔记12：VTM5.0中的estIntraPredChromaQT（）函数的学习

对于学习视频编码中帧内预测部分的同学来说，estIntraPredChromaQT（）函数可以说是必须掌握的一个很重要的函数，该函数是帧内色度预测的第一主函数入口，和帧内亮度预测的第一主入口函数estIntraPredLumaQT（）函数相对应，关于亮度预测的入口函数的代码细节在我之前的博客已经详细讲过啦，这里就直接给出链接：H.266/VVC代码学习笔记8：VTM5.0中帧内亮度预测函数——estIntraPredLumaQT()函数，都非常重要。两个函数的作用非常类似，就是在各自颜色分量的所有预测候选模式中通过RDcost选择出最优的预测模式，不过相对于亮度分量来说，色度预测模式的RDcost过程要简单很多，下面我详细讲一下在目前的VTM5.0版本中该函数的具体实现细节。

该函数的主要流程如下：
一、首先进行一些初始化的工作，定义一些变量，检查一些可能发生的异常；

二、初始化候选模式列表，用getIntraChromaCandModes（）函数对色度候选模式列表进行填充，将八种色度模式填充进去

三、判断当前的CU是否满足继续划分的条件，如果满足则跳入splitCurrArea（）函数继续进行树划分。然后为划分好的TU分配储存空间，用于接下来模式筛选做准备。

四、开始第一轮STAD的预选，对LM、Planar、DM模式不进行SATD，只对传统角度模式和MDLM进行SATD预选， satdModeList[]中存放每个进行SATD的色度预测模式，satdSortedCost[]中存放每种进行SATD的模式最后的SATD代价，该过程具体的流程如下：
   ①首选初始化用于SATD预选择的列表，再确定可用的模式列表，初始化帧内模式通道（Cb分量和Cr分量）
   ② 调用xGetLumaRecPixels()函数为MDLM模式做准备，进行亮度下采样操作
   ③对每个需要进行SATD预选的模式进行循环遍历，计算每种模式的SATD代价，并放入Cost列表satdSortedCost中去
   ④然后从小到大排序每种进行SATD的模式的代价，更新satdSortedCost列表
   ⑤最后淘汰掉SATD较大的两种模式

五、SATD预选结束后，再开始计算RDcost的细选操作，对SATD之后剩余的那些模式循环遍历进行进一步的RDCost细选，细选的具体步骤如下：
   ①若当前预测块CCLM模式禁用，则跳过进行下一个模式的计算；若MDLM模式也不可用，则跳过进行下一个模式的计算。
   ②编码各种数据之前，重置上下文模型，然后进入帧内色度预测模式的第二入口函数： xRecurIntraChromaCodingQT（），这个函数的代码细节我在之前的博客已经讲过啦，链接如下：H.266/VVC代码学习笔记11：VTM5.0中的xRecurIntraChromaCodingQT（）函数。
   ③然后得到当前遍历到的模式的码率和失真，计算对应的RDcost，如果当前模式的RDcost比历史最优的RDcost还小，则当前的模式代替历史最优并且保存RDcost最优的模式。

六、最后对CbCr分量循环遍历，保存各自预测块最优的预测值，残差值，模式号等预测信息，得到最终得到最优的预测模式uiBestMode，并且赋给 pu.intraDir[1]中

以上就是色度RDcost的主要流程，关于具体的代码细节如下，代码中有我详细的注释，有兴趣的同学可以跟进去详细学习一下，有不懂或者不太清楚的地方可以私信我或者给我留言

void IntraSearch::estIntraPredChromaQT( CodingUnit &cu, Partitioner &partitioner, const double maxCostAllowed )
{
  const ChromaFormat format   = cu.chromaFormat;
  const uint32_t    numberValidComponents = getNumberValidComponents(format);
  CodingStructure &cs = *cu.cs;
  const TempCtx ctxStart  ( m_CtxCache, m_CABACEstimator->getCtx() );

  cs.setDecomp( cs.area.Cb(), false );

  double    bestCostSoFar = maxCostAllowed;
  bool      lumaUsesISP   = !CS::isDualITree( *cu.cs ) && cu.ispMode;
  PartSplit ispType       = lumaUsesISP ? CU::getISPType( cu, COMPONENT_Y ) : TU_NO_ISP;
  CHECK( cu.ispMode && bestCostSoFar < 0, "bestCostSoFar must be positive!" );

  auto &pu = *cu.firstPU;
  //对所有的色度候选模式进行循环遍历，然后每一次遍历都计算相应的残差的变化量化，并且再反量化反变换计算重建值，用于RDcost的计算，
  {
    uint32_t       uiBestMode = 0;
    Distortion uiBestDist = 0;
    double     dBestCost = MAX_DOUBLE;

    //----- init mode list ----
    //初始化候选模式列表
    {
      uint32_t  uiMinMode = 0;
      uint32_t  uiMaxMode = NUM_CHROMA_MODE;//最大的色度预测候选模式，目前总共为8个

      //----- 检查色度候选模式 -----
      uint32_t chromaCandModes[ NUM_CHROMA_MODE ];//色度预测模式的候选列表
      PU::getIntraChromaCandModes( pu, chromaCandModes );//这里对chromaCandModes列表进行填充，将八种色度模式填充进去

      // create a temporary CS
      //构造当前CS的临时缓存BUFF
      CodingStructure &saveCS = *m_pSaveCS[0];
      saveCS.pcv      = cs.pcv;
      saveCS.picture  = cs.picture;
      saveCS.area.repositionTo( cs.area );
      saveCS.clearTUs();

      if( !CS::isDualITree( cs ) && cu.ispMode )//如果当前的CS亮度色度不单独划分且当前CU不使用ISP模式
      {
        saveCS.clearCUs();
        saveCS.clearPUs();
      }

      if( CS::isDualITree( cs ) )//如果当前CS色度单独划分
      {
        if( partitioner.canSplit( TU_MAX_TR_SPLIT, cs ) )
        {
          partitioner.splitCurrArea( TU_MAX_TR_SPLIT, cs );//满足划分条件时候继续对当前CU进行划分

          do
          {
            cs.addTU( CS::getArea( cs, partitioner.currArea(), partitioner.chType ), partitioner.chType ).depth = partitioner.currTrDepth;
          } while( partitioner.nextPart( cs ) );

          partitioner.exitCurrSplit();
        }
        else
        cs.addTU( CS::getArea( cs, partitioner.currArea(), partitioner.chType ), partitioner.chType );
      }

      std::vector<TransformUnit*> orgTUs;

      if( lumaUsesISP )//亮度块使用ISP模式
      {
        CodingUnit& auxCU = saveCS.addCU( cu, partitioner.chType );
        auxCU.ispMode = cu.ispMode;
        saveCS.sps = cu.cs->sps;
        saveCS.addPU( *cu.firstPU, partitioner.chType );
      }


      // create a store for the TUs
      //为划分好的TU分配储存空间
      for( const auto &ptu : cs.tus )
      {
        // for split TUs in HEVC, add the TUs without Chroma parts for correct setting of Cbfs
        if( lumaUsesISP || pu.contains( *ptu, CHANNEL_TYPE_CHROMA ) )
        {
          saveCS.addTU( *ptu, partitioner.chType );
          orgTUs.push_back( ptu );
        }
      }
      if( lumaUsesISP )
      {
        saveCS.clearCUs();
      }
      // SATD pre-selecting.
      //第一轮STAD的预选，对LM、Planar、DM模式不进行SATD，只对传统角度模式和MDLM进行SATD预选
      int satdModeList[NUM_CHROMA_MODE];//定义要进行SATD粗选的模式列表，里面存放每个进行SATD的色度预测模式
      int64_t satdSortedCost[NUM_CHROMA_MODE];//定义SATD代价列表，里面存放每种进行SATD的模式最后的SATD代价

      //这里初始化用于SATD预选择的列表
      for (int i = 0; i < NUM_CHROMA_MODE; i++)
      {
        satdSortedCost[i] = 0; // for the mode not pre-select by SATD, do RDO by default, so set the initial value 0.
        satdModeList[i] = 0;
      }

      bool modeIsEnable[NUM_INTRA_MODE + 1]; // use intra mode idx to check whether enable
      //确定可用的模式列表
      for (int i = 0; i < NUM_INTRA_MODE + 1; i++)
      {
        modeIsEnable[i] = 1;
      }

      DistParam distParam;
#if JVET_N0329_IBC_SEARCH_IMP
      const bool useHadamard = !cu.transQuantBypass;
#else
      const bool useHadamard = true;
#endif
      pu.intraDir[1] = MDLM_L_IDX; // temporary assigned, just to indicate this is a MDLM mode. for luma down-sampling operation.
      //临时分配，只是为了表明这是一个MDLM模式。 用于亮度下采样操作。

      initIntraPatternChType(cu, pu.Cb());//初始化帧内模式通道，即色度通道类型，分配PU区域
      initIntraPatternChType(cu, pu.Cr());

      //为MDLM模式做准备，进行亮度下采样操作
      xGetLumaRecPixels(pu, pu.Cb());


      //这里对每个需要进行SATD预选的模式进行循环遍历，计算每种模式的SATD代价，并放入Cost列表
      for (int idx = uiMinMode; idx <= uiMaxMode - 1; idx++)
      {
        int mode = chromaCandModes[idx];
        satdModeList[idx] = mode;
        //对LM模式不进行SATD
        if (PU::isLMCMode(mode) && !PU::isLMCModeEnabled(pu, mode))
        {
          continue;
        }
        //对LM、Planar、DM模式不进行SATD，只对角度模式和MDLM进行SATD粗选
        if ((mode == LM_CHROMA_IDX) || (mode == PLANAR_IDX) || (mode == DM_CHROMA_IDX)) // only pre-check regular modes and MDLM modes, not including DM ,Planar, and LM
        {
          continue;
        }
        pu.intraDir[1] = mode; // 临时分配，用于SATD检查.
        //选中的SATD候选模式放入帧内最终选中的色度预测模式intraDir中

        int64_t sad = 0;
        CodingStructure& cs = *(pu.cs);//定义编码单元

        CompArea areaCb = pu.Cb();//定义Cb块的编码区域
        PelBuf orgCb = cs.getOrgBuf(areaCb);//获取Cb块的原始YUV值
        PelBuf predCb = cs.getPredBuf(areaCb);//获取Cb块的预测YUV值

        m_pcRdCost->setDistParam(distParam, orgCb, predCb, pu.cs->sps->getBitDepth(CHANNEL_TYPE_CHROMA), COMPONENT_Cb, useHadamard);
        distParam.applyWeight = false;
        
        if (PU::isLMCMode(mode))
        {
          //Cb块的CCLM预测
          predIntraChromaLM(COMPONENT_Cb, predCb, pu, areaCb, mode);
        }
        else
        {
          //Cb块的角度预测
          initPredIntraParams(pu, pu.Cb(), *pu.cs->sps);
          predIntraAng(COMPONENT_Cb, predCb, pu);
        }

        sad += distParam.distFunc(distParam);

        CompArea areaCr = pu.Cr();//定义Cr块的编码区域
        PelBuf orgCr = cs.getOrgBuf(areaCr);//获取Cr块的原始YUV值
        PelBuf predCr = cs.getPredBuf(areaCr);//获取Cr块的预测YUV值

        m_pcRdCost->setDistParam(distParam, orgCr, predCr, pu.cs->sps->getBitDepth(CHANNEL_TYPE_CHROMA), COMPONENT_Cr, useHadamard);
        distParam.applyWeight = false;

        if (PU::isLMCMode(mode))
        {
          //Cr块的CCLM预测
          predIntraChromaLM(COMPONENT_Cr, predCr, pu, areaCr, mode);
        }
        else
        {
          //Cr块的角度预测
          initPredIntraParams(pu, pu.Cr(), *pu.cs->sps);
          predIntraAng(COMPONENT_Cr, predCr, pu);
        }
        sad += distParam.distFunc(distParam);//CbCr分量预测结束的时候，计算总的SATD代价
        satdSortedCost[idx] = sad;//将每种模式的SATD代价放入satdSortedCost列表
      }
      //sort the mode based on the cost from small to large.
      //从小到大排序每种进行SATD的模式的代价，更新satdSortedCost列表
      int tempIdx = 0;
      int64_t tempCost = 0;
      for (int i = uiMinMode; i <= uiMaxMode - 1; i++)
      {
        for (int j = i + 1; j <= uiMaxMode - 1; j++)
        {
          if (satdSortedCost[j] < satdSortedCost[i])
          {
            tempIdx = satdModeList[i];
            satdModeList[i] = satdModeList[j];
            satdModeList[j] = tempIdx;

            tempCost = satdSortedCost[i];
            satdSortedCost[i] = satdSortedCost[j];
            satdSortedCost[j] = tempCost;

          }
        }
      }
      //SATD后要被淘汰的模式数量为2。即去掉SATD较大的两种模式
      int reducedModeNumber = 2; // reduce the number of chroma modes
      for (int i = 0; i < reducedModeNumber; i++)
      {
        modeIsEnable[satdModeList[uiMaxMode - 1 - i]] = 0; // disable the last reducedModeNumber modes
        //去掉satdSortedCost中的后两个模式
      }

      //保存失真
      Distortion baseDist = cs.dist;//为RDcost定义最优的失真


      //这里开始计算RDcost的细选操作，对SATD之后剩余的那些模式循环遍历进行进一步的RDCost细选
      for (uint32_t uiMode = uiMinMode; uiMode < uiMaxMode; uiMode++)
      {
        const int chromaIntraMode = chromaCandModes[uiMode];//从色度预测模式的候选列表取出对应的模式赋给chromaIntraMode
        //若当前预测块CCLM模式禁用，则跳过进行下一个模式的计算
        if( PU::isLMCMode( chromaIntraMode ) && ! PU::isLMCModeEnabled( pu, chromaIntraMode ) )
        {
          continue;
        }
        //若当前预测块CCLM模式不可用，且MDLM模式也不可用，则跳过进行下一个模式的计算
        if (!modeIsEnable[chromaIntraMode] && PU::isLMCModeEnabled(pu, chromaIntraMode)) // when CCLM is disable, then MDLM is disable. not use satd checking
        {
          continue;
        }
        cs.setDecomp( pu.Cb(), false );
        cs.dist = baseDist;
        //----- restore context models -----
        //编码各种数据之前，重置上下文模型
        m_CABACEstimator->getCtx() = ctxStart;

        //----- chroma coding -----
        pu.intraDir[1] = chromaIntraMode;

        //帧内色度预测模式的函数入口
        xRecurIntraChromaCodingQT( cs, partitioner, bestCostSoFar, ispType );

        if( lumaUsesISP && cs.dist == MAX_UINT )
        {
          continue;
        }

        if (cs.pps->getUseTransformSkip())
        {
          m_CABACEstimator->getCtx() = ctxStart;
        }

        uint64_t fracBits   = xGetIntraFracBitsQT( cs, partitioner, false, true, -1, ispType );//每种预测模式预测后的码率
        Distortion uiDist = cs.dist;//每种预测模式预测后的失真
        double    dCost   = m_pcRdCost->calcRdCost( fracBits, uiDist - baseDist );//计算最终的RDcost

        //----- compare -----
        //循环比较每种模式的RDcost，保存RDcost最小的模式
        if( dCost < dBestCost )//如果当前模式的RDcost比历史最优的RDcost还小，则当前的模式代替历史最优
        {
          if( lumaUsesISP && dCost < bestCostSoFar )
          {
            bestCostSoFar = dCost;
          }
          for( uint32_t i = getFirstComponentOfChannel( CHANNEL_TYPE_CHROMA ); i < numberValidComponents; i++ )
          {
            const CompArea &area = pu.blocks[i];

            saveCS.getRecoBuf     ( area ).copyFrom( cs.getRecoBuf   ( area ) );
#if KEEP_PRED_AND_RESI_SIGNALS
            saveCS.getPredBuf     ( area ).copyFrom( cs.getPredBuf   ( area ) );
            saveCS.getResiBuf     ( area ).copyFrom( cs.getResiBuf   ( area ) );
#endif
            saveCS.getPredBuf     ( area ).copyFrom( cs.getPredBuf   (area ) );
            cs.picture->getPredBuf( area ).copyFrom( cs.getPredBuf   (area ) );
            cs.picture->getRecoBuf( area ).copyFrom( cs.getRecoBuf( area ) );

            for( uint32_t j = 0; j < saveCS.tus.size(); j++ )
            {
              saveCS.tus[j]->copyComponentFrom( *orgTUs[j], area.compID );
            }
          }

          dBestCost  = dCost;//最优的代价
          uiBestDist = uiDist;//最优的失真
          uiBestMode = chromaIntraMode;//保存RDcost最优的模式
        }
      }
      //对CbCr分量循环遍历，保存各自预测块最优的预测值，残差值，模式号等预测信息
      for( uint32_t i = getFirstComponentOfChannel( CHANNEL_TYPE_CHROMA ); i < numberValidComponents; i++ )
      {
        const CompArea &area = pu.blocks[i];

        cs.getRecoBuf         ( area ).copyFrom( saveCS.getRecoBuf( area ) );
#if KEEP_PRED_AND_RESI_SIGNALS
        cs.getPredBuf         ( area ).copyFrom( saveCS.getPredBuf( area ) );
        cs.getResiBuf         ( area ).copyFrom( saveCS.getResiBuf( area ) );
#endif
        cs.getPredBuf         ( area ).copyFrom( saveCS.getPredBuf( area ) );
        cs.picture->getPredBuf( area ).copyFrom( cs.getPredBuf    ( area ) );

        cs.picture->getRecoBuf( area ).copyFrom( cs.    getRecoBuf( area ) );

        for( uint32_t j = 0; j < saveCS.tus.size(); j++ )
        {
          orgTUs[ j ]->copyComponentFrom( *saveCS.tus[ j ], area.compID );
        }
      }
    }

    pu.intraDir[1] = uiBestMode;//最终得到最优的预测模式
    cs.dist        = uiBestDist;
  }

  //----- restore context models -----
  //更新上下文模型
  m_CABACEstimator->getCtx() = ctxStart;
  if( lumaUsesISP && bestCostSoFar >= maxCostAllowed )
  {
    cu.ispMode = 0;
  }
}