view python/cvutils.py @ 99:e7dc5a780f09

added conversion functions and homography computation
author Nicolas Saunier <nicolas.saunier@polymtl.ca>
date Tue, 12 Jul 2011 17:01:59 -0400
parents b85912ab4064
children 2a3cafcf5faf
line wrap: on
line source

#! /usr/bin/env python
'''Image/Video utilities'''

import Image, ImageDraw # PIL
import cv
from sys import stdout

#import aggdraw # agg on top of PIL (antialiased drawing)
from moving import Point
#import utils

__metaclass__ = type

def drawLines(filename, origins, destinations, w = 1, resultFilename='image.png'):
    '''Draws lines over the image '''
    
    img = Image.open(filename)

    draw = ImageDraw.Draw(img)
    #draw = aggdraw.Draw(img)
    #pen = aggdraw.Pen("red", width)
    for p1, p2 in zip(origins, destinations):
        draw.line([p1.x, p1.y, p2.x, p2.y], width = w, fill = (256,0,0))
        #draw.line([p1.x, p1.y, p2.x, p2.y], pen)
    del draw 

    #out = utils.openCheck(resultFilename)
    img.save(resultFilename)

def computeHomography(srcPoints, dstPoints, method=0, ransacReprojThreshold=0.0):
    '''Returns the homography matrix mapping from srcPoints to dstPoints (dimension Nx2)'''
    cvSrcPoints = arrayToCvMat(srcPoints);
    cvDstPoints = arrayToCvMat(dstPoints);
    H = cv.CreateMat(3, 3, cv.CV_64FC1)
    cv.FindHomography(cvSrcPoints, cvDstPoints, H, method, ransacReprojThreshold)
    return H

def cvMatToArray(cvmat):
    '''Converts an OpenCV CvMat to numpy array.'''
    from numpy.core.multiarray import zeros
    a = zeros((cvmat.rows, cvmat.cols))#array([[0.0]*cvmat.width]*cvmat.height)
    for i in xrange(cvmat.rows):
        for j in xrange(cvmat.cols):
            a[i,j] = cvmat[i,j]
    return a

def arrayToCvMat(a, t = cv.CV_64FC1):
    '''Converts a numpy array to an OpenCV CvMat, with default type CV_64FC1.'''
    cvmat = cv.CreateMat(a.shape[0], a.shape[1], t)
    for i in range(cvmat.rows):
        for j in range(cvmat.cols):
            cvmat[i,j] = a[i,j]
    return cvmat

def printCvMat(cvmat, out = stdout):
    '''Prints the cvmat to out'''
    for i in xrange(cvmat.rows):
        for j in xrange(cvmat.cols):
            out.write('{0} '.format(cvmat[i,j]))
        out.write('\n')

def projectArray(homography, points):
    '''Returns the coordinates of the projected points (format 2xN points)
    through homography'''
    from numpy.core._dotblas import dot
    from numpy.core.multiarray import array
    from numpy.lib.function_base import append

    if points.shape[0] != 2:
        raise Exception('points of dimension {0} {1}'.format(points.shape[0], points.shape[1]))

    if (homography!=None) and homography.size>0:
        augmentedPoints = append(points,[[1]*points.shape[1]], 0)
        prod = dot(homography, augmentedPoints)
        return prod[0:2]/prod[2]
    else:
        return p

def project(homography, p):
    '''Returns the coordinates of the projection of the point p
    through homography'''
    from numpy.core.multiarray import array
    return projectArray(homography, array([[p[0]],p[1]]))

def projectTrajectory(homography, trajectory):
    '''Projects a series of points in the format
    [[x1, x2, ...],
    [y1, y2, ...]]'''
    from numpy.core.multiarray import array
    return projectArray(homography, array(trajectory))

def invertHomography(homography):
    'Returns an inverted homography'
    from numpy.linalg.linalg import inv
    invH = inv(homography)
    invH /= invH[2,2]
    return invH

class FourWayIntersection:
    '''Simple class for simple intersection outline'''
    def __init__(self, dimension, coordX, coordY):
        self.dimension = dimension
        self.coordX = coordX
        self.coordY = coordY

    def plot(self, options = 'k'):
        from matplotlib.pyplot import plot, axis
    
        minX = min(self.dimension[0])
        maxX = max(self.dimension[0])
        minY = min(self.dimension[1])
        maxY = max(self.dimension[1])
        
        plot([minX, self.coordX[0], self.coordX[0]], [self.coordY[0], self.coordY[0], minY],options)
        plot([self.coordX[1], self.coordX[1], maxX], [minY, self.coordY[0], self.coordY[0]],options)
        plot([minX, self.coordX[0], self.coordX[0]], [self.coordY[1], self.coordY[1], maxY],options)
        plot([self.coordX[1], self.coordX[1], maxX], [maxY, self.coordY[1], self.coordY[1]],options)
        axis('equal')