Update app.py

app.py

@@ -1,183 +1,140 @@
 import os
 import cv2
-import gradio as gr
 import numpy as np
 import torch
-import torch.nn as nn
 from torchvision import transforms
-import torch.nn.functional as F
 from PIL import Image
-import warnings
-warnings.filterwarnings("ignore")
+import gradio as gr
 
-# Initialize device
+# Set up device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Clone repository and setup model
+# Clone model repository if needed
 if not os.path.exists("DIS"):
-    os.system("git clone --depth 1 https://github.com/xuebinqin/DIS")
+    os.system("git clone https://github.com/xuebinqin/DIS")
     os.system("mv DIS/IS-Net/* .")
 
 # Import model components
-from data_loader_cache import normalize, im_reader, im_preprocess
 from models import ISNetDIS
+from data_loader_cache import normalize
 
-# Setup model directory
+# Create model directory
 os.makedirs("saved_models", exist_ok=True)
 if os.path.exists("isnet.pth"):
     os.rename("isnet.pth", "saved_models/isnet.pth")
 
-class GOSNormalize:
+# 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, image):
-        return normalize(image, self.mean, self.std)
-
-transform = transforms.Compose([GOSNormalize()])
-
-def load_image(im_path, hypar):
-    im = im_reader(im_path)
-    im, im_shp = im_preprocess(im, hypar["cache_size"])
-    im = torch.divide(im, 255.0)
-    shape = torch.from_numpy(np.array(im_shp))
-    return transform(im).unsqueeze(0), shape.unsqueeze(0)
-
-def build_model(hypar, device):
-    net = hypar["model"]
-    if hypar["model_digit"] == "half":
-        net.half()
-        for layer in net.modules():
-            if isinstance(layer, nn.BatchNorm2d):
-                layer.float()
-    
-    net.to(device)
     
-    model_path = os.path.join(hypar["model_path"], hypar["restore_model"])
-    if os.path.exists(model_path):
-        state_dict = torch.load(model_path, map_location=device)
-        net.load_state_dict(state_dict)
-    
-    net.eval()
-    return net
+    def __call__(self, img):
+        return normalize(img, self.mean, self.std)
 
-def predict(net, inputs_val, shapes_val, hypar, device):
-    with torch.no_grad():
-        inputs_val = inputs_val.type(torch.float16 if hypar["model_digit"] == "half" else torch.float32)
-        inputs_val = inputs_val.to(device)
-        ds_val = net(inputs_val)[0]
-        pred_val = ds_val[0][0,:,:,:]
-        
-        pred_val = F.interpolate(
-            pred_val.unsqueeze(0).unsqueeze(0),
-            size=(shapes_val[0][0], shapes_val[0][1]),
-            mode='bilinear',
-            align_corners=False
-        ).squeeze()
-        
-        pred_val = (pred_val - pred_val.min()) / (pred_val.max() - pred_val.min() + 1e-8)
-        
-        return (pred_val.cpu().numpy() * 255).astype(np.uint8)
+transform = transforms.Compose([ImageNormalizer()])
 
