Merge branch 'sao'

build/VS2010/HEVC_encoder.vcxproj

@@ -23,9 +23,7 @@
     <Keyword>Win32Proj</Keyword>
     <RootNamespace>HEVC_encoder</RootNamespace>
   </PropertyGroup>
-
   <Import Project="$(VCTargetsPath)\Microsoft.Cpp.Default.props" />
-
   <PropertyGroup Condition="'$(Configuration)|$(Platform)'=='Debug|Win32'" Label="Configuration">
     <ConfigurationType>StaticLibrary</ConfigurationType>
     <UseDebugLibraries>true</UseDebugLibraries>
@@ -42,9 +40,7 @@
     <ConfigurationType>StaticLibrary</ConfigurationType>
     <UseDebugLibraries>false</UseDebugLibraries>
   </PropertyGroup>
-
   <Import Project="$(VCTargetsPath)\Microsoft.Cpp.props" />
-
   <ImportGroup Label="ExtensionSettings">
     <Import Project="..\..\yasm\vsyasm.props" />
   </ImportGroup>
@@ -88,6 +84,7 @@
     <ClCompile Include="..\..\src\intra.c" />
     <ClCompile Include="..\..\src\nal.c" />
     <ClCompile Include="..\..\src\picture.c" />
+    <ClCompile Include="..\..\src\sao.c" />
     <ClCompile Include="..\..\src\search.c" />
     <ClCompile Include="..\..\src\transform.c" />
   </ItemGroup>
@@ -104,6 +101,7 @@
     <ClInclude Include="..\..\src\intra.h" />
     <ClInclude Include="..\..\src\nal.h" />
     <ClInclude Include="..\..\src\picture.h" />
+    <ClInclude Include="..\..\src\sao.h" />
     <ClInclude Include="..\..\src\search.h" />
     <ClInclude Include="..\..\src\transform.h" />
     <ClInclude Include="..\..\src\x64\test64.h" />

build/VS2010/HEVC_encoder.vcxproj.filters

@@ -69,6 +69,9 @@
     <ClCompile Include="..\..\src\debug.c">
       <Filter>Source Files</Filter>
     </ClCompile>
+    <ClCompile Include="..\..\src\sao.c">
+      <Filter>Source Files</Filter>
+    </ClCompile>
   </ItemGroup>
   <ItemGroup>
     <ClInclude Include="..\..\src\global.h">
@@ -119,6 +122,9 @@
     <ClInclude Include="..\..\src\debug.h">
       <Filter>Header Files</Filter>
     </ClInclude>
+    <ClInclude Include="..\..\src\sao.h">
+      <Filter>Header Files</Filter>
+    </ClInclude>
   </ItemGroup>
   <ItemGroup>
     <YASM Include="..\..\src\x86\test.asm">

src/cabac.c

@@ -11,6 +11,7 @@
 
 #include "cabac.h"
 
+#include <assert.h>
 #include <stdlib.h>
 #include <stdio.h>
 #include <math.h>
