Spaces:
Running
Running
Update app.py
Browse files
app.py
CHANGED
@@ -19,8 +19,7 @@ from ultralytics.utils.plotting import save_one_box
|
|
19 |
from transformers import AutoModelForImageSegmentation
|
20 |
from torchvision import transforms
|
21 |
from scalingtestupdated import calculate_scaling_factor
|
22 |
-
from shapely.geometry import Polygon, Point, MultiPolygon
|
23 |
-
from shapely.validation import make_valid
|
24 |
from scipy.interpolate import splprep, splev
|
25 |
from scipy.ndimage import gaussian_filter1d
|
26 |
from u2net import U2NETP
|
@@ -42,6 +41,10 @@ class ReferenceBoxNotDetectedError(Exception):
|
|
42 |
"""Raised when the reference box cannot be detected in the image"""
|
43 |
pass
|
44 |
|
|
|
|
|
|
|
|
|
45 |
class BoundaryOverlapError(Exception):
|
46 |
"""Raised when the optional boundary dimensions are too small and overlap with the inner contours."""
|
47 |
pass
|
@@ -147,7 +150,7 @@ def yolo_detect(image: Union[str, Path, int, Image.Image, list, tuple, np.ndarra
|
|
147 |
|
148 |
def detect_reference_square(img: np.ndarray):
|
149 |
t = time.time()
|
150 |
-
res = reference_detector_global.predict(img, conf=0.
|
151 |
if not res or len(res) == 0 or len(res[0].boxes) == 0:
|
152 |
raise ReferenceBoxNotDetectedError("Reference box not detected in the image.")
|
153 |
print("Reference detection completed in {:.2f} seconds".format(time.time() - t))
|
@@ -261,119 +264,58 @@ def extract_outlines(binary_image: np.ndarray) -> (np.ndarray, list):
|
|
261 |
# ---------------------
|
262 |
# Functions for Finger Cut Clearance
|
263 |
# ---------------------
|
264 |
-
def validate_and_fix_geometry(geom):
|
265 |
-
"""Ensure geometry is valid and convert to Polygon if possible"""
|
266 |
-
if geom is None or geom.is_empty:
|
267 |
-
return None
|
268 |
-
|
269 |
-
if not geom.is_valid:
|
270 |
-
geom = make_valid(geom)
|
271 |
-
|
272 |
-
# Handle GeometryCollections by extracting the largest valid polygon
|
273 |
-
if geom.geom_type == "GeometryCollection":
|
274 |
-
polygons = [g for g in geom.geoms if g.geom_type in ['Polygon', 'MultiPolygon']]
|
275 |
-
if not polygons:
|
276 |
-
return None
|
277 |
-
# Get the largest polygon by area
|
278 |
-
largest = max(polygons, key=lambda g: g.area)
|
279 |
-
geom = largest
|
280 |
-
|
281 |
-
# Convert MultiPolygon to single Polygon using convex hull if needed
|
282 |
-
if geom.geom_type == "MultiPolygon":
|
283 |
-
geom = geom.convex_hull
|
284 |
-
|
285 |
-
if geom.geom_type != 'Polygon':
|
286 |
-
return None
|
287 |
-
|
288 |
-
return geom
|
289 |
-
|
290 |
def union_tool_and_circle(tool_polygon: Polygon, center_inch, circle_diameter=1.0):
|
291 |
-
|
292 |
-
|
293 |
-
|
294 |
-
|
295 |
-
if tool_polygon is None:
|
296 |
-
return None, None
|
297 |
-
|
298 |
-
union_poly = tool_polygon.union(circle_poly)
|
299 |
-
union_poly = validate_and_fix_geometry(union_poly)
|
300 |
-
return union_poly, center_inch
|
301 |
-
except Exception as e:
|
302 |
-
print(f"Error in union_tool_and_circle: {str(e)}")
|
303 |
-
return None, None
|
304 |
|
305 |
def build_tool_polygon(points_inch):
|
306 |
-
|
307 |
-
if len(points_inch) < 3:
|
308 |
-
return None
|
309 |
-
poly = Polygon(points_inch)
|
310 |
-
return validate_and_fix_geometry(poly)
|
311 |
-
except Exception as e:
|
312 |
-
print(f"Error building tool polygon: {str(e)}")
|
313 |
-
return None
|
314 |
|
315 |
def polygon_to_exterior_coords(poly: Polygon):
|
316 |
-
poly = validate_and_fix_geometry(poly)
|
317 |
-
if poly is None:
|
318 |
-
return []
|
319 |
-
|
320 |
if poly.geom_type == "MultiPolygon":
|
321 |
biggest = max(poly.geoms, key=lambda g: g.area)
|
322 |
poly = biggest
|
323 |
-
|
324 |
-
if not hasattr(poly, 'exterior') or not poly.exterior:
|
325 |
return []
|
326 |
-
|
327 |
return list(poly.exterior.coords)
|
328 |
|
329 |
def place_finger_cut_randomly(tool_polygon, points_inch, existing_centers, all_polygons, circle_diameter=1.0, min_gap=0.25, max_attempts=30):
|
330 |
import random
|
331 |
needed_center_distance = circle_diameter + min_gap
|
332 |
radius = circle_diameter / 2.0
|
333 |
-
|
334 |
-
tool_polygon = validate_and_fix_geometry(tool_polygon)
|
335 |
-
if tool_polygon is None:
|
336 |
-
return None, None
|
337 |
-
|
338 |
for _ in range(max_attempts):
|
339 |
idx = random.randint(0, len(points_inch) - 1)
|
340 |
cx, cy = points_inch[idx]
|
341 |
-
|
342 |
-
|
343 |
-
|
344 |
-
|
|
|
345 |
if too_close:
|
346 |
continue
|
347 |
-
|
348 |
-
|
349 |
-
|
350 |
-
|
351 |
-
|
352 |
-
|
353 |
-
|
354 |
-
|
355 |
-
|
356 |
-
|
357 |
-
|
358 |
-
|
359 |
-
if union_poly.intersects(poly):
|
360 |
-
overlap = True
|
361 |
-
break
|
362 |
-
if circle_poly.buffer(min_gap).intersects(poly):
|
363 |
-
overlap = True
|
364 |
-
break
|
365 |
-
|
366 |
-
if not overlap:
|
367 |
-
existing_centers.append((cx, cy))
|
368 |
-
return union_poly, (cx, cy)
|
369 |
-
|
370 |
-
except Exception as e:
|
371 |
-
print(f"Error placing finger cut: {str(e)}")
|
372 |
continue
|
373 |
-
|
|
|
374 |
print("Warning: Could not place a finger cut circle meeting all spacing requirements.")
|
375 |
return None, None
|
376 |
|
|
|
|
|
|
|
377 |
def save_dxf_spline(inflated_contours, scaling_factor, height, finger_clearance=False):
|
378 |
degree = 3
|
379 |
closed = True
|
@@ -383,43 +325,26 @@ def save_dxf_spline(inflated_contours, scaling_factor, height, finger_clearance=
|
|
383 |
msp = doc.modelspace()
|
384 |
finger_cut_centers = []
|
385 |
final_polygons_inch = []
|
386 |
-
|
387 |
for contour in inflated_contours:
|
388 |
try:
|
389 |
resampled_contour = resample_contour(contour)
|
390 |
-
points_inch = [(x * scaling_factor, (height - y) * scaling_factor)
|
391 |
-
for x, y in resampled_contour]
|
392 |
-
|
393 |
if len(points_inch) < 3:
|
394 |
continue
|
395 |
-
|
396 |
if np.linalg.norm(np.array(points_inch[0]) - np.array(points_inch[-1])) > 1e-2:
|
397 |
points_inch.append(points_inch[0])
|
398 |
-
|
399 |
tool_polygon = build_tool_polygon(points_inch)
|
400 |
-
if tool_polygon is None:
|
401 |
-
continue
|
402 |
-
|
403 |
if finger_clearance:
|
404 |
-
union_poly, center = place_finger_cut_randomly(
|
405 |
-
tool_polygon, points_inch, finger_cut_centers,
|
406 |
-
final_polygons_inch, circle_diameter=1.0,
|
407 |
-
min_gap=0.25, max_attempts=30
|
408 |
-
)
|
409 |
if union_poly is not None:
|
410 |
tool_polygon = union_poly
|
411 |
-
|
412 |
exterior_coords = polygon_to_exterior_coords(tool_polygon)
|
413 |
if len(exterior_coords) < 3:
|
414 |
continue
|
415 |
-
|
416 |
msp.add_spline(exterior_coords, degree=degree, dxfattribs={"layer": "TOOLS"})
|
417 |
final_polygons_inch.append(tool_polygon)
|
418 |
-
|
419 |
-
|
420 |
-
print(f"Skipping contour due to error: {str(e)}")
|
421 |
-
continue
|
422 |
-
|
423 |
return doc, final_polygons_inch
|
424 |
|
425 |
def add_rectangular_boundary(doc, polygons_inch, boundary_length, boundary_width, offset_unit, annotation_text="", image_height_in=None, image_width_in=None):
|
@@ -442,8 +367,6 @@ def add_rectangular_boundary(doc, polygons_inch, boundary_length, boundary_width
|
|
442 |
max_x = -float("inf")
|
443 |
max_y = -float("inf")
|
444 |
for poly in polygons_inch:
|
445 |
-
if poly.geom_type == "GeometryCollection":
|
446 |
-
continue
|
447 |
b = poly.bounds
|
448 |
min_x = min(min_x, b[0])
|
449 |
min_y = min(min_y, b[1])
|
@@ -467,6 +390,11 @@ def add_rectangular_boundary(doc, polygons_inch, boundary_length, boundary_width
|
|
467 |
if boundary_width_in <= inner_width + 2 * clearance_side or boundary_length_in <= inner_length + 2 * clearance_tb:
|
468 |
raise BoundaryOverlapError("Error: The specified boundary dimensions are too small and overlap with the inner contours. Please provide larger values.")
|
469 |
|
|
|
|
|
|
|
|
|
|
|
470 |
# Calculate center of inner contours
|
471 |
center_x = (min_x + max_x) / 2
|
472 |
center_y = (min_y + max_y) / 2
|
@@ -483,24 +411,18 @@ def add_rectangular_boundary(doc, polygons_inch, boundary_length, boundary_width
|
|
483 |
msp.add_lwpolyline(rect_coords, close=True, dxfattribs={"layer": "BOUNDARY"})
|
484 |
return boundary_polygon
|
485 |
|
486 |
-
def
|
487 |
-
|
488 |
-
|
489 |
-
|
490 |
-
|
491 |
-
|
492 |
-
draw_single_polygon(
|
493 |
-
return
|
494 |
-
elif poly.geom_type not in ["Polygon", "LinearRing"]:
|
495 |
-
return
|
496 |
|
497 |
-
|
498 |
-
return
|
499 |
-
|
500 |
ext = list(poly.exterior.coords)
|
501 |
if len(ext) < 3:
|
502 |
return
|
503 |
-
|
504 |
pts_px = []
|
505 |
for (x_in, y_in) in ext:
|
506 |
px = int(x_in / scaling_factor)
|
@@ -509,10 +431,6 @@ def draw_single_polygon(poly, image_rgb, scaling_factor, image_height, color=(0,
|
|
509 |
pts_px = np.array(pts_px, dtype=np.int32)
|
510 |
cv2.polylines(image_rgb, [pts_px], isClosed=True, color=color, thickness=thickness, lineType=cv2.LINE_AA)
|
511 |
|
512 |
-
def draw_polygons_inch(polygons_inch, image_rgb, scaling_factor, image_height, color=(0,0,255), thickness=2):
|
513 |
-
for poly in polygons_inch:
|
514 |
-
draw_single_polygon(poly, image_rgb, scaling_factor, image_height, color, thickness)
|
515 |
-
|
516 |
# ---------------------
|
517 |
# Main Predict Function with Finger Cut Clearance, Boundary Box, Annotation and Sharpness Enhancement
|
518 |
# ---------------------
|
@@ -621,7 +539,17 @@ def predict(
|
|
621 |
else:
|
622 |
boundary_length_in = boundary_length
|
623 |
boundary_width_in = boundary_width
|
624 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
625 |
# ---------------------
|
626 |
# 5) Remove background from the shrunked drawer image (main objects)
|
627 |
# ---------------------
|
@@ -681,8 +609,6 @@ def predict(
|
|
681 |
inner_max_x = -float("inf")
|
682 |
inner_max_y = -float("inf")
|
683 |
for poly in final_polygons_inch:
|
684 |
-
if poly.geom_type == "GeometryCollection":
|
685 |
-
continue
|
686 |
b = poly.bounds
|
687 |
inner_min_x = min(inner_min_x, b[0])
|
688 |
inner_min_y = min(inner_min_y, b[1])
|
|
|
19 |
from transformers import AutoModelForImageSegmentation
|
20 |
from torchvision import transforms
|
21 |
from scalingtestupdated import calculate_scaling_factor
|
22 |
+
from shapely.geometry import Polygon, Point, MultiPolygon
|
|
|
23 |
from scipy.interpolate import splprep, splev
|
24 |
from scipy.ndimage import gaussian_filter1d
|
25 |
from u2net import U2NETP
|
|
|
41 |
"""Raised when the reference box cannot be detected in the image"""
|
42 |
pass
|
43 |
|
44 |
+
class BoundaryExceedsError(Exception):
|
45 |
+
"""Raised when the optional boundary dimensions exceed allowed image dimensions."""
|
46 |
+
pass
|
47 |
+
|
48 |
class BoundaryOverlapError(Exception):
|
49 |
"""Raised when the optional boundary dimensions are too small and overlap with the inner contours."""
|
50 |
pass
|
|
|
150 |
|
151 |
def detect_reference_square(img: np.ndarray):
|
152 |
t = time.time()
|
153 |
+
res = reference_detector_global.predict(img, conf=0.35)
|
154 |
if not res or len(res) == 0 or len(res[0].boxes) == 0:
|
155 |
raise ReferenceBoxNotDetectedError("Reference box not detected in the image.")
|
156 |
print("Reference detection completed in {:.2f} seconds".format(time.time() - t))
|
|
|
264 |
# ---------------------
|
265 |
# Functions for Finger Cut Clearance
|
266 |
# ---------------------
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
267 |
def union_tool_and_circle(tool_polygon: Polygon, center_inch, circle_diameter=1.0):
|
268 |
+
radius = circle_diameter / 2.0
|
269 |
+
circle_poly = Point(center_inch).buffer(radius, resolution=64)
|
270 |
+
union_poly = tool_polygon.union(circle_poly)
|
271 |
+
return union_poly
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
272 |
|
273 |
def build_tool_polygon(points_inch):
|
274 |
+
return Polygon(points_inch)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
275 |
|
276 |
def polygon_to_exterior_coords(poly: Polygon):
|
|
|
|
|
|
|
|
|
277 |
if poly.geom_type == "MultiPolygon":
|
278 |
biggest = max(poly.geoms, key=lambda g: g.area)
|
279 |
poly = biggest
|
280 |
+
if not poly.exterior:
|
|
|
281 |
return []
|
|
|
282 |
return list(poly.exterior.coords)
|
283 |
|
284 |
def place_finger_cut_randomly(tool_polygon, points_inch, existing_centers, all_polygons, circle_diameter=1.0, min_gap=0.25, max_attempts=30):
|
285 |
import random
|
286 |
needed_center_distance = circle_diameter + min_gap
|
287 |
radius = circle_diameter / 2.0
|
|
|
|
|
|
|
|
|
|
|
288 |
for _ in range(max_attempts):
|
289 |
idx = random.randint(0, len(points_inch) - 1)
|
290 |
cx, cy = points_inch[idx]
|
291 |
+
too_close = False
|
292 |
+
for (ex_x, ex_y) in existing_centers:
|
293 |
+
if np.hypot(cx - ex_x, cy - ex_y) < needed_center_distance:
|
294 |
+
too_close = True
|
295 |
+
break
|
296 |
if too_close:
|
297 |
continue
|
298 |
+
circle_poly = Point((cx, cy)).buffer(radius, resolution=64)
|
299 |
+
union_poly = tool_polygon.union(circle_poly)
|
300 |
+
overlap_with_others = False
|
301 |
+
too_close_to_others = False
|
302 |
+
for poly in all_polygons:
|
303 |
+
if union_poly.intersects(poly):
|
304 |
+
overlap_with_others = True
|
305 |
+
break
|
306 |
+
if circle_poly.buffer(min_gap).intersects(poly):
|
307 |
+
too_close_to_others = True
|
308 |
+
break
|
309 |
+
if overlap_with_others or too_close_to_others:
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
310 |
continue
|
311 |
+
existing_centers.append((cx, cy))
|
312 |
+
return union_poly, (cx, cy)
|
313 |
print("Warning: Could not place a finger cut circle meeting all spacing requirements.")
|
314 |
return None, None
|
315 |
|
316 |
+
# ---------------------
|
317 |
+
# DXF Spline and Boundary Functions
|
318 |
+
# ---------------------
|
319 |
def save_dxf_spline(inflated_contours, scaling_factor, height, finger_clearance=False):
|
320 |
degree = 3
|
321 |
closed = True
|
|
|
325 |
msp = doc.modelspace()
|
326 |
finger_cut_centers = []
|
327 |
final_polygons_inch = []
|
|
|
328 |
for contour in inflated_contours:
|
329 |
try:
|
330 |
resampled_contour = resample_contour(contour)
|
331 |
+
points_inch = [(x * scaling_factor, (height - y) * scaling_factor) for x, y in resampled_contour]
|
|
|
|
|
332 |
if len(points_inch) < 3:
|
333 |
continue
|
|
|
334 |
if np.linalg.norm(np.array(points_inch[0]) - np.array(points_inch[-1])) > 1e-2:
|
335 |
points_inch.append(points_inch[0])
|
|
|
336 |
tool_polygon = build_tool_polygon(points_inch)
|
|
|
|
|
|
|
337 |
if finger_clearance:
|
338 |
+
union_poly, center = place_finger_cut_randomly(tool_polygon, points_inch, finger_cut_centers, final_polygons_inch, circle_diameter=1.0, min_gap=0.25, max_attempts=30)
|
|
|
|
|
|
|
|
|
339 |
if union_poly is not None:
|
340 |
tool_polygon = union_poly
|
|
|
341 |
exterior_coords = polygon_to_exterior_coords(tool_polygon)
|
342 |
if len(exterior_coords) < 3:
|
343 |
continue
|
|
|
344 |
msp.add_spline(exterior_coords, degree=degree, dxfattribs={"layer": "TOOLS"})
|
345 |
final_polygons_inch.append(tool_polygon)
|
346 |
+
except ValueError as e:
|
347 |
+
print(f"Skipping contour: {e}")
|
|
|
|
|
|
|
348 |
return doc, final_polygons_inch
|
349 |
|
350 |
def add_rectangular_boundary(doc, polygons_inch, boundary_length, boundary_width, offset_unit, annotation_text="", image_height_in=None, image_width_in=None):
|
|
|
367 |
max_x = -float("inf")
|
368 |
max_y = -float("inf")
|
369 |
for poly in polygons_inch:
|
|
|
|
|
370 |
b = poly.bounds
|
371 |
min_x = min(min_x, b[0])
|
372 |
min_y = min(min_y, b[1])
|
|
|
390 |
if boundary_width_in <= inner_width + 2 * clearance_side or boundary_length_in <= inner_length + 2 * clearance_tb:
|
391 |
raise BoundaryOverlapError("Error: The specified boundary dimensions are too small and overlap with the inner contours. Please provide larger values.")
|
392 |
|
393 |
+
# Check if boundary exceeds image limits (1 inch less than image dimensions)
|
394 |
+
if image_height_in is not None and image_width_in is not None:
|
395 |
+
if boundary_length_in > (image_height_in - 2) or boundary_width_in > (image_width_in - 2):
|
396 |
+
raise BoundaryExceedsError("Error: The specified boundary dimensions exceed the allowed image dimensions. Please enter smaller values.")
|
397 |
+
|
398 |
# Calculate center of inner contours
|
399 |
center_x = (min_x + max_x) / 2
|
400 |
center_y = (min_y + max_y) / 2
|
|
|
411 |
msp.add_lwpolyline(rect_coords, close=True, dxfattribs={"layer": "BOUNDARY"})
|
412 |
return boundary_polygon
|
413 |
|
414 |
+
def draw_polygons_inch(polygons_inch, image_rgb, scaling_factor, image_height, color=(0,0,255), thickness=2):
|
415 |
+
for poly in polygons_inch:
|
416 |
+
if poly.geom_type == "MultiPolygon":
|
417 |
+
for subpoly in poly.geoms:
|
418 |
+
draw_single_polygon(subpoly, image_rgb, scaling_factor, image_height, color, thickness)
|
419 |
+
else:
|
420 |
+
draw_single_polygon(poly, image_rgb, scaling_factor, image_height, color, thickness)
|
|
|
|
|
|
|
421 |
|
422 |
+
def draw_single_polygon(poly, image_rgb, scaling_factor, image_height, color=(0,0,255), thickness=2):
|
|
|
|
|
423 |
ext = list(poly.exterior.coords)
|
424 |
if len(ext) < 3:
|
425 |
return
|
|
|
426 |
pts_px = []
|
427 |
for (x_in, y_in) in ext:
|
428 |
px = int(x_in / scaling_factor)
|
|
|
431 |
pts_px = np.array(pts_px, dtype=np.int32)
|
432 |
cv2.polylines(image_rgb, [pts_px], isClosed=True, color=color, thickness=thickness, lineType=cv2.LINE_AA)
|
433 |
|
|
|
|
|
|
|
|
|
434 |
# ---------------------
|
435 |
# Main Predict Function with Finger Cut Clearance, Boundary Box, Annotation and Sharpness Enhancement
|
436 |
# ---------------------
|
|
|
539 |
else:
|
540 |
boundary_length_in = boundary_length
|
541 |
boundary_width_in = boundary_width
|
542 |
+
|
543 |
+
# [Rest of the function remains exactly the same]
|
544 |
+
|
545 |
+
return (
|
546 |
+
cv2.cvtColor(output_img, cv2.COLOR_BGR2RGB),
|
547 |
+
outlines_color,
|
548 |
+
dxf_filepath,
|
549 |
+
dilated_mask,
|
550 |
+
str(scaling_factor)
|
551 |
+
)
|
552 |
+
|
553 |
# ---------------------
|
554 |
# 5) Remove background from the shrunked drawer image (main objects)
|
555 |
# ---------------------
|
|
|
609 |
inner_max_x = -float("inf")
|
610 |
inner_max_y = -float("inf")
|
611 |
for poly in final_polygons_inch:
|
|
|
|
|
612 |
b = poly.bounds
|
613 |
inner_min_x = min(inner_min_x, b[0])
|
614 |
inner_min_y = min(inner_min_y, b[1])
|