Update app.py

app.py

@@ -1,41 +1,76 @@
-import gradio as gr
-import numpy as np
 import cv2
+import numpy as np
+import torch
 from PIL import Image
+import gradio as gr
+from huggingface_hub import hf_hub_download
 
-def smooth_edges(image: Image.Image):
-    # Convert to numpy array
-    img_np = np.array(image.convert("RGB"))
-
-    # Create a grayscale version to detect edges
-    gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
-
-    # Detect edges using Canny
-    edges = cv2.Canny(gray, threshold1=50, threshold2=150)
-
-    # Dilate edges a bit to make the mask thicker
-    kernel = np.ones((3, 3), np.uint8)
-    dilated = cv2.dilate(edges, kernel, iterations=1)
-
-    # Create a mask with blur
-    mask = cv2.GaussianBlur(dilated, (11, 11), 0)
-
-    # Convert mask to 3-channel
-    mask_3c = cv2.cvtColor(mask, cv2.COLOR_GRAY2RGB) / 255.0
+# Load MODNet (PyTorch version)
+MODNET_REPO = "ZHTX/modnet"
+MODNET_FILE = "modnet_photographic_portrait_matting.ckpt"
 
-    # Apply soft blending of edges
-    softened = (img_np * (1 - mask_3c) + cv2.GaussianBlur(img_np, (15, 15), 0) * mask_3c).astype(np.uint8)
+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')))
+    modnet.eval()
+except Exception as e:
+    print(f"Error loading MODNet: {e}")
+    modnet = None
 
-    return Image.fromarray(softened)
+def refine_with_modnet(input_image, bg_color="#FFFFFF", threshold=0.1):
+    """Refine alpha matte using MODNet"""
+    if modnet is None:
+        raise gr.Error("MODNet model failed to load")
+    
+    # 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
+    
+    # Inference
+    with torch.no_grad():
+        _, _, matte = modnet(img, True)
+    
+    # Process output
+    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)
+    
+    # Apply refined matte
+    refined = Image.fromarray(matte).resize(input_image.size)
+    result = Image.composite(input_image, bg, refined)
+    
+    return refined, result
 
-# Interface
-iface = gr.Interface(
-    fn=smooth_edges,
-    inputs=gr.Image(type="pil", label="Upload Image (even non-transparent!)"),
-    outputs=gr.Image(type="pil", label="Softened Edges"),
-    title="Edge Smoother v2",
-    description="Smooth the edges of your subject — works for normal and background-removed images."
-)
+# 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.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")
+            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],
+        outputs=[matte_output, final_output]
+    )
 
 if __name__ == "__main__":
-    iface.launch()
+    demo.launch()