Update app.py

app.py

@@ -1,6 +1,124 @@
-torch>=1.10.0
-torchvision>=0.11.0
-opencv-python>=4.5.0
-numpy>=1.21.0
-Pillow>=9.0.0
-gradio>=3.0
+import os
+import cv2
+import numpy as np
+import torch
+from torchvision import transforms
+from PIL import Image
+import gradio as gr
+
+# Set up device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Clone model repository if needed
+if not os.path.exists("DIS"):
+    os.system("git clone https://github.com/xuebinqin/DIS")
+    os.system("mv DIS/IS-Net/* .")
+
+# Create model directory and move weights
+os.makedirs("saved_models", exist_ok=True)
+if os.path.exists("isnet.pth"):
+    os.rename("isnet.pth", "saved_models/isnet.pth")
+
+# Custom normalize function to replace skimage dependency
+def normalize(image, mean, std):
+    """Normalize image with mean and std"""
+    if isinstance(mean, (int, float)):
+        mean = [mean] * image.shape[0]
+    if isinstance(std, (int, float)):
+        std = [std] * image.shape[0]
+    
+    image = image.clone()
+    for t, m, s in zip(image, mean, std):
+        t.sub_(m).div_(s)
+    return image
+
+# Define image preprocessing
+class ImageNormalizer:
+    def __init__(self, mean=[0.5, 0.5, 0.5], std=[1.0, 1.0, 1.0]):
+        self.mean = mean
+        self.std = std
+    
+    def __call__(self, img):
+        return normalize(img, self.mean, self.std)
+
+transform = transforms.Compose([ImageNormalizer()])
+
+# Load model
+from models import ISNetDIS
+
+model = ISNetDIS().to(device)
+model_path = "saved_models/isnet.pth"
+if os.path.exists(model_path):
+    model.load_state_dict(torch.load(model_path, map_location=device))
+model.eval()
+
+def process_image(input_image):
+    """Process an image through the segmentation model"""
+    try:
+        # Convert Gradio input to usable image
+        if isinstance(input_image, str):
+            image_path = input_image
+        else:
+            image_path = input_image.name
+        
+        # Read image with OpenCV (replaces skimage)
+        img = cv2.imread(image_path)
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        
+        # Convert to tensor and normalize
+        img_tensor = torch.from_numpy(img).float().permute(2, 0, 1) / 255.0
+        img_tensor = transform(img_tensor).unsqueeze(0).to(device)
+        
+        # Get prediction
+        with torch.no_grad():
+            pred = model(img_tensor)[0][0]
+            pred = torch.sigmoid(pred[0])
+            pred = (pred - pred.min()) / (pred.max() - pred.min() + 1e-8)
+            mask = (pred.cpu().numpy() * 255).astype(np.uint8)
+        
+        # Create output images
+        original_img = Image.open(image_path).convert("RGB")
+        mask_img = Image.fromarray(mask).convert("L")
+        transparent_img = original_img.copy()
+        transparent_img.putalpha(mask_img)
+        
+        return transparent_img, mask_img
+    
+    except Exception as e:
+        raise gr.Error(f"Error processing image: {str(e)}")
+
+# Gradio interface
+title = "Image Background Removal"
+description = "Upload an image to automatically remove the background"
+
+with gr.Blocks() as app:
+    gr.Markdown(f"## {title}")
+    gr.Markdown(description)
+    
+    with gr.Row():
+        with gr.Column():
+            image_input = gr.Image(type="filepath", label="Input Image")
+            submit_btn = gr.Button("Process", variant="primary")
+        
+        with gr.Column():
+            transparent_output = gr.Image(label="Result with Transparency", type="pil")
+            mask_output = gr.Image(label="Segmentation Mask", type="pil")
+    
+    # Add examples if files exist
+    example_files = [f for f in ["robot.png", "ship.png"] if os.path.exists(f)]
+    if example_files:
+        gr.Examples(
+            examples=[[f] for f in example_files],
+            inputs=image_input,
+            outputs=[transparent_output, mask_output],
+            fn=process_image,
+            cache_examples=True
+        )
+    
+    submit_btn.click(
+        fn=process_image,
+        inputs=image_input,
+        outputs=[transparent_output, mask_output]
+    )
+
+app.launch(server_name="0.0.0.0", server_port=7860)