Update app.py

app.py

@@ -159,27 +159,23 @@ def detect_reference_square(img: np.ndarray):
         res[0].cpu().boxes.xyxy[0]
     )
 
-# Use U2NETP for reference background removal.
-def remove_bg_u2netp(image: np.ndarray) -> np.ndarray:
+# 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)
-    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")
+    input_images = transform_image_global(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.
+        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.
 def remove_bg(image: np.ndarray) -> np.ndarray:
     t = time.time()
     image_pil = Image.fromarray(image)
@@ -191,7 +187,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 background removal completed in {:.2f} seconds".format(time.time() - t))
+    print("BiRefNet (object) background removal completed in {:.2f} seconds".format(time.time() - t))
     return result
 
 def make_square(img: np.ndarray):
@@ -473,7 +469,6 @@ 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
@@ -495,9 +490,9 @@ def predict(
     # ---------------------
     t = time.time()
     reference_obj_img = make_square(reference_obj_img)
-    reference_square_mask = remove_bg_u2netp(reference_obj_img)
+    # Use BiRefNet for reference background removal instead of U2NETP:
+    reference_square_mask = remove_bg_reference(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))
@@ -570,7 +565,6 @@ 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")
 
     # ---------------------
@@ -579,16 +573,12 @@ 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))
@@ -633,14 +623,8 @@ 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
@@ -660,27 +644,8 @@ 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,
@@ -691,14 +656,11 @@ 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
@@ -734,5 +696,4 @@ if __name__ == "__main__":
             ["./Test21.jpg", 0.075, "inches", "Yes", "Yes", 300.0, 200.0, "Tool2"]
         ]
     )
-    iface.launch(share=True)
-    
+    iface.launch(share=True)