Spaces:
Running
Running
Update app.py
Browse files
app.py
CHANGED
@@ -43,7 +43,7 @@ class DrawerNotDetectedError(Exception):
|
|
43 |
pass
|
44 |
|
45 |
class ReferenceBoxNotDetectedError(Exception):
|
46 |
-
"""Raised when the
|
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, "
|
73 |
if not os.path.exists(reference_model_path):
|
74 |
print("Caching YOLO reference model to", reference_model_path)
|
75 |
-
shutil.copy("
|
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, "
|
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.
|
159 |
if not res or len(res) == 0 or len(res[0].boxes) == 0:
|
160 |
-
raise ReferenceBoxNotDetectedError("Reference
|
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 |
-
|
291 |
needed_center_distance = circle_diameter + min_gap
|
292 |
radius = circle_diameter / 2.0
|
293 |
-
|
294 |
-
|
295 |
-
|
296 |
-
|
297 |
-
|
298 |
-
|
299 |
-
|
300 |
-
|
301 |
-
|
302 |
-
|
303 |
-
|
304 |
-
|
305 |
-
|
306 |
-
|
307 |
-
|
308 |
-
|
309 |
-
|
310 |
-
|
311 |
-
|
312 |
-
|
313 |
-
|
314 |
-
|
315 |
-
|
316 |
-
|
317 |
-
|
318 |
-
|
319 |
-
|
320 |
-
|
321 |
-
|
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=
|
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 |
-
|
424 |
-
|
425 |
-
|
426 |
-
|
427 |
-
|
|
|
|
|
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
|
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
|
541 |
-
print("Using default scaling factor of 0
|
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=
|
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",
|
785 |
-
["./Test21.jpg", 0.075, "inches", "Yes", "Yes",
|
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)
|