OpenCVとPythonと私 vol.6

#!/usr/bin/python

import cv, numpy
from opencv import highgui

import sys
import getopt

filename = None
debug    = False
rows     = 0
columns  = 0
span     = 0
sourceFile = None
destinationFile = 'result.png'

def main () :

        global filename
        global debug
        global rows
        global columns
        global span
        global sourceFile
        global destinationFile

        try:
                opts, args = getopt.getopt( sys.argv[1:], "hdr:c:s:u:o:",["help","debug","rows","columns","span","undistort","output"])
        except getopt.GetoptError, err :
                print str( err )
                usage()
                sys.exit(2)

        if len( args ) != 0 :
                filename = args[:]
        else :
                usage()
                sys.exit(0)

        for o, a in opts :
                if o in ( "-h", "--help" ) :
                        usage()
                        sys.exit(0)
                elif o in ( "-d", "--debug" ) :
                        debug = True
                elif o in ( "-r", "--rows" ) :
                        rows = int( a )
                elif o in ( "-c", "--columns" ) :
                        columns = int( a )
                elif o in ( "-s", "--span" ) :
                        span = int( a )
                elif o in ( "-u", "--undistort" ) :
                        sourceFile = a
                elif o in ( "-o", "--output" ) :
                        destinationFile = a
                else :
                        assert False, "unhandled option"

def debugout ( str ) :
        if debug == True :
                print str
   
def usage () :
        print "Usage: undistortTest.py [OPTION] FILE"
        print " "
        print " -h, --help      Display this message."
        print " -d, --debug     Output debug log."
        print " -r, --rows      Row counts of chessboard's corners"
        print " -c, --columns   Column counts of chessboard's corners"
        print " -s, --span      Length of between 2 corners"
        print " -u, --undistort Specify distorted image."
        print " -o, --output    Set undistorted output file name."
       
if __name__ == "__main__" :

        main()

        objectPointsList =
        imagePointsList =

        for element in filename :
                '''
                debugout( "Input file name is " + element )
                '''
                img = cv.LoadImage( element , cv.CV_LOAD_IMAGE_GRAYSCALE )
                corners = cv.FindChessboardCorners( img, ( columns, rows ) )

                for corner in corners[1] :
                        imagePointsList.append( list( corner ) )

                ytmp = range( columns )
                ytmp.reverse()
                for x in range( rows ) :
                        for y in ytmp :
                                objectPoint = []
                                objectPoint.append( x * span )
                                objectPoint.append( y * span )
                                objectPoint.append( 0 )
                                objectPointsList.append( objectPoint )
                                del objectPoint

        imagePointsArray = numpy.array( imagePointsList, 'float' )
        imagePoints = cv.fromarray( imagePointsArray )

        objectPointsArray = numpy.array( objectPointsList, 'float' )
        objectPoints = cv.fromarray( objectPointsArray )

        cameraMatrix = cv.CreateMat( 3, 3, cv.CV_32FC1 )
        distortionCoeffs = cv.CreateMat( 1, 4, cv.CV_32FC1 )
        rotationVector = cv.CreateMat( len( filename ), 3, cv.CV_32FC1 )
        translationVector = cv.CreateMat( len( filename ), 3, cv.CV_32FC1 )

        pointCountsArray = rows * columns * numpy.ones( ( len( filename ), 1 ) ,'int32')
        pointCounts = cv.fromarray( pointCountsArray )

        try:
                cv.CalibrateCamera2( objectPoints, imagePoints, \
                                             pointCounts ,\
                                             ( img.height, img.width ) ,\
                                             cameraMatrix, distortionCoeffs,\
                                             rotationVector, translationVector, 0 )
        except cv.error, err :
                print str( err )
                sys.exit(2)

        distortedImage = cv.LoadImageM( sourceFile )
        undistortedImage = cv.LoadImageM( sourceFile )

        cv.Undistort2( distortedImage, undistortedImage, cameraMatrix, distortionCoeffs )

        cv.NamedWindow( "Before", cv.CV_WINDOW_AUTOSIZE )
        cv.NamedWindow( "After", cv.CV_WINDOW_AUTOSIZE )
        cv.MoveWindow( "Before", 100, 100)
        cv.MoveWindow( "After", 100, 200 )

        cv.ShowImage( "Before", distortedImage )
        cv.ShowImage( "After", undistortedImage )

        cv.SaveImage( destinationFile, undistortedImage )

        print "Hit any key."
        cv.WaitKey( 0 )

        cv.DestroyWindow( "Before" )
        cv.DestroyWindow( "After" )