Segmentation cleanup, adding {set, get}_segment_id functions.

Change-Id: I55c2688e06ae5d7dfccc1b1983f233ab1c7978db

2 files changed, 71 insertions, 76 deletions

diff --git a/vp9/decoder/vp9_decodemv.c b/vp9/decoder/vp9_decodemv.c
index 9b3cc03aa..2a8521a14 100644
--- a/vp9/decoder/vp9_decodemv.c
+++ b/vp9/decoder/vp9_decodemv.c
@@ -73,14 +73,10 @@ static MB_PREDICTION_MODE read_uv_mode(vp9_reader *r, const vp9_prob *p) {
   return (MB_PREDICTION_MODE)treed_read(r, vp9_uv_mode_tree, p);
 }
 
-// This function reads the current macro block's segnent id from the bitstream
-// It should only be called if a segment map update is indicated.
-static void read_mb_segid(vp9_reader *r, MB_MODE_INFO *mi, MACROBLOCKD *xd) {
-  if (xd->segmentation_enabled && xd->update_mb_segmentation_map) {
-    const vp9_prob *const p = xd->mb_segment_tree_probs;
-    mi->segment_id = vp9_read(r, p[0]) ? 2 + vp9_read(r, p[2])
-                                       : vp9_read(r, p[1]);
-  }
+static int read_mb_segid(vp9_reader *r, MACROBLOCKD *xd) {
+  const vp9_prob *const p = xd->mb_segment_tree_probs;
+  return vp9_read(r, p[0]) ? 2 + vp9_read(r, p[2])
+                           :     vp9_read(r, p[1]);
 }
 
 // This function reads the current macro block's segnent id from the bitstream
@@ -98,6 +94,52 @@ static int read_mb_segid_except(vp9_reader *r,
              :     (pred_seg_id >= 2 ? vp9_read(r, p[1]) : (pred_seg_id == 0));
 }
 
+static void set_segment_id(VP9_COMMON *cm, MB_MODE_INFO *mbmi,
+                           int mb_row, int mb_col, int segment_id) {
+  const int mb_index = mb_row * cm->mb_cols + mb_col;
+  const BLOCK_SIZE_TYPE sb_type = mbmi->sb_type;
+  if (sb_type) {
+    const int bw = 1 << mb_width_log2(sb_type);
+    const int bh = 1 << mb_height_log2(sb_type);
+    const int ymbs = MIN(cm->mb_rows - mb_row, bh);
+    const int xmbs = MIN(cm->mb_cols - mb_col, bw);
+    int x, y;
+
+    for (y = 0; y < ymbs; y++) {
+      for (x = 0; x < xmbs; x++) {
+        const int index = mb_index + (y * cm->mb_cols + x);
+        cm->last_frame_seg_map[index] = segment_id;
+      }
+    }
+  } else {
+    cm->last_frame_seg_map[mb_index] = segment_id;
+  }
+}
+
+static int get_segment_id(VP9_COMMON *cm, MB_MODE_INFO *mbmi,
+                          int mb_row, int mb_col) {
+  const int mb_index = mb_row * cm->mb_cols + mb_col;
+  const BLOCK_SIZE_TYPE sb_type = mbmi->sb_type;
+  if (sb_type) {
+    const int bw = 1 << mb_width_log2(sb_type);
+    const int bh = 1 << mb_height_log2(sb_type);
+    const int ymbs = MIN(cm->mb_rows - mb_row, bh);
+    const int xmbs = MIN(cm->mb_cols - mb_col, bw);
+    int segment_id = INT_MAX;
+    int x, y;
+
+    for (y = 0; y < ymbs; y++) {
+      for (x = 0; x < xmbs; x++) {
+        const int index = mb_index + (y * cm->mb_cols + x);
+        segment_id = MIN(segment_id, cm->last_frame_seg_map[index]);
+      }
+    }
+    return segment_id;
+  } else {
+    return cm->last_frame_seg_map[mb_index];
+  }
+}
+
 extern const int vp9_i8x8_block[4];
 static void kfread_modes(VP9D_COMP *pbi, MODE_INFO *m,
                          int mb_row, int mb_col,
@@ -105,30 +147,14 @@ static void kfread_modes(VP9D_COMP *pbi, MODE_INFO *m,
   VP9_COMMON *const cm = &pbi->common;
   MACROBLOCKD *const xd = &pbi->mb;
   const int mis = cm->mode_info_stride;
-  const int map_index = mb_row * cm->mb_cols + mb_col;
   m->mbmi.ref_frame = INTRA_FRAME;
 
   // Read the Macroblock segmentation map if it is being updated explicitly
   // this frame (reset to 0 by default).
   m->mbmi.segment_id = 0;
-  if (xd->update_mb_segmentation_map) {
-    read_mb_segid(r, &m->mbmi, xd);
-    if (m->mbmi.sb_type) {
-      const int bw = 1 << mb_width_log2(m->mbmi.sb_type);
-      const int bh = 1 << mb_height_log2(m->mbmi.sb_type);
-      const int ymbs = MIN(cm->mb_rows - mb_row, bh);
-      const int xmbs = MIN(cm->mb_cols - mb_col, bw);
-      int x, y;
-
-      for (y = 0; y < ymbs; y++) {
-        for (x = 0; x < xmbs; x++) {
-          const int index = y * cm->mb_cols + x;
-          cm->last_frame_seg_map[map_index + index] =  m->mbmi.segment_id;
-        }
-      }
-    } else {
-      cm->last_frame_seg_map[map_index] = m->mbmi.segment_id;
-    }
+  if (xd->segmentation_enabled && xd->update_mb_segmentation_map) {
+    m->mbmi.segment_id = read_mb_segid(r, xd);
+    set_segment_id(cm, &m->mbmi, mb_row, mb_col, m->mbmi.segment_id);
   }
 
   m->mbmi.mb_skip_coeff = vp9_segfeature_active(&pbi->mb, m->mbmi.segment_id,
@@ -545,44 +571,12 @@ static void read_mb_segment_id(VP9D_COMP *pbi,
                                read_mb_segid_except(r, cm, xd, mb_row, mb_col);
       } else {
         // Normal unpredicted coding mode
-        read_mb_segid(r, mbmi, xd);
+        mbmi->segment_id = read_mb_segid(r, xd);
       }
 
-      if (mbmi->sb_type) {
-        const int bw = 1 << mb_width_log2(mbmi->sb_type);
-        const int bh = 1 << mb_height_log2(mbmi->sb_type);
-        const int ymbs = MIN(cm->mb_rows - mb_row, bh);
-        const int xmbs = MIN(cm->mb_cols - mb_col, bw);
-        int x, y;
-
-        for (y = 0; y < ymbs; y++) {
-          for (x = 0; x < xmbs; x++) {
-            const int index = y * cm->mb_cols + x;
-            cm->last_frame_seg_map[mb_index + index] = mbmi->segment_id;
-          }
-        }
-      } else {
-        cm->last_frame_seg_map[mb_index] = mbmi->segment_id;
-      }
+      set_segment_id(cm, mbmi, mb_row, mb_col, mbmi->segment_id);
     } else {
-      if (mbmi->sb_type) {
-        const int bw = 1 << mb_width_log2(mbmi->sb_type);
-        const int bh = 1 << mb_height_log2(mbmi->sb_type);
-        const int ymbs = MIN(cm->mb_rows - mb_row, bh);
-        const int xmbs = MIN(cm->mb_cols - mb_col, bw);
-        unsigned segment_id = -1;
-        int x, y;
-
-        for (y = 0; y < ymbs; y++) {
-          for (x = 0; x < xmbs; x++) {
-            segment_id = MIN(segment_id,
-                cm->last_frame_seg_map[mb_index + x + y * cm->mb_cols]);
-          }
-        }
-        mbmi->segment_id = segment_id;
-      } else {
-        mbmi->segment_id = cm->last_frame_seg_map[mb_index];
-      }
+      mbmi->segment_id = get_segment_id(cm, mbmi, mb_row, mb_col);
     }
   } else {
     // The encoder explicitly sets the segment_id to 0
diff --git a/vp9/decoder/vp9_decodframe.c b/vp9/decoder/vp9_decodframe.c
index f8ef6c030..7f958801d 100644
--- a/vp9/decoder/vp9_decodframe.c
+++ b/vp9/decoder/vp9_decodframe.c
@@ -1146,42 +1146,43 @@ static void update_frame_size(VP9D_COMP *pbi) {
 static void setup_segmentation(VP9_COMMON *pc, MACROBLOCKD *xd, vp9_reader *r) {
   int i, j;
 
+  xd->update_mb_segmentation_map = 0;
+  xd->update_mb_segmentation_data = 0;
+
   xd->segmentation_enabled = vp9_read_bit(r);
   if (xd->segmentation_enabled) {
-    // Read whether or not the segmentation map is being explicitly updated
-    // this frame.
+    // Segmentation map update
     xd->update_mb_segmentation_map = vp9_read_bit(r);
-
     if (xd->update_mb_segmentation_map) {
-      // Which macro block level features are enabled. Read the probs used to
-      // decode the segment id for each macro block.
       for (i = 0; i < MB_FEATURE_TREE_PROBS; i++)
-        xd->mb_segment_tree_probs[i] = vp9_read_bit(r) ? vp9_read_prob(r) : 255;
+        xd->mb_segment_tree_probs[i] = vp9_read_bit(r) ? vp9_read_prob(r)
+                                                       : MAX_PROB;
 
-      // Read the prediction probs needed to decode the segment id
       pc->temporal_update = vp9_read_bit(r);
       if (pc->temporal_update) {
         const vp9_prob *p = xd->mb_segment_tree_probs;
-        vp9_prob *p_mod = xd->mb_segment_mispred_tree_probs;
+        vp9_prob *mispred_p = xd->mb_segment_mispred_tree_probs;
 
         const int c0 =        p[0]  *        p[1];
         const int c1 =        p[0]  * (256 - p[1]);
         const int c2 = (256 - p[0]) *        p[2];
         const int c3 = (256 - p[0]) * (256 - p[2]);
 
-        p_mod[0] = get_binary_prob(c1, c2 + c3);
-        p_mod[1] = get_binary_prob(c0, c2 + c3);
-        p_mod[2] = get_binary_prob(c0 + c1, c3);
-        p_mod[3] = get_binary_prob(c0 + c1, c2);
+        mispred_p[0] = get_binary_prob(c1, c2 + c3);
+        mispred_p[1] = get_binary_prob(c0, c2 + c3);
+        mispred_p[2] = get_binary_prob(c0 + c1, c3);
+        mispred_p[3] = get_binary_prob(c0 + c1, c2);
 
         for (i = 0; i < PREDICTION_PROBS; i++)
-          pc->segment_pred_probs[i] = vp9_read_bit(r) ? vp9_read_prob(r) : 255;
+          pc->segment_pred_probs[i] = vp9_read_bit(r) ? vp9_read_prob(r)
+                                                      : MAX_PROB;
       } else {
         for (i = 0; i < PREDICTION_PROBS; i++)
-          pc->segment_pred_probs[i] = 255;
+          pc->segment_pred_probs[i] = MAX_PROB;
       }
     }
 
+    // Segmentation data update
     xd->update_mb_segmentation_data = vp9_read_bit(r);
     if (xd->update_mb_segmentation_data) {
       xd->mb_segment_abs_delta = vp9_read_bit(r);