Camera Horizon Calibration Script

import cv2
    import numpy as np
    import glob
    import pickle
    import time
    from picamera2 import Picamera2
    
    class PiCarX_CameraCalibrator:
      #This is to ensure that our program is detecting a 7x7 chessboard pattern with 6x6 corners
        def __init__(self, pattern_size=(6, 6), calibration_file='camera_calibration.pkl'):
            self.pattern_size = pattern_size
            self.calibration_file = calibration_file
            self.camera_matrix = None
            self.dist_coeffs = None
            self.calibrated = False
            self.tilt_angles = []
    
            self.objp = np.zeros((pattern_size[0] * pattern_size[1], 3), np.float32)
            self.objp[:, :2] = np.mgrid[0:pattern_size[0], 0:pattern_size[1]].T.reshape(-1, 2)
    
      #Take 20 images in order to detect the corners
        def capture_calibration_images(self, num_images=20, save_path='calibration_images/'):
            picam2 = Picamera2()
            config = picam2.create_still_configuration(main={"size": (640, 480)})
            picam2.configure(config)
            picam2.start()
            time.sleep(1)
    
            print(f"Capturing {num_images} calibration images...")
    
            for i in range(num_images):
                input(f"Press Enter to capture image {i+1}/{num_images}...")
                image = picam2.capture_array()
                cv2.imwrite(f"{save_path}calib_{i:02d}.jpg", image)
                print(f"Saved {save_path}calib_{i:02d}.jpg")
    
            picam2.close()
            cv2.destroyAllWindows()
    
      #Show the pattern found on the chessboard that can be used to calibrate the camera
        def calibrate_camera(self, image_paths='calibration_images/*.jpg'):
            obj_points = []
            img_points = []
    
            images = glob.glob(image_paths)
            if not images:
                raise FileNotFoundError(f"No images found at {image_paths}")
    
            for fname in images:
                img = cv2.imread(fname)
                gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    
                ret, corners = cv2.findChessboardCorners(gray, self.pattern_size, None)
    
                if ret:
                    obj_points.append(self.objp)
                    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
                    corners_refined = cv2.cornerSubPix(gray, corners, (11, 11), (-1, -1), criteria)
                    img_points.append(corners_refined)
    
                    cv2.drawChessboardCorners(img, self.pattern_size, corners_refined, ret)
                    cv2.imshow('Chessboard Corners', img)
                    cv2.waitKey(500)
    
            cv2.destroyAllWindows()
    
            if not obj_points:
                raise ValueError("No chessboard patterns found in images!")
    
            ret, self.camera_matrix, self.dist_coeffs, rvecs, tvecs = cv2.calibrateCamera(
                obj_points, img_points, gray.shape[::-1], None, None)
    
            self.tilt_angles = [self._rotation_vector_to_euler_angles(rvec) for rvec in rvecs]
            
            #Roll -> Tilt, Pitch -> pan, Yaw -> Horizon
            print("Camera Tilt Angles (Roll, Pitch, Yaw in degrees):")
            for idx, (roll, pitch, yaw) in enumerate(self.tilt_angles):
                print(f"Image {idx+1}: Roll={roll:.2f}, Pitch={pitch:.2f}, Yaw={yaw:.2f}")
    
            self.avg_roll = np.mean([angle[0] for angle in self.tilt_angles])
            print(f"\nAverage roll angle: {self.avg_roll:.2f}° (positive = rotated clockwise)")
            
            avg_pitch = np.mean([angle[1] for angle in self.tilt_angles])
            print(f"\nAverage pitch angle: {avg_pitch:.2f}°")
            print("(Negative value means camera is pointing slightly downward)")
            print("You may want to manually adjust the camera tilt if needed")
    
            mean_error = self._calculate_reprojection_error(obj_points, img_points, rvecs, tvecs)
            print(f"\nCalibration complete! Reprojection error: {mean_error:.5f} pixels")
    
            self.calibrated = True
            self._save_calibration()
    
        def _save_calibration(self):
            calibration_data = {
                'camera_matrix': self.camera_matrix,
                'dist_coeffs': self.dist_coeffs,
                'tilt_angles': self.tilt_angles
            }
            with open(self.calibration_file, 'wb') as f:
                pickle.dump(calibration_data, f)
            print(f"\nCalibration data saved to {self.calibration_file}")
    
        def load_calibration(self):
            with open(self.calibration_file, 'rb') as f:
                calibration_data = pickle.load(f)
            self.camera_matrix = calibration_data['camera_matrix']
            self.dist_coeffs = calibration_data['dist_coeffs']
            self.tilt_angles = calibration_data.get('tilt_angles', [])
            self.calibrated = True
            print("Calibration data loaded successfully")
    
        def undistort_image(self, image):
            if not self.calibrated:
                raise RuntimeError("Camera not calibrated! Call calibrate_camera() first")
    
            h, w = image.shape[:2]
            new_camera_mtx, roi = cv2.getOptimalNewCameraMatrix(
                self.camera_matrix, self.dist_coeffs, (w, h), 0, (w, h))
    
            dst = cv2.undistort(image, self.camera_matrix, self.dist_coeffs, None, new_camera_mtx)
            x, y, w, h = roi
            dst = dst[y:y + h, x:x + w]
            return dst
                        
        def find_best_roll(self, image, angle_range=(-15, 15), step=0.5):
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
            edges = cv2.Canny(gray, 50, 150, apertureSize=3)
    
            lines = cv2.HoughLines(edges, 1, np.pi / 180, threshold=100)
    
            if lines is None:
                print("No lines detected!")
                return 0.0
    
            angles = []
            for rho, theta in lines[:, 0]:
                angle = (theta * 180 / np.pi)  # Convert from radians to degrees
                if angle < -90:
                    angle += 180
                if angle > 90:
                    angle -= 180
                angles.append(angle)	
    
            if len(angles) == 0:
                print("No valid near-horizontal lines found!")
                return 0.0
    
            avg_angle = np.mean(angles)
            return avg_angle
    
                  
    if __name__ == "__main__":
        calibrator = PiCarX_CameraCalibrator()
    
        # Step 1: Uncomment to capture new calibration images
        #calibrator.capture_calibration_images(num_images=15)
    
        # Step 2: Calibrate camera
        # calibrator.calibrate_camera()
    
        # Step 3: Test undistortion
        picam2 = Picamera2()
        config = picam2.create_still_configuration(main={"size": (640, 480)})
        picam2.configure(config)
        picam2.start()
        time.sleep(10)
    
        # Capture image
        image = picam2.capture_array()
        cv2.imshow("Original Image", image)
        
        # After capturing
        print("Checking captured image...")
        if image is None:
            raise ValueError("Image is None — failed to capture anything.")
        print("Image shape:", image.shape)
        print("Image dtype:", image.dtype)
        print("Image type:", type(image))
    
        if len(image.shape) != 3 or image.shape[2] != 3:
           raise ValueError(f"Bad image shape: {image.shape}")
    
        if not isinstance(image, np.ndarray):
            raise ValueError(f"Image is not a numpy array: {type(image)}")
    
        if image.dtype != np.uint8: 
           print(f"Fixing image dtype {image.dtype} to uint8")
           image = image.astype(np.uint8)
    
        print("Image is valid, proceeding to process.")
    
    
        picam2.close()
    
        # Find best roll correction
        best_roll_angle = calibrator.find_best_roll(image)
    
        print(f"Best roll correction angle: {best_roll_angle:.2f} degrees")
    
        # Correct the image 
        h, w = image.shape[:2]
        center = (w // 2, h // 2)
        M = cv2.getRotationMatrix2D(center, angle= -best_roll_angle, scale=1.0) 
        leveled_image = cv2.warpAffine(image, M, (w, h), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE)
        
        #By running this code multiple times through multiple distances we found that the horizon tilt in our camera is -2.5
        print("the tilt in our camera is -2.5")
    
        # Show
        cv2.imshow("Leveled Image", leveled_image)
        cv2.waitKey(0)
        cv2.destroyAllWindows()