cinepakenc: add an option to further refine the selected vectors
[nihav.git] / nihav-commonfmt / src / codecs / cinepakenc.rs
CommitLineData
a178c22c
KS
1use nihav_core::codecs::*;
2use nihav_core::io::byteio::*;
3use nihav_codec_support::vq::*;
4
5#[derive(Default,Clone,Copy,PartialEq,Debug)]
6struct YUVCode {
7 y: [u8; 4],
8 u: u8,
9 v: u8,
10}
11impl VQElement for YUVCode {
12 fn dist(&self, rval: Self) -> u32 {
13 let mut ysum = 0;
14 for (y0, y1) in self.y.iter().zip(rval.y.iter()) {
15 let yd = i32::from(*y0) - i32::from(*y1);
16 ysum += yd * yd;
17 }
18 let ud = i32::from(self.u) - i32::from(rval.u);
19 let vd = i32::from(self.v) - i32::from(rval.v);
20 (ysum + ud * ud + vd * vd) as u32
21 }
22 fn min_cw() -> Self { YUVCode { y: [0; 4], u: 0, v: 0 } }
23 fn max_cw() -> Self { YUVCode { y: [255; 4], u: 255, v: 255 } }
24 fn min(&self, rval: Self) -> Self {
25 let mut ycode = YUVCode::default();
26 for i in 0..4 {
27 ycode.y[i] = self.y[i].min(rval.y[i]);
28 }
29 ycode.u = self.u.min(rval.u);
30 ycode.v = self.v.min(rval.v);
31 ycode
32 }
33 fn max(&self, rval: Self) -> Self {
34 let mut ycode = YUVCode::default();
35 for i in 0..4 {
36 ycode.y[i] = self.y[i].max(rval.y[i]);
37 }
38 ycode.u = self.u.max(rval.u);
39 ycode.v = self.v.max(rval.v);
40 ycode
41 }
42 fn num_components() -> usize { 6 }
43 fn sort_by_component(arr: &mut [Self], component: usize) {
44 let mut counts = [0; 256];
45 for entry in arr.iter() {
46 let idx = match component {
47 0 | 1 | 2 | 3 => entry.y[component],
48 4 => entry.u,
49 _ => entry.v,
50 } as usize;
51 counts[idx] += 1;
52 }
53 let mut offs = [0; 256];
54 for i in 0..255 {
55 offs[i + 1] = offs[i] + counts[i];
56 }
57 let mut dst = vec![YUVCode::default(); arr.len()];
58 for entry in arr.iter() {
59 let idx = match component {
60 0 | 1 | 2 | 3 => entry.y[component],
61 4 => entry.u,
62 _ => entry.v,
63 } as usize;
64 dst[offs[idx]] = *entry;
65 offs[idx] += 1;
66 }
67 arr.copy_from_slice(dst.as_slice());
68 }
69 fn max_dist_component(min: &Self, max: &Self) -> usize {
70 let mut comp = 0;
71 let mut diff = 0;
72 for i in 0..4 {
73 let d = u32::from(max.y[i]) - u32::from(min.y[i]);
74 if d > diff {
75 diff = d;
76 comp = i;
77 }
78 }
79 let ud = u32::from(max.u) - u32::from(min.u);
80 if ud > diff {
81 diff = ud;
82 comp = 4;
83 }
84 let vd = u32::from(max.v) - u32::from(min.v);
85 if vd > diff {
86 comp = 5;
87 }
88 comp
89 }
90}
91
92#[derive(Default)]
93struct YUVCodeSum {
94 ysum: [u64; 4],
95 usum: u64,
96 vsum: u64,
97 count: u64,
98}
99
100impl VQElementSum<YUVCode> for YUVCodeSum {
101 fn zero() -> Self { Self::default() }
102 fn add(&mut self, rval: YUVCode, count: u64) {
103 for i in 0..4 {
104 self.ysum[i] += u64::from(rval.y[i]) * count;
105 }
106 self.usum += u64::from(rval.u) * count;
107 self.vsum += u64::from(rval.v) * count;
108 self.count += count;
109 }
110 fn get_centroid(&self) -> YUVCode {
111 if self.count != 0 {
112 let mut ycode = YUVCode::default();
113 for i in 0..4 {
114 ycode.y[i] = ((self.ysum[i] + self.count / 2) / self.count) as u8;
115 }
116 ycode.u = ((self.usum + self.count / 2) / self.count) as u8;
117 ycode.v = ((self.vsum + self.count / 2) / self.count) as u8;
118 ycode
119 } else {
120 YUVCode::default()
121 }
122 }
123}
124
125struct RNG {
126 seed: u32,
127}
128
129impl RNG {
130 fn new() -> Self { Self { seed: 0x12345678 } }
131 fn next(&mut self) -> u8 {
132 let mut x = self.seed;
133 x ^= x.wrapping_shl(13);
134 x ^= x >> 17;
135 self.seed = x;
136 (self.seed >> 24) as u8
137 }
138 fn fill_entry(&mut self, entry: &mut YUVCode) {
139 for y in entry.y.iter_mut() {
140 *y = self.next();
141 }
142 entry.u = self.next();
143 entry.v = self.next();
144 }
145}
146
147const GRAY_FORMAT: NAPixelFormaton = NAPixelFormaton {
148 model: ColorModel::YUV(YUVSubmodel::YUVJ),
149 components: 1,
150 comp_info: [Some(NAPixelChromaton{h_ss: 0, v_ss: 0, packed: false, depth: 8, shift: 0, comp_offs: 0, next_elem: 1}), None, None, None, None],
151 elem_size: 1,
152 be: true,
153 alpha: false,
154 palette: false,
155 };
156
157struct MaskWriter {
158 masks: Vec<u32>,
159 mask: u32,
160 pos: u8,
161}
162
163impl MaskWriter {
164 fn new() -> Self {
165 Self {
166 masks: Vec::new(),
167 mask: 0,
168 pos: 0,
169 }
170 }
171 fn reset(&mut self) {
37952415 172 self.masks.clear();
a178c22c
KS
173 self.mask = 0;
174 self.pos = 0;
175 }
176 fn put_v1(&mut self) {
177 self.mask <<= 1;
178 self.pos += 1;
179 if self.pos == 32 {
180 self.flush();
181 }
182 }
183 fn put_v4(&mut self) {
184 self.mask <<= 1;
185 self.mask |= 1;
186 self.pos += 1;
187 if self.pos == 32 {
188 self.flush();
189 }
190 }
191 fn put_inter(&mut self, skip: bool) {
192 self.mask <<= 1;
193 self.mask |= !skip as u32;
194 self.pos += 1;
195 if self.pos == 32 {
196 self.flush();
197 }
198 }
199 fn flush(&mut self) {
200 self.masks.push(self.mask);
201 self.mask = 0;
202 self.pos = 0;
203 }
204 fn end(&mut self) {
205 if self.pos == 0 { return; }
206 while self.pos < 32 {
207 self.mask <<= 1;
208 self.pos += 1;
209 }
210 self.flush();
211 }
212}
213
3c406629
KS
214#[derive(Clone,Copy,PartialEq)]
215enum QuantMode {
216 ELBG,
42ef5325 217 Fast,
3c406629
KS
218 MedianCut,
219}
220
61cab15b 221impl std::string::ToString for QuantMode {
3c406629
KS
222 fn to_string(&self) -> String {
223 match *self {
224 QuantMode::ELBG => "elbg".to_string(),
42ef5325 225 QuantMode::Fast => "fast".to_string(),
3c406629
KS
226 QuantMode::MedianCut => "mediancut".to_string(),
227 }
228 }
229}
230
6064de24
KS
231#[derive(Clone,Copy,PartialEq)]
232enum CodingMode {
233 Skip,
234 V1,
235 V4
236}
237
a178c22c
KS
238struct CinepakEncoder {
239 stream: Option<NAStreamRef>,
240 lastfrm: Option<NAVideoBufferRef<u8>>,
241 pkt: Option<NAPacket>,
242 frmcount: u8,
3c406629
KS
243 key_int: u8,
244 qmode: QuantMode,
a178c22c 245 quality: u8,
6064de24 246 refine: bool,
a178c22c 247 nstrips: usize,
c6f14420 248 force_v1: bool,
9a57fda1 249 cur_strip: usize,
a178c22c
KS
250 v1_entries: Vec<YUVCode>,
251 v4_entries: Vec<YUVCode>,
9a57fda1
KS
252 v1_cb: Vec<[YUVCode; 256]>,
253 v4_cb: Vec<[YUVCode; 256]>,
254 v1_cur_cb: Vec<[YUVCode; 256]>,
255 v4_cur_cb: Vec<[YUVCode; 256]>,
592d2889
KS
256 v1_len: usize,
257 v4_len: usize,
a178c22c
KS
258 v1_idx: Vec<u8>,
259 v4_idx: Vec<u8>,
260 grayscale: bool,
261 rng: RNG,
262 masks: MaskWriter,
263 skip_dist: Vec<u32>,
42ef5325 264 fst_bins: [Vec<YUVCode>; 4],
6064de24
KS
265 v1_cand: Vec<YUVCode>,
266 v4_cand: Vec<YUVCode>,
267 cmode: Vec<CodingMode>,
a178c22c
KS
268}
269
270fn avg4(a: u8, b: u8, c: u8, d: u8) -> u8 {
271 ((u16::from(a) + u16::from(b) + u16::from(c) + u16::from(d) + 3) >> 2) as u8
272}
273
42ef5325
KS
274fn variance(a: u8, mean: u8) -> u32 {
275 if a >= mean {
276 u32::from(a - mean) * u32::from(a - mean)
277 } else {
278 u32::from(mean - a) * u32::from(mean - a)
279 }
280}
281
a178c22c
KS
282fn patch_size(bw: &mut ByteWriter, pos: u64) -> EncoderResult<()> {
283 let size = bw.tell() - pos;
284 bw.seek(SeekFrom::Current(-((size + 3) as i64)))?;
285 bw.write_u24be((size + 4) as u32)?;
286 bw.seek(SeekFrom::End(0))?;
287 Ok(())
288}
289
7460b53e
KS
290fn elbg_quant(entries: &[YUVCode], codebook: &mut [YUVCode]) -> usize {
291 let cb_len = quantise_median_cut::<YUVCode, YUVCodeSum>(entries, codebook);
292 if cb_len < codebook.len() {
293 cb_len
294 } else {
295 let mut elbg: ELBG<YUVCode, YUVCodeSum> = ELBG::new(codebook);
296 elbg.quantise(entries, codebook)
297 }
298}
299
42ef5325
KS
300fn quant_fast(bins: &mut [Vec<YUVCode>; 4], entries: &[YUVCode], codebook: &mut [YUVCode]) -> usize {
301 for bin in bins.iter_mut() {
302 bin.clear();
303 }
304 for &entry in entries.iter() {
305 let y_avg = avg4(entry.y[0], entry.y[1], entry.y[2], entry.y[3]);
306 let dist = entry.y.iter().fold(0u32, |acc, &x| acc + variance(x, y_avg));
307 let ilog = if dist == 0 { 0 } else { 32 - dist.leading_zeros() };
308 let bin = match ilog {
309 0..=3 => &mut bins[0],
310 4..=7 => &mut bins[1],
311 8..=11 => &mut bins[2],
312 _ => &mut bins[3],
313 };
314 bin.push(entry);
315 }
316 let mut free_cw = codebook.len();
317 let mut entries_left = entries.len();
318 let mut offset = 0;
319 for bin in bins.iter() {
320 if bin.is_empty() {
321 continue;
322 }
323 if free_cw == 0 || entries_left == 0 {
324 break;
325 }
326 let target = (free_cw * bin.len() + entries_left - 1) / entries_left;
327 let cur_len = elbg_quant(bin, &mut codebook[offset..][..target]);
328 offset += cur_len;
329 free_cw -= cur_len;
330 entries_left -= bin.len();
331 }
332 offset
333}
334
a178c22c
KS
335impl CinepakEncoder {
336 fn new() -> Self {
337 Self {
338 stream: None,
339 pkt: None,
340 lastfrm: None,
341 frmcount: 0,
42ef5325 342 qmode: QuantMode::Fast,
3c406629 343 key_int: 25,
a178c22c 344 quality: 0,
6064de24 345 refine: false,
a178c22c 346 nstrips: 2,
c6f14420 347 force_v1: false,
9a57fda1 348 cur_strip: 0,
a178c22c
KS
349 v1_entries: Vec::new(),
350 v4_entries: Vec::new(),
9a57fda1
KS
351 v1_cb: Vec::with_capacity(2),
352 v4_cb: Vec::with_capacity(2),
353 v1_cur_cb: Vec::with_capacity(2),
354 v4_cur_cb: Vec::with_capacity(2),
592d2889
KS
355 v1_len: 0,
356 v4_len: 0,
a178c22c
KS
357 grayscale: false,
358 rng: RNG::new(),
359 v1_idx: Vec::new(),
360 v4_idx: Vec::new(),
361 masks: MaskWriter::new(),
362 skip_dist: Vec::new(),
42ef5325 363 fst_bins: [Vec::new(), Vec::new(), Vec::new(), Vec::new()],
6064de24
KS
364 v1_cand: Vec::new(),
365 v4_cand: Vec::new(),
366 cmode: Vec::new(),
a178c22c
KS
367 }
368 }
369 fn read_strip(&mut self, in_frm: &NAVideoBuffer<u8>, start: usize, end: usize) {
370 let ystride = in_frm.get_stride(0);
371 let mut yoff = in_frm.get_offset(0) + start * ystride;
372 let ustride = in_frm.get_stride(1);
373 let mut uoff = in_frm.get_offset(1) + start / 2 * ustride;
374 let vstride = in_frm.get_stride(2);
375 let mut voff = in_frm.get_offset(2) + start / 2 * vstride;
376 let (width, _) = in_frm.get_dimensions(0);
377 let data = in_frm.get_data();
37952415
KS
378 self.v1_entries.clear();
379 self.v4_entries.clear();
a178c22c
KS
380 for _ in (start..end).step_by(4) {
381 for x in (0..width).step_by(4) {
382 let mut yblk = [0; 16];
383 let mut ublk = [128; 4];
384 let mut vblk = [128; 4];
385 for j in 0..4 {
386 for i in 0..4 {
387 yblk[i + j * 4] = data[yoff + x + i + j * ystride];
388 }
389 }
390 if !self.grayscale {
391 for j in 0..2 {
392 for i in 0..2 {
393 ublk[i + j * 2] = data[uoff + x / 2 + i + j * ustride];
394 vblk[i + j * 2] = data[voff + x / 2 + i + j * vstride];
395 }
396 }
397 }
398 self.v1_entries.push(YUVCode {
399 y: [avg4(yblk[ 0], yblk[ 1], yblk[ 4], yblk[ 5]),
400 avg4(yblk[ 2], yblk[ 3], yblk[ 6], yblk[ 7]),
401 avg4(yblk[ 8], yblk[ 9], yblk[12], yblk[13]),
402 avg4(yblk[10], yblk[11], yblk[14], yblk[15])],
403 u: avg4(ublk[0], ublk[1], ublk[2], ublk[3]),
404 v: avg4(vblk[0], vblk[1], vblk[2], vblk[3]),
405 });
406 for i in 0..4 {
407 let yidx = (i & 1) * 2 + (i & 2) * 4;
408 self.v4_entries.push(YUVCode {
409 y: [ yblk[yidx], yblk[yidx + 1], yblk[yidx + 4], yblk[yidx + 5] ],
410 u: ublk[i],
411 v: vblk[i],
412 });
413 }
414 }
415 yoff += ystride * 4;
416 uoff += ustride * 2;
417 voff += vstride * 2;
418 }
419 }
592d2889 420 fn find_nearest(codebook: &[YUVCode], code: YUVCode) -> (u8, u32) {
a178c22c
KS
421 let mut min_dist = std::u32::MAX;
422 let mut idx = 0;
423 for (i, cw) in codebook.iter().enumerate() {
424 let dist = cw.dist(code);
425 if dist < min_dist {
426 min_dist = dist;
427 idx = i;
428 if dist == 0 {
429 break;
430 }
431 }
432 }
433 (idx as u8, min_dist)
434 }
9f7faaf3 435 fn can_update_cb(new_cb: &[YUVCode], old_cb: &[YUVCode], cb_size: usize) -> bool {
a178c22c
KS
436 let mut skip_count = 0;
437 for (new, old) in new_cb.iter().zip(old_cb.iter()) {
438 if new == old {
439 skip_count += 1;
440 }
441 }
9f7faaf3
KS
442 let full_size = cb_size * new_cb.len();
443 let upd_size = cb_size * (new_cb.len() - skip_count) + (new_cb.len() + 31) / 32 * 4;
a178c22c
KS
444 upd_size < full_size
445 }
9f7faaf3 446 fn write_cb(bw: &mut ByteWriter, mut id: u8, new_cb: &[YUVCode], old_cb: &[YUVCode], grayscale: bool, update: bool, num_elem: usize) -> EncoderResult<()> {
a178c22c
KS
447 if grayscale {
448 id |= 4;
449 }
450 if update {
451 id |= 1;
452 }
453 bw.write_byte(id)?;
454 bw.write_u24be(0)?;
455 let chunk_pos = bw.tell();
456 if !update {
9f7faaf3 457 for entry in new_cb.iter().take(num_elem) {
a178c22c
KS
458 bw.write_buf(&entry.y)?;
459 if !grayscale {
460 bw.write_byte(entry.u ^ 0x80)?;
461 bw.write_byte(entry.v ^ 0x80)?;
462 }
463 }
464 } else {
9f7faaf3
KS
465 let mut end = num_elem;
466 for (i, (ncw, ocw)) in new_cb.iter().zip(old_cb.iter()).enumerate().take(num_elem).rev() {
a178c22c
KS
467 if ncw == ocw {
468 end = i;
469 } else {
470 break;
471 }
472 }
473 for i in (0..end).step_by(32) {
474 let mut mask = 0;
475 for j in 0..32 {
476 mask <<= 1;
477 if new_cb[i + j] != old_cb[i + j] {
478 mask |= 1;
479 }
480 }
481 bw.write_u32be(mask)?;
482 for j in 0..32 {
483 if new_cb[i + j] == old_cb[i + j] { continue; }
484 bw.write_buf(&new_cb[i + j].y)?;
485 if !grayscale {
486 bw.write_byte(new_cb[i + j].u ^ 0x80)?;
487 bw.write_byte(new_cb[i + j].v ^ 0x80)?;
488 }
489 }
490 }
491 }
492 patch_size(bw, chunk_pos)?;
493 Ok(())
494 }
495 fn render_stripe(&mut self, intra: bool, start: usize, end: usize) {
496 if let Some(ref mut dst_frm) = self.lastfrm {
497 let ystride = dst_frm.get_stride(0);
498 let mut yoff = dst_frm.get_offset(0) + start * ystride;
499 let ustride = dst_frm.get_stride(1);
500 let mut uoff = dst_frm.get_offset(1) + start / 2 * ustride;
501 let vstride = dst_frm.get_stride(2);
502 let mut voff = dst_frm.get_offset(2) + start / 2 * vstride;
503 let (width, _) = dst_frm.get_dimensions(0);
504 let data = dst_frm.get_data_mut().unwrap();
505 let mut miter = self.masks.masks.iter();
506 let mut v1_iter = self.v1_idx.iter();
507 let mut v4_iter = self.v4_idx.iter();
508 let mut cur_mask = 0;
509 let mut cur_bit = 0;
510 for _ in (start..end).step_by(4) {
511 for x in (0..width).step_by(4) {
512 if cur_bit == 0 {
61cab15b 513 if !intra || !self.v1_idx.is_empty() {
a178c22c
KS
514 cur_mask = *miter.next().unwrap();
515 } else {
516 cur_mask = 0xFFFFFFFF;
517 }
518 cur_bit = 1 << 31;
519 }
520 if !intra {
521 if (cur_mask & cur_bit) == 0 {
522 cur_bit >>= 1;
523 continue;
524 }
525 cur_bit >>= 1;
526 if cur_bit == 0 {
527 cur_mask = *miter.next().unwrap();
528 cur_bit = 1 << 31;
529 }
530 }
531 if (cur_mask & cur_bit) == 0 {
532 let idx = *v1_iter.next().unwrap() as usize;
9a57fda1 533 let cb = &self.v1_cur_cb[self.cur_strip][idx];
a178c22c
KS
534
535 let mut coff = yoff + x;
536 data[coff] = cb.y[0]; data[coff + 1] = cb.y[0];
537 data[coff + 2] = cb.y[1]; data[coff + 3] = cb.y[1];
538 coff += ystride;
539 data[coff] = cb.y[0]; data[coff + 1] = cb.y[0];
540 data[coff + 2] = cb.y[1]; data[coff + 3] = cb.y[1];
541 coff += ystride;
542 data[coff] = cb.y[2]; data[coff + 1] = cb.y[2];
543 data[coff + 2] = cb.y[3]; data[coff + 3] = cb.y[3];
544 coff += ystride;
545 data[coff] = cb.y[2]; data[coff + 1] = cb.y[2];
546 data[coff + 2] = cb.y[3]; data[coff + 3] = cb.y[3];
547
548 if !self.grayscale {
549 let mut coff = uoff + x / 2;
550 data[coff] = cb.u; data[coff + 1] = cb.u;
551 coff += ustride;
552 data[coff] = cb.u; data[coff + 1] = cb.u;
553
554 let mut coff = voff + x / 2;
555 data[coff] = cb.v; data[coff + 1] = cb.v;
556 coff += vstride;
557 data[coff] = cb.v; data[coff + 1] = cb.v;
558 }
559 } else {
560 let idx0 = *v4_iter.next().unwrap() as usize;
9a57fda1 561 let cb0 = &self.v4_cur_cb[self.cur_strip][idx0];
a178c22c 562 let idx1 = *v4_iter.next().unwrap() as usize;
9a57fda1 563 let cb1 = &self.v4_cur_cb[self.cur_strip][idx1];
a178c22c 564 let idx2 = *v4_iter.next().unwrap() as usize;
9a57fda1 565 let cb2 = &self.v4_cur_cb[self.cur_strip][idx2];
a178c22c 566 let idx3 = *v4_iter.next().unwrap() as usize;
9a57fda1 567 let cb3 = &self.v4_cur_cb[self.cur_strip][idx3];
a178c22c
KS
568
569 let mut coff = yoff + x;
570 data[coff] = cb0.y[0]; data[coff + 1] = cb0.y[1];
571 data[coff + 2] = cb1.y[0]; data[coff + 3] = cb1.y[1];
572 coff += ystride;
573 data[coff] = cb0.y[2]; data[coff + 1] = cb0.y[3];
574 data[coff + 2] = cb1.y[2]; data[coff + 3] = cb1.y[3];
575 coff += ystride;
576 data[coff] = cb2.y[0]; data[coff + 1] = cb2.y[1];
577 data[coff + 2] = cb3.y[0]; data[coff + 3] = cb3.y[1];
578 coff += ystride;
579 data[coff] = cb2.y[2]; data[coff + 1] = cb2.y[3];
580 data[coff + 2] = cb3.y[2]; data[coff + 3] = cb3.y[3];
581
582 if !self.grayscale {
583 let mut coff = uoff + x / 2;
584 data[coff] = cb0.u; data[coff + 1] = cb1.u;
585 coff += ustride;
586 data[coff] = cb2.u; data[coff + 1] = cb3.u;
587
588 let mut coff = voff + x / 2;
589 data[coff] = cb0.v; data[coff + 1] = cb1.v;
590 coff += vstride;
591 data[coff] = cb2.v; data[coff + 1] = cb3.v;
592 }
593 }
594 cur_bit >>= 1;
595 }
596 yoff += ystride * 4;
597 uoff += ustride * 2;
598 voff += vstride * 2;
599 }
600 } else {
601 unreachable!();
602 }
603 }
604 fn calc_skip_dist(&mut self, in_frm: &NAVideoBuffer<u8>, start: usize, end: usize) {
37952415 605 self.skip_dist.clear();
a178c22c
KS
606 if let Some(ref ref_frm) = self.lastfrm {
607 let rystride = ref_frm.get_stride(0);
608 let mut ryoff = ref_frm.get_offset(0) + start * rystride;
609 let rustride = ref_frm.get_stride(1);
610 let mut ruoff = ref_frm.get_offset(1) + start / 2 * rustride;
611 let rvstride = ref_frm.get_stride(2);
612 let mut rvoff = ref_frm.get_offset(2) + start / 2 * rvstride;
613 let (width, _) = ref_frm.get_dimensions(0);
614 let rdata = ref_frm.get_data();
615
616 let iystride = in_frm.get_stride(0);
617 let mut iyoff = in_frm.get_offset(0) + start * iystride;
618 let iustride = in_frm.get_stride(1);
619 let mut iuoff = in_frm.get_offset(1) + start / 2 * iustride;
620 let ivstride = in_frm.get_stride(2);
621 let mut ivoff = in_frm.get_offset(2) + start / 2 * ivstride;
622 let idata = in_frm.get_data();
623
624 for _ in (start..end).step_by(4) {
625 for x in (0..width).step_by(4) {
626 let mut dist = 0;
627 let mut roff = ryoff + x;
628 let mut ioff = iyoff + x;
629 for _ in 0..4 {
630 for i in 0..4 {
631 let d = i32::from(rdata[roff + i]) - i32::from(idata[ioff + i]);
632 dist += d * d;
633 }
634 roff += rystride;
635 ioff += iystride;
636 }
637 if !self.grayscale {
638 let mut roff = ruoff + x / 2;
639 let mut ioff = iuoff + x / 2;
640 let ud = i32::from(rdata[roff]) - i32::from(idata[ioff]);
641 dist += ud * ud;
642 let ud = i32::from(rdata[roff + 1]) - i32::from(idata[ioff + 1]);
643 dist += ud * ud;
644 roff += rustride; ioff += iustride;
645 let ud = i32::from(rdata[roff]) - i32::from(idata[ioff]);
646 dist += ud * ud;
647 let ud = i32::from(rdata[roff + 1]) - i32::from(idata[ioff + 1]);
648 dist += ud * ud;
649
650 let mut roff = rvoff + x / 2;
651 let mut ioff = ivoff + x / 2;
652 let vd = i32::from(rdata[roff]) - i32::from(idata[ioff]);
653 dist += vd * vd;
654 let vd = i32::from(rdata[roff + 1]) - i32::from(idata[ioff + 1]);
655 dist += vd * vd;
656 roff += rvstride; ioff += ivstride;
657 let vd = i32::from(rdata[roff]) - i32::from(idata[ioff]);
658 dist += vd * vd;
659 let vd = i32::from(rdata[roff + 1]) - i32::from(idata[ioff + 1]);
660 dist += vd * vd;
661 }
662 self.skip_dist.push(dist as u32);
663 }
664
665 iyoff += iystride * 4;
666 iuoff += iustride * 2;
667 ivoff += ivstride * 2;
668 ryoff += rystride * 4;
669 ruoff += rustride * 2;
670 rvoff += rvstride * 2;
671 }
672 } else {
673 unreachable!();
674 }
675 }
3c406629
KS
676 fn quant_vectors(&mut self) {
677 match self.qmode {
678 QuantMode::ELBG => {
7460b53e
KS
679 self.v1_len = elbg_quant(&self.v1_entries, &mut self.v1_cur_cb[self.cur_strip]);
680 self.v4_len = if !self.force_v1 {
681 elbg_quant(&self.v4_entries, &mut self.v4_cur_cb[self.cur_strip])
c6f14420
KS
682 } else {
683 0
684 };
3c406629 685 },
42ef5325
KS
686 QuantMode::Fast => {
687 for bin in self.fst_bins.iter_mut() {
688 bin.clear();
689 }
690 self.v1_len = quant_fast(&mut self.fst_bins, &self.v1_entries, &mut self.v1_cur_cb[self.cur_strip]);
691 self.v4_len = if !self.force_v1 {
692 quant_fast(&mut self.fst_bins, &self.v4_entries, &mut self.v4_cur_cb[self.cur_strip])
693 } else {
694 0
695 };
696 },
3c406629 697 QuantMode::MedianCut => {
9a57fda1 698 self.v1_len = quantise_median_cut::<YUVCode, YUVCodeSum>(&self.v1_entries, &mut self.v1_cur_cb[self.cur_strip]);
c6f14420
KS
699 if !self.force_v1 {
700 self.v4_len = quantise_median_cut::<YUVCode, YUVCodeSum>(&self.v4_entries, &mut self.v4_cur_cb[self.cur_strip]);
701 } else {
702 self.v4_len = 0;
703 }
3c406629
KS
704 },
705 };
592d2889 706
9a57fda1
KS
707 for e in self.v1_cur_cb[self.cur_strip].iter_mut().skip(self.v1_len) { *e = YUVCode::default(); }
708 for e in self.v4_cur_cb[self.cur_strip].iter_mut().skip(self.v4_len) { *e = YUVCode::default(); }
3c406629 709 }
6064de24
KS
710 fn refine_vectors(&mut self) {
711 match self.qmode {
712 QuantMode::ELBG => {
713 self.v1_len = if !self.v1_cand.is_empty() {
714 elbg_quant(&self.v1_cand, &mut self.v1_cur_cb[self.cur_strip])
715 } else {
716 0
717 };
718 self.v4_len = if !self.force_v1 && !self.v4_cand.is_empty() {
719 elbg_quant(&self.v4_cand, &mut self.v4_cur_cb[self.cur_strip])
720 } else {
721 0
722 };
723 },
724 QuantMode::Fast => {
725 for bin in self.fst_bins.iter_mut() {
726 bin.clear();
727 }
728 self.v1_len = if !self.v1_cand.is_empty() {
729 quant_fast(&mut self.fst_bins, &self.v1_cand, &mut self.v1_cur_cb[self.cur_strip])
730 } else {
731 0
732 };
733 self.v4_len = if !self.force_v1 && !self.v4_cand.is_empty() {
734 quant_fast(&mut self.fst_bins, &self.v4_cand, &mut self.v4_cur_cb[self.cur_strip])
735 } else {
736 0
737 };
738 },
739 QuantMode::MedianCut => {
740 self.v1_len = if !self.v1_cand.is_empty() {
741 quantise_median_cut::<YUVCode, YUVCodeSum>(&self.v1_cand, &mut self.v1_cur_cb[self.cur_strip])
742 } else {
743 0
744 };
745 if !self.force_v1 && !self.v4_cand.is_empty() {
746 self.v4_len = quantise_median_cut::<YUVCode, YUVCodeSum>(&self.v4_cand, &mut self.v4_cur_cb[self.cur_strip]);
747 } else {
748 self.v4_len = 0;
749 }
750 },
751 };
752
753 for e in self.v1_cur_cb[self.cur_strip].iter_mut().skip(self.v1_len) { *e = YUVCode::default(); }
754 for e in self.v4_cur_cb[self.cur_strip].iter_mut().skip(self.v4_len) { *e = YUVCode::default(); }
755 }
a178c22c
KS
756 fn encode_intra(&mut self, bw: &mut ByteWriter, in_frm: &NAVideoBuffer<u8>) -> EncoderResult<bool> {
757 let (width, height) = in_frm.get_dimensions(0);
758 let mut strip_h = (height / self.nstrips + 3) & !3;
759 if strip_h == 0 {
760 self.nstrips = 1;
761 strip_h = height;
762 }
763 let mut start_line = 0;
764 let mut end_line = strip_h;
765
766 bw.write_byte(0)?; // intra flag
767 bw.write_u24be(0)?; // frame size
768 let frame_data_pos = bw.tell();
769 bw.write_u16be(width as u16)?;
770 bw.write_u16be(height as u16)?;
771 bw.write_u16be(self.nstrips as u16)?;
772
9a57fda1
KS
773 self.cur_strip = 0;
774 for entry in self.v1_cb[self.cur_strip].iter_mut() {
a178c22c
KS
775 self.rng.fill_entry(entry);
776 }
9a57fda1 777 for entry in self.v4_cb[self.cur_strip].iter_mut() {
a178c22c
KS
778 self.rng.fill_entry(entry);
779 }
780 while start_line < height {
781 self.read_strip(in_frm, start_line, end_line);
782
9a57fda1
KS
783 if self.cur_strip > 0 {
784 self.v1_cb[self.cur_strip] = self.v1_cb[self.cur_strip - 1];
785 self.v4_cb[self.cur_strip] = self.v4_cb[self.cur_strip - 1];
786 }
3c406629 787 self.quant_vectors();
a178c22c 788 if self.grayscale {
9a57fda1 789 for cw in self.v1_cur_cb[self.cur_strip].iter_mut() {
a178c22c
KS
790 cw.u = 128;
791 cw.v = 128;
792 }
9a57fda1 793 for cw in self.v4_cur_cb[self.cur_strip].iter_mut() {
a178c22c
KS
794 cw.u = 128;
795 cw.v = 128;
796 }
797 }
798
37952415
KS
799 self.v1_idx.clear();
800 self.v4_idx.clear();
a178c22c
KS
801 self.masks.reset();
802
6064de24
KS
803 self.cmode.clear();
804 self.v1_cand.clear();
805 self.v4_cand.clear();
a178c22c 806 for (v1_entry, v4_entries) in self.v1_entries.iter().zip(self.v4_entries.chunks(4)) {
9a57fda1 807 let (v1_idx, v1_dist) = Self::find_nearest(&self.v1_cur_cb[self.cur_strip][..self.v1_len], *v1_entry);
c6f14420 808 if v1_dist == 0 || self.force_v1 {
6064de24
KS
809 if !self.refine {
810 self.masks.put_v1();
811 self.v1_idx.push(v1_idx);
812 } else {
813 self.cmode.push(CodingMode::V1);
814 self.v1_cand.push(*v1_entry);
815 }
a178c22c
KS
816 continue;
817 }
9a57fda1
KS
818 let (v40_idx, v40_dist) = Self::find_nearest(&self.v4_cur_cb[self.cur_strip][..self.v4_len], v4_entries[0]);
819 let (v41_idx, v41_dist) = Self::find_nearest(&self.v4_cur_cb[self.cur_strip][..self.v4_len], v4_entries[1]);
820 let (v42_idx, v42_dist) = Self::find_nearest(&self.v4_cur_cb[self.cur_strip][..self.v4_len], v4_entries[2]);
821 let (v43_idx, v43_dist) = Self::find_nearest(&self.v4_cur_cb[self.cur_strip][..self.v4_len], v4_entries[3]);
a178c22c 822 if v40_dist + v41_dist + v42_dist + v43_dist > v1_dist {
6064de24
KS
823 if !self.refine {
824 self.masks.put_v4();
825 self.v4_idx.push(v40_idx);
826 self.v4_idx.push(v41_idx);
827 self.v4_idx.push(v42_idx);
828 self.v4_idx.push(v43_idx);
829 } else {
830 self.cmode.push(CodingMode::V4);
831 self.v4_cand.extend_from_slice(v4_entries);
832 }
a178c22c 833 } else {
6064de24
KS
834 if !self.refine {
835 self.masks.put_v1();
836 self.v1_idx.push(v1_idx);
837 } else {
838 self.cmode.push(CodingMode::V1);
839 self.v1_cand.push(*v1_entry);
840 }
841 }
842 }
843 if self.refine {
844 self.refine_vectors();
845 let mut v1_src = self.v1_cand.iter();
846 let mut v4_src = self.v4_cand.chunks_exact(4);
847 for &cmode in self.cmode.iter() {
848 match cmode {
849 CodingMode::Skip => unreachable!(),
850 CodingMode::V1 => {
851 let v1_entry = v1_src.next().unwrap();
852 let (v1_idx, _) = Self::find_nearest(&self.v1_cur_cb[self.cur_strip][..self.v1_len], *v1_entry);
853 self.masks.put_v1();
854 self.v1_idx.push(v1_idx);
855 },
856 CodingMode::V4 => {
857 let v4_entries = v4_src.next().unwrap();
858 let (v40_idx, _) = Self::find_nearest(&self.v4_cur_cb[self.cur_strip][..self.v4_len], v4_entries[0]);
859 let (v41_idx, _) = Self::find_nearest(&self.v4_cur_cb[self.cur_strip][..self.v4_len], v4_entries[1]);
860 let (v42_idx, _) = Self::find_nearest(&self.v4_cur_cb[self.cur_strip][..self.v4_len], v4_entries[2]);
861 let (v43_idx, _) = Self::find_nearest(&self.v4_cur_cb[self.cur_strip][..self.v4_len], v4_entries[3]);
862
863 self.masks.put_v4();
864 self.v4_idx.push(v40_idx);
865 self.v4_idx.push(v41_idx);
866 self.v4_idx.push(v42_idx);
867 self.v4_idx.push(v43_idx);
868 },
869 };
a178c22c
KS
870 }
871 }
872 self.masks.end();
873
874 let mut is_intra_strip = start_line == 0;
875 let (upd_v1, upd_v4) = if !is_intra_strip {
876 let cb_size = if self.grayscale { 4 } else { 6 };
9f7faaf3
KS
877 (Self::can_update_cb(&self.v1_cur_cb[self.cur_strip][..self.v1_len], &self.v1_cb[self.cur_strip][..self.v1_len], cb_size),
878 Self::can_update_cb(&self.v4_cur_cb[self.cur_strip][..self.v4_len], &self.v4_cb[self.cur_strip][..self.v4_len], cb_size))
a178c22c
KS
879 } else {
880 (false, false)
881 };
882 if !is_intra_strip && !upd_v1 && !upd_v4 {
883 is_intra_strip = true;
884 }
885 bw.write_byte(if is_intra_strip { 0x10 } else { 0x11 })?;
886 bw.write_u24be(0)?; // strip size
887 let strip_data_pos = bw.tell();
888 bw.write_u16be(0)?; // yoff
889 bw.write_u16be(0)?; // xoff
890 bw.write_u16be((end_line - start_line) as u16)?;
891 bw.write_u16be(width as u16)?;
892
9f7faaf3
KS
893 Self::write_cb(bw, 0x20, &self.v4_cur_cb[self.cur_strip], &self.v4_cb[self.cur_strip], self.grayscale, upd_v4, self.v4_len)?;
894 Self::write_cb(bw, 0x22, &self.v1_cur_cb[self.cur_strip], &self.v1_cb[self.cur_strip], self.grayscale, upd_v1, self.v1_len)?;
a178c22c
KS
895
896 self.render_stripe(true, start_line, end_line);
897
61cab15b 898 if self.v4_idx.is_empty() {
a178c22c 899 bw.write_byte(0x32)?;
1047e983
KS
900 bw.write_u24be((self.v1_idx.len() + 4) as u32)?;
901 bw.write_buf(self.v1_idx.as_slice())?;
a178c22c
KS
902 } else {
903 bw.write_byte(0x30)?;
904 bw.write_u24be(0)?;
905 let chunk_pos = bw.tell();
906 let mut v1_pos = 0;
907 let mut v4_pos = 0;
908 for _ in 0..32 {
909 self.v1_idx.push(0);
910 self.v4_idx.push(0);
911 self.v4_idx.push(0);
912 self.v4_idx.push(0);
913 self.v4_idx.push(0);
914 }
915 for mask in self.masks.masks.iter() {
916 bw.write_u32be(*mask)?;
917 for j in (0..32).rev() {
918 if (mask & (1 << j)) == 0 {
919 bw.write_byte(self.v1_idx[v1_pos])?;
920 v1_pos += 1;
921 } else {
922 bw.write_byte(self.v4_idx[v4_pos])?;
923 bw.write_byte(self.v4_idx[v4_pos + 1])?;
924 bw.write_byte(self.v4_idx[v4_pos + 2])?;
925 bw.write_byte(self.v4_idx[v4_pos + 3])?;
926 v4_pos += 4;
927 }
928 }
929 }
930 patch_size(bw, chunk_pos)?;
931 }
932
933 patch_size(bw, strip_data_pos)?;
934
9a57fda1
KS
935 self.v1_cb[self.cur_strip].copy_from_slice(&self.v1_cur_cb[self.cur_strip]);
936 self.v4_cb[self.cur_strip].copy_from_slice(&self.v4_cur_cb[self.cur_strip]);
a178c22c
KS
937 start_line = end_line;
938 end_line = (end_line + strip_h).min(height);
9a57fda1
KS
939
940 self.cur_strip += 1;
a178c22c
KS
941 }
942 patch_size(bw, frame_data_pos)?;
943 Ok(true)
944 }
945 fn encode_inter(&mut self, bw: &mut ByteWriter, in_frm: &NAVideoBuffer<u8>) -> EncoderResult<bool> {
946 let (width, height) = in_frm.get_dimensions(0);
947 let mut strip_h = (height / self.nstrips + 3) & !3;
948 if strip_h == 0 {
949 self.nstrips = 1;
950 strip_h = height;
951 }
952 let mut start_line = 0;
953 let mut end_line = strip_h;
954
955 bw.write_byte(1)?; // intra flag
956 bw.write_u24be(0)?; // frame size
957 let frame_data_pos = bw.tell();
958 bw.write_u16be(width as u16)?;
959 bw.write_u16be(height as u16)?;
960 bw.write_u16be(self.nstrips as u16)?;
961
9a57fda1 962 self.cur_strip = 0;
a178c22c
KS
963 while start_line < height {
964 self.read_strip(in_frm, start_line, end_line);
965 self.calc_skip_dist(in_frm, start_line, end_line);
966
3c406629 967 self.quant_vectors();
a178c22c 968 if self.grayscale {
9a57fda1 969 for cw in self.v1_cur_cb[self.cur_strip].iter_mut() {
a178c22c
KS
970 cw.u = 128;
971 cw.v = 128;
972 }
9a57fda1 973 for cw in self.v4_cur_cb[self.cur_strip].iter_mut() {
a178c22c
KS
974 cw.u = 128;
975 cw.v = 128;
976 }
977 }
978
37952415
KS
979 self.v1_idx.clear();
980 self.v4_idx.clear();
a178c22c
KS
981 self.masks.reset();
982
6064de24
KS
983 self.cmode.clear();
984 self.v1_cand.clear();
985 self.v4_cand.clear();
986
a178c22c
KS
987 let mut skip_iter = self.skip_dist.iter();
988 for (v1_entry, v4_entries) in self.v1_entries.iter().zip(self.v4_entries.chunks(4)) {
989 let skip_dist = *skip_iter.next().unwrap();
990 if skip_dist == 0 {
6064de24
KS
991 if !self.refine {
992 self.masks.put_inter(true);
993 } else {
994 self.cmode.push(CodingMode::Skip);
995 }
a178c22c
KS
996 continue;
997 }
9a57fda1 998 let (v1_idx, v1_dist) = Self::find_nearest(&self.v1_cur_cb[self.cur_strip][..self.v1_len], *v1_entry);
a178c22c 999 if skip_dist < v1_dist {
6064de24
KS
1000 if !self.refine {
1001 self.masks.put_inter(true);
1002 } else {
1003 self.cmode.push(CodingMode::Skip);
1004 }
a178c22c 1005 continue;
6064de24 1006 } else if !self.refine {
a178c22c
KS
1007 self.masks.put_inter(false);
1008 }
c6f14420 1009 if v1_dist == 0 || self.force_v1 {
6064de24
KS
1010 if !self.refine {
1011 self.masks.put_v1();
1012 self.v1_idx.push(v1_idx);
1013 } else {
1014 self.cmode.push(CodingMode::V1);
1015 self.v1_cand.push(*v1_entry);
1016 }
a178c22c
KS
1017 continue;
1018 }
9a57fda1
KS
1019 let (v40_idx, v40_dist) = Self::find_nearest(&self.v4_cur_cb[self.cur_strip][..self.v4_len], v4_entries[0]);
1020 let (v41_idx, v41_dist) = Self::find_nearest(&self.v4_cur_cb[self.cur_strip][..self.v4_len], v4_entries[1]);
1021 let (v42_idx, v42_dist) = Self::find_nearest(&self.v4_cur_cb[self.cur_strip][..self.v4_len], v4_entries[2]);
1022 let (v43_idx, v43_dist) = Self::find_nearest(&self.v4_cur_cb[self.cur_strip][..self.v4_len], v4_entries[3]);
a178c22c 1023 if v40_dist + v41_dist + v42_dist + v43_dist > v1_dist {
6064de24
KS
1024 if !self.refine {
1025 self.masks.put_v4();
1026 self.v4_idx.push(v40_idx);
1027 self.v4_idx.push(v41_idx);
1028 self.v4_idx.push(v42_idx);
1029 self.v4_idx.push(v43_idx);
1030 } else {
1031 self.cmode.push(CodingMode::V4);
1032 self.v4_cand.extend_from_slice(v4_entries);
1033 }
a178c22c 1034 } else {
6064de24
KS
1035 if !self.refine {
1036 self.masks.put_v1();
1037 self.v1_idx.push(v1_idx);
1038 } else {
1039 self.cmode.push(CodingMode::V1);
1040 self.v1_cand.push(*v1_entry);
1041 }
1042 }
1043 }
1044 if self.refine {
1045 self.refine_vectors();
1046 let mut v1_src = self.v1_cand.iter();
1047 let mut v4_src = self.v4_cand.chunks_exact(4);
1048 for &cmode in self.cmode.iter() {
1049 match cmode {
1050 CodingMode::Skip => {
1051 self.masks.put_inter(true);
1052 },
1053 CodingMode::V1 => {
1054 let v1_entry = v1_src.next().unwrap();
1055 let (v1_idx, _) = Self::find_nearest(&self.v1_cur_cb[self.cur_strip][..self.v1_len], *v1_entry);
1056 self.masks.put_inter(false);
1057 self.masks.put_v1();
1058 self.v1_idx.push(v1_idx);
1059 },
1060 CodingMode::V4 => {
1061 let v4_entries = v4_src.next().unwrap();
1062 let (v40_idx, _) = Self::find_nearest(&self.v4_cur_cb[self.cur_strip][..self.v4_len], v4_entries[0]);
1063 let (v41_idx, _) = Self::find_nearest(&self.v4_cur_cb[self.cur_strip][..self.v4_len], v4_entries[1]);
1064 let (v42_idx, _) = Self::find_nearest(&self.v4_cur_cb[self.cur_strip][..self.v4_len], v4_entries[2]);
1065 let (v43_idx, _) = Self::find_nearest(&self.v4_cur_cb[self.cur_strip][..self.v4_len], v4_entries[3]);
1066
1067 self.masks.put_inter(false);
1068 self.masks.put_v4();
1069 self.v4_idx.push(v40_idx);
1070 self.v4_idx.push(v41_idx);
1071 self.v4_idx.push(v42_idx);
1072 self.v4_idx.push(v43_idx);
1073 },
1074 };
a178c22c
KS
1075 }
1076 }
1077 self.masks.end();
1078
1079 let (upd_v1, upd_v4) = {
1080 let cb_size = if self.grayscale { 4 } else { 6 };
9f7faaf3
KS
1081 (Self::can_update_cb(&self.v1_cur_cb[self.cur_strip][..self.v1_len], &self.v1_cb[self.cur_strip][..self.v1_len], cb_size),
1082 Self::can_update_cb(&self.v4_cur_cb[self.cur_strip][..self.v4_len], &self.v4_cb[self.cur_strip][..self.v4_len], cb_size))
a178c22c
KS
1083 };
1084 bw.write_byte(0x11)?;
1085 bw.write_u24be(0)?; // strip size
1086 let strip_data_pos = bw.tell();
1087 bw.write_u16be(0)?; // yoff
1088 bw.write_u16be(0)?; // xoff
1089 bw.write_u16be((end_line - start_line) as u16)?;
1090 bw.write_u16be(width as u16)?;
1091
9f7faaf3
KS
1092 Self::write_cb(bw, 0x20, &self.v4_cur_cb[self.cur_strip], &self.v4_cb[self.cur_strip], self.grayscale, upd_v4, self.v4_len)?;
1093 Self::write_cb(bw, 0x22, &self.v1_cur_cb[self.cur_strip], &self.v1_cb[self.cur_strip], self.grayscale, upd_v1, self.v1_len)?;
a178c22c
KS
1094
1095 self.render_stripe(false, start_line, end_line);
1096
1097 bw.write_byte(0x31)?;
1098 bw.write_u24be(0)?;
1099 let chunk_pos = bw.tell();
1100 let mut v1_pos = 0;
1101 let mut v4_pos = 0;
1102 for _ in 0..32 {
1103 self.v1_idx.push(0);
1104 self.v4_idx.push(0);
1105 self.v4_idx.push(0);
1106 self.v4_idx.push(0);
1107 self.v4_idx.push(0);
1108 }
1109 let mut skip = true;
1110 for mask in self.masks.masks.iter() {
1111 bw.write_u32be(*mask)?;
10dedf0d 1112 if *mask == 0 && skip { continue; }
a178c22c
KS
1113 let mut bit = 1 << 31;
1114 while bit > 0 {
1115 if skip {
1116 skip = (mask & bit) == 0;
1117 bit >>= 1;
1118 } else {
1119 if (mask & bit) == 0 {
1120 bw.write_byte(self.v1_idx[v1_pos])?;
1121 v1_pos += 1;
1122 } else {
1123 bw.write_byte(self.v4_idx[v4_pos])?;
1124 bw.write_byte(self.v4_idx[v4_pos + 1])?;
1125 bw.write_byte(self.v4_idx[v4_pos + 2])?;
1126 bw.write_byte(self.v4_idx[v4_pos + 3])?;
1127 v4_pos += 4;
1128 }
1129 bit >>= 1;
1130 skip = true;
1131 }
1132 }
1133 }
1134 patch_size(bw, chunk_pos)?;
1135
1136 patch_size(bw, strip_data_pos)?;
1137
9a57fda1
KS
1138 self.v1_cb[self.cur_strip].copy_from_slice(&self.v1_cur_cb[self.cur_strip]);
1139 self.v4_cb[self.cur_strip].copy_from_slice(&self.v4_cur_cb[self.cur_strip]);
a178c22c
KS
1140 start_line = end_line;
1141 end_line = (end_line + strip_h).min(height);
9a57fda1
KS
1142
1143 self.cur_strip += 1;
a178c22c
KS
1144 }
1145 patch_size(bw, frame_data_pos)?;
1146 Ok(true)
1147 }
1148}
1149
1150impl NAEncoder for CinepakEncoder {
1151 fn negotiate_format(&self, encinfo: &EncodeParameters) -> EncoderResult<EncodeParameters> {
1152 match encinfo.format {
1153 NACodecTypeInfo::None => {
6f263099
KS
1154 Ok(EncodeParameters {
1155 format: NACodecTypeInfo::Video(NAVideoInfo::new(0, 0, true, YUV420_FORMAT)),
1156 ..Default::default()
1157 })
a178c22c 1158 },
61cab15b 1159 NACodecTypeInfo::Audio(_) => Err(EncoderError::FormatError),
a178c22c
KS
1160 NACodecTypeInfo::Video(vinfo) => {
1161 let pix_fmt = if vinfo.format == GRAY_FORMAT { GRAY_FORMAT } else { YUV420_FORMAT };
4abaf99e 1162 let outinfo = NAVideoInfo::new((vinfo.width + 3) & !3, (vinfo.height + 3) & !3, false, pix_fmt);
d722ffe9 1163 let mut ofmt = *encinfo;
a178c22c
KS
1164 ofmt.format = NACodecTypeInfo::Video(outinfo);
1165 Ok(ofmt)
1166 }
1167 }
1168 }
2757a028 1169 fn get_capabilities(&self) -> u64 { 0 }
a178c22c
KS
1170 fn init(&mut self, stream_id: u32, encinfo: EncodeParameters) -> EncoderResult<NAStreamRef> {
1171 match encinfo.format {
1172 NACodecTypeInfo::None => Err(EncoderError::FormatError),
1173 NACodecTypeInfo::Audio(_) => Err(EncoderError::FormatError),
1174 NACodecTypeInfo::Video(vinfo) => {
1175 if vinfo.format != YUV420_FORMAT && vinfo.format != GRAY_FORMAT {
1176 return Err(EncoderError::FormatError);
1177 }
1178 if ((vinfo.width | vinfo.height) & 3) != 0 {
1179 return Err(EncoderError::FormatError);
1180 }
1181 if (vinfo.width | vinfo.height) >= (1 << 16) {
1182 return Err(EncoderError::FormatError);
1183 }
1184
1185 let out_info = NAVideoInfo::new(vinfo.width, vinfo.height, false, vinfo.format);
61cab15b 1186 let info = NACodecInfo::new("cinepak", NACodecTypeInfo::Video(out_info), None);
a480a0de 1187 let mut stream = NAStream::new(StreamType::Video, stream_id, info, encinfo.tb_num, encinfo.tb_den, 0);
2ff56201
KS
1188 stream.set_num(stream_id as usize);
1189 let stream = stream.into_ref();
a178c22c
KS
1190
1191 self.stream = Some(stream.clone());
1192 self.quality = encinfo.quality;
1193 self.grayscale = vinfo.format != YUV420_FORMAT;
1194 let num_blocks = vinfo.width / 2 * vinfo.height / 2;
1195 self.v1_entries = Vec::with_capacity(num_blocks);
1196 self.v4_entries = Vec::with_capacity(num_blocks * 4);
1197 self.v1_idx = Vec::with_capacity(num_blocks);
1198 self.v4_idx = Vec::with_capacity(num_blocks * 4);
1199 self.skip_dist = Vec::with_capacity(vinfo.width / 4 * vinfo.height / 4);
1200
1201 let buf = alloc_video_buffer(out_info, 2)?;
1202 self.lastfrm = Some(buf.get_vbuf().unwrap());
c8db9313 1203
a178c22c
KS
1204 Ok(stream)
1205 },
1206 }
1207 }
1208 fn encode(&mut self, frm: &NAFrame) -> EncoderResult<()> {
1209 let buf = frm.get_buffer();
1210 if let Some(ref vbuf) = buf.get_vbuf() {
8c3f822a
KS
1211 if self.nstrips == 0 {
1212 let (w, h) = vbuf.get_dimensions(0);
1213 self.nstrips = ((((w * h) >> 4) + 1200) / 2400).max(1).min(3);
1214 let strip_h = ((h + self.nstrips - 1) / self.nstrips + 3) & !3;
1215 self.nstrips = (h + strip_h - 1) / strip_h;
1216 }
9a57fda1
KS
1217 let cur_strips = self.v1_cb.len();
1218 if cur_strips != self.nstrips {
1219 self.frmcount = 0;
1220 }
1221 if cur_strips < self.nstrips {
1222 for _ in cur_strips..self.nstrips {
1223 self.v1_cb.push([YUVCode::default(); 256]);
1224 self.v4_cb.push([YUVCode::default(); 256]);
1225 self.v1_cur_cb.push([YUVCode::default(); 256]);
1226 self.v4_cur_cb.push([YUVCode::default(); 256]);
1227 }
1228 }
1229
a178c22c
KS
1230 let mut dbuf = Vec::with_capacity(4);
1231 let mut gw = GrowableMemoryWriter::new_write(&mut dbuf);
1232 let mut bw = ByteWriter::new(&mut gw);
1233 let is_intra = if self.frmcount == 0 {
1234 self.encode_intra(&mut bw, vbuf)?
1235 } else {
1236 self.encode_inter(&mut bw, vbuf)?
1237 };
1238 self.pkt = Some(NAPacket::new(self.stream.clone().unwrap(), frm.ts, is_intra, dbuf));
1239 self.frmcount += 1;
3c406629 1240 if self.frmcount == self.key_int {
a178c22c
KS
1241 self.frmcount = 0;
1242 }
1243 Ok(())
1244 } else {
1245 Err(EncoderError::InvalidParameters)
1246 }
1247 }
1248 fn get_packet(&mut self) -> EncoderResult<Option<NAPacket>> {
1249 let mut npkt = None;
1250 std::mem::swap(&mut self.pkt, &mut npkt);
1251 Ok(npkt)
1252 }
1253 fn flush(&mut self) -> EncoderResult<()> {
1254 self.frmcount = 0;
1255 Ok(())
1256 }
1257}
1258
3c406629
KS
1259const ENCODER_OPTS: &[NAOptionDefinition] = &[
1260 NAOptionDefinition {
ee0ca773 1261 name: KEYFRAME_OPTION, description: KEYFRAME_OPTION_DESC,
3c406629
KS
1262 opt_type: NAOptionDefinitionType::Int(Some(0), Some(128)) },
1263 NAOptionDefinition {
1264 name: "nstrips", description: "Number of strips per frame (0 - automatic)",
1265 opt_type: NAOptionDefinitionType::Int(Some(0), Some(16)) },
1266 NAOptionDefinition {
1267 name: "quant_mode", description: "Quantisation mode",
42ef5325 1268 opt_type: NAOptionDefinitionType::String(Some(&["elbg", "fast", "mediancut"])) },
c6f14420
KS
1269 NAOptionDefinition {
1270 name: "force_v1", description: "Force coarse (V1-only) mode",
1271 opt_type: NAOptionDefinitionType::Bool },
6064de24
KS
1272 NAOptionDefinition {
1273 name: "refine", description: "Try to improve coded picture",
1274 opt_type: NAOptionDefinitionType::Bool },
3c406629
KS
1275];
1276
a178c22c 1277impl NAOptionHandler for CinepakEncoder {
3c406629
KS
1278 fn get_supported_options(&self) -> &[NAOptionDefinition] { ENCODER_OPTS }
1279 fn set_options(&mut self, options: &[NAOption]) {
1280 for option in options.iter() {
3c406629
KS
1281 for opt_def in ENCODER_OPTS.iter() {
1282 if opt_def.check(option).is_ok() {
1283 match option.name {
ee0ca773 1284 KEYFRAME_OPTION => {
3c406629
KS
1285 if let NAValue::Int(intval) = option.value {
1286 self.key_int = intval as u8;
1287 }
1288 },
1289 "nstrips" => {
1290 if let NAValue::Int(intval) = option.value {
1291 self.nstrips = intval as usize;
1292 }
1293 },
1294 "quant_mode" => {
405cec9e
KS
1295 if let NAValue::String(ref strval) = option.value {
1296 match strval.as_str() {
3c406629 1297 "elbg" => self.qmode = QuantMode::ELBG,
42ef5325 1298 "fast" => self.qmode = QuantMode::Fast,
3c406629
KS
1299 "mediancut" => self.qmode = QuantMode::MedianCut,
1300 _ => {},
1301 };
1302 }
1303 },
c6f14420
KS
1304 "force_v1" => {
1305 if let NAValue::Bool(val) = option.value {
1306 self.force_v1 = val;
1307 }
1308 },
6064de24
KS
1309 "refine" => {
1310 if let NAValue::Bool(val) = option.value {
1311 self.refine = val;
1312 }
1313 },
3c406629
KS
1314 _ => {},
1315 };
1316 }
1317 }
1318 }
1319 }
1320 fn query_option_value(&self, name: &str) -> Option<NAValue> {
1321 match name {
8211e0aa 1322 KEYFRAME_OPTION => Some(NAValue::Int(i64::from(self.key_int))),
3c406629
KS
1323 "nstrips" => Some(NAValue::Int(self.nstrips as i64)),
1324 "quant_mode" => Some(NAValue::String(self.qmode.to_string())),
c6f14420 1325 "force_v1" => Some(NAValue::Bool(self.force_v1)),
6064de24 1326 "refine" => Some(NAValue::Bool(self.refine)),
3c406629
KS
1327 _ => None,
1328 }
1329 }
a178c22c
KS
1330}
1331
1332pub fn get_encoder() -> Box<dyn NAEncoder + Send> {
1333 Box::new(CinepakEncoder::new())
1334}
1335
1336#[cfg(test)]
1337mod test {
1338 use nihav_core::codecs::*;
1339 use nihav_core::demuxers::*;
1340 use nihav_core::muxers::*;
1341 use crate::*;
1342 use nihav_codec_support::test::enc_video::*;
1343
1344 #[test]
1345 fn test_cinepak_encoder() {
1346 let mut dmx_reg = RegisteredDemuxers::new();
1347 generic_register_all_demuxers(&mut dmx_reg);
1348 let mut dec_reg = RegisteredDecoders::new();
78fb6560 1349 generic_register_all_decoders(&mut dec_reg);
a178c22c
KS
1350 let mut mux_reg = RegisteredMuxers::new();
1351 generic_register_all_muxers(&mut mux_reg);
1352 let mut enc_reg = RegisteredEncoders::new();
1353 generic_register_all_encoders(&mut enc_reg);
1354
886cde48 1355 // sample: https://samples.mplayerhq.hu/V-codecs/UCOD/TalkingHead_352x288.avi
a178c22c
KS
1356 let dec_config = DecoderTestParams {
1357 demuxer: "avi",
1358 in_name: "assets/Misc/TalkingHead_352x288.avi",
1359 stream_type: StreamType::Video,
1360 limit: Some(2),
1361 dmx_reg, dec_reg,
1362 };
1363 let enc_config = EncoderTestParams {
1364 muxer: "avi",
1365 enc_name: "cinepak",
1366 out_name: "cinepak.avi",
1367 mux_reg, enc_reg,
1368 };
1369 let dst_vinfo = NAVideoInfo {
1370 width: 0,
1371 height: 0,
1372 format: YUV420_FORMAT,
1373 flipped: true,
6bc499a0 1374 bits: 12,
a178c22c
KS
1375 };
1376 let enc_params = EncodeParameters {
1377 format: NACodecTypeInfo::Video(dst_vinfo),
1378 quality: 0,
1379 bitrate: 0,
1380 tb_num: 0,
1381 tb_den: 0,
1382 flags: 0,
1383 };
42ef5325
KS
1384 let enc_options = &[
1385 NAOption { name: "quant_mode", value: NAValue::String("mediancut".to_string()) },
1386 ];
1387 //test_encoding_to_file(&dec_config, &enc_config, enc_params, enc_options);
1388 test_encoding_md5(&dec_config, &enc_config, enc_params, enc_options,
9f7faaf3 1389 &[0x1d4690c8, 0x3b15b4b3, 0xc2df3c7b, 0x1a25b159]);
a178c22c
KS
1390 }
1391}