H.266/VVC代码学习：色度预测estIntraPredChromaQT函数

VVC中色度预测一共有八种模式，分别为PLANAR，VER， HOR，DC，LM_CHROMA，MDLM_L，MDLM_T，DM_CHROMA，其中LM_CHROMA，MDLM_L和MDLM_T分别表示CCLM模式，DM_CHROMA为从亮度分量获取到的模式。

VTM中色度预测最佳模式的选择是通过estIntraPredChromaQT函数进行的，其主要流程如下：

1. 初始化色度模式列表，利用getIntraChromaCandModes函数获取色度模式列表，并初始化PU、TU
2. 利用SATD对DC、Ver、Hor、LM_L、LM_T五种模式进行粗选，禁用色度模式列表中SATD最大的两个模式
3. 利用RD Cost对剩余的六种模式进行细选，主要是通过xRecurIntraChromaCodingQT进行变换、量化后计算RD Cost，选出最低的RD Cost
4. 保存最佳色度模式。

1、初始化色度模式列表

在estIntraPredChromaQT函数中执行getIntraChromaCandModes函数进行色度模式列表初始化的工作，将当前列表初始化为PLANAR，VER， HOR，DC，LM_CHROMA，MDLM_L，MDLM_T，DM_CHROMA。

其中DM（derived mode）模式是一种色度预测特有的模式。在进行色度预测的时候，亮度预测已经完成，此时当前块的色度预测模式可以参考已经完成的亮度预测模式。对于亮度块和色度块相同划分方式，则直接使用亮度块的预测模式；对于亮度块和色度块不同的划分方式时，使用色度块对应的亮度块的中心位置处的亮度预测模式作为当前色度块的预测模式。

注意，此时存在一个问题，即如果色度块对应亮度块的预测模式是色度模式列表前四种预测模式之一（PLANAR，VER， HOR，DC），那么后续的RD模式选择就会重复进行一次选择。为了解决该问题，色度列表会进行微调，将四种传统模式与DM预测出的模式冲突的位置换成66，用以占位。获取色度模式列表不单单在预测过程中会启用，在CABAC编解码中也会启用，用于编码色度信息。色度列表如下表，其中横行的亮度指的是色度块对应亮度块中间块的预测方向，第一纵行表示色度模式列表的索引，中间列表每个纵行都是对应其亮度方向的一个色度模式列表。

Chroma prediction mode index	Corresponding luma intra prediction mode
	0	50	18	1	X ( 0  <=  X  <=  66 )
0	66	0	0	0	0
1	50	66	50	50	50
2	18	18	66	18	18
3	1	1	1	66	1
4	0	50	18	1	X
5	81	81	81	81	81
6	82	82	82	82	82
7	83	83	83	83	83

getIntraChromaCandModes函数如下所示。

void PU::getIntraChromaCandModes( const PredictionUnit &pu, unsigned modeList[NUM_CHROMA_MODE] )
{
  {
    modeList[  0 ] = PLANAR_IDX;
    modeList[  1 ] = VER_IDX;
    modeList[  2 ] = HOR_IDX;
    modeList[  3 ] = DC_IDX;
    modeList[4] = LM_CHROMA_IDX;
    modeList[5] = MDLM_L_IDX;
    modeList[6] = MDLM_T_IDX;
    modeList[7] = DM_CHROMA_IDX;

#if JVET_R0350_MIP_CHROMA_444_SINGLETREE
    // If Direct Mode is MIP, mode cannot be already in the list.
    // 如果亮度和色度相同划分方式且色度对应亮度块是MIP模式且是4:4:4视频格式，不需要对列表进行微调
    if (isDMChromaMIP(pu))
    {
      return;
    }
#endif
    const uint32_t lumaMode = getCoLocatedIntraLumaMode(pu);//获取对应亮度块的预测模式
    for( int i = 0; i < 4; i++ )
    {
      if( lumaMode == modeList[i] )//当DM预测的模式与四种传统的色度预测模式相同时
      {
        modeList[i] = VDIA_IDX;//使用斜对角，即模式66替代传统的色度预测模式，起到占位作用，选中几率很小
        break;
      }
    }
  }
}

