changed the variables names and fixed some functions

app.py

@@ -1,10 +1,8 @@
-import cv2
 import gradio as gr
-import numpy as np
-from distortion import generate_matrix, get_select_coords, correct_image, track, allign_tab
+from distortion import generate_matrix, corner_tab, correct_image, slider, color_tab
 
 
-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.
+CHESSBOARD_TUTORIAL = '''##### 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.
 
@@ -13,16 +11,17 @@ Quisque vitae lacus ac nunc hendrerit consequat. Donec rhoncus ultrices erat, vi
 '''
 
 with gr.Blocks() as demo:
-    title = """<p><h1 align="center" style="font-size: 36px;">Auto Allign Quadrilateral Workpiece</h1></p>"""
+    title = """<p><h1 align="center" style="font-size: 36px;">Auto Align Quadrilateral Workpieces</h1></p>"""
     gr.HTML(title)
-    with gr.Tab("Allign"):
-        
+    with gr.Tab("Color Alignment"):
+        title = """<p><h1 align="center" style="font-size: 24px;">We will use this tab to align the workpieces after masking wit slider inputs</h1></p>"""
+        gr.HTML(title)
         with gr.Row():
             with gr.Column():
-                image_input = gr.Image(label="Unmasked Image", type="filepath")
-                image_output = gr.Image(label="Masked Image", type="filepath")
-                file_input = gr.File(label="Select O-Matrix", type="file")
-                third_image = gr.Image(label="Detected Shape")
+                unmasked_distorted_image = gr.Image(label="Unmasked + Distorted Image", type="filepath")
+                masked_distorted_image = gr.Image(label="Masked + Distorted Image", type="filepath")
+                camera_matrix = gr.File(label="Camera Matrix", type="file")
+                detected_shape = gr.Image(label="Detected Shape")
                 
             with gr.Column():
                 first_slider = gr.Slider(0,255, value=0, label="1. Slider")
@@ -31,54 +30,50 @@ with gr.Blocks() as demo:
                 forth_slider = gr.Slider(0,255, value=255,label="4. Slider")
                 fifth_slider = gr.Slider(0,255, value=255,label="5. Slider")
                 sixth_slider = gr.Slider(0,255, value=255,label="6. Slider")
-                current_values = gr.Textbox("Current vallues of sliders")
-                download_dxf = gr.Button("Generate DXF File")
-                last_output = gr.File(label="Generated dxf file", type="file")
+                generate_dxf = gr.Button("Generate DXF File")
+                dxf_file = gr.File(label="Generated DXF file", type="file")
                 
-        download_dxf.click(fn=allign_tab, inputs=[image_output, file_input], outputs=[last_output, third_image])
+        generate_dxf.click(fn=color_tab, inputs=[masked_distorted_image, camera_matrix], outputs=[dxf_file, detected_shape])
 
-        
-        first_slider.change(track, inputs=[image_input, first_slider, second_slider, third_slider, forth_slider, fifth_slider, sixth_slider], outputs=[image_output, current_values])
-        second_slider.change(track, inputs=[image_input, first_slider, second_slider, third_slider, forth_slider, fifth_slider, sixth_slider], outputs=[image_output, current_values])
-        third_slider.change(track, inputs=[image_input, first_slider, second_slider, third_slider, forth_slider, fifth_slider, sixth_slider], outputs=[image_output, current_values])
-        forth_slider.change(track, inputs=[image_input, first_slider, second_slider, third_slider, forth_slider, fifth_slider, sixth_slider], outputs=[image_output, current_values])
-        fifth_slider.change(track, inputs=[image_input, first_slider, second_slider, third_slider, forth_slider, fifth_slider, sixth_slider], outputs=[image_output, current_values])
-        sixth_slider.change(track, inputs=[image_input, first_slider, second_slider, third_slider, forth_slider, fifth_slider, sixth_slider], outputs=[image_output, current_values])
+        first_slider.change(slider, inputs=[unmasked_distorted_image, first_slider, second_slider, third_slider, forth_slider, fifth_slider, sixth_slider], outputs=[masked_distorted_image])
+        second_slider.change(slider, inputs=[unmasked_distorted_image, first_slider, second_slider, third_slider, forth_slider, fifth_slider, sixth_slider], outputs=[masked_distorted_image])
+        third_slider.change(slider, inputs=[unmasked_distorted_image, first_slider, second_slider, third_slider, forth_slider, fifth_slider, sixth_slider], outputs=[masked_distorted_image])
+        forth_slider.change(slider, inputs=[unmasked_distorted_image, first_slider, second_slider, third_slider, forth_slider, fifth_slider, sixth_slider], outputs=[masked_distorted_image])
+        fifth_slider.change(slider, inputs=[unmasked_distorted_image, first_slider, second_slider, third_slider, forth_slider, fifth_slider, sixth_slider], outputs=[masked_distorted_image])
+        sixth_slider.change(slider, inputs=[unmasked_distorted_image, first_slider, second_slider, third_slider, forth_slider, fifth_slider, sixth_slider], outputs=[masked_distorted_image])
 
 
-    with gr.Tab("Pick Corners"):
+    with gr.Tab("Corner Allignment"):
         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")
