]>
Commit | Line | Data |
---|---|---|
8bde0eae KS |
1 | use nihav_core::codecs::*; |
2 | use nihav_core::io::byteio::*; | |
3 | use nihav_core::io::bitreader::*; | |
4 | use nihav_core::io::codebook::*; | |
5 | ||
6 | #[repr(u8)] | |
7 | enum TM2StreamType { | |
8 | CHigh = 0, | |
9 | CLow, | |
10 | LHigh, | |
11 | LLow, | |
12 | Update, | |
13 | Motion, | |
14 | BlockType, | |
15 | Num | |
16 | } | |
17 | ||
18 | #[repr(u8)] | |
19 | #[derive(Debug,Clone,Copy)] | |
20 | enum TM2BlockType { | |
21 | HiRes, | |
22 | MedRes, | |
23 | LowRes, | |
24 | NullRes, | |
25 | Update, | |
26 | Still, | |
27 | Motion | |
28 | } | |
29 | ||
30 | const TM2_BLOCK_TYPES: [TM2BlockType; 7] = [ | |
31 | TM2BlockType::HiRes, TM2BlockType::MedRes, TM2BlockType::LowRes, TM2BlockType::NullRes, | |
32 | TM2BlockType::Update, TM2BlockType::Still, TM2BlockType::Motion | |
33 | ]; | |
34 | ||
35 | trait ReadLenEsc { | |
36 | fn read_len_esc(&mut self) -> DecoderResult<usize>; | |
37 | } | |
38 | ||
39 | const TM2_ESCAPE: usize = 0x80000000; | |
40 | ||
41 | impl<'a> ReadLenEsc for ByteReader<'a> { | |
42 | fn read_len_esc(&mut self) -> DecoderResult<usize> { | |
43 | let len = self.read_u32le()? as usize; | |
44 | if len == TM2_ESCAPE { | |
45 | let len2 = self.read_u32le()? as usize; | |
46 | Ok(len2) | |
47 | } else { | |
48 | Ok(len) | |
49 | } | |
50 | } | |
51 | } | |
52 | ||
53 | struct HuffDef { | |
54 | val_bits: u8, | |
55 | max_bits: u8, | |
56 | nelems: usize, | |
57 | } | |
58 | ||
59 | impl HuffDef { | |
60 | fn read(&mut self, br: &mut BitReader, codes: &mut Vec<FullCodebookDesc<u8>>, prefix: u32, len: u8) -> DecoderResult<()> { | |
61 | validate!(len <= self.max_bits); | |
62 | if !br.read_bool()? { | |
63 | validate!(codes.len() < self.nelems); | |
64 | let sym = br.read(self.val_bits)? as u8; | |
65 | codes.push(FullCodebookDesc { code: prefix, bits: len, sym }); | |
66 | } else { | |
67 | self.read(br, codes, (prefix << 1) | 0, len + 1)?; | |
68 | self.read(br, codes, (prefix << 1) | 1, len + 1)?; | |
69 | } | |
70 | Ok(()) | |
71 | } | |
72 | } | |
73 | ||
74 | struct HuffTree { | |
75 | cb: Option<Codebook<u8>>, | |
76 | sym0: u8, | |
77 | } | |
78 | ||
79 | impl HuffTree { | |
80 | fn new() -> Self { | |
81 | Self { cb: None, sym0: 0 } | |
82 | } | |
83 | } | |
84 | ||
85 | const TM2_MAX_DELTAS: usize = 64; | |
86 | ||
87 | struct TM2Stream { | |
88 | tokens: Vec<u8>, | |
89 | deltas: [i32; TM2_MAX_DELTAS], | |
90 | pos: usize, | |
91 | } | |
92 | ||
93 | impl Default for TM2Stream { | |
94 | fn default() -> Self { | |
95 | Self { | |
96 | tokens: Vec::new(), | |
97 | deltas: [0; TM2_MAX_DELTAS], | |
98 | pos: 0, | |
99 | } | |
100 | } | |
101 | } | |
102 | ||
103 | impl TM2Stream { | |
104 | fn read_header(&mut self, src: &[u8], br: &mut ByteReader) -> DecoderResult<()> { | |
105 | self.tokens.truncate(0); | |
106 | self.pos = 0; | |
107 | ||
108 | let len = br.read_u32le()? as usize; | |
109 | let endpos = br.tell() + (len as u64) * 4; | |
110 | if len == 0 { | |
111 | return Ok(()); | |
112 | } | |
113 | let ntoks = br.read_u32le()? as usize; | |
114 | validate!(ntoks < (1 << 24)); | |
115 | if (ntoks & 1) != 0 { | |
116 | let dlen = br.read_len_esc()?; | |
117 | if (dlen as i32) > 0 { | |
118 | let rest_size = (endpos - br.tell()) as usize; | |
119 | let skip_size = self.read_deltas(&src[br.tell() as usize..][..rest_size])?; | |
120 | validate!(skip_size == dlen * 4); | |
121 | br.read_skip(skip_size)?; | |
122 | } | |
123 | } | |
124 | let _len = br.read_len_esc()?; | |
125 | let _algo = br.read_u32le()?; | |
126 | ||
127 | let mut htree = HuffTree::new(); | |
128 | let rest_size = (endpos - br.tell()) as usize; | |
129 | let skip_size = self.read_huff_tree(&src[br.tell() as usize..][..rest_size], &mut htree)?; | |
130 | br.read_skip(skip_size)?; | |
131 | ||
132 | let len = br.read_u32le()? as usize; | |
133 | validate!(br.tell() + (len as u64) * 4 <= endpos); | |
134 | if len > 0 { | |
135 | self.tokens.reserve(ntoks >> 1); | |
136 | let rest_size = (endpos - br.tell()) as usize; | |
137 | let skip_size = self.read_tokens(&src[br.tell() as usize..][..rest_size], &htree, ntoks >> 1)?; | |
138 | br.read_skip(skip_size)?; | |
139 | } else { | |
140 | self.tokens.resize(ntoks >> 1, htree.sym0); | |
141 | } | |
142 | ||
143 | ||
144 | let pos = br.tell(); | |
145 | validate!(pos <= endpos); | |
146 | let toskip = endpos - pos; | |
147 | br.read_skip(toskip as usize)?; | |
148 | ||
149 | Ok(()) | |
150 | } | |
151 | fn read_deltas(&mut self, src: &[u8]) -> DecoderResult<usize> { | |
152 | let mut br = BitReader::new(src, src.len(), BitReaderMode::LE32MSB); | |
153 | let coded_deltas = br.read(9)? as usize; | |
154 | let bits = br.read(5)? as u8; | |
155 | validate!((coded_deltas <= TM2_MAX_DELTAS) && (bits > 0)); | |
156 | let mask = 1 << (bits - 1); | |
157 | let bias = 1 << bits; | |
158 | self.deltas = [0; TM2_MAX_DELTAS]; | |
159 | for i in 0..coded_deltas { | |
160 | let val = br.read(bits)?; | |
161 | if (val & mask) != 0 { | |
162 | self.deltas[i] = (val as i32) - bias; | |
163 | } else { | |
164 | self.deltas[i] = val as i32; | |
165 | } | |
166 | } | |
167 | ||
168 | Ok(((br.tell() + 31) >> 5) << 2) | |
169 | } | |
170 | fn read_huff_tree(&mut self, src: &[u8], htree: &mut HuffTree) -> DecoderResult<usize> { | |
171 | let mut br = BitReader::new(src, src.len(), BitReaderMode::LE32MSB); | |
172 | ||
173 | let val_bits = br.read(5)? as u8; | |
174 | let max_bits = br.read(5)? as u8; | |
175 | let min_bits = br.read(5)? as u8; | |
176 | let nelems = br.read(17)? as usize; | |
177 | validate!(val_bits > 0 && val_bits <= 6); | |
178 | validate!(nelems > 0); | |
179 | validate!((max_bits < 25) && (min_bits <= max_bits)); | |
180 | ||
181 | let mut codes: Vec<FullCodebookDesc<u8>> = Vec::with_capacity(nelems); | |
182 | let mut hdef = HuffDef { val_bits, max_bits, nelems }; | |
183 | hdef.read(&mut br, &mut codes, 0, 0)?; | |
184 | htree.sym0 = codes[0].sym; | |
185 | if nelems > 1 { | |
186 | let mut cr = FullCodebookDescReader::new(codes); | |
187 | htree.cb = Some(Codebook::new(&mut cr, CodebookMode::MSB)?); | |
188 | } | |
189 | ||
190 | Ok(((br.tell() + 31) >> 5) << 2) | |
191 | } | |
192 | fn read_tokens(&mut self, src: &[u8], htree: &HuffTree, ntoks: usize) -> DecoderResult<usize> { | |
193 | let mut br = BitReader::new(src, src.len(), BitReaderMode::LE32MSB); | |
194 | ||
195 | if let Some(ref cb) = htree.cb { | |
196 | for _ in 0..ntoks { | |
197 | let tok = br.read_cb(cb)?; | |
198 | self.tokens.push(tok); | |
199 | } | |
200 | } | |
201 | ||
202 | Ok(((br.tell() + 31) >> 5) << 2) | |
203 | } | |
204 | ||
205 | fn get_block_type(&mut self) -> DecoderResult<u8> { | |
206 | validate!(self.pos < self.tokens.len()); | |
207 | let res = self.tokens[self.pos]; | |
208 | self.pos += 1; | |
209 | Ok(res) | |
210 | } | |
211 | fn get_token(&mut self) -> DecoderResult<i32> { | |
212 | validate!(self.pos < self.tokens.len()); | |
213 | let idx = self.tokens[self.pos] as usize; | |
214 | validate!(idx < TM2_MAX_DELTAS); | |
215 | self.pos += 1; | |
216 | Ok(self.deltas[idx]) | |
217 | } | |
218 | } | |
219 | ||
220 | #[derive(Default)] | |
221 | struct DeltaState { | |
222 | dy: [i32; 4], | |
223 | dc: [[i32; 2]; 2], | |
224 | } | |
225 | ||
226 | impl DeltaState { | |
227 | fn apply_y(&mut self, dst: &mut [u8], mut yoff: usize, ystride: usize, ydeltas: &[i32; 16], last: &mut [i32]) { | |
228 | for y in 0..4 { | |
229 | let mut d = self.dy[y]; | |
230 | for x in 0..4 { | |
231 | d += ydeltas[x + y * 4]; | |
232 | last[x] += d; | |
233 | dst[yoff + x] = last[x].max(0).min(255) as u8; | |
234 | } | |
235 | self.dy[y] = d; | |
236 | yoff += ystride; | |
237 | } | |
238 | } | |
239 | fn apply_c(&mut self, dst: &mut [i16], mut coff: usize, cstride: usize, cdeltas: &[i32; 4], idx: usize, last: &mut [i32]) { | |
240 | for y in 0..2 { | |
241 | let mut d = self.dc[idx][y]; | |
242 | for x in 0..2 { | |
243 | d += cdeltas[x + y * 2]; | |
244 | last[x] += d; | |
245 | dst[coff + x] = last[x] as i16; | |
246 | } | |
247 | self.dc[idx][y] = d; | |
248 | coff += cstride; | |
249 | } | |
250 | } | |
251 | fn interpolate_y_low(&mut self, last: &mut [i32]) { | |
252 | let dsum = self.dy[0] + self.dy[1] + self.dy[2] + self.dy[3]; | |
253 | last[1] = (last[0] - dsum + last[2]) >> 1; | |
254 | last[3] = (last[2] + last[4]) >> 1; | |
255 | ||
256 | let t0 = self.dy[0] + self.dy[1]; | |
257 | let t1 = self.dy[2] + self.dy[3]; | |
258 | self.dy[0] = t0 >> 1; | |
259 | self.dy[1] = t0 - (t0 >> 1); | |
260 | self.dy[2] = t1 >> 1; | |
261 | self.dy[3] = t1 - (t1 >> 1); | |
262 | } | |
263 | fn interpolate_y_null(&mut self, last: &mut [i32]) { | |
264 | let dsum = self.dy[0] + self.dy[1] + self.dy[2] + self.dy[3]; | |
265 | let left = last[0] - dsum; | |
266 | let right = last[4]; | |
267 | let diff = right - left; | |
268 | last[1] = left + (diff >> 2); | |
269 | last[2] = left + (diff >> 1); | |
270 | last[3] = right - (diff >> 2); | |
271 | ||
272 | let mut sum = left; | |
273 | self.dy[0] = (left + (dsum >> 2)) - sum; | |
274 | sum += self.dy[0]; | |
275 | self.dy[1] = (left + (dsum >> 1)) - sum; | |
276 | sum += self.dy[1]; | |
277 | self.dy[2] = (left + dsum - (dsum >> 2)) - sum; | |
278 | sum += self.dy[2]; | |
279 | self.dy[3] = (left + dsum) - sum; | |
280 | } | |
281 | fn interpolate_c(&mut self, idx: usize, last: &mut [i32]) { | |
282 | let dsum = self.dc[idx][0] + self.dc[idx][1]; | |
283 | let l = (last[0] + last[2] - dsum) >> 1; | |
284 | self.dc[idx][0] = dsum >> 1; | |
285 | self.dc[idx][1] = dsum - (dsum >> 1); | |
286 | last[1] = l; | |
287 | } | |
288 | fn recalc_y(&mut self, dst: &[u8], yoff: usize, ystride: usize, last: &mut [i32]) { | |
289 | let src = &dst[yoff+3..]; | |
290 | self.dy[0] = (src[ystride * 0] as i32) - last[3]; | |
291 | self.dy[1] = (src[ystride * 1] as i32) - (src[ystride * 0] as i32); | |
292 | self.dy[2] = (src[ystride * 2] as i32) - (src[ystride * 1] as i32); | |
293 | self.dy[3] = (src[ystride * 3] as i32) - (src[ystride * 2] as i32); | |
294 | let src = &dst[yoff + 3 * ystride..]; | |
295 | for x in 0..4 { | |
296 | last[x] = src[x] as i32; | |
297 | } | |
298 | } | |
299 | fn recalc_c(&mut self, dst: &[i16], coff: usize, cstride: usize, idx: usize, last: &mut [i32]) { | |
300 | self.dc[idx][0] = (dst[coff + 1] as i32) - last[1]; | |
301 | self.dc[idx][1] = (dst[coff + 1 + cstride] as i32) - (dst[coff + 1] as i32); | |
302 | last[0] = dst[coff + cstride + 0] as i32; | |
303 | last[1] = dst[coff + cstride + 1] as i32; | |
304 | } | |
305 | } | |
306 | ||
307 | #[derive(Default)] | |
308 | struct TM2Frame { | |
309 | ydata: Vec<u8>, | |
310 | udata: Vec<i16>, | |
311 | vdata: Vec<i16>, | |
312 | ystride: usize, | |
313 | cstride: usize, | |
314 | } | |
315 | ||
316 | impl TM2Frame { | |
317 | fn alloc(width: usize, height: usize) -> Self { | |
318 | let ystride = (width + 3) & !3; | |
319 | let ysize = ystride * ((height + 3) & !3); | |
320 | let mut ydata = Vec::with_capacity(ysize); | |
321 | ydata.resize(ysize, 0); | |
322 | let cstride = ystride >> 1; | |
323 | let csize = cstride * (((height + 3) & !3) >> 1); | |
324 | let mut udata = Vec::with_capacity(csize); | |
325 | udata.resize(csize, 0); | |
326 | let mut vdata = Vec::with_capacity(csize); | |
327 | vdata.resize(csize, 0); | |
328 | Self { ydata, udata, vdata, ystride, cstride } | |
329 | } | |
330 | } | |
331 | ||
332 | #[derive(Default)] | |
333 | struct TM2Decoder { | |
2422d969 | 334 | info: NACodecInfoRef, |
8bde0eae KS |
335 | streams: [TM2Stream; TM2StreamType::Num as usize], |
336 | width: usize, | |
337 | height: usize, | |
338 | cur_frame: TM2Frame, | |
339 | prev_frame: TM2Frame, | |
340 | } | |
341 | ||
342 | impl TM2Decoder { | |
343 | fn new() -> Self { Self::default() } | |
344 | fn decode_blocks(&mut self) -> DecoderResult<bool> { | |
345 | let ydst = &mut self.cur_frame.ydata; | |
346 | let udst = &mut self.cur_frame.udata; | |
347 | let vdst = &mut self.cur_frame.vdata; | |
348 | let ystride = self.cur_frame.ystride; | |
349 | let cstride = self.cur_frame.cstride; | |
350 | let mut offs: [usize; 2] = [0; 2]; | |
351 | let mut is_intra = true; | |
352 | ||
353 | let bw = self.width >> 2; | |
354 | let bh = self.height >> 2; | |
355 | validate!(self.streams[TM2StreamType::BlockType as usize].tokens.len() == bw * bh); | |
356 | ||
357 | let mut ydeltas: [i32; 16] = [0; 16]; | |
358 | let mut cdeltas: [[i32; 4]; 2] = [[0; 4]; 2]; | |
359 | let mut lasty: Vec<i32> = Vec::with_capacity(self.width + 1); | |
360 | lasty.resize(self.width + 1, 0); | |
361 | let mut lastu: Vec<i32> = Vec::with_capacity(self.width/2 + 1); | |
362 | lastu.resize(self.width/2 + 1, 0); | |
363 | let mut lastv: Vec<i32> = Vec::with_capacity(self.width/2 + 1); | |
364 | lastv.resize(self.width/2 + 1, 0); | |
365 | for by in 0..bh { | |
366 | let mut dstate = DeltaState::default(); | |
367 | for bx in 0..bw { | |
368 | let bidx = self.streams[TM2StreamType::BlockType as usize].get_block_type()? as usize; | |
369 | validate!(bidx < TM2_BLOCK_TYPES.len()); | |
370 | let btype = TM2_BLOCK_TYPES[bidx]; | |
371 | match btype { | |
372 | TM2BlockType::HiRes => { | |
373 | for i in 0..4 { | |
374 | cdeltas[0][i] = self.streams[TM2StreamType::CHigh as usize].get_token()?; | |
375 | cdeltas[1][i] = self.streams[TM2StreamType::CHigh as usize].get_token()?; | |
376 | } | |
377 | dstate.apply_c(udst, offs[1] + bx * 2, cstride, &cdeltas[0], 0, &mut lastu[bx*2+1..]); | |
378 | dstate.apply_c(vdst, offs[1] + bx * 2, cstride, &cdeltas[1], 1, &mut lastv[bx*2+1..]); | |
379 | for i in 0..4*4 { | |
380 | ydeltas[i] = self.streams[TM2StreamType::LHigh as usize].get_token()?; | |
381 | } | |
382 | dstate.apply_y(ydst, offs[0] + bx * 4, ystride, &ydeltas, &mut lasty[bx*4+1..]); | |
383 | }, | |
384 | TM2BlockType::MedRes => { | |
385 | cdeltas = [[0; 4]; 2]; | |
386 | cdeltas[0][0] = self.streams[TM2StreamType::CLow as usize].get_token()?; | |
387 | cdeltas[1][0] = self.streams[TM2StreamType::CLow as usize].get_token()?; | |
388 | dstate.interpolate_c(0, &mut lastu[bx*2..]); | |
389 | dstate.apply_c(udst, offs[1] + bx * 2, cstride, &cdeltas[0], 0, &mut lastu[bx*2+1..]); | |
390 | dstate.interpolate_c(1, &mut lastv[bx*2..]); | |
391 | dstate.apply_c(vdst, offs[1] + bx * 2, cstride, &cdeltas[1], 1, &mut lastv[bx*2+1..]); | |
392 | for i in 0..4*4 { | |
393 | ydeltas[i] = self.streams[TM2StreamType::LHigh as usize].get_token()?; | |
394 | } | |
395 | dstate.apply_y(ydst, offs[0] + bx * 4, ystride, &ydeltas, &mut lasty[bx*4+1..]); | |
396 | }, | |
397 | TM2BlockType::LowRes => { | |
398 | cdeltas = [[0; 4]; 2]; | |
399 | cdeltas[0][0] = self.streams[TM2StreamType::CLow as usize].get_token()?; | |
400 | cdeltas[1][0] = self.streams[TM2StreamType::CLow as usize].get_token()?; | |
401 | dstate.interpolate_c(0, &mut lastu[bx*2..]); | |
402 | dstate.apply_c(udst, offs[1] + bx * 2, cstride, &cdeltas[0], 0, &mut lastu[bx*2+1..]); | |
403 | dstate.interpolate_c(1, &mut lastv[bx*2..]); | |
404 | dstate.apply_c(vdst, offs[1] + bx * 2, cstride, &cdeltas[1], 1, &mut lastv[bx*2+1..]); | |
405 | ydeltas = [0; 16]; | |
406 | ydeltas[ 0] = self.streams[TM2StreamType::LLow as usize].get_token()?; | |
407 | ydeltas[ 2] = self.streams[TM2StreamType::LLow as usize].get_token()?; | |
408 | ydeltas[ 8] = self.streams[TM2StreamType::LLow as usize].get_token()?; | |
409 | ydeltas[10] = self.streams[TM2StreamType::LLow as usize].get_token()?; | |
410 | dstate.interpolate_y_low(&mut lasty[bx*4..]); | |
411 | dstate.apply_y(ydst, offs[0] + bx * 4, ystride, &ydeltas, &mut lasty[bx*4+1..]); | |
412 | }, | |
413 | TM2BlockType::NullRes => { | |
414 | cdeltas = [[0; 4]; 2]; | |
415 | dstate.interpolate_c(0, &mut lastu[bx*2..]); | |
416 | dstate.apply_c(udst, offs[1] + bx * 2, cstride, &cdeltas[0], 0, &mut lastu[bx*2+1..]); | |
417 | dstate.interpolate_c(1, &mut lastv[bx*2..]); | |
418 | dstate.apply_c(vdst, offs[1] + bx * 2, cstride, &cdeltas[1], 1, &mut lastv[bx*2+1..]); | |
419 | ydeltas = [0; 16]; | |
420 | dstate.interpolate_y_null(&mut lasty[bx*4..]); | |
421 | dstate.apply_y(ydst, offs[0] + bx * 4, ystride, &ydeltas, &mut lasty[bx*4+1..]); | |
422 | }, | |
423 | TM2BlockType::Update => { | |
424 | is_intra = false; | |
425 | ||
426 | let mut coff = offs[1] + bx * 2; | |
427 | let usrc = &self.prev_frame.udata; | |
428 | let vsrc = &self.prev_frame.vdata; | |
429 | for _ in 0..2 { | |
430 | for x in 0..2 { | |
431 | let du = self.streams[TM2StreamType::Update as usize].get_token()?; | |
432 | let dv = self.streams[TM2StreamType::Update as usize].get_token()?; | |
433 | udst[coff + x] = usrc[coff + x] + (du as i16); | |
434 | vdst[coff + x] = vsrc[coff + x] + (dv as i16); | |
435 | } | |
436 | coff += cstride; | |
437 | } | |
438 | dstate.recalc_c(udst, offs[1] + bx * 2, cstride, 0, &mut lastu[bx*2+1..]); | |
439 | dstate.recalc_c(vdst, offs[1] + bx * 2, cstride, 1, &mut lastv[bx*2+1..]); | |
440 | let mut yoff = offs[0] + bx * 4; | |
441 | let ysrc = &self.prev_frame.ydata; | |
442 | for _ in 0..4 { | |
443 | for x in 0..4 { | |
444 | let dy = self.streams[TM2StreamType::Update as usize].get_token()?; | |
445 | ydst[yoff + x] = ((ysrc[yoff + x] as i32) + dy) as u8; | |
446 | } | |
447 | yoff += ystride; | |
448 | } | |
449 | dstate.recalc_y(ydst, offs[0] + bx * 4, ystride, &mut lasty[bx*4+1..]); | |
450 | }, | |
451 | TM2BlockType::Still => { | |
452 | is_intra = false; | |
453 | ||
454 | let mut coff = offs[1] + bx * 2; | |
455 | let usrc = &self.prev_frame.udata; | |
456 | let vsrc = &self.prev_frame.vdata; | |
457 | for _ in 0..2 { | |
458 | for x in 0..2 { | |
459 | udst[coff + x] = usrc[coff + x]; | |
460 | vdst[coff + x] = vsrc[coff + x]; | |
461 | } | |
462 | coff += cstride; | |
463 | } | |
464 | dstate.recalc_c(udst, offs[1] + bx * 2, cstride, 0, &mut lastu[bx*2+1..]); | |
465 | dstate.recalc_c(vdst, offs[1] + bx * 2, cstride, 1, &mut lastv[bx*2+1..]); | |
466 | let mut yoff = offs[0] + bx * 4; | |
467 | let ysrc = &self.prev_frame.ydata; | |
468 | for _ in 0..4 { | |
469 | for x in 0..4 { | |
470 | ydst[yoff + x] = ysrc[yoff + x]; | |
471 | } | |
472 | yoff += ystride; | |
473 | } | |
474 | dstate.recalc_y(ydst, offs[0] + bx * 4, ystride, &mut lasty[bx*4+1..]); | |
475 | }, | |
476 | TM2BlockType::Motion => { | |
477 | is_intra = false; | |
478 | ||
479 | let mx = self.streams[TM2StreamType::Motion as usize].get_token()?; | |
480 | let my = self.streams[TM2StreamType::Motion as usize].get_token()?; | |
481 | let xpos = (((bx as i32) * 4) + mx).max(0).min((self.width - 4) as i32) as usize; | |
482 | let ypos = (((by as i32) * 4) + my).max(0).min((self.height - 4) as i32) as usize; | |
483 | let mut coff = offs[1] + bx * 2; | |
484 | let mut csoff = (xpos >> 1) + (ypos >> 1) * cstride; | |
485 | let usrc = &self.prev_frame.udata; | |
486 | let vsrc = &self.prev_frame.vdata; | |
487 | for _ in 0..2 { | |
488 | for x in 0..2 { | |
489 | udst[coff + x] = usrc[csoff + x]; | |
490 | vdst[coff + x] = vsrc[csoff + x]; | |
491 | } | |
492 | coff += cstride; | |
493 | csoff += cstride; | |
494 | } | |
495 | dstate.recalc_c(udst, offs[1] + bx * 2, cstride, 0, &mut lastu[bx*2+1..]); | |
496 | dstate.recalc_c(vdst, offs[1] + bx * 2, cstride, 1, &mut lastv[bx*2+1..]); | |
497 | let mut yoff = offs[0] + bx * 4; | |
498 | let mut ysoff = xpos + ypos * ystride; | |
499 | let ysrc = &self.prev_frame.ydata; | |
500 | for _ in 0..4 { | |
501 | for x in 0..4 { | |
502 | ydst[yoff + x] = ysrc[ysoff + x]; | |
503 | } | |
504 | yoff += ystride; | |
505 | ysoff += ystride; | |
506 | } | |
507 | dstate.recalc_y(ydst, offs[0] + bx * 4, ystride, &mut lasty[bx*4+1..]); | |
508 | }, | |
509 | }; | |
510 | } | |
511 | offs[0] += ystride * 4; | |
512 | offs[1] += cstride * 2; | |
513 | } | |
514 | ||
515 | Ok(is_intra) | |
516 | } | |
517 | fn output_frame(&mut self, buf: &mut NAVideoBuffer<u8>) { | |
518 | let fmt = buf.get_info().get_format(); | |
519 | let offs = [fmt.get_chromaton(0).unwrap().get_offset() as usize, | |
520 | fmt.get_chromaton(1).unwrap().get_offset() as usize, | |
521 | fmt.get_chromaton(2).unwrap().get_offset() as usize]; | |
522 | let stride = buf.get_stride(0); | |
1a967e6b | 523 | let data = buf.get_data_mut().unwrap(); |
8bde0eae KS |
524 | let dst = data.as_mut_slice(); |
525 | ||
526 | let mut off = 0; | |
527 | let mut ysrc = 0; | |
528 | let mut csrc = 0; | |
529 | for y in 0..self.height { | |
530 | let out = &mut dst[off..]; | |
531 | for (x, pic) in out.chunks_exact_mut(3).take(self.width).enumerate() { | |
532 | let y = self.cur_frame.ydata[ysrc + x] as i16; | |
533 | let u = self.cur_frame.udata[csrc + (x >> 1)]; | |
534 | let v = self.cur_frame.vdata[csrc + (x >> 1)]; | |
535 | pic[offs[0]] = (y + u).max(0).min(255) as u8; | |
536 | pic[offs[1]] = y.max(0).min(255) as u8; | |
537 | pic[offs[2]] = (y + v).max(0).min(255) as u8; | |
538 | } | |
539 | off += stride; | |
540 | ysrc += self.cur_frame.ystride; | |
541 | if (y & 1) != 0 { | |
542 | csrc += self.cur_frame.cstride; | |
543 | } | |
544 | } | |
545 | } | |
546 | } | |
547 | ||
548 | impl NADecoder for TM2Decoder { | |
01613464 | 549 | fn init(&mut self, _supp: &mut NADecoderSupport, info: NACodecInfoRef) -> DecoderResult<()> { |
8bde0eae KS |
550 | if let NACodecTypeInfo::Video(vinfo) = info.get_properties() { |
551 | let myinfo = NACodecTypeInfo::Video(NAVideoInfo::new(vinfo.get_width(), vinfo.get_height(), false, YUV410_FORMAT)); | |
552 | self.width = vinfo.get_width(); | |
553 | self.height = vinfo.get_height(); | |
554 | self.cur_frame = TM2Frame::alloc(self.width, self.height); | |
555 | self.prev_frame = TM2Frame::alloc(self.width, self.height); | |
2422d969 | 556 | self.info = NACodecInfo::new_ref(info.get_name(), myinfo, info.get_extradata()).into_ref(); |
8bde0eae KS |
557 | Ok(()) |
558 | } else { | |
559 | Err(DecoderError::InvalidData) | |
560 | } | |
561 | } | |
01613464 | 562 | fn decode(&mut self, _supp: &mut NADecoderSupport, pkt: &NAPacket) -> DecoderResult<NAFrameRef> { |
8bde0eae KS |
563 | let src = pkt.get_buffer(); |
564 | validate!(src.len() >= 40 + (TM2StreamType::Num as usize) * 4 + 4); | |
565 | let mut mr = MemoryReader::new_read(&src); | |
566 | let mut br = ByteReader::new(&mut mr); | |
567 | ||
568 | let magic = br.read_u32be()?; | |
569 | validate!(magic == 0x100 || magic == 0x101); | |
570 | br.read_skip(36)?; | |
571 | for str in self.streams.iter_mut() { | |
572 | str.read_header(&src, &mut br)?; | |
573 | } | |
574 | ||
575 | let myinfo = NAVideoInfo::new(self.width, self.height, false, RGB24_FORMAT); | |
576 | let bufret = alloc_video_buffer(myinfo, 2); | |
577 | if let Err(_) = bufret { return Err(DecoderError::InvalidData); } | |
b70cc006 | 578 | let bufinfo = bufret.unwrap(); |
8bde0eae KS |
579 | let mut buf = bufinfo.get_vbuf().unwrap(); |
580 | ||
581 | let is_intra = self.decode_blocks()?; | |
582 | self.output_frame(&mut buf); | |
583 | std::mem::swap(&mut self.cur_frame, &mut self.prev_frame); | |
584 | ||
585 | let mut frm = NAFrame::new_from_pkt(pkt, self.info.clone(), bufinfo); | |
586 | frm.set_keyframe(is_intra); | |
587 | frm.set_frame_type(if is_intra { FrameType::I } else { FrameType::P }); | |
171860fc | 588 | Ok(frm.into_ref()) |
8bde0eae KS |
589 | } |
590 | } | |
591 | ||
6011e201 | 592 | pub fn get_decoder() -> Box<dyn NADecoder> { |
8bde0eae KS |
593 | Box::new(TM2Decoder::new()) |
594 | } | |
595 | ||
596 | #[cfg(test)] | |
597 | mod test { | |
598 | use nihav_core::codecs::RegisteredDecoders; | |
599 | use nihav_core::demuxers::RegisteredDemuxers; | |
600 | use nihav_core::test::dec_video::*; | |
601 | use crate::codecs::duck_register_all_codecs; | |
602 | use nihav_commonfmt::demuxers::generic_register_all_demuxers; | |
603 | #[test] | |
604 | fn test_tm2() { | |
605 | let mut dmx_reg = RegisteredDemuxers::new(); | |
606 | generic_register_all_demuxers(&mut dmx_reg); | |
607 | let mut dec_reg = RegisteredDecoders::new(); | |
608 | duck_register_all_codecs(&mut dec_reg); | |
609 | ||
610 | test_file_decoding("avi", "assets/Duck/tm20.avi", Some(16), true, false, None/*Some("tm2")*/, &dmx_reg, &dec_reg); | |
611 | } | |
612 | } |