1use crate::api::internal::InterConfig;
2use crate::config::EncoderConfig;
3use crate::context::{BlockOffset, FrameBlocks, TileBlockOffset};
4use crate::cpu_features::CpuFeatureLevel;
5use crate::dist::get_satd;
6use crate::encoder::{
7 FrameInvariants, FrameState, Sequence, IMPORTANCE_BLOCK_SIZE,
8};
9use crate::frame::{AsRegion, PlaneOffset};
10use crate::me::{estimate_tile_motion, RefMEStats};
11use crate::partition::{get_intra_edges, BlockSize};
12use crate::predict::{IntraParam, PredictionMode};
13use crate::tiling::{Area, PlaneRegion, TileRect};
14use crate::transform::TxSize;
15use crate::util::Aligned;
16use crate::Pixel;
17use rayon::iter::*;
18use std::sync::Arc;
19use v_frame::frame::Frame;
20use v_frame::pixel::CastFromPrimitive;
21use v_frame::plane::Plane;
22
23pub(crate) const IMP_BLOCK_MV_UNITS_PER_PIXEL: i64 = 8;
24pub(crate) const IMP_BLOCK_SIZE_IN_MV_UNITS: i64 =
25 IMPORTANCE_BLOCK_SIZE as i64 * IMP_BLOCK_MV_UNITS_PER_PIXEL;
26pub(crate) const IMP_BLOCK_AREA_IN_MV_UNITS: i64 =
27 IMP_BLOCK_SIZE_IN_MV_UNITS * IMP_BLOCK_SIZE_IN_MV_UNITS;
28
29#[profiling::function]
30pub(crate) fn estimate_intra_costs<T: Pixel>(
31 temp_plane: &mut Plane<T>, frame: &Frame<T>, bit_depth: usize,
32 cpu_feature_level: CpuFeatureLevel,
33) -> Box<[u32]> {
34 let plane = &frame.planes[0];
35 let plane_after_prediction = temp_plane;
36
37 let bsize = BlockSize::from_width_and_height(
38 IMPORTANCE_BLOCK_SIZE,
39 IMPORTANCE_BLOCK_SIZE,
40 );
41 let tx_size = bsize.tx_size();
42
43 let h_in_imp_b = plane.cfg.height / IMPORTANCE_BLOCK_SIZE;
44 let w_in_imp_b = plane.cfg.width / IMPORTANCE_BLOCK_SIZE;
45 let mut intra_costs = Vec::with_capacity(h_in_imp_b * w_in_imp_b);
46
47 for y in 0..h_in_imp_b {
48 for x in 0..w_in_imp_b {
49 let plane_org = plane.region(Area::Rect {
50 x: (x * IMPORTANCE_BLOCK_SIZE) as isize,
51 y: (y * IMPORTANCE_BLOCK_SIZE) as isize,
52 width: IMPORTANCE_BLOCK_SIZE,
53 height: IMPORTANCE_BLOCK_SIZE,
54 });
55
56 // TODO: other intra prediction modes.
57 let mut edge_buf = Aligned::uninit_array();
58 let edge_buf = get_intra_edges(
59 &mut edge_buf,
60 &plane.as_region(),
61 TileBlockOffset(BlockOffset { x, y }),
62 0,
63 0,
64 bsize,
65 PlaneOffset {
66 x: (x * IMPORTANCE_BLOCK_SIZE) as isize,
67 y: (y * IMPORTANCE_BLOCK_SIZE) as isize,
68 },
69 TxSize::TX_8X8,
70 bit_depth,
71 Some(PredictionMode::DC_PRED),
72 false,
73 IntraParam::None,
74 );
75
76 let mut plane_after_prediction_region = plane_after_prediction
77 .region_mut(Area::Rect {
78 x: (x * IMPORTANCE_BLOCK_SIZE) as isize,
79 y: (y * IMPORTANCE_BLOCK_SIZE) as isize,
80 width: IMPORTANCE_BLOCK_SIZE,
81 height: IMPORTANCE_BLOCK_SIZE,
82 });
83
84 PredictionMode::DC_PRED.predict_intra(
85 TileRect {
86 x: x * IMPORTANCE_BLOCK_SIZE,
87 y: y * IMPORTANCE_BLOCK_SIZE,
88 width: IMPORTANCE_BLOCK_SIZE,
89 height: IMPORTANCE_BLOCK_SIZE,
90 },
91 &mut plane_after_prediction_region,
92 tx_size,
93 bit_depth,
94 &[], // Not used by DC_PRED
95 IntraParam::None,
96 None, // Not used by DC_PRED
97 &edge_buf,
98 cpu_feature_level,
99 );
100
101 let plane_after_prediction_region =
102 plane_after_prediction.region(Area::Rect {
103 x: (x * IMPORTANCE_BLOCK_SIZE) as isize,
104 y: (y * IMPORTANCE_BLOCK_SIZE) as isize,
105 width: IMPORTANCE_BLOCK_SIZE,
106 height: IMPORTANCE_BLOCK_SIZE,
107 });
108
109 let intra_cost = get_satd(
110 &plane_org,
111 &plane_after_prediction_region,
112 bsize.width(),
113 bsize.height(),
114 bit_depth,
115 cpu_feature_level,
116 );
117
118 intra_costs.push(intra_cost);
119 }
120 }
121
122 intra_costs.into_boxed_slice()
123}
124
125#[profiling::function]
126pub(crate) fn estimate_importance_block_difference<T: Pixel>(
127 frame: Arc<Frame<T>>, ref_frame: Arc<Frame<T>>,
128) -> f64 {
129 let plane_org = &frame.planes[0];
130 let plane_ref = &ref_frame.planes[0];
131 let h_in_imp_b = plane_org.cfg.height / IMPORTANCE_BLOCK_SIZE;
132 let w_in_imp_b = plane_org.cfg.width / IMPORTANCE_BLOCK_SIZE;
133
134 let mut imp_block_costs = 0;
135
136 (0..h_in_imp_b).for_each(|y| {
137 (0..w_in_imp_b).for_each(|x| {
138 // Coordinates of the top-left corner of the reference block, in MV
139 // units.
140 let region_org = plane_org.region(Area::Rect {
141 x: (x * IMPORTANCE_BLOCK_SIZE) as isize,
142 y: (y * IMPORTANCE_BLOCK_SIZE) as isize,
143 width: IMPORTANCE_BLOCK_SIZE,
144 height: IMPORTANCE_BLOCK_SIZE,
145 });
146
147 let region_ref = plane_ref.region(Area::Rect {
148 x: (x * IMPORTANCE_BLOCK_SIZE) as isize,
149 y: (y * IMPORTANCE_BLOCK_SIZE) as isize,
150 width: IMPORTANCE_BLOCK_SIZE,
151 height: IMPORTANCE_BLOCK_SIZE,
152 });
153
154 let sum_8x8_block = |region: &PlaneRegion<T>| {
155 region
156 .rows_iter()
157 .map(|row| {
158 // 16-bit precision is sufficient for an 8px row, as IMPORTANCE_BLOCK_SIZE * (2^12 - 1) < 2^16 - 1,
159 // so overflow is not possible
160 row.iter().map(|pixel| u16::cast_from(*pixel)).sum::<u16>() as i64
161 })
162 .sum::<i64>()
163 };
164
165 let histogram_org_sum = sum_8x8_block(&region_org);
166 let histogram_ref_sum = sum_8x8_block(&region_ref);
167
168 let count = (IMPORTANCE_BLOCK_SIZE * IMPORTANCE_BLOCK_SIZE) as i64;
169
170 let mean = (((histogram_org_sum + count / 2) / count)
171 - ((histogram_ref_sum + count / 2) / count))
172 .abs();
173
174 imp_block_costs += mean as u64;
175 });
176 });
177
178 imp_block_costs as f64 / (w_in_imp_b * h_in_imp_b) as f64
179}
180
181#[profiling::function]
182pub(crate) fn estimate_inter_costs<T: Pixel>(
183 frame: Arc<Frame<T>>, ref_frame: Arc<Frame<T>>, bit_depth: usize,
184 mut config: EncoderConfig, sequence: Arc<Sequence>, buffer: RefMEStats,
185) -> f64 {
186 config.low_latency = true;
187 config.speed_settings.multiref = false;
188 let inter_cfg = InterConfig::new(&config);
189 let last_fi = FrameInvariants::new_key_frame(
190 Arc::new(config),
191 sequence,
192 0,
193 Box::new([]),
194 );
195 let mut fi = FrameInvariants::new_inter_frame(
196 &last_fi,
197 &inter_cfg,
198 0,
199 1,
200 2,
201 false,
202 Box::new([]),
203 )
204 .unwrap();
205
206 // Compute the motion vectors.
207 let mut fs = FrameState::new_with_frame_and_me_stats_and_rec(
208 &fi,
209 Arc::clone(&frame),
210 buffer,
211 // We do not use this field, so we can avoid the expensive allocation
212 Arc::new(Frame {
213 planes: [
214 Plane::new(0, 0, 0, 0, 0, 0),
215 Plane::new(0, 0, 0, 0, 0, 0),
216 Plane::new(0, 0, 0, 0, 0, 0),
217 ],
218 }),
219 );
220 compute_motion_vectors(&mut fi, &mut fs, &inter_cfg);
221
222 // Estimate inter costs
223 let plane_org = &frame.planes[0];
224 let plane_ref = &ref_frame.planes[0];
225 let h_in_imp_b = plane_org.cfg.height / IMPORTANCE_BLOCK_SIZE;
226 let w_in_imp_b = plane_org.cfg.width / IMPORTANCE_BLOCK_SIZE;
227 let stats = &fs.frame_me_stats.read().expect("poisoned lock")[0];
228 let bsize = BlockSize::from_width_and_height(
229 IMPORTANCE_BLOCK_SIZE,
230 IMPORTANCE_BLOCK_SIZE,
231 );
232
233 let mut inter_costs = 0;
234 (0..h_in_imp_b).for_each(|y| {
235 (0..w_in_imp_b).for_each(|x| {
236 let mv = stats[y * 2][x * 2].mv;
237
238 // Coordinates of the top-left corner of the reference block, in MV
239 // units.
240 let reference_x = x as i64 * IMP_BLOCK_SIZE_IN_MV_UNITS + mv.col as i64;
241 let reference_y = y as i64 * IMP_BLOCK_SIZE_IN_MV_UNITS + mv.row as i64;
242
243 let region_org = plane_org.region(Area::Rect {
244 x: (x * IMPORTANCE_BLOCK_SIZE) as isize,
245 y: (y * IMPORTANCE_BLOCK_SIZE) as isize,
246 width: IMPORTANCE_BLOCK_SIZE,
247 height: IMPORTANCE_BLOCK_SIZE,
248 });
249
250 let region_ref = plane_ref.region(Area::Rect {
251 x: reference_x as isize / IMP_BLOCK_MV_UNITS_PER_PIXEL as isize,
252 y: reference_y as isize / IMP_BLOCK_MV_UNITS_PER_PIXEL as isize,
253 width: IMPORTANCE_BLOCK_SIZE,
254 height: IMPORTANCE_BLOCK_SIZE,
255 });
256
257 inter_costs += get_satd(
258 &region_org,
259 &region_ref,
260 bsize.width(),
261 bsize.height(),
262 bit_depth,
263 fi.cpu_feature_level,
264 ) as u64;
265 });
266 });
267 inter_costs as f64 / (w_in_imp_b * h_in_imp_b) as f64
268}
269
270#[profiling::function]
271pub(crate) fn compute_motion_vectors<T: Pixel>(
272 fi: &mut FrameInvariants<T>, fs: &mut FrameState<T>, inter_cfg: &InterConfig,
273) {
274 let mut blocks: FrameBlocks = FrameBlocks::new(cols:fi.w_in_b, rows:fi.h_in_b);
275 fi.sequence
276 .tiling
277 .tile_iter_mut(fs, &mut blocks)
278 .collect::<Vec<_>>()
279 .into_par_iter()
280 .for_each(|mut ctx: TileContextMut<'_, T>| {
281 let ts: &mut TileStateMut<'_, T> = &mut ctx.ts;
282 estimate_tile_motion(fi, ts, inter_cfg);
283 });
284}
285