+                distorted_image = gr.Image(label="Distorted Image", type="filepath")
+                yml_input = gr.File(label="Camera Matrix", type="file")
+                correct_button = gr.Button("Correct Image")
                 
             with gr.Column():
-                title = """<p><h1 align="center" style="font-size: 24px;">Then, let's mark the corners of the Workpiece</h1></p>"""
+                title = """<p><h1 align="center" style="font-size: 24px;">Then mark the corners of the Workpiece</h1></p>"""
                 gr.HTML(title)
-                corrected_img = gr.Image(label="Corrected Image")   
+                corrected_image = gr.Image(label="Corrected Image")   
                 dxf_output = gr.File(type="file")
                 
-        corrected_img.select(get_select_coords, [corrected_img], dxf_output)
-        correct_button.click(fn=correct_image, inputs=[distorted_input, yml_input], outputs=corrected_img)
+        corrected_image.select(corner_tab, [corrected_image], dxf_output)
+        correct_button.click(fn=correct_image, inputs=[distorted_image, yml_input], outputs=corrected_image)
  
-    with gr.Tab("O-Matrix"):
+    with gr.Tab("Matrix Generation"):
         with gr.Row():
             with gr.Column():
-                distorted_image = gr.Image(label="Selected Image", type="filepath")
-
+                distorted_image = gr.Image(label="Chessboard 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)
+                    chess_vert_input = gr.Textbox(label="Number of squares in vertical direction")
+                    chess_horz_input = gr.Textbox(label="Number of squares in horizontal direction")
+                    
+                yaml_button = gr.Button("Generate the YML File")
+                yaml_file = gr.File(label="YML File")
+                yaml_button.click(generate_matrix, inputs=[distorted_image, chess_vert_input, chess_horz_input], outputs= yaml_file)
+            
             with gr.Column():
-                gr.Markdown(SHARED_UI_WARNING)
-
+                gr.Markdown(CHESSBOARD_TUTORIAL)
 
 demo.launch(debug=True)

distortion.py

@@ -6,48 +6,89 @@ from pathlib import Path
 
 coordis =  []
 
-def save_coefficients(mtx, dist, path):
+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_coefficients_yaml(path):
+
+def load_matrix(path):
     """Load camera matrix and distortion coefficients from file."""
-    cv_file = cv2.FileStorage(path.name, cv2.FILE_STORAGE_READ)
+    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_coefficients_yaml(yaml)
-    dst = cv2.undistort(image, mtx, dist, None, None)
+    mtx, dist = load_matrix(yaml.name)
+    dst = cv2.undistort(image, mtx, dist, None, None)    
     return dst
 
 
+def color_tab(file_path, yaml):
 
-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]
+    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)
 
-def get_select_coords(img, evt: gr.SelectData):
+    # 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)
-        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})
@@ -57,7 +98,13 @@ def get_select_coords(img, evt: gr.SelectData):
         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_file = dwg.saveas("output.dxf") 
+        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"
@@ -70,7 +117,6 @@ def generate_matrix(filename, board_vert, board_horz):
 
     # Define the checkerboard pattern size and criteria for corner detection
     CHECKERBOARD = (int(board_vert)-1, int(board_horz)-1)
-    #CHECKERBOARD = (9, 7)
     criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
 
     # Initialize object points and image points arrays
@@ -101,19 +147,17 @@ def generate_matrix(filename, board_vert, board_horz):
     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_coefficients(mtx, dist, f"{filename_stem}.yml")
+        save_matrix(mtx, dist, f"{filename_stem}.yml")
         return f"{filename_stem}.yml"
     except:
         print("Please check the Chessboard Dimensions")
         
         
-def track(img, h1, s1, v1, h2, s2, v2):
+def slider(img, h1, s1, v1, h2, s2, v2):
     # Load the image
     img = cv2.imread(img)
 
-
     while(1):
-        #hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
         img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
 
         # Create the NumPy arrays
@@ -124,70 +168,9 @@ def track(img, h1, s1, v1, h2, s2, v2):
         lower_red = np.array([int(x) for x in lower_red])
         upper_red = np.array([int(x) for x in upper_red])
         
-        total = np.array([h1, s1, v1, h2, s2, v2])
 
         mask = cv2.inRange(img, lower_red, upper_red)
         res = cv2.bitwise_and(img,img, mask= mask)
-        #hsv = cv2.cvtColor(res, cv2.COLOR_BGR2HSV)
-    
-        return res, total
-    
-def allign_tab(file_path, yaml):
-    
-
-    coordinates = []
-    # Load the masked image
-    img = cv2.imread(file_path)
-    
-    mtx, dist = load_coefficients_yaml(yaml)
-    
-    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)
-    print(approx)
-    
-    # Draw the polygon on the original image
-    cv2.polylines(img, [approx], True, (0, 255, 0), thickness=2)
+        return res
     
-    # cv2.imshow('image', img)
-    # cv2.waitKey(0)
-    # cv2.destroyAllWindows()
-        
-    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])
-        #print(f"({x}, {y})")
-        
-        
-    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))
-    
-
-    dwf_file = doc.saveas("output.dxf") 
-    return "output.dxf", img