hevc: Use parsed VUI colorimetry in avcodec

Also limit the parsed VUI values to known valid ones.

Signed-off-by: Vittorio Giovara <vittorio.giovara@gmail.com>

2 files changed, 24 insertions, 0 deletions

diff --git a/libavcodec/hevc.c b/libavcodec/hevc.c
index 276f3fdfbe..e01db16868 100644
--- a/libavcodec/hevc.c
+++ b/libavcodec/hevc.c
@@ -401,6 +401,22 @@ static int set_sps(HEVCContext *s, const HEVCSPS *sps)
     s->avctx->sample_aspect_ratio = sps->vui.sar;
     s->avctx->has_b_frames        = sps->temporal_layer[sps->max_sub_layers - 1].num_reorder_pics;
 
+    if (sps->vui.video_signal_type_present_flag)
+        s->avctx->color_range = sps->vui.video_full_range_flag ? AVCOL_RANGE_JPEG
+                                                               : AVCOL_RANGE_MPEG;
+    else
+        s->avctx->color_range = AVCOL_RANGE_MPEG;
+
+    if (sps->vui.colour_description_present_flag) {
+        s->avctx->color_primaries = sps->vui.colour_primaries;
+        s->avctx->color_trc       = sps->vui.transfer_characteristic;
+        s->avctx->colorspace      = sps->vui.matrix_coeffs;
+    } else {
+        s->avctx->color_primaries = AVCOL_PRI_UNSPECIFIED;
+        s->avctx->color_trc       = AVCOL_TRC_UNSPECIFIED;
+        s->avctx->colorspace      = AVCOL_SPC_UNSPECIFIED;
+    }
+
     ff_hevc_pred_init(&s->hpc,     sps->bit_depth);
     ff_hevc_dsp_init (&s->hevcdsp, sps->bit_depth);
     ff_videodsp_init (&s->vdsp,    sps->bit_depth);
diff --git a/libavcodec/hevc_ps.c b/libavcodec/hevc_ps.c
index 15f73be9bf..9c0d0ee753 100644
--- a/libavcodec/hevc_ps.c
+++ b/libavcodec/hevc_ps.c
@@ -453,6 +453,14 @@ static void decode_vui(HEVCContext *s, HEVCSPS *sps)
             vui->colour_primaries        = get_bits(gb, 8);
             vui->transfer_characteristic = get_bits(gb, 8);
             vui->matrix_coeffs           = get_bits(gb, 8);
+
+            // Set invalid values to "unspecified"
+            if (vui->colour_primaries >= AVCOL_PRI_NB)
+                vui->colour_primaries = AVCOL_PRI_UNSPECIFIED;
+            if (vui->transfer_characteristic >= AVCOL_TRC_NB)
+                vui->transfer_characteristic = AVCOL_TRC_UNSPECIFIED;
+            if (vui->matrix_coeffs >= AVCOL_SPC_NB)
+                vui->matrix_coeffs = AVCOL_SPC_UNSPECIFIED;
         }
     }