Upload 11 files

.gitattributes

@@ -33,3 +33,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/Test20.jpg filter=lfs diff=lfs merge=lfs -text
+examples/Test21.jpg filter=lfs diff=lfs merge=lfs -text
+examples/Test22.jpg filter=lfs diff=lfs merge=lfs -text
+examples/Test23.jpg filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,279 @@
+import os
+from pathlib import Path
+from typing import List, Union
+from PIL import Image
+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
+
+
+
+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:
+    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")
+
+    # Prediction
+    with torch.no_grad():
+        preds = birefnet(input_images)[-1].sigmoid().cpu()
+    pred = preds[0].squeeze()
+
+    # Show Results
+    pred_pil = 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))
+    
+    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:
+    # 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 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_SIMPLE
+    )
+
+    # 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
+    )  # White color for outlines
+
+    return cv2.bitwise_not(outline_image), smoothed_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(outlines):
+#     upper_range_tuple = (200)
+#     lower_range_tuple = (0)
+
+#     doc = ezdxf.new('R2010')
+#     msp = doc.modelspace() 
+#     masked_jpg = cv2.inRange(outlines,lower_range_tuple, upper_range_tuple)
+
+#     for i in range(0,masked_jpg.shape[0]):
+#         for j in range(0,masked_jpg.shape[1]):
+#             if masked_jpg[i][j] == 255:
+#                 msp.add_line((j,masked_jpg.shape[0] - i), (j,masked_jpg.shape[0] - i))
+
+#     doc.saveas("./outputs/out.dxf")
+#     return "./outputs/out.dxf"
+
+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)
+    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 predict(image):
+    drawer_img = yolo_detect(image, ["box"])
+    shrunked_img = 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)
+    try:
+        scaling_factor = calculate_scaling_factor(
+            reference_image_path="./Reference_ScalingBox.jpg",
+            target_image=reference_obj_img_scaled,
+            feature_detector="SIFT",
+        )
+    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
+    )
+    # Scale the object mask according to scaling factor
+    # objects_mask_scaled = scale_image(objects_mask, scaling_factor)
+    Image.fromarray(objects_mask).save("./outputs/scaled_mask_new.jpg")
+    outlines, contours = extract_outlines(objects_mask)
+    dxf = to_dxf(contours)
+
+    return outlines, dxf, objects_mask, scaling_factor, reference_obj_img_scaled
+
+
+
+if __name__ == "__main__":
+    os.makedirs("./outputs", exist_ok=True)
+
+    ifer = gr.Interface(
+        fn=predict,
+        inputs=[gr.Image(label="Input Image")],
+        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")
+        ],
+        examples=["./examples/Test20.jpg", "./examples/Test21.jpg", "./examples/Test22.jpg", "./examples/Test23.jpg"]
+    )
+    ifer.launch(share=True)
+

last.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8ecf93886616e47bcbd997c9149521eab864aea3c4fa9ff48a95ab23d8ecf51e
+size 6254691

requirements.txt

@@ -0,0 +1,6 @@
+transformers
+ultralytics
+ezdxf
+gradio
+kornia
+timm

scalingtestupdated.py

