Spaces:
Running
Running
Upload 2 files
Browse files
app.py
ADDED
@@ -0,0 +1,356 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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 |
+
|
22 |
+
birefnet = AutoModelForImageSegmentation.from_pretrained(
|
23 |
+
"zhengpeng7/BiRefNet", trust_remote_code=True
|
24 |
+
)
|
25 |
+
|
26 |
+
device = "cpu"
|
27 |
+
torch.set_float32_matmul_precision(["high", "highest"][0])
|
28 |
+
|
29 |
+
birefnet.to(device)
|
30 |
+
birefnet.eval()
|
31 |
+
transform_image = transforms.Compose(
|
32 |
+
[
|
33 |
+
transforms.Resize((1024, 1024)),
|
34 |
+
transforms.ToTensor(),
|
35 |
+
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
|
36 |
+
]
|
37 |
+
)
|
38 |
+
|
39 |
+
def remove_bg(image: np.ndarray) -> np.ndarray:
|
40 |
+
|
41 |
+
image = Image.fromarray(image)
|
42 |
+
input_images = transform_image(image).unsqueeze(0).to("cpu")
|
43 |
+
|
44 |
+
# Prediction
|
45 |
+
with torch.no_grad():
|
46 |
+
preds = birefnet(input_images)[-1].sigmoid().cpu()
|
47 |
+
pred = preds[0].squeeze()
|
48 |
+
|
49 |
+
# Show Results
|
50 |
+
pred_pil: Image = transforms.ToPILImage()(pred)
|
51 |
+
print(pred_pil)
|
52 |
+
# Scale proportionally with max length to 1024 for faster showing
|
53 |
+
scale_ratio = 1024 / max(image.size)
|
54 |
+
scaled_size = (int(image.size[0] * scale_ratio), int(image.size[1] * scale_ratio))
|
55 |
+
|
56 |
+
return np.array(pred_pil.resize(scaled_size))
|
57 |
+
|
58 |
+
def make_square(img: np.ndarray):
|
59 |
+
# Get dimensions
|
60 |
+
height, width = img.shape[:2]
|
61 |
+
|
62 |
+
# Find the larger dimension
|
63 |
+
max_dim = max(height, width)
|
64 |
+
|
65 |
+
# Calculate padding
|
66 |
+
pad_height = (max_dim - height) // 2
|
67 |
+
pad_width = (max_dim - width) // 2
|
68 |
+
|
69 |
+
# Handle odd dimensions
|
70 |
+
pad_height_extra = max_dim - height - 2 * pad_height
|
71 |
+
pad_width_extra = max_dim - width - 2 * pad_width
|
72 |
+
|
73 |
+
# Create padding with edge colors
|
74 |
+
if len(img.shape) == 3: # Color image
|
75 |
+
# Pad the image
|
76 |
+
padded = np.pad(
|
77 |
+
img,
|
78 |
+
(
|
79 |
+
(pad_height, pad_height + pad_height_extra),
|
80 |
+
(pad_width, pad_width + pad_width_extra),
|
81 |
+
(0, 0),
|
82 |
+
),
|
83 |
+
mode="edge",
|
84 |
+
)
|
85 |
+
else: # Grayscale image
|
86 |
+
padded = np.pad(
|
87 |
+
img,
|
88 |
+
(
|
89 |
+
(pad_height, pad_height + pad_height_extra),
|
90 |
+
(pad_width, pad_width + pad_width_extra),
|
91 |
+
),
|
92 |
+
mode="edge",
|
93 |
+
)
|
94 |
+
|
95 |
+
return padded
|
96 |
+
|
97 |
+
def exclude_scaling_box(
|
98 |
+
image: np.ndarray,
|
99 |
+
bbox: np.ndarray,
|
100 |
+
orig_size: tuple,
|
101 |
+
processed_size: tuple,
|
102 |
+
expansion_factor: float = 1.5,
|
103 |
+
) -> np.ndarray:
|
104 |
+
# Unpack the bounding box
|
105 |
+
x_min, y_min, x_max, y_max = map(int, bbox)
|
106 |
+
|
107 |
+
# Calculate scaling factors
|
108 |
+
scale_x = processed_size[1] / orig_size[1] # Width scale
|
109 |
+
scale_y = processed_size[0] / orig_size[0] # Height scale
|
110 |
+
|
111 |
+
# Adjust bounding box coordinates
|
112 |
+
x_min = int(x_min * scale_x)
|
113 |
+
x_max = int(x_max * scale_x)
|
114 |
+
y_min = int(y_min * scale_y)
|
115 |
+
y_max = int(y_max * scale_y)
|
116 |
+
|
117 |
+
# Calculate expanded box coordinates
|
118 |
+
box_width = x_max - x_min
|
119 |
+
box_height = y_max - y_min
|
120 |
+
expanded_x_min = max(0, int(x_min - (expansion_factor - 1) * box_width / 2))
|
121 |
+
expanded_x_max = min(
|
122 |
+
image.shape[1], int(x_max + (expansion_factor - 1) * box_width / 2)
|
123 |
+
)
|
124 |
+
expanded_y_min = max(0, int(y_min - (expansion_factor - 1) * box_height / 2))
|
125 |
+
expanded_y_max = min(
|
126 |
+
image.shape[0], int(y_max + (expansion_factor - 1) * box_height / 2)
|
127 |
+
)
|
128 |
+
|
129 |
+
# Black out the expanded region
|
130 |
+
image[expanded_y_min:expanded_y_max, expanded_x_min:expanded_x_max] = 0
|
131 |
+
|
132 |
+
return image
|
133 |
+
|
134 |
+
def resample_contour(contour):
|
135 |
+
# Get all the parameters at the start:
|
136 |
+
num_points = 1000
|
137 |
+
smoothing_factor = 5
|
138 |
+
spline_degree = 3 # Typically k=3 for cubic spline
|
139 |
+
|
140 |
+
smoothed_x_sigma = 1
|
141 |
+
smoothed_y_sigma = 1
|
142 |
+
|
143 |
+
# Ensure contour has enough points
|
144 |
+
if len(contour) < spline_degree + 1:
|
145 |
+
raise ValueError(f"Contour must have at least {spline_degree + 1} points, but has {len(contour)} points.")
|
146 |
+
|
147 |
+
contour = contour[:, 0, :]
|
148 |
+
|
149 |
+
tck, _ = splprep([contour[:, 0], contour[:, 1]], s=smoothing_factor)
|
150 |
+
u = np.linspace(0, 1, num_points)
|
151 |
+
resampled_points = splev(u, tck)
|
152 |
+
|
153 |
+
smoothed_x = gaussian_filter1d(resampled_points[0], sigma=smoothed_x_sigma)
|
154 |
+
smoothed_y = gaussian_filter1d(resampled_points[1], sigma=smoothed_y_sigma)
|
155 |
+
|
156 |
+
return np.array([smoothed_x, smoothed_y]).T
|
157 |
+
|
158 |
+
def save_dxf_spline(inflated_contours, scaling_factor, height):
|
159 |
+
degree = 3
|
160 |
+
closed = True
|
161 |
+
|
162 |
+
doc = ezdxf.new(units=0)
|
163 |
+
doc.units = ezdxf.units.IN
|
164 |
+
doc.header["$INSUNITS"] = ezdxf.units.IN
|
165 |
+
|
166 |
+
msp = doc.modelspace()
|
167 |
+
|
168 |
+
for contour in inflated_contours:
|
169 |
+
try:
|
170 |
+
resampled_contour = resample_contour(contour)
|
171 |
+
points = [
|
172 |
+
(x * scaling_factor, (height - y) * scaling_factor)
|
173 |
+
for x, y in resampled_contour
|
174 |
+
]
|
175 |
+
if len(points) >= 3:
|
176 |
+
if np.linalg.norm(np.array(points[0]) - np.array(points[-1])) > 1e-2:
|
177 |
+
points.append(points[0])
|
178 |
+
|
179 |
+
spline = msp.add_spline(points, degree=degree)
|
180 |
+
spline.closed = closed
|
181 |
+
except ValueError as e:
|
182 |
+
print(f"Skipping contour: {e}")
|
183 |
+
|
184 |
+
dxf_filepath = os.path.join("./outputs", "out.dxf")
|
185 |
+
doc.saveas(dxf_filepath)
|
186 |
+
return dxf_filepath
|
187 |
+
|
188 |
+
def extract_outlines(binary_image: np.ndarray) -> np.ndarray:
|
189 |
+
"""
|
190 |
+
Extracts and draws the outlines of masks from a binary image.
|
191 |
+
Args:
|
192 |
+
binary_image: Grayscale binary image where white represents masks and black is the background.
|
193 |
+
Returns:
|
194 |
+
Image with outlines drawn.
|
195 |
+
"""
|
196 |
+
# Detect contours from the binary image
|
197 |
+
contours, _ = cv2.findContours(
|
198 |
+
binary_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
|
199 |
+
)
|
200 |
+
|
201 |
+
outline_image = np.zeros_like(binary_image)
|
202 |
+
|
203 |
+
# Draw the contours on the blank image
|
204 |
+
cv2.drawContours(
|
205 |
+
outline_image, contours, -1, (255), thickness=1
|
206 |
+
) # White color for outlines
|
207 |
+
|
208 |
+
return cv2.bitwise_not(outline_image), contours
|
209 |
+
|
210 |
+
def to_dxf(contours):
|
211 |
+
doc = ezdxf.new()
|
212 |
+
msp = doc.modelspace()
|
213 |
+
|
214 |
+
for contour in contours:
|
215 |
+
points = [(point[0][0], point[0][1]) for point in contour]
|
216 |
+
msp.add_lwpolyline(points, close=True) # Add a polyline for each contour
|
217 |
+
|
218 |
+
doc.saveas("./outputs/out.dxf")
|
219 |
+
return "./outputs/out.dxf"
|
220 |
+
|
221 |
+
def smooth_contours(contour):
|
222 |
+
epsilon = 0.01 * cv2.arcLength(contour, True) # Adjust factor (e.g., 0.01)
|
223 |
+
return cv2.approxPolyDP(contour, epsilon, True)
|
224 |
+
|
225 |
+
|
226 |
+
def scale_image(image: np.ndarray, scale_factor: float) -> np.ndarray:
|
227 |
+
"""
|
228 |
+
Resize image by scaling both width and height by the same factor.
|
229 |
+
|
230 |
+
Args:
|
231 |
+
image: Input numpy image
|
232 |
+
scale_factor: Factor to scale the image (e.g., 0.5 for half size, 2 for double size)
|
233 |
+
|
234 |
+
Returns:
|
235 |
+
np.ndarray: Resized image
|
236 |
+
"""
|
237 |
+
if scale_factor <= 0:
|
238 |
+
raise ValueError("Scale factor must be positive")
|
239 |
+
|
240 |
+
current_height, current_width = image.shape[:2]
|
241 |
+
|
242 |
+
# Calculate new dimensions
|
243 |
+
new_width = int(current_width * scale_factor)
|
244 |
+
new_height = int(current_height * scale_factor)
|
245 |
+
|
246 |
+
# Choose interpolation method based on whether we're scaling up or down
|
247 |
+
interpolation = cv2.INTER_AREA if scale_factor < 1 else cv2.INTER_CUBIC
|
248 |
+
|
249 |
+
# Resize image
|
250 |
+
resized_image = cv2.resize(
|
251 |
+
image, (new_width, new_height), interpolation=interpolation
|
252 |
+
)
|
253 |
+
|
254 |
+
return resized_image
|
255 |
+
|
256 |
+
def detect_reference_square(img) -> np.ndarray:
|
257 |
+
box_detector = YOLO("./best1.pt")
|
258 |
+
res = box_detector.predict(img, conf=0.05)
|
259 |
+
del box_detector
|
260 |
+
return save_one_box(res[0].cpu().boxes.xyxy, res[0].orig_img, save=False), res[
|
261 |
+
0
|
262 |
+
].cpu().boxes.xyxy[0]
|
263 |
+
|
264 |
+
|
265 |
+
def resize_img(img: np.ndarray, resize_dim):
|
266 |
+
return np.array(Image.fromarray(img).resize(resize_dim))
|
267 |
+
|
268 |
+
def predict(image, offset_inches):
|
269 |
+
try:
|
270 |
+
reference_obj_img, scaling_box_coords = detect_reference_square(image)
|
271 |
+
except Exception as e:
|
272 |
+
raise gr.Error(f"Unable to DETECT COIN, please take another picture with different magnification level! Error: {e}")
|
273 |
+
|
274 |
+
reference_obj_img = make_square(reference_obj_img)
|
275 |
+
reference_square_mask = remove_bg(reference_obj_img)
|
276 |
+
reference_square_mask = resize_img(reference_square_mask, (reference_obj_img.shape[1], reference_obj_img.shape[0]))
|
277 |
+
|
278 |
+
try:
|
279 |
+
scaling_factor = calculate_scaling_factor(
|
280 |
+
reference_image_path="./coin.png",
|
281 |
+
target_image=reference_square_mask,
|
282 |
+
feature_detector="ORB",
|
283 |
+
)
|
284 |
+
except ZeroDivisionError:
|
285 |
+
scaling_factor = None
|
286 |
+
print("Error calculating scaling factor: Division by zero")
|
287 |
+
except Exception as e:
|
288 |
+
scaling_factor = None
|
289 |
+
print(f"Error calculating scaling factor: {e}")
|
290 |
+
|
291 |
+
# Default to a scaling factor of 1.0 if calculation fails
|
292 |
+
if scaling_factor is None or scaling_factor == 0:
|
293 |
+
scaling_factor = 1.0
|
294 |
+
print("Using default scaling factor of 1.0 due to calculation error")
|
295 |
+
|
296 |
+
orig_size = image.shape[:2]
|
297 |
+
objects_mask = remove_bg(image)
|
298 |
+
processed_size = objects_mask.shape[:2]
|
299 |
+
|
300 |
+
objects_mask = exclude_scaling_box(
|
301 |
+
objects_mask,
|
302 |
+
scaling_box_coords,
|
303 |
+
orig_size,
|
304 |
+
processed_size,
|
305 |
+
expansion_factor=1.5,
|
306 |
+
)
|
307 |
+
objects_mask = resize_img(objects_mask, (image.shape[1], image.shape[0]))
|
308 |
+
|
309 |
+
# Ensure offset_inches is valid
|
310 |
+
if scaling_factor != 0:
|
311 |
+
offset_pixels = (offset_inches / scaling_factor) * 2 + 1
|
312 |
+
else:
|
313 |
+
offset_pixels = 1 # Default value in case of invalid scaling factor
|
314 |
+
|
315 |
+
dilated_mask = cv2.dilate(objects_mask, np.ones((int(offset_pixels), int(offset_pixels)), np.uint8))
|
316 |
+
|
317 |
+
Image.fromarray(dilated_mask).save("./outputs/scaled_mask_new.jpg")
|
318 |
+
outlines, contours = extract_outlines(dilated_mask)
|
319 |
+
shrunked_img_contours = cv2.drawContours(image, contours, -1, (0, 0, 255), thickness=2)
|
320 |
+
dxf = save_dxf_spline(contours, scaling_factor, processed_size[0])
|
321 |
+
|
322 |
+
return (
|
323 |
+
shrunked_img_contours,
|
324 |
+
outlines,
|
325 |
+
dxf,
|
326 |
+
dilated_mask,
|
327 |
+
scaling_factor,
|
328 |
+
)
|
329 |
+
|
330 |
+
if __name__ == "__main__":
|
331 |
+
os.makedirs("./outputs", exist_ok=True)
|
332 |
+
|
333 |
+
ifer = gr.Interface(
|
334 |
+
fn=predict,
|
335 |
+
inputs=[
|
336 |
+
gr.Image(label="Input Image"),
|
337 |
+
gr.Number(label="Offset value for Mask(inches)", value=0.075),
|
338 |
+
],
|
339 |
+
outputs=[
|
340 |
+
gr.Image(label="Ouput Image"),
|
341 |
+
gr.Image(label="Outlines of Objects"),
|
342 |
+
gr.File(label="DXF file"),
|
343 |
+
gr.Image(label="Mask"),
|
344 |
+
gr.Textbox(
|
345 |
+
label="Scaling Factor(mm)",
|
346 |
+
placeholder="Every pixel is equal to mentioned number in inches",
|
347 |
+
),
|
348 |
+
],
|
349 |
+
examples=[
|
350 |
+
["./examples/Test20.jpg", 0.075],
|
351 |
+
["./examples/Test21.jpg", 0.075],
|
352 |
+
["./examples/Test22.jpg", 0.075],
|
353 |
+
["./examples/Test23.jpg", 0.075],
|
354 |
+
],
|
355 |
+
)
|
356 |
+
ifer.launch(share=True)
|
best1.pt
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:bd8be4cb96fec63fafd7036e2ee48f84085b01c0d3442c9fc42c3563a088e309
|
3 |
+
size 16087700
|