mlbench123 commited on
Commit
8cead57
·
verified ·
1 Parent(s): c1657df

Delete app.py

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