Update app.py

app.py

@@ -2,6 +2,7 @@ import os
 from pathlib import Path
 from typing import List, Union
 from PIL import Image
+import ezdxf.units
 import numpy as np
 import torch
 from torchvision import transforms
@@ -14,7 +15,25 @@ import ezdxf
 import gradio as gr
 import gc
 from scalingtestupdated import calculate_scaling_factor
+from scipy.interpolate import splprep, splev
+from scipy.ndimage import gaussian_filter1d
 
+birefnet = AutoModelForImageSegmentation.from_pretrained(
+    "zhengpeng7/BiRefNet", trust_remote_code=True
+)
+
+device = "cpu"
+torch.set_float32_matmul_precision(["high", "highest"][0])
+
+birefnet.to(device)
+birefnet.eval()
+transform_image = transforms.Compose(
+    [
+        transforms.Resize((1024, 1024)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+    ]
+)
 
 
 def yolo_detect(
@@ -36,23 +55,6 @@ def yolo_detect(
 
 
 def remove_bg(image: np.ndarray) -> np.ndarray:
-    birefnet = AutoModelForImageSegmentation.from_pretrained(
-    "zhengpeng7/BiRefNet", trust_remote_code=True
-    )
-
-    device = "cpu"
-    torch.set_float32_matmul_precision(["high", "highest"][0])
-
-    birefnet.to(device)
-    birefnet.eval()
-    transform_image = transforms.Compose(
-        [
-            transforms.Resize((1024, 1024)),
-            transforms.ToTensor(),
-            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
-        ]
-    )
-
     image = Image.fromarray(image)
     input_images = transform_image(image).unsqueeze(0).to("cpu")
 
@@ -62,16 +64,62 @@ def remove_bg(image: np.ndarray) -> np.ndarray:
     pred = preds[0].squeeze()
 
     # Show Results
-    pred_pil = transforms.ToPILImage()(pred)
+    pred_pil: Image = transforms.ToPILImage()(pred)
+    print(pred_pil)
     # Scale proportionally with max length to 1024 for faster showing
     scale_ratio = 1024 / max(image.size)
     scaled_size = (int(image.size[0] * scale_ratio), int(image.size[1] * scale_ratio))
-    
-    del birefnet
 
     return np.array(pred_pil.resize(scaled_size))
 
-def exclude_scaling_box(image: np.ndarray, bbox: np.ndarray, orig_size: tuple, processed_size: tuple, expansion_factor: float = 1.5) -> np.ndarray:
+
+def make_square(img: np.ndarray):
+    # Get dimensions
+    height, width = img.shape[:2]
+
+    # Find the larger dimension
+    max_dim = max(height, width)
+
+    # Calculate padding
+    pad_height = (max_dim - height) // 2
+    pad_width = (max_dim - width) // 2
+
+    # Handle odd dimensions
+    pad_height_extra = max_dim - height - 2 * pad_height
+    pad_width_extra = max_dim - width - 2 * pad_width
+
+    # Create padding with edge colors
+    if len(img.shape) == 3:  # Color image
+        # Pad the image
+        padded = np.pad(
+            img,
+            (
+                (pad_height, pad_height + pad_height_extra),
+                (pad_width, pad_width + pad_width_extra),
+                (0, 0),
+            ),
+            mode="edge",
+        )
+    else:  # Grayscale image
+        padded = np.pad(
+            img,
+            (
+                (pad_height, pad_height + pad_height_extra),
+                (pad_width, pad_width + pad_width_extra),
+            ),
+            mode="edge",
+        )
+
+    return padded
+
+
+def exclude_scaling_box(
+    image: np.ndarray,
+    bbox: np.ndarray,
+    orig_size: tuple,
+    processed_size: tuple,
+    expansion_factor: float = 1.5,
+) -> np.ndarray:
     # Unpack the bounding box
     x_min, y_min, x_max, y_max = map(int, bbox)
 
@@ -89,9 +137,13 @@ def exclude_scaling_box(image: np.ndarray, bbox: np.ndarray, orig_size: tuple, p
     box_width = x_max - x_min
     box_height = y_max - y_min
     expanded_x_min = max(0, int(x_min - (expansion_factor - 1) * box_width / 2))
-    expanded_x_max = min(image.shape[1], int(x_max + (expansion_factor - 1) * box_width / 2))
+    expanded_x_max = min(
+        image.shape[1], int(x_max + (expansion_factor - 1) * box_width / 2)
+    )
     expanded_y_min = max(0, int(y_min - (expansion_factor - 1) * box_height / 2))
-    expanded_y_max = min(image.shape[0], int(y_max + (expansion_factor - 1) * box_height / 2))
+    expanded_y_max = min(
+        image.shape[0], int(y_max + (expansion_factor - 1) * box_height / 2)
+    )
 
     # Black out the expanded region
     image[expanded_y_min:expanded_y_max, expanded_x_min:expanded_x_max] = 0
@@ -99,6 +151,61 @@ def exclude_scaling_box(image: np.ndarray, bbox: np.ndarray, orig_size: tuple, p
     return image
 
 
+def resample_contour(contour):
+    # ---------------------------------------------------------------------------------------- #
+    # Get all the parameters at the start:
+    num_points = 1000
+    smoothing_factor = 5
+
+    smoothed_x_sigma = 1
+    smoothed_y_sigma = 1
+    # ---------------------------------------------------------------------------------------- #
+    contour = contour[:, 0, :]
+
+    tck, _ = splprep([contour[:, 0], contour[:, 1]], s=smoothing_factor)
+
+    u = np.linspace(0, 1, num_points)
+    resampled_points = splev(u, tck)
+
+    smoothed_x = gaussian_filter1d(resampled_points[0], sigma=smoothed_x_sigma)
+    smoothed_y = gaussian_filter1d(resampled_points[1], sigma=smoothed_y_sigma)
+
+    return np.array([smoothed_x, smoothed_y]).T
+
+
+def save_dxf_spline(inflated_contours, scaling_factor, height):
+    # ---------------------------------------------------------------------------------------- #
+    # Get all the parameters at the start:
+    degree = 3
+    closed = True
+    # ---------------------------------------------------------------------------------------- #
+
+    doc = ezdxf.new(units=0)
+    doc.units = ezdxf.units.IN
+    doc.header['$INSUNITS'] = ezdxf.units.IN
+
+    msp = doc.modelspace()
+
+    for contour in inflated_contours:
+        resampled_contour = resample_contour(contour)
+        points = [
+            (x * scaling_factor, (height - y)* scaling_factor)
+            for x, y in resampled_contour
+        ]
+        if len(points) >= 3:
+            # Manually Closing the Contour in case it hasn't been closed by the contours before.
+            if np.linalg.norm(np.array(points[0]) - np.array(points[-1])) > 1e-2:
+                points.append(points[0])
+
+            spline = msp.add_spline(points, degree=degree)
+            spline.closed = closed
+
+    # Step 14: Save the DXF file
+    dxf_filepath = os.path.join("./outputs", "out.dxf")
+    doc.saveas(dxf_filepath)
+    return dxf_filepath
+
+
 def extract_outlines(binary_image: np.ndarray) -> np.ndarray:
     """
     Extracts and draws the outlines of masks from a binary image.
@@ -109,26 +216,21 @@ def extract_outlines(binary_image: np.ndarray) -> np.ndarray:
     """
     # Detect contours from the binary image
     contours, _ = cv2.findContours(
-        binary_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
+        binary_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
     )
 
+    # smooth_contours_list = []
+    # for contour in contours:
+    #     smooth_contours_list.append(smooth_contours(contour))
     # Create a blank image to draw contours
     outline_image = np.zeros_like(binary_image)
-    # Smooth the contours
-    smoothed_contours = []
-    for contour in contours:
-        # Calculate epsilon for approxPolyDP
-        epsilon = 0.002 * cv2.arcLength(contour, True)
-        # Approximate the contour with fewer points
-        smoothed_contour = cv2.approxPolyDP(contour, epsilon, True)
-        smoothed_contours.append(smoothed_contour)
 
     # Draw the contours on the blank image
     cv2.drawContours(
-        outline_image, smoothed_contours, -1, (255), thickness=1
+        outline_image, contours, -1, (255), thickness=1
     )  # White color for outlines
 
-    return cv2.bitwise_not(outline_image), smoothed_contours
+    return cv2.bitwise_not(outline_image), contours
 
 
 def shrink_bbox(image: np.ndarray, shrink_factor: float):
@@ -162,7 +264,7 @@ def shrink_bbox(image: np.ndarray, shrink_factor: float):
 #     lower_range_tuple = (0)
 
 #     doc = ezdxf.new('R2010')
-#     msp = doc.modelspace() 
+#     msp = doc.modelspace()
 #     masked_jpg = cv2.inRange(outlines,lower_range_tuple, upper_range_tuple)
 
 #     for i in range(0,masked_jpg.shape[0]):
@@ -173,6 +275,7 @@ def shrink_bbox(image: np.ndarray, shrink_factor: float):
 #     doc.saveas("./outputs/out.dxf")
 #     return "./outputs/out.dxf"
 
+
 def to_dxf(contours):
     doc = ezdxf.new()
     msp = doc.modelspace()
@@ -180,10 +283,11 @@ def to_dxf(contours):
     for contour in contours:
         points = [(point[0][0], point[0][1]) for point in contour]
         msp.add_lwpolyline(points, close=True)  # Add a polyline for each contour
-    
+
     doc.saveas("./outputs/out.dxf")
     return "./outputs/out.dxf"
 
+
 def smooth_contours(contour):
     epsilon = 0.01 * cv2.arcLength(contour, True)  # Adjust factor (e.g., 0.01)
     return cv2.approxPolyDP(contour, epsilon, True)
@@ -224,40 +328,68 @@ def detect_reference_square(img) -> np.ndarray:
     box_detector = YOLO("./last.pt")
     res = box_detector.predict(img)
     del box_detector
-    return save_one_box(res[0].cpu().boxes.xyxy, res[0].orig_img, save=False), res[0].cpu().boxes.xyxy[0
-                                                                                                       ]
+    return save_one_box(res[0].cpu().boxes.xyxy, res[0].orig_img, save=False), res[
+        0
+    ].cpu().boxes.xyxy[0]
 
-def predict(image):
+
+def resize_img(img: np.ndarray, resize_dim):
+    return np.array(Image.fromarray(img).resize(resize_dim))
+
+
+def predict(image, offset_inches):
     drawer_img = yolo_detect(image, ["box"])
-    shrunked_img = shrink_bbox(drawer_img, 0.8)
+    shrunked_img = make_square(shrink_bbox(drawer_img, 0.8))
     # Detect the scaling reference square
     reference_obj_img, scaling_box_coords = detect_reference_square(shrunked_img)
-    reference_obj_img_scaled = shrink_bbox(reference_obj_img, 1.2)
+    # reference_obj_img_scaled = shrink_bbox(reference_obj_img, 1.2)
+    # make the image sqaure so it does not effect the size of objects
+    reference_obj_img = make_square(reference_obj_img)
+    reference_square_mask = remove_bg(reference_obj_img)
+
+    # make the mask same size as org image
+    reference_square_mask = resize_img(
+        reference_square_mask, (reference_obj_img.shape[1], reference_obj_img.shape[0])
+    )
+
     try:
         scaling_factor = calculate_scaling_factor(
             reference_image_path="./Reference_ScalingBox.jpg",
-            target_image=reference_obj_img_scaled,
-            feature_detector="SIFT",
+            target_image=reference_square_mask,
+            feature_detector="ORB",
         )
     except:
         scaling_factor = 1.0
+
     # Save original size before `remove_bg` processing
     orig_size = shrunked_img.shape[:2]
     # Generate foreground mask and save its size
     objects_mask = remove_bg(shrunked_img)
+
     processed_size = objects_mask.shape[:2]
     # Exclude scaling box region from objects mask
     objects_mask = exclude_scaling_box(
-        objects_mask, scaling_box_coords, orig_size, processed_size, expansion_factor=3.0
+        objects_mask,
+        scaling_box_coords,
+        orig_size,
+        processed_size,
+        expansion_factor=3.0,
+    )
+    objects_mask = resize_img(
+        objects_mask, (shrunked_img.shape[1], shrunked_img.shape[0])
+    )
+    offset_pixels = offset_inches / scaling_factor
+    dilated_mask = cv2.dilate(
+        objects_mask, np.ones((int(offset_pixels), int(offset_pixels)), np.uint8)
     )
-    # Scale the object mask according to scaling factor
-    objects_mask_scaled = scale_image(objects_mask, scaling_factor)
-    Image.fromarray(objects_mask_scaled).save("./outputs/scaled_mask_new.jpg")
-    outlines, contours = extract_outlines(objects_mask_scaled)
-    dxf = to_dxf(contours)
 
-    return outlines, dxf, objects_mask, scaling_factor, reference_obj_img_scaled
+    # Scale the object mask according to scaling factor
+    # objects_mask_scaled = scale_image(objects_mask, scaling_factor)
+    Image.fromarray(dilated_mask).save("./outputs/scaled_mask_new.jpg")
+    outlines, contours = extract_outlines(dilated_mask)
+    dxf = save_dxf_spline(contours, scaling_factor, processed_size[0])
 
+    return outlines, dxf, dilated_mask, scaling_factor, reference_obj_img
 
 
 if __name__ == "__main__":
@@ -265,15 +397,22 @@ if __name__ == "__main__":
 
     ifer = gr.Interface(
         fn=predict,
-        inputs=[gr.Image(label="Input Image")],
+        inputs=[gr.Image(label="Input Image"), gr.Number(label="Offset value for Mask(inches)", value=0.075)],
         outputs=[
             gr.Image(label="Ouput Image"),
             gr.File(label="DXF file"),
             gr.Image(label="Mask"),
-            gr.Textbox(label="Scaling Factor(mm)", placeholder="Every pixel is equal to mentioned number in mm(milimeter)"),
-            gr.Image(label="Image used for calculating scaling factor")
+            gr.Textbox(
+                label="Scaling Factor(mm)",
+                placeholder="Every pixel is equal to mentioned number in mm(milimeter)",
+            ),
+            gr.Image(label="Image used for calculating scaling factor"),
+        ],
+        examples=[
+            ["./examples/Test20.jpg", 0.075],
+            ["./examples/Test21.jpg", 0.075],
+            ["./examples/Test22.jpg", 0.075],
+            ["./examples/Test23.jpg", 0.075],
         ],
-        examples=["./examples/Test20.jpg", "./examples/Test21.jpg", "./examples/Test22.jpg", "./examples/Test23.jpg"]
     )
     ifer.launch(share=True)
-