Spaces:
Running
Running
File size: 6,514 Bytes
3a66ca9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 |
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).
"""
contours, _ = cv2.findContours(
target_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
)
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 = 0.5 / avg_square_size_px # mm per pixel
return scaling_factor #, square_height_px, square_width_px, roi_binary
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=12.7,
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("./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}", box_img)
print("File: ",f"./outputs/{idx}_{file}")
try:
scaling_factor = calculate_scaling_factor(
reference_image_path="./Reference_ScalingBox.jpg",
target_image=box_img,
known_square_size_mm=12.7,
feature_detector="ORB",
debug=False,
roi_margin=90,
)
# cv2.imwrite(f"./outputs/{idx}_binary_{file}", roi_binary)
# Square size in mm
# square_size_mm = 12.7
# # Compute the calculated scaling factors and compare
# calculated_scaling_factor = square_size_mm / height_px
# discrepancy = abs(calculated_scaling_factor - scaling_factor)
# import pprint
# pprint.pprint({
# "height_px": height_px,
# "width_px": width_px,
# "given_scaling_factor": scaling_factor,
# "calculated_scaling_factor": calculated_scaling_factor,
# "discrepancy": discrepancy,
# })
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}")
|