import cv2
import numpy as np
from matplotlib import pyplot as plt

def drawlines(img1,img2,lines,pts1,pts2):
	''' img1 - image on which we draw the epilines for the points in img2
	lines - corresponding epilines '''
	r,c = img1.shape[:2]
	img1 = cv2.cvtColor(img1,cv2.COLOR_GRAY2BGR)
	img2 = cv2.cvtColor(img2,cv2.COLOR_GRAY2BGR)
	for r,pt1,pt2 in zip(lines,pts1,pts2):
		color = tuple(np.random.randint(0,255,3).tolist())
		x0,y0 = map(int, [0, -r[2]/r[1] ])
		x1,y1 = map(int, [c, -(r[2]+r[0]*c)/r[1] ])
		img1 = cv2.line(img1, (x0,y0), (x1,y1), color,1)
		img1 = cv2.circle(img1,tuple(pt1),5,color,-1)
		img2 = cv2.circle(img2,tuple(pt2),5,color,-1)
	return img1,img2



SOURCE_IMAGE1='../tea05.jpg'
SOURCE_IMAGE2='../tea08.jpg'

## képek beolvasása
img1 = cv2.imread(SOURCE_IMAGE1);
img2 = cv2.imread(SOURCE_IMAGE2);

## a képet szürkeárnyalatossá konvertáljuk
gray_img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
gray_img2 = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)

## jellemzőpontok detektálása 
surf = cv2.xfeatures2d.SURF_create()
keypoints1 = surf.detect(gray_img1, None)
keypoints2 = surf.detect(gray_img2, None)

## kulcspont leírók számítása
keypoints1, descriptors1 = surf.compute(gray_img1, keypoints1)
keypoints2, descriptors2 = surf.compute(gray_img2, keypoints2)

## pontpárok keresése
# FLANN parameterek
FLANN_INDEX_KDTREE = 0
index_params = dict(algorithm = FLANN_INDEX_KDTREE, trees = 5)
search_params = dict(checks=50)   # or pass empty dictionary

flann = cv2.FlannBasedMatcher(index_params,search_params)

matches = flann.knnMatch(descriptors1,descriptors2,k=2) ## ez kNN-alapú, 
                                                        ## minden pontnak két lehetséges párja lehet

# csak a jó párosításokat tároljuk el, amelyek átmentek a Lowe-teszten
good = []

for m,n in matches:
    if m.distance < 0.7*n.distance:
        good.append(m)

points1 = []
points2 = []

for m in good:
	points1.append(keypoints1[m.queryIdx].pt)
	points2.append(keypoints2[m.trainIdx].pt)
points1, points2 = np.float32((points1, points2))

draw_params = dict(matchColor = (0,255,0),
                   singlePointColor = (255,0,0),
#                   matchesMask = matchesMask,
                   flags = 0)

matching_img = cv2.drawMatchesKnn(img1,keypoints1,img2,keypoints2,matches[:300],None)
cv2.imwrite("matching_image.png", matching_img)

## A fundamentális mátrix meghatározása
F, F_mask = cv2.findFundamentalMat(points1, points2, cv2.FM_8POINT)

print("A fundamentális mátrix:")
print(F)

## Az epipoláris egyenesek meghatározása
lines1 = cv2.computeCorrespondEpilines(points2.reshape(-1,1,2), 2,F)
lines1 = lines1.reshape(-1,3)
img5,img6 = drawlines(gray_img1,gray_img2,lines1,points1,points2)
cv2.imwrite('lines1.png', img5)

lines2 = cv2.computeCorrespondEpilines(points1.reshape(-1,1,2), 1,F)
lines2 = lines2.reshape(-1,3)
img5,img6 = drawlines(gray_img2,gray_img1,lines2,points2,points1)
cv2.imwrite('lines2.png', img5)