Python sript which detects cpecial LED marker and determines a position.

live_pos.py

# Launch script on the lab computer(Manual)
# 1)log in (standart lab password)
# 2) launch command line
# 3) command 'idle'
# 4) go to recent files and open finalScript)
# 5) plug in camera
# 6) put on 950 nm  ir-filter
# 7) run script
# 8) Good Job!
import cv2
import numpy as np
import time

video = cv2.VideoCapture(1)
video.set(4,960)
video.set(3,1280)

# pattern noise filtering


def ok(a,arr):
    for index,i in enumerate(arr):
        if((i[0]-a[0])**2+(i[1]-a[1])**2)<20:
            if(i[2]>a[2]):
                return
            else:
                arr.pop(index)
                arr.append(a)
                return
    arr.append(a)
    return

def order(pts):
    answer = np.zeros((6, 2), dtype="float32")
    s = pts.sum(axis=1)
    # the top-left point will have the smallest sum, whereas
    # the bottom-right point will have the largest sum
    answer[0] = pts[np.argmin(s)]
    answer[5] = pts[np.argmax(s)]


    # now, compute the difference between the points, the
    # top-right point will have the smallest difference,
    # whereas the bottom-left will have the largest difference
    diff = np.diff(pts, axis=1)
    answer[3] = pts[np.argmin(diff)]
    answer[1] = pts[np.argmax(diff)]
    # the middle point will have the smallest sum, whereas
    # the middle-bottom is the largesr
    pred_middle = (answer[0]+answer[5])/2
    pred_bot = (answer[3]+answer[5])/2

    nearest = 1000000
    point = pts[0]

    #middle
    for i in pts:
        sm = (i[0]-pred_middle[0])**2+(i[1]-pred_middle[1])**2
        if sm<nearest:
            nearest = sm
            point = i
    answer[2] =point


    nearest = 1000000
    point = pts[0]

    #bot-middle
    for i in pts:
        sm = (i[0]-pred_bot[0])**2+(i[1]-pred_bot[1])**2
        if sm<nearest:
            nearest = sm
            point = i
    answer[4] =point
    return answer


def check_line(a,b,c):
    m = (a[1]-b[1])/(a[0]-b[0])
    koef = m*a[0]+a[1]

    # check third dot
    value = c[0]*m+koef - c[1]
    if(c<=0.1):
        return true
    
    
    

def marker_points(image,COEF = 0.3):
    # image should be in rgb
    gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    (minVal, maxVal, minLoc, maxLoc) = cv2.minMaxLoc(gray_image)
    bound = maxVal*COEF
    raw_pixels = np.argwhere(gray_image >= bound)
    values = np.reshape(gray_image[gray_image>=bound],(len(raw_pixels),1))
    raw_pixels =  np.concatenate((raw_pixels,values),1)
    median = np.median(raw_pixels, axis=0)
    pixels = []
    # clear from crap pixels
    for i in raw_pixels:
        if((i[0]-median[0])**2+(i[1]-median[1])**2)<600:
            ok(i,pixels)
    # show?
    #print 'marker time'
    #print time.time()-stamp
    if len(pixels)==6:
            #cv2.imshow("IMAGE", image)
            #for i in pixels:
#                cv2.circle(image, (i[1],i[0]), 3, (0, 255, 0), -1)
#                cv2.imshow("IMAGE", image)
#    print pixels
#            cv2.waitKey(0)
#            cv2.destroyAllWindows()
            #return np.array(pixels,dtype="double")
            return np.array(pixels,dtype="double")[:,:-1]
    else:
        return pixels

        
    #Safety version
    #elif len(pixels)>6:
        #return marker_points(image,COEF+0.05)
    #else:
#        return marker_points(image,COEF-0.05)


def pose_estimation(image):
    #stamp =  time.time()
    model_points = np.array([
                            (0.0, 0.0, 0.0),
                            (0.0, 120.0, 0.0),                    
                            (60.0, 60.0, 0.0),
                            (120.0, 0.0, 0.0),
                            (120.0, 60.0, 0.0), 
                            (120.0, 120.0, 0.0),
                            ])

    #reorganize dots
    size = image.shape
    stamp =  time.time()
    image_points = marker_points(image)
    if len(image_points)!=6:
        return None
    #print image_points
    
    focal_length = size[1]
    center = (size[1]/2, size[0]/2)
    camera_matrix = np.array(
                             [[focal_length, 0, center[0]],
                             [0, focal_length, center[1]],
                             [0, 0, 1]], dtype = "double"
                             )
    camera_matrix = np.array(
                             [[1.10616906e+03*1.05,0.00000000e+00,6.35668602e+02],
                             [0.00000000e+00,1.10819769e+03*1.05,4.96993394e+02],
                             [0, 0, 1]], dtype = "double"
                             )
    
    #print "Camera Matrix :\n {0}".format(camera_matrix)
    image_points = np.array(image_points,dtype = "double")
    image_points = order(image_points)
    image_points[:,[0,1]] = image_points[:,[1,0]]
    #dist_coeffs = np.zeros((4,1)) # Assuming no lens distortion
    dist_coeffs = np.array([-0.20523672,0.11735183,-0.00085103,-0.0005939,0.00980464])# need to recalibrate everything
    (success, rotation_vector, translation_vector) = cv2.solvePnP(model_points, image_points, camera_matrix, dist_coeffs)#, flags=cv2.CV_ITERATIVE)
    #print 'translation:'
    #x-right y-down
    #a = image_points[0]
    #b = image_points[1]
    #c = image_points[2]
    #distance = ((a[0]-b[0])**2+(a[1]-b[1])**2)**0.5
    #one = 12./distance
    #ans = [(c[1]-640)*one,(c[0]-480)*one,translation_vector[2]]
    translation_vector[0]+=60
    translation_vector[1]+=60
    #return ans
    return translation_vector



def live_pos():
    while True:
        ret, image = video.read()
        coords = pose_estimation(image)
        if coords is not None:
            font = cv2.FONT_HERSHEY_SIMPLEX
            cv2.putText(image,', '.join([str(x) for x in coords]),(150,150), font, 1,(255,255,255),2,cv2.LINE_AA)
            cv2.imshow("Image", image)
        else:
            cv2.imshow("Image", image)
        
        k = cv2.waitKey(50) & 0xFF
        if k == -27:
            break
    
    video.release()
    cv2.destroyAllWindows()

live_pos()