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"))
+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))
+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()
@@ -121,14 +121,14 @@ def unload_and_reload_models():
     )
     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_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_birefnet
+    u2net_global = new_u2net
     print("Models reloaded in {:.2f} seconds".format(time.time() - start_time))
 
 # ---------------------
@@ -160,24 +160,24 @@ def detect_reference_square(img: np.ndarray):
     )
 
 # 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
+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:
@@ -492,7 +492,7 @@ def predict(
     # ---------------------
     t = time.time()
     reference_obj_img = make_square(reference_obj_img)
-    reference_square_mask = remove_bg(reference_obj_img)
+    reference_square_mask = remove_bg_u2netp(reference_obj_img)
     print("Reference image processing completed in {:.2f} seconds".format(time.time() - t))
 
     t = time.time()