data/frames/gen-frames.py

   1 #!/usr/bin/python
   2
   3 import copy
   4 import glob
   5 import gzip
   6 import itertools
   7 import operator
   8 import re
   9 import struct
  10 import sys
  11
  12 class FormatRGB:
  13         def __init__(self, name, mapping):
  14                 self.name = name
  15                 self.mapping = copy.copy(mapping)
  16                 self.depth = sum(v[1] for v in mapping.values())
  17
  18                 if self.mapping.has_key('a'):
  19                         self.alpha_bits = self.mapping['a'][1]
  20                 elif self.mapping.has_key('x'):
  21                         self.alpha_bits = self.mapping['x'][1]
  22                 else:
  23                         self.alpha_bits = 0
  24
  25         def alpha(self, alpha):
  26                 if self.mapping.has_key('a'):
  27                         return (alpha >> (8 - self.mapping['a'][1])) << self.mapping['a'][0]
  28                 elif self.mapping.has_key('x'):
  29                         return (alpha >> (8 - self.mapping['x'][1])) << self.mapping['x'][0]
  30                 else:
  31                         return 0
  32
  33         def generate_8(self, width, height, rgb):
  34                 output = []
  35                 for i in xrange(width * height):
  36                         r, g, b = rgb[i]
  37                         pixel = ((r >> (8 - self.mapping['r'][1])) << self.mapping['r'][0]) \
  38                               | ((g >> (8 - self.mapping['g'][1])) << self.mapping['g'][0]) \
  39                               | ((b >> (8 - self.mapping['b'][1])) << self.mapping['b'][0])
  40                         output.append(struct.pack('B', pixel))
  41                 return ''.join(output)
  42
  43         def generate_16(self, width, height, rgb, alpha):
  44                 output = []
  45                 for i in xrange(width * height):
  46                         r, g, b = rgb[i]
  47                         pixel = ((r >> (8 - self.mapping['r'][1])) << self.mapping['r'][0]) \
  48                               | ((g >> (8 - self.mapping['g'][1])) << self.mapping['g'][0]) \
  49                               | ((b >> (8 - self.mapping['b'][1])) << self.mapping['b'][0]) \
  50                               | self.alpha(alpha)
  51                         output.append(struct.pack('<H', pixel))
  52                 return ''.join(output)
  53
  54         def generate_24(self, width, height, rgb):
  55                 output = []
  56                 for i in xrange(width * height):
  57                         r, g, b = rgb[i]
  58                         pixel = ((r >> (8 - self.mapping['r'][1])) << self.mapping['r'][0]) \
  59                               | ((g >> (8 - self.mapping['g'][1])) << self.mapping['g'][0]) \
  60                               | ((b >> (8 - self.mapping['b'][1])) << self.mapping['b'][0])
  61                         output.append(struct.pack('<I', pixel)[0:3])
  62                 return ''.join(output)
  63
  64         def generate_32(self, width, height, rgb, alpha):
  65                 output = []
  66                 for i in xrange(width * height):
  67                         r, g, b = rgb[i]
  68                         pixel = ((r >> (8 - self.mapping['r'][1])) << self.mapping['r'][0]) \
  69                               | ((g >> (8 - self.mapping['g'][1])) << self.mapping['g'][0]) \
  70                               | ((b >> (8 - self.mapping['b'][1])) << self.mapping['b'][0]) \
  71                               | self.alpha(alpha)
  72                         output.append(struct.pack('<I', pixel))
  73                 return ''.join(output)
  74
  75         def convert(self, width, height, rgb, alpha):
  76                 if self.depth == 8:
  77                         return self.generate_8(width, height, rgb)
  78                 elif self.depth == 16:
  79                         return self.generate_16(width, height, rgb, alpha)
  80                 elif self.depth == 24:
  81                         return self.generate_24(width, height, rgb)
  82                 elif self.depth == 32:
  83                         return self.generate_32(width, height, rgb, alpha)
  84                 else:
  85                         raise RuntimeError, 'Invalid depth %s' % self.depth
  86
  87         def bin(self, bins, val):
  88                 bins[val >> 2] += 1
  89
  90         def histogram(self, width, height, rgb):
  91                 rgb_min = [255, 255, 255]
  92                 rgb_max = [0, 0, 0]
  93                 rgb_sum = [0, 0, 0]
  94                 rgb_bins = [[0] * 64, [0] * 64, [0] * 64]
  95
  96                 for i in xrange(width * height):
  97                         pixel = rgb[i]
  98
  99                         rgb_min = map(min, pixel, rgb_min)
 100                         rgb_max = map(max, pixel, rgb_max)
 101                         rgb_sum = map(operator.add, pixel, rgb_sum)
 102
 103                         map(self.bin, rgb_bins, pixel)
 104
 105                 output = []
 106                 for i in xrange(len(rgb_min)):
 107                         output.append(struct.pack('BBBB', rgb_min[i], 0, rgb_max[i], 0))
 108                 output.append(struct.pack('<3I', *rgb_sum))
 109                 for i in xrange(len(rgb_bins)):
 110                         output.append(struct.pack('<64I', *rgb_bins[i]))
 111
 112                 return ''.join(output)
 113
 114         def compose(self, ninputs, width, height, rgb):
 115                 output = [(0, 0, 0)] * (width * height)
 116                 offset = 50
 117
 118                 for input in xrange(ninputs):
 119                         length = width - offset
 120                         for y in xrange(height - offset):
 121                                 dst_offset = (y + offset) * width + offset
 122                                 src_offset = y * width
 123                                 output[dst_offset:dst_offset+length] = rgb[src_offset:src_offset+length]
 124                         offset += 50
 125
 126                 return ''.join(chr(d) for d in list(itertools.chain.from_iterable(output)))
 127
 128
 129 class FormatYUVPacked:
 130         def __init__(self, name, mapping):
 131                 self.name = name
 132                 self.mapping = copy.copy(mapping)
 133
 134         def convert(self, width, height, yuv):
 135                 output = []
 136                 for i in xrange(width * height / 2):
 137                         pixel = yuv[i*4:(i+1)*4]
 138                         pixel = (pixel[self.mapping[0]], pixel[self.mapping[1]],
 139                                  pixel[self.mapping[2]], pixel[self.mapping[3]])
 140                         output.extend(pixel)
 141                 return ''.join(output)
 142
 143         def bin(self, bins, val):
 144                 bins[val >> 2] += 1
 145
 146         def histogram(self, width, height, yuv):
 147                 vyu_min = [255, 255, 255]
 148                 vyu_max = [0, 0, 0]
 149                 vyu_sum = [0, 0, 0]
 150                 vyu_bins = [[0] * 64, [0] * 64, [0] * 64]
 151
 152                 for y in xrange(height):
 153                         for x in xrange(width / 2):
 154                                 offset = y * width * 2 + x * 4
 155                                 u0 = ord(yuv[offset])
 156                                 y0 = ord(yuv[offset+1])
 157                                 v0 = ord(yuv[offset+2])
 158                                 y1 = ord(yuv[offset+3])
 159
 160                                 if x != width / 2 - 1:
 161                                         u2 = ord(yuv[offset+4])
 162                                         v2 = ord(yuv[offset+6])
 163                                         u1 = (u0 + u2) / 2
 164                                         v1 = (v0 + v2) / 2
 165                                 else:
 166                                         u1 = u0
 167                                         v1 = v1
 168
 169                                 for vyu in ((v0, y0, u0), (v1, y1, u1)):
 170                                         vyu_min = map(min, vyu, vyu_min)
 171                                         vyu_max = map(max, vyu, vyu_max)
 172                                         vyu_sum = map(operator.add, vyu, vyu_sum)
 173
 174                                         map(self.bin, vyu_bins, vyu)
 175
 176                 output = []
 177                 for i in xrange(len(vyu_min)):
 178                         output.append(struct.pack('BBBB', vyu_min[i], 0, vyu_max[i], 0))
 179                 output.append(struct.pack('<3I', *vyu_sum))
 180                 for i in xrange(len(vyu_bins)):
 181                         output.append(struct.pack('<64I', *vyu_bins[i]))
 182
 183                 return ''.join(output)
 184
 185         def compose(self, ninputs, width, height, yuv):
 186                 output = ['\0'] * (width * height * 2)
 187                 offset = 50
 188
 189                 for input in xrange(ninputs):
 190                         length = (width - offset) * 2
 191                         for y in xrange(height - offset):
 192                                 dst_offset = ((y + offset) * width + offset) * 2
 193                                 src_offset = y * width * 2
 194                                 output[dst_offset:dst_offset+length] = yuv[src_offset:src_offset+length]
 195                         offset += 50
 196
 197                 return ''.join(output)
 198
 199
 200 class FormatNV:
 201         def __init__(self, name, hsub, vsub, mapping):
 202                 self.name = name
 203                 self.hsub = hsub
 204                 self.vsub = vsub
 205                 self.mapping = copy.copy(mapping)
 206
 207         def convert(self, width, height, yuv):
 208                 output = []
 209
 210                 for i in xrange(width * height):
 211                         output.append(yuv[2*i+1])
 212
 213                 for y in xrange(height / self.vsub):
 214                         for x in xrange(width / 2):
 215                                 offset = (y * self.vsub * width * 2) + x * 4
 216                                 uv = (yuv[offset], yuv[offset+2])
 217                                 uv = (uv[self.mapping[0]], uv[self.mapping[1]])
 218                                 output.extend(uv)
 219
 220                 return ''.join(output)
 221
 222
 223 class FormatYUVPlanar:
 224         def __init__(self, name, hsub, vsub, mapping):
 225                 self.name = name
 226                 self.hsub = hsub
 227                 self.vsub = vsub
 228                 self.mapping = copy.copy(mapping)
 229
 230         def convert(self, width, height, yuv):
 231                 output = []
 232
 233                 for i in xrange(width * height):
 234                         output.append(yuv[2*i+1])
 235
 236                 for y in xrange(height / self.vsub):
 237                         for x in xrange(width / 2):
 238                                 offset = (y * self.vsub * width * 2) + x * 4
 239                                 u = yuv[offset + self.mapping[0] * 2]
 240                                 output.append(u)
 241
 242                 for y in xrange(height / self.vsub):
 243                         for x in xrange(width / 2):
 244                                 offset = (y * self.vsub * width * 2) + x * 4
 245                                 v = yuv[offset + self.mapping[1] * 2]
 246                                 output.append(v)
 247
 248                 return ''.join(output)
 249
 250
 251 formats_rgb = {
 252         'rgb332': FormatRGB('rgb332',  {'r': (5, 3), 'g': (2, 3), 'b': (0, 2)}),
 253         'rgb565': FormatRGB('rgb565',  {'r': (11, 5), 'g': (5, 6), 'b': (0, 5)}),
 254         'bgr24':  FormatRGB('bgr24',   {'r': (16, 8), 'g': (8, 8),  'b': (0, 8)}),
 255         'rgb24':  FormatRGB('rgb24',   {'r': (0, 8),  'g': (8, 8),  'b': (16, 8)}),
 256 }
 257
 258 formats_argb = {
 259         'argb555': FormatRGB('argb555', {'a': (15, 1), 'r': (10, 5), 'g': (5, 5), 'b': (0, 5)}),
 260         'xrgb555': FormatRGB('xrgb555', {'x': (15, 1), 'r': (10, 5), 'g': (5, 5), 'b': (0, 5)}),
 261         'abgr32':  FormatRGB('abgr32',  {'a': (24, 8), 'r': (16, 8), 'g': (8, 8),  'b': (0, 8)}),
 262         'argb32':  FormatRGB('argb32',  {'a': (0, 8),  'r': (8, 8),  'g': (16, 8), 'b': (24, 8)}),
 263         'xbgr32':  FormatRGB('xbgr32',  {'x': (24, 8), 'r': (16, 8), 'g': (8, 8),  'b': (0, 8)}),
 264         'xrgb32':  FormatRGB('xrgb32',  {'x': (0, 8),  'r': (8, 8),  'g': (16, 8), 'b': (24, 8)}),
 265 }
 266
 267 formats_yuv = {
 268         'uyvy':    FormatYUVPacked('uyvy', (0, 1, 2, 3)),
 269         'vyuy':    FormatYUVPacked('vyuy', (2, 1, 0, 3)),
 270         'yuyv':    FormatYUVPacked('yuyv', (1, 0, 3, 2)),
 271         'yvyu':    FormatYUVPacked('yvyu', (1, 2, 3, 0)),
 272         'nv12m':   FormatNV('nv12m', 2, 2, (0, 1)),
 273         'nv21m':   FormatNV('nv21m', 2, 2, (1, 0)),
 274         'nv16m':   FormatNV('nv16m', 2, 1, (0, 1)),
 275         'nv61m':   FormatNV('nv61m', 2, 1, (1, 0)),
 276         'yuv420m': FormatYUVPlanar('yuv420m', 2, 2, (0, 1)),
 277 }
 278
 279 resolutions = ((640, 480), (1024, 768))
 280
 281 def main(argv):
 282         re_fname = re.compile('frame-([a-z]*)-([0-9]*)x([0-9]*).([a-z]*).gz')
 283
 284         for fname in glob.glob('*.gz'):
 285                 match = re_fname.match(fname)
 286                 if not match:
 287                         continue
 288
 289                 typ = match.group(1)
 290                 res = (int(match.group(2)), int(match.group(3)))
 291                 fmt = match.group(4)
 292
 293                 if fmt == 'rgb':
 294                         rgb = gzip.open(fname, 'rb').read()
 295                         rgb = [struct.unpack('BBB', rgb[i*3:(i+1)*3]) for i in xrange(len(rgb) / 3)]
 296
 297                         for format in formats_rgb.values():
 298                                 bin_fname = 'frame-%s-%s-%ux%u.bin' % (typ, format.name, res[0], res[1])
 299                                 print 'Generating %s' % bin_fname
 300                                 bin = format.convert(res[0], res[1], rgb, 0)
 301                                 file(bin_fname, 'wb').write(bin)
 302
 303                         for format in formats_argb.values():
 304                                 for alpha in (0, 100, 200, 255):
 305                                         if format.alpha_bits == 1 and alpha not in (0, 255):
 306                                                 continue
 307                                         bin_fname = 'frame-%s-%s-%ux%u-alpha%u.bin' % (typ, format.name, res[0], res[1], alpha)
 308                                         print 'Generating %s' % bin_fname
 309                                         bin = format.convert(res[0], res[1], rgb, alpha)
 310                                         file(bin_fname, 'wb').write(bin)
 311
 312                         format = formats_rgb['rgb24']
 313
 314                         bin_fname = 'histo-%s-%s-%ux%u.bin' % (typ, format.name, res[0], res[1])
 315                         print 'Generating %s' % bin_fname
 316                         bin = format.histogram(res[0], res[1], rgb)
 317                         file(bin_fname, 'wb').write(bin)
 318
 319                         if typ == 'reference' and res[0] == 1024 and res[1] == 768:
 320                                 for ninputs in xrange(1, 6):
 321                                         bin_fname = 'frame-composed-%u-%s-%ux%u.bin' % (ninputs, format.name, res[0], res[1])
 322                                         print 'Generating %s' % bin_fname
 323                                         bin = format.compose(ninputs, res[0], res[1], rgb)
 324                                         file(bin_fname, 'wb').write(bin)
 325
 326                 elif fmt == 'yuv':
 327                         yuv = gzip.open(fname, 'rb').read()
 328
 329                         for format in formats_yuv.values():
 330                                 bin_fname = 'frame-%s-%s-%ux%u.bin' % (typ, format.name, res[0], res[1])
 331                                 print 'Generating %s' % bin_fname
 332                                 bin = format.convert(res[0], res[1], yuv)
 333                                 file(bin_fname, 'wb').write(bin)
 334
 335                         format = formats_yuv['uyvy']
 336                         bin_fname = 'histo-%s-%s-%ux%u.bin' % (typ, format.name, res[0], res[1])
 337                         print 'Generating %s' % bin_fname
 338                         bin = format.histogram(res[0], res[1], yuv)
 339                         file(bin_fname, 'wb').write(bin)
 340
 341                         if typ == 'reference' and res[0] == 1024 and res[1] == 768:
 342                                 for ninputs in xrange(1, 6):
 343                                         bin_fname = 'frame-composed-%u-%s-%ux%u.bin' % (ninputs, format.name, res[0], res[1])
 344                                         print 'Generating %s' % bin_fname
 345                                         bin = format.compose(ninputs, res[0], res[1], yuv)
 346                                         file(bin_fname, 'wb').write(bin)
 347
 348
 349 if __name__ == '__main__':
 350         sys.exit(main(sys.argv))