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

今天抽空讲一下目前VTM5.0中的编码端的xRecurIntraChromaCodingQT（）函数的代码细节，该函数是帧内色度预测编码的第二入口函数，里面有所有的帧内色度预测模式的底层的实现细节，根据所传递进来的模式号，对应特定的预测模式进行预测编码。

该函数的主要流程如下：
一、 一些初始化操作，定义一些缓存空间及变量（在代码里有详细备注）
二、 通过getFinalIntraMode（）函数获取当前块所遍历到的色度候选模式号，首先同时对Cb预测块和Cr预测块同时进行初始化（注意：这里和VTM4.0以及之前的版本有很大不同，以前的版本中对于Cb和Cr分量是在该函数中进行两次循环遍历，在同一预测模式下对两个颜色分量的预测块进行两次循环分别预测编码，两个分量有其独立的残差！现如今在VTM5.0中由于联合色度残差模式编码技术的引入，使得此处不需要对两个颜色分量遍历两次，而直接在同一时刻同步处理，在后面的代码中可以看到），获取到模式号后，进行对应模式的预测。

三、 两个颜色分量预测块预测结束之后，对Cb和Cr分量的预测块进行循环，在xIntraCodingTUBlock（）函数中分别求出各自分量独立的残差，里面对残差进行量化以编码，解码反量化反变换的，并且通过失真和码率计算率失真代价（RDCost），并且保存下来，用于后面和联合残差进行RDcost的比较，选出最优的残差编码模式。
关于xIntraCodingTUBlock（）函数的代码细节我在之前的博客已将讲过了，这里我直接给出链接:
H.266/VVC代码学习笔记10：xIntraCodingTUBlock（）函数

四、 如果联合残差模式可用，同样在xIntraCodingTUBlock（）函数中计算CbCr的联合残差，并计算联合残差对应的Cost，和CbCr各自的Cost之和进行比较，选取最优的残差编码模式，这里就是联合色度模式和原始的独立色度模式相互竞争。

xRecurIntraChromaCodingQT（）函数的大致流程就是以上四点，关于代码的细节我都有详细的中文注释，有兴趣的同学可以跟进去看一下：