@@ -0,0 +1,167 @@
+import cv2
+import numpy as np
+import os
+import argparse
+from typing import Union
+from matplotlib import pyplot as plt
+
+
+class ScalingSquareDetector:
+    def __init__(self, feature_detector="ORB", debug=False):
+        """
+        Initialize the detector with the desired feature matching algorithm.
+        :param feature_detector: "ORB" or "SIFT" (default is "ORB").
+        :param debug: If True, saves intermediate images for debugging.
+        """
+        self.feature_detector = feature_detector
+        self.debug = debug
+        self.detector = self._initialize_detector()
+
+    def _initialize_detector(self):
+        """
+        Initialize the chosen feature detector.
+        :return: OpenCV detector object.
+        """
+        if self.feature_detector.upper() == "SIFT":
+            return cv2.SIFT_create()
+        elif self.feature_detector.upper() == "ORB":
+            return cv2.ORB_create()
+        else:
+            raise ValueError("Invalid feature detector. Choose 'ORB' or 'SIFT'.")
+
+    def find_scaling_square(
+        self, reference_image_path, target_image, known_size_mm, roi_margin=30
+    ):
+        """
+        Detect the scaling square in the target image based on the reference image.
+        :param reference_image_path: Path to the reference image of the square.
+        :param target_image_path: Path to the target image containing the square.
+        :param known_size_mm: Physical size of the square in millimeters.
+        :param roi_margin: Margin to expand the ROI around the detected square (in pixels).
+        :return: Scaling factor (mm per pixel).
+        """
+        target_image = cv2.cvtColor(target_image, cv2.COLOR_RGB2GRAY)
+
+        roi = target_image.copy()
+        # Find contours in the ROI
+        roi_blurred = cv2.GaussianBlur(roi, (5, 5), 0)
+        _, roi_binary = cv2.threshold(
+            roi_blurred, 128, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU
+        )
+        contours, _ = cv2.findContours(
+            roi_binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
+        )
+
+        if not contours:
+            raise ValueError("No contours found in the cropped ROI.")
+
+        # Select the largest square-like contour
+        largest_square = None
+        largest_square_area = 0
+        for contour in contours:
+            x_c, y_c, w_c, h_c = cv2.boundingRect(contour)
+            aspect_ratio = w_c / float(h_c)
+            if 0.9 <= aspect_ratio <= 1.1:
+                peri = cv2.arcLength(contour, True)
+                approx = cv2.approxPolyDP(contour, 0.02 * peri, True)
+                if len(approx) == 4:
+                    area = cv2.contourArea(contour)
+                    if area > largest_square_area:
+                        largest_square = contour
+                        largest_square_area = area
+
+        if largest_square is None:
+            raise ValueError("No square-like contour found in the ROI.")
+
+        # Draw the largest contour on the original image
+        target_image_color = cv2.cvtColor(target_image, cv2.COLOR_GRAY2BGR)
+        cv2.drawContours(
+            target_image_color, largest_square, -1, (255, 0, 0), 3
+        )
+
+        if self.debug:
+            cv2.imwrite("largest_contour.jpg", target_image_color)
+
+        # Calculate the bounding rectangle of the largest contour
+        x, y, w, h = cv2.boundingRect(largest_square)
+        square_width_px = w
+        square_height_px = h
+
+        # Calculate the scaling factor
+        avg_square_size_px = (square_width_px + square_height_px) / 2
+        scaling_factor = known_size_mm / avg_square_size_px  # mm per pixel
+
+        return scaling_factor
+
+    def draw_debug_images(self, output_folder):
+        """
+        Save debug images if enabled.
+        :param output_folder: Directory to save debug images.
+        """
+        if self.debug:
+            if not os.path.exists(output_folder):
+                os.makedirs(output_folder)
+            debug_images = ["largest_contour.jpg"]
+            for img_name in debug_images:
+                if os.path.exists(img_name):
+                    os.rename(img_name, os.path.join(output_folder, img_name))
+
+
+def calculate_scaling_factor(
+    reference_image_path,
+    target_image,
+    known_square_size_mm=9.0,
+    feature_detector="ORB",
+    debug=False,
+    roi_margin=30,
+):
+    # Initialize detector
+    detector = ScalingSquareDetector(feature_detector=feature_detector, debug=debug)
+
+    # Find scaling square and calculate scaling factor
+    scaling_factor = detector.find_scaling_square(
+        reference_image_path=reference_image_path,
+        target_image=target_image,
+        known_size_mm=known_square_size_mm,
+        roi_margin=roi_margin,
+    )
+
+    # Save debug images
+    if debug:
+        detector.draw_debug_images("debug_outputs")
+
+    return scaling_factor
+
+
+# Example usage:
+if __name__ == "__main__":
+    import os
+    from PIL import Image
+    from ultralytics import YOLO
+    from app import yolo_detect, shrink_bbox
+    from ultralytics.utils.plotting import save_one_box
+
+    for idx, file in enumerate(os.listdir("./sample_images")):
+        img = np.array(Image.open(os.path.join("./sample_images", file)))
+        img = yolo_detect(img, ['box'])
+        model = YOLO("./runs/detect/train/weights/last.pt")
+        res = model.predict(img, conf=0.6)
+        
+        box_img = save_one_box(res[0].cpu().boxes.xyxy, im=res[0].orig_img, save=False)
+        img = shrink_bbox(box_img, 1.20)
+        cv2.imwrite(f"./outputs/{idx}_{file}", img)
+        try:
+
+            scaling_factor = calculate_scaling_factor(
+                reference_image_path="./Reference_ScalingBox.jpg",
+                target_image=img,
+                known_square_size_mm=9.0,
+                feature_detector="ORB",
+                debug=False,
+                roi_margin=90,
+            )
+            print(f"Scaling Factor (mm per pixel): {scaling_factor:.6f}")
+        except Exception as e:
+            from traceback import print_exc
+            print(print_exc())
+            print(f"Error: {e}")