scalingtestupdated.py

import cv2
import numpy as np
import os
import argparse
from typing import Union, Tuple
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, target_image, known_size_mm, roi_margin=30
    ):
        """
        Detect the scaling square in the target image based on the reference image.
        :param target_image: Binary 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).
        """
        contours, _ = cv2.findContours(
            target_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
        )

        if not contours:
            raise ValueError("No contours found in the target image.")

        # Select the largest contour (assuming it's the reference object)
        print(f"No of contours: {len(contours)}")
        largest_contour = max(contours, key=cv2.contourArea)

        # Draw the largest contour on the original image for debugging
        target_image_color = cv2.cvtColor(target_image, cv2.COLOR_GRAY2BGR)
        cv2.drawContours(
            target_image_color, [largest_contour], -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_contour)
        square_width_px = w
        square_height_px = h
        
        print(f"Reference object size: {known_size_mm} mm")
        print(f"Detected width: {square_width_px} px")
        print(f"Detected height: {square_height_px} px")

        # Calculate the scaling factor using average of width and height
        avg_square_size_px = (square_width_px + square_height_px) / 2
        print(f"Average square size: {avg_square_size_px} px")
        
        scaling_factor = known_size_mm / avg_square_size_px  # mm per pixel
        print(f"Calculated scaling factor: {scaling_factor:.6f} 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(
    target_image,
    reference_obj_size_mm,
    feature_detector="ORB",
    debug=False,
    roi_margin=30,
) -> float:
    """
    Calculate scaling factor from reference object in image.
    
    :param target_image: Input image (numpy array)
    :param reference_obj_size_mm: Known size of reference object in millimeters
    :param feature_detector: Feature detector to use ("ORB" or "SIFT")
    :param debug: Enable debug output
    :param roi_margin: ROI margin in pixels
    :return: Scaling factor in mm per pixel
    """
    # Initialize detector
    detector = ScalingSquareDetector(feature_detector=feature_detector, debug=debug)

    # Find scaling square and calculate scaling factor
    scaling_factor = detector.find_scaling_square(
        target_image=target_image,
        known_size_mm=reference_obj_size_mm,
        roi_margin=roi_margin,
    )

    # Save debug images
    if debug:
        detector.draw_debug_images("debug_outputs")

    return scaling_factor


def convert_units(value: float, from_unit: str, to_unit: str) -> float:
    """
    Convert between mm and inches.
    
    :param value: Value to convert
    :param from_unit: Source unit ("mm" or "inches")
    :param to_unit: Target unit ("mm" or "inches")
    :return: Converted value
    """
    if from_unit == to_unit:
        return value
    
    if from_unit == "inches" and to_unit == "mm":
        return value * 25.4
    elif from_unit == "mm" and to_unit == "inches":
        return value / 25.4
    else:
        raise ValueError(f"Unsupported unit conversion: {from_unit} to {to_unit}")


def calculate_scaling_factor_with_units(
    target_image,
    reference_obj_size: float,
    reference_unit: str = "mm",
    output_unit: str = "mm",
    feature_detector="ORB",
    debug=False,
    roi_margin=30,
) -> Tuple[float, str]:
    """
    Calculate scaling factor with proper unit handling.
    
    :param target_image: Input image (numpy array)
    :param reference_obj_size: Known size of reference object
    :param reference_unit: Unit of reference object size ("mm" or "inches")
    :param output_unit: Desired unit for scaling factor ("mm" or "inches")
    :param feature_detector: Feature detector to use ("ORB" or "SIFT")
    :param debug: Enable debug output
    :param roi_margin: ROI margin in pixels
    :return: Tuple of (scaling_factor, unit_string)
    """
    # Convert reference size to mm for internal calculation
    reference_size_mm = convert_units(reference_obj_size, reference_unit, "mm")
    
    # Calculate scaling factor in mm/px
    scaling_factor_mm = calculate_scaling_factor(
        target_image=target_image,
        reference_obj_size_mm=reference_size_mm,
        feature_detector=feature_detector,
        debug=debug,
        roi_margin=roi_margin,
    )
    
    # Convert scaling factor to desired output unit
    if output_unit == "inches":
        scaling_factor = scaling_factor_mm / 25.4  # Convert mm/px to inches/px
        unit_string = "inches per pixel"
    else:
        scaling_factor = scaling_factor_mm
        unit_string = "mm per pixel"
    
    print(f"Final scaling factor: {scaling_factor:.6f} {unit_string}")
    
    return scaling_factor, unit_string


