X-Git-Url: https://git.nihav.org/?a=blobdiff_plain;f=nihav-itu%2Fsrc%2Fcodecs%2Fh264%2Fslice.rs;h=ad1d244afe5cfd84b29fbe0060697971491c0558;hb=HEAD;hp=a4ce4bf0271c92817151ef4b1039bbb948e260b0;hpb=696e4e20bf7167121352f247893370cb83213d5c;p=nihav.git

diff --git a/nihav-itu/src/codecs/h264/slice.rs b/nihav-itu/src/codecs/h264/slice.rs
index a4ce4bf..ad1d244 100644
--- a/nihav-itu/src/codecs/h264/slice.rs
+++ b/nihav-itu/src/codecs/h264/slice.rs
@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use nihav_core::codecs::{DecoderResult, DecoderError};
 use nihav_core::frame::FrameType;
 use nihav_core::io::bitreader::*;
@@ -18,23 +20,14 @@ pub enum SliceType {
 
 impl SliceType {
     pub fn is_intra(self) -> bool {
-        match self {
-            SliceType::I | SliceType::SI => true,
-            _ => false,
-        }
+        matches!(self, SliceType::I | SliceType::SI)
     }
     pub fn is_p(self) -> bool {
-        match self {
-            SliceType::P | SliceType::SP => true,
-            _ => false,
-        }
+        matches!(self, SliceType::P | SliceType::SP)
     }
     pub fn is_b(self) -> bool { self == SliceType::B }
     pub fn is_s(self) -> bool {
-        match self {
-            SliceType::SI | SliceType::SP => true,
-            _ => false,
-        }
+        matches!(self, SliceType::SI | SliceType::SP)
     }
     pub fn to_frame_type(self) -> FrameType {
         match self {
@@ -50,14 +43,22 @@ const SLICE_TYPES: [SliceType; 10] = [
     SliceType::P, SliceType::B, SliceType::I, SliceType::SP, SliceType::SI,
 ];
 
-#[derive(Clone,Copy)]
+#[derive(Clone,Copy,Default)]
 pub struct WeightInfo {
     pub luma_weighted:                      bool,
     pub luma_weight:                        i8,
     pub luma_offset:                        i8,
+    pub luma_shift:                         u8,
     pub chroma_weighted:                    bool,
     pub chroma_weight:                      [i8; 2],
     pub chroma_offset:                      [i8; 2],
+    pub chroma_shift:                       u8,
+}
+
+impl WeightInfo {
+    pub fn is_weighted(&self) -> bool {
+        self.luma_weighted || self.chroma_weighted
+    }
 }
 
 #[derive(Clone,Copy)]
@@ -117,6 +118,35 @@ pub struct SliceHeader {
     pub slice_group_change_cycle:           u32,
 }
 
+pub const DEF_WEIGHT_INFO: WeightInfo = WeightInfo {
+    luma_weighted:                      false,
+    luma_weight:                        0,
+    luma_offset:                        0,
+    luma_shift:                         0,
+    chroma_weighted:                    false,
+    chroma_weight:                      [0; 2],
+    chroma_offset:                      [0; 2],
+    chroma_shift:                       0,
+};
+
+impl SliceHeader {
+    pub fn get_weight(&self, list_id: u8, idx: usize) -> WeightInfo {
+        if list_id == 0 {
+            if idx < self.num_ref_idx_l0_active {
+                self.weights_l0[idx]
+            } else {
+                DEF_WEIGHT_INFO
+            }
+        } else {
+            if idx < self.num_ref_idx_l1_active {
+                self.weights_l1[idx]
+            } else {
+                DEF_WEIGHT_INFO
+            }
+        }
+    }
+}
+
 pub fn parse_slice_header_minimal(br: &mut BitReader) -> DecoderResult<(usize, SliceType)> {
     let first_mb_in_slice                           = br.read_ue()? as usize;
     let stype                                       = br.read_ue_lim(SLICE_TYPES.len() as u32 - 1)?;
@@ -124,7 +154,9 @@ pub fn parse_slice_header_minimal(br: &mut BitReader) -> DecoderResult<(usize, S
     Ok((first_mb_in_slice, slice_type))
 }
 
-pub fn parse_slice_header(br: &mut BitReader, sps_arr: &[SeqParameterSet], pps_arr: &[PicParameterSet], is_idr: bool, nal_ref_idc: u8) -> DecoderResult<SliceHeader> {
+#[allow(clippy::cognitive_complexity)]
+#[allow(clippy::manual_range_contains)]
+pub fn parse_slice_header(br: &mut BitReader, sps_arr: &[Arc<SeqParameterSet>], pps_arr: &[Arc<PicParameterSet>], is_idr: bool, nal_ref_idc: u8) -> DecoderResult<SliceHeader> {
     let mut hdr: SliceHeader = unsafe { std::mem::zeroed() };
 
     hdr.first_mb_in_slice                           = br.read_ue()? as usize;
@@ -198,6 +230,15 @@ pub fn parse_slice_header(br: &mut BitReader, sps_arr: &[SeqParameterSet], pps_a
     if (pps.weighted_pred && hdr.slice_type.is_p()) ||
         (pps.weighted_bipred_idc == 1 && hdr.slice_type.is_b()) {
         parse_pred_weight_table(&mut hdr, br)?;
+    } else {
+        for weight in hdr.weights_l0[..hdr.num_ref_idx_l0_active].iter_mut() {
+            weight.luma_weighted = false;
+            weight.chroma_weighted = false;
+        }
+        for weight in hdr.weights_l1[..hdr.num_ref_idx_l1_active].iter_mut() {
+            weight.luma_weighted = false;
+            weight.chroma_weighted = false;
+        }
     }
     if nal_ref_idc != 0 {
         if is_idr {
@@ -345,6 +386,8 @@ fn parse_pred_weight_table(hdr: &mut SliceHeader, br: &mut BitReader) -> Decoder
             validate!(offset >= -128 && offset <= 127);
             weight.luma_offset = offset as i8;
         }
+        weight.luma_shift = hdr.luma_log2_weight_denom;
+
         weight.chroma_weighted                      = br.read_bool()?;
         if weight.chroma_weighted {
             for i in 0..2 {
@@ -356,6 +399,7 @@ fn parse_pred_weight_table(hdr: &mut SliceHeader, br: &mut BitReader) -> Decoder
                 weight.chroma_offset[i] = offset as i8;
             }
         }
+        weight.chroma_shift = hdr.chroma_log2_weight_denom;
     }
     for weight in hdr.weights_l1[..hdr.num_ref_idx_l1_active].iter_mut() {
         weight.luma_weighted                        = br.read_bool()?;
@@ -367,6 +411,8 @@ fn parse_pred_weight_table(hdr: &mut SliceHeader, br: &mut BitReader) -> Decoder
             validate!(offset >= -128 && offset <= 127);
             weight.luma_offset = offset as i8;
         }
+        weight.luma_shift = hdr.luma_log2_weight_denom;
+
         weight.chroma_weighted                      = br.read_bool()?;
         if weight.chroma_weighted {
             for i in 0..2 {
@@ -378,6 +424,7 @@ fn parse_pred_weight_table(hdr: &mut SliceHeader, br: &mut BitReader) -> Decoder
                 weight.chroma_offset[i] = offset as i8;
             }
         }
+        weight.chroma_shift = hdr.chroma_log2_weight_denom;
     }
 
     Ok(())