view python/cvutils.py @ 100:2a3cafcf5faf

added function to compute the translation between two images
author Nicolas Saunier <nicolas.saunier@polymtl.ca>
date Tue, 12 Jul 2011 19:59:25 -0400
parents e7dc5a780f09
children abfc54c097d4
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

def computeTranslation(img1, img2, img1Points, maxTranslation, minNMatches, windowSize = (5,5), level = 5, criteria = (cv.CV_TERMCRIT_EPS, 0, 0.01)):
    '''Computes the translation between of img2 with respect to img1
    (loaded using OpenCV)
    img1Points are used to compute the translation

    TODO add diagnostic if data is all over the place, and it most likely is not a translation (eg zoom)'''
    from numpy.core.multiarray import zeros
    from numpy.lib.function_base import median

    (img2Points, status, track_error) = cv.CalcOpticalFlowPyrLK(img1, img2, zeros((img1.rows,img1.cols+8)), zeros((img1.rows,img1.cols+8)), img1Points, windowSize, level, criteria, 0)

    deltaX = []
    deltaY = []
    for (k, (p1,p2)) in enumerate(zip(img1Points, img2Points)):
        if status[k] == 1:
            dx = p2[0]-p1[0]
            dy = p2[1]-p1[1]
            d = dx**2 + dy**2
            if d < maxTranslation:
                deltaX.append(dx)
                deltaY.append(dy)
    if len(deltaX) >= 10:
        return [median(deltaX), median(deltaY)]
    else:
        return None

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')