Further work on Segmentation Experiment:

This check in includes quite a lot of clean up and refactoring.

Most of the analysis and set up for the different coding options for the
segment map (currently simple distribution based coding or temporaly
predicted coding), has been moved to one location (the function
choose_segmap_coding_method() in segmenation.c). This code was previously
scattered around in various locations making integration with other
experiments and modification / debug more difficult.

Currently the functionality is as it was with the exception that the
prediction probabilities are now only transmitted when the temporal
prediction mode is selected.

There is still quite a bit more clean up work that will be possible
when the #ifdef is removed. Also at that time I may rename and alter
the sense of macroblock based variable "segment_flag" which indicates
(1 that the segmnet id is not predicted vs 0 that it is predicted).

I also intend to experiment with a spatial prediction mode that can be
used when coding a key frame segment map or in cases where temporal
prediction does not work well but there is spatial correlation.

In a later check in when the ifdefs have gone I may also move the call
to choose_segmap_coding_method() to just before where the bitsream is
packed (currently it is in vp8_encode_frame()) to further reduce the
possibility of clashes with other experiments and prevent it being called
on each itteration of the recode loop.

Change-Id: I3d4aba2a2826ec21f367678d5b07c1d1c36db168

9 files changed, 252 insertions, 215 deletions

diff --git a/vp8/common/blockd.h b/vp8/common/blockd.h
index 690e0f8c8..c477d2d6d 100644
--- a/vp8/common/blockd.h
+++ b/vp8/common/blockd.h
@@ -32,6 +32,7 @@ void vpx_log(const char *format, ...);
 #define DCPREDCNTTHRESH 3
 
 #define MB_FEATURE_TREE_PROBS   3
+#define SEGMENT_PREDICTION_PROBS 3
 
 #define MAX_MB_SEGMENTS         4
 
@@ -187,6 +188,7 @@ typedef struct
     unsigned char mb_skip_coeff;                                /* does this mb has coefficients at all, 1=no coefficients, 0=need decode tokens */
     unsigned char need_to_clamp_mvs;
     unsigned char segment_id;                  /* Which set of segmentation parameters should be used for this MB */
+
 } MB_MODE_INFO;
 
 typedef struct
@@ -258,11 +260,14 @@ typedef struct MacroBlockD
 
     /* Per frame flags that define which MB level features (such as quantizer or loop filter level) */
     /* are enabled and when enabled the proabilities used to decode the per MB flags in MB_MODE_INFO */
+
+    // Probability Tree used to code Segment number
+    vp8_prob mb_segment_tree_probs[MB_FEATURE_TREE_PROBS];
+
 #if CONFIG_SEGMENTATION
-    vp8_prob mb_segment_tree_probs[MB_FEATURE_TREE_PROBS + 3];         // Probability Tree used to code Segment number
+    // Context probabilities when using predictive coding of segment id
+    vp8_prob mb_segment_pred_probs[SEGMENT_PREDICTION_PROBS];
     unsigned char temporal_update;
-#else
-    vp8_prob mb_segment_tree_probs[MB_FEATURE_TREE_PROBS];
 #endif
 
     // Segment features
diff --git a/vp8/decoder/decodemv.c b/vp8/decoder/decodemv.c
index e13fd6f09..416865cae 100644
--- a/vp8/decoder/decodemv.c
+++ b/vp8/decoder/decodemv.c
@@ -18,10 +18,6 @@
 //#if CONFIG_SEGFEATURES
 #include "vp8/common/seg_common.h"
 
-#if CONFIG_SEGMENTATION
-#include "vp8/common/seg_common.h"
-#endif
-
 #if CONFIG_DEBUG
 #include <assert.h>
 #endif
