H.266/VVC代码学习：xIntraCodingTUBlock函数

xIntraCodingTUBlock函数主要是进行帧内编码并计算相应的失真，用于上层函数对最佳预测模式和最佳变换模式的选择。

xIntraCodingTUBlock函数流程如下：

1. 初始化使用模式
2. 进行预测并保存预测信号
3. 由预测信号和原始信号获得残差信号
4. 进行变换和量化
5. 进行反变换
6. 重建
7. 更新失真distortion

代码如下所示：

//帧内编码TU block
void IntraSearch::xIntraCodingTUBlock(TransformUnit &tu, const ComponentID &compID, const bool &checkCrossCPrediction, Distortion& ruiDist, const int &default0Save1Load2, uint32_t* numSig, std::vector* trModes, const bool loadTr)
{
  if (!tu.blocks[compID].valid())
  {
    return;
  }

  CodingStructure &cs                       = *tu.cs;
  m_pcRdCost->setChromaFormat(cs.sps->getChromaFormatIdc());

  const CompArea      &area                 = tu.blocks[compID];
  const SPS           &sps                  = *cs.sps;
  const PPS           &pps                  = *cs.pps;

  const ChannelType    chType               = toChannelType(compID);
  const int            bitDepth             = sps.getBitDepth(chType);

  PelBuf         piOrg                      = cs.getOrgBuf    (area);
  PelBuf         piPred                     = cs.getPredBuf   (area);
  PelBuf         piResi                     = cs.getResiBuf   (area);
  PelBuf         piOrgResi                  = cs.getOrgResiBuf(area);
  PelBuf         piReco                     = cs.getRecoBuf   (area);

  const PredictionUnit &pu                  = *cs.getPU(area.pos(), chType);
  const uint32_t           uiChFinalMode        = PU::getFinalIntraMode(pu, chType);

  const bool           bUseCrossCPrediction = pps.getPpsRangeExtension().getCrossComponentPredictionEnabledFlag() && isChroma( compID ) && PU::isChromaIntraModeCrossCheckMode( pu ) && checkCrossCPrediction;
  const bool           ccUseRecoResi        = m_pcEncCfg->getUseReconBasedCrossCPredictionEstimate();
#if !JVET_O0502_ISP_CLEANUP
  const bool           ispSplitIsAllowed    = sps.getUseISP() && CU::canUseISP( *tu.cu, compID );
#endif


  //===== init availability pattern =====
#if JVET_O0105_ICT
  CHECK( tu.jointCbCr && compID == COMPONENT_Cr, "wrong combination of compID and jointCbCr" );
#endif
  bool jointCbCr = tu.jointCbCr && compID == COMPONENT_Cb;

  if ( compID == COMPONENT_Y )
  {
  PelBuf sharedPredTS( m_pSharedPredTransformSkip[compID], area );
  if( default0Save1Load2 != 2 )
  {
#if JVET_O0502_ISP_CLEANUP
#if JVET_O0106_ISP_4xN_PREDREG_FOR_1xN_2xN
    bool predRegDiffFromTB = CU::isPredRegDiffFromTB(*tu.cu, compID);
    bool firstTBInPredReg = CU::isFirstTBInPredReg(*tu.cu, compID, area);
    CompArea areaPredReg(COMPONENT_Y, tu.chromaFormat, area);
#endif
    if (tu.cu->ispMode && isLuma(compID))
    {
#if JVET_O0106_ISP_4xN_PREDREG_FOR_1xN_2xN
      if (predRegDiffFromTB)
      {
        if (firstTBInPredReg)
        {
          CU::adjustPredArea(areaPredReg);
          initIntraPatternChTypeISP(*tu.cu, areaPredReg, piReco);
        }
      }
      else
#endif
        initIntraPatternChTypeISP(*tu.cu, area, piReco);
    }
    else
    {
      initIntraPatternChType(*tu.cu, area);
    }
#else
#if JVET_O0106_ISP_4xN_PREDREG_FOR_1xN_2xN
    bool predRegDiffFromTB = CU::isPredRegDiffFromTB(*tu.cu, compID);
    bool firstTBInPredReg = CU::isFirstTBInPredReg(*tu.cu, compID, area);
    CompArea areaPredReg(COMPONENT_Y, tu.chromaFormat, area);
    if (predRegDiffFromTB)
    {
      if (firstTBInPredReg)
      {
        CU::adjustPredArea(areaPredReg);
        initIntraPatternChType(*tu.cu, areaPredReg);
      }
    }
    else
#endif
      initIntraPatternChType(*tu.cu, area);
#endif

    //===== get prediction signal =====
    if( compID != COMPONENT_Y && PU::isLMCMode( uiChFinalMode ) )
    {
      {
        xGetLumaRecPixels( pu, area );
      }
      predIntraChromaLM( compID, piPred, pu, area, uiChFinalMode );
    }
    else
    {
      if( PU::isMIP( pu, chType ) )
      {
        predIntraMip( compID, piPred, pu );
      }
      else
      {
#if JVET_O0106_ISP_4xN_PREDREG_FOR_1xN_2xN
        if (predRegDiffFromTB)
        {
          if (firstTBInPredReg)
          {
            PelBuf piPredReg = cs.getPredBuf(areaPredReg);
            predIntraAng(compID, piPredReg, pu);
          }
        }
        else
#endif
          predIntraAng(compID, piPred, pu);
      }
    }


    // save prediction
    if( default0Save1Load2 == 1 )
    {
      sharedPredTS.copyFrom( piPred );
    }
  }
  else
  {
    // load prediction
    piPred.copyFrom( sharedPredTS );
  }
  }


  DTRACE( g_trace_ctx, D_PRED, "@(%4d,%4d) [%2dx%2d] IMode=%d\n", tu.lx(), tu.ly(), tu.lwidth(), tu.lheight(), uiChFinalMode );
  //DTRACE_PEL_BUF( D_PRED, piPred, tu, tu.cu->predMode, COMPONENT_Y );

  const Slice           &slice = *cs.slice;
  bool flag = slice.getLmcsEnabledFlag() && (slice.isIntra() || (!slice.isIntra() && m_pcReshape->getCTUFlag()));
#if JVET_O0105_ICT
  if (isLuma(compID))
  {
#else
  if (flag && slice.getLmcsChromaResidualScaleFlag() && isChroma(compID))
  {
    const Area area = tu.Y().valid() ? tu.Y() : Area(recalcPosition(tu.chromaFormat, tu.chType, CHANNEL_TYPE_LUMA, tu.blocks[tu.chType].pos()), recalcSize(tu.chromaFormat, tu.chType, CHANNEL_TYPE_LUMA, tu.blocks[tu.chType].size()));
    const CompArea &areaY = CompArea(COMPONENT_Y, tu.chromaFormat, area );
#if JVET_O1109_UNFIY_CRS
    int adj = m_pcReshape->calculateChromaAdjVpduNei(tu, areaY);
#else
    PelBuf piPredY;
    piPredY = cs.picture->getPredBuf(areaY);
    const Pel avgLuma = piPredY.computeAvg();
    int adj = m_pcReshape->calculateChromaAdj(avgLuma);
#endif
    tu.setChromaAdj(adj);
  }
#endif
  //===== get residual signal =====
  piResi.copyFrom( piOrg  );
  if (slice.getLmcsEnabledFlag() && m_pcReshape->getCTUFlag() && compID == COMPONENT_Y)
  {
    CompArea      tmpArea(COMPONENT_Y, area.chromaFormat, Position(0, 0), area.size());
    PelBuf tmpPred = m_tmpStorageLCU.getBuf(tmpArea);
    tmpPred.copyFrom(piPred);
    piResi.rspSignal(m_pcReshape->getFwdLUT());
    piResi.subtract(tmpPred);
  }
  else
  piResi.subtract( piPred );

  if (pps.getPpsRangeExtension().getCrossComponentPredictionEnabledFlag() && isLuma(compID))
  {
    piOrgResi.copyFrom (piResi);
  }

  if (bUseCrossCPrediction)
  {
    if (xCalcCrossComponentPredictionAlpha(tu, compID, ccUseRecoResi) == 0)
    {
      return;
    }
    CrossComponentPrediction::crossComponentPrediction(tu, compID, cs.getResiBuf(tu.Y()), piResi, piResi, false);
  }
#if JVET_O0105_ICT
  }
#endif

  //===== transform and quantization =====
  //变换和量化
  //--- init rate estimation arrays for RDOQ ---
  //初始化RDOQ的速率估计数组
  //--- transform and quantization           ---
  TCoeff uiAbsSum = 0;

  const QpParam cQP(tu, compID);

#if RDOQ_CHROMA_LAMBDA
  m_pcTrQuant->selectLambda(compID);
#endif

  flag =flag && (tu.blocks[compID].width*tu.blocks[compID].height > 4);
  //获得色度残差伸缩系数
  if (flag && isChroma(compID) && slice.getLmcsChromaResidualScaleFlag() )
  {
    int cResScaleInv = tu.getChromaAdj();
#if JVET_O0429_CRS_LAMBDA_FIX
    double cResScale = (double)(1 << CSCALE_FP_PREC) / (double)cResScaleInv;
#else
    double cResScale = round((double)(1 << CSCALE_FP_PREC) / (double)cResScaleInv);
#endif
    m_pcTrQuant->setLambda(m_pcTrQuant->getLambda() / (cResScale*cResScale));
#if !JVET_O0105_ICT
    if ( !jointCbCr ) // Joint CbCr signal is to be scaled in the case of joint chroma
    piResi.scaleSignal(cResScaleInv, 1, tu.cu->cs->slice->clpRng(compID));
#endif
  }

  const CompArea &crArea = tu.blocks     [ COMPONENT_Cr ];
  PelBuf          crOrg  = cs.getOrgBuf  ( crArea );
  PelBuf          crPred = cs.getPredBuf ( crArea );
  PelBuf          crResi = cs.getResiBuf ( crArea );
  PelBuf          crReco = cs.getRecoBuf ( crArea );

  if ( jointCbCr )
  {
#if !JVET_O0105_ICT
    // Get Cr prediction and residual
    crResi.copyFrom( crOrg  );
    crResi.subtract( crPred );

    // Create joint residual and store it for Cb component: jointResi = (cbResi - crResi)/2
    piResi.subtractAndHalve( crResi );

    // Scale the joint signal
    if ( flag && slice.getLmcsChromaResidualScaleFlag() )
      piResi.scaleSignal(tu.getChromaAdj(), 1, tu.cu->cs->slice->clpRng(compID));
#endif
    // Lambda is loosened for the joint mode with respect to single modes as the same residual is used for both chroma blocks
    //由于两个色块使用了相同的残差，因此相对于单个模式，对联合模式的Lambda是松散的
#if JVET_O0105_ICT
    const int    absIct = abs( TU::getICTMode(tu) );
    const double lfact  = ( absIct == 1 || absIct == 3 ? 0.8 : 0.5 );
    m_pcTrQuant->setLambda( lfact * m_pcTrQuant->getLambda() );
#else
    m_pcTrQuant->setLambda( 0.60 * m_pcTrQuant->getLambda() );
#endif
  }

#if JVET_O0105_ICT
#if JVET_O0376_SPS_JOINTCBCR_FLAG
  if ( sps.getJointCbCrEnabledFlag() && isChroma(compID) && (tu.cu->cs->slice->getSliceQp() > 18) )
  {
    m_pcTrQuant->setLambda( 1.3 * m_pcTrQuant->getLambda() );
  }
#else
  if( isChroma(compID) && tu.cu->cs->slice->getSliceQp() > 18 )
  {
    m_pcTrQuant->setLambda( 1.3 * m_pcTrQuant->getLambda() );
  }
#endif
#else
#if JVET_O0376_SPS_JOINTCBCR_FLAG
  else if ( sps.getJointCbCrEnabledFlag() && isChroma(compID) && (tu.cu->cs->slice->getSliceQp() > 18) )
#else
  else if ( isChroma(compID) && tu.cu->cs->slice->getSliceQp() > 18 )
#endif
    m_pcTrQuant->setLambda( 1.10 * m_pcTrQuant->getLambda() );
#endif

#if JVET_O0105_ICT
  if( isLuma(compID) )//亮度块变换和量化
  {
#endif
#if JVET_O0502_ISP_CLEANUP
    if (trModes)
    {
      m_pcTrQuant->transformNxN(tu, compID, cQP, trModes, CU::isIntra(*tu.cu) ? m_pcEncCfg->getIntraMTSMaxCand() : m_pcEncCfg->getInterMTSMaxCand());
      tu.mtsIdx = trModes->at(0).first;
    }
    m_pcTrQuant->transformNxN(tu, compID, cQP, uiAbsSum, m_CABACEstimator->getCtx(), loadTr);
#else
  double diagRatio = 0, horVerRatio = 0;
  if( trModes )
  {
    m_pcTrQuant->transformNxN( tu, compID, cQP, trModes, CU::isIntra( *tu.cu ) ? m_pcEncCfg->getIntraMTSMaxCand() : m_pcEncCfg->getInterMTSMaxCand(), ispSplitIsAllowed ? &diagRatio : nullptr, ispSplitIsAllowed ? &horVerRatio : nullptr );
    tu.mtsIdx = trModes->at(0).first;
  }
  m_pcTrQuant->transformNxN( tu, compID, cQP, uiAbsSum, m_CABACEstimator->getCtx(), loadTr, &diagRatio, &horVerRatio );
    if ( !tu.cu->ispMode && isLuma(compID) && ispSplitIsAllowed && tu.mtsIdx == MTS_DCT2_DCT2 && ispSplitIsAllowed )
  {
    m_intraModeDiagRatio        .push_back(diagRatio);
    m_intraModeHorVerRatio      .push_back(horVerRatio);
    m_intraModeTestedNormalIntra.push_back((int)uiChFinalMode);
  }
#endif
  DTRACE( g_trace_ctx, D_TU_ABS_SUM, "%d: comp=%d, abssum=%d\n", DTRACE_GET_COUNTER( g_trace_ctx, D_TU_ABS_SUM ), compID, uiAbsSum );

#if JVET_O0502_ISP_CLEANUP
  if (tu.cu->ispMode && isLuma(compID) && CU::isISPLast(*tu.cu, area, area.compID) && CU::allLumaCBFsAreZero(*tu.cu))
  {
    // ISP has to have at least one non-zero CBF
    ruiDist = MAX_INT;
    return;
  }
#endif
  //--- inverse transform ---
  //反变换
  if (uiAbsSum > 0)
  {
    m_pcTrQuant->invTransformNxN(tu, compID, piResi, cQP);
  }
  else
  {
    piResi.fill(0);
  }
#if JVET_O0105_ICT
  } // isLuma(compID) 

  else // 色度块
  {
    int         codedCbfMask  = 0;
    ComponentID codeCompId    = (tu.jointCbCr ? (tu.jointCbCr >> 1 ? COMPONENT_Cb : COMPONENT_Cr) : compID);
    const QpParam qpCbCr(tu, codeCompId);

    if( tu.jointCbCr )
    {
      ComponentID otherCompId = ( codeCompId==COMPONENT_Cr ? COMPONENT_Cb : COMPONENT_Cr );
      tu.getCoeffs( otherCompId ).fill(0); // do we need that?
      TU::setCbfAtDepth (tu, otherCompId, tu.depth, false );
    }
    PelBuf& codeResi = ( codeCompId == COMPONENT_Cr ? crResi : piResi );
    uiAbsSum = 0;
    m_pcTrQuant->transformNxN(tu, codeCompId, qpCbCr, uiAbsSum, m_CABACEstimator->getCtx());
    DTRACE( g_trace_ctx, D_TU_ABS_SUM, "%d: comp=%d, abssum=%d\n", DTRACE_GET_COUNTER( g_trace_ctx, D_TU_ABS_SUM ), codeCompId, uiAbsSum );
    if( uiAbsSum > 0 )
    {
      m_pcTrQuant->invTransformNxN(tu, codeCompId, codeResi, qpCbCr);
      codedCbfMask += ( codeCompId == COMPONENT_Cb ? 2 : 1 );
    }
    else
    {
      codeResi.fill(0);
    }

    if( tu.jointCbCr )
    {
      if( tu.jointCbCr == 3 && codedCbfMask == 2 )
      {
        codedCbfMask = 3;
        TU::setCbfAtDepth (tu, COMPONENT_Cr, tu.depth, true );
      }
      if( tu.jointCbCr != codedCbfMask )
      {
        ruiDist = std::numeric_limits::max();
        return;
      }
      m_pcTrQuant->invTransformICT( tu, piResi, crResi );
      uiAbsSum = codedCbfMask;
    }
  } //色度块
#endif

  //===== reconstruction =====
  if ( flag && uiAbsSum > 0 && isChroma(compID) && slice.getLmcsChromaResidualScaleFlag() )
  {
    piResi.scaleSignal(tu.getChromaAdj(), 0, tu.cu->cs->slice->clpRng(compID));
#if JVET_O0105_ICT
    if( jointCbCr )
    {
      crResi.scaleSignal(tu.getChromaAdj(), 0, tu.cu->cs->slice->clpRng(COMPONENT_Cr));
    }
#endif
  }
  if (bUseCrossCPrediction)
  {
    CrossComponentPrediction::crossComponentPrediction(tu, compID, cs.getResiBuf(tu.Y()), piResi, piResi, true);
#if JVET_O0105_ICT
    if( jointCbCr )
    {
      CrossComponentPrediction::crossComponentPrediction(tu, COMPONENT_Cr, cs.getResiBuf(tu.Y()), crResi, crResi, true);
    }
#endif
  }

  if (slice.getLmcsEnabledFlag() && m_pcReshape->getCTUFlag() && compID == COMPONENT_Y)
  {
    CompArea      tmpArea(COMPONENT_Y, area.chromaFormat, Position(0,0), area.size());
    PelBuf tmpPred = m_tmpStorageLCU.getBuf(tmpArea);
    tmpPred.copyFrom(piPred);
    piReco.reconstruct(tmpPred, piResi, cs.slice->clpRng(compID));
  }
  else
#if JVET_O0105_ICT
  {
    piReco.reconstruct(piPred, piResi, cs.slice->clpRng( compID ));
    if( jointCbCr )
    {
      crReco.reconstruct(crPred, crResi, cs.slice->clpRng( COMPONENT_Cr ));
    }
  }
#else
  piReco.reconstruct(piPred, piResi, cs.slice->clpRng( compID ));
#endif

#if !JVET_O0105_ICT
  if ( jointCbCr )
  {
    // Cr uses negative of the signalled Cb residual
    if (uiAbsSum > 0)
      crResi.copyAndNegate( piResi );
    else
      crResi.fill(0);

    tu.getCoeffs(COMPONENT_Cr).fill(0);

    // Set cbf also for Cr
    TU::setCbfAtDepth (tu, COMPONENT_Cr, tu.depth, uiAbsSum > 0 ? true : false);

    // Cr reconstruction and its contribution to the total error
    crReco.reconstruct(crPred, crResi, cs.slice->clpRng( COMPONENT_Cr ));

#if WCG_EXT
    if ( m_pcEncCfg->getLumaLevelToDeltaQPMapping().isEnabled() ||
        (m_pcEncCfg->getReshaper()
          && slice.getLmcsEnabledFlag()
         && (m_pcReshape->getCTUFlag() || (isChroma(compID) && m_pcEncCfg->getReshapeIntraCMD()))))
    {
      const CPelBuf orgLuma = cs.getOrgBuf( cs.area.blocks[COMPONENT_Y] );
      ruiDist += m_pcRdCost->getDistPart( crOrg, crReco, bitDepth, COMPONENT_Cr, DF_SSE_WTD, &orgLuma );
    }
    else
#endif
    {
      ruiDist += m_pcRdCost->getDistPart( crOrg, crReco, bitDepth, COMPONENT_Cr, DF_SSE );
    }
  }
#endif

  //===== update distortion =====
#if WCG_EXT
  if (m_pcEncCfg->getLumaLevelToDeltaQPMapping().isEnabled() || (m_pcEncCfg->getReshaper()
    && slice.getLmcsEnabledFlag() && (m_pcReshape->getCTUFlag() || (isChroma(compID) && m_pcEncCfg->getReshapeIntraCMD()))))
  {
    const CPelBuf orgLuma = cs.getOrgBuf( cs.area.blocks[COMPONENT_Y] );
    if (compID == COMPONENT_Y  && !(m_pcEncCfg->getLumaLevelToDeltaQPMapping().isEnabled()))
    {
      CompArea      tmpArea1(COMPONENT_Y, area.chromaFormat, Position(0, 0), area.size());
      PelBuf tmpRecLuma = m_tmpStorageLCU.getBuf(tmpArea1);
      tmpRecLuma.copyFrom(piReco);
      tmpRecLuma.rspSignal(m_pcReshape->getInvLUT());
      ruiDist += m_pcRdCost->getDistPart(piOrg, tmpRecLuma, sps.getBitDepth(toChannelType(compID)), compID, DF_SSE_WTD, &orgLuma);
    }
    else
#if JVET_O0105_ICT
    {
      ruiDist += m_pcRdCost->getDistPart(piOrg, piReco, bitDepth, compID, DF_SSE_WTD, &orgLuma);
      if( jointCbCr )
      {
        ruiDist += m_pcRdCost->getDistPart(crOrg, crReco, bitDepth, COMPONENT_Cr, DF_SSE_WTD, &orgLuma);
      }
    }
#else
      ruiDist += m_pcRdCost->getDistPart(piOrg, piReco, bitDepth, compID, DF_SSE_WTD, &orgLuma);
#endif
  }
  else
#endif
  {
    ruiDist += m_pcRdCost->getDistPart( piOrg, piReco, bitDepth, compID, DF_SSE );
#if JVET_O0105_ICT
    if( jointCbCr )
    {
      ruiDist += m_pcRdCost->getDistPart( crOrg, crReco, bitDepth, COMPONENT_Cr, DF_SSE );
    }
#endif
  }
}