Spaces:
Paused
Paused
| # 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 yolo_detect( | |
| # image: Union[str, Path, int, Image.Image, list, tuple, np.ndarray, torch.Tensor], | |
| # classes: List[str], | |
| # ) -> np.ndarray: | |
| # drawer_detector = YOLOWorld("yolov8x-worldv2.pt") | |
| # drawer_detector.set_classes(classes) | |
| # results: List[Results] = drawer_detector.predict(image) | |
| # boxes = [] | |
| # for result in results: | |
| # boxes.append( | |
| # save_one_box(result.cpu().boxes.xyxy, im=result.orig_img, save=False) | |
| # ) | |
| # del drawer_detector | |
| # return boxes[0] | |
| # 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.2, | |
| # ) -> 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 | |
| # ) | |
| # # 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) | |
| # # 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 shrink_bbox(image: np.ndarray, shrink_factor: float): | |
| # """ | |
| # Crops the central 80% of the image, maintaining proportions for non-square images. | |
| # Args: | |
| # image: Input image as a NumPy array. | |
| # Returns: | |
| # Cropped image as a NumPy array. | |
| # """ | |
| # height, width = image.shape[:2] | |
| # center_x, center_y = width // 2, height // 2 | |
| # # Calculate 80% dimensions | |
| # new_width = int(width * shrink_factor) | |
| # new_height = int(height * shrink_factor) | |
| # # Determine the top-left and bottom-right points for cropping | |
| # x1 = max(center_x - new_width // 2, 0) | |
| # y1 = max(center_y - new_height // 2, 0) | |
| # x2 = min(center_x + new_width // 2, width) | |
| # y2 = min(center_y + new_height // 2, height) | |
| # # Crop the image | |
| # cropped_image = image[y1:y2, x1:x2] | |
| # return cropped_image | |
| # 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("./last.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: | |
| # drawer_img = yolo_detect(image, ["box"]) | |
| # shrunked_img = make_square(shrink_bbox(drawer_img, 0.90)) | |
| # except: | |
| # raise gr.Error("Unable to DETECT DRAWER, please take another picture with different magnification level!") | |
| # # Detect the scaling reference square | |
| # try: | |
| # reference_obj_img, scaling_box_coords = detect_reference_square(shrunked_img) | |
| # except: | |
| # raise gr.Error("Unable to DETECT REFERENCE BOX, please take another picture with different magnification level!") | |
| # # 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_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") | |
| # # 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=1.2, | |
| # ) | |
| # objects_mask = resize_img( | |
| # objects_mask, (shrunked_img.shape[1], shrunked_img.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) | |
| # ) | |
| # # 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) | |
| # shrunked_img_contours = cv2.drawContours( | |
| # shrunked_img, contours, -1, (0, 0, 255), thickness=2 | |
| # ) | |
| # dxf = save_dxf_spline(contours, scaling_factor, processed_size[0]) | |
| # return ( | |
| # cv2.cvtColor(shrunked_img_contours, cv2.COLOR_BGR2RGB), | |
| # 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) | |
| 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 | |
| # --- DOCUMENTATION STRINGS (Drawer Detection App) --- | |
| GUIDELINE_SETUP = """ | |
| ## 1. Quick Start Guide: Setup and DXF Generation | |
| This application analyzes an image of items inside a drawer, calculates scaling, and outputs a manufacturing-ready DXF file with offsets applied. | |
| 1. **Upload Image:** Upload a clear image of the drawer area, ensuring the items and the scaling reference box are visible. | |
| 2. **Set Offset:** Enter the desired offset value in **inches**. This determines the clearance around the contour (e.g., 0.075 inches is the default). | |
| 3. **Run:** Click the **"Submit"** button (or run using an example). | |
| 4. **Review & Download:** Review the resulting images (Contoured Output, Outlines, Mask) and download the generated **DXF file**. | |
| """ | |
| GUIDELINE_INPUT = """ | |
| ## 2. Expected Inputs and Preprocessing | |
| | Input Field | Purpose | Requirement | | |
| | :--- | :--- | :--- | | |
| | **Input Image** | A high-resolution image of the drawer containing the objects to be contoured. | Must show the items and the reference scaling box clearly. | | |
| | **Offset value (inches)** | The physical distance (clearance) to be added around the detected contours for manufacturing tolerance. | Input must be a positive number (float). Default is 0.075 inches. | | |
| """ | |
| GUIDELINE_OUTPUT = """ | |
| ## 3. Expected Outputs (Manufacturing Results) | |
| The application provides five key outputs: | |
| 1. **Ouput Image:** The original cropped drawer image overlaid with the final, offset contours (blue lines). | |
| 2. **Outlines of Objects:** A grayscale image showing only the final, smoothed contour lines. | |
| 3. **DXF file (Downloadable):** The primary output. This file contains scaled 2D spline geometry (in inches) based on the calculated contours, ready for CAD or CNC machines. | |
| 4. **Mask:** The raw, dilated binary mask used to generate the contours. | |
| 5. **Scaling Factor (Textbox):** The calculated ratio (in pixels per inch) used to accurately convert pixel dimensions into real-world units for the DXF file. | |
| """ | |
| # ---------------------------------------------------- | |
| # END GUIDELINE DEFINITIONS | |
| # ---------------------------------------------------- | |
| 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( | |
| image: Union[str, Path, int, Image.Image, list, tuple, np.ndarray, torch.Tensor], | |
| classes: List[str], | |
| ) -> np.ndarray: | |
| drawer_detector = YOLOWorld("yolov8x-worldv2.pt") | |
| drawer_detector.set_classes(classes) | |
| results: List[Results] = drawer_detector.predict(image) | |
| boxes = [] | |
| for result in results: | |
| boxes.append( | |
| save_one_box(result.cpu().boxes.xyxy, im=result.orig_img, save=False) | |
| ) | |
| del drawer_detector | |
| gc.collect() # Ensure memory is cleared | |
| return boxes[0] | |
| 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) | |
| # 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.2, | |
| ) -> 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 | |
| ) | |
| # Create a blank image to draw contours | |
| 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 shrink_bbox(image: np.ndarray, shrink_factor: float): | |
| """ | |
| Crops the central portion of the image. | |
| """ | |
| height, width = image.shape[:2] | |
| center_x, center_y = width // 2, height // 2 | |
| # Calculate dimensions | |
| new_width = int(width * shrink_factor) | |
| new_height = int(height * shrink_factor) | |
| # Determine the top-left and bottom-right points for cropping | |
| x1 = max(center_x - new_width // 2, 0) | |
| y1 = max(center_y - new_height // 2, 0) | |
| x2 = min(center_x + new_width // 2, width) | |
| y2 = min(center_y + new_height // 2, height) | |
| # Crop the image | |
| cropped_image = image[y1:y2, x1:x2] | |
| return cropped_image | |
| 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. | |
| """ | |
| 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("./last.pt") | |
| res = box_detector.predict(img, conf=0.05) | |
| del box_detector | |
| gc.collect() | |
| 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: | |
| drawer_img = yolo_detect(image, ["box"]) | |
| shrunked_img = make_square(shrink_bbox(drawer_img, 0.90)) | |
| except: | |
| raise gr.Error("Unable to DETECT DRAWER, please take another picture with different magnification level!") | |
| # Detect the scaling reference square | |
| try: | |
| reference_obj_img, scaling_box_coords = detect_reference_square(shrunked_img) | |
| except: | |
| raise gr.Error("Unable to DETECT REFERENCE BOX, please take another picture with different magnification level!") | |
| # 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]) | |
| ) | |
| scaling_factor = 1.0 | |
| try: | |
| scaling_factor = calculate_scaling_factor( | |
| reference_image_path="./Reference_ScalingBox.jpg", | |
| target_image=reference_square_mask, | |
| feature_detector="ORB", | |
| ) | |
| except ZeroDivisionError: | |
| print("Error calculating scaling factor: Division by zero") | |
| except Exception as e: | |
| print(f"Error calculating scaling factor: {e}") | |
| # Default to a scaling factor of 1.0 if calculation fails or is 0 | |
| 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") | |
| # 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=1.2, | |
| ) | |
| objects_mask = resize_img( | |
| objects_mask, (shrunked_img.shape[1], shrunked_img.shape[0]) | |
| ) | |
| # Ensure offset_inches is valid | |
| # Calculate pixel dilation amount: (offset_inches / scaling_factor) * 2 + 1 | |
| # We use 1.0 / scaling_factor because scaling_factor is px/inch. | |
| if scaling_factor > 0: | |
| # Convert inches to pixels | |
| offset_pixels = int(offset_inches / scaling_factor * 2) + 1 | |
| else: | |
| offset_pixels = 1 | |
| # Dilate mask for offset | |
| dilated_mask = cv2.dilate( | |
| objects_mask, np.ones((offset_pixels, 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( | |
| shrunked_img, contours, -1, (0, 0, 255), thickness=2 | |
| ) | |
| dxf = save_dxf_spline(contours, scaling_factor, processed_size[0]) | |
| return ( | |
| cv2.cvtColor(shrunked_img_contours, cv2.COLOR_BGR2RGB), | |
| outlines, | |
| dxf, | |
| dilated_mask, | |
| scaling_factor, | |
| ) | |
| if __name__ == "__main__": | |
| os.makedirs("./outputs", exist_ok=True) | |
| # Use gr.Blocks to allow for the structured guideline accordion | |
| with gr.Blocks(title="Drawer Contouring and DXF Generator") as demo: | |
| gr.Markdown("<h1 style='text-align: center;'>Drawer Contouring and DXF Generator (YOLO + BiRefNet)</h1>") | |
| gr.Markdown("Tool for generating scaled manufacturing contours from an input image.") | |
| # 1. Guidelines Section | |
| with gr.Accordion(" Tips & User Guidelines", open=False): | |
| gr.Markdown(GUIDELINE_SETUP) | |
| gr.Markdown("---") | |
| gr.Markdown(GUIDELINE_INPUT) | |
| gr.Markdown("---") | |
| gr.Markdown(GUIDELINE_OUTPUT) | |
| # 2. Main Interface | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| gr.Markdown("## Step 1: Upload an Image") | |
| input_image = gr.Image(label="1. Input Image", type="numpy") | |
| gr.Markdown("## Step 2: Set the offset value (Optional) ") | |
| offset_input = gr.Number(label="2. Offset value for Mask (inches)", value=0.075) | |
| gr.Markdown("## Step 3: Click the button ") | |
| submit_button = gr.Button(" Process and Generate DXF", variant="primary") | |
| with gr.Column(scale=2): | |
| gr.Markdown("## Results") | |
| scaling_output = gr.Textbox( | |
| label="Scaling Factor (pixels/inch)", | |
| placeholder="Calculated conversion rate", | |
| ) | |
| output_image = gr.Image(label="Output Image (Contours Drawn)") | |
| with gr.Row(): | |
| output_outlines = gr.Image(label="Outlines of Objects") | |
| output_mask = gr.Image(label="Final Dilated Mask") | |
| dxf_file = gr.File(label="DXF file (Download)") | |
| # 3. Examples Section | |
| gr.Markdown("## Examples ") | |
| gr.Examples( | |
| examples=[ | |
| ["./examples/Test20.jpg", 0.075], | |
| ["./examples/Test21.jpg", 0.075], | |
| ["./examples/Test22.jpg", 0.075], | |
| ["./examples/Test23.jpg", 0.075], | |
| ], | |
| inputs=[input_image, offset_input], | |
| outputs=[output_image, output_outlines, dxf_file, output_mask, scaling_output], | |
| fn=predict, | |
| cache_examples=False, | |
| label="Example Images (Click to load and run)", | |
| ) | |
| # Event Handler | |
| submit_button.click( | |
| fn=predict, | |
| inputs=[input_image, offset_input], | |
| outputs=[output_image, output_outlines, dxf_file, output_mask, scaling_output], | |
| ) | |
| demo.queue() | |
| demo.launch(share=True) |