Source

hakanardo / video / video.py

Full commit
from array import array
from cfile import cfile
import os, re, sys
import itertools
from math import sqrt

class Image(object):
    typecode = 'B'
    def __init__(self, w, h, data=None):
        self.width = w
        self.height = h
        if data is None:
            self.data = array(self.typecode, (0,)) * (w * h)
        else:
            self.data = array(self.typecode, data)

    def new(self):
        return Image(self.width, self.height)

    def _idx(self, (x, y)):
        if 0 <= x < self.width and 0 <= y < self.height:
            return y * self.width + x
        raise IndexError

    def __setitem__(self, idx, val):
        self.data[self._idx(idx)] = val

    def __getitem__(self, idx):
        return self.data[self._idx(idx)]

    def pixels(self, border=0, xborder=None, yborder=None, dx=1, dy=1):
        if xborder is None:
            xborder = border
        if yborder is None:
            yborder = border
        for y in range(yborder, self.height - yborder, dy):
            for x in range(xborder, self.width - xborder, dx):
                yield x, y

    def blocks(self, w, h, dx=1, dy=1, pad=None):
        if pad is None:
            for y in range(0, self.height - h + 1, dy):
                for x in range(0, self.width - w + 1, dy):
                    yield SubImage(self, w, h, x, y, pad)
        else:
            xoff = w / 2
            yoff = h / 2
            for y in range(0, self.height, dy):
                for x in range(0, self.width, dy):
                    yield SubImage(self, w, h, x-xoff, y-yoff, pad)

    def __iter__(self):
        for p in self.data:
            yield p

    def force(self):
        pass

    def __add__(self, other):
        return Add(self, other)
    __radd__=__add__

    def __mul__(self, other):
        return Mul(self, other)
    __rmul__ = __mul__

    def __div__(self, other):
        return Div(self, other)
    def __rdiv__(self, other):
        return Div(other, self)
    
    def __sub__(self, other):
        return Sub(self, other)
    def __rsub__(self, other):
        return Sub(other, self)
    
    def __pow__(self, other):
        return Pow(self, other)

class RGBImage(Image):
    pass

class VirtualImage(Image):
    def __init__(self, w, h):
        self.width = w
        self.height = h
        self.data = None

    def force(self):
        raise NotImplementedError

    def force_if_needed(self):
        if self.data is None:
            self.force()

    def __setitem__(self, idx, val):
        self.force_if_needed()
        return Image.__setitem__(self, idx, val)

    def __getitem__(self, idx):
        self.force_if_needed()
        return Image.__getitem__(self, idx)

    def __iter__(self):
        raise NotImplementedError

        
class Pixelwise(VirtualImage):
    typecode = 'd'
    def __init__(self, *args):
        w = set([])
        h = set([])
        self.args = []
        for a in args:
            if isinstance(a, Image):
                w.add(a.width)
                h.add(a.height)
            else:
                a = itertools.cycle((a,))
            self.args.append(a)
        assert len(w) == len(h) == 1
        self.width = w.pop()
        self.height = h.pop()
        self.data = None
        
    def __iter__(self):
        if self.data is None:
            for p in zip(*self.args):
                yield self.op(*p)
        else:
            for p in self.data:
                yield p

    def force(self):
        data = array(self.typecode, (0,)) * (self.width * self.height)
        for i, p in enumerate(self):
            data[i] = p
        self.data = data

    def op(self, p):
        raise NotImplementedError

class Blockwise(Pixelwise):
    size = (3, 3)
    pad = 0
    dx = dy = 1
    def __iter__(self):
        if self.data is None:
            inputs = [a.blocks(self.size[0], self.size[1], dx=self.dx, dy=self.dy, pad=self.pad) for a in self.args]
            for blocks in zip(*inputs):
                yield self.op(*blocks)
        else:
            for p in self.data:
                yield p

class Sqrt(Pixelwise):
    def op(self, p):
        return sqrt(p)

class Add(Pixelwise):
    def op(self, a, b):
        return a + b

class Sub(Pixelwise):
    def op(self, a, b):
        return a - b

class Mul(Pixelwise):
    def op(self, a, b):
        return a * b

class Div(Pixelwise):
    def op(self, a, b):
        return a / b

class Pow(Pixelwise):
    def op(self, a, b):
        return a ** b

class UInt8(Pixelwise):
    typecode = 'B'
    def op(self, p):
        return int(round(p))

class SobelDx(Blockwise):
    def op(self, blk):
        return (1*blk[0, 0] - 1*blk[2, 0] +
                2*blk[0, 1] - 2*blk[2, 1] +
                1*blk[0, 2] - 1*blk[2, 2])

class SobelDy(Blockwise):
    def op(self, blk):
        return ( 1*blk[0, 0] + 2*blk[1, 0] + 1*blk[2, 0] +
                -1*blk[0, 2] - 2*blk[1, 2] - 1*blk[2, 2])

def SobelMagnitude(a):
    return SobelDx(a)**2 + SobelDy(a)**2

    
class SubImage(Image):
    def __init__(self, image, width, height, dx, dy, pad=None):
        self.image = image
        self.dx = dx
        self.dy = dy
        self.width = width
        self.height = height
        self.pad = pad

    def __getitem__(self, (x, y)):
        if 0 <= x < self.width and 0 <= y < self.height:
            try:
                return self.image[x + self.dx, y + self.dy]
            except IndexError:
                if self.pad is None:
                    raise
                else:
                    return self.pad
        raise IndexError



#so = cfile('glview.c', ldflags='-lglut -lGLU -lGL', cflags='-g')
so = cfile('xvview.c', ldflags='-lX11 -lXv', cflags='-g')
viewer = so.xvview
from ctypes import c_int, c_void_p, c_char
viewer.argtypes = [c_int, c_int, c_void_p, c_int, c_char]

def view(img, scale=False):
    if scale:
        img.force()
        ma, mi = max(img.data), min(img.data)
        if ma != mi:
            img = (img - mi) / (ma - mi) * 255
    if img.typecode != 'B':
        img = UInt8(img)
    img.force()
    viewer(img.width, img.height, img.data.buffer_info()[0],
           isinstance(img, RGBImage), img.data.typecode)
    
def mplayer(fn='tv://'):
    f = os.popen('mplayer -really-quiet -noframedrop '+
                 '-vo yuv4mpeg:file=/dev/stdout 2>/dev/null </dev/null ' + fn)
    hdr = f.readline()
    m = re.search('W(\d+) H(\d+)', hdr)
    w, h = int(m.group(1)), int(m.group(2))
    while True:
        hdr = f.readline()
        if hdr != 'FRAME\n':
            break
        img = Image(w, h)
        img.data = array('B')
        img.data.fromfile(f, w*h)
        f.read(w*h/2) # Color data
        yield img