Change q_segmentation to spatial_segmentation

Adds spatial prediction to standard segmentation instead of a separate
segmentation feature.
When using temporal prediction skipped blocks are flagged as mispredicted.

Change-Id: I0c32281286d3fbce66d339c9247bcc6516f37a63
diff --git a/av1/encoder/bitstream.c b/av1/encoder/bitstream.c
index d914fc3..1b75dc9 100644
--- a/av1/encoder/bitstream.c
+++ b/av1/encoder/bitstream.c
@@ -615,7 +615,7 @@
 }
 #endif  // CONFIG_LV_MAP
 
-#if CONFIG_Q_SEGMENTATION
+#if CONFIG_SPATIAL_SEGMENTATION
 static int neg_interleave(int x, int ref, int max) {
   const int diff = x - ref;
   if (!ref) return x;
@@ -639,16 +639,16 @@
   }
 }
 
-static void write_q_segment_id(const AV1_COMMON *cm, int skip,
-                               const MB_MODE_INFO *const mbmi, aom_writer *w,
-                               const struct segmentation *seg,
-                               struct segmentation_probs *segp,
-                               BLOCK_SIZE bsize, int mi_row, int mi_col) {
+static void write_segment_id(AV1_COMP *cpi, const MB_MODE_INFO *const mbmi,
+                             aom_writer *w, const struct segmentation *seg,
+                             struct segmentation_probs *segp, int mi_row,
+                             int mi_col, int skip) {
+  AV1_COMMON *const cm = &cpi->common;
   int prev_ul = 0; /* Top left segment_id */
   int prev_l = 0;  /* Current left segment_id */
   int prev_u = 0;  /* Current top segment_id */
 
-  if (!seg->q_lvls) return;
+  if (!seg->enabled || !seg->update_map) return;
 
   MODE_INFO *const mi = cm->mi + mi_row * cm->mi_stride + mi_col;
   int tinfo = mi->mbmi.boundary_info;
@@ -656,39 +656,47 @@
   int left = (!(tinfo & TILE_LEFT_BOUNDARY)) && ((mi_col - 1) >= 0);
 
   if (above && left)
-    prev_ul =
-        get_segment_id(cm, cm->q_seg_map, BLOCK_4X4, mi_row - 1, mi_col - 1);
+    prev_ul = get_segment_id(cm, cm->current_frame_seg_map, BLOCK_4X4,
+                             mi_row - 1, mi_col - 1);
 
   if (above)
-    prev_u = get_segment_id(cm, cm->q_seg_map, BLOCK_4X4, mi_row - 1, mi_col);
+    prev_u = get_segment_id(cm, cm->current_frame_seg_map, BLOCK_4X4,
+                            mi_row - 1, mi_col - 0);
 
   if (left)
-    prev_l = get_segment_id(cm, cm->q_seg_map, BLOCK_4X4, mi_row, mi_col - 1);
+    prev_l = get_segment_id(cm, cm->current_frame_seg_map, BLOCK_4X4,
+                            mi_row - 0, mi_col - 1);
 
-  int cdf_num = pick_q_seg_cdf(prev_ul, prev_u, prev_l);
-  int pred = pick_q_seg_pred(prev_ul, prev_u, prev_l);
+  int cdf_num = pick_spatial_seg_cdf(prev_ul, prev_u, prev_l);
+  int pred = pick_spatial_seg_pred(prev_ul, prev_u, prev_l);
 
   if (skip) {
-    set_q_segment_id(cm, cm->q_seg_map, mbmi->sb_type, mi_row, mi_col, pred);
+    set_spatial_segment_id(cm, cm->current_frame_seg_map, mbmi->sb_type, mi_row,
+                           mi_col, pred);
+    set_spatial_segment_id(cm, cpi->segmentation_map, mbmi->sb_type, mi_row,
+                           mi_col, pred);
+    /* mbmi is read only but we need to update segment_id */
+    ((MB_MODE_INFO *)mbmi)->segment_id = pred;
     return;
   }
 
-  int coded_id = neg_interleave(mbmi->q_segment_id, pred, seg->q_lvls);
+  int coded_id =
+      neg_interleave(mbmi->segment_id, pred, cm->last_active_segid + 1);
 
-  aom_cdf_prob *pred_cdf = segp->q_seg_cdf[cdf_num];
+  aom_cdf_prob *pred_cdf = segp->spatial_pred_seg_cdf[cdf_num];
   aom_write_symbol(w, coded_id, pred_cdf, 8);
 
-  set_q_segment_id(cm, cm->q_seg_map, bsize, mi_row, mi_col,
-                   mbmi->q_segment_id);
+  set_spatial_segment_id(cm, cm->current_frame_seg_map, mbmi->sb_type, mi_row,
+                         mi_col, mbmi->segment_id);
 }
-#endif
-
+#else
 static void write_segment_id(aom_writer *w, const struct segmentation *seg,
                              struct segmentation_probs *segp, int segment_id) {
   if (seg->enabled && seg->update_map) {
     aom_write_symbol(w, segment_id, segp->tree_cdf, MAX_SEGMENTS);
   }
 }
+#endif
 
 #define WRITE_REF_BIT(bname, pname) \
   aom_write_symbol(w, bname, av1_get_pred_cdf_##pname(cm, xd), 2)
@@ -1229,6 +1237,77 @@
 #endif
 }
 
+static void write_inter_segment_id(AV1_COMP *cpi, aom_writer *w,
+                                   const struct segmentation *const seg,
+                                   struct segmentation_probs *const segp,
+                                   int mi_row, int mi_col, int skip,
+                                   int preskip) {
+  MACROBLOCKD *const xd = &cpi->td.mb.e_mbd;
+  const MODE_INFO *mi = xd->mi[0];
+  const MB_MODE_INFO *const mbmi = &mi->mbmi;
+#if CONFIG_SPATIAL_SEGMENTATION
+  AV1_COMMON *const cm = &cpi->common;
+#else
+  (void)mi_row;
+  (void)mi_col;
+  (void)skip;
+  (void)preskip;
+#endif
+
+  if (seg->update_map) {
+#if CONFIG_SPATIAL_SEGMENTATION
+    if (preskip) {
+      if (!cm->preskip_segid) return;
+    } else {
+      if (cm->preskip_segid) return;
+      if (skip) {
+        int prev_segid = mbmi->segment_id;
+        write_segment_id(cpi, mbmi, w, seg, segp, mi_row, mi_col, 0);
+
+        if (seg->temporal_update) {
+          const int pred_flag = mbmi->seg_id_predicted;
+          const int pred_context = av1_get_pred_context_seg_id(xd);
+          unsigned(*temporal_predictor_count)[2] = cm->counts.seg.pred;
+          unsigned *t_unpred_seg_counts = cm->counts.seg.tree_mispred;
+
+          temporal_predictor_count[pred_context][pred_flag]--;
+          if (!pred_flag) t_unpred_seg_counts[prev_segid]--;
+
+          ((MB_MODE_INFO *)mbmi)->seg_id_predicted = 0;
+          temporal_predictor_count[pred_context][0]--;
+          t_unpred_seg_counts[mbmi->segment_id]--;
+        }
+        return;
+      }
+    }
+#endif
+    if (seg->temporal_update) {
+      const int pred_flag = mbmi->seg_id_predicted;
+      aom_cdf_prob *pred_cdf = av1_get_pred_cdf_seg_id(segp, xd);
+      aom_write_symbol(w, pred_flag, pred_cdf, 2);
+      if (!pred_flag) {
+#if CONFIG_SPATIAL_SEGMENTATION
+        write_segment_id(cpi, mbmi, w, seg, segp, mi_row, mi_col, 0);
+#else
+        write_segment_id(w, seg, segp, mbmi->segment_id);
+#endif
+      }
+#if CONFIG_SPATIAL_SEGMENTATION
+      if (pred_flag) {
+        set_spatial_segment_id(cm, cm->current_frame_seg_map, mbmi->sb_type,
+                               mi_row, mi_col, mbmi->segment_id);
+      }
+#endif
+    } else {
+#if CONFIG_SPATIAL_SEGMENTATION
+      write_segment_id(cpi, mbmi, w, seg, segp, mi_row, mi_col, 0);
+#else
+      write_segment_id(w, seg, segp, mbmi->segment_id);
+#endif
+    }
+  }
+}
+
 static void pack_inter_mode_mvs(AV1_COMP *cpi, const int mi_row,
                                 const int mi_col, aom_writer *w) {
   AV1_COMMON *const cm = &cpi->common;
@@ -1251,16 +1330,7 @@
   (void)mi_row;
   (void)mi_col;
 
-  if (seg->update_map) {
-    if (seg->temporal_update) {
-      const int pred_flag = mbmi->seg_id_predicted;
-      aom_cdf_prob *pred_cdf = av1_get_pred_cdf_seg_id(segp, xd);
-      aom_write_symbol(w, pred_flag, pred_cdf, 2);
-      if (!pred_flag) write_segment_id(w, seg, segp, segment_id);
-    } else {
-      write_segment_id(w, seg, segp, segment_id);
-    }
-  }
+  write_inter_segment_id(cpi, w, seg, segp, mi_row, mi_col, 0, 1);
 
 #if CONFIG_EXT_SKIP
   write_skip_mode(cm, xd, segment_id, mi, w);
@@ -1275,9 +1345,9 @@
   }
 #endif  // CONFIG_EXT_SKIP
 
-#if CONFIG_Q_SEGMENTATION
-  write_q_segment_id(cm, skip, mbmi, w, seg, segp, bsize, mi_row, mi_col);
-#endif  // CONFIG_Q_SEGMENTATION
+#if CONFIG_SPATIAL_SEGMENTATION
+  write_inter_segment_id(cpi, w, seg, segp, mi_row, mi_col, skip, 0);
+#endif
 
   write_cdef(cm, w, skip, mi_col, mi_row);
 
@@ -1570,12 +1640,13 @@
 }
 #endif  // CONFIG_INTRABC
 
