example_frames.py

# coding: utf-8

# In[1]:

#imports
import cv2
import numpy as np
from os import listdir
from os.path import isfile, join
import glob
import os
import time
import math
from collections import defaultdict

# In[2]:

criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)

def nothing(x):
    pass


# In[ ]:

cv2.destroyAllWindows()


# In[3]:

visualize = True
def detect(img):
    img_gs = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
    img_gs = clahe.apply(img_gs)
    img_gs_f = np.float32(img_gs)
    dst = cv2.cornerHarris(img_gs_f,2,3,0.04)
    #result is dilated for marking the corners, not important
    dst = cv2.dilate(dst,None)
    
    coords = np.transpose(np.nonzero(dst>0.01*dst.max()))
    # Threshold for an optimal value, it may vary depending on the image.
    #if visualize:
        #img[dst>0.01*dst.max()]=[0,0,255]
    #print coords
    return (coords,img_gs_f, img_gs)

def dist2(a, b):
    return (a[0]-b[0])*(a[0]-b[0]) + (a[1]-b[1])*(a[1]-b[1])

def subimage(image, center, ts, tc, width, height):
    v_x = (tc, ts)
    v_y = (-ts,tc)
    s_x = center[0] - v_x[0] * (width / 2) - v_y[0] * (height / 2)
    s_y = center[1] - v_x[1] * (width / 2) - v_y[1] * (height / 2)

    mapping = np.array([[v_x[0],v_y[0], s_x],
                        [v_x[1],v_y[1], s_y]])


    return cv2.warpAffine(
        image,
        mapping,
        (int(width), int(height)),
        flags=cv2.WARP_INVERSE_MAP,
        borderMode=cv2.BORDER_REPLICATE)

#get image of a line between two points with given width
def subimage2(image, c1, c2, width):
    w = c1[0]-c2[0]
    h = c1[1]-c2[1]
    center = topoint((c1+c2)/2)
    c = math.sqrt(dist2(c1, c2))
    ts = -w/c
    tc = h/c
    #print (topoint(c1), topoint(c2), center, ts, tc, w, h, c)
    return subimage(image, center, ts, tc, width, c)

def circleCoords(coords, mask_zero, radius = 2, maxSize = 22):
    mask = mask_zero[:,:,0].copy()
    for c in coords:
        cv2.circle(mask, (c[1],c[0]), radius,255,-1)
    ret, markers, stats, centroids = cv2.connectedComponentsWithStats(mask)
    if visualize:#for visualization only
        max_ = np.amax(markers)
        markers *= 255/max_
        markers = markers.astype(np.uint8)
        markers_col = cv2.applyColorMap(markers, cv2.COLORMAP_JET)
    candidates = []
    for (i,c) in enumerate(centroids):
        if  (stats[i,cv2.CC_STAT_WIDTH] < maxSize and stats[i,cv2.CC_STAT_HEIGHT] < maxSize and
            (stats[i,cv2.CC_STAT_WIDTH] > 4*radius or stats[i,cv2.CC_STAT_HEIGHT] >4*radius)):
            if visualize:
                cv2.circle(markers_col, (int(c[0]),int(c[1])), maxSize/2,(0,0,255),1)
            candidates += [c]
    if visualize:#for visualization only
        cv2.imshow('markers_col', markers_col)
    return candidates
    

def topoint(x):
    return (int(x[0]),int(x[1]))
    

def findGraph(candidates, img_gs, max_mean=40):
    mask = np.zeros_like(img_gs)
    graph = defaultdict(lambda:[])
    if visualize:
        cv2.imshow('clahe', img_gs)
    
    for (i,c) in enumerate(candidates):
        for j in range(i):
            sub = subimage2(img_gs, c, candidates[j],10)
            min_ = np.amin(sub, axis = 1)
            
            
            if min_.mean() < max_mean:
                if visualize:
                    cv2.line(mask, topoint(c), topoint(candidates[j]),255,1)
                graph[i] += [j]
                graph[j] += [i]
    if visualize:
        cv2.imshow('graph', mask)
    return graph, mask

#http://stackoverflow.com/a/246063
def CrossProductZ(a,b):
    return a[0] * b[1] - a[1] * b[0]

def Orientation(a, b, c):
    return CrossProductZ(a, b) + CrossProductZ(b, c) + CrossProductZ(c, a)

def OrientationC(g, a, b, c):
    return Orientation(g[a],g[b],g[c])

def add1(a):
    return np.append(a,[1])

