import numpy as np
import cv2
import glob

chess_points = np.zeros((6*7,3), np.float32)
size = (7,6)
square_size = 1

point_index = 0

#Costruiamo l'array di punti mondo della scacchiera, ovviamente è uguale per ogni immagine, visto che il sistema di riferimento è sempre posto sullo stesso vertice della scacchiera
for y in np.arange(0,size[1], square_size):
    for x in np.arange(0,size[0], square_size):
        chess_points[point_index] = np.asarray([x,y,0])
        point_index = point_index + 1

world_points = [] 
image_points = [] 

images = glob.glob('*.jpg')
image_size = None
used_images = []
for fname in images:
    img = cv2.imread(fname)
    gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)

    #Troviamo i punti immagine
    ret, corners = cv2.findChessboardCorners(gray, size)

    if ret == True:
        used_images.append(img)
        image_size = gray.shape[::-1]
        image_points.append(corners)
        world_points.append(chess_points)

    img = cv2.drawChessboardCorners(img, size, corners,ret)
    cv2.imshow('img',img)
    cv2.waitKey(1000)

#Calibriamo
ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(world_points, image_points, image_size, None, None)
print(f'{ret} - {mtx} - {dist} - {rvecs} - {tvecs}')


tot_error = 0

for i in range(len(world_points)):
    #Proiettiamo i punti mondo sull'immagine, usando la calibrazione appena calcolata
    image_points_proj, _ = cv2.projectPoints(world_points[i], rvecs[i], tvecs[i], mtx, dist)

    #Calcoliamo l'errore di riproiezione
    error = cv2.norm(image_points[i],image_points_proj, cv2.NORM_L2)
    tot_error += error
    img = cv2.drawChessboardCorners(used_images[i], size, image_points_proj,True)
    cv2.imshow('img',img)
    cv2.waitKey(1000)


print (f'total error:  {tot_error}')


cv2.destroyAllWindows()