@@ -281,20 +282,52 @@ void cabac_write_unary_max_symbol(cabac_data *data, cabac_ctx *ctx, uint32_t sym
 {
   int8_t code_last = max_symbol > symbol;
 
+  assert(symbol <= max_symbol);
+
   if (!max_symbol) return;
   
   data->ctx = &ctx[0];
-  cabac_encode_bin(data, symbol ? 1 : 0);
+  CABAC_BIN(data, symbol ? 1 : 0, "ums");
   
   if (!symbol) return;
   
   while (--symbol) {
     data->ctx = &ctx[offset];
-    cabac_encode_bin(data, 1);
+    CABAC_BIN(data, 1, "ums");
   }
   if (code_last) {
     data->ctx = &ctx[offset];
-    cabac_encode_bin(data, 0);
+    CABAC_BIN(data, 0, "ums");
+  }
+}
+
+/**
+ * This can be used for Truncated Rice binarization with cRiceParam=0.
+ */
+void cabac_write_unary_max_symbol_ep(cabac_data *data, unsigned symbol, unsigned max_symbol)
+{
+  /*if (symbol == 0) {
+    CABAC_BIN_EP(data, 0, "ums_ep");
+  } else {
+    // Make a bit-string of (symbol) times 1 and a single 0, except when
+    // symbol == max_symbol.
+    unsigned bins = ((1 << symbol) - 1) << (symbol < max_symbol);
+    CABAC_BINS_EP(data, bins, symbol + (symbol < max_symbol), "ums_ep");
+  }*/
+  
+  int8_t code_last = max_symbol > symbol;
+
+  assert(symbol <= max_symbol);
+
+  CABAC_BIN_EP(data, symbol ? 1 : 0, "ums_ep");
+  
+  if (!symbol) return;
+  
+  while (--symbol) {
+    CABAC_BIN_EP(data, 1, "ums_ep");
+  }
+  if (code_last) {
+    CABAC_BIN_EP(data, 0, "ums_ep");
   }
 }
 

src/cabac.h

@@ -55,6 +55,7 @@ void cabac_write_ep_ex_golomb(cabac_data *data, uint32_t symbol,
 void cabac_write_unary_max_symbol(cabac_data *data, cabac_ctx *ctx, 
                                   uint32_t symbol, int32_t offset, 
                                   uint32_t max_symbol);
+void cabac_write_unary_max_symbol_ep(cabac_data *data, unsigned symbol, unsigned max_symbol);
 
 
 // Macros

src/context.c

@@ -19,6 +19,8 @@
 
 
 // CONTEXTS
+cabac_ctx g_sao_merge_flag_model;
+cabac_ctx g_sao_type_idx_model;
 cabac_ctx g_split_flag_model[3]; //!< \brief split flag context models
 cabac_ctx g_intra_mode_model;    //!< \brief intra mode context models
 cabac_ctx g_chroma_pred_model[2];
@@ -74,6 +76,9 @@ void init_contexts(encoder_control *encoder, int8_t slice)
   uint16_t i;
 
   // Initialize contexts
+  ctx_init(&g_sao_merge_flag_model, encoder->QP, INIT_SAO_MERGE_FLAG[slice]);
+  ctx_init(&g_sao_type_idx_model, encoder->QP, INIT_SAO_TYPE_IDX[slice]);
+
   ctx_init(&g_cu_merge_flag_ext_model, encoder->QP, INIT_MERGE_FLAG_EXT[slice][0]);
   ctx_init(&g_cu_merge_idx_ext_model, encoder->QP, INIT_MERGE_IDX_EXT[slice][0]);
   ctx_init(&g_cu_pred_mode_model, encoder->QP, INIT_PRED_MODE[slice][0]);

src/context.h

@@ -37,6 +37,8 @@ int32_t context_get_sig_ctx_inc(int32_t pattern_sig_ctx,uint32_t scan_idx,int32_
                                 int32_t pos_y,int32_t block_type,int32_t width, int8_t texture_type);
 
 // CONTEXTS
+extern cabac_ctx g_sao_merge_flag_model;
+extern cabac_ctx g_sao_type_idx_model;
 extern cabac_ctx g_split_flag_model[3];
 extern cabac_ctx g_intra_mode_model;
 extern cabac_ctx g_chroma_pred_model[2];
@@ -65,6 +67,9 @@ extern cabac_ctx g_mvp_idx_model[2];
 extern cabac_ctx g_cu_qt_root_cbf_model;
 #define CNU 154
 
+static const uint8_t INIT_SAO_MERGE_FLAG[3] = { 153, 153, 153 };
+static const uint8_t INIT_SAO_TYPE_IDX[3] = { 160, 185, 200 };
+
 static const uint8_t 
 INIT_QT_ROOT_CBF[3][1] = 
 {

src/encmain.c

@@ -150,7 +150,7 @@ int main(int argc, char *argv[])
   encoder->beta_offset_div2  = 0;
   encoder->tc_offset_div2    = 0;
   // SAO
-  encoder->sao_enable = 0;
+  encoder->sao_enable = 1;
 
   init_encoder_input(&encoder->in, input, cfg->width, cfg->height);
 

src/encoder.c

@@ -27,6 +27,7 @@
 #include "inter.h"
 #include "filter.h"
 #include "search.h"
+#include "sao.h"
 
 int16_t g_lambda_cost[55];
 uint32_t* g_sig_last_scan[3][7];
@@ -320,10 +321,6 @@ void encode_one_frame(encoder_control* encoder)
     bitstream_clear_buffer(encoder->stream);
   }  
   
-  // Filtering
-  if(encoder->deblock_enable) {
-    filter_deblock(encoder);
-  }
   // Calculate checksum
   add_checksum(encoder);
 }
@@ -708,6 +705,8 @@ void encode_VUI(encoder_control* encoder)
 
 void encode_slice_header(encoder_control* encoder)
 {
+  picture *cur_pic = encoder->in.cur_pic;
+
 #ifdef _DEBUG
   printf("=========== Slice ===========\n");
 #endif
@@ -752,8 +751,8 @@ void encode_slice_header(encoder_control* encoder)
     //end if
   //end if
   if (encoder->sao_enable) {
-      WRITE_U(encoder->stream, 1,1, "slice_sao_luma_flag");
+    WRITE_U(encoder->stream, cur_pic->slice_sao_luma_flag, 1, "slice_sao_luma_flag");
-      WRITE_U(encoder->stream, 0,1, "slice_sao_chroma_flag");
+    WRITE_U(encoder->stream, cur_pic->slice_sao_chroma_flag, 1, "slice_sao_chroma_flag");
   }
     
   if (encoder->in.cur_pic->slicetype != SLICE_I) {
@@ -771,10 +770,141 @@ void encode_slice_header(encoder_control* encoder)
     //WRITE_U(encoder->stream, 1, 1, "alignment");
 }
 
+
+void encode_sao_color(encoder_control *encoder, sao_info *sao, color_index color_i)
+{
+  picture *pic = encoder->in.cur_pic;
+
+  // Skip colors with no SAO.
+  if (color_i == COLOR_Y && !pic->slice_sao_luma_flag) {
+    return;
+  } else if (!pic->slice_sao_chroma_flag) {
+    return;
+  }
+
+  if (color_i != COLOR_V) {
+    //CABAC_BIN(&cabac, sao->type, "sao_type_idx");
+    // TR cMax=2
+    // HM codes only the first bin with context.
+    //cabac_write_unary_max_symbol(&cabac, &g_sao_type_idx_model, sao->type, 0, 2);
+    cabac.ctx = &g_sao_type_idx_model;
+    CABAC_BIN(&cabac, sao->type == 0 ? 0 : 1, "sao_type_idx");
+    if (sao->type == SAO_TYPE_BAND) {
+      CABAC_BIN_EP(&cabac, 0, "sao_type_idx_ep");
+    } else if (sao->type == SAO_TYPE_EDGE) {
+      CABAC_BIN_EP(&cabac, 1, "sao_type_idx_ep");
+    }
+  }
+
+  if (sao->type != SAO_TYPE_NONE) {
+    sao_eo_cat i;
+  
+    // TR cMax=7 (for 8bit), cRiseParam=0
+    for (i = SAO_EO_CAT1; i <= SAO_EO_CAT2; ++i) {
+      assert(sao->offsets[i] >= 0);
+      cabac_write_unary_max_symbol_ep(&cabac, sao->offsets[i], SAO_ABS_OFFSET_MAX);
+    }
+    for (i = SAO_EO_CAT3; i <= SAO_EO_CAT4; ++i) {
+      assert(sao->offsets[i] <= 0);
+      cabac_write_unary_max_symbol_ep(&cabac, -sao->offsets[i], SAO_ABS_OFFSET_MAX);
+    }
+
+    if (sao->type == SAO_TYPE_BAND) {
+      for (i = SAO_EO_CAT1; i < SAO_EO_CAT4; ++i) {
+        // Parahprasing spec: "If offset_sign is equal to 0, offsetSign is set
+        // equal to 1. Otherwise to -1."
+        // follows: >=0 is coded as 0, <0 is coded as 1
+        // FL cMax=1 (1 bit)
+        CABAC_BIN_EP(&cabac, sao->offsets[i] >= 0 ? 0 : 1, "sao_offset_sign");
+      }
+      // TODO: sao_band_position
+      // FL cMax=31 (6 bits)
+    } else if (color_i != COLOR_V) {
+      // FL cMax=3 (2 bits)
+      CABAC_BINS_EP(&cabac, sao->eo_class, 2, "sao_eo_class");
+    }
+  }
+}
+
+void encode_sao_merge_flags(encoder_control *encoder, sao_info *sao,
+                            unsigned x_ctb, unsigned y_ctb)
+{
+  // SAO merge flags are not present if merge candidate is not in the same
+  // slice AND tile, but there isn't any such segmentation right now.
+  assert(!USE_SLICES && !USE_TILES);
+
+  // SAO merge flags are not present for the first row and column.
+  if (x_ctb > 0) {
+    cabac.ctx = &g_sao_merge_flag_model;
+    CABAC_BIN(&cabac, sao->merge_left_flag ? 1 : 0, "sao_merge_left_flag");
+  }
+  if (y_ctb > 0 && !sao->merge_left_flag) {
+    cabac.ctx = &g_sao_merge_flag_model;
+    CABAC_BIN(&cabac, sao->merge_up_flag ? 1 : 0, "sao_merge_up_flag");
+  }
+}
+
+/**
+ * \brief Stub that encodes all LCU's as none type.
+ */
+void encode_sao(encoder_control *encoder, unsigned x_lcu, uint16_t y_lcu,
+                sao_info *sao_luma, sao_info *sao_chroma)
+{
+  unsigned sao_type[3] = {SAO_TYPE_NONE, SAO_TYPE_NONE, SAO_TYPE_NONE};
+  picture *pic = encoder->in.cur_pic;
+  
+  // TODO: transmit merge flags outside sao_info
+  encode_sao_merge_flags(encoder, sao_luma, x_lcu, y_lcu);
+
+  // If SAO is merged, nothing else needs to be coded.
+  if (!sao_luma->merge_left_flag && !sao_luma->merge_up_flag) {
+    encode_sao_color(encoder, sao_luma, COLOR_Y);
+    encode_sao_color(encoder, sao_chroma, COLOR_U);
+    encode_sao_color(encoder, sao_chroma, COLOR_V);
+  }
+}
+
 void encode_slice_data(encoder_control* encoder)
 {
   uint16_t x_ctb, y_ctb;
+  picture *pic = encoder->in.cur_pic;
   
+  // Filtering
+  if(encoder->deblock_enable) {
+    filter_deblock(encoder);
+  }
+
+  if (encoder->sao_enable) {
+    pixel *new_y_data = MALLOC(pixel, pic->width * pic->height);
+    pixel *new_u_data = MALLOC(pixel, (pic->width * pic->height) >> 2);
+    pixel *new_v_data = MALLOC(pixel, (pic->width * pic->height) >> 2);
+    memcpy(new_y_data, pic->y_recdata, sizeof(pixel) * pic->width * pic->height);
+    memcpy(new_u_data, pic->u_recdata, sizeof(pixel) * (pic->width * pic->height) >> 2);
+    memcpy(new_v_data, pic->v_recdata, sizeof(pixel) * (pic->width * pic->height) >> 2);
+
+    for (y_ctb = 0; y_ctb < encoder->in.height_in_lcu; y_ctb++) {
+      for (x_ctb = 0; x_ctb < encoder->in.width_in_lcu; x_ctb++) {
+        unsigned stride = encoder->in.width_in_lcu;
+        sao_info *sao_luma = &pic->sao_luma[y_ctb * stride + x_ctb];
+        sao_info *sao_chroma = &pic->sao_chroma[y_ctb * stride + x_ctb];
+        init_sao_info(sao_luma);
+        init_sao_info(sao_chroma);
+
+        sao_search_luma(encoder->in.cur_pic, x_ctb, y_ctb, sao_luma);
+        sao_search_chroma(encoder->in.cur_pic, x_ctb, y_ctb, sao_chroma);
+        // sao_do_merge(encoder, x_ctb, y_ctb, sao_luma, sao_chroma);
+        // sao_do_rdo(encoder, x_ctb, y_ctb, sao_luma, sao_chroma);
+        sao_reconstruct(pic, new_y_data, x_ctb, y_ctb, sao_luma, COLOR_Y);
+        sao_reconstruct(pic, new_u_data, x_ctb, y_ctb, sao_chroma, COLOR_U);
+        sao_reconstruct(pic, new_v_data, x_ctb, y_ctb, sao_chroma, COLOR_V);
+      }
+    }
+
+    free(new_y_data);
+    free(new_u_data);
+    free(new_v_data);
+  }
+
   init_contexts(encoder,encoder->in.cur_pic->slicetype);
 
   // Loop through every LCU in the slice
@@ -785,6 +915,15 @@ void encode_slice_data(encoder_control* encoder)
       uint8_t last_cu_x = (x_ctb == (encoder->in.width_in_lcu - 1)) ? 1 : 0;
       uint8_t depth = 0;
 
+      if (encoder->sao_enable) {
+        picture *pic = encoder->in.cur_pic;
+        unsigned stride = encoder->in.width_in_lcu;
+        sao_info sao_luma = pic->sao_luma[y_ctb * stride + x_ctb];
+        sao_info sao_chroma = pic->sao_chroma[y_ctb * stride + x_ctb];
+
+        encode_sao(encoder, x_ctb, y_ctb, &sao_luma, &sao_chroma);
+      }
+
       // Recursive function for looping through all the sub-blocks
       encode_coding_tree(encoder, x_ctb << MAX_DEPTH, y_ctb << MAX_DEPTH, depth);
 

src/global.h

@@ -65,6 +65,9 @@ typedef int16_t coefficient;
 
 /* END OF CONFIG VARIABLES */
 
+#define LCU_LUMA_SIZE (LCU_WIDTH * LCU_WIDTH)
+#define LCU_CHROMA_SIZE (LCU_WIDTH * LCU_WIDTH >> 2)
+
 #define MAX_REF_PIC_COUNT 5
 
 #define AMVP_MAX_NUM_CANDS 2
@@ -80,11 +83,20 @@ typedef int16_t coefficient;
 #define NO_SCU_IN_LCU(no_lcu) ((no_lcu) << MAX_DEPTH)
 #define WITHIN(val, min_val, max_val) ((min_val) <= (val) && (val) <= (max_val))
 
+#define LOG2_LCU_WIDTH 6
+// CU_TO_PIXEL = y * lcu_width * pic_width + x * lcu_width
+#define CU_TO_PIXEL(x, y, depth, width) (((y) << (LOG2_LCU_WIDTH - (depth))) * (width) \
+                                         + ((x) << (LOG2_LCU_WIDTH - (depth))))
+//#define SIGN3(x) ((x) > 0) ? +1 : ((x) == 0 ? 0 : -1)
+#define SIGN3(x) (((x) > 0) - ((x) < 0))
+
 #define VERSION_STRING "0.2               "
 #define VERSION 0.2
 
 //#define VERBOSE 1
 
+#define SAO_ABS_OFFSET_MAX ((1 << (MIN(BIT_DEPTH, 10) - 5)) - 1)
+
 
 #define SIZE_2Nx2N 0
 #define SIZE_2NxN  1
@@ -92,6 +104,12 @@ typedef int16_t coefficient;
 #define SIZE_NxN   3
 #define SIZE_NONE  15
 
+// These are for marking incomplete implementations that break if slices or
+// tiles are used with asserts. They should be set to 1 if they are ever
+// implemented.
+#define USE_SLICES 0
+#define USE_TILES 0
+
 /* Inlining functions */
 #ifdef _MSC_VER /* Visual studio */
   #define INLINE __forceinline
@@ -110,4 +128,9 @@ typedef int16_t coefficient;
 #define FREE_POINTER(pointer) { free(pointer); pointer = NULL; }
 #define MOVE_POINTER(dst_pointer,src_pointer) { dst_pointer = src_pointer; src_pointer = NULL; }
 
+typedef struct {
+  int x;
+  int y;
+} vector2d;
+
 #endif

src/picture.c

@@ -16,6 +16,8 @@
 #include <stdlib.h>
 #include <math.h>
 
+#include "sao.h"
+
 
 #define PSNRMAX (255.0 * 255.0)
 
@@ -66,6 +68,40 @@ void picture_set_block_residual(picture *pic, uint32_t x_scu, uint32_t y_scu,
   }
 }
 
+/**
+ * \brief BLock Image Transfer from one buffer to another.
+ *
+ * It's a stupidly simple loop that copies pixels.
+ *
+ * \param orig  Start of the originating buffer.
+ * \param dst  Start of the destination buffer.
+ * \param width  Width of the copied region.
+ * \param height  Height of the copied region.
+ * \param orig_stride  Width of a row in the originating buffer.
+ * \param dst_stride  Width of a row in the destination buffer.
+ *
+ * This should be inlined, but it's defined here for now to see if Visual
+ * Studios LTCG will inline it.
+ */
+void picture_blit_pixels(const pixel *orig, pixel *dst,
+                         unsigned width, unsigned height,
+                         unsigned orig_stride, unsigned dst_stride)
+{
+  unsigned y, x;
+
+  const pixel *borig = orig;
+  const pixel *bdst = dst;
+
+  for (y = 0; y < height; ++y) {
+    for (x = 0; x < width; ++x) {
+      dst[x] = orig[x];
+    }
+    // Move pointers to the next row.
+    orig += orig_stride;
+    dst += dst_stride;
+  }
+}
+
 /**
  * \brief Set block coded status
  * \param pic    picture to use
@@ -270,6 +306,11 @@ picture *picture_init(int32_t width, int32_t height,
   pic->coeff_y = NULL; pic->coeff_u = NULL; pic->coeff_v = NULL;
   pic->pred_y = NULL; pic->pred_u = NULL; pic->pred_v = NULL;
 
+  pic->slice_sao_luma_flag = 1;
+  pic->slice_sao_chroma_flag = 1;
+  pic->sao_luma = MALLOC(sao_info, width_in_lcu * height_in_lcu);
+  pic->sao_chroma = MALLOC(sao_info, width_in_lcu * height_in_lcu);
+
   return pic;
 }
 
@@ -309,6 +350,9 @@ int picture_destroy(picture *pic)
   FREE_POINTER(pic->pred_u);
   FREE_POINTER(pic->pred_v);
 
+  FREE_POINTER(pic->sao_luma);
+  FREE_POINTER(pic->sao_chroma);
+
   return 1;
 }
 

src/picture.h

@@ -14,6 +14,9 @@
 
 #include "global.h"
 
+//#include "sao.h"
+struct sao_info_struct;
+
 
 //////////////////////////////////////////////////////////////////////////
 // CONSTANTS
@@ -103,6 +106,10 @@ typedef struct
   cu_info** cu_array;     //!< \brief Info for each CU at each depth.
   uint8_t type;
   uint8_t slicetype;
+  uint8_t slice_sao_luma_flag;
+  uint8_t slice_sao_chroma_flag;
+  struct sao_info_struct *sao_luma;   //!< \brief Array of sao parameters for every LCU.
+  struct sao_info_struct *sao_chroma;   //!< \brief Array of sao parameters for every LCU.
 } picture;
 
 /**
@@ -130,6 +137,10 @@ void picture_set_block_split(picture *pic, uint32_t x_scu, uint32_t y_scu,
                              uint8_t depth, int8_t split);
 void picture_set_block_skipped(picture *pic, uint32_t x_scu, uint32_t y_scu,
                                 uint8_t depth, int8_t skipped);
+void picture_blit_pixels(const pixel* orig, pixel *dst,
+                         unsigned width, unsigned height,
+                         unsigned orig_stride, unsigned dst_stride);
+
 picture_list * picture_list_init(int size);
 int picture_list_resize(picture_list *list, int size);
 int picture_list_destroy(picture_list *list);

src/sao.c

@@ -0,0 +1,362 @@
+/**
+ * \file
+ * 
+ * \author Marko Viitanen ( fador@iki.fi ), 
+ *         Tampere University of Technology,
+ *         Department of Pervasive Computing.
+ * \author Ari Koivula ( ari@koivu.la ), 
+ *         Tampere University of Technology,
+ *         Department of Pervasive Computing.
+ */
+
+#include "sao.h"
+
+#include <string.h>
+
+#include "picture.h"
+
+
+
+void init_sao_info(sao_info *sao) {
+  sao->type = SAO_TYPE_NONE;
+  sao->merge_left_flag = 0;
+  sao->merge_up_flag = 0;
+}
+
+// Mapping of edge_idx values to eo-classes.
+static const unsigned g_sao_eo_idx_to_eo_category[] = { 1, 2, 0, 3, 4 };
+// Mapping relationships between a, b and c to eo_idx.
+#define EO_IDX(a, b, c) (2 + SIGN3((c) - (a)) + SIGN3((c) - (b)))
+
+/**
+ * \param orig_data  Original pixel data. 64x64 for luma, 32x32 for chroma.
+ * \param rec_data  Reconstructed pixel data. 64x64 for luma, 32x32 for chroma.
+ * \param dir_offsets
+ * \param is_chroma  0 for luma, 1 for chroma. Indicates 
+ */
+void calc_sao_edge_dir(const pixel *orig_data, const pixel *rec_data,
+                       int eo_class, int block_width, int block_height,
+                       int cat_sum_cnt[2][NUM_SAO_EDGE_CATEGORIES])
+{
+  int y, x;
+  vector2d a_ofs = g_sao_edge_offsets[eo_class][0];
+  vector2d b_ofs = g_sao_edge_offsets[eo_class][1];
+  // Arrays orig_data and rec_data are quarter size for chroma.
+
+  // Don't sample the edge pixels because this function doesn't have access to
+  // their neighbours.
+  for (y = 1; y < block_height - 1; ++y) {
+    for (x = 1; x < block_width - 1; ++x) {
+      const pixel *c_data = &rec_data[y * block_width + x];
+      pixel a = c_data[a_ofs.y * block_width + a_ofs.x];
+      pixel c = c_data[0];
+      pixel b = c_data[b_ofs.y * block_width + b_ofs.x];
+      
+      int eo_idx = EO_IDX(a, b, c);
+      int eo_cat = g_sao_eo_idx_to_eo_category[eo_idx];
+
+      cat_sum_cnt[0][eo_cat] += orig_data[y * block_width + x] - c;
+      cat_sum_cnt[1][eo_cat] += 1;
+    }
+  }
+}
+
+void sao_reconstruct_color(const pixel *rec_data, pixel *new_rec_data, const sao_info *sao, 
+                           int stride, int new_stride, int block_width, int block_height)
+{
+  int y, x;
+  vector2d a_ofs = g_sao_edge_offsets[sao->eo_class][0];
+  vector2d b_ofs = g_sao_edge_offsets[sao->eo_class][1];
+  // Arrays orig_data and rec_data are quarter size for chroma.
+
+  // Don't sample the edge pixels because this function doesn't have access to
+  // their neighbours.
+  for (y = 0; y < block_height; ++y) {
+    for (x = 0; x < block_width; ++x) {
+      const pixel *c_data = &rec_data[y * stride + x];
+      pixel *new_data = &new_rec_data[y * new_stride + x];
+      pixel a = c_data[a_ofs.y * stride + a_ofs.x];
+      pixel c = c_data[0];
+      pixel b = c_data[b_ofs.y * stride + b_ofs.x];
+      
+      int eo_idx = EO_IDX(a, b, c);
+      int eo_cat = g_sao_eo_idx_to_eo_category[eo_idx];
+
+      new_data[0] = CLIP(0, (1 << BIT_DEPTH) - 1, c_data[0] + sao->offsets[eo_cat]);
+    }
+  }
+}
+
+/**
+ * \brief Calculate dimensions of the buffer used by sao reconstruction.
+ *
+ * This function calculates 4 vectors that can be used to make the temporary
+ * buffers required by sao_reconstruct_color.
+ *
+ * Vector block is the area affected by sao. Vectors tr and br are top-left
+ * margin and bottom-right margin, which contain pixels that are not modified
+ * by the reconstruction of this LCU but are needed by the reconstruction.
+ * Vector rec is the offset from the CU to the required pixel area.
+ *
+ * The margins are always either 0 or 1, depending on the direction of the
+ * edge offset class.
+ *
+ * This also takes into account borders of the picture and non-LCU sized
+ * CU's at the bottom and right of the picture.
+ * 
+ * \ CU + rec
+ *  +------+
+ *  |\ tl  |
+ *  | +--+ |
+ *  | |\ block
+ *  | | \| |
+ *  | +--+ |
+ *  |     \ br
+ *  +------+
+ *
+ * \param pic  Picture.
+ * \param sao  Sao parameters.
+ * \param rec  Top-left corner of the LCU, modified to be top-left corner of 
+ */
+void sao_calc_block_dims(const picture *pic, color_index color_i, 
+                         const sao_info *sao, vector2d *rec, 
+                         vector2d *tl, vector2d *br, vector2d *block)
+{
+  vector2d a_ofs = g_sao_edge_offsets[sao->eo_class][0];
+  vector2d b_ofs = g_sao_edge_offsets[sao->eo_class][1];
+  const int is_chroma = (color_i != COLOR_Y ? 1 : 0);
+  int width = pic->width >> is_chroma;
+  int height = pic->height >> is_chroma;
+  int block_width = LCU_WIDTH >> is_chroma;
+
+  // Handle top and left.
+  if (rec->y == 0) {
+    tl->y = 0;
+    if (a_ofs.y == -1 || b_ofs.y == -1) {
+      block->y -= 1;
+      tl->y += 1;
+    }
+  }
+  if (rec->x == 0) {
+    tl->x = 0;
+    if (a_ofs.x == -1 || b_ofs.x == -1) {
+      block->x -= 1;
+      tl->x += 1;
+    }
+  }
+
+  // Handle right and bottom, taking care of non-LCU sized CUs.
+  if (rec->y + block_width >= height) {
+    br->y = 0;
+    if (rec->y + block_width >= height) {
+      block->y = height - rec->y;
+    }
+    if (a_ofs.y == 1 || b_ofs.y == 1) {
+      block->y -= 1;
+      br->y += 1;
+    }
+  }
+  if (rec->x + block_width >= width) {
+    br->x = 0;
+    if (rec->x + block_width > width) {
+      block->x = width - rec->x;
+    }
+    if (a_ofs.x == 1 || b_ofs.x == 1) {
+      block->x -= 1;
+      br->x += 1;
+    }
+  }
+
+  rec->y = (rec->y == 0 ? 0 : -1);
+  rec->x = (rec->x == 0 ? 0 : -1);
+}
+
+void sao_reconstruct(picture *pic, const pixel *old_rec, 
+                     unsigned x_ctb, unsigned y_ctb, 
+                     const sao_info *sao, color_index color_i)
+{
+  const int is_chroma = (color_i != COLOR_Y ? 1 : 0);
+  const int pic_stride = pic->width >> is_chroma;
+  const int lcu_stride = LCU_WIDTH >> is_chroma;
+  const int buf_stride = lcu_stride + 2;
+
+  pixel *recdata = (color_i == COLOR_Y ? pic->y_recdata : 
+                    (color_i == COLOR_U ? pic->u_recdata : pic->v_recdata));
+  pixel buf_rec[(LCU_WIDTH + 2) * (LCU_WIDTH + 2)];
+  pixel new_rec[LCU_WIDTH * LCU_WIDTH];
+  // Calling CU_TO_PIXEL with depth 1 is the same as using block size of 32.
+  pixel *lcu_rec = &recdata[CU_TO_PIXEL(x_ctb, y_ctb, is_chroma, pic_stride)];
+  const pixel *old_lcu_rec = &old_rec[CU_TO_PIXEL(x_ctb, y_ctb, is_chroma, pic_stride)];
+
+  vector2d ofs;
+  vector2d tl = { 1, 1 };
+  vector2d br = { 1, 1 };
+  vector2d block = { LCU_WIDTH, LCU_WIDTH };
+
+  if (sao->type == SAO_TYPE_NONE) {
+    return;
+  }
+
+  ofs.x = x_ctb * lcu_stride;
+  ofs.y = y_ctb * lcu_stride;
+  block.x = lcu_stride;
+  block.y = lcu_stride;
+  sao_calc_block_dims(pic, color_i, sao, &ofs, &tl, &br, &block);
+
+  // Data to tmp buffer.
+  picture_blit_pixels(&old_lcu_rec[ofs.y * pic_stride + ofs.x], 
+                      buf_rec,
+                      tl.x + block.x + br.x,
+                      tl.y + block.y + br.y,
+                      pic_stride, buf_stride);
+
+  sao_reconstruct_color(&buf_rec[tl.y * buf_stride + tl.x], 
+                        &new_rec[(ofs.y + tl.y) * lcu_stride + ofs.x + tl.x],
+                        sao, 
+                        buf_stride, lcu_stride,
+                        block.x, block.y);
+
+  // Copy reconstructed block from tmp buffer to rec image.
+  picture_blit_pixels(&new_rec[(tl.y + ofs.y) * lcu_stride + (tl.x + ofs.x)], 
+                      &lcu_rec[(tl.y + ofs.y) * pic_stride + (tl.x + ofs.x)],
+                      block.x, block.y, lcu_stride, pic_stride);
+}
+
+
+
+void sao_search_best_mode(const pixel *data[], const pixel *recdata[], 
+                          int block_width, int block_height,
+                          unsigned buf_cnt,
+                          sao_info *sao_out)
+{
+  sao_eo_class edge_class;
+  // This array is used to calculate the mean offset used to minimize distortion.
+  int cat_sum_cnt[2][NUM_SAO_EDGE_CATEGORIES];
+  memset(cat_sum_cnt, 0, sizeof(int) * 2 * NUM_SAO_EDGE_CATEGORIES);
+
+  sao_out->ddistortion = INT_MAX;
+
+  for (edge_class = SAO_EO0; edge_class <= SAO_EO3; ++edge_class) {
+    int edge_offset[NUM_SAO_EDGE_CATEGORIES];
+    int sum_ddistortion = 0;
+    sao_eo_cat edge_cat;
+    unsigned i = 0;
+
+    // Call calc_sao_edge_dir once for luma and twice for chroma.
+    for (i = 0; i < buf_cnt; ++i) {
+      calc_sao_edge_dir(data[i], recdata[i], edge_class,
+                        block_width, block_height, cat_sum_cnt);
+    }
+    
+    for (edge_cat = SAO_EO_CAT1; edge_cat <= SAO_EO_CAT4; ++edge_cat) {
+      int cat_sum = cat_sum_cnt[0][edge_cat];
+      int cat_cnt = cat_sum_cnt[1][edge_cat];
+      
+      // The optimum offset can be calculated by getting the minima of the
+      // fast ddistortion estimation formula. The minima is the mean error
+      // and we round that to the nearest integer.
+      int offset = 0;
+      if (cat_cnt != 0) {
+        offset = (cat_sum + (cat_cnt >> 1)) / cat_cnt;
+        offset = CLIP(-SAO_ABS_OFFSET_MAX, SAO_ABS_OFFSET_MAX, offset);
+      }
+
+      // Sharpening edge offsets can't be encoded, so set them to 0 here.
+      if (edge_cat >= SAO_EO_CAT1 && edge_cat <= SAO_EO_CAT2 && offset < 0) {
+        offset = 0;
+      }
+      if (edge_cat >= SAO_EO_CAT3 && edge_cat <= SAO_EO_CAT4 && offset > 0) {
+        offset = 0;
+      }
+
+      edge_offset[edge_cat] = offset;
+      // The ddistortion is amount by which the SSE of data changes. It should
+      // be negative for all categories, if offset was chosen correctly.
+      // ddistortion = N * h^2 - 2 * h * E, where N is the number of samples 
+      // and E is the sum of errors.
+      // It basically says that all pixels that are not improved by offset
+      // increase increase SSE by h^2 and all pixels that are improved by
+      // offset decrease SSE by h*E.
+      sum_ddistortion += cat_cnt * offset * offset - 2 * offset * cat_sum;
+    }
+    // SAO is not applied for category 0.
+    edge_offset[SAO_EO_CAT0] = 0;
+
+    // Choose the offset class that offers the least error after offset.
+    if (sum_ddistortion < sao_out->ddistortion) {
+      sao_out->eo_class = edge_class;
+      sao_out->ddistortion = sum_ddistortion;
+      memcpy(sao_out->offsets, edge_offset, sizeof(int) * NUM_SAO_EDGE_CATEGORIES);
+    }
+  }
+}
+
+ void sao_search_chroma(const picture *pic, unsigned x_ctb, unsigned y_ctb, sao_info *sao)
+{
+  pixel orig_u[LCU_CHROMA_SIZE];
+  pixel rec_u[LCU_CHROMA_SIZE];
+  pixel orig_v[LCU_CHROMA_SIZE];
+  pixel rec_v[LCU_CHROMA_SIZE];
+  pixel *orig[2] = { orig_u, orig_v };
+  pixel *rec[2] = { rec_u, rec_v };
+  pixel *u_data = &pic->u_data[CU_TO_PIXEL(x_ctb, y_ctb, 1, pic->width / 2)];
+  pixel *u_recdata = &pic->u_recdata[CU_TO_PIXEL(x_ctb, y_ctb, 1, pic->width / 2)];
+  pixel *v_data = &pic->v_data[CU_TO_PIXEL(x_ctb, y_ctb, 1, pic->width / 2)];
+  pixel *v_recdata = &pic->v_recdata[CU_TO_PIXEL(x_ctb, y_ctb, 1, pic->width / 2)];
+  int block_width  = (LCU_WIDTH / 2);
+  int block_height = (LCU_WIDTH / 2);
+
+  if (x_ctb * (LCU_WIDTH / 2) + (LCU_WIDTH / 2) >= (unsigned)pic->width / 2) {
+    block_width = (pic->width - x_ctb * LCU_WIDTH) / 2;
+  }
+  if (y_ctb * (LCU_WIDTH / 2) + (LCU_WIDTH / 2) >= (unsigned)pic->height / 2) {
+    block_height = (pic->height - y_ctb * LCU_WIDTH) / 2;
+  }
+
+  sao->type = SAO_TYPE_EDGE;
+
+  // Fill temporary buffers with picture data.
+  // These buffers are needed only until we switch to a LCU based data
+  // structure for pixels. Then we can give pointers directly to that structure
+  // without making copies.
+  picture_blit_pixels(u_data, orig_u, block_width, block_height,
+                      pic->width / 2, LCU_WIDTH / 2);
+  picture_blit_pixels(v_data, orig_v, block_width, block_height, 
+                      pic->width / 2, LCU_WIDTH / 2);
+  picture_blit_pixels(u_recdata, rec_u, block_width, block_height,
+                      pic->width / 2, LCU_WIDTH / 2);
+  picture_blit_pixels(v_recdata, rec_v, block_width, block_height,
+                      pic->width / 2, LCU_WIDTH / 2);
+
+  sao_search_best_mode(orig, rec, block_width, block_height, 2, sao);
+}
+
+void sao_search_luma(const picture *pic, unsigned x_ctb, unsigned y_ctb, sao_info *sao)
+{
+  pixel orig_y[LCU_LUMA_SIZE];
+  pixel rec_y[LCU_LUMA_SIZE];
+  pixel *orig[1] = { orig_y };
+  pixel *rec[1] = { rec_y };
+  pixel *y_data = &pic->y_data[CU_TO_PIXEL(x_ctb, y_ctb, 0, pic->width)];
+  pixel *y_recdata = &pic->y_recdata[CU_TO_PIXEL(x_ctb, y_ctb, 0, pic->width)];
+  int block_width = LCU_WIDTH;
+  int block_height = LCU_WIDTH;
+
+  if (x_ctb * LCU_WIDTH + LCU_WIDTH >= (unsigned)pic->width) {
+    block_width = pic->width - x_ctb * LCU_WIDTH;
+  }
+  if (y_ctb * LCU_WIDTH + LCU_WIDTH >= (unsigned)pic->height) {
+    block_height = pic->height - y_ctb * LCU_WIDTH;
+  }
+
+  sao->type = SAO_TYPE_EDGE;
+
+  // Fill temporary buffers with picture data.
+  // These buffers are needed only until we switch to a LCU based data
+  // structure for pixels. Then we can give pointers directly to that structure
+  // without making copies.
+  picture_blit_pixels(y_data, orig_y, block_width, block_height, pic->width, LCU_WIDTH);
+  picture_blit_pixels(y_recdata, rec_y, block_width, block_height, pic->width, LCU_WIDTH);
+
+  sao_search_best_mode(orig, rec, block_width, block_height, 1, sao);
+}

src/sao.h

@@ -0,0 +1,54 @@
+#ifndef SAO_H_
+#define SAO_H_
+/**
+ * \file
+ * \brief Coding Unit (CU) and picture data related functions.
+ * 
+ * \author Marko Viitanen ( fador@iki.fi ), 
+ *         Tampere University of Technology,
+ *         Department of Pervasive Computing.
+ * \author Ari Koivula ( ari@koivu.la ), 
+ *         Tampere University of Technology,
+ *         Department of Pervasive Computing.
+ */
+
+#include "global.h"
+#include "picture.h"
+
+
+typedef enum { COLOR_Y = 0, COLOR_U = 1, COLOR_V = 2, NUM_COLORS } color_index;
+typedef enum { SAO_TYPE_NONE = 0, SAO_TYPE_BAND, SAO_TYPE_EDGE } sao_type;
+typedef enum { SAO_EO0 = 0, SAO_EO1, SAO_EO2, SAO_EO3, SAO_NUM_EO } sao_eo_class;
+typedef enum { SAO_EO_CAT0 = 0, SAO_EO_CAT1, SAO_EO_CAT2, SAO_EO_CAT3, SAO_EO_CAT4, NUM_SAO_EDGE_CATEGORIES } sao_eo_cat;
+
+// Offsets of a and b in relation to c.
+// dir_offset[dir][a or b]
+// |       |   a   | a     |     a |
+// | a c b |   c   |   c   |   c   |
+// |       |   b   |     b | b     |
+static const vector2d g_sao_edge_offsets[SAO_NUM_EO][2] = { 
+  { { -1, 0 }, { 1, 0 } },
+  { { 0, -1 }, { 0, 1 } },
+  { { -1, -1 }, { 1, 1 } },
+  { { 1, -1 }, { -1, 1 } }
+};
+
+
+typedef struct sao_info_struct {
+  sao_type type;
+  sao_eo_class eo_class;
+  int ddistortion;
+  int merge_left_flag;
+  int merge_up_flag;
+  int offsets[NUM_SAO_EDGE_CATEGORIES];
+} sao_info;
+
+
+void init_sao_info(sao_info *sao);
+void sao_search_chroma(const picture *pic, unsigned x_ctb, unsigned y_ctb, sao_info *sao);
+void sao_search_luma(const picture *pic, unsigned x_ctb, unsigned y_ctb, sao_info *sao);
+void sao_reconstruct(picture *pic, const pixel *old_rec, 
+                     unsigned x_ctb, unsigned y_ctb, 
+                     const sao_info *sao, color_index color_i);
+
+#endif

src/search.c

@@ -35,11 +35,6 @@
   && (x) + (block_width) <= (width) \
   && (y) + (block_height) <= (height))
 
-typedef struct {
-  int x;
-  int y;
-} vector2d;
-
 /** 
  * This is used in the hexagon_search to select 3 points to search.
  *