def findHouse(graph, mask_, img, candidates, primary=5, secondary=4, name='_X', color=(0,0,255),img_gs=None):
    if visualize:
        mask = mask_.copy()
    for i in graph: #subgraph localization
        if len(graph[i]) != 2:
            continue
        [a,b] = graph[i]
        if len(graph[a]) != primary or len(graph[a]) != primary:
            continue
        ok = 1
        for c in graph[a]:
            if c != b and c!= i and len(graph[c]) != secondary:
                ok = 0
                break
        if ok == 0:
            continue
        if OrientationC(candidates,i,a,b)<0:
            (a,b) = (b,a)
        for x in [a]+graph[a]:
            for y in graph[x]:
                if x<y:
                    if visualize:
                        cv2.line(mask, topoint(candidates[x]), topoint(candidates[y]),255,3)
                    cv2.line(img, topoint(candidates[x]), topoint(candidates[y]),color,1)
        used = {i: 1, a:1, b:1}
        
        x = 0
        s = 0
        if name == '_X':#find middle point
            for j in filter(lambda x: x not in used, graph[a]):
                s_ = 0
                for k in graph[j]:
                    s_ += dist2(candidates[j],candidates[k])
                if x == 0 or s_ < s:
                    x = j
                    s = s_
            used[x] = 1
        cd = filter(lambda x: x not in used, graph[a])
        if len(cd) != 2:
            #print "cd has len %d" %len(cd)
            return (None, None)
        [c,d] = cd
        if OrientationC(candidates,x,c,d)<0:
            (c,d) = (d,c)
    
        points_imdg = [candidates[__] for __ in [a,b,c,d]]#,x,i
        if name == '_X':
            points_img = [candidates[__] for __ in [a,b,c,d,x]]#,x,i
            new_3d = np.float32([[100,0,0],[0,0,0],[100,100,0],[0,100,0],[50,50,0]])
        else:
            points_img = [candidates[__] for __ in [a,b,c,d]]#,x,i
            new_3d = np.float32([[100,0,0],[0,0,0],[100,100,0],[0,100,0]])
        p_img = np.array(points_img, np.float32)
        cv2.cornerSubPix(img_gs,p_img,(11,11),(-1,-1),criteria)
        #cv2.circle(img, topoint(p_img[5]), 5, (255,0,0),-1)
        cv2.circle(img, topoint(p_img[0]), 5, (0,0,255),-1)#tr
        cv2.circle(img, topoint(p_img[1]), 5, (0,255,0),-1)#tl
        #cv2.circle(img, topoint(p_img[4]), 5, (255,255,255),-1)
        cv2.circle(img, topoint(p_img[2]), 5, (0,255,255),-1)#br
        cv2.circle(img, topoint(p_img[3]), 5, (255,0,255),-1)#bl
        return draw3d(img,p_img,new_3d)
    #   
    return (None,None)

goal_3d = np.float32([[50,50,-1]])
def draw3d(img, pts, new_3d):
    fx = 0.5 + cv2.getTrackbarPos('focal', 'dst') / 50.0
    h, w = img.shape[:2]
    K = np.float64([[fx*w, 0, 0.5*(w-1)],
                    [0, fx*w, 0.5*(h-1)],
                    [0.0,0.0,      1.0]])
    dist_coef = np.zeros(4)
    ret, rvec, tvec = cv2.solvePnP(new_3d, pts, K, dist_coef)
    goal_3d[0,2] = -1 * cv2.getTrackbarPos('height', 'dst')
    verts = cv2.projectPoints(goal_3d, rvec, tvec, K, dist_coef)[0].reshape(-1, 2)
    cv2.circle(img, topoint(verts[0]), 5, (255,255,255),-1)
    for p in pts:
        cv2.line(img, topoint(p), topoint(verts[0]), (255,255,255),2, -1)

def detectHouse(frame, house_type='_X'):
    (coords_,img_gs_f, img_gs) = detect(frame)
    candidates = circleCoords(coords_, np.zeros_like(frame))
    graph, mask = findGraph(candidates, img_gs,max_mean=cv2.getTrackbarPos('max_mean', 'dst'))
    findHouse(graph, mask, frame, candidates, 5, 4, '_X', img_gs=img_gs)
    #findHouse(graph, mask, frame, candidates, 3, 2, '_N', img_gs=img_gs)
    cv2.imshow('dst',frame)


#scanned pictures example
mypath = 'frames'
cv2.namedWindow('dst')
cv2.createTrackbar('focal', 'dst', 25, 50, nothing)
cv2.createTrackbar('height', 'dst', 50, 200, nothing)
cv2.createTrackbar('max_mean', 'dst', 40, 200, nothing)

for f in glob.glob(os.path.join(mypath,'*.jpg')):
    frame_orig = cv2.imread(f)
    while True:
        frame = frame_orig.copy()
        detectHouse(frame)
        key = cv2.waitKey(100)
        if key & 0xFF == ord('q') or key & 0xFF == ord(' '):
            break
    if key & 0xFF == ord('q'):
        break