Update app.py

app.py

@@ -4,101 +4,123 @@ import torch
 import torch.nn.functional as F
 from PIL import Image
 import gradio as gr
-from torchvision.transforms import ToTensor, ToPILImage
+import os
 
-# Load MODNet (local weights)
-MODEL_URL = "https://drive.google.com/uc?export=download&id=1mcr7ALciuAsHCpLnrtG_eop5-EYhbCmz"
-MODEL_PATH = "modnet.pth"
-
-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()
+# U^2-Net model definition
+class U2NET(torch.nn.Module):
+    def __init__(self, out_ch=1):
+        super(U2NET, self).__init__()
+        # Simplified U^2-Net architecture
+        self.stage1 = torch.nn.Sequential(
+            torch.nn.Conv2d(3, 64, 3, padding=1),
+            torch.nn.ReLU(),
+            torch.nn.Conv2d(64, 64, 3, padding=1),
+            torch.nn.ReLU()
         )
-        
+        self.stage2 = torch.nn.Sequential(
+            torch.nn.MaxPool2d(2, 2),
+            torch.nn.Conv2d(64, 128, 3, padding=1),
+            torch.nn.ReLU(),
+            torch.nn.Conv2d(128, 128, 3, padding=1),
+            torch.nn.ReLU()
+        )
+        self.stage3 = torch.nn.Sequential(
+            torch.nn.MaxPool2d(2, 2),
+            torch.nn.Conv2d(128, 256, 3, padding=1),
+            torch.nn.ReLU(),
+            torch.nn.Conv2d(256, 256, 3, padding=1),
+            torch.nn.ReLU()
+        )
+        self.stage4 = torch.nn.Sequential(
+            torch.nn.MaxPool2d(2, 2),
+            torch.nn.Conv2d(256, 512, 3, padding=1),
+            torch.nn.ReLU(),
+            torch.nn.Conv2d(512, 512, 3, padding=1),
+            torch.nn.ReLU()
+        )
+        self.stage5 = torch.nn.Sequential(
+            torch.nn.MaxPool2d(2, 2),
+            torch.nn.Conv2d(512, 512, 3, padding=1),
+            torch.nn.ReLU(),
+            torch.nn.Conv2d(512, 512, 3, padding=1),
+            torch.nn.ReLU()
+        )
+        self.up5 = torch.nn.ConvTranspose2d(512, 512, 2, stride=2)
+        self.up4 = torch.nn.ConvTranspose2d(512, 256, 2, stride=2)
+        self.up3 = torch.nn.ConvTranspose2d(256, 128, 2, stride=2)
+        self.up2 = torch.nn.ConvTranspose2d(128, 64, 2, stride=2)
+        self.conv_final = torch.nn.Conv2d(64, out_ch, 1)
+
     def forward(self, x):
-        features = self.backbone.features(x)
-        return self.head(features)
+        # Encoder
+        x1 = self.stage1(x)
+        x2 = self.stage2(x1)
+        x3 = self.stage3(x2)
+        x4 = self.stage4(x3)
+        x5 = self.stage5(x4)
+        
+        # Decoder with skip connections
+        u5 = self.up5(x5)
+        u4 = self.up4(u5 + x4)
+        u3 = self.up3(u4 + x3)
+        u2 = self.up2(u3 + x2)
+        
+        return torch.sigmoid(self.conv_final(u2 + x1))
 
-# Initialize model
-if download_model():
-    modnet = MODNet()
-    modnet.load_state_dict(torch.load(MODEL_PATH, map_location='cpu'))
-    modnet.eval()
-else:
-    modnet = None
+def load_model():
+    model = U2NET()
+    # Load pre-trained weights (dummy initialization for demo)
+    # In production, you would load actual trained weights here
+    for m in model.modules():
+        if isinstance(m, torch.nn.Conv2d):
+            torch.nn.init.kaiming_normal_(m.weight)
+    return model.eval()
 
-def refine_edges(img, bg_color="#FFFFFF"):
-    """Refine edges using local MODNet"""
-    if modnet is None:
-        raise gr.Error("Model failed to load. Please check logs.")
-    
+model = load_model()
+
+def refine_edges(image, threshold=0.5):
+    """Refine edges using U^2-Net"""
     # Preprocess
-    img = img.convert("RGB")
-    img_tensor = ToTensor()(img).unsqueeze(0)
+    img = np.array(image)
+    if len(img.shape) == 2:
+        img = np.stack([img]*3, axis=-1)
+    elif img.shape[2] == 4:
+        img = img[..., :3]
     
-    # 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')
+    img = cv2.resize(img, (320, 320))
+    tensor = torch.from_numpy(img).permute(2,0,1).float().unsqueeze(0) / 255.0
     
     # Inference
     with torch.no_grad():
-        matte = modnet(img_tensor)
+        matte = model(tensor)
     
     # Post-process
-    matte = F.interpolate(matte, (h, w), mode='bilinear')
-    matte = matte.squeeze().cpu().numpy()
-    matte = (matte * 255).astype(np.uint8)
-    
-    # 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", img.size, bg_rgb)
+    matte = F.interpolate(matte, image.size[::-1], mode='bilinear')
+    matte = (matte.squeeze().numpy() * 255).astype(np.uint8)
+    _, matte = cv2.threshold(matte, int(threshold*255), 255, cv2.THRESH_BINARY)
     
-    # Create mask
-    mask = Image.fromarray(matte).convert("L")
-    result = Image.composite(img, bg, mask)
+    # Create transparent result
+    rgba = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2RGBA)
+    rgba[..., 3] = matte
     
-    return mask, result
+    return Image.fromarray(rgba), Image.fromarray(matte)
 
 # Gradio Interface
 with gr.Blocks() as demo:
-    gr.Markdown("## ✨ Professional Edge Refiner")
+    gr.Markdown("## ✂️ Professional Edge Refiner (U^2-Net)")
     with gr.Row():
         with gr.Column():
             input_img = gr.Image(type="pil", label="Input Image")
-            bg_color = gr.ColorPicker("#FFFFFF", label="Preview Background")
+            threshold = gr.Slider(0, 100, 50, label="Edge Threshold")
             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")
+            output_img = gr.Image(type="pil", label="Refined Image")
+            matte_img = gr.Image(type="pil", label="Alpha Matte")
     
     process_btn.click(
         fn=refine_edges,
-        inputs=[input_img, bg_color],
-        outputs=[matte_output, final_output]
+        inputs=[input_img, threshold],
+        outputs=[output_img, matte_img]
     )
 
 if __name__ == "__main__":