Nighty3912 commited on
Commit
dca9b7a
·
verified ·
1 Parent(s): 95bf645

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +52 -60
app.py CHANGED
@@ -43,7 +43,7 @@ class DrawerNotDetectedError(Exception):
43
  pass
44
 
45
  class ReferenceBoxNotDetectedError(Exception):
46
- """Raised when the Reference coin cannot be detected in the image"""
47
  pass
48
 
49
  class BoundaryOverlapError(Exception):
@@ -69,10 +69,10 @@ print("YOLOWorld model loaded in {:.2f} seconds".format(time.time() - start_time
69
 
70
  print("Loading YOLO reference model...")
71
  start_time = time.time()
72
- reference_model_path = os.path.join(CACHE_DIR, "coin_det.pt")
73
  if not os.path.exists(reference_model_path):
74
  print("Caching YOLO reference model to", reference_model_path)
75
- shutil.copy("coin_det.pt", reference_model_path)
76
  reference_detector_global = YOLO(reference_model_path)
77
  print("YOLO reference model loaded in {:.2f} seconds".format(time.time() - start_time))
78
 
@@ -120,7 +120,7 @@ def unload_and_reload_models():
120
  gc.collect()
121
  new_drawer_detector = YOLOWorld(os.path.join(CACHE_DIR, "yolov8x-worldv2.pt"))
122
  new_drawer_detector.set_classes(["box"])
123
- new_reference_detector = YOLO(os.path.join(CACHE_DIR, "coin_det.pt"))
124
  new_birefnet = AutoModelForImageSegmentation.from_pretrained(
125
  "zhengpeng7/BiRefNet", trust_remote_code=True, cache_dir=CACHE_DIR
126
  )
@@ -155,9 +155,9 @@ def yolo_detect(image: Union[str, Path, int, Image.Image, list, tuple, np.ndarra
155
 
156
  def detect_reference_square(img: np.ndarray):
157
  t = time.time()
158
- res = reference_detector_global.predict(img, conf=0.3)
159
  if not res or len(res) == 0 or len(res[0].boxes) == 0:
160
- raise ReferenceBoxNotDetectedError("Reference Coin not detected in the image.")
161
  print("Reference detection completed in {:.2f} seconds".format(time.time() - t))
162
  return (
163
  save_one_box(res[0].cpu().boxes.xyxy, res[0].orig_img, save=False),
@@ -286,51 +286,42 @@ def polygon_to_exterior_coords(poly: Polygon):
286
  return []
287
  return list(poly.exterior.coords)
288
 
289
-
290
- def place_finger_cut_adjusted(tool_polygon, points_inch, existing_centers, all_polygons, circle_diameter=1.0, min_gap=0.25, max_attempts=30): #1st best
291
  needed_center_distance = circle_diameter + min_gap
292
  radius = circle_diameter / 2.0
293
- import random
294
- for _ in range(max_attempts):
295
- idx = random.randint(0, len(points_inch) - 1)
296
- cx, cy = points_inch[idx]
297
-
298
- # Check if this point is too close to an existing center
299
- too_close = any(np.hypot(cx - ex_x, cy - ex_y) < needed_center_distance for ex_x, ex_y in existing_centers)
300
- if too_close:
301
- continue
302
-
303
- # Create the finger cut circle and try adding it to the tool
304
- circle_poly = Point((cx, cy)).buffer(radius, resolution=64)
305
- union_poly = tool_polygon.union(circle_poly)
306
-
307
- # Check for overlap and spacing with other tools
308
- overlap_with_others = False
309
- too_close_to_others = False
310
-
311
- for poly in all_polygons:
312
- if poly.equals(tool_polygon):
313
- continue # Skip comparing the tool to itself
314
-
315
- if union_poly.intersects(poly):
316
- overlap_with_others = True
317
- break
318
-
319
- if circle_poly.buffer(min_gap).intersects(poly):
320
- too_close_to_others = True
321
- break
322
-
323
- if overlap_with_others or too_close_to_others:
324
- continue
325
-
326
- existing_centers.append((cx, cy))
327
- return union_poly, (cx, cy)
328
-
329
  print("Warning: Could not place a finger cut circle meeting all spacing requirements.")
330
  return None, None
331
 
332
-
333
-
334
  # ---------------------
335
  # DXF Spline and Boundary Functions
336
  # ---------------------
@@ -353,7 +344,7 @@ def save_dxf_spline(inflated_contours, scaling_factor, height, finger_clearance=
353
  points_inch.append(points_inch[0])
354
  tool_polygon = build_tool_polygon(points_inch)
355
  if finger_clearance:
356
- union_poly, center = place_finger_cut_adjusted(tool_polygon, points_inch, finger_cut_centers, final_polygons_inch, circle_diameter=1.0, min_gap=0.25, max_attempts=100)
357
  if union_poly is not None:
358
  tool_polygon = union_poly
359
  exterior_coords = polygon_to_exterior_coords(tool_polygon)
@@ -420,11 +411,13 @@ def add_rectangular_boundary(doc, polygons_inch, boundary_length, boundary_width
420
  msp.add_lwpolyline(rect_coords, close=True, dxfattribs={"layer": "BOUNDARY"})
421
 
422
  text_top = boundary_polygon.bounds[1] + 1
423
- too_small = boundary_width_in < inner_width + 2 * clearance_side or boundary_length_in < inner_length + 2 * clearance_tb
424
- if too_small:
425
- raise BoundaryOverlapError("Error: The specified boundary dimensions are too small and overlap with the inner contours. Please provide larger values.")
426
- if annotation_text.strip() and text_top > min_y - 0.75:
427
- raise TextOverlapError("Error: The text is too close to the inner contours. Please increase boundary length.")
 
 
428
  return boundary_polygon
429
 
430
  def draw_polygons_inch(polygons_inch, image_rgb, scaling_factor, image_height, color=(0,0,255), thickness=2):
@@ -508,7 +501,7 @@ def predict(
508
  try:
509
  t = time.time()
510
  reference_obj_img, scaling_box_coords = detect_reference_square(shrunked_img)
511
- print("Reference coin detection completed in {:.2f} seconds".format(time.time() - t))
512
  except ReferenceBoxNotDetectedError as e:
513
  return None, None, None, None, f"Error: {str(e)}"
514
 
@@ -518,15 +511,14 @@ def predict(
518
  t = time.time()
519
  reference_obj_img = make_square(reference_obj_img)
520
  reference_square_mask = remove_bg_u2netp(reference_obj_img)
521
- reference_square_mask= resize_img(reference_square_mask,(reference_obj_img.shape[1],reference_obj_img.shape[0]))
522
  print("Reference image processing completed in {:.2f} seconds".format(time.time() - t))
523
 
524
  t = time.time()
525
  try:
526
  cv2.imwrite("mask.jpg", cv2.cvtColor(reference_obj_img, cv2.COLOR_RGB2GRAY))
527
  scaling_factor = calculate_scaling_factor(
 
528
  target_image=reference_square_mask,
529
- reference_obj_size_mm=0.955,
530
  feature_detector="ORB",
531
  )
532
  except ZeroDivisionError:
@@ -537,8 +529,8 @@ def predict(
537
  print(f"Error calculating scaling factor: {e}")
538
 
539
  if scaling_factor is None or scaling_factor == 0:
540
- scaling_factor = 0.7
541
- print("Using default scaling factor of 0.7 due to calculation error")
542
  gc.collect()
543
  print("Scaling factor determined: {}".format(scaling_factor))
544
 
@@ -572,7 +564,7 @@ def predict(
572
  objects_mask = remove_bg(shrunked_img)
573
  processed_size = objects_mask.shape[:2]
574
 
575
- objects_mask = exclude_scaling_box(objects_mask, scaling_box_coords, orig_size, processed_size, expansion_factor=1.2)
576
  objects_mask = resize_img(objects_mask, (shrunked_img.shape[1], shrunked_img.shape[0]))
577
  del scaling_box_coords
578
  gc.collect()
@@ -781,8 +773,8 @@ if __name__ == "__main__":
781
  gr.Textbox(label="Scaling Factor (inches/pixel)")
782
  ],
783
  examples=[
784
- ["./Test20.jpg", 0.075, "inches", "No", "No", 30.0, 20.0, "MyTool"],
785
- ["./Test21.jpg", 0.075, "inches", "Yes", "Yes", 30.0, 20.0, "Tool2"]
786
  ]
787
  )
788
  iface.launch(share=True)
 
43
  pass
44
 
45
  class ReferenceBoxNotDetectedError(Exception):
46
+ """Raised when the reference box cannot be detected in the image"""
47
  pass
48
 
49
  class BoundaryOverlapError(Exception):
 
69
 
70
  print("Loading YOLO reference model...")
71
  start_time = time.time()
72
+ reference_model_path = os.path.join(CACHE_DIR, "best.pt")
73
  if not os.path.exists(reference_model_path):
74
  print("Caching YOLO reference model to", reference_model_path)
75
+ shutil.copy("best.pt", reference_model_path)
76
  reference_detector_global = YOLO(reference_model_path)
77
  print("YOLO reference model loaded in {:.2f} seconds".format(time.time() - start_time))
78
 
 
120
  gc.collect()
121
  new_drawer_detector = YOLOWorld(os.path.join(CACHE_DIR, "yolov8x-worldv2.pt"))
122
  new_drawer_detector.set_classes(["box"])
123
+ new_reference_detector = YOLO(os.path.join(CACHE_DIR, "best.pt"))
124
  new_birefnet = AutoModelForImageSegmentation.from_pretrained(
125
  "zhengpeng7/BiRefNet", trust_remote_code=True, cache_dir=CACHE_DIR
126
  )
 
155
 
156
  def detect_reference_square(img: np.ndarray):
157
  t = time.time()
158
+ res = reference_detector_global.predict(img, conf=0.15)
159
  if not res or len(res) == 0 or len(res[0].boxes) == 0:
160
+ raise ReferenceBoxNotDetectedError("Reference box not detected in the image.")
161
  print("Reference detection completed in {:.2f} seconds".format(time.time() - t))
162
  return (
163
  save_one_box(res[0].cpu().boxes.xyxy, res[0].orig_img, save=False),
 
286
  return []
287
  return list(poly.exterior.coords)
288
 
289
+ def place_finger_cut_adjusted(tool_polygon, points_inch, existing_centers, all_polygons, circle_diameter=1.0, min_gap=0.25, max_attempts=30):
290
+ import random
291
  needed_center_distance = circle_diameter + min_gap
292
  radius = circle_diameter / 2.0
293
+ attempts = 0
294
+ indices = list(range(len(points_inch)))
295
+ random.shuffle(indices) # Shuffle indices for randomness
296
+
297
+ for i in indices:
298
+ if attempts >= max_attempts:
299
+ break
300
+ cx, cy = points_inch[i]
301
+ # Try small adjustments around the chosen candidate
302
+ for dx in np.linspace(-0.1, 0.1, 5):
303
+ for dy in np.linspace(-0.1, 0.1, 5):
304
+ candidate_center = (cx + dx, cy + dy)
305
+ # Check distance from already placed centers
306
+ if any(np.hypot(candidate_center[0] - ex, candidate_center[1] - ey) < needed_center_distance for ex, ey in existing_centers):
307
+ continue
308
+ circle_poly = Point(candidate_center).buffer(radius, resolution=64)
309
+ union_poly = tool_polygon.union(circle_poly)
310
+ overlap = False
311
+ # Check against other tool polygons for overlap or proximity issues
312
+ for poly in all_polygons:
313
+ if union_poly.intersects(poly) or circle_poly.buffer(min_gap).intersects(poly):
314
+ overlap = True
315
+ break
316
+ if overlap:
317
+ continue
318
+ # If candidate passes, accept it
319
+ existing_centers.append(candidate_center)
320
+ return union_poly, candidate_center
321
+ attempts += 1
 
 
 
 
 
 
 
322
  print("Warning: Could not place a finger cut circle meeting all spacing requirements.")
323
  return None, None
324
 
 
 
325
  # ---------------------
326
  # DXF Spline and Boundary Functions
327
  # ---------------------
 
344
  points_inch.append(points_inch[0])
345
  tool_polygon = build_tool_polygon(points_inch)
346
  if finger_clearance:
347
+ union_poly, center = place_finger_cut_adjusted(tool_polygon, points_inch, finger_cut_centers, final_polygons_inch, circle_diameter=1.0, min_gap=0.25, max_attempts=30)
348
  if union_poly is not None:
349
  tool_polygon = union_poly
350
  exterior_coords = polygon_to_exterior_coords(tool_polygon)
 
411
  msp.add_lwpolyline(rect_coords, close=True, dxfattribs={"layer": "BOUNDARY"})
412
 
413
  text_top = boundary_polygon.bounds[1] + 1
414
+ if (annotation_text.strip()==0):
415
+ if boundary_width_in <= inner_width + 2 * clearance_side or boundary_length_in <= inner_length + 2 * clearance_tb:
416
+ raise BoundaryOverlapError("Error: The specified boundary dimensions are too small and overlap with the inner contours. Please provide larger values.")
417
+ else:
418
+ if text_top > (min_y - 0.75):
419
+ raise TextOverlapError("Error: The Text is overlapping the inner contours of the object.")
420
+
421
  return boundary_polygon
422
 
423
  def draw_polygons_inch(polygons_inch, image_rgb, scaling_factor, image_height, color=(0,0,255), thickness=2):
 
501
  try:
502
  t = time.time()
503
  reference_obj_img, scaling_box_coords = detect_reference_square(shrunked_img)
504
+ print("Reference square detection completed in {:.2f} seconds".format(time.time() - t))
505
  except ReferenceBoxNotDetectedError as e:
506
  return None, None, None, None, f"Error: {str(e)}"
507
 
 
511
  t = time.time()
512
  reference_obj_img = make_square(reference_obj_img)
513
  reference_square_mask = remove_bg_u2netp(reference_obj_img)
 
514
  print("Reference image processing completed in {:.2f} seconds".format(time.time() - t))
515
 
516
  t = time.time()
517
  try:
518
  cv2.imwrite("mask.jpg", cv2.cvtColor(reference_obj_img, cv2.COLOR_RGB2GRAY))
519
  scaling_factor = calculate_scaling_factor(
520
+ reference_image_path="./Reference_ScalingBox.jpg",
521
  target_image=reference_square_mask,
 
522
  feature_detector="ORB",
523
  )
524
  except ZeroDivisionError:
 
529
  print(f"Error calculating scaling factor: {e}")
530
 
531
  if scaling_factor is None or scaling_factor == 0:
532
+ scaling_factor = 1.0
533
+ print("Using default scaling factor of 1.0 due to calculation error")
534
  gc.collect()
535
  print("Scaling factor determined: {}".format(scaling_factor))
536
 
 
564
  objects_mask = remove_bg(shrunked_img)
565
  processed_size = objects_mask.shape[:2]
566
 
567
+ objects_mask = exclude_scaling_box(objects_mask, scaling_box_coords, orig_size, processed_size, expansion_factor=2)
568
  objects_mask = resize_img(objects_mask, (shrunked_img.shape[1], shrunked_img.shape[0]))
569
  del scaling_box_coords
570
  gc.collect()
 
773
  gr.Textbox(label="Scaling Factor (inches/pixel)")
774
  ],
775
  examples=[
776
+ ["./Test20.jpg", 0.075, "inches", "No", "No", 300.0, 200.0, "MyTool"],
777
+ ["./Test21.jpg", 0.075, "inches", "Yes", "Yes", 300.0, 200.0, "Tool2"]
778
  ]
779
  )
780
  iface.launch(share=True)