Update app.py

app.py

@@ -1,34 +1,34 @@
 import os
 import cv2
 import gradio as gr
-from PIL import Image
 import numpy as np
 import torch
-from torch.autograd import Variable
+import torch.nn as nn
 from torchvision import transforms
 import torch.nn.functional as F
+from PIL import Image
 import warnings
 warnings.filterwarnings("ignore")
 
 # Initialize device
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Clone repository if not exists
+# Clone repository and setup model
 if not os.path.exists("DIS"):
-    os.system("git clone https://github.com/xuebinqin/DIS")
+    os.system("git clone --depth 1 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 data_loader_cache import normalize, im_reader, im_preprocess
 from models import ISNetDIS
 
 # Setup model directory
 os.makedirs("saved_models", exist_ok=True)
 if os.path.exists("isnet.pth"):
-    os.system("mv isnet.pth saved_models/")
+    os.rename("isnet.pth", "saved_models/isnet.pth")
 
 class GOSNormalize:
-    def __init__(self, mean=[0.5,0.5,0.5], std=[1.0,1.0,1.0]):
+    def __init__(self, mean=[0.5, 0.5, 0.5], std=[1.0, 1.0, 1.0]):
         self.mean = mean
         self.std = std
 
@@ -54,32 +54,31 @@ def build_model(hypar, device):
     
     net.to(device)
     
-    if hypar["restore_model"]:
-        net.load_state_dict(torch.load(
-            os.path.join(hypar["model_path"], hypar["restore_model"]),
-            map_location=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 predict(net, inputs_val, shapes_val, hypar, device):
-    net.eval()
-    inputs_val = inputs_val.type(torch.FloatTensor if hypar["model_digit"] == "full" else torch.HalfTensor)
-    inputs_val_v = Variable(inputs_val, requires_grad=False).to(device)
-    ds_val = net(inputs_val_v)[0]
-    pred_val = ds_val[0][0,:,:,:]
-    
-    pred_val = torch.squeeze(F.interpolate(
-        torch.unsqueeze(pred_val, 0),
-        size=(shapes_val[0][0], shapes_val[0][1]),
-        mode='bilinear'
-    ))
-    
-    pred_val = (pred_val - pred_val.min()) / (pred_val.max() - pred_val.min())
-    
-    if device == 'cuda':
-        torch.cuda.empty_cache()
-    return (pred_val.detach().cpu().numpy() * 255).astype(np.uint8)
+    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)
 
 # Model configuration
 hypar = {
@@ -99,11 +98,12 @@ net = build_model(hypar, device)
 
 def process_image(image):
     try:
-        if isinstance(image, str):
-            image_path = image
-        else:
-            image_path = image.name
-            
+        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}")
+        
         image_tensor, orig_size = load_image(image_path, hypar)
         mask = predict(net, image_tensor, orig_size, hypar, device)
         
@@ -118,27 +118,43 @@ def process_image(image):
         raise gr.Error(f"Error processing image: {str(e)}")
 
 # Interface setup
-title = "Image Segmentation Demo"
-description = "Upload an image to extract its foreground"
+title = "DIS Image Segmentation"
+description = """
+Highly Accurate Dichotomous Image Segmentation
+<br>GitHub: [xuebinqin/DIS](https://github.com/xuebinqin/DIS)
+"""
 
+# Prepare examples
 examples = []
-if os.path.exists("robot.png"):
-    examples.append(["robot.png"])
-if os.path.exists("ship.png"):
-    examples.append(["ship.png"])
+for example_file in ["robot.png", "ship.png"]:
+    if os.path.exists(example_file):
+        examples.append([example_file])
 
-with gr.Blocks() as app:
+# Create Gradio interface
+with gr.Blocks(title=title) as app:
     gr.Markdown(f"## {title}")
     gr.Markdown(description)
     
     with gr.Row():
         with gr.Column():
-            input_image = gr.Image(type="filepath", label="Input Image")
-            submit_btn = gr.Button("Process")
+            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")
-            output_mask = gr.Image(label="Segmentation Mask", type="pil")
+            output_rgba = gr.Image(
+                label="Transparent Background",
+                type="pil",
+                height=400
+            )
+            output_mask = gr.Image(
+                label="Segmentation Mask",
+                type="pil",
+                height=400
+            )
     
     if examples:
         gr.Examples(
@@ -146,18 +162,22 @@ with gr.Blocks() as app:
             inputs=input_image,
             outputs=[output_rgba, output_mask],
             fn=process_image,
-            cache_examples=True
+            cache_examples=True,
+            label="Example Images"
         )
     
     submit_btn.click(
         fn=process_image,
         inputs=input_image,
-        outputs=[output_rgba, output_mask]
+        outputs=[output_rgba, output_mask],
+        api_name="predict"
     )
 
+# Launch application
 if __name__ == "__main__":
     app.launch(
         server_name="0.0.0.0",
         server_port=7860,
-        show_error=True
+        show_error=True,
+        share=False
     )