distortion.py

import cv2
import numpy as np
import ezdxf
import gradio as gr
from pathlib import Path

coordis =  []

def save_matrix(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_matrix(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 correct_image(image, yaml):
    image = cv2.imread(image)
    mtx, dist = load_matrix(yaml.name)
    dst = cv2.undistort(image, mtx, dist, None, None)    
    return dst


def color_tab(file_path, yaml):

    coordinates = []
    img = cv2.imread(file_path)
    
    mtx, dist = load_matrix(yaml.name)
    
    img = cv2.undistort(img, mtx, dist, None, None)
    
    height, width, _ = img.shape
    
    # Convert the image to grayscale
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    
    # Apply a threshold to convert the grayscale image into a binary image
    _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY)
    
    # Find the contours in the binary image
    contours, hierarchy = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    # Select the contour with the largest area (the object we want to extract the corners from)
    contour = max(contours, key=cv2.contourArea)
    
    # Approximate the contour with a polygon
    epsilon = 0.01 * cv2.arcLength(contour, True)
    approx = cv2.approxPolyDP(contour, epsilon, True)

    # Draw the polygon on the original image
    cv2.polylines(img, [approx], True, (0, 255, 0), thickness=2)
  
        
    doc = ezdxf.new("R2010", setup=True)
    msp = doc.modelspace()

    # Print the coordinates of the corners
    for corner in approx:
        x, y = corner[0]
        coordinates.append([x,y])
        
    #This method of adding line is brute force and needs to be changed, find a way to generalize for polygons
    msp.add_line((coordinates[0][0], -coordinates[0][1]), (coordinates[1][0], -coordinates[1][1]))
    msp.add_line((coordinates[1][0], -coordinates[1][1]), (coordinates[2][0], -coordinates[2][1]))
    msp.add_line((coordinates[2][0], -coordinates[2][1]), (coordinates[3][0], -coordinates[3][1]))
    msp.add_line((coordinates[3][0], -coordinates[3][1]), (coordinates[0][0], -coordinates[0][1]))
    msp.add_line((0,0), (0, -height))
    msp.add_line((0, -height), (width, -height))
    msp.add_line((width, -height), (width, 0))
    msp.add_line((width, 0), (0,0))

    doc.saveas("output.dxf") 
    return "output.dxf", img

def corner_tab(img, evt: gr.SelectData):
    
    height, width, _ = img.shape

    row, col = evt.index
    coordis.append([row, col])
    
    if len(coordis) == 4 :
        coordinates = np.array(coordis)
        dwg = ezdxf.new("R2010")
        msp = dwg.modelspace()
        dwg.layers.new(name="greeny green lines", dxfattribs={"color": 3})
        
        msp.add_line((coordinates[0][0], -coordinates[0][1]), (coordinates[1][0], -coordinates[1][1]))
        msp.add_line((coordinates[1][0], -coordinates[1][1]), (coordinates[2][0], -coordinates[2][1]))
        msp.add_line((coordinates[2][0], -coordinates[2][1]), (coordinates[3][0], -coordinates[3][1]))
        msp.add_line((coordinates[3][0], -coordinates[3][1]), (coordinates[0][0], -coordinates[0][1]))
        
        msp.add_line((0,0), (0, -height))
        msp.add_line((0, -height), (width, -height))
        msp.add_line((width, -height), (width, 0))
        msp.add_line((width, 0), (0,0))
        
        
        dwg.saveas("output.dxf") 
        coordis.clear()

        return "output.dxf"


def generate_matrix(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)-1, int(board_horz)-1)
    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]
    try:
        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_matrix(mtx, dist, f"{filename_stem}.yml")
        return f"{filename_stem}.yml"
    except:
        print("Please check the Chessboard Dimensions")
        
        
def slider(img, h1, s1, v1, h2, s2, v2):
    # Load the image
    img = cv2.imread(img)

    while(1):
        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)

        # Create the NumPy arrays
        lower_red = np.array([h1, s1, v1])
        upper_red = np.array([h2, s2, v2])
        
        # Convert every element to integer using int() function
        lower_red = np.array([int(x) for x in lower_red])
        upper_red = np.array([int(x) for x in upper_red])
        

        mask = cv2.inRange(img, lower_red, upper_red)
        res = cv2.bitwise_and(img,img, mask= mask)
    
        return res