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;
  }
}