ChromaCbfs IntraSearch::xRecurIntraChromaCodingQT( CodingStructure &cs, Partitioner& partitioner, const double bestCostSoFar, const PartSplit ispType )
{
  UnitArea currArea                   = partitioner.currArea();
  const bool keepResi                 = cs.sps->getUseLMChroma() || KEEP_PRED_AND_RESI_SIGNALS;
  if( !currArea.Cb().valid() ) return ChromaCbfs( false );


  TransformUnit &currTU               = *cs.getTU( currArea.chromaPos(), CHANNEL_TYPE_CHROMA );//获取当前TU
  const PredictionUnit &pu            = *cs.getPU( currArea.chromaPos(), CHANNEL_TYPE_CHROMA );//获取当前PU

  bool lumaUsesISP                    = !CS::isDualITree( cs ) && currTU.cu->ispMode;
  uint32_t     currDepth                  = partitioner.currTrDepth;
  const PPS &pps                      = *cs.pps;
  ChromaCbfs cbfs                     ( false );

  if (currDepth == currTU.depth)
  {
    if (!currArea.Cb().valid() || !currArea.Cr().valid())
    {
      return cbfs;
    }


    CodingStructure &saveCS = *m_pSaveCS[1];
    saveCS.pcv      = cs.pcv;
    saveCS.picture  = cs.picture;
    saveCS.area.repositionTo( cs.area );
    saveCS.initStructData( MAX_INT, false, true );

    if( !CS::isDualITree( cs ) && currTU.cu->ispMode )
    {
      saveCS.clearCUs();
      CodingUnit& auxCU = saveCS.addCU( *currTU.cu, partitioner.chType );
      auxCU.ispMode = currTU.cu->ispMode;
      saveCS.sps = currTU.cs->sps;
      saveCS.clearPUs();
      saveCS.addPU( *currTU.cu->firstPU, partitioner.chType );
    }
    
    //定义一个TU的临时缓存
    TransformUnit &tmpTU = saveCS.addTU(currArea, partitioner.chType);

    cs.setDecomp(currArea.Cb(), true); // set in advance (required for Cb2/Cr2 in 4:2:2 video)

    const unsigned      numTBlocks  = ::getNumberValidTBlocks( *cs.pcv );//获取有效变换块的数量，总共3个，分别是YCbCr分量
    
#if JVET_N0054_JOINT_CHROMA
    CompArea&  cbArea         = currTU.blocks[COMPONENT_Cb];//定义Cb预测区域
    CompArea&  crArea         = currTU.blocks[COMPONENT_Cr];//定义Cr预测区域
    double     bestCostCb     = MAX_DOUBLE;//Cb预测区域最优的代价
    double     bestCostCr     = MAX_DOUBLE;//Cr预测区域最优的代价
    Distortion bestDistCb     = 0;//Cb预测区域最优的失真
    Distortion bestDistCr     = 0;//Cr预测区域最优的失真
    int        maxModesTested = 0;
    bool       earlyExitISP   = false;
    
    TempCtx ctxStartTU( m_CtxCache );
    TempCtx ctxStart  ( m_CtxCache );
    TempCtx ctxBest   ( m_CtxCache );
   

    ctxStartTU       = m_CABACEstimator->getCtx();
    currTU.jointCbCr = 0;//jointCbCr联合模式标志位初始化为0
    
    // Do predictions here to avoid repeating the "default0Save1Load2" stuff
    //这里的预测是为了避免重复“default0Save1Load2”
    int  predMode   = PU::getFinalIntraMode( pu, CHANNEL_TYPE_CHROMA );//获取当前PU的预测模式，DM模式会从亮度中心位置那里获取一个角度模式
    
    PelBuf piPredCb = cs.getPredBuf(cbArea);//定义Cb预测块的缓存区
    PelBuf piPredCr = cs.getPredBuf(crArea);//定义Cr预测块的缓存区
    
    initIntraPatternChType( *currTU.cu, cbArea);//初始化Cb分量通道
    initIntraPatternChType( *currTU.cu, crArea);//初始化Cr分量通道
    
    if( PU::isLMCMode( predMode ) )
    {
      //分别对CbCr分量进行CCLM预测
      xGetLumaRecPixels( pu, cbArea );
      predIntraChromaLM( COMPONENT_Cb, piPredCb, pu, cbArea, predMode );
      predIntraChromaLM( COMPONENT_Cr, piPredCr, pu, crArea, predMode );
    }
    else
    {
      //分别对CbCr分量进行角度预测
      predIntraAng( COMPONENT_Cb, piPredCb, pu);
      predIntraAng( COMPONENT_Cr, piPredCr, pu);
    }
#endif
    //对Cb和Cr分量的预测块进行循环，分别求出各自独立的残差以及代价，并且保存下来，用于后面和联合残差进行RDcost的比较，选出最优的残差编码模式
    for( uint32_t c = COMPONENT_Cb; c < numTBlocks; c++)
    {
      const ComponentID compID  = ComponentID(c);
      const CompArea&   area    = currTU.blocks[compID];//获取当前分量的区域

      double     dSingleCost    = MAX_DOUBLE;//给单个色度分量一个最大的Cost
      int        bestModeId     = 0;//最优的模式号
#if !JVET_N0054_JOINT_CHROMA
      Distortion singleDistC    = 0;
#endif
      Distortion singleDistCTmp = 0;//单个色度分量的失真临时缓存
      double     singleCostTmp  = 0;//单个色度分量的代价临时缓存

      //如果是DM模式，为ture
      const bool checkCrossComponentPrediction = PU::isChromaIntraModeCrossCheckMode( pu ) && pps.getPpsRangeExtension().getCrossComponentPredictionEnabledFlag() && TU::getCbf( currTU, COMPONENT_Y );

      const int  crossCPredictionModesToTest = checkCrossComponentPrediction ? 2 : 1;
      const int  totalModesToTest            = crossCPredictionModesToTest;
#if JVET_N0054_JOINT_CHROMA
      const bool isOneMode                   = false;//是否为第一个模式，初始化默认为false
      maxModesTested                         = totalModesToTest > maxModesTested ? totalModesToTest : maxModesTested;
#else
      const bool isOneMode                   = (totalModesToTest == 1);
#endif
      
      int currModeId = 0;//当前块的模式号为0
      int default0Save1Load2 = 0;

#if !JVET_N0054_JOINT_CHROMA
      TempCtx ctxStart  ( m_CtxCache );
      TempCtx ctxBest   ( m_CtxCache );
#endif

      if (!isOneMode)
      {
        ctxStart = m_CABACEstimator->getCtx();
      }

      {
        for (int crossCPredictionModeId = 0; crossCPredictionModeId < crossCPredictionModesToTest; crossCPredictionModeId++)
        {
          currTU.compAlpha    [compID] = 0;

          currModeId++;

          const bool isFirstMode = (currModeId == 1);
#if JVET_N0054_JOINT_CHROMA
          const bool isLastMode  = false; // 是否为最后一个模式，默认为false。Always store output to saveCS and tmpTU
#else
          const bool isLastMode  = (currModeId == totalModesToTest); // currModeId is indexed from 1

          if (isOneMode)
          {
            default0Save1Load2 = 0;
          }
          else if (!isOneMode && (crossCPredictionModeId == 0))
          {
            default0Save1Load2 = 1; //save prediction on first mode
          }
          else
          {
            default0Save1Load2 = 2; //load it on subsequent modes
          }
#endif
          
          if (!isFirstMode) // if not first mode to be tested
          {
            m_CABACEstimator->getCtx() = ctxStart;
          }

          singleDistCTmp = 0;

          //里面获取预测后的残差信号，进行变换量化
          xIntraCodingTUBlock( currTU, compID, crossCPredictionModeId != 0, singleDistCTmp, default0Save1Load2 );

          if( ( ( crossCPredictionModeId == 1 ) && ( currTU.compAlpha[compID] == 0 ) ) ) //In order not to code TS flag when cbf is zero, the case for TS with cbf being zero is forbidden.
          {
            singleCostTmp = MAX_DOUBLE;
          }
          else if( lumaUsesISP && bestCostSoFar != MAX_DOUBLE && c == COMPONENT_Cb )
          {
            uint64_t fracBitsTmp = xGetIntraFracBitsQTSingleChromaComponent( cs, partitioner, ComponentID( c ) );
            singleCostTmp = m_pcRdCost->calcRdCost( fracBitsTmp, singleDistCTmp );//singleCostTmp保存的是对应单个色度分量的Cost
            if( isOneMode || ( !isOneMode && !isLastMode ) )
            {
              m_CABACEstimator->getCtx() = ctxStart;
            }
          }
          else if( !isOneMode )
          {
            uint64_t fracBitsTmp = xGetIntraFracBitsQTChroma( currTU, compID );
            singleCostTmp = m_pcRdCost->calcRdCost( fracBitsTmp, singleDistCTmp );
          }

          if( singleCostTmp < dSingleCost )//如果单个分量的Cost在正常范围内
          {
#if JVET_N0054_JOINT_CHROMA
            dSingleCost = singleCostTmp;
            bestModeId  = currModeId;
            
            //将各个分量残差的Cost和Dist分别保存下来
            if ( c == COMPONENT_Cb )
            {
              bestCostCb = singleCostTmp;
              bestDistCb = singleDistCTmp;
            }
            else
            {
              bestCostCr = singleCostTmp;
              bestDistCr = singleDistCTmp;
            }
#else
            dSingleCost = singleCostTmp;
            singleDistC = singleDistCTmp;
            bestModeId  = currModeId;
#endif
            
            if( !isLastMode )
            {
#if KEEP_PRED_AND_RESI_SIGNALS
              saveCS.getPredBuf   (area).copyFrom(cs.getPredBuf   (area));
              saveCS.getOrgResiBuf(area).copyFrom(cs.getOrgResiBuf(area));
#endif
              //保存预测值和残差值
              saveCS.getPredBuf   (area).copyFrom(cs.getPredBuf   (area));
              if( keepResi )
              {
                saveCS.getResiBuf (area).copyFrom(cs.getResiBuf   (area));
              }
              saveCS.getRecoBuf   (area).copyFrom(cs.getRecoBuf   (area));

              tmpTU.copyComponentFrom(currTU, compID);

              ctxBest = m_CABACEstimator->getCtx();
            }
          }
        }
      }

      if( lumaUsesISP && dSingleCost > bestCostSoFar && c == COMPONENT_Cb )
      {
        //亮度的Cost和CbdeCost之和已经大于最优的Cost，因此不需要去测试计算Cr的Cost
        cs.dist = MAX_UINT;
        m_CABACEstimator->getCtx() = ctxStart;
#if JVET_N0054_JOINT_CHROMA
        earlyExitISP               = true;
#endif
        break;
        //return cbfs;
      }

#if JVET_N0054_JOINT_CHROMA

      //使用Cr和Cb的单独编码的一个组件完成，如果仍需要Cr编码，只需切换到最佳Cb上下文即可
      if ( c == COMPONENT_Cb && bestModeId < totalModesToTest)
      {
        m_CABACEstimator->getCtx() = ctxBest;
        //需要用Cb的Cbf来估计Cr Cbf的成本
        currTU.copyComponentFrom(tmpTU, COMPONENT_Cb); // Cbf of Cb is needed to estimate cost for Cr Cbf
      }
#else
      if (bestModeId < totalModesToTest)
      {
#if KEEP_PRED_AND_RESI_SIGNALS
        cs.getPredBuf   (area).copyFrom(saveCS.getPredBuf   (area));
        cs.getOrgResiBuf(area).copyFrom(saveCS.getOrgResiBuf(area));
#endif
        cs.getPredBuf   (area).copyFrom(saveCS.getPredBuf   (area));
        if( keepResi )
        {
          cs.getResiBuf (area).copyFrom(saveCS.getResiBuf   (area));
        }
        cs.getRecoBuf   (area).copyFrom(saveCS.getRecoBuf   (area));

        currTU.copyComponentFrom(tmpTU, compID);

        m_CABACEstimator->getCtx() = ctxBest;
      }

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

      cbfs.cbf(compID) = TU::getCbf(currTU, compID);

      cs.dist += singleDistC;
#endif // not JVET_N0054_JOINT_CHROMA
    }
    
#if JVET_N0054_JOINT_CHROMA
    //这里是对联合残差进行计算以及计算对应的Cost，和CbCr各自的Cost之和进行比较，这里就是联合色度模式和原始的独立色度模式相互竞争
    if ( !earlyExitISP )
    {
      //测试使用联合色度编码
      double     bestCostCbCr   = bestCostCb + bestCostCr;//最优的色度残差代价，初始化为Cb和Cr块各自的最优Cost之和
      Distortion bestDistCbCr   = bestDistCb + bestDistCr;//最优的色度失真，初始化为Cb和Cr块各自的最优Dist之和
      int        bestJointCbCr  = 0;//最优的联合色度
      bool       checkJointCbCr = TU::getCbf(tmpTU, COMPONENT_Cb) || TU::getCbf(tmpTU, COMPONENT_Cr);
      
      if ( checkJointCbCr )//如果是联合色度残差模式
      {
        Distortion distTmp = 0;
        
        currTU.jointCbCr               = 1;//jointCbCr设置为1
        currTU.compAlpha[COMPONENT_Cb] = 0;
        
        m_CABACEstimator->getCtx() = ctxStartTU;
        
        xIntraCodingTUBlock( currTU, COMPONENT_Cb, false, distTmp, 0 );//重新计算对应色度候选模式预测的联合残差，保存在Cb预测块作为Cb的残差，对该联合残差进行变换量化编码、解码反量化反变换，求重建值得到失真以及码率
        
        uint64_t bits  = xGetIntraFracBitsQTChroma( currTU, COMPONENT_Cb );//只计算Cb的比特流，比特流中是联合残差
        double costTmp = m_pcRdCost->calcRdCost( bits, distTmp );//计算联合残差的Cost
        
        if( costTmp < bestCostCbCr )//如果成立，则联合色度模式更优，否则原始的独立色度模式更优
        {
          bestCostCbCr  = costTmp;
          bestDistCbCr  = distTmp;
          bestJointCbCr = 1;
        }
      }
      
      // Retrieve the best CU data (unless it was the very last one tested)
      //检索最佳CU数据（除非是最后测试的数据）
      if ( !(maxModesTested == 1 && !checkJointCbCr) && bestJointCbCr == 0 )
      {
#if KEEP_PRED_AND_RESI_SIGNALS
        cs.getPredBuf   (cbArea).copyFrom(saveCS.getPredBuf   (cbArea));
        cs.getOrgResiBuf(cbArea).copyFrom(saveCS.getOrgResiBuf(cbArea));
        cs.getPredBuf   (crArea).copyFrom(saveCS.getPredBuf   (crArea));
        cs.getOrgResiBuf(crArea).copyFrom(saveCS.getOrgResiBuf(crArea));
#endif
        cs.getPredBuf   (cbArea).copyFrom(saveCS.getPredBuf   (cbArea));
        cs.getPredBuf   (crArea).copyFrom(saveCS.getPredBuf   (crArea));

        if( keepResi )
        {
          cs.getResiBuf (cbArea).copyFrom(saveCS.getResiBuf   (cbArea));
          cs.getResiBuf (crArea).copyFrom(saveCS.getResiBuf   (crArea));
        }
        cs.getRecoBuf   (cbArea).copyFrom(saveCS.getRecoBuf   (cbArea));
        cs.getRecoBuf   (crArea).copyFrom(saveCS.getRecoBuf   (crArea));
        
        currTU.copyComponentFrom(tmpTU, COMPONENT_Cb);
        currTU.copyComponentFrom(tmpTU, COMPONENT_Cr);
        
        m_CABACEstimator->getCtx() = ctxBest;
      }
      
      // Copy results to the picture structures
      cs.picture->getRecoBuf(cbArea).copyFrom(cs.getRecoBuf(cbArea));
      cs.picture->getRecoBuf(crArea).copyFrom(cs.getRecoBuf(crArea));
      cs.picture->getPredBuf(cbArea).copyFrom(cs.getPredBuf(cbArea));
      cs.picture->getPredBuf(crArea).copyFrom(cs.getPredBuf(crArea));
      
      cbfs.cbf(COMPONENT_Cb) = TU::getCbf(currTU, COMPONENT_Cb);
      cbfs.cbf(COMPONENT_Cr) = TU::getCbf(currTU, COMPONENT_Cr);
      
      currTU.jointCbCr = cbfs.cbf(COMPONENT_Cb) ? bestJointCbCr : 0;
      cs.dist         += bestDistCbCr;
    }
#endif // JVET_N0054_JOINT_CHROMA
  }
  else
  {
    unsigned    numValidTBlocks   = ::getNumberValidTBlocks( *cs.pcv );
    ChromaCbfs  SplitCbfs         ( false );

    if( partitioner.canSplit( TU_MAX_TR_SPLIT, cs ) )
    {
      partitioner.splitCurrArea( TU_MAX_TR_SPLIT, cs );
    }
    else if( currTU.cu->ispMode )
    {
      partitioner.splitCurrArea( ispType, cs );
    }
    else
      THROW( "Implicit TU split not available" );

    do
    {
      ChromaCbfs subCbfs = xRecurIntraChromaCodingQT( cs, partitioner, bestCostSoFar, ispType );

      for( uint32_t ch = COMPONENT_Cb; ch < numValidTBlocks; ch++ )
      {
        const ComponentID compID = ComponentID( ch );
        SplitCbfs.cbf( compID ) |= subCbfs.cbf( compID );
      }
    } while( partitioner.nextPart( cs ) );

    partitioner.exitCurrSplit();

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

      cbfs.Cb |= SplitCbfs.Cb;
      cbfs.Cr |= SplitCbfs.Cr;

      if( !lumaUsesISP )
      {
        for( auto &ptu : cs.tus )
        {
          if( currArea.Cb().contains( ptu->Cb() ) || ( !ptu->Cb().valid() && currArea.Y().contains( ptu->Y() ) ) )
          {
            TU::setCbfAtDepth( *ptu, COMPONENT_Cb, currDepth, SplitCbfs.Cb );
            TU::setCbfAtDepth( *ptu, COMPONENT_Cr, currDepth, SplitCbfs.Cr );
          }
        }
      }
    }
  }

  return cbfs;
}