Spaces:
Running
Running
Upload 8 files
Browse files- README.md +5 -7
- app.py +438 -0
- best.pt +3 -0
- coin.png +0 -0
- convert.py +68 -0
- requirements.txt +7 -0
- scalingtestupdated.py +180 -0
- yolov8x-worldv2.pt +3 -0
README.md
CHANGED
@@ -1,14 +1,12 @@
|
|
1 |
---
|
2 |
-
title: Contours
|
3 |
-
emoji:
|
4 |
-
colorFrom:
|
5 |
-
colorTo:
|
6 |
sdk: gradio
|
7 |
-
sdk_version: 5.
|
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
|