uint32_t PU::getCoLocatedIntraLumaMode(const PredictionUnit &pu)
{
  return PU::getIntraDirLuma(PU::getCoLocatedLumaPU(pu));
}

uint32_t PU::getIntraDirLuma( const PredictionUnit &pu )
{
  if (isMIP(pu))
  {
    return PLANAR_IDX;
  }
  else
  {
    return pu.intraDir[CHANNEL_TYPE_LUMA];
  }
}

#if JVET_R0350_MIP_CHROMA_444_SINGLETREE
const PredictionUnit &PU::getCoLocatedLumaPU(const PredictionUnit &pu)
{
  Position              topLeftPos = pu.blocks[pu.chType].lumaPos();//亮度块的左上角
  Position              refPos     = topLeftPos.offset(pu.blocks[pu.chType].lumaSize().width  >> 1,
                                                       pu.blocks[pu.chType].lumaSize().height >> 1);//亮度块的中心位置
  const PredictionUnit &lumaPU     = pu.cu->isSepTree() ? *pu.cs->picture->cs->getPU(refPos, CHANNEL_TYPE_LUMA)//亮度和色度单独划分时，获取亮度块中间的预测模式
                                                        : *pu.cs->getPU(topLeftPos, CHANNEL_TYPE_LUMA);//亮度和色度相同划分模式，获取亮度块左上角的预测模式

  return lumaPU;
}

 2、进行RD 模式选择

