Upload app.py

app.py

@@ -0,0 +1,416 @@
+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
+from ultralytics import YOLOWorld, YOLO
+from ultralytics.engine.results import Results
+from ultralytics.utils.plotting import save_one_box
+from transformers import AutoModelForImageSegmentation
+import cv2
+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
+import json
+
+
+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 remove_bg(image: np.ndarray) -> np.ndarray:
+
+    image = Image.fromarray(image)
+    input_images = transform_image(image).unsqueeze(0).to("cpu")
+
+    # Prediction
+    with torch.no_grad():
+        preds = birefnet(input_images)[-1].sigmoid().cpu()
+    pred = preds[0].squeeze()
+
+    # Show Results
+    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))
+    print(f"scaled size {scaled_size}")
+
+    return np.array(pred_pil.resize(scaled_size))
+
+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)
+
+    # Calculate scaling factors
+    scale_x = processed_size[1] / orig_size[1]  # Width scale
+    scale_y = processed_size[0] / orig_size[0]  # Height scale
+
+    # Adjust bounding box coordinates
+    x_min = int(x_min * scale_x)
+    x_max = int(x_max * scale_x)
+    y_min = int(y_min * scale_y)
+    y_max = int(y_max * scale_y)
+
+    # Calculate expanded box coordinates
+    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_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)
+    )
+
+    # Black out the expanded region
+    image[expanded_y_min:expanded_y_max, expanded_x_min:expanded_x_max] = 0
+
+    return image
+
+def resample_contour(contour):
+    # Get all the parameters at the start:
+    num_points = 1000
+    smoothing_factor = 5
+    spline_degree = 3  # Typically k=3 for cubic spline
+
+    smoothed_x_sigma = 1
+    smoothed_y_sigma = 1
+
+    # Ensure contour has enough points
+    if len(contour) < spline_degree + 1:
+        raise ValueError(f"Contour must have at least {spline_degree + 1} points, but has {len(contour)} points.")
+
+    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):
+    degree = 3
+    closed = True
+
+    # Create a new DXF document with millimeters as the unit
+    doc = ezdxf.new(units=ezdxf.units.MM)
+    doc.units = ezdxf.units.MM  # Ensure units are millimeters
+    doc.header["$INSUNITS"] = ezdxf.units.MM  # Set insertion units to millimeters
+
+    msp = doc.modelspace()
+
+    for contour in inflated_contours:
+        try:
+            resampled_contour = resample_contour(contour)
+            points = [
+                (x * scaling_factor, (height - y) * scaling_factor)
+                for x, y in resampled_contour
+            ]
+            if len(points) >= 3:
+                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
+
+        except ValueError as e:
+            print(f"Skipping contour: {e}")
+
+    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.
+    Args:
+        binary_image: Grayscale binary image where white represents masks and black is the background.
+    Returns:
+        Image with outlines drawn.
+    """
+    # Detect contours from the binary image
+    contours, _ = cv2.findContours(
+        binary_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
+    )
+
+    outline_image = np.zeros_like(binary_image)
+
+    # Draw the contours on the blank image
+    cv2.drawContours(
+        outline_image, contours, -1, (255), thickness=1
+    )  # White color for outlines
+
+    return cv2.bitwise_not(outline_image), contours
+
+def to_dxf(contours):
+    # Create a new DXF document with millimeters as the unit
+    doc = ezdxf.new(units=ezdxf.units.MM)
+    doc.units = ezdxf.units.MM  # Ensure units are millimeters
+    doc.header["$INSUNITS"] = ezdxf.units.MM  # Set insertion units to millimeters)
+    msp = doc.modelspace()
+
+    try:
+        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
+    except Exception as e:
+        raise gr.Error(f"Unable to generate DXF: {e}")
+
+    output_path = "./outputs/out2.dxf"
+    doc.saveas(output_path)
+    return output_path
+
+def smooth_contours(contour):
+    epsilon = 0.01 * cv2.arcLength(contour, True)  # Adjust factor (e.g., 0.01)
+    return cv2.approxPolyDP(contour, epsilon, True)
+
+
+def scale_image(image: np.ndarray, scale_factor: float) -> np.ndarray:
+    """
+    Resize image by scaling both width and height by the same factor.
+    Args:
+        image: Input numpy image
+        scale_factor: Factor to scale the image (e.g., 0.5 for half size, 2 for double size)
+    Returns:
+        np.ndarray: Resized image
+    """
+    if scale_factor <= 0:
+        raise ValueError("Scale factor must be positive")
+
+    current_height, current_width = image.shape[:2]
+
+    # Calculate new dimensions
+    new_width = int(current_width * scale_factor)
+    new_height = int(current_height * scale_factor)
+
+    # Choose interpolation method based on whether we're scaling up or down
+    interpolation = cv2.INTER_AREA if scale_factor < 1 else cv2.INTER_CUBIC
+
+    # Resize image
+    resized_image = cv2.resize(
+        image, (new_width, new_height), interpolation=interpolation
+    )
+
+    return resized_image
+
+def detect_reference_square(img) -> np.ndarray:
+    box_detector = YOLO("./best1.pt")
+    res = box_detector.predict(img, conf=0.05)
+    del box_detector
+    return save_one_box(res[0].cpu().boxes.xyxy, res[0].orig_img, save=False), res[
+        0
+    ].cpu().boxes.xyxy[0]
+
+
+def resize_img(img: np.ndarray, resize_dim):
+    return np.array(Image.fromarray(img).resize(resize_dim))
+
+
+# Load selected language file
+def load_language(lang_code):
+    with open(f'./locales/{lang_code}.json', 'r') as f:
+        return json.load(f)
+
+def predict(image, offset, coin_size_mm):
+    # Load the selected language
+    # translations = load_language(language)
+
+    if offset <= -1:
+        raise gr.Error("Offset Value Can't be too SMALL for OBJECT's Mask Dilation")
+
+    try:
+        reference_obj_img, scaling_box_coords = detect_reference_square(image)
+    except:
+        raise gr.Error("Unable to detect the COIN. Please try again with different magnification.")
+
+    reference_obj_img = make_square(reference_obj_img)
+
+    reference_square_mask = remove_bg(reference_obj_img)
+
+    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="./coin.jpg",
+            target_image=reference_square_mask,
+            reference_obj_size_mm = coin_size_mm,
+            feature_detector="ORB",
+        )
+    except Exception as e:
+        scaling_factor = None
+        print(f"Error calculating scaling factor: {e}")
+
+    # Default to a scaling factor if calculation fails
+    if scaling_factor is None or scaling_factor == 0:
+        scaling_factor = 0.07
+        print("Using default scaling factor due to calculation error")
+
+    orig_size = image.shape[:2]
+    objects_mask = remove_bg(image)
+    processed_size = objects_mask.shape[:2]
+
+    objects_mask = exclude_scaling_box(
+        objects_mask,
+        scaling_box_coords,
+        orig_size,
+        processed_size,
+        expansion_factor=1.2,
+    )
+    objects_mask = resize_img(objects_mask, (image.shape[1], image.shape[0]))
+    
+    # Ensure offset_inches is valid
+    if scaling_factor != 0:
+        offset_pixels = (float(offset) / float(scaling_factor)) * 2 + 1
+    else:
+        offset_pixels = 1  # Default value in case of invalid scaling factor
+
+    dilated_mask = cv2.dilate(objects_mask, np.ones((int(offset_pixels), int(offset_pixels)), np.uint8))
+
+    Image.fromarray(dilated_mask).save("./outputs/scaled_mask_new.jpg")
+    outlines, contours = extract_outlines(dilated_mask)
+    shrunked_img_contours = cv2.drawContours(image, contours, -1, (0, 0, 255), thickness=2)
+    dxf = save_dxf_spline(contours, scaling_factor, processed_size[0])
+
+    return (
+        shrunked_img_contours,
+        outlines,
+        dxf,
+        dilated_mask,
+        scaling_factor,
+    )
+
+if __name__ == "__main__": 
+    os.makedirs("./outputs", exist_ok=True)
+
+    # Language selector in UI
+    ifer = gr.Interface(
+        fn=predict,
+        inputs=[
+            gr.Image(label="Input Image", type="numpy"),
+            gr.Number(label="Offset value for Mask(mm)", value=0.15),
+            gr.Number(label="Diameter of reference coin(mm). Adjust according to coin.", value=22),
+            # gr.Dropdown(["en", "nl"], label="Language", value="en")  # Default English
+        ],
+        outputs=[
+            gr.Image(label="Output Image"),
+            gr.Image(label="Outlines of Objects"),
+            gr.File(label="DXF file"),
+            gr.Image(label="Mask"),
+            gr.Textbox(
+                label="Scaling Factor(mm)",
+                placeholder="Every pixel is equal to mentioned number in millimeters",
+            ),
+        ],
+        examples=[
+            ["./examples/Test20.jpg", 0.15],
+            ["./examples/Test21.jpg", 0.15],
+            ["./examples/Test22.jpg", 0.15],
+            ["./examples/Test23.jpg", 0.15],
+        ],
+    )
+    ifer.launch(share=True)
+
+
+# if __name__ == "__main__":
+#     os.makedirs("./outputs", exist_ok=True)
+
+#     ifer = gr.Interface(
+#         fn=predict,
+#         inputs=[
+#             gr.Image(label="Input Image"),
+#             gr.Number(label="Offset value for Mask(mm)", value=2),
+#         ],
+#         outputs=[
+#             gr.Image(label="Ouput Image"),
+#             gr.Image(label="Outlines of Objects"),
+#             gr.File(label="DXF file"),
+#             gr.Image(label="Mask"),
+#             gr.Textbox(
+#                 label="Scaling Factor(mm)",
+#                 placeholder="Every pixel is equal to mentioned number in millimeters",
+#             ),
+#         ],
+#         examples=[
+#             ["./examples/Test20.jpg", 2],
+#             ["./examples/Test21.jpg", 2],
+#             ["./examples/Test22.jpg", 2],
+#             ["./examples/Test23.jpg", 2],
+#         ],
+#     )
+#     ifer.launch(share=True)
+
+
+