app.py

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


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))

    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

    doc = ezdxf.new(units=0)
    doc.units = ezdxf.units.IN
    doc.header["$INSUNITS"] = ezdxf.units.IN

    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):
    doc = ezdxf.new()
    msp = doc.modelspace()

    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)


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))

def predict(image, offset_inches):
    try:
        reference_obj_img, scaling_box_coords = detect_reference_square(image)
    except Exception as e:
        raise gr.Error(f"Unable to DETECT COIN, please take another picture with different magnification level! Error: {e}")

    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.png",
            target_image=reference_square_mask,
            feature_detector="ORB",
        )
    except ZeroDivisionError:
        scaling_factor = None
        print("Error calculating scaling factor: Division by zero")
    except Exception as e:
        scaling_factor = None
        print(f"Error calculating scaling factor: {e}")

    # Default to a scaling factor of 1.0 if calculation fails
    if scaling_factor is None or scaling_factor == 0:
        scaling_factor = 1.0
        print("Using default scaling factor of 1.0 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.5,
    )
    objects_mask = resize_img(objects_mask, (image.shape[1], image.shape[0]))
    
    # Ensure offset_inches is valid
    if scaling_factor != 0:
        offset_pixels = (offset_inches / 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)

    ifer = gr.Interface(
        fn=predict,
        inputs=[
            gr.Image(label="Input Image"),
            gr.Number(label="Offset value for Mask(inches)", value=0.075),
        ],
        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 inches",
            ),
        ],
        examples=[
            ["./examples/Test20.jpg", 0.075],
            ["./examples/Test21.jpg", 0.075],
            ["./examples/Test22.jpg", 0.075],
            ["./examples/Test23.jpg", 0.075],
        ],
    )
    ifer.launch(share=True)