代码和注释如下：

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   = !cu.isSepTree() && cu.ispMode;//亮度使用ISP标志
  PartSplit ispType       = lumaUsesISP ? CU::getISPType( cu, COMPONENT_Y ) : TU_NO_ISP;//ISP类型
  CHECK( cu.ispMode && bestCostSoFar < 0, "bestCostSoFar must be positive!" );

  auto &pu = *cu.firstPU;
  
  {//代码块
    uint32_t       uiBestMode = 0;//最佳模式
    Distortion uiBestDist = 0;//最佳失真
    double     dBestCost = MAX_DOUBLE;//最佳Cost
    int32_t bestBDPCMMode = 0;

    //----- init mode list ----
    {//初始化模式列表（列出空表）
      int32_t  uiMinMode = 0;
      int32_t  uiMaxMode = NUM_CHROMA_MODE;
      //----- check chroma modes -----
      //获取色度模式（填满表）
      uint32_t chromaCandModes[ NUM_CHROMA_MODE ];
      PU::getIntraChromaCandModes( pu, chromaCandModes );

      // create a temporary CS
      // 创建一个临时CS
      CodingStructure &saveCS = *m_pSaveCS[0];
      saveCS.pcv      = cs.pcv;
      saveCS.picture  = cs.picture;
      saveCS.area.repositionTo( cs.area );
      saveCS.clearTUs();

      if( !cu.isSepTree() && cu.ispMode )
      {
        saveCS.clearCUs();
        saveCS.clearPUs();
      }
      //如果是色度单独划分的话，将色度进行划分，并把TU存好
      if( cu.isSepTree() )
      {
        if( partitioner.canSplit( TU_MAX_TR_SPLIT, cs ) )
        {
          partitioner.splitCurrArea( TU_MAX_TR_SPLIT, cs );

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

      if( lumaUsesISP )// 如果亮度是ISP，记录当前cu，pu信息
      {
        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进orgTUs中以备后续处理
      for( const auto &ptu : cs.tus )
      {
        // for split TUs in HEVC, add the TUs without Chroma parts for correct setting of Cbfs
        // 对于HEVC中的拆分TU，请添加不包含色度部分的TU，以正确设置Cbfs
        if( lumaUsesISP || pu.contains( *ptu, CHANNEL_TYPE_CHROMA ) )
        {
          saveCS.addTU( *ptu, partitioner.chType );
          orgTUs.push_back( ptu );
        }
      }
      if( lumaUsesISP )
      {
        saveCS.clearCUs();
      }
      /*************进行第一次选择，按DC、Ver、Hor、LM_L、LM_T进行预测，按SATD排序，去掉两个SATD最大的模式**************/
      // SATD pre-selecting.SATD预选
      int satdModeList[NUM_CHROMA_MODE];//色度基于SATD的模式列表
      int64_t satdSortedCost[NUM_CHROMA_MODE];//保存SATD对应Cost
      for (int i = 0; i < NUM_CHROMA_MODE; i++)
      {
        //对于不是由SATD预先选择的模式，默认情况下执行RDO，因此设置初始值0。
        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使用帧内模式idx检查是否启用
      for (int i = 0; i < NUM_INTRA_MODE + 1; i++)
      {
        modeIsEnable[i] = 1;
      }
      DistParam distParamSad;
      DistParam distParamSatd;
      // 这里将pu模式设置为MDLM模式，为了在下面对luma重建信号进行下采样
      pu.intraDir[1] = MDLM_L_IDX; // temporary assigned, just to indicate this is a MDLM mode. for luma down-sampling operation.

      initIntraPatternChType(cu, pu.Cb());// 获取Cb分量参考像素
      initIntraPatternChType(cu, pu.Cr());// 获取Cr分量参考像素
      xGetLumaRecPixels(pu, pu.Cb());// CCLM预测模式亮度重建参考像素获取
      
      /***********进行第一次选择，利用SATD进行粗选DC VER HOR LM_T LM_L五种模式*******************/
      for (int idx = uiMinMode; idx <= uiMaxMode - 1; idx++)//遍历所有模式
      {
        int mode = chromaCandModes[idx];
        satdModeList[idx] = mode;
        if (PU::isLMCMode(mode) && !PU::isLMCModeEnabled(pu, mode))
        {
          continue;
        }
        //仅预检查常规模式和MDLM模式，不包括DM、Planar和LM
        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; // temporary assigned, for SATD checking.

        int64_t sad = 0;
        int64_t sadCb = 0;
        int64_t satdCb = 0;
        int64_t sadCr = 0;
        int64_t satdCr = 0;
        CodingStructure& cs = *(pu.cs);
        //计算Cb分量的SAD和SATD
        CompArea areaCb = pu.Cb();
        PelBuf orgCb = cs.getOrgBuf(areaCb);//获取Cb原始值
        PelBuf predCb = cs.getPredBuf(areaCb);//获取Cb预测值
        m_pcRdCost->setDistParam(distParamSad, orgCb, predCb, pu.cs->sps->getBitDepth(CHANNEL_TYPE_CHROMA), COMPONENT_Cb, false);
        m_pcRdCost->setDistParam(distParamSatd, orgCb, predCb, pu.cs->sps->getBitDepth(CHANNEL_TYPE_CHROMA), COMPONENT_Cb, true);
        distParamSad.applyWeight = false;
        distParamSatd.applyWeight = false;
        if (PU::isLMCMode(mode))// CCLM，LM_T，LM_L模式，进行跨分量预测
        {
          predIntraChromaLM(COMPONENT_Cb, predCb, pu, areaCb, mode);
        }
        else//传统角度预测
        {
          initPredIntraParams(pu, pu.Cb(), *pu.cs->sps);
          predIntraAng(COMPONENT_Cb, predCb, pu);
        }
        sadCb = distParamSad.distFunc(distParamSad) * 2;//计算Cb的SAD
        satdCb = distParamSatd.distFunc(distParamSatd);//计算Cb的SATD
        sad += std::min(sadCb, satdCb);//选取Cb的SAD和SATD的最小值
        //计算Cr分量的SAD和SATD
        CompArea areaCr = pu.Cr();
        PelBuf orgCr = cs.getOrgBuf(areaCr);
        PelBuf predCr = cs.getPredBuf(areaCr);
        m_pcRdCost->setDistParam(distParamSad, orgCr, predCr, pu.cs->sps->getBitDepth(CHANNEL_TYPE_CHROMA), COMPONENT_Cr, false);
        m_pcRdCost->setDistParam(distParamSatd, orgCr, predCr, pu.cs->sps->getBitDepth(CHANNEL_TYPE_CHROMA), COMPONENT_Cr, true);
        distParamSad.applyWeight = false;
        distParamSatd.applyWeight = false;
        if (PU::isLMCMode(mode))
        {
          predIntraChromaLM(COMPONENT_Cr, predCr, pu, areaCr, mode);
        }
        else
        {
          initPredIntraParams(pu, pu.Cr(), *pu.cs->sps);
          predIntraAng(COMPONENT_Cr, predCr, pu);
        }
        sadCr = distParamSad.distFunc(distParamSad) * 2;
        satdCr = distParamSatd.distFunc(distParamSatd);
        sad += std::min(sadCr, satdCr);
        satdSortedCost[idx] = sad;//保存Cb和Cr的SATD
      }//第一次选择
      // sort the mode based on the cost from small to large.
      // 根据成本从小到大对模式进行排序。
      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;

          }
        }
      }
      int reducedModeNumber = 2; // reduce the number of chroma modes 减少色度模式的数量
      for (int i = 0; i < reducedModeNumber; i++)
      { //禁用最后reducedModelNumber个模式，即将SATD最大的两个模式排除
        modeIsEnable[satdModeList[uiMaxMode - 1 - i]] = 0; // disable the last reducedModeNumber modes
      }

      // save the dist
      // 保存失真
      Distortion baseDist = cs.dist;
      bool testBDPCM = true;
      testBDPCM = testBDPCM && CU::bdpcmAllowed(cu, COMPONENT_Cb) && cu.ispMode == 0 && cu.mtsFlag == 0 && cu.lfnstIdx == 0;
      
      /***********进行第二次选择，把SATD小的三种和未比较的三种共六种模式，编码后计算RDcost *******************/
      for (int32_t uiMode = uiMinMode - (2 * int(testBDPCM)); uiMode < uiMaxMode; uiMode++)
      {
        int chromaIntraMode = chromaCandModes[uiMode];

        if (uiMode < 0)//如果是BDPCM模式
        {
            cu.bdpcmModeChroma = -uiMode;
#if JVET_Q0110_Q0785_CHROMA_BDPCM_420 
            chromaIntraMode = cu.bdpcmModeChroma == 2 ? chromaCandModes[1] : chromaCandModes[2];
#else
            chromaIntraMode = chromaCandModes[0];
#endif
        }
        else
        {
            cu.bdpcmModeChroma = 0;
        if( PU::isLMCMode( chromaIntraMode ) && ! PU::isLMCModeEnabled( pu, chromaIntraMode ) )
        {
          continue;
        }
        //当CCLM被禁用时，MDLM被禁用，不使用satd检查
        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;
        //计算变换量化，反变换反量化，计算重建之后distortion
        xRecurIntraChromaCodingQT( cs, partitioner, bestCostSoFar, ispType );
        if( lumaUsesISP && cs.dist == MAX_UINT )
        {
          continue;
        }

        if (cs.sps->getTransformSkipEnabledFlag())
        {
          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 );//计算RD Cost

        //----- compare -----
        if( dCost < dBestCost )
        {
          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;
          bestBDPCMMode = cu.bdpcmModeChroma;
        }
      }//第二次选择
      // 将最优信号，更新在picture buffer
      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;//保存最佳失真
    cu.bdpcmModeChroma = bestBDPCMMode;
  }//代码块

  //----- restore context models -----
  m_CABACEstimator->getCtx() = ctxStart;
  if( lumaUsesISP && bestCostSoFar >= maxCostAllowed )
  {
    cu.ispMode = 0;
  }
}