H.266/VVC代码学习笔记15：VTM6.0中的xCheckRDCostMergeTriangle2Nx2N（）函数

VTM6.0中xCheckRDCostMergeTriangle2Nx2N（）函数是帧间预测技术中Merge模式中TPM技术的主函数的入口，想要了解TPM在代码中的实现，这个函数是啃定了，我只是大概看了一下该函数的大致上的流程以及其中三角Merge列表的构建、加权预测、60种三角形组合模式的率失真代价选择的大致过程。今天将所看的代码分享给大家，如过代码中有理解错的地方，还请指正，谢谢大家~，关于三角预测的原理部分可以参考博客：H.266/VVC相关技术学习笔记：帧间预测中的TPM技术（Triangle partition Mode）

还是老套路，用文字过一下该函数的大致流程，再将加了注释的代码附在后面：
1、首先进行一些初始化的操作，定义一些变量、存放预测值以及残差的buffer、RD代价列表之类

2、定义三角Merge单向预测列表，是从常规的Merge列表中派生出来的，然后调用getTriangleMergeCandidates（）函数，获取三角形预测的单向Merge候选列表，关于该函数的细节，请参考博客： H.266/VVC代码学习笔记16：VTM6.0中的getTriangleMergeCandidates（）函数

3、遍历6个三角Merge候选 ，运动补偿出每个候选的单向预测值。

4、遍历60种组合模式（2×6×5），进行亮度分量的SAD——Cost，进行粗选操作，通过调用weightedTriangleBlk（）函数计算亮度分量下的每种三角组合下的加权预测值，关于该函数细节具体参考博客：H.266/VVC代码学习笔记17：VTM6.0中的xWeightedTriangleBlk（）函数
每遍历一个组合模式，需要去通过调用updateCandList（）函数更新三角模式的RD候选模式列表triangleRdModeList以及三角代价列表tianglecandCostList。遍历更新完之后，triangleRdModeList列表中的候选的Cost是从小到大排序的，RD列表的最大长度为3，也就是从60组合中RDCost出3个较优的候选

5、使用SATD-cost限制候选列表的数量，在上一步SAD粗选完之后对RD候选列表的容量左进一步的限制，减少RD候选列表的范围，在3个里再缩减，可以缩减为0，1，2，3中任意一个可能

6、对那些经过SAD以及进一步数量限制筛选出的较优的候选进行遍历，逐个计算出相应色度分量的加权预测值。（注意：在之前用亮度分量进行筛选，筛选完之后再去对相应的色度进行加权。个人认为考虑到两方面的原因：一是亮度分量的可靠性更强，用亮度筛选出的候选比较真实可靠。二是为了节省复杂度，只需要用亮度分量去进行筛选就足够了，不需要再去求色度分量）

7、开始正式的HAD_Cost的细选过程。遍历需要进行SATD细选的列表中的每一个候选，进行SATD_Cost的比较，得到cost最小的那个候选作为最优的三角模式。然后再去与其余的Merge模式去进行率失真代价的比较，选出最优的一种三角组合模式，再在外层与其他Merge模式一起去竞争。这其中调用xEncodeInterResidual（）函数进行帧间残差编码，这与之前的xCheckRDCostMerge2Nx2N（）函数的学习中所调用的一样：H.266/VVC代码学习笔记14：xCheckRDCostMerge2Nx2N（）函数

以上就是xCheckRDCostMergeTriangle2Nx2N（）函数的大致流程，关于这部分代码的详细注释附在下面，供大家参考使用：

