Update app.py

app.py

@@ -1,74 +1,103 @@
 import cv2
 import numpy as np
 import torch
+import torch.nn.functional as F
 from PIL import Image
 import gradio as gr
-from huggingface_hub import hf_hub_download
+from torchvision.transforms import ToTensor, ToPILImage
 
-# Load MODNet (PyTorch version)
-MODNET_REPO = "ZHTX/modnet"
-MODNET_FILE = "modnet_photographic_portrait_matting.ckpt"
+# Load MODNet (local weights)
+MODEL_URL = "https://drive.google.com/uc?export=download&id=1mcr7ALciuAsHCpLnrtG_eop5-EYhbCmz"
+MODEL_PATH = "modnet.pth"
 
-try:
-    model_path = hf_hub_download(repo_id=MODNET_REPO, filename=MODNET_FILE)
-    modnet = torch.hub.load('ZHTX/modnet', 'modnet', pretrained=False)
-    modnet.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
+def download_model():
+    import requests
+    import os
+    if not os.path.exists(MODEL_PATH):
+        print("Downloading MODNet weights...")
+        try:
+            response = requests.get(MODEL_URL, stream=True)
+            with open(MODEL_PATH, 'wb') as f:
+                for chunk in response.iter_content(chunk_size=1024):
+                    if chunk:
+                        f.write(chunk)
+            print("Download complete!")
+        except Exception as e:
+            print(f"Download failed: {e}")
+            return False
+    return True
+
+class MODNet(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.backbone = torch.hub.load('pytorch/vision:v0.10.0', 'mobilenet_v2', pretrained=True)
+        self.head = torch.nn.Sequential(
+            torch.nn.Conv2d(1280, 1, kernel_size=3, padding=1),
+            torch.nn.Sigmoid()
+        )
+        
+    def forward(self, x):
+        features = self.backbone.features(x)
+        return self.head(features)
+
+# Initialize model
+if download_model():
+    modnet = MODNet()
+    modnet.load_state_dict(torch.load(MODEL_PATH, map_location='cpu'))
     modnet.eval()
-except Exception as e:
-    print(f"Error loading MODNet: {e}")
+else:
     modnet = None
 
-def refine_with_modnet(input_image, bg_color="#FFFFFF", threshold=0.1):
-    """Refine alpha matte using MODNet"""
+def refine_edges(img, bg_color="#FFFFFF"):
+    """Refine edges using local MODNet"""
     if modnet is None:
-        raise gr.Error("MODNet model failed to load")
+        raise gr.Error("Model failed to load. Please check logs.")
     
-    # Convert input
-    img = np.array(input_image.convert("RGB"))
-    img = cv2.resize(img, (512, 512), interpolation=cv2.INTER_AREA)
-    img = torch.from_numpy(img).permute(2,0,1).unsqueeze(0).float() / 255.0
+    # Preprocess
+    img = img.convert("RGB")
+    img_tensor = ToTensor()(img).unsqueeze(0)
+    
+    # Resize to nearest multiple of 32
+    h, w = img_tensor.shape[2], img_tensor.shape[3]
+    new_h = h - h % 32
+    new_w = w - w % 32
+    img_tensor = F.interpolate(img_tensor, (new_h, new_w), mode='area')
     
     # Inference
     with torch.no_grad():
-        _, _, matte = modnet(img, True)
+        matte = modnet(img_tensor)
     
-    # Process output
+    # Post-process
+    matte = F.interpolate(matte, (h, w), mode='bilinear')
     matte = matte.squeeze().cpu().numpy()
     matte = (matte * 255).astype(np.uint8)
-    matte = cv2.threshold(matte, int(threshold*255), 255, cv2.THRESH_BINARY)[1]
     
     # Composite with background
     bg_color = bg_color.lstrip('#')
     bg_rgb = tuple(int(bg_color[i:i+2], 16) for i in (0, 2, 4))
-    bg = Image.new("RGB", input_image.size, bg_rgb)
+    bg = Image.new("RGB", img.size, bg_rgb)
     
-    # Apply refined matte
-    refined = Image.fromarray(matte).resize(input_image.size)
-    result = Image.composite(input_image, bg, refined)
+    # Create mask
+    mask = Image.fromarray(matte).convert("L")
+    result = Image.composite(img, bg, mask)
     
-    return refined, result
+    return mask, result
 
 # Gradio Interface
-with gr.Blocks(title="🔍 MODNet Edge Refiner") as demo:
-    gr.Markdown("""
-    ## 🔍 MODNet Professional Edge Refinement
-    Uses AI to perfectly refine hair/fur edges from trimmed images
-    """)
-    
+with gr.Blocks() as demo:
+    gr.Markdown("## ✨ Professional Edge Refiner")
     with gr.Row():
         with gr.Column():
-            input_img = gr.Image(type="pil", label="Trimmed Input")
-            bg_color = gr.ColorPicker("#FFFFFF", label="Background Color")
-            threshold = gr.Slider(0, 100, 10, label="Edge Threshold")
+            input_img = gr.Image(type="pil", label="Input Image")
+            bg_color = gr.ColorPicker("#FFFFFF", label="Preview Background")
             process_btn = gr.Button("Refine Edges", variant="primary")
-        
         with gr.Column():
             matte_output = gr.Image(label="Refined Alpha Matte", type="pil")
             final_output = gr.Image(label="Composited Result", type="pil")
     
     process_btn.click(
-        fn=refine_with_modnet,
-        inputs=[input_img, bg_color, threshold],
+        fn=refine_edges,
+        inputs=[input_img, bg_color],
         outputs=[matte_output, final_output]
     )