diff options
Diffstat (limited to 'libardrone/arvideo.py')
-rw-r--r-- | libardrone/arvideo.py | 579 |
1 files changed, 579 insertions, 0 deletions
diff --git a/libardrone/arvideo.py b/libardrone/arvideo.py new file mode 100644 index 0000000..4ae011c --- /dev/null +++ b/libardrone/arvideo.py @@ -0,0 +1,579 @@ +#!/usr/bin/env python + +# Copyright (c) 2011 Bastian Venthur +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. + + +""" +Video decoding for the AR.Drone. + +This library uses psyco to speed-up the decoding process. It is however written +in a way that it works also without psyco installed. On the author's +development machine the speed up is from 2FPS w/o psyco to > 20 FPS w/ psyco. +""" + + +import array +import cProfile +import datetime +import struct +import sys + +try: + import psyco +except ImportError: + print "Please install psyco for better video decoding performance." + + +# from zig-zag back to normal +ZIG_ZAG_POSITIONS = array.array('B', + ( 0, 1, 8, 16, 9, 2, 3, 10, + 17, 24, 32, 25, 18, 11, 4, 5, + 12, 19, 26, 33, 40, 48, 41, 34, + 27, 20, 13, 6, 7, 14, 21, 28, + 35, 42, 49, 56, 57, 50, 43, 36, + 29, 22, 15, 23, 30, 37, 44, 51, + 58, 59, 52, 45, 38, 31, 39, 46, + 53, 60, 61, 54, 47, 55, 62, 63)) + +# Inverse quantization +IQUANT_TAB = array.array('B', + ( 3, 5, 7, 9, 11, 13, 15, 17, + 5, 7, 9, 11, 13, 15, 17, 19, + 7, 9, 11, 13, 15, 17, 19, 21, + 9, 11, 13, 15, 17, 19, 21, 23, + 11, 13, 15, 17, 19, 21, 23, 25, + 13, 15, 17, 19, 21, 23, 25, 27, + 15, 17, 19, 21, 23, 25, 27, 29, + 17, 19, 21, 23, 25, 27, 29, 31)) + +# Used for upscaling the 8x8 b- and r-blocks to 16x16 +SCALE_TAB = array.array('B', + ( 0, 0, 1, 1, 2, 2, 3, 3, + 0, 0, 1, 1, 2, 2, 3, 3, + 8, 8, 9, 9, 10, 10, 11, 11, + 8, 8, 9, 9, 10, 10, 11, 11, + 16, 16, 17, 17, 18, 18, 19, 19, + 16, 16, 17, 17, 18, 18, 19, 19, + 24, 24, 25, 25, 26, 26, 27, 27, + 24, 24, 25, 25, 26, 26, 27, 27, + + 4, 4, 5, 5, 6, 6, 7, 7, + 4, 4, 5, 5, 6, 6, 7, 7, + 12, 12, 13, 13, 14, 14, 15, 15, + 12, 12, 13, 13, 14, 14, 15, 15, + 20, 20, 21, 21, 22, 22, 23, 23, + 20, 20, 21, 21, 22, 22, 23, 23, + 28, 28, 29, 29, 30, 30, 31, 31, + 28, 28, 29, 29, 30, 30, 31, 31, + + 32, 32, 33, 33, 34, 34, 35, 35, + 32, 32, 33, 33, 34, 34, 35, 35, + 40, 40, 41, 41, 42, 42, 43, 43, + 40, 40, 41, 41, 42, 42, 43, 43, + 48, 48, 49, 49, 50, 50, 51, 51, + 48, 48, 49, 49, 50, 50, 51, 51, + 56, 56, 57, 57, 58, 58, 59, 59, + 56, 56, 57, 57, 58, 58, 59, 59, + + 36, 36, 37, 37, 38, 38, 39, 39, + 36, 36, 37, 37, 38, 38, 39, 39, + 44, 44, 45, 45, 46, 46, 47, 47, + 44, 44, 45, 45, 46, 46, 47, 47, + 52, 52, 53, 53, 54, 54, 55, 55, + 52, 52, 53, 53, 54, 54, 55, 55, + 60, 60, 61, 61, 62, 62, 63, 63, + 60, 60, 61, 61, 62, 62, 63, 63)) + +# Count leading zeros look up table +CLZLUT = array.array('B', + (8, 7, 6, 6, 5, 5, 5, 5, 4, 4, 4, 4, 4, 4, 4, 4, + 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, + 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, + 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0)) + +# Map pixels from four 8x8 blocks to one 16x16 +MB_TO_GOB_MAP = array.array('B', + [ 0, 1, 2, 3, 4, 5, 6, 7, + 16, 17, 18, 19, 20, 21, 22, 23, + 32, 33, 34, 35, 36, 37, 38, 39, + 48, 49, 50, 51, 52, 53, 54, 55, + 64, 65, 66, 67, 68, 69, 70, 71, + 80, 81, 82, 83, 84, 85, 86, 87, + 96, 97, 98, 99, 100, 101, 102, 103, + 112, 113, 114, 115, 116, 117, 118, 119, + 8, 9, 10, 11, 12, 13, 14, 15, + 24, 25, 26, 27, 28, 29, 30, 31, + 40, 41, 42, 43, 44, 45, 46, 47, + 56, 57, 58, 59, 60, 61, 62, 63, + 72, 73, 74, 75, 76, 77, 78, 79, + 88, 89, 90, 91, 92, 93, 94, 95, + 104, 105, 106, 107, 108, 109, 110, 111, + 120, 121, 122, 123, 124, 125, 126, 127, + 128, 129, 130, 131, 132, 133, 134, 135, + 144, 145, 146, 147, 148, 149, 150, 151, + 160, 161, 162, 163, 164, 165, 166, 167, + 176, 177, 178, 179, 180, 181, 182, 183, + 192, 193, 194, 195, 196, 197, 198, 199, + 208, 209, 210, 211, 212, 213, 214, 215, + 224, 225, 226, 227, 228, 229, 230, 231, + 240, 241, 242, 243, 244, 245, 246, 247, + 136, 137, 138, 139, 140, 141, 142, 143, + 152, 153, 154, 155, 156, 157, 158, 159, + 168, 169, 170, 171, 172, 173, 174, 175, + 184, 185, 186, 187, 188, 189, 190, 191, + 200, 201, 202, 203, 204, 205, 206, 207, + 216, 217, 218, 219, 220, 221, 222, 223, + 232, 233, 234, 235, 236, 237, 238, 239, + 248, 249, 250, 251, 252, 253, 254, 255]) +MB_ROW_MAP = array.array('B', [i / 16 for i in MB_TO_GOB_MAP]) +MB_COL_MAP = array.array('B', [i % 16 for i in MB_TO_GOB_MAP]) + +# An array of zeros. It is much faster to take the zeros from here than to +# generate a new list when needed. +ZEROS = array.array('i', [0 for i in range(256)]) + +# Constants needed for the inverse discrete cosine transform. +FIX_0_298631336 = 2446 +FIX_0_390180644 = 3196 +FIX_0_541196100 = 4433 +FIX_0_765366865 = 6270 +FIX_0_899976223 = 7373 +FIX_1_175875602 = 9633 +FIX_1_501321110 = 12299 +FIX_1_847759065 = 15137 +FIX_1_961570560 = 16069 +FIX_2_053119869 = 16819 +FIX_2_562915447 = 20995 +FIX_3_072711026 = 25172 +CONST_BITS = 13 +PASS1_BITS = 1 +F1 = CONST_BITS - PASS1_BITS - 1 +F2 = CONST_BITS - PASS1_BITS +F3 = CONST_BITS + PASS1_BITS + 3 + +# tuning parameter for get_block +TRIES = 16 +MASK = 2**(TRIES*32)-1 +SHIFT = 32*(TRIES-1) + + +def _first_half(data): + """Helper function used to precompute the zero values in a 12 bit datum. + """ + # data has to be 12 bits wide + streamlen = 0 + # count the zeros + zerocount = CLZLUT[data >> 4]; + data = (data << (zerocount + 1)) & 0b111111111111 + streamlen += zerocount + 1 + # get number of remaining bits to read + toread = 0 if zerocount <= 1 else zerocount - 1 + additional = data >> (12 - toread) + data = (data << toread) & 0b111111111111 + streamlen += toread + # add as many zeros to out_list as indicated by additional bits + # if zerocount is 0, tmp = 0 else the 1 merged with additional bits + tmp = 0 if zerocount == 0 else (1 << toread) | additional + return [streamlen, tmp] + + +def _second_half(data): + """Helper function to precompute the nonzeror values in a 15 bit datum. + """ + # data has to be 15 bits wide + streamlen = 0 + zerocount = CLZLUT[data >> 7] + data = (data << (zerocount + 1)) & 0b111111111111111 + streamlen += zerocount + 1 + # 01 == EOB + eob = False + if zerocount == 1: + eob = True + return [streamlen, None, eob] + # get number of remaining bits to read + toread = 0 if zerocount == 0 else zerocount - 1 + additional = data >> (15 - toread) + data = (data << toread) & 0b111111111111111 + streamlen += toread + tmp = (1 << toread) | additional + # get one more bit for the sign + tmp = -tmp if data >> (15 - 1) else tmp + tmp = int(tmp) + streamlen += 1 + return [streamlen, tmp, eob] + + +# Precompute all 12 and 15 bit values for the entropy decoding process +FH = [_first_half(i) for i in range(2**12)] +SH = [_second_half(i) for i in range(2**15)] + + +class BitReader(object): + """Bitreader. Given a stream of data, it allows to read it bitwise.""" + + def __init__(self, packet): + self.packet = packet + self.offset = 0 + self.bits_left = 0 + self.chunk = 0 + self.read_bits = 0 + + def read(self, nbits, consume=True): + """Read nbits and return the integervalue of the read bits. + + If consume is False, it behaves like a 'peek' method (ie it reads the + bits but does not consume them. + """ + # Read enough bits into chunk so we have at least nbits available + while nbits > self.bits_left: + try: + self.chunk = (self.chunk << 32) | struct.unpack_from('<I', self.packet, self.offset)[0] + except struct.error: + self.chunk <<= 32 + self.offset += 4 + self.bits_left += 32 + # Get the first nbits bits from chunk (and remove them from chunk) + shift = self.bits_left - nbits + res = self.chunk >> shift + if consume: + self.chunk -= res << shift + self.bits_left -= nbits + self.read_bits += nbits + return res + + def align(self): + """Byte align the data stream.""" + shift = (8 - self.read_bits) % 8 + self.read(shift) + + +def inverse_dct(block): + """Inverse discrete cosine transform. + """ + workspace = ZEROS[0:64] + data = ZEROS[0:64] + for pointer in range(8): + if (block[pointer + 8] == 0 and block[pointer + 16] == 0 and + block[pointer + 24] == 0 and block[pointer + 32] == 0 and + block[pointer + 40] == 0 and block[pointer + 48] == 0 and + block[pointer + 56] == 0): + dcval = block[pointer] << PASS1_BITS + for i in range(8): + workspace[pointer + i*8] = dcval + continue + + z2 = block[pointer + 16] + z3 = block[pointer + 48] + z1 = (z2 + z3) * FIX_0_541196100 + tmp2 = z1 + z3 * -FIX_1_847759065 + tmp3 = z1 + z2 * FIX_0_765366865 + z2 = block[pointer] + z3 = block[pointer + 32] + tmp0 = (z2 + z3) << CONST_BITS + tmp1 = (z2 - z3) << CONST_BITS + tmp10 = tmp0 + tmp3 + tmp13 = tmp0 - tmp3 + tmp11 = tmp1 + tmp2 + tmp12 = tmp1 - tmp2 + tmp0 = block[pointer + 56] + tmp1 = block[pointer + 40] + tmp2 = block[pointer + 24] + tmp3 = block[pointer + 8] + z1 = tmp0 + tmp3 + z2 = tmp1 + tmp2 + z3 = tmp0 + tmp2 + z4 = tmp1 + tmp3 + z5 = (z3 + z4) * FIX_1_175875602 + tmp0 *= FIX_0_298631336 + tmp1 *= FIX_2_053119869 + tmp2 *= FIX_3_072711026 + tmp3 *= FIX_1_501321110 + z1 *= -FIX_0_899976223 + z2 *= -FIX_2_562915447 + z3 *= -FIX_1_961570560 + z4 *= -FIX_0_390180644 + z3 += z5 + z4 += z5 + tmp0 += z1 + z3 + tmp1 += z2 + z4 + tmp2 += z2 + z3 + tmp3 += z1 + z4 + workspace[pointer + 0] = ((tmp10 + tmp3 + (1 << F1)) >> F2) + workspace[pointer + 56] = ((tmp10 - tmp3 + (1 << F1)) >> F2) + workspace[pointer + 8] = ((tmp11 + tmp2 + (1 << F1)) >> F2) + workspace[pointer + 48] = ((tmp11 - tmp2 + (1 << F1)) >> F2) + workspace[pointer + 16] = ((tmp12 + tmp1 + (1 << F1)) >> F2) + workspace[pointer + 40] = ((tmp12 - tmp1 + (1 << F1)) >> F2) + workspace[pointer + 24] = ((tmp13 + tmp0 + (1 << F1)) >> F2) + workspace[pointer + 32] = ((tmp13 - tmp0 + (1 << F1)) >> F2) + + for pointer in range(0, 64, 8): + z2 = workspace[pointer + 2] + z3 = workspace[pointer + 6] + z1 = (z2 + z3) * FIX_0_541196100 + tmp2 = z1 + z3 * -FIX_1_847759065 + tmp3 = z1 + z2 * FIX_0_765366865 + tmp0 = (workspace[pointer] + workspace[pointer + 4]) << CONST_BITS + tmp1 = (workspace[pointer] - workspace[pointer + 4]) << CONST_BITS + tmp10 = tmp0 + tmp3 + tmp13 = tmp0 - tmp3 + tmp11 = tmp1 + tmp2 + tmp12 = tmp1 - tmp2 + tmp0 = workspace[pointer + 7] + tmp1 = workspace[pointer + 5] + tmp2 = workspace[pointer + 3] + tmp3 = workspace[pointer + 1] + z1 = tmp0 + tmp3 + z2 = tmp1 + tmp2 + z3 = tmp0 + tmp2 + z4 = tmp1 + tmp3 + z5 = (z3 + z4) * FIX_1_175875602 + tmp0 *= FIX_0_298631336 + tmp1 *= FIX_2_053119869 + tmp2 *= FIX_3_072711026 + tmp3 *= FIX_1_501321110 + z1 *= -FIX_0_899976223 + z2 *= -FIX_2_562915447 + z3 *= -FIX_1_961570560 + z4 *= -FIX_0_390180644 + z3 += z5 + z4 += z5 + tmp0 += z1 + z3 + tmp1 += z2 + z4 + tmp2 += z2 + z3 + tmp3 += z1 + z4 + data[pointer + 0] = (tmp10 + tmp3) >> F3 + data[pointer + 7] = (tmp10 - tmp3) >> F3 + data[pointer + 1] = (tmp11 + tmp2) >> F3 + data[pointer + 6] = (tmp11 - tmp2) >> F3 + data[pointer + 2] = (tmp12 + tmp1) >> F3 + data[pointer + 5] = (tmp12 - tmp1) >> F3 + data[pointer + 3] = (tmp13 + tmp0) >> F3 + data[pointer + 4] = (tmp13 - tmp0) >> F3 + + return data + + +def get_pheader(bitreader): + """Read the picture header. + + Returns the width and height of the image. + """ + bitreader.align() + psc = bitreader.read(22) + assert(psc == 0b0000000000000000100000) + pformat = bitreader.read(2) + assert(pformat != 0b00) + if pformat == 1: + # CIF + width, height = 88, 72 + else: + # VGA + width, height = 160, 120 + presolution = bitreader.read(3) + assert(presolution != 0b000) + # double resolution presolution-1 times + width = width << presolution - 1 + height = height << presolution - 1 + #print "width/height:", width, height + ptype = bitreader.read(3) + pquant = bitreader.read(5) + pframe = bitreader.read(32) + return width, height + + +def get_mb(bitreader, picture, width, offset): + """Get macro block. + + This method does not return data but modifies the picture parameter in + place. + """ + mbc = bitreader.read(1) + if mbc == 0: + mbdesc = bitreader.read(8) + assert(mbdesc >> 7 & 1) + if mbdesc >> 6 & 1: + mbdiff = bitreader.read(2) + y = get_block(bitreader, mbdesc & 1) + y.extend(get_block(bitreader, mbdesc >> 1 & 1)) + y.extend(get_block(bitreader, mbdesc >> 2 & 1)) + y.extend(get_block(bitreader, mbdesc >> 3 & 1)) + cb = get_block(bitreader, mbdesc >> 4 & 1) + cr = get_block(bitreader, mbdesc >> 5 & 1) + # ycbcr to rgb + for i in range(256): + j = SCALE_TAB[i] + Y = y[i] - 16 + B = cb[j] - 128 + R = cr[j] - 128 + r = (298 * Y + 409 * R + 128) >> 8 + g = (298 * Y - 100 * B - 208 * R + 128) >> 8 + b = (298 * Y + 516 * B + 128) >> 8 + r = 0 if r < 0 else r + r = 255 if r > 255 else r + g = 0 if g < 0 else g + g = 255 if g > 255 else g + b = 0 if b < 0 else b + b = 255 if b > 255 else b + # re-order the pixels + row = MB_ROW_MAP[i] + col = MB_COL_MAP[i] + picture[offset + row*width + col] = ''.join((chr(r), chr(g), chr(b))) + else: + print "mbc was not zero" + + +def get_block(bitreader, has_coeff): + """Read a 8x8 block from the data stream. + + This method takes care of the huffman-, RLE, zig-zag and idct and returns a + list of 64 ints. + """ + # read the first 10 bits in a 16 bit datum + out_list = ZEROS[0:64] + out_list[0] = int(bitreader.read(10)) * IQUANT_TAB[0] + if not has_coeff: + return inverse_dct(out_list) + i = 1 + while 1: + _ = bitreader.read(32*TRIES, False) + streamlen = 0 + ####################################################################### + for j in range(TRIES): + data = (_ << streamlen) & MASK + data >>= SHIFT + + l, tmp = FH[data >> 20] + streamlen += l + data = (data << l) & 0xffffffff + i += tmp + + l, tmp, eob = SH[data >> 17] + streamlen += l + if eob: + bitreader.read(streamlen) + return inverse_dct(out_list) + j = ZIG_ZAG_POSITIONS[i] + out_list[j] = tmp*IQUANT_TAB[j] + i += 1 + ####################################################################### + bitreader.read(streamlen) + return inverse_dct(out_list) + + +def get_gob(bitreader, picture, slicenr, width): + """Read a group of blocks. + + The method does not return data, the picture parameter is modified in place + instead. + """ + # the first gob has a special header + if slicenr > 0: + bitreader.align() + gobsc = bitreader.read(22) + if gobsc == 0b0000000000000000111111: + print "weeeee" + return False + elif (not (gobsc & 0b0000000000000000100000) or + (gobsc & 0b1111111111111111000000)): + print "Got wrong GOBSC, aborting.", bin(gobsc) + return False + _ = bitreader.read(5) + offset = slicenr*16*width + for i in range(width / 16): + get_mb(bitreader, picture, width, offset+16*i) + + +def read_picture(data): + """Convert an AR.Drone image packet to rgb-string. + + Returns: width, height, image and time to decode the image + """ + bitreader = BitReader(data) + t = datetime.datetime.now() + width, height = get_pheader(bitreader) + slices = height / 16 + blocks = width / 16 + image = [0 for i in range(width*height)] + for i in range(0, slices): + get_gob(bitreader, image, i, width) + bitreader.align() + eos = bitreader.read(22) + assert(eos == 0b0000000000000000111111) + t2 = datetime.datetime.now() + return width, height, ''.join(image), (t2 - t).microseconds / 1000000. + + +try: + psyco.bind(BitReader) + psyco.bind(get_block) + psyco.bind(get_gob) + psyco.bind(get_mb) + psyco.bind(inverse_dct) + psyco.bind(read_picture) +except NameError: + print "Unable to bind video decoding methods with psyco. Proceeding anyways, but video decoding will be slow!" + + +def main(): + fh = open('framewireshark.raw', 'r') + #fh = open('videoframe.raw', 'r') + data = fh.read() + fh.close() + runs = 20 + t = 0 + for i in range(runs): + print '.', + width, height, image, ti = read_picture(data) + #show_image(image, width, height) + t += ti + print + print 'avg time:\t', t / runs, 'sec' + print 'avg fps:\t', 1 / (t / runs), 'fps' + if 'image' in sys.argv: + import pygame + pygame.init() + W, H = 320, 240 + screen = pygame.display.set_mode((W, H)) + surface = pygame.image.fromstring(image, (width, height), 'RGB') + screen.blit(surface, (0, 0)) + pygame.display.flip() + raw_input() + + +if __name__ == '__main__': + if 'profile' in sys.argv: + cProfile.run('main()') + else: + main() + |