## A forráskód az OpenCV függvénykönyvtár hivatalos tutorialja alapján 
## készült. 
## Dr. Németh Gábor 
## Szegedi Tudományegyetem

import numpy as np
import cv2
import cv2.fisheye
import glob
import os

IMAGES_DIR='fisheye_15mm_grid22mm'
gridsize = 22 ## laptop:22, monitor 37 mm egy négyzet a mintán
SENSOR_WIDTH=22.9    #  canon eos 50d = 22.3mm
SENSOR_HEIGHT=14.9   #  canon eos 50d = 14.9mm

#megállási feltétel a subpixel kereséshez
criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)

#valós világbeli pontok
real_world_points = np.zeros( (1,6*9, 3), np.float32) 

# egy 6x9-os minta lesz, amely tuladonképpen egy rács, a Z koordináta mindenhol 0
real_world_points[0,:,:2] = np.mgrid[0:6*gridsize:gridsize,0:9*gridsize:gridsize].T.reshape(-1,2) 

# ezekben a tömbökben tároljuk a pontokat minden fényképhez
objpoints = [] # a 3D valós világbeli pontoknak
imgpoints = [] # a képen detektált pontoknak

images = glob.glob(IMAGES_DIR+'/Chessboard*.jpg')
print(images)
img = cv2.imread(IMAGES_DIR+'/Chessboard_01.jpg')
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
imgsize = gray.shape[::-1]

## -----------------
## képek beolvasása
## -----------------
image_counter = 0;
# végighaladunk a képeken
print('reading images:'); 
for fname in images:
	print(fname)
	img = cv2.imread(fname)
	gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) # a képet szürkeárnyalatossá konvertáljuk
	
	## Megkeressük a sarokpontokat a sakktáblán
	ret, corners = cv2.findChessboardCorners(gray, (6,9), None, cv2.CALIB_CB_FAST_CHECK | cv2.CALIB_CB_NORMALIZE_IMAGE | cv2.CALIB_CB_ADAPTIVE_THRESH );
	
	## ha találtunk (ret == True), akkor 
	## az objpoints tömbhöz hozzáadjuk a valós világbeli pontrácsot
	## az imgpointshoz pedig a képen detektált sarokpontokat
	print(ret)
	if ret == True:
		objpoints.append(real_world_points)
		corners2 = cv2.cornerSubPix(gray, corners, (11,11), (-1,-1), criteria)
		imgpoints.append(corners2)
		
		# kirajzoljuk a képre a megtalált sarokpontokat 
		##print('rajzolunk')
		img = cv2.drawChessboardCorners(img, (6,9), corners2, ret)
		cv2.imwrite('corners_' + os.path.basename(fname), img )
		#cv2.imshow('img', img)
		#cv2.waitKey(500);
		
#cv2.destroyAllWindows();


## --------------------
## Kalibráció 
## --------------------
calibration_flags = cv2.fisheye.CALIB_RECOMPUTE_EXTRINSIC+cv2.fisheye.CALIB_CHECK_COND+cv2.fisheye.CALIB_FIX_SKEW
# rms:   visszatérési érték
# K:   kamera mátrix
# D:  disztorziós együtthatók
# rvecs: forgatás mátrix paraméterek
# tvecs: eltolás vektor paraméterek
K = np.zeros((3,3))
D = np.zeros((4,1))
N_OK = len(objpoints)
rvecs = [np.zeros((1,1,3), dtype=np.float64) for i in range(N_OK)]
tvecs = [np.zeros((1,1,3), dtype=np.float64) for i in range(N_OK)]
rms, _, _, _, _ = cv2.fisheye.calibrate(objpoints, imgpoints, imgsize, K, D, rvecs, tvecs, calibration_flags, (cv2.TERM_CRITERIA_EPS+cv2.TERM_CRITERIA_MAX_ITER, 30, 1e-6))

fovx, fovy, focalLength, principalPoint, aspectRatio = cv2.calibrationMatrixValues(K,imgsize, SENSOR_WIDTH,SENSOR_HEIGHT)
## kb 5mm-es apertúra mellett kapom meg a valódi értékeket

#print('Calibration result: ' + ret)
#if ret == True:
print('K = ' +str(K))
print('D = ' + str(D))
print('rms = ' + str(rms))
print('focalLength = ' + str(focalLength) )
print('fovx = ' + str(fovx) )
print('fovy = ' + str(fovy) )
print('aspectRatio = ' + str(aspectRatio) )