mlbench123 commited on
Commit
a88b356
·
verified ·
1 Parent(s): 155b415

Upload 8 files

Browse files
Files changed (8) hide show
  1. README.md +5 -7
  2. app.py +438 -0
  3. best.pt +3 -0
  4. coin.png +0 -0
  5. convert.py +68 -0
  6. requirements.txt +7 -0
  7. scalingtestupdated.py +180 -0
  8. yolov8x-worldv2.pt +3 -0
README.md CHANGED
@@ -1,14 +1,12 @@
1
  ---
2
- title: Contours Extraction
3
- emoji: 🔥
4
- colorFrom: yellow
5
- colorTo: blue
6
  sdk: gradio
7
- sdk_version: 5.12.0
8
  app_file: app.py
9
  pinned: false
10
- license: mit
11
- short_description: This model extracts contours of objects
12
  ---
13
 
14
  Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
 
1
  ---
2
+ title: Detect Contours
3
+ emoji: 🐢
4
+ colorFrom: green
5
+ colorTo: yellow
6
  sdk: gradio
7
+ sdk_version: 5.8.0
8
  app_file: app.py
9
  pinned: false
 
 
10
  ---
11
 
12
  Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
app.py ADDED
@@ -0,0 +1,438 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ from pathlib import Path
3
+ from typing import List, Union
4
+ from PIL import Image
5
+ import ezdxf.units
6
+ import numpy as np
7
+ import torch
8
+ from torchvision import transforms
9
+ from ultralytics import YOLOWorld, YOLO
10
+ from ultralytics.engine.results import Results
11
+ from ultralytics.utils.plotting import save_one_box
12
+ from transformers import AutoModelForImageSegmentation
13
+ import cv2
14
+ import ezdxf
15
+ import gradio as gr
16
+ import gc
17
+ from scalingtestupdated import calculate_scaling_factor
18
+ from scipy.interpolate import splprep, splev
19
+ from scipy.ndimage import gaussian_filter1d
20
+
21
+ birefnet = AutoModelForImageSegmentation.from_pretrained(
22
+ "zhengpeng7/BiRefNet", trust_remote_code=True
23
+ )
24
+
25
+ device = "cpu"
26
+ torch.set_float32_matmul_precision(["high", "highest"][0])
27
+
28
+ birefnet.to(device)
29
+ birefnet.eval()
30
+ transform_image = transforms.Compose(
31
+ [
32
+ transforms.Resize((1024, 1024)),
33
+ transforms.ToTensor(),
34
+ transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
35
+ ]
36
+ )
37
+
38
+
39
+ def yolo_detect(
40
+ image: Union[str, Path, int, Image.Image, list, tuple, np.ndarray, torch.Tensor],
41
+ classes: List[str],
42
+ ) -> np.ndarray:
43
+ drawer_detector = YOLOWorld("yolov8x-worldv2.pt")
44
+ drawer_detector.set_classes(classes)
45
+ results: List[Results] = drawer_detector.predict(image)
46
+ boxes = []
47
+ for result in results:
48
+ boxes.append(
49
+ save_one_box(result.cpu().boxes.xyxy, im=result.orig_img, save=False)
50
+ )
51
+
52
+ del drawer_detector
53
+
54
+ return boxes[0]
55
+
56
+
57
+ def remove_bg(image: np.ndarray) -> np.ndarray:
58
+ image = Image.fromarray(image)
59
+ input_images = transform_image(image).unsqueeze(0).to("cpu")
60
+
61
+ # Prediction
62
+ with torch.no_grad():
63
+ preds = birefnet(input_images)[-1].sigmoid().cpu()
64
+ pred = preds[0].squeeze()
65
+
66
+ # Show Results
67
+ pred_pil: Image = transforms.ToPILImage()(pred)
68
+ print(pred_pil)
69
+ # Scale proportionally with max length to 1024 for faster showing
70
+ scale_ratio = 1024 / max(image.size)
71
+ scaled_size = (int(image.size[0] * scale_ratio), int(image.size[1] * scale_ratio))
72
+
73
+ return np.array(pred_pil.resize(scaled_size))
74
+
75
+
76
+ def make_square(img: np.ndarray):
77
+ # Get dimensions
78
+ height, width = img.shape[:2]
79
+
80
+ # Find the larger dimension
81
+ max_dim = max(height, width)
82
+
83
+ # Calculate padding
84
+ pad_height = (max_dim - height) // 2
85
+ pad_width = (max_dim - width) // 2
86
+
87
+ # Handle odd dimensions
88
+ pad_height_extra = max_dim - height - 2 * pad_height
89
+ pad_width_extra = max_dim - width - 2 * pad_width
90
+
91
+ # Create padding with edge colors
92
+ if len(img.shape) == 3: # Color image
93
+ # Pad the image
94
+ padded = np.pad(
95
+ img,
96
+ (
97
+ (pad_height, pad_height + pad_height_extra),
98
+ (pad_width, pad_width + pad_width_extra),
99
+ (0, 0),
100
+ ),
101
+ mode="edge",
102
+ )
103
+ else: # Grayscale image
104
+ padded = np.pad(
105
+ img,
106
+ (
107
+ (pad_height, pad_height + pad_height_extra),
108
+ (pad_width, pad_width + pad_width_extra),
109
+ ),
110
+ mode="edge",
111
+ )
112
+
113
+ return padded
114
+
115
+
116
+ def exclude_scaling_box(
117
+ image: np.ndarray,
118
+ bbox: np.ndarray,
119
+ orig_size: tuple,
120
+ processed_size: tuple,
121
+ expansion_factor: float = 1.5,
122
+ ) -> np.ndarray:
123
+ # Unpack the bounding box
124
+ x_min, y_min, x_max, y_max = map(int, bbox)
125
+
126
+ # Calculate scaling factors
127
+ scale_x = processed_size[1] / orig_size[1] # Width scale
128
+ scale_y = processed_size[0] / orig_size[0] # Height scale
129
+
130
+ # Adjust bounding box coordinates
131
+ x_min = int(x_min * scale_x)
132
+ x_max = int(x_max * scale_x)
133
+ y_min = int(y_min * scale_y)
134
+ y_max = int(y_max * scale_y)
135
+
136
+ # Calculate expanded box coordinates
137
+ box_width = x_max - x_min
138
+ box_height = y_max - y_min
139
+ expanded_x_min = max(0, int(x_min - (expansion_factor - 1) * box_width / 2))
140
+ expanded_x_max = min(
141
+ image.shape[1], int(x_max + (expansion_factor - 1) * box_width / 2)
142
+ )
143
+ expanded_y_min = max(0, int(y_min - (expansion_factor - 1) * box_height / 2))
144
+ expanded_y_max = min(
145
+ image.shape[0], int(y_max + (expansion_factor - 1) * box_height / 2)
146
+ )
147
+
148
+ # Black out the expanded region
149
+ image[expanded_y_min:expanded_y_max, expanded_x_min:expanded_x_max] = 0
150
+
151
+ return image
152
+
153
+
154
+ def resample_contour(contour):
155
+ # ---------------------------------------------------------------------------------------- #
156
+ # Get all the parameters at the start:
157
+ num_points = 1000
158
+ smoothing_factor = 5
159
+
160
+ smoothed_x_sigma = 1
161
+ smoothed_y_sigma = 1
162
+ # ---------------------------------------------------------------------------------------- #
163
+ contour = contour[:, 0, :]
164
+
165
+ tck, _ = splprep([contour[:, 0], contour[:, 1]], s=smoothing_factor)
166
+
167
+ u = np.linspace(0, 1, num_points)
168
+ resampled_points = splev(u, tck)
169
+
170
+ smoothed_x = gaussian_filter1d(resampled_points[0], sigma=smoothed_x_sigma)
171
+ smoothed_y = gaussian_filter1d(resampled_points[1], sigma=smoothed_y_sigma)
172
+
173
+ return np.array([smoothed_x, smoothed_y]).T
174
+
175
+
176
+ def save_dxf_spline(inflated_contours, scaling_factor, height):
177
+ # ---------------------------------------------------------------------------------------- #
178
+ # Get all the parameters at the start:
179
+ degree = 3
180
+ closed = True
181
+ # ---------------------------------------------------------------------------------------- #
182
+
183
+ doc = ezdxf.new(units=0)
184
+ doc.units = ezdxf.units.IN
185
+ doc.header["$INSUNITS"] = ezdxf.units.IN
186
+
187
+ msp = doc.modelspace()
188
+
189
+ for contour in inflated_contours:
190
+ resampled_contour = resample_contour(contour)
191
+ points = [
192
+ (x * scaling_factor, (height - y) * scaling_factor)
193
+ for x, y in resampled_contour
194
+ ]
195
+ if len(points) >= 3:
196
+ # Manually Closing the Contour in case it hasn't been closed by the contours before.
197
+ if np.linalg.norm(np.array(points[0]) - np.array(points[-1])) > 1e-2:
198
+ points.append(points[0])
199
+
200
+ spline = msp.add_spline(points, degree=degree)
201
+ spline.closed = closed
202
+
203
+ # Step 14: Save the DXF file
204
+ dxf_filepath = os.path.join("./outputs", "out.dxf")
205
+ doc.saveas(dxf_filepath)
206
+ return dxf_filepath
207
+
208
+
209
+ def extract_outlines(binary_image: np.ndarray) -> np.ndarray:
210
+ """
211
+ Extracts and draws the outlines of masks from a binary image.
212
+ Args:
213
+ binary_image: Grayscale binary image where white represents masks and black is the background.
214
+ Returns:
215
+ Image with outlines drawn.
216
+ """
217
+ # Detect contours from the binary image
218
+ contours, _ = cv2.findContours(
219
+ binary_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
220
+ )
221
+
222
+ # smooth_contours_list = []
223
+ # for contour in contours:
224
+ # smooth_contours_list.append(smooth_contours(contour))
225
+ # Create a blank image to draw contours
226
+ outline_image = np.zeros_like(binary_image)
227
+
228
+ # Draw the contours on the blank image
229
+ cv2.drawContours(
230
+ outline_image, contours, -1, (255), thickness=1
231
+ ) # White color for outlines
232
+
233
+ return cv2.bitwise_not(outline_image), contours
234
+
235
+
236
+ def shrink_bbox(image: np.ndarray, shrink_factor: float):
237
+ """
238
+ Crops the central 80% of the image, maintaining proportions for non-square images.
239
+ Args:
240
+ image: Input image as a NumPy array.
241
+ Returns:
242
+ Cropped image as a NumPy array.
243
+ """
244
+ height, width = image.shape[:2]
245
+ center_x, center_y = width // 2, height // 2
246
+
247
+ # Calculate 80% dimensions
248
+ new_width = int(width * shrink_factor)
249
+ new_height = int(height * shrink_factor)
250
+
251
+ # Determine the top-left and bottom-right points for cropping
252
+ x1 = max(center_x - new_width // 2, 0)
253
+ y1 = max(center_y - new_height // 2, 0)
254
+ x2 = min(center_x + new_width // 2, width)
255
+ y2 = min(center_y + new_height // 2, height)
256
+
257
+ # Crop the image
258
+ cropped_image = image[y1:y2, x1:x2]
259
+ return cropped_image
260
+
261
+
262
+ # def to_dxf(outlines):
263
+ # upper_range_tuple = (200)
264
+ # lower_range_tuple = (0)
265
+
266
+ # doc = ezdxf.new('R2010')
267
+ # msp = doc.modelspace()
268
+ # masked_jpg = cv2.inRange(outlines,lower_range_tuple, upper_range_tuple)
269
+
270
+ # for i in range(0,masked_jpg.shape[0]):
271
+ # for j in range(0,masked_jpg.shape[1]):
272
+ # if masked_jpg[i][j] == 255:
273
+ # msp.add_line((j,masked_jpg.shape[0] - i), (j,masked_jpg.shape[0] - i))
274
+
275
+ # doc.saveas("./outputs/out.dxf")
276
+ # return "./outputs/out.dxf"
277
+
278
+
279
+ def to_dxf(contours):
280
+ doc = ezdxf.new()
281
+ msp = doc.modelspace()
282
+
283
+ for contour in contours:
284
+ points = [(point[0][0], point[0][1]) for point in contour]
285
+ msp.add_lwpolyline(points, close=True) # Add a polyline for each contour
286
+
287
+ doc.saveas("./outputs/out.dxf")
288
+ return "./outputs/out.dxf"
289
+
290
+
291
+ def smooth_contours(contour):
292
+ epsilon = 0.01 * cv2.arcLength(contour, True) # Adjust factor (e.g., 0.01)
293
+ return cv2.approxPolyDP(contour, epsilon, True)
294
+
295
+
296
+ def scale_image(image: np.ndarray, scale_factor: float) -> np.ndarray:
297
+ """
298
+ Resize image by scaling both width and height by the same factor.
299
+
300
+ Args:
301
+ image: Input numpy image
302
+ scale_factor: Factor to scale the image (e.g., 0.5 for half size, 2 for double size)
303
+
304
+ Returns:
305
+ np.ndarray: Resized image
306
+ """
307
+ if scale_factor <= 0:
308
+ raise ValueError("Scale factor must be positive")
309
+
310
+ current_height, current_width = image.shape[:2]
311
+
312
+ # Calculate new dimensions
313
+ new_width = int(current_width * scale_factor)
314
+ new_height = int(current_height * scale_factor)
315
+
316
+ # Choose interpolation method based on whether we're scaling up or down
317
+ interpolation = cv2.INTER_AREA if scale_factor < 1 else cv2.INTER_CUBIC
318
+
319
+ # Resize image
320
+ resized_image = cv2.resize(
321
+ image, (new_width, new_height), interpolation=interpolation
322
+ )
323
+
324
+ return resized_image
325
+
326
+
327
+ def detect_reference_square(img) -> np.ndarray:
328
+ box_detector = YOLO("./best.pt")
329
+ res = box_detector.predict(img)
330
+ del box_detector
331
+ return save_one_box(res[0].cpu().boxes.xyxy, res[0].orig_img, save=False), res[
332
+ 0
333
+ ].cpu().boxes.xyxy[0]
334
+
335
+
336
+ def resize_img(img: np.ndarray, resize_dim):
337
+ return np.array(Image.fromarray(img).resize(resize_dim))
338
+
339
+
340
+ def predict(image, offset_inches):
341
+ try:
342
+ drawer_img = yolo_detect(image, ["box"])
343
+ shrunked_img = make_square(shrink_bbox(drawer_img, 0.8))
344
+ except:
345
+ raise gr.Error("Unable to DETECT DRAWER, please take another picture with different magnification level!")
346
+
347
+ # Detect the scaling reference square
348
+ try:
349
+ reference_obj_img, scaling_box_coords = detect_reference_square(shrunked_img)
350
+ except:
351
+ raise gr.Error("Unable to DETECT REFERENCE BOX, please take another picture with different magnification level!")
352
+
353
+ # reference_obj_img_scaled = shrink_bbox(reference_obj_img, 1.2)
354
+ # make the image sqaure so it does not effect the size of objects
355
+ reference_obj_img = make_square(reference_obj_img)
356
+ reference_square_mask = remove_bg(reference_obj_img)
357
+
358
+ # make the mask same size as org image
359
+ reference_square_mask = resize_img(
360
+ reference_square_mask, (reference_obj_img.shape[1], reference_obj_img.shape[0])
361
+ )
362
+
363
+ try:
364
+ scaling_factor = calculate_scaling_factor(
365
+ reference_image_path="./coin.png",
366
+ target_image=reference_square_mask,
367
+ feature_detector="ORB",
368
+ )
369
+ except:
370
+ scaling_factor = 1.0
371
+
372
+ # Save original size before `remove_bg` processing
373
+ orig_size = shrunked_img.shape[:2]
374
+ # Generate foreground mask and save its size
375
+ objects_mask = remove_bg(shrunked_img)
376
+
377
+ processed_size = objects_mask.shape[:2]
378
+ # Exclude scaling box region from objects mask
379
+ objects_mask = exclude_scaling_box(
380
+ objects_mask,
381
+ scaling_box_coords,
382
+ orig_size,
383
+ processed_size,
384
+ expansion_factor=3.0,
385
+ )
386
+ objects_mask = resize_img(
387
+ objects_mask, (shrunked_img.shape[1], shrunked_img.shape[0])
388
+ )
389
+ offset_pixels = (offset_inches / scaling_factor) * 2 + 1
390
+ dilated_mask = cv2.dilate(
391
+ objects_mask, np.ones((int(offset_pixels), int(offset_pixels)), np.uint8)
392
+ )
393
+
394
+ # Scale the object mask according to scaling factor
395
+ # objects_mask_scaled = scale_image(objects_mask, scaling_factor)
396
+ Image.fromarray(dilated_mask).save("./outputs/scaled_mask_new.jpg")
397
+ outlines, contours = extract_outlines(dilated_mask)
398
+ shrunked_img_contours = cv2.drawContours(
399
+ shrunked_img, contours, -1, (0, 0, 255), thickness=2
400
+ )
401
+ dxf = save_dxf_spline(contours, scaling_factor, processed_size[0])
402
+
403
+ return (
404
+ cv2.cvtColor(shrunked_img_contours, cv2.COLOR_BGR2RGB),
405
+ outlines,
406
+ dxf,
407
+ dilated_mask,
408
+ scaling_factor,
409
+ )
410
+
411
+
412
+ if __name__ == "__main__":
413
+ os.makedirs("./outputs", exist_ok=True)
414
+
415
+ ifer = gr.Interface(
416
+ fn=predict,
417
+ inputs=[
418
+ gr.Image(label="Input Image"),
419
+ gr.Number(label="Offset value for Mask(inches)", value=0.075),
420
+ ],
421
+ outputs=[
422
+ gr.Image(label="Ouput Image"),
423
+ gr.Image(label="Outlines of Objects"),
424
+ gr.File(label="DXF file"),
425
+ gr.Image(label="Mask"),
426
+ gr.Textbox(
427
+ label="Scaling Factor(mm)",
428
+ placeholder="Every pixel is equal to mentioned number in inches",
429
+ ),
430
+ ],
431
+ examples=[
432
+ ["./examples/Test20.jpg", 0.075],
433
+ ["./examples/Test21.jpg", 0.075],
434
+ ["./examples/Test22.jpg", 0.075],
435
+ ["./examples/Test23.jpg", 0.075],
436
+ ],
437
+ )
438
+ ifer.launch(share=True)
best.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:016663c7243bbaf34fe923ddec534fb32bf558efa7b326f6a3b9adcb581de29c
3
+ size 6209625
coin.png ADDED
convert.py ADDED
@@ -0,0 +1,68 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # import ezdxf
2
+
3
+ # # Load the DXF file
4
+ # doc = ezdxf.readfile("out.dxf")
5
+
6
+ # # Iterate through all entities in the modelspace
7
+ # for entity in doc.modelspace():
8
+ # entity_type = entity.dxftype() # Get the entity type
9
+ # print(f"Entity Type: {entity_type}")
10
+
11
+ # # Handle different entity types
12
+ # if entity_type == "LINE":
13
+ # print(f"Start: {entity.dxf.start}, End: {entity.dxf.end}")
14
+ # elif entity_type == "CIRCLE":
15
+ # print(f"Center: {entity.dxf.center}, Radius: {entity.dxf.radius}")
16
+ # elif entity_type == "ARC":
17
+ # print(f"Center: {entity.dxf.center}, Radius: {entity.dxf.radius}, Start Angle: {entity.dxf.start_angle}, End Angle: {entity.dxf.end_angle}")
18
+ # elif entity_type == "SPLINE":
19
+ # if entity.control_points:
20
+ # print(f"Control Points: {entity.control_points}")
21
+ # elif entity.fit_points:
22
+ # print(f"Fit Points: {entity.fit_points}")
23
+ # elif entity.knots:
24
+ # print(f"Knots: {entity.knots}")
25
+ # else:
26
+ # print("No control, fit, or knot points found for this SPLINE.")
27
+ # else:
28
+ # print(f"No specific handler for entity type: {entity_type}")
29
+
30
+ import numpy as np
31
+ import ezdxf
32
+
33
+ # Load the DXF file
34
+ doc = ezdxf.readfile("out.dxf")
35
+
36
+ def calculate_distance(p1, p2):
37
+ """Calculate the distance between two points."""
38
+ return np.linalg.norm(np.array(p1) - np.array(p2))
39
+
40
+ def process_fit_points(fit_points):
41
+ """Process fit points to calculate distances and bounding box."""
42
+ distances = []
43
+ for i in range(len(fit_points) - 1):
44
+ distances.append(calculate_distance(fit_points[i], fit_points[i + 1]))
45
+
46
+ # Calculate perimeter
47
+ perimeter = sum(distances)
48
+
49
+ # Calculate bounding box
50
+ fit_points_np = np.array(fit_points)
51
+ min_x, min_y = np.min(fit_points_np[:, :2], axis=0)
52
+ max_x, max_y = np.max(fit_points_np[:, :2], axis=0)
53
+
54
+ return {
55
+ "distances": distances,
56
+ "perimeter": perimeter,
57
+ "bounding_box": (min_x, min_y, max_x, max_y)
58
+ }
59
+
60
+ # Iterate through all entities in the modelspace
61
+ for entity in doc.modelspace():
62
+ if entity.dxftype() == "SPLINE" and entity.fit_points:
63
+ print(f"Entity Type: SPLINE")
64
+ fit_points = entity.fit_points
65
+ results = process_fit_points(fit_points)
66
+ print(f"Perimeter: {results['perimeter']}")
67
+ print(f"Bounding Box: {results['bounding_box']}")
68
+ print(f"Distances: {results['distances'][:]}... (showing first 5)")
requirements.txt ADDED
@@ -0,0 +1,7 @@
 
 
 
 
 
 
 
 
1
+ transformers
2
+ ultralytics==8.3.9
3
+ ezdxf
4
+ gradio
5
+ kornia
6
+ timm
7
+ einops
scalingtestupdated.py ADDED
@@ -0,0 +1,180 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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, reference_image_path, 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
+
44
+ contours, _ = cv2.findContours(
45
+ target_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
46
+ )
47
+
48
+ if not contours:
49
+ raise ValueError("No contours found in the cropped ROI.")
50
+
51
+ # # Select the largest square-like contour
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
+ # 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 image
70
+ target_image_color = cv2.cvtColor(target_image, cv2.COLOR_GRAY2BGR)
71
+ cv2.drawContours(
72
+ target_image_color, largest_square, -1, (255, 0, 0), 3
73
+ )
74
+
75
+ # if self.debug:
76
+ cv2.imwrite("largest_contour.jpg", target_image_color)
77
+
78
+ # Calculate the bounding rectangle of the largest contour
79
+ x, y, w, h = cv2.boundingRect(largest_square)
80
+ square_width_px = w
81
+ square_height_px = h
82
+
83
+ # Calculate the scaling factor
84
+ avg_square_size_px = (square_width_px + square_height_px) / 2
85
+ scaling_factor = 0.5 / avg_square_size_px # mm per pixel
86
+
87
+ return scaling_factor #, square_height_px, square_width_px, roi_binary
88
+
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
+
103
+ def calculate_scaling_factor(
104
+ reference_image_path,
105
+ target_image,
106
+ known_square_size_mm=22,
107
+ feature_detector="ORB",
108
+ debug=False,
109
+ roi_margin=30,
110
+ ):
111
+ # Initialize detector
112
+ detector = ScalingSquareDetector(feature_detector=feature_detector, debug=debug)
113
+
114
+ # Find scaling square and calculate scaling factor
115
+ 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 images
123
+ if debug:
124
+ detector.draw_debug_images("debug_outputs")
125
+
126
+ return scaling_factor
127
+
128
+
129
+ # Example usage:
130
+ if __name__ == "__main__":
131
+ import os
132
+ from PIL import Image
133
+ from ultralytics import YOLO
134
+ from app import yolo_detect, shrink_bbox
135
+ from ultralytics.utils.plotting import save_one_box
136
+
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("./best.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="./coin.png",
152
+ target_image=box_img,
153
+ known_square_size_mm=22,
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 mm
161
+ # square_size_mm = 12.7
162
+
163
+ # # Compute the calculated scaling factors and compare
164
+ # calculated_scaling_factor = square_size_mm / height_px
165
+ # discrepancy = abs(calculated_scaling_factor - scaling_factor)
166
+ # import pprint
167
+ # 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_exc
179
+ print(print_exc())
180
+ print(f"Error: {e}")
yolov8x-worldv2.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:41e771bfbbb8894dd857f3fef7cac3b3578dffd49fd3547101efa6a606a02a0e
3
+ size 146355704