# Paper size configurations (in mm and inches)
PAPER_SIZES = {
    "A4": {"width_mm": 210, "height_mm": 297, "width_inches": 8.27, "height_inches": 11.69},
    "A3": {"width_mm": 297, "height_mm": 420, "width_inches": 11.69, "height_inches": 16.54},
    "US Letter": {"width_mm": 215.9, "height_mm": 279.4, "width_inches": 8.5, "height_inches": 11.0}
}


def calculate_paper_scaling_factor(
    paper_contour: np.ndarray, 
    paper_size: str, 
    output_unit: str = "mm"
) -> Tuple[float, str]:
    """
    Calculate scaling factor based on detected paper dimensions with proper unit handling.
    
    :param paper_contour: Detected paper contour
    :param paper_size: Paper size identifier ("A4", "A3", "US Letter")
    :param output_unit: Desired unit for scaling factor ("mm" or "inches")
    :return: Tuple of (scaling_factor, unit_string)
    """
    # Get paper dimensions in the desired unit
    if output_unit == "inches":
        expected_width = PAPER_SIZES[paper_size]["width_inches"]
        expected_height = PAPER_SIZES[paper_size]["height_inches"]
        unit_string = "inches per pixel"
    else:
        expected_width = PAPER_SIZES[paper_size]["width_mm"]
        expected_height = PAPER_SIZES[paper_size]["height_mm"]
        unit_string = "mm per pixel"
    
    # Calculate bounding rectangle of paper contour
    rect = cv2.boundingRect(paper_contour)
    detected_width_px = rect[2]
    detected_height_px = rect[3]
    
    # Calculate scaling factors for both dimensions
    scale_x = expected_width / detected_width_px
    scale_y = expected_height / detected_height_px
    
    # Use the minimum scale to ensure the object fits within expected dimensions
    scaling_factor = min(scale_x, scale_y)
    
    print(f"Paper detection: {detected_width_px}x{detected_height_px} px")
    print(f"Expected paper size: {expected_width}x{expected_height} {output_unit}")
    print(f"Scale X: {scale_x:.6f}, Scale Y: {scale_y:.6f}")
    print(f"Final scaling factor: {scaling_factor:.6f} {unit_string}")
    
    return scaling_factor, unit_string


# Example usage:
if __name__ == "__main__":
    import os
    from PIL import Image
    
    # Test with different units
    sample_dir = "./sample_images"
    if os.path.exists(sample_dir):
        for idx, file in enumerate(os.listdir(sample_dir)):
            if file.lower().endswith(('.jpg', '.jpeg', '.png')):
                img_path = os.path.join(sample_dir, file)
                img = np.array(Image.open(img_path))
                
                # Convert to grayscale if needed
                if len(img.shape) == 3:
                    img_gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
                else:
                    img_gray = img
                
                print(f"\nProcessing: {file}")
                
                try:
                    # Test with mm units
                    scaling_factor_mm, unit_mm = calculate_scaling_factor_with_units(
                        target_image=img_gray,
                        reference_obj_size=20.0,  # 20mm reference object
                        reference_unit="mm",
                        output_unit="mm",
                        feature_detector="ORB",
                        debug=False,
                        roi_margin=90,
                    )
                    
                    # Test with inch units
                    scaling_factor_inch, unit_inch = calculate_scaling_factor_with_units(
                        target_image=img_gray,
                        reference_obj_size=0.787,  # ~20mm in inches
                        reference_unit="inches",
                        output_unit="inches",
                        feature_detector="ORB",
                        debug=False,
                        roi_margin=90,
                    )
                    
                    print(f"MM scaling: {scaling_factor_mm:.6f} {unit_mm}")
                    print(f"Inch scaling: {scaling_factor_inch:.6f} {unit_inch}")
                    
                    # Verify conversion consistency
                    converted_mm_to_inch = scaling_factor_mm / 25.4
                    print(f"Converted mm to inch: {converted_mm_to_inch:.6f}")
                    print(f"Difference: {abs(scaling_factor_inch - converted_mm_to_inch):.8f}")
                    
                except Exception as e:
                    print(f"Error processing {file}: {e}")
    else:
        print(f"Sample directory {sample_dir} not found")