Finished bugs related to second tab

app.py

@@ -1,20 +1,54 @@
+import cv2
+import ezdxf
 import gradio as gr
-from distortion import main
+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]
 
 
-coords = [[0,0]]
+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
-    coords.append([row, col])
+    coordis.append([row, col])
     selected_coords.value = f"Selected Coordinates: ({row}, {col})"
-    return coords[-1]
-            
+    
+    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.
 
@@ -41,12 +75,22 @@ with gr.Blocks() as demo:
         
     with gr.Tab("Pick Corners"):
         with gr.Row():
-
-            input_img = gr.Image(label="Select Image")
-            selected_coords = gr.Textbox(label="Selected Coordinates")
-
-        input_img.select(get_select_coords, [input_img], selected_coords)
-
+            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():
@@ -59,7 +103,7 @@ with gr.Blocks() as demo:
                 error_box = gr.Textbox(label="Error", visible=False)
                 
                 yaml_file = gr.File()
-                download_dxf.click(main, inputs=[distorted_image, chess_vert_input, chess_horz_input], outputs= yaml_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)
 

distortion.py

@@ -21,7 +21,7 @@ def load_coefficients(path):
     return [camera_matrix, dist_matrix]
 
 
-def main(filename, board_vert, board_horz):
+def generate_matrix(filename, board_vert, board_horz):
     """Main function to calibrate camera and undistort image."""
 
     filename_stem = Path(filename).stem