src/enc/partitioning/partition_score_func_area.h

// Copyright 2019 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// -----------------------------------------------------------------------------
//
// Block position/size scoring functions.
//
// Author: Yannis Guyon (yguyon@google.com)

#ifndef WP2_ENC_PARTITION_SCORE_FUNC_AREA_H_
#define WP2_ENC_PARTITION_SCORE_FUNC_AREA_H_

#include "src/common/integral.h"
#include "src/common/lossy/block.h"
#include "src/common/lossy/block_size.h"
#include "src/common/lossy/block_size_io.h"
#include "src/common/lossy/predictor.h"
#include "src/common/symbols.h"
#include "src/dec/wp2_dec_i.h"
#include "src/enc/partitioning/partition_score_func_block.h"
#include "src/enc/symbols_enc.h"
#include "src/enc/wp2_enc_i.h"
#include "src/utils/csp.h"
#include "src/utils/plane.h"
#include "src/utils/vector.h"
#include "src/wp2/encode.h"

namespace WP2 {

//------------------------------------------------------------------------------

// Encodes all the blocks in an area to estimate a rate-distortion score.
// Can only be used with a AreaRDOptPartitioner.
class AreaScoreFunc : public BlockScoreFunc {
 public:
  // Dimensions in pixels of the areas.
  AreaScoreFunc(uint32_t area_width, uint32_t area_height);
  WP2Status Init(const EncoderConfig& config, const Rectangle& tile_rect,
                 const YUVPlane& yuv, const GlobalParams& gparams,
                 const ProgressRange& progress) override;
  const Rectangle& GetCurrentArea() const { return area_; }

  // Unused.
  WP2Status ComputeScore(const Block&, const ProgressRange&,
                         float* const) override;

  // Returns the 'blocks' and 'score' of encoding the 'area_' with the default
  // partitioning (multipass).
  WP2Status GetAreaDefaultScore(VectorNoCtor<Block>* const blocks,
                                float* const score);
  // Returns the 'score' of encoding the 'area_' with all 'blocks' having the
  // same 'block_size'.
  WP2Status GetAreaGridScore(BlockSize block_size,
                             VectorNoCtor<Block>* const blocks,
                             float* const score);

  // Sets a block as final. After choosing the partition of the current 'area_',
  // call this for each block in it.
  WP2Status Use(const Block& block) override;

 protected:
  // Returns the distortion and the bitstream size of encoding 'area_blocks'.
  WP2Status GetDistoRate(Vector<CodedBlock>* const area_blocks,
                         BlockScorer* const scorer, float* const disto,
                         float* const rate) const;

  // Setup a new zone to focus on for partitioning.
  // 'area_x' and 'area_y' are coordinates in pixels within the tile.
  virtual WP2Status BeginArea(uint32_t area_x, uint32_t area_y);

  // Returns the 'area_blocks' as if they were selected by MultiScoreFunc.
  WP2Status GetAreaDefaultPartition(
      Vector<CodedBlock>* const area_blocks) const;

  // Defined by 'config_' and not modified outside Init().
  VectorNoCtor<Block> default_partition_;  // Generated by MultiScoreFunc.

  // Objects for the current zone.
  const uint32_t area_width_, area_height_;  // In pixels. Can be less on edges.
  Rectangle area_;  // Boundaries in pixels (no pad) of the region to consider
                    // for the RD-opt. Only use the surroundings as context for
                    // prediction. Next area begins when this area is complete.
  FrontMgrArea area_front_mgr_;    // Used to browse blocks area by area.
  const FrontMgrArea::Comp comp_;  // Used to group blocks per area.

  // Cached default partitioning stats of the current 'area_' for comparison.
  float default_disto_ = -1.f, default_rate_ = -1.f;

 private:  // Visual debug.
  WP2Status RegisterScoreForVDebug(BlockSize grid_size,
                                   const Vector<CodedBlock>& area_blocks,
                                   float score, float disto, float rate) const;
};

//------------------------------------------------------------------------------

// Similar to AreaScoreFunc but the default partitioning is computed and
// compared for each block inside every area instead of once per area.
// Can only be used with a SubAreaRDOptPartitioner.
class SubAreaScoreFunc : public AreaScoreFunc {
 public:
  SubAreaScoreFunc(uint32_t area_width, uint32_t area_height)
      : AreaScoreFunc(area_width, area_height) {}

  WP2Status ComputeScore(const Block& block, const ProgressRange& progress,
                         float* const score) override;
  WP2Status Use(const Block& block) override;

 protected:
  WP2Status BeginArea(uint32_t area_x, uint32_t area_y) override;

  WP2Status GetAreaRemainingDefaultPartition(
      const ProgressRange& progress,
      Vector<CodedBlock>* const area_remaining_blocks) const;

  Vector<CodedBlock> area_used_blocks_;  // Final encoded blocks in this 'area_'

  // Cached default partitioning stats of the current block for comparison.
  Block default_block_;  // BLK_LAST if undefined.
  float default_block_disto_ = 0.f, default_block_rate_ = 0.f;

 private:  // Visual debug.
  WP2Status RegisterScoreForVDebug(
      const Block& block, const Vector<CodedBlock>& area_remaining_blocks,
      float score, float disto, float rate) const;
};

//------------------------------------------------------------------------------

}  // namespace WP2

#endif  // WP2_ENC_PARTITION_SCORE_FUNC_AREA_H_