import numpy as np
import cv2
import os

def load_images_folder(images_folder):
    images = []
    files = os.listdir(images_folder)
    for file in files:
        img = cv2.imread(os.path.join(images_folder, file))
        images.append(img)
    return images

def get_chessboard_points(chessboard_shape, dx, dy):
    return [[(i%chessboard_shape[0])*dx, (i//chessboard_shape[0])*dy, 0] for i in range(np.prod(chessboard_shape))]

images = load_images_folder('/Users/hugomcp/Desktop/input images_calibration/')

corners = [cv2.findChessboardCorners(i, (8,6)) for i in images]

valid_corners = [cor[1] for cor in corners if cor[0]]

num_valid_images = len(valid_corners)

real_points = get_chessboard_points((8, 6), 30, 30)

object_points = np.asarray([real_points for i in range(num_valid_images)], dtype=np.float32)

image_points = np.asarray(valid_corners, dtype=np.float32)

rms, intrinsics, dist_coeffs, rvecs, tvecs = cv2.calibrateCamera(object_points, image_points, images[1].shape[0:2], None, None)

extrinsics = list(map(lambda rvec, tvec: np.hstack((cv2.Rodrigues(rvec)[0], tvec)), rvecs, tvecs))

np.savez('calib_left', intrinsic=intrinsics, extrinsic=extrinsics)

# Lets print some outputs
print("Corners standard intrinsics:\n",intrinsics)
print("Corners standard dist_coefs:\n", dist_coeffs)
print("rms:\n", rms)