ammariii08 commited on
Commit
3850225
·
verified ·
1 Parent(s): c427817

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +62 -136
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, GeometryCollection
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.30)
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
- try:
292
- radius = circle_diameter / 2.0
293
- circle_poly = Point(center_inch).buffer(radius, resolution=64)
294
- tool_polygon = validate_and_fix_geometry(tool_polygon)
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
- try:
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
- # Check distance to existing centers
343
- too_close = any(np.hypot(cx - ex_x, cy - ex_y) < needed_center_distance
344
- for (ex_x, ex_y) in existing_centers)
 
345
  if too_close:
346
  continue
347
-
348
- try:
349
- circle_poly = Point((cx, cy)).buffer(radius, resolution=64)
350
- union_poly = tool_polygon.union(circle_poly)
351
- union_poly = validate_and_fix_geometry(union_poly)
352
-
353
- if union_poly is None:
354
- continue
355
-
356
- # Check for overlaps
357
- overlap = False
358
- for poly in all_polygons:
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
- except Exception as e:
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 draw_single_polygon(poly, image_rgb, scaling_factor, image_height, color=(0,0,255), thickness=2):
487
- # Handle different geometry types
488
- if poly.geom_type == "GeometryCollection":
489
- return
490
- elif poly.geom_type == "MultiPolygon":
491
- for geom in poly.geoms:
492
- draw_single_polygon(geom, image_rgb, scaling_factor, image_height, color, thickness)
493
- return
494
- elif poly.geom_type not in ["Polygon", "LinearRing"]:
495
- return
496
 
497
- if not hasattr(poly, 'exterior') or not poly.exterior:
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])