@@ -434,7 +430,7 @@ static void read_mb_modes_mv(VP8D_COMP *pbi, MODE_INFO *mi, MB_MODE_INFO *mbmi,
     MACROBLOCKD *const xd  = & pbi->mb;
 
 #if CONFIG_SEGMENTATION
-    int sum;
+    int pred_context;
     int index = mb_row * pbi->common.mb_cols + mb_col;
 #endif
     int_mv *const mv = & mbmi->mv;
@@ -465,14 +461,16 @@ static void read_mb_modes_mv(VP8D_COMP *pbi, MODE_INFO *mi, MB_MODE_INFO *mbmi,
 #if CONFIG_SEGMENTATION
                 if (xd->temporal_update)
                 {
-                    sum = 0;
+                    pred_context = 0;
 
                     if (mb_col != 0)
-                        sum += (mi-1)->mbmi.segment_flag;
+                        pred_context += (mi-1)->mbmi.segment_flag;
                     if (mb_row != 0)
-                        sum += (mi-pbi->common.mb_cols)->mbmi.segment_flag;
+                        pred_context +=
+                            (mi-pbi->common.mb_cols)->mbmi.segment_flag;
 
-                    if (vp8_read(bc, xd->mb_segment_tree_probs[3+sum]) == 0)
+                    if (vp8_read(bc,
+                                 xd->mb_segment_pred_probs[pred_context]) == 0)
                     {
                         mbmi->segment_id = pbi->segmentation_map[index];
                         mbmi->segment_flag = 0;
diff --git a/vp8/decoder/decodframe.c b/vp8/decoder/decodframe.c
index 5ba499155..c572b0714 100644
--- a/vp8/decoder/decodframe.c
+++ b/vp8/decoder/decodframe.c
@@ -1061,18 +1061,38 @@ int vp8_decode_frame(VP8D_COMP *pbi)
         if (xd->update_mb_segmentation_map)
         {
             /* Which macro block level features are enabled */
-            vpx_memset(xd->mb_segment_tree_probs, 255, sizeof(xd->mb_segment_tree_probs));
-#if CONFIG_SEGMENTATION
-            /* Read the probs used to decode the segment id for each macro block. */
-            for (i = 0; i < MB_FEATURE_TREE_PROBS+3; i++)
-#else
+            vpx_memset(xd->mb_segment_tree_probs, 255,
+                       sizeof(xd->mb_segment_tree_probs));
+            vpx_memset(xd->mb_segment_pred_probs, 255,
+                       sizeof(xd->mb_segment_pred_probs));
+
+            // Read the probs used to decode the segment id for each macro
+            // block.
             for (i = 0; i < MB_FEATURE_TREE_PROBS; i++)
-#endif
             {
-                /* If not explicitly set value is defaulted to 255 by memset above */
+                // If not explicitly set value is defaulted to 255 by
+                //memset above
                 if (vp8_read_bit(bc))
-                    xd->mb_segment_tree_probs[i] = (vp8_prob)vp8_read_literal(bc, 8);
+                    xd->mb_segment_tree_probs[i] =
+                        (vp8_prob)vp8_read_literal(bc, 8);
             }
+#if CONFIG_SEGMENTATION
+            // If predictive coding of segment map is enabled read the
+            // prediction probabilities.
+            if ( xd->temporal_update )
+            {
+                // Read the prediction probs needed to decode the segment id
+                // when predictive coding enabled
+                for (i = 0; i < SEGMENT_PREDICTION_PROBS; i++)
+                {
+                    // If not explicitly set value is defaulted to 255 by
+                    // memset above
+                    if (vp8_read_bit(bc))
+                        xd->mb_segment_pred_probs[i] =
+                            (vp8_prob)vp8_read_literal(bc, 8);
+                }
+            }
+#endif
         }
     }
 
diff --git a/vp8/encoder/bitstream.c b/vp8/encoder/bitstream.c
index ff8fb8a48..eda40f375 100644
--- a/vp8/encoder/bitstream.c
+++ b/vp8/encoder/bitstream.c
@@ -64,10 +64,6 @@ extern unsigned int active_section;
 
 #ifdef MODE_STATS
 int count_mb_seg[4] = { 0, 0, 0, 0 };
-#if CONFIG_SEGMENTATION
-int segment_modes_intra[MAX_MB_SEGMENTS] = { 0, 0, 0, 0 };
-int segment_modes_inter[MAX_MB_SEGMENTS] = { 0, 0, 0, 0 };
-#endif
 #endif
 
 
@@ -945,7 +941,7 @@ static void pack_inter_mode_mvs(VP8_COMP *const cpi)
     MACROBLOCKD *xd = &cpi->mb.e_mbd;
 #if CONFIG_SEGMENTATION
     int i;
-    int sum;
+    int pred_context;
     int index = 0;
 #endif
     const int *const rfct = cpi->count_mb_ref_frame_usage;
@@ -1042,37 +1038,30 @@ static void pack_inter_mode_mvs(VP8_COMP *const cpi)
             active_section = 9;
 #endif
 
-#ifdef MODE_STATS
-#if CONFIG_SEGMENTATION
-            segment_modes_inter[segment_id]++;
-#endif
-#endif
             if (cpi->mb.e_mbd.update_mb_segmentation_map)
             {
 #if CONFIG_SEGMENTATION
                 if (xd->temporal_update)
                 {
-                    sum = 0;
+                    pred_context = 0;
                     if (mb_col != 0)
-                        sum +=  (m-1)->mbmi.segment_flag;
+                        pred_context +=  (m-1)->mbmi.segment_flag;
                     if (mb_row != 0)
-                        sum += (m-pc->mb_cols)->mbmi.segment_flag;
+                        pred_context += (m-pc->mb_cols)->mbmi.segment_flag;
 
                     if (m->mbmi.segment_flag == 0)
                     {
-                        vp8_write(w,0,xd->mb_segment_tree_probs[3+sum]);
+                        vp8_write(w,0,xd->mb_segment_pred_probs[pred_context]);
                     }
                     else
                     {
-                        vp8_write(w,1,xd->mb_segment_tree_probs[3+sum]);
+                        vp8_write(w,1,xd->mb_segment_pred_probs[pred_context]);
                         write_mb_segid(w, mi, &cpi->mb.e_mbd);
-                        cpi->segmentation_map[index] = segment_id;
                     }
                 }
                 else
                 {
                     write_mb_segid(w, mi, &cpi->mb.e_mbd);
-                    cpi->segmentation_map[index] = segment_id;
                 }
                 index++;
 #else
@@ -1268,21 +1257,12 @@ static void write_kfmodes(VP8_COMP *cpi)
             const int ym = m->mbmi.mode;
             int segment_id = m->mbmi.segment_id;
 
-#ifdef MODE_STATS
-#if CONFIG_SEGMENTATION
-            segment_modes_intra[segment_id]++;
-#endif
-#endif
-
             if (cpi->mb.e_mbd.update_mb_segmentation_map)
             {
 #if CONFIG_SEGMENTATION
-                write_mb_segid(bc, &m->mbmi, &cpi->mb.e_mbd);
-                cpi->segmentation_map[index] = segment_id;
                 index++;
-#else
-                write_mb_segid(bc, &m->mbmi, &cpi->mb.e_mbd);
 #endif
+                write_mb_segid(bc, &m->mbmi, &cpi->mb.e_mbd);
             }
 
 //#if CONFIG_SEGFEATURES
@@ -2050,12 +2030,9 @@ void vp8_pack_bitstream(VP8_COMP *cpi, unsigned char *dest, unsigned long *size)
 
         if (xd->update_mb_segmentation_map)
         {
- #if CONFIG_SEGMENTATION
-            // Write the probs used to decode the segment id for each macro block.
-            for (i = 0; i < MB_FEATURE_TREE_PROBS+3; i++)
-#else
+            // Send the tree probabilities used to decode unpredicted
+            // macro-block segments
             for (i = 0; i < MB_FEATURE_TREE_PROBS; i++)
-#endif
             {
                 int Data = xd->mb_segment_tree_probs[i];
 
@@ -2067,6 +2044,25 @@ void vp8_pack_bitstream(VP8_COMP *cpi, unsigned char *dest, unsigned long *size)
                 else
                     vp8_write_bit(bc, 0);
             }
+#if CONFIG_SEGMENTATION
+            // If predictive coding of segment map is enabled send the
+            // prediction probabilities.
+            if ( xd->temporal_update )
+            {
+                for (i = 0; i < SEGMENT_PREDICTION_PROBS; i++)
+                {
+                    int Data = xd->mb_segment_pred_probs[i];
+
+                    if (Data != 255)
+                    {
+                        vp8_write_bit(bc, 1);
+                        vp8_write_literal(bc, Data, 8);
+                    }
+                    else
+                        vp8_write_bit(bc, 0);
+                }
+            }
+#endif
         }
     }
 
diff --git a/vp8/encoder/encodeframe.c b/vp8/encoder/encodeframe.c
index b48a5f687..4cde41df4 100644
--- a/vp8/encoder/encodeframe.c
+++ b/vp8/encoder/encodeframe.c
@@ -578,9 +578,7 @@ void encode_mb_row(VP8_COMP *cpi,
     int recon_y_stride = cm->yv12_fb[ref_fb_idx].y_stride;
     int recon_uv_stride = cm->yv12_fb[ref_fb_idx].uv_stride;
     int map_index = (mb_row * cpi->common.mb_cols);
-#if CONFIG_SEGMENTATION
-    int sum;
-#endif
+
 #if CONFIG_MULTITHREAD
     const int nsync = cpi->mt_sync_range;
     const int rightmost_col = cm->mb_cols - 1;
@@ -768,43 +766,7 @@ void encode_mb_row(VP8_COMP *cpi,
         recon_yoffset += 16;
         recon_uvoffset += 8;
 
-#if CONFIG_SEGMENTATION
-       //cpi->segmentation_map[mb_row * cm->mb_cols + mb_col] =  xd->mbmi.segment_id;
-        if (cm->frame_type == KEY_FRAME)
-        {
-            segment_counts[xd->mode_info_context->mbmi.segment_id]++;
-        }
-        else
-        {
-            sum = 0;
-            if (mb_col != 0)
-                sum += (xd->mode_info_context-1)->mbmi.segment_flag;
-            if (mb_row != 0)
-                sum += (xd->mode_info_context-cm->mb_cols)->mbmi.segment_flag;
-
-            if ( xd->mode_info_context->mbmi.segment_id ==
-                 cpi->last_segmentation_map[(mb_row*cm->mb_cols) + mb_col] )
-            {
-                xd->mode_info_context->mbmi.segment_flag = 0;
-            }
-            else
-                xd->mode_info_context->mbmi.segment_flag = 1;
-
-            if (xd->mode_info_context->mbmi.segment_flag == 0)
-            {
-                segment_counts[SEEK_SAMEID + sum]++;
-                segment_counts[10]++;
-            }
-            else
-            {
-                segment_counts[SEEK_DIFFID + sum]++;
-                segment_counts[11]++;
-                //calculate individual segment ids
-                segment_counts[xd->mode_info_context->mbmi.segment_id] ++;
-            }
-        }
-        segment_counts[SEEK_SEGID + xd->mode_info_context->mbmi.segment_id] ++;
-#else
+#if !CONFIG_SEGMENTATION
         segment_counts[xd->mode_info_context->mbmi.segment_id] ++;
 #endif
         // skip to next mb
@@ -948,12 +910,7 @@ void vp8_encode_frame(VP8_COMP *cpi)
     MACROBLOCKD *const xd = & x->e_mbd;
 
     TOKENEXTRA *tp = cpi->tok;
-
-#if CONFIG_SEGMENTATION
-    int segment_counts[MAX_MB_SEGMENTS + SEEK_SEGID];
-#else
     int segment_counts[MAX_MB_SEGMENTS];
-#endif
     int totalrate;
 
 
@@ -1144,7 +1101,7 @@ void vp8_encode_frame(VP8_COMP *cpi)
 
 #if CONFIG_SEGMENTATION
         // Select the coding strategy for the segment map (temporal or spatial)
-        choose_segmap_coding_method( cpi, segment_counts );
+        choose_segmap_coding_method( cpi );
 #else
         tot_count = segment_counts[0] + segment_counts[1] + segment_counts[2] + segment_counts[3];
         count1 = segment_counts[0] + segment_counts[1];
@@ -1159,17 +1116,14 @@ void vp8_encode_frame(VP8_COMP *cpi)
         if (count2 > 0)
             xd->mb_segment_tree_probs[2] = (segment_counts[2] * 255) /count2;
 
-#endif
         // Zero probabilities not allowed
-#if CONFIG_SEGMENTATION
-            for (i = 0; i < MB_FEATURE_TREE_PROBS+3; i++)
-#else
-            for (i = 0; i < MB_FEATURE_TREE_PROBS; i++)
+        for (i = 0; i < MB_FEATURE_TREE_PROBS; i++)
+
+        {
+            if (xd->mb_segment_tree_probs[i] == 0)
+                xd->mb_segment_tree_probs[i] = 1;
+        }
 #endif
-            {
-                if (xd->mb_segment_tree_probs[i] == 0)
-                    xd->mb_segment_tree_probs[i] = 1;
-            }
     }
 
     // 256 rate units to the bit
diff --git a/vp8/encoder/onyx_if.c b/vp8/encoder/onyx_if.c
index 8dbb0059b..991b469d2 100644
--- a/vp8/encoder/onyx_if.c
+++ b/vp8/encoder/onyx_if.c
@@ -152,10 +152,6 @@ extern int b_modes[10]  ;
 extern int inter_y_modes[10] ;
 extern int inter_uv_modes[4] ;
 extern unsigned int inter_b_modes[15];
-#if CONFIG_SEGMENTATION
-extern int segment_modes_intra[MAX_MB_SEGMENTS];
-extern int segment_modes_inter[MAX_MB_SEGMENTS];
-#endif
 #endif
 
 extern void (*vp8_short_fdct4x4)(short *input, short *output, int pitch);
@@ -2529,9 +2525,6 @@ void vp8_remove_compressor(VP8_PTR *ptr)
                 fprintf(f, "\n");
 
             }
-#if CONFIG_SEGMENTATION
-            fprintf(f, "Segments:%8d, %8d, %8d, %8d\n", segment_modes_intra[0], segment_modes_intra[1], segment_modes_intra[2], segment_modes_intra[3]);
-#endif
 
             fprintf(f, "Modes in Inter Frames:\n");
             fprintf(f, "Y: %8d, %8d, %8d, %8d, %8d, %8d, %8d, %8d, %8d, %8d\n",
@@ -2550,10 +2543,6 @@ void vp8_remove_compressor(VP8_PTR *ptr)
             }
             fprintf(f, "P:%8d, %8d, %8d, %8d\n", count_mb_seg[0], count_mb_seg[1], count_mb_seg[2], count_mb_seg[3]);
             fprintf(f, "PB:%8d, %8d, %8d, %8d\n", inter_b_modes[LEFT4X4], inter_b_modes[ABOVE4X4], inter_b_modes[ZERO4X4], inter_b_modes[NEW4X4]);
-
-#if CONFIG_SEGMENTATION
-            fprintf(f, "Segments:%8d, %8d, %8d, %8d\n", segment_modes_inter[0], segment_modes_inter[1], segment_modes_inter[2], segment_modes_inter[3]);
-#endif
             fclose(f);
         }
 #endif
diff --git a/vp8/encoder/onyx_int.h b/vp8/encoder/onyx_int.h
index 8bcf65acc..c35dfed08 100644
--- a/vp8/encoder/onyx_int.h
+++ b/vp8/encoder/onyx_int.h
@@ -56,12 +56,6 @@
 #define VP8_TEMPORAL_ALT_REF 1
 #endif
 
-#if CONFIG_SEGMENTATION
-#define SEEK_SEGID 12
-#define SEEK_SAMEID 4
-#define SEEK_DIFFID 7
-#endif
-
 typedef struct
 {
     int kf_indicated;
@@ -226,11 +220,7 @@ typedef struct
 typedef struct
 {
     MACROBLOCK  mb;
-#if CONFIG_SEGMENTATION
-    int segment_counts[MAX_MB_SEGMENTS + 8];
-#else
     int segment_counts[MAX_MB_SEGMENTS];
-#endif
     int totalrate;
 } MB_ROW_COMP;
 
diff --git a/vp8/encoder/segmentation.c b/vp8/encoder/segmentation.c
index 16dd74f04..2eb0ce408 100644
--- a/vp8/encoder/segmentation.c
+++ b/vp8/encoder/segmentation.c
@@ -116,123 +116,209 @@ void vp8_set_segment_data(VP8_PTR ptr,
 }
 
 #if CONFIG_SEGMENTATION
-void choose_segmap_coding_method( VP8_COMP *cpi,
-                                  int * segment_counts )
+// Based on set of segment counts calculate a probability tree
+void calc_segtree_probs( MACROBLOCKD * xd,
+                         int * segcounts,
+                         vp8_prob * segment_tree_probs )
 {
-    VP8_COMMON *const cm = & cpi->common;
-    MACROBLOCKD *const xd = & cpi->mb.e_mbd;
-
+    int count1,count2;
     int tot_count;
     int i;
-    int count1,count2,count3,count4;
-    int prob[3];
-    int new_cost, original_cost;
 
-    // Select the coding strategy for the segment map (temporal or spatial)
-    tot_count = segment_counts[12] + segment_counts[13] +
-                segment_counts[14] + segment_counts[15];
-    count1 = segment_counts[12] + segment_counts[13];
-    count2 = segment_counts[14] + segment_counts[15];
+    // Blank the strtucture to start with
+    vpx_memset(segment_tree_probs, 0, sizeof(segment_tree_probs));
+
+    // Total count for all segments
+    count1 = segcounts[0] + segcounts[1];
+    count2 = segcounts[2] + segcounts[3];
+    tot_count = count1 + count2;
 
+    // Work out probabilities of each segment
     if (tot_count)
-        prob[0] = (count1 * 255) / tot_count;
+        segment_tree_probs[0] = (count1 * 255) / tot_count;
+    if (count1 > 0)
+        segment_tree_probs[1] = (segcounts[0] * 255) / count1;
+    if (count2 > 0)
+        segment_tree_probs[2] = (segcounts[2] * 255) / count2;
+
+    // Clamp probabilities to minimum allowed value
+    for (i = 0; i < MB_FEATURE_TREE_PROBS; i++)
+    {
+        if (segment_tree_probs[i] == 0)
+            segment_tree_probs[i] = 1;
+    }
+}
+
+// Based on set of segment counts and probabilities calculate a cost estimate
+int cost_segmap( MACROBLOCKD * xd,
+                 int * segcounts,
+                 vp8_prob * probs )
+{
+    int cost;
+    int count1,count2;
+
+    // Cost the top node of the tree
+    count1 = segcounts[0] + segcounts[1];
+    count2 = segcounts[2] + segcounts[3];
+    cost = count1 * vp8_cost_zero(probs[0]) +
+           count2 * vp8_cost_one(probs[0]);
 
+    // Now add the cost of each individual segment branch
     if (count1 > 0)
-        prob[1] = (segment_counts[12] * 255) /count1;
+        cost += segcounts[0] * vp8_cost_zero(probs[1]) +
+                segcounts[1] * vp8_cost_one(probs[1]);
 
     if (count2 > 0)
-        prob[2] = (segment_counts[14] * 255) /count2;
+        cost += segcounts[2] * vp8_cost_zero(probs[2]) +
+                segcounts[3] * vp8_cost_one(probs[2]) ;
 
-    if (cm->frame_type != KEY_FRAME)
-    {
-        tot_count = segment_counts[4] + segment_counts[7];
-        if (tot_count)
-            xd->mb_segment_tree_probs[3] = (segment_counts[4] * 255)/tot_count;
+    return cost;
 
-        tot_count = segment_counts[5] + segment_counts[8];
-        if (tot_count)
-            xd->mb_segment_tree_probs[4] = (segment_counts[5] * 255)/tot_count;
+}
 
-        tot_count = segment_counts[6] + segment_counts[9];
-        if (tot_count)
-            xd->mb_segment_tree_probs[5] = (segment_counts[6] * 255)/tot_count;
-    }
+void choose_segmap_coding_method( VP8_COMP *cpi )
+{
+    VP8_COMMON *const cm = & cpi->common;
+    MACROBLOCKD *const xd = & cpi->mb.e_mbd;
+
+    int i;
+    int tot_count;
+    int no_pred_cost;
+    int t_pred_cost = INT_MAX;
+    int pred_context;
 
-    tot_count = segment_counts[0] + segment_counts[1] +
-                segment_counts[2] + segment_counts[3];
-    count3 = segment_counts[0] + segment_counts[1];
-    count4 = segment_counts[2] + segment_counts[3];
+    int mb_row, mb_col;
+    int segmap_index = 0;
+    unsigned char segment_id;
 
-    if (tot_count)
-        xd->mb_segment_tree_probs[0] = (count3 * 255) / tot_count;
+    int temporal_predictor_count[SEGMENT_PREDICTION_PROBS][2];
+    int no_pred_segcounts[MAX_MB_SEGMENTS];
+    int t_unpred_seg_counts[MAX_MB_SEGMENTS];
 
-    if (count3 > 0)
-        xd->mb_segment_tree_probs[1] = (segment_counts[0] * 255) /count3;
+    vp8_prob no_pred_tree[MB_FEATURE_TREE_PROBS];
+    vp8_prob t_pred_tree[MB_FEATURE_TREE_PROBS];
+    vp8_prob t_nopred_prob[SEGMENT_PREDICTION_PROBS];
 
-    if (count4 > 0)
-        xd->mb_segment_tree_probs[2] = (segment_counts[2] * 255) /count4;
+    vpx_memset(no_pred_segcounts, 0, sizeof(no_pred_segcounts));
+    vpx_memset(t_unpred_seg_counts, 0, sizeof(t_unpred_seg_counts));
+    vpx_memset(temporal_predictor_count, 0, sizeof(temporal_predictor_count));
 
-    for (i = 0; i < MB_FEATURE_TREE_PROBS+3; i++)
+    // First of all generate stats regarding how well the last segment map
+    // predicts this one
+
+    // Initialize macroblod decoder mode info context for to the first mb
+    // in the frame
+    xd->mode_info_context = cm->mi;
+
+    for (mb_row = 0; mb_row < cm->mb_rows; mb_row++)
     {
-        if (xd->mb_segment_tree_probs[i] == 0)
-            xd->mb_segment_tree_probs[i] = 1;
-    }
+        for (mb_col = 0; mb_col < cm->mb_cols; mb_col++)
+        {
+            segment_id = xd->mode_info_context->mbmi.segment_id;
 
-    original_cost = count1 * vp8_cost_zero(prob[0]) +
-                    count2 * vp8_cost_one(prob[0]);
+            // Count the number of hits on each segment with no prediction
+            no_pred_segcounts[segment_id]++;
 
-    if (count1 > 0)
-        original_cost += segment_counts[12] * vp8_cost_zero(prob[1]) +
-                         segment_counts[13] * vp8_cost_one(prob[1]);
+            // Temporal prediction not allowed on key frames
+            if (cm->frame_type != KEY_FRAME)
+            {
+                // Get temporal prediction context
+                pred_context = 0;
+                if (mb_col != 0)
+                    pred_context +=
+                        (xd->mode_info_context-1)->mbmi.segment_flag;
+                if (mb_row != 0)
+                    pred_context +=
+                        (xd->mode_info_context-cm->mb_cols)->mbmi.segment_flag;
+
+                // Test to see if the last frame segment id at the same
+                // locationcorrectly predicts the segment_id for this MB.
+                // Update the prediction flag and count as appropriate;
+                if ( segment_id == cpi->last_segmentation_map[segmap_index] )
+                {
+                    //xd->mode_info_context->mbmi.segment_predicted = 1;
+                    xd->mode_info_context->mbmi.segment_flag = 0;
+                    temporal_predictor_count[pred_context][0]++;
+                }
+                else
+                {
+                    //xd->mode_info_context->mbmi.segment_predicted = 0;
+                    xd->mode_info_context->mbmi.segment_flag = 1;
+                    temporal_predictor_count[pred_context][1]++;
 
-    if (count2 > 0)
-        original_cost += segment_counts[14] * vp8_cost_zero(prob[2]) +
-                         segment_counts[15] * vp8_cost_one(prob[2]) ;
+                    // Update the "undpredicted" segment count
+                    t_unpred_seg_counts[segment_id]++;
+                }
+            }
 
-    new_cost = 0;
+            // Step on to the next mb
+            xd->mode_info_context++;
 
+            // Step on to the next entry in the segment maps
+            segmap_index++;
+        }
+
+        // this is to account for the border in mode_info_context
+        xd->mode_info_context++;
+    }
+
+    // Work out probability tree for coding segments without prediction
+    // and the cost.
+    calc_segtree_probs( xd, no_pred_segcounts, no_pred_tree );
+    no_pred_cost = cost_segmap( xd, no_pred_segcounts, no_pred_tree );
+
+    // Key frames cannot use temporal prediction
     if (cm->frame_type != KEY_FRAME)
     {
-        new_cost = segment_counts[4] *
-                        vp8_cost_zero(xd->mb_segment_tree_probs[3]) +
-                   segment_counts[7] *
-                        vp8_cost_one(xd->mb_segment_tree_probs[3]);
-
-        new_cost += segment_counts[5] *
-                        vp8_cost_zero(xd->mb_segment_tree_probs[4]) +
-                    segment_counts[8] *
-                        vp8_cost_one(xd->mb_segment_tree_probs[4]);
-
-        new_cost += segment_counts[6] *
-                        vp8_cost_zero(xd->mb_segment_tree_probs[5]) +
-                    segment_counts[9] *
-                        vp8_cost_one (xd->mb_segment_tree_probs[5]);
-    }
+        // Work out probability tree for coding those segments not
+        // predicted using the temporal method and the cost.
+        calc_segtree_probs( xd, t_unpred_seg_counts, t_pred_tree );
+        t_pred_cost = cost_segmap( xd, t_unpred_seg_counts, t_pred_tree );
 
-    if (tot_count > 0)
-        new_cost += count3 * vp8_cost_zero(xd->mb_segment_tree_probs[0]) +
-                    count4 * vp8_cost_one(xd->mb_segment_tree_probs[0]);
+        // Add in the cost of the signalling for each prediction context
+        for ( i = 0; i < SEGMENT_PREDICTION_PROBS; i++ )
+        {
+            tot_count = temporal_predictor_count[i][0] +
+                        temporal_predictor_count[i][1];
 
-    if (count3 > 0)
-        new_cost += segment_counts[0] *
-                        vp8_cost_zero(xd->mb_segment_tree_probs[1]) +
-                    segment_counts[1] *
-                        vp8_cost_one(xd->mb_segment_tree_probs[1]);
+            // Work out the context probabilities for the segment
+            // prediction flag
+            if ( tot_count )
+            {
+                t_nopred_prob[i] = ( temporal_predictor_count[i][0] * 255 ) /
+                                   tot_count;
 
-    if (count4 > 0)
-        new_cost += segment_counts[2] *
-                        vp8_cost_zero(xd->mb_segment_tree_probs[2]) +
-                    segment_counts[3] *
-                        vp8_cost_one(xd->mb_segment_tree_probs[2]) ;
+                // Clamp to minimum allowed value
+                if ( t_nopred_prob[i] < 1 )
+                    t_nopred_prob[i] = 1;
+            }
+            else
+                t_nopred_prob[i] = 1;
+
+            // Add in the predictor signaling cost
+            t_pred_cost += ( temporal_predictor_count[i][0] *
+                               vp8_cost_zero(t_nopred_prob[i]) ) +
+                           ( temporal_predictor_count[i][1] *
+                               vp8_cost_one(t_nopred_prob[i]) );
+        }
+    }
 
-    if (new_cost < original_cost)
-        xd->temporal_update = 1;
+    // Now choose which coding method to use.
+    if ( t_pred_cost < no_pred_cost )
+    {
+         xd->temporal_update = 1;
+         vpx_memcpy( xd->mb_segment_tree_probs,
+                     t_pred_tree, sizeof(t_pred_tree) );
+         vpx_memcpy( &xd->mb_segment_pred_probs,
+                     t_nopred_prob, sizeof(t_nopred_prob) );
+    }
     else
     {
-        xd->temporal_update = 0;
-        xd->mb_segment_tree_probs[0] = prob[0];
-        xd->mb_segment_tree_probs[1] = prob[1];
-        xd->mb_segment_tree_probs[2] = prob[2];
+         xd->temporal_update = 0;
+         vpx_memcpy( xd->mb_segment_tree_probs,
+                     no_pred_tree, sizeof(no_pred_tree) );
+         //vpx_memcpy( &xd->mb_segment_pred_probs,
+         //            t_nopred_prob, sizeof(t_nopred_prob) );
     }
 }
 #endif
diff --git a/vp8/encoder/segmentation.h b/vp8/encoder/segmentation.h
index b6394fd09..f3f2b9c6f 100644
--- a/vp8/encoder/segmentation.h
+++ b/vp8/encoder/segmentation.h
@@ -39,8 +39,7 @@ extern void vp8_set_segmentation_map(VP8_PTR ptr, unsigned char *segmentation_ma
 extern void vp8_set_segment_data(VP8_PTR ptr, signed char *feature_data, unsigned char abs_delta);
 
 #if CONFIG_SEGMENTATION
-extern void choose_segmap_coding_method( VP8_COMP *cpi,
-                                         int * segment_counts );
+extern void choose_segmap_coding_method( VP8_COMP *cpi );
 #endif
 
 #endif /* __INC_SEGMENTATION_H__ */