robocup ball detector

simple_robocup_cv.py

import cv2.cv as cv

class SimpleBallRecognition:
    
    def __init__(self, image):
        self.bgrImage = image
        self.ycrcbImage = cv.CreateImage(cv.GetSize(self.bgrImage), cv.IPL_DEPTH_8U, 3)
        self.resultImage = cv.CreateImage(cv.GetSize(self.bgrImage), cv.IPL_DEPTH_8U, 3)
        self.tmpGrayImage = cv.CreateImage(cv.GetSize(self.bgrImage), cv.IPL_DEPTH_8U, 1)
        cv.CvtColor(self.bgrImage, self.ycrcbImage, cv.CV_BGR2YCrCb)
        
        self.selectedBGRColor = None
        self.selectedYCrCbColor = None
        
        cv.SetMouseCallback('Source', self.onMouse)
        
        self.yThreshold = 58;
        self.crThreshold = 20;
        self.cbThreshold = 28;
        
        cv.CreateTrackbar('Y-Threshold', 'Threshold', self.yThreshold, 255, self.yThresholdChanged);
        cv.CreateTrackbar('Cr-Threshold', 'Threshold', self.crThreshold, 255, self.crThresholdChanged);
        cv.CreateTrackbar('Cb-Threshold', 'Threshold', self.cbThreshold, 255, self.cbThresholdChanged);
        
        self.findBall()
      
    def onMouse(self, event, x, y, flag, param):
        if event == cv.CV_EVENT_LBUTTONDOWN:
            print "Mouse Click on (" + str(x) + "," + str(y) + ")"
            self.selectedBGRColor = cv.Get2D(self.bgrImage, y, x)
            self.selectedYCrCbColor = cv.Get2D(self.ycrcbImage, y, x)
            self.findBall()
        
    def yThresholdChanged(self, sliderPos):
        self.yThreshold = sliderPos
        print "New Y Threshold: " + str(self.yThreshold)
        self.findBall()
    
    def crThresholdChanged(self, sliderPos):
        self.crThreshold = sliderPos
        print "New Cr Threshold: " + str(self.crThreshold)
        self.findBall()
        
    def cbThresholdChanged(self, sliderPos):
        self.cbThreshold = sliderPos
        print "New Cb Threshold: " + str(self.cbThreshold)
        self.findBall()
        
    def findBall(self):
        cv.Zero(self.resultImage)
        cv.Zero(self.tmpGrayImage)
        if self.selectedYCrCbColor != None:
            cv.InRangeS(self.ycrcbImage,
                       cv.Scalar(self.selectedYCrCbColor[0] - self.yThreshold,
                                 self.selectedYCrCbColor[1] - self.crThreshold,
                                 self.selectedYCrCbColor[2] - self.cbThreshold, 0),
                       cv.Scalar(self.selectedYCrCbColor[0] + self.yThreshold,
                                 self.selectedYCrCbColor[1] + self.crThreshold,
                                 self.selectedYCrCbColor[2] + self.cbThreshold, 0), self.tmpGrayImage)
            cv.Copy(self.bgrImage, self.resultImage, self.tmpGrayImage)
            contour = cv.FindContours(self.tmpGrayImage, cv.CreateMemStorage())
            while contour != None:
                if cv.ContourArea(contour) > 10*10:
                    rect = cv.BoundingRect(contour)
                    if abs(rect[2] - rect[3]) < 4:
                        cv.Rectangle(self.resultImage, (rect[0], rect[1]), (rect[0] + rect[2], rect[1] + rect[3]), cv.CV_RGB(255, 0, 0));
                        print "Ball found in: (" + str(rect[0] + rect[2]/2) + "," + str(rect[1] + rect[3]/2) + ")"
                contour = contour.h_next()
        cv.ShowImage('Result', self.resultImage)
    
if __name__ == '__main__':
    image = cv.LoadImage('test.jpg')
    
    cv.NamedWindow('Threshold')
    cv.NamedWindow('Source')
    cv.NamedWindow('Result')
    
    cv.ShowImage('Source', image)
    
    SimpleBallRecognition(image)

    cv.WaitKey(0)
    cv.DestroyAllWindows()