mlbench123 commited on
Commit
fab9725
·
verified ·
1 Parent(s): 0d118b9

Upload 2 files

Browse files
Files changed (2) hide show
  1. app.py +356 -0
  2. best1.pt +3 -0
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