Update app.py

app.py

@@ -71,18 +71,18 @@ if not os.path.exists(reference_model_path):
 reference_detector_global = YOLO(reference_model_path)
 print("YOLO reference model loaded in {:.2f} seconds".format(time.time() - start_time))
 
-print("Loading U²-Net model for reference background removal (U2NETP)...")
-start_time = time.time()
-u2net_model_path = os.path.join(CACHE_DIR, "u2netp.pth")
-if not os.path.exists(u2net_model_path):
-    print("Caching U²-Net model to", u2net_model_path)
-    shutil.copy("u2netp.pth", u2net_model_path)
-u2net_global = U2NETP(3, 1)
-u2net_global.load_state_dict(torch.load(u2net_model_path, map_location="cpu"))
-device = "cpu"
-u2net_global.to(device)
-u2net_global.eval()
-print("U²-Net model loaded in {:.2f} seconds".format(time.time() - start_time))
+# print("Loading U²-Net model for reference background removal (U2NETP)...")
+# start_time = time.time()
+# u2net_model_path = os.path.join(CACHE_DIR, "u2netp.pth")
+# if not os.path.exists(u2net_model_path):
+#     print("Caching U²-Net model to", u2net_model_path)
+#     shutil.copy("u2netp.pth", u2net_model_path)
+# u2net_global = U2NETP(3, 1)
+# u2net_global.load_state_dict(torch.load(u2net_model_path, map_location="cpu"))
+# device = "cpu"
+# u2net_global.to(device)
+# u2net_global.eval()
+# print("U²-Net model loaded in {:.2f} seconds".format(time.time() - start_time))
 
 print("Loading BiRefNet model...")
 start_time = time.time()
@@ -119,16 +119,16 @@ def unload_and_reload_models():
     new_birefnet = AutoModelForImageSegmentation.from_pretrained(
         "zhengpeng7/BiRefNet", trust_remote_code=True, cache_dir=CACHE_DIR
     )
-    new_birefnet.to(device)
-    new_birefnet.eval()
-    new_u2net = U2NETP(3, 1)
-    new_u2net.load_state_dict(torch.load(os.path.join(CACHE_DIR, "u2netp.pth"), map_location="cpu"))
-    new_u2net.to(device)
-    new_u2net.eval()
+    # new_birefnet.to(device)
+    # new_birefnet.eval()
+    # new_u2net = U2NETP(3, 1)
+    # new_u2net.load_state_dict(torch.load(os.path.join(CACHE_DIR, "u2netp.pth"), map_location="cpu"))
+    # new_u2net.to(device)
+    # new_u2net.eval()
     drawer_detector_global = new_drawer_detector
     reference_detector_global = new_reference_detector
     birefnet_global = new_birefnet
-    u2net_global = new_u2net
+    u2net_global = new_birefnet
     print("Models reloaded in {:.2f} seconds".format(time.time() - start_time))
 
 # ---------------------
@@ -159,23 +159,27 @@ def detect_reference_square(img: np.ndarray):
         res[0].cpu().boxes.xyxy[0]
     )
 
-# For reference background removal, we now use BiRefNet.
-def remove_bg_reference(image: np.ndarray) -> np.ndarray:
-    # Use the same BiRefNet method as for the main object.
-    t = time.time()
-    image_pil = Image.fromarray(image)
-    input_images = transform_image_global(image_pil).unsqueeze(0).to("cpu")
-    with torch.no_grad():
-        preds = birefnet_global(input_images)[-1].sigmoid().cpu()
-    pred = preds[0].squeeze()
-    pred_pil = transforms.ToPILImage()(pred)
-    scale_ratio = 1024 / max(image_pil.size)
-    scaled_size = (int(image_pil.size[0] * scale_ratio), int(image_pil.size[1] * scale_ratio))
-    result = np.array(pred_pil.resize(scaled_size))
-    print("BiRefNet (reference) background removal completed in {:.2f} seconds".format(time.time() - t))
-    return result
-
-# The main background removal for objects still uses BiRefNet.
+# Use U2NETP for reference background removal.
+# def remove_bg_u2netp(image: np.ndarray) -> np.ndarray:
+#     t = time.time()
+#     image_pil = Image.fromarray(image)
+#     transform_u2netp = transforms.Compose([
+#         transforms.Resize((320, 320)),
+#         transforms.ToTensor(),
+#         transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+#     ])
+#     input_tensor = transform_u2netp(image_pil).unsqueeze(0).to("cpu")
+#     with torch.no_grad():
+#         outputs = u2net_global(input_tensor)
+#     pred = outputs[0]
+#     pred = (pred - pred.min()) / (pred.max() - pred.min() + 1e-8)
+#     pred_np = pred.squeeze().cpu().numpy()
+#     pred_np = cv2.resize(pred_np, (image_pil.width, image_pil.height))
+#     pred_np = (pred_np * 255).astype(np.uint8)
+#     print("U2NETP background removal completed in {:.2f} seconds".format(time.time() - t))
+#     return pred_np
+
+# Use BiRefNet for main object background removal.
 def remove_bg(image: np.ndarray) -> np.ndarray:
     t = time.time()
     image_pil = Image.fromarray(image)
