app.py

import cv2
import ezdxf
import gradio as gr
import numpy as np
from distortion import generate_matrix



def load_coefficients(path):
    """Load camera matrix and distortion coefficients from file."""
    cv_file = cv2.FileStorage(path.name, 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_coefficients(yaml)
    dst = cv2.undistort(image, mtx, dist, None, None)
    return dst



coordis =  []

def get_select_coords(img, evt: gr.SelectData):

    row, col = evt.index
    coordis.append([row, col])
    selected_coords.value = f"Selected Coordinates: ({row}, {col})"
    
    if len(coordis) == 4 :
        coordinates = np.array(coordis)
        print("shape of second array:", coordinates.shape)
        print(coordinates)
        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]))
        
        dwg.saveas("output.dxf") 
        coordis.clear()

        return "DXF File has been saved successfully"  
                

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.

Pellentesque nunc turpis, porta ut accumsan eget, iaculis nec odio. Praesent fringilla a sem sed elementum. Proin orci justo, rutrum et feugiat eleifend, auctor sed odio. Maecenas posuere urna tortor, ut euismod ligula mattis sed. Sed ipsum velit, pretium sed lacinia sed, placerat eget urna. Mauris lobortis mi vitae odio luctus viverra sed eget urna. Pellentesque blandit semper felis sed congue. Aenean congue enim id euismod finibus.

Aliquam finibus accumsan orci, vitae pellentesque ex fringilla a. Fusce fermentum, arcu sit amet dapibus lacinia, lacus enim maximus ipsum, vitae consectetur ligula odio non dolor. Donec sed felis eget odio dignissim sagittis. Pellentesque non nunc hendrerit odio mattis finibus. Quisque tellus massa, aliquet a condimentum ut, faucibus vel leo. Aliquam quis dui sed felis viverra fringilla. Vestibulum volutpat sem ut tortor mollis facilisis. Mauris pharetra elit a purus pulvinar, eget interdum orci feugiat. Vestibulum auctor gravida ante. Curabitur a elit efficitur, vestibulum neque quis, imperdiet turpis. Donec enim magna, rutrum ut finibus sit amet, mollis non magna. Ut quis magna velit. Mauris varius dolor in mauris faucibus, dignissim blandit nibh dapibus. Quisque dictum risus sit amet est finibus, a feugiat enim imperdiet.
Quisque vitae lacus ac nunc hendrerit consequat. Donec rhoncus ultrices erat, vitae commodo quam volutpat et.
'''

with gr.Blocks() as demo:
    title = """<p><h1 align="center" style="font-size: 36px;">Auto Allign Quadrilateral Workpiece</h1></p>"""
    gr.HTML(title)
    with gr.Tab("Allign"):
        with gr.Row():
            with gr.Column():
                image_input = gr.Image(label="Select Image")
                file_input = gr.File(label="Select O-Matrix")
                upload_button = gr.Button("Upload")
                mask_input = gr.Button("Initiate masking (Temporary)") #will call the opencv screen
            with gr.Column():
                image_output = gr.Image()
                download_dxf = gr.Button("Download the DXF File")
                
        
    with gr.Tab("Pick Corners"):
        with gr.Row():
            with gr.Column():
                title = """<p><h1 align="center" style="font-size: 24px;">First, let's correct the distorted Image</h1></p>"""
                gr.HTML(title)
                distorted_input = gr.Image(label="Select distorted Image", type="filepath")
                yml_input = gr.File(label="Select O-Matrix (.yml file)", type="file")
                correct_button = gr.Button("Correct the Image")
                
            with gr.Column():
                title = """<p><h1 align="center" style="font-size: 24px;">Then, let's mark the corners of the Workpiece</h1></p>"""
                gr.HTML(title)
                corrected_img = gr.Image(label="Corrected Image")   
                selected_coords = gr.Textbox()
                
        corrected_img.select(get_select_coords, [corrected_img], selected_coords)
        correct_button.click(fn=correct_image, inputs=[distorted_input, yml_input], outputs=corrected_img)
 
    with gr.Tab("O-Matrix"):
        with gr.Row():
            with gr.Column():
                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")
                error_box = gr.Textbox(label="Error", visible=False)
                
                yaml_file = gr.File()
                download_dxf.click(generate_matrix, inputs=[distorted_image, chess_vert_input, chess_horz_input], outputs= yaml_file)
            with gr.Column():
                gr.Markdown(SHARED_UI_WARNING)


demo.launch(debug=True)