-static void write_mb_modes_kf(AV1_COMMON *cm, MACROBLOCKD *xd,
+static void write_mb_modes_kf(AV1_COMP *cpi, MACROBLOCKD *xd,
 #if CONFIG_INTRABC
                               const MB_MODE_INFO_EXT *mbmi_ext,
 #endif  // CONFIG_INTRABC
                               const int mi_row, const int mi_col,
                               aom_writer *w) {
+  AV1_COMMON *const cm = &cpi->common;
   FRAME_CONTEXT *ec_ctx = xd->tile_ctx;
   const struct segmentation *const seg = &cm->seg;
   struct segmentation_probs *const segp = &ec_ctx->seg;
@@ -1587,12 +1658,18 @@
   (void)mi_row;
   (void)mi_col;
 
+#if CONFIG_SPATIAL_SEGMENTATION
+  if (cm->preskip_segid && seg->update_map)
+    write_segment_id(cpi, mbmi, w, seg, segp, mi_row, mi_col, 0);
+#else
   if (seg->update_map) write_segment_id(w, seg, segp, mbmi->segment_id);
+#endif
 
   const int skip = write_skip(cm, xd, mbmi->segment_id, mi, w);
 
-#if CONFIG_Q_SEGMENTATION
-  write_q_segment_id(cm, skip, mbmi, w, seg, segp, bsize, mi_row, mi_col);
+#if CONFIG_SPATIAL_SEGMENTATION
+  if (!cm->preskip_segid && seg->update_map)
+    write_segment_id(cpi, mbmi, w, seg, segp, mi_row, mi_col, skip);
 #endif
 
   write_cdef(cm, w, skip, mi_col, mi_row);
@@ -1837,7 +1914,7 @@
                               ((mi_row & MAX_MIB_MASK) << TX_UNIT_HIGH_LOG2);
     }
 #endif  // CONFIG_INTRABC
-    write_mb_modes_kf(cm, xd,
+    write_mb_modes_kf(cpi, xd,
 #if CONFIG_INTRABC
                       cpi->td.mb.mbmi_ext,
 #endif  // CONFIG_INTRABC
@@ -2709,6 +2786,10 @@
     }
   }
 
+#if CONFIG_SPATIAL_SEGMENTATION
+  cm->preskip_segid = 0;
+#endif
+
   // Segmentation data
   aom_wb_write_bit(wb, seg->update_data);
   if (seg->update_data) {
@@ -2719,6 +2800,10 @@
         if (active) {
           const int data = get_segdata(seg, i, j);
           const int data_max = av1_seg_feature_data_max(j);
+#if CONFIG_SPATIAL_SEGMENTATION
+          cm->preskip_segid |= j >= SEG_LVL_REF_FRAME;
+          cm->last_active_segid = i;
+#endif
 
           if (av1_is_segfeature_signed(j)) {
             encode_unsigned_max(wb, abs(data), data_max);
@@ -2732,32 +2817,6 @@
   }
 }
 
-#if CONFIG_Q_SEGMENTATION
-static void encode_q_segmentation(AV1_COMMON *cm,
-                                  struct aom_write_bit_buffer *wb) {
-  int i;
-  struct segmentation *seg = &cm->seg;
-
-  for (i = 0; i < MAX_SEGMENTS; i++) {
-    if (segfeature_active(seg, i, SEG_LVL_ALT_Q)) {
-      seg->q_lvls = 0;
-      return;
-    }
-  }
-
-  aom_wb_write_bit(wb, !!seg->q_lvls);
-  if (!seg->q_lvls) return;
-
-  encode_unsigned_max(wb, seg->q_lvls, MAX_SEGMENTS);
-
-  for (i = 0; i < seg->q_lvls; i++) {
-    const int val = seg->q_delta[i];
-    encode_unsigned_max(wb, abs(val), MAXQ);
-    aom_wb_write_bit(wb, val < 0);
-  }
-}
-#endif
-
 static void write_tx_mode(AV1_COMMON *cm, TX_MODE *mode,
                           struct aom_write_bit_buffer *wb) {
   if (cm->all_lossless) {
@@ -3931,9 +3990,6 @@
   encode_loopfilter(cm, wb);
   encode_quantization(cm, wb);
   encode_segmentation(cm, xd, wb);
-#if CONFIG_Q_SEGMENTATION
-  encode_q_segmentation(cm, wb);
-#endif
   {
     int delta_q_allowed = 1;
 #if !CONFIG_EXT_DELTA_Q
@@ -3945,9 +4001,6 @@
         segment_quantizer_active = 1;
       }
     }
-#if CONFIG_Q_SEGMENTATION
-    segment_quantizer_active |= !!seg->q_lvls;
-#endif
     delta_q_allowed = !segment_quantizer_active;
 #endif
 
@@ -4283,9 +4336,6 @@
   encode_loopfilter(cm, wb);
   encode_quantization(cm, wb);
   encode_segmentation(cm, xd, wb);
-#if CONFIG_Q_SEGMENTATION
-  encode_q_segmentation(cm, wb);
-#endif
   {
     int delta_q_allowed = 1;
 #if !CONFIG_EXT_DELTA_Q
@@ -4297,9 +4347,6 @@
         segment_quantizer_active = 1;
       }
     }
-#if CONFIG_Q_SEGMENTATION
-    segment_quantizer_active |= !!seg->q_lvls;
-#endif
     delta_q_allowed = !segment_quantizer_active;
 #endif