added distortion correction

app.py

@@ -1,8 +1,17 @@
 import numpy as np
 import gradio as gr
+from distortion import main
+import os
+
+
+
 
 coords = [[0,0]]
 
+import gradio as gr
+import cv2
+
+
 
 SHARED_UI_WARNING = f'''##### Lorem ipsum dolor sit amet, consectetur adipiscing elit. Sed tempus dignissim odio, at elementum erat vulputate sit amet. Vestibulum sodales viverra fermentum. In ac hendrerit dolor, vitae mattis odio. Maecenas suscipit consectetur suscipit. Curabitur sodales dui eget neque venenatis tincidunt. In sed libero mi. Nam sollicitudin metus urna, sit amet sagittis ex laoreet sed.
 
@@ -44,13 +53,18 @@ with gr.Blocks() as demo:
     with gr.Tab("O-Matrix"):
         with gr.Row():
             with gr.Column():
-                image_input = gr.Image(label="Select Image")
-                mask_input = gr.Button("Upload the Image with Chess Grid")    #will call the opencv screen
-                download_dxf = gr.Button("Download the YAML File") #make visible once process is done
-            with gr.Column():
-                gr.Markdown(SHARED_UI_WARNING)
+                distorted_image = gr.Image(label="Selected Image", type="filepath")
+
+                with gr.Row():
+                    chess_vert_input = gr.Textbox(label="Enter the number of squares in vertical direction")
+                    chess_horz_input = gr.Textbox(label="Enter the number of squares in horizontal direction")
+                download_dxf = gr.Button("Generate the YAML File")
+                yaml_file = gr.File()
                 
                 
+                download_dxf.click(main, inputs=[distorted_image, chess_vert_input, chess_horz_input], outputs= yaml_file)
+            with gr.Column():
+                gr.Markdown(SHARED_UI_WARNING)
 
 
-demo.launch()
+demo.launch(debug=True)

distortion.py

@@ -0,0 +1,72 @@
+import cv2
+import numpy as np
+from pathlib import Path
+
+
+def save_coefficients(mtx, dist, path):
+    """Save camera matrix and distortion coefficients to file."""
+    cv_file = cv2.FileStorage(path, cv2.FILE_STORAGE_WRITE)
+    cv_file.write('K', mtx)
+    cv_file.write('D', dist)
+    cv_file.release()
+
+
+def load_coefficients(path):
+    """Load camera matrix and distortion coefficients from file."""
+    cv_file = cv2.FileStorage(path, cv2.FILE_STORAGE_READ)
+    camera_matrix = cv_file.getNode('K').mat()
+    dist_matrix = cv_file.getNode('D').mat()
+    cv_file.release()
+    return [camera_matrix, dist_matrix]
+
+
+def main(filename, board_vert, board_horz):
+    """Main function to calibrate camera and undistort image."""
+
+    filename_stem = Path(filename).stem
+
+    # Define the checkerboard pattern size and criteria for corner detection
+    CHECKERBOARD = (int(board_vert), int(board_horz))
+    #CHECKERBOARD = (9, 7)
+    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
+
+    # Initialize object points and image points arrays
+    object_points = []  # 3D points in real world space
+    image_points = []  # 2D points in image plane
+
+    # Create the object points for the chessboard corners
+    objectp3d = np.zeros((1, CHECKERBOARD[0] * CHECKERBOARD[1], 3), np.float32)
+    objectp3d[0, :, :2] = np.mgrid[0:CHECKERBOARD[0], 0:CHECKERBOARD[1]].T.reshape(-1, 2)
+
+    # Load the image and convert to grayscale
+    image = cv2.imread(filename)
+    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+
+    # Find chessboard corners
+    ret, corners = cv2.findChessboardCorners(gray, CHECKERBOARD, cv2.CALIB_CB_ADAPTIVE_THRESH
+                                             + cv2.CALIB_CB_FAST_CHECK + cv2.CALIB_CB_NORMALIZE_IMAGE)
+
+    # If corners found, add object points and image points to arrays
+    if ret:
+        object_points.append(objectp3d)
+        corners2 = cv2.cornerSubPix(gray, corners, (11, 11), (-1, -1), criteria)
+        image_points.append(corners2)
+        image = cv2.drawChessboardCorners(image, CHECKERBOARD, corners2, ret)
+
+    # Calibrate camera using object points and image points
+    h, w = gray.shape[:2]
+    ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(object_points, image_points, gray.shape[::-1], None, None)
+
+    # Save camera matrix and distortion coefficients to file
+    save_coefficients(mtx, dist, f"{filename_stem}.yml")
+
+    # Load camera matrix and distortion coefficients from file
+    mtx, dist = load_coefficients(f'{filename_stem}.yml')
+
+    # Undistort the original image using camera matrix and distortion coefficients
+    original = cv2.imread(filename)
+    dst = cv2.undistort(original, mtx, dist, None, None)
+
+    # Save the undistorted image
+    cv2.imwrite(f'undist_{filename_stem}.jpg', dst)
+    return f"{filename_stem}.yml"