-# Model configuration
-hypar = {
+# Load and configure model
+model_config = {
     "model_path": "saved_models",
-    "restore_model": "isnet.pth",
-    "interm_sup": False,
-    "model_digit": "full",
-    "seed": 0,
-    "cache_size": [1024, 1024],
+    "model_file": "isnet.pth",
     "input_size": [1024, 1024],
-    "crop_size": [1024, 1024],
-    "model": ISNetDIS()
+    "device": device
 }
 
-# Initialize model
-net = build_model(hypar, device)
+model = ISNetDIS().to(device)
+if os.path.exists(f"{model_config['model_path']}/{model_config['model_file']}"):
+    model.load_state_dict(
+        torch.load(
+            f"{model_config['model_path']}/{model_config['model_file']}",
+            map_location=device
+        )
+    )
+model.eval()
 
-def process_image(image):
+def process_image(input_image):
+    """Process an image through the segmentation model"""
     try:
-        image_path = image if isinstance(image, str) else image.name
-        
-        # Verify image exists
-        if not os.path.exists(image_path):
-            raise FileNotFoundError(f"Image file not found: {image_path}")
+        # Convert Gradio input to usable image path
+        if hasattr(input_image, 'name'):
+            image_path = input_image.name
+        else:
+            image_path = input_image
         
-        image_tensor, orig_size = load_image(image_path, hypar)
-        mask = predict(net, image_tensor, orig_size, hypar, device)
+        # Read and preprocess image
+        img = cv2.imread(image_path)
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        img = torch.from_numpy(img).float().permute(2, 0, 1) / 255.0
+        img = transform(img).unsqueeze(0).to(device)
         
-        mask_img = Image.fromarray(mask).convert('L')
-        rgb_img = Image.open(image_path).convert("RGB")
+        # Get prediction
+        with torch.no_grad():
+            pred = model(img)[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)
         
-        rgba_img = rgb_img.copy()
-        rgba_img.putalpha(mask_img)
+        # 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 rgba_img, mask_img
+        return transparent_img, mask_img
+    
     except Exception as e:
         raise gr.Error(f"Error processing image: {str(e)}")
 
-# Interface setup
-title = "DIS Image Segmentation"
+# Gradio interface setup
+title = "Image Background Removal"
 description = """
-Highly Accurate Dichotomous Image Segmentation
-<br>GitHub: [xuebinqin/DIS](https://github.com/xuebinqin/DIS)
+Upload an image to automatically remove the background using DIS (Dichotomous Image Segmentation).
+<br>Model from: <a href="https://github.com/xuebinqin/DIS">xuebinqin/DIS</a>
 """
 
-# Prepare examples
+# Check for example images
 examples = []
-for example_file in ["robot.png", "ship.png"]:
-    if os.path.exists(example_file):
-        examples.append([example_file])
+for img_file in ["robot.png", "ship.png"]:
+    if os.path.exists(img_file):
+        examples.append([img_file])
 
-# Create Gradio interface
-with gr.Blocks(title=title) as app:
+# Create interface
+with gr.Blocks() as app:
     gr.Markdown(f"## {title}")
     gr.Markdown(description)
     
     with gr.Row():
-        with gr.Column():
-            input_image = gr.Image(
-                type="filepath",
-                label="Input Image",
-                height=400
-            )
-            submit_btn = gr.Button("Process", variant="primary")
-        
-        with gr.Column():
-            output_rgba = gr.Image(
-                label="Transparent Background",
-                type="pil",
-                height=400
-            )
-            output_mask = gr.Image(
-                label="Segmentation Mask",
-                type="pil",
-                height=400
-            )
+        input_col = gr.Column()
+        output_col = gr.Column()
+    
+    with input_col:
+        image_input = gr.Image(type="filepath", label="Upload Image")
+        submit_btn = gr.Button("Remove Background", variant="primary")
+    
+    with output_col:
+        transparent_output = gr.Image(label="Transparent Result", type="pil")
+        mask_output = gr.Image(label="Segmentation Mask", type="pil")
     
     if examples:
         gr.Examples(
             examples=examples,
-            inputs=input_image,
-            outputs=[output_rgba, output_mask],
+            inputs=image_input,
+            outputs=[transparent_output, mask_output],
             fn=process_image,
             cache_examples=True,
-            label="Example Images"
+            label="Try Example Images"
         )
     
     submit_btn.click(
         fn=process_image,
-        inputs=input_image,
-        outputs=[output_rgba, output_mask],
-        api_name="predict"
+        inputs=image_input,
+        outputs=[transparent_output, mask_output]
     )
 
-# Launch application
+# Launch the app
 if __name__ == "__main__":
     app.launch(
         server_name="0.0.0.0",
         server_port=7860,
-        show_error=True,
         share=False
     )