MLBench/Drawer_Detection

1import cv22import numpy as np3import os4import argparse5from typing import Union6from matplotlib import pyplot as plt7 8 9class ScalingSquareDetector:10    def __init__(self, feature_detector="ORB", debug=False):11        """12        Initialize the detector with the desired feature matching algorithm.13        :param feature_detector: "ORB" or "SIFT" (default is "ORB").14        :param debug: If True, saves intermediate images for debugging.15        """16        self.feature_detector = feature_detector17        self.debug = debug18        self.detector = self._initialize_detector()19 20    def _initialize_detector(self):21        """22        Initialize the chosen feature detector.23        :return: OpenCV detector object.24        """25        if self.feature_detector.upper() == "SIFT":26            return cv2.SIFT_create()27        elif self.feature_detector.upper() == "ORB":28            return cv2.ORB_create()29        else:30            raise ValueError("Invalid feature detector. Choose 'ORB' or 'SIFT'.")31 32    def find_scaling_square(33        self, reference_image_path, target_image, known_size_mm, roi_margin=3034    ):35        """36        Detect the scaling square in the target image based on the reference image.37        :param reference_image_path: Path to the reference image of the square.38        :param target_image_path: Path to the target image containing the square.39        :param known_size_mm: Physical size of the square in millimeters.40        :param roi_margin: Margin to expand the ROI around the detected square (in pixels).41        :return: Scaling factor (mm per pixel).42        """43        44        contours, _ = cv2.findContours(45            target_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE46        )47 48        if not contours:49            raise ValueError("No contours found in the cropped ROI.")50 51        # # Select the largest square-like contour52        largest_square = None53        largest_square_area = 054        for contour in contours:55            x_c, y_c, w_c, h_c = cv2.boundingRect(contour)56            aspect_ratio = w_c / float(h_c)57            if 0.9 <= aspect_ratio <= 1.1:58                peri = cv2.arcLength(contour, True)59                approx = cv2.approxPolyDP(contour, 0.02 * peri, True)60                if len(approx) == 4:61                    area = cv2.contourArea(contour)62                    if area > largest_square_area:63                        largest_square = contour64                        largest_square_area = area65 66        # if largest_square is None:67        #     raise ValueError("No square-like contour found in the ROI.")68 69        # Draw the largest contour on the original image70        target_image_color = cv2.cvtColor(target_image, cv2.COLOR_GRAY2BGR)71        cv2.drawContours(72            target_image_color, largest_square, -1, (255, 0, 0), 373        )74 75        # if self.debug:76        cv2.imwrite("largest_contour.jpg", target_image_color)77 78        # Calculate the bounding rectangle of the largest contour79        x, y, w, h = cv2.boundingRect(largest_square)80        square_width_px = w81        square_height_px = h82 83        # Calculate the scaling factor84        avg_square_size_px = (square_width_px + square_height_px) / 285        scaling_factor = 0.5 / avg_square_size_px  # mm per pixel86 87        return scaling_factor #, square_height_px, square_width_px, roi_binary88 89    def draw_debug_images(self, output_folder):90        """91        Save debug images if enabled.92        :param output_folder: Directory to save debug images.93        """94        if self.debug:95            if not os.path.exists(output_folder):96                os.makedirs(output_folder)97            debug_images = ["largest_contour.jpg"]98            for img_name in debug_images:99                if os.path.exists(img_name):100                    os.rename(img_name, os.path.join(output_folder, img_name))101 102 103def calculate_scaling_factor(104    reference_image_path,105    target_image,106    known_square_size_mm=12.7,107    feature_detector="ORB",108    debug=False,109    roi_margin=30,110):111    # Initialize detector112    detector = ScalingSquareDetector(feature_detector=feature_detector, debug=debug)113 114    # Find scaling square and calculate scaling factor115    scaling_factor = detector.find_scaling_square(116        reference_image_path=reference_image_path,117        target_image=target_image,118        known_size_mm=known_square_size_mm,119        roi_margin=roi_margin,120    )121 122    # Save debug images123    if debug:124        detector.draw_debug_images("debug_outputs")125 126    return scaling_factor127 128 129# Example usage:130if __name__ == "__main__":131    import os132    from PIL import Image133    from ultralytics import YOLO134    from app import yolo_detect, shrink_bbox135    from ultralytics.utils.plotting import save_one_box136 137    for idx, file in enumerate(os.listdir("./sample_images")):138        img = np.array(Image.open(os.path.join("./sample_images", file)))139        img = yolo_detect(img, ['box'])140        model = YOLO("./last.pt")141        res = model.predict(img, conf=0.6)142        143        box_img = save_one_box(res[0].cpu().boxes.xyxy, im=res[0].orig_img, save=False)144        # img = shrink_bbox(box_img, 1.20)145        cv2.imwrite(f"./outputs/{idx}_{file}", box_img)146        147        print("File: ",f"./outputs/{idx}_{file}")148        try:149 150            scaling_factor = calculate_scaling_factor(151                reference_image_path="./Reference_ScalingBox.jpg",152                target_image=box_img,153                known_square_size_mm=12.7,154                feature_detector="ORB",155                debug=False,156                roi_margin=90,157            )158            # cv2.imwrite(f"./outputs/{idx}_binary_{file}", roi_binary)159            160            # Square size in mm161            # square_size_mm = 12.7162 163            # # Compute the calculated scaling factors and compare164            # calculated_scaling_factor = square_size_mm / height_px165            # discrepancy = abs(calculated_scaling_factor - scaling_factor)166            # import pprint167            # pprint.pprint({168            #     "height_px": height_px,169            #     "width_px": width_px,170            #     "given_scaling_factor": scaling_factor,171            #     "calculated_scaling_factor": calculated_scaling_factor,172            #     "discrepancy": discrepancy,173            # })174 175 176            print(f"Scaling Factor (mm per pixel): {scaling_factor:.6f}")177        except Exception as e:178            from traceback import print_exc179            print(print_exc())180            print(f"Error: {e}")181