@@ -187,7 +191,7 @@ def remove_bg(image: np.ndarray) -> np.ndarray:
     scale_ratio = 1024 / max(image_pil.size)
     scaled_size = (int(image_pil.size[0] * scale_ratio), int(image_pil.size[1] * scale_ratio))
     result = np.array(pred_pil.resize(scaled_size))
-    print("BiRefNet (object) background removal completed in {:.2f} seconds".format(time.time() - t))
+    print("BiRefNet background removal completed in {:.2f} seconds".format(time.time() - t))
     return result
 
 def make_square(img: np.ndarray):
@@ -469,6 +473,7 @@ def predict(
         print("Drawer detection completed in {:.2f} seconds".format(time.time() - t))
     except DrawerNotDetectedError as e:
         return None, None, None, None, f"Error: {str(e)}"
+    # Ensure that shrunked_img is defined only after successful detection.
     t = time.time()
     shrunked_img = make_square(shrink_bbox(drawer_img, 0.90))
     del drawer_img
@@ -490,9 +495,9 @@ def predict(
     # ---------------------
     t = time.time()
     reference_obj_img = make_square(reference_obj_img)
-    # Use BiRefNet for reference background removal instead of U2NETP:
-    reference_square_mask = remove_bg_reference(reference_obj_img)
+    reference_square_mask = remove_bg(reference_obj_img)
     print("Reference image processing completed in {:.2f} seconds".format(time.time() - t))
+
     t = time.time()
     try:
         cv2.imwrite("mask.jpg", cv2.cvtColor(reference_obj_img, cv2.COLOR_RGB2GRAY))
@@ -565,6 +570,7 @@ def predict(
     del objects_mask
     gc.collect()
     print("Mask dilation completed in {:.2f} seconds".format(time.time() - t))
+
     Image.fromarray(dilated_mask).save("./outputs/scaled_mask_new.jpg")
 
     # ---------------------
@@ -573,12 +579,16 @@ def predict(
     t = time.time()
     outlines, contours = extract_outlines(dilated_mask)
     print("Outline extraction completed in {:.2f} seconds".format(time.time() - t))
+
     output_img = shrunked_img.copy()
     del shrunked_img
     gc.collect()
+
     t = time.time()
     use_finger_clearance = True if finger_clearance.lower() == "yes" else False
-    doc, final_polygons_inch = save_dxf_spline(contours, scaling_factor, processed_size[0], finger_clearance=use_finger_clearance)
+    doc, final_polygons_inch = save_dxf_spline(
+        contours, scaling_factor, processed_size[0], finger_clearance=use_finger_clearance
+    )
     del contours
     gc.collect()
     print("DXF generation completed in {:.2f} seconds".format(time.time() - t))
@@ -623,8 +633,14 @@ def predict(
         text_x = (inner_min_x + inner_max_x) / 2.0  
         text_height_dxf = 0.5  
         text_y_dxf = inner_min_y - 0.125 - text_height_dxf  
-        text_entity = msp.add_text(annotation_text.strip(),
-                                   dxfattribs={"height": text_height_dxf, "layer": "ANNOTATION", "style": "Bold"})
+        text_entity = msp.add_text(
+            annotation_text.strip(),
+            dxfattribs={
+                "height": text_height_dxf,
+                "layer": "ANNOTATION",
+                "style": "Bold"
+            }
+        )
         text_entity.dxf.insert = (text_x, text_y_dxf)
 
     # Save the DXF
@@ -644,8 +660,27 @@ def predict(
         text_y_in = inner_min_y - 0.125 - text_height_cv  
         text_y_img = int(processed_size[0] - (text_y_in / scaling_factor))
         org = (text_x_img - int(len(annotation_text.strip()) * 6), text_y_img)
-        cv2.putText(output_img, annotation_text.strip(), org, cv2.FONT_HERSHEY_SIMPLEX, 1.3, (0, 0, 255), 3, cv2.LINE_AA)
-        cv2.putText(new_outlines, annotation_text.strip(), org, cv2.FONT_HERSHEY_SIMPLEX, 1.3, (0, 0, 255), 3, cv2.LINE_AA)
+
+        cv2.putText(
+            output_img,
+            annotation_text.strip(),
+            org,
+            cv2.FONT_HERSHEY_SIMPLEX,
+            1.3,
+            (0, 0, 255),
+            3,
+            cv2.LINE_AA
+        )
+        cv2.putText(
+            new_outlines,
+            annotation_text.strip(),
+            org,
+            cv2.FONT_HERSHEY_SIMPLEX,
+            1.3,
+            (0, 0, 255),
+            3,
+            cv2.LINE_AA
+        )
 
     # Restore brightness for display purposes:
     # Since we reduced brightness by 0.5 during preprocessing,
@@ -656,11 +691,14 @@ def predict(
 
     outlines_color = cv2.cvtColor(new_outlines, cv2.COLOR_BGR2RGB)
     print("Total prediction time: {:.2f} seconds".format(time.time() - overall_start))
-    return (cv2.cvtColor(output_img, cv2.COLOR_BGR2RGB),
-            outlines_color,
-            dxf_filepath,
-            dilated_mask,
-            str(scaling_factor))
+
+    return (
+        cv2.cvtColor(output_img, cv2.COLOR_BGR2RGB),
+        outlines_color,
+        dxf_filepath,
+        dilated_mask,
+        str(scaling_factor)
+    )
 
 # ---------------------
 # Gradio Interface