void EncCu::xCheckRDCostMergeTriangle2Nx2N( CodingStructure *&tempCS, CodingStructure *&bestCS, Partitioner &partitioner, const EncTestMode& encTestMode )
{
  const Slice &slice = *tempCS->slice;
  const SPS &sps = *tempCS->sps;

  if (slice.getMaxNumTriangleCand() < 2)
    return;

  CHECK( slice.getSliceType() != B_SLICE, "Triangle mode is only applied to B-slices" );

  tempCS->initStructData( encTestMode.qp, encTestMode.lossless );//初始化数据结构

  bool trianglecandHasNoResidual[TRIANGLE_MAX_NUM_CANDS];//60个三角形预测候选是否有残差
  for( int mergeCand = 0; mergeCand < TRIANGLE_MAX_NUM_CANDS; mergeCand++ )
  {
    trianglecandHasNoResidual[mergeCand] = false;//都初始化为false
  }

  bool bestIsSkip = false;

  uint8_t                                         numTriangleCandidate   = TRIANGLE_MAX_NUM_CANDS;//三角候选数量为60,P1预测块有6种MV的选择*P2预测块有5种MV的选择*两种分区划分模式（对角和反对角线）=60
  uint8_t                                         triangleNumMrgSATDCand = TRIANGLE_MAX_NUM_SATD_CANDS;//三角预测用于SATD比较的候选数量最大为3
  PelUnitBuf                                      triangleBuffer[TRIANGLE_MAX_NUM_UNI_CANDS];//三角形的缓存单元数最大为6，三角预测的单向Merge列表长度为6
  PelUnitBuf                                      triangleWeightedBuffer[TRIANGLE_MAX_NUM_CANDS];//三角加权预测值的缓存为60个
  static_vector<uint8_t, TRIANGLE_MAX_NUM_CANDS> triangleRdModeList;//三角形预测的RD的模式列表，首先粗选选出的组合候选放入到最终的RDCost候选列表中
  static_vector<double,  TRIANGLE_MAX_NUM_CANDS> tianglecandCostList;//三角形预测的候选代价列表，首先粗选选出的组合候选的代价放入到最终的代价列表中
  uint8_t                                         numTriangleCandComb = slice.getMaxNumTriangleCand() * (slice.getMaxNumTriangleCand() - 1) * 2;//三角预测模式的组合数量：6*5*2=60


  DistParam distParam;
  const bool useHadamard = !encTestMode.lossless && !tempCS->slice->getDisableSATDForRD();
  m_pcRdCost->setDistParam( distParam, tempCS->getOrgBuf().Y(), m_acMergeBuffer[0].Y(), sps.getBitDepth( CHANNEL_TYPE_LUMA ), COMPONENT_Y, useHadamard );

  const UnitArea localUnitArea( tempCS->area.chromaFormat, Area( 0, 0, tempCS->area.Y().width, tempCS->area.Y().height) );//开辟编码CU的空间

  const double sqrtLambdaForFirstPass = m_pcRdCost->getMotionLambda(encTestMode.lossless);

  MergeCtx triangleMrgCtx;//定义三角Merge单向预测列表，是从常规的Merge列表中派生出来的
  {
    CodingUnit cu( tempCS->area );
    cu.cs       = tempCS;
    cu.predMode = MODE_INTER;//预测模式为帧间模式
    cu.slice    = tempCS->slice;//指向当前的Slice
    cu.tileIdx          = tempCS->picture->brickMap->getBrickIdxRsMap( tempCS->area.lumaPos() );
    cu.triangle = true;//是三角预测模式
    cu.mmvdSkip = false;//不是MMVDSkip模式
    cu.GBiIdx   = GBI_DEFAULT;//双向预测权重索引


    PredictionUnit pu( tempCS->area );//为当前的CU开辟临时缓存空间
    pu.cu = &cu;
    pu.cs = tempCS;
    pu.regularMergeFlag = false;//不是regular模式

    //获取三角形预测的单向Merge候选列表
    PU::getTriangleMergeCandidates( pu, triangleMrgCtx );

    const uint8_t maxNumTriangleCand = pu.cs->slice->getMaxNumTriangleCand();//三角预测的最大候选数

    //遍历6个三角Merge候选
    for (uint8_t mergeCand = 0; mergeCand < maxNumTriangleCand; mergeCand++)
    {
      triangleBuffer[mergeCand] = m_acMergeBuffer[mergeCand].getBuf(localUnitArea);
      triangleMrgCtx.setMergeInfo( pu, mergeCand );//设置三角Merge候选的初始信息
      PU::spanMotionInfo( pu, triangleMrgCtx );

      if( m_pcEncCfg->getMCTSEncConstraint() && ( !( MCTSHelper::checkMvBufferForMCTSConstraint( pu ) ) ) )
      {
        // Do not use this mode
        tempCS->initStructData( encTestMode.qp, encTestMode.lossless );//初始化数据结构
        return;
      }
      m_pcInterSearch->motionCompensation( pu, triangleBuffer[mergeCand] );//运动补偿获得每个候选的单向预测值
    }
  }

  triangleNumMrgSATDCand = min(triangleNumMrgSATDCand, numTriangleCandComb);//需要进行SATDCost比较的三角Merge候选的数量

  {
    CodingUnit &cu      = tempCS->addCU( tempCS->area, partitioner.chType );//定义CU

    partitioner.setCUData( cu );//初始化CU
    cu.slice            = tempCS->slice;//定义片的临时缓存
    cu.tileIdx          = tempCS->picture->brickMap->getBrickIdxRsMap( tempCS->area.lumaPos() );
    cu.skip             = false;
    cu.predMode         = MODE_INTER;//帧间模式
    cu.transQuantBypass = encTestMode.lossless;
    cu.chromaQpAdj      = cu.transQuantBypass ? 0 : m_cuChromaQpOffsetIdxPlus1;
    cu.qp               = encTestMode.qp;
    cu.triangle         = true;//三角模式为真
    cu.mmvdSkip         = false;//不使用Skip模式
    cu.GBiIdx           = GBI_DEFAULT;//双向预测权重索引为默认值

    PredictionUnit &pu  = tempCS->addPU( cu, partitioner.chType );//添加PU

    if( abs(g_aucLog2[cu.lwidth()] - g_aucLog2[cu.lheight()]) >= 2 )//如果当前CU的宽高比大于等于2，则只有一种划分方向，共30种模式
    {
      numTriangleCandidate = 30;
    }
    else
    {
      numTriangleCandidate = TRIANGLE_MAX_NUM_CANDS;//最大三角候选的数量为60
    }

    numTriangleCandidate = min(numTriangleCandidate, numTriangleCandComb);//从三角候选数量和三角的结合候选数量中挑一个最小的


    //遍历60种候选，进行SAD——Cost，进行粗选操作
    for( uint8_t mergeCand = 0; mergeCand < numTriangleCandidate; mergeCand++ )
    {
      bool    splitDir = m_triangleModeTest[mergeCand].m_splitDir;//获取三角的划分方向
      uint8_t candIdx0 = m_triangleModeTest[mergeCand].m_candIdx0;//获取第一个三角分区的候选索引，6个中选一个
      uint8_t candIdx1 = m_triangleModeTest[mergeCand].m_candIdx1;//获取第二个三角分区的候选索引，剩下的5个中选一个


      //将上面三个变量传给当前PU
      pu.triangleSplitDir = splitDir;
      pu.triangleMergeIdx0 = candIdx0;
      pu.triangleMergeIdx1 = candIdx1;

      pu.mergeFlag = true;//Merge模式可用
      pu.regularMergeFlag = false;//regularMerge模式不可用


      triangleWeightedBuffer[mergeCand] = m_acTriangleWeightedBuffer[mergeCand].getBuf( localUnitArea );//三角权重的临时缓存
      triangleBuffer[candIdx0] = m_acMergeBuffer[candIdx0].getBuf( localUnitArea );//第一个分区的预测值的临时缓存
      triangleBuffer[candIdx1] = m_acMergeBuffer[candIdx1].getBuf( localUnitArea );//第二个分区的预测值的临时缓存

      //亮度分量加权过程
      m_pcInterSearch->weightedTriangleBlk( pu, splitDir, CHANNEL_TYPE_LUMA, triangleWeightedBuffer[mergeCand], triangleBuffer[candIdx0], triangleBuffer[candIdx1] );
      
      
      distParam.cur = triangleWeightedBuffer[mergeCand].Y();//当前组合模式的三角亮度权重

      Distortion uiSad = distParam.distFunc( distParam );//计算当前组合模式失真的SAD

      uint32_t uiBitsCand = m_triangleIdxBins[splitDir][candIdx0][candIdx1];//当前组合模式的总比特数（残差+头信息+语法元素）

      double cost = (double)uiSad + (double)uiBitsCand * sqrtLambdaForFirstPass;//计算SAD代价

      updateCandList( mergeCand, cost, triangleRdModeList, tianglecandCostList
        , triangleNumMrgSATDCand );//更新RD候选列表，将cost较小的放入列表中
    }
    //遍历更新完之后，triangleRdModeList列表中的候选的Cost是从小到大排序的，RD列表的最大长度为3，也就是从60组合中RDCost出3个较优的候选


    // limit number of candidates using SATD-costs
    //使用SATD-cost限制候选列表的数量，在上一步SAD粗选完之后对RD候选列表的容量左进一步的限制，减少RD候选列表的范围，在3个里再缩减，可以缩减为0，1，2，3中任意一个可能
    for( uint8_t i = 0; i < triangleNumMrgSATDCand; i++ )
    {
      //如果当前RD候选的Cost大于一定的阈值，则就跳出，将前面那些在Cost阈值范围之内的候选保留在列表中，后面那些超出Cost阈值的候选被排除在外
      if( tianglecandCostList[i] > MRG_FAST_RATIO * tianglecandCostList[0] || tianglecandCostList[i] > getMergeBestSATDCost() )
      {
        triangleNumMrgSATDCand = i;//STAD候选的数量保留为那些在阈值之内的候选数量
        break;
      }
    }

    // perform chroma weighting process
    //执行色度加权过程

    //对那些经过SAD以及进一步数量限制筛选出的较优的候选进行遍历（注意：用亮度分量进行筛选，筛选完之后再去对相应的色度进行加权。个人认为考虑到两方面的原因：一是亮度分量的可靠性更强，用亮度筛选出的候选比较真实可靠。二是为了节省复杂度，只需要用亮度分量去进行筛选就足够了，不需要再去求色度分量）
    for( uint8_t i = 0; i < triangleNumMrgSATDCand; i++ )
    {
      uint8_t  mergeCand = triangleRdModeList[i];//RdModeList列表中的第几个候选
      bool     splitDir  = m_triangleModeTest[mergeCand].m_splitDir;//划分方向
      uint8_t  candIdx0  = m_triangleModeTest[mergeCand].m_candIdx0;//分区一的候选
      uint8_t  candIdx1  = m_triangleModeTest[mergeCand].m_candIdx1;//分区二的候选

      pu.triangleSplitDir = splitDir;
      pu.triangleMergeIdx0 = candIdx0;
      pu.triangleMergeIdx1 = candIdx1;
      pu.mergeFlag = true;//标志Merge模式可用
      pu.regularMergeFlag = false;//标志regular_Merge模式不可用

      //色度分量加权过程
      m_pcInterSearch->weightedTriangleBlk( pu, splitDir, CHANNEL_TYPE_CHROMA, triangleWeightedBuffer[mergeCand], triangleBuffer[candIdx0], triangleBuffer[candIdx1] );
    }

    //再初始化数据结构
    tempCS->initStructData( encTestMode.qp, encTestMode.lossless );
  }
  //重新获取三角预测模式的数量，缩减之后的triangleNumMrgSATDCand和RD列表的尺寸中取最小
  triangleNumMrgSATDCand = min(triangleNumMrgSATDCand, (uint8_t)triangleRdModeList.size());

  m_bestModeUpdated = tempCS->useDbCost = bestCS->useDbCost = false;//是否是最优的Cost
  
  
  //这里开始正式的HAD_Cost的细选过程
  {
    uint8_t iteration;
    uint8_t iterationBegin = 0;//迭代从0开始
    if (encTestMode.lossless)//如果测试的候选模式无损，则迭代一次
    {
      iteration = 1;
    }
    else//否则迭代两次
    {
      iteration = 2;
    }
    //外层的迭代次数
    for (uint8_t noResidualPass = iterationBegin; noResidualPass < iteration; ++noResidualPass)
    {
      //遍历需要进行SATD细选的列表中的每一个候选，进行SATD_Cost的比较，得到cost最小的那个候选作为最优的三角模式。然后再去与其余的Merge模式去进行率失真代价的比较
      for( uint8_t mrgHADIdx = 0; mrgHADIdx < triangleNumMrgSATDCand; mrgHADIdx++ )
      {
        uint8_t mergeCand = triangleRdModeList[mrgHADIdx];

        if ( ( (noResidualPass != 0) && trianglecandHasNoResidual[mergeCand] )
          || ( (noResidualPass == 0) && bestIsSkip ) )
        {
          continue;
        }

        bool    splitDir = m_triangleModeTest[mergeCand].m_splitDir;
        uint8_t candIdx0 = m_triangleModeTest[mergeCand].m_candIdx0;
        uint8_t candIdx1 = m_triangleModeTest[mergeCand].m_candIdx1;

        CodingUnit &cu = tempCS->addCU(tempCS->area, partitioner.chType);

        partitioner.setCUData(cu);//设置CU的数据
        cu.slice = tempCS->slice;
        cu.tileIdx          = tempCS->picture->brickMap->getBrickIdxRsMap( tempCS->area.lumaPos() );
        cu.skip = false;
        cu.predMode = MODE_INTER;
        cu.transQuantBypass = encTestMode.lossless;
        cu.chromaQpAdj = cu.transQuantBypass ? 0 : m_cuChromaQpOffsetIdxPlus1;
        cu.qp = encTestMode.qp;
        cu.triangle = true;
        cu.mmvdSkip = false;
        cu.GBiIdx   = GBI_DEFAULT;
        PredictionUnit &pu = tempCS->addPU(cu, partitioner.chType);

        pu.triangleSplitDir = splitDir;
        pu.triangleMergeIdx0 = candIdx0;
        pu.triangleMergeIdx1 = candIdx1;
        pu.mergeFlag = true;
        pu.regularMergeFlag = false;

        //重新设置三角预测运动信息
        PU::spanTriangleMotionInfo(pu, triangleMrgCtx, splitDir, candIdx0, candIdx1 );

        if( m_pcEncCfg->getMCTSEncConstraint() && ( !( MCTSHelper::checkMvBufferForMCTSConstraint( *cu.firstPU ) ) ) )
        {
          // Do not use this mode
          tempCS->initStructData( encTestMode.qp, encTestMode.lossless );
          return;
        }

        //将之前得到的加权后的预测值直接拷贝过来。用于下面的残差编码（其中就进行了HAD的cost的计算以及比较）
        tempCS->getPredBuf().copyFrom( triangleWeightedBuffer[mergeCand] );//得到三角预测的最终预测值

        //编码帧间残差
        xEncodeInterResidual( tempCS, bestCS, partitioner, encTestMode, noResidualPass, ( noResidualPass == 0 ? &trianglecandHasNoResidual[mergeCand] : NULL ) );

        if (m_pcEncCfg->getUseFastDecisionForMerge() && !bestIsSkip)
        {
          bestIsSkip = bestCS->getCU(partitioner.chType)->rootCbf == 0;
        }
        tempCS->initStructData(encTestMode.qp, encTestMode.lossless);
      }// end loop mrgHADIdx
    }
  }
  if ( m_bestModeUpdated && bestCS->cost != MAX_DOUBLE )//选出最优的一种三角组合模式，再在外层与其他Merge模式一起去竞争
  {
    xCalDebCost( *bestCS, partitioner );
  }
}