#!/usr/bin/python
"""
homecart.py

Home a Racetrack cart using vision.

By: Ralph Stirling
Based on vision example fitellipse.py by:
Original C implementation by:  Denis Burenkov.
Python implementation by: Roman Stanchak, James Bowman
"""

import sys
import urllib2
import random
import cv2.cv as cv
import time
import hal

def contour_iterator(contour):
    while contour:
        yield contour
        contour = contour.h_next()

class FitEllipse:

    def __init__(self, source_image, slider_pos):
        self.source_image = source_image
        offset, located = self.process_image(116)
        self.offset = offset
        self.located = located

    def process_image(self, slider_pos): 
        """
        This function finds contours, draws them and their approximation by ellipses.
        """
        stor = cv.CreateMemStorage()
        
        # Create the destination images
        image02 = cv.CloneImage(self.source_image)
        cv.Zero(image02)
        image04 = cv.CreateImage(cv.GetSize(self.source_image), cv.IPL_DEPTH_8U, 3)
        cv.Zero(image04)

        # Threshold the source image. This needful for cv.FindContours().
        cv.Threshold(self.source_image, image02, slider_pos, 255, cv.CV_THRESH_BINARY)

        # Find all contours.
        cont = cv.FindContours(image02,
            stor,
            cv.CV_RETR_EXTERNAL,
            cv.CV_CHAIN_APPROX_NONE,
            (0, 0))

        for c in contour_iterator(cont):
            # Number of points must be more than or equal to 6 for cv.FitEllipse2
            if len(c) >= 6:
                # Copy the contour into an array of (x,y)s
                PointArray2D32f = cv.CreateMat(1, len(c), cv.CV_32FC2)
                for (i, (x, y)) in enumerate(c):
                    PointArray2D32f[0, i] = (x, y)
                
                # Draw the current contour in gray
                gray = cv.CV_RGB(100, 100, 100)
                cv.DrawContours(image04, c, gray, gray,0,1,8,(0,0))
                
                # Fits ellipse to current contour.
                (center, size, angle) = cv.FitEllipse2(PointArray2D32f)
                w=size[0]
                h=size[1]

                # see if this is the desired size of circle
                if (w>280 and w<290 and h>280 and h<290):

                    # convert offset to inches from center
                    ppi = (w + h)/1.25     # .625in diameter
                    off = (240.0 - center[1]) / ppi
                    #print "x=",center[0], "y=",center[1], "w=",size[0], "h=",size[1],"ppi=",ppi,"off=",off
                    return off, True
# didn't find any
        return 0.0, False

if __name__ == '__main__':

    h = hal.component("homecart")
    h.newpin("offset", hal.HAL_FLOAT, hal.HAL_OUT)
    h.newpin("located", hal.HAL_BIT, hal.HAL_OUT)
    h.ready()

    capture = cv.CaptureFromCAM(0)

    while True:
        img = cv.QueryFrame(capture)

        gimg = cv.CreateImage((640,480), cv.IPL_DEPTH_8U, 1)

        cv.CvtColor(img, gimg, cv.CV_RGB2GRAY)

        fe = FitEllipse(gimg, 116)
        h['offset'] = fe.offset
        h['located'] = fe.located
        #print "off=",h['offset'],"located=",h['located']