mlbench123 commited on
Commit
4441cfe
·
verified ·
1 Parent(s): ff93bb4

Delete scalingtestupdated.py

Browse files
Files changed (1) hide show
  1. scalingtestupdated.py +0 -185
scalingtestupdated.py DELETED
@@ -1,185 +0,0 @@
1
- import cv2
2
- import numpy as np
3
- import os
4
- import argparse
5
- from typing import Union
6
- from matplotlib import pyplot as plt
7
-
8
-
9
- class 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_detector
17
- self.debug = debug
18
- 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, target_image, known_size_mm, roi_margin=30
34
- ):
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
- contours, _ = cv2.findContours(
44
- target_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
45
- )
46
-
47
- if not contours:
48
- raise ValueError("No contours found in the cropped ROI.")
49
-
50
- # # Select the largest square-like contour
51
- print(f"No of contours: {len(contours)}")
52
- largest_square = None
53
- # largest_square_area = 0
54
- # 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 = contour
64
- # largest_square_area = area
65
-
66
- for contour in contours:
67
- largest_square = contour
68
-
69
- # if largest_square is None:
70
- # raise ValueError("No square-like contour found in the ROI.")
71
-
72
- # Draw the largest contour on the original image
73
- target_image_color = cv2.cvtColor(target_image, cv2.COLOR_GRAY2BGR)
74
- cv2.drawContours(
75
- target_image_color, largest_square, -1, (255, 0, 0), 3
76
- )
77
-
78
- # if self.debug:
79
- cv2.imwrite("largest_contour.jpg", target_image_color)
80
-
81
- # Calculate the bounding rectangle of the largest contour
82
- x, y, w, h = cv2.boundingRect(largest_square)
83
- square_width_px = w
84
- square_height_px = h
85
- print(f"Reference object size: {known_size_mm} mm")
86
- print(f"width: {square_width_px} px")
87
- print(f"height: {square_height_px} px")
88
-
89
- # Calculate the scaling factor
90
- avg_square_size_px = (square_width_px + square_height_px) / 2
91
- print(f"avg square size: {avg_square_size_px} px")
92
- scaling_factor = known_size_mm / avg_square_size_px # mm per pixel
93
- print(f"scaling factor: {scaling_factor} mm per pixel")
94
-
95
- return scaling_factor #, square_height_px, square_width_px, roi_binary
96
-
97
- def draw_debug_images(self, output_folder):
98
- """
99
- Save debug images if enabled.
100
- :param output_folder: Directory to save debug images.
101
- """
102
- if self.debug:
103
- if not os.path.exists(output_folder):
104
- os.makedirs(output_folder)
105
- debug_images = ["largest_contour.jpg"]
106
- for img_name in debug_images:
107
- if os.path.exists(img_name):
108
- os.rename(img_name, os.path.join(output_folder, img_name))
109
-
110
-
111
- def calculate_scaling_factor(
112
- target_image,
113
- reference_obj_size_mm,
114
- feature_detector="ORB",
115
- debug=False,
116
- roi_margin=30,
117
- ):
118
- # Initialize detector
119
- detector = ScalingSquareDetector(feature_detector=feature_detector, debug=debug)
120
-
121
- # Find scaling square and calculate scaling factor
122
- scaling_factor = detector.find_scaling_square(
123
- target_image=target_image,
124
- known_size_mm=reference_obj_size_mm,
125
- roi_margin=roi_margin,
126
- )
127
-
128
- # Save debug images
129
- if debug:
130
- detector.draw_debug_images("debug_outputs")
131
-
132
- return scaling_factor
133
-
134
-
135
- # Example usage:
136
- if __name__ == "__main__":
137
- import os
138
- from PIL import Image
139
- from ultralytics import YOLO
140
- from app import yolo_detect, shrink_bbox
141
- from ultralytics.utils.plotting import save_one_box
142
-
143
- for idx, file in enumerate(os.listdir("./sample_images")):
144
- img = np.array(Image.open(os.path.join("./sample_images", file)))
145
- img = yolo_detect(img, ['box'])
146
- model = YOLO("./best.pt")
147
- res = model.predict(img, conf=0.6)
148
-
149
- box_img = save_one_box(res[0].cpu().boxes.xyxy, im=res[0].orig_img, save=False)
150
- # img = shrink_bbox(box_img, 1.20)
151
- cv2.imwrite(f"./outputs/{idx}_{file}", box_img)
152
-
153
- print("File: ",f"./outputs/{idx}_{file}")
154
- try:
155
-
156
- scaling_factor = calculate_scaling_factor(
157
- target_image=box_img,
158
- known_square_size_mm=20,
159
- feature_detector="ORB",
160
- debug=False,
161
- roi_margin=90,
162
- )
163
- # cv2.imwrite(f"./outputs/{idx}_binary_{file}", roi_binary)
164
-
165
- # Square size in mm
166
- # square_size_mm = 12.7
167
-
168
- # # Compute the calculated scaling factors and compare
169
- # calculated_scaling_factor = square_size_mm / height_px
170
- # discrepancy = abs(calculated_scaling_factor - scaling_factor)
171
- # import pprint
172
- # pprint.pprint({
173
- # "height_px": height_px,
174
- # "width_px": width_px,
175
- # "given_scaling_factor": scaling_factor,
176
- # "calculated_scaling_factor": calculated_scaling_factor,
177
- # "discrepancy": discrepancy,
178
- # })
179
-
180
-
181
- print(f"Scaling Factor (mm per pixel): {scaling_factor:.6f}")
182
- except Exception as e:
183
- from traceback import print_exc
184
- print(print_exc())
185
- print(f"Error: {e}")