Update app.py

app.py

@@ -1,51 +1,41 @@
+import os
 import cv2
 import gradio as gr
-import os
 from PIL import Image
 import numpy as np
 import torch
 from torch.autograd import Variable
 from torchvision import transforms
 import torch.nn.functional as F
-import gdown
-import matplotlib.pyplot as plt
 import warnings
 warnings.filterwarnings("ignore")
 
-# Clean up any previous runs
-if os.path.exists("DIS"):
-    os.system("rm -rf DIS")
+# Initialize device
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
-# Clone and setup the model
-os.system("git clone https://github.com/xuebinqin/DIS")
-os.system("mv DIS/IS-Net/* .")
+# Clone repository if not exists
+if not os.path.exists("DIS"):
+    os.system("git clone https://github.com/xuebinqin/DIS")
+    os.system("mv DIS/IS-Net/* .")
 
-# Project imports
+# Import model components
 from data_loader_cache import normalize, im_reader, im_preprocess 
-from models import *
+from models import ISNetDIS
 
-# Device configuration
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
+# Setup model directory
+os.makedirs("saved_models", exist_ok=True)
+if os.path.exists("isnet.pth"):
+    os.system("mv isnet.pth saved_models/")
 
-# Setup model directory and weights
-if not os.path.exists("saved_models"):
-    os.makedirs("saved_models", exist_ok=True)
-    if os.path.exists("isnet.pth"):
-        os.system("mv isnet.pth saved_models/")
-
-class GOSNormalize(object):
-    '''
-    Normalize the Image using torch.transforms
-    '''
-    def __init__(self, mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]):
+class GOSNormalize:
+    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):
-        image = normalize(image, self.mean, self.std)
-        return image
+        return normalize(image, self.mean, self.std)
 
-transform = transforms.Compose([GOSNormalize([0.5,0.5,0.5],[1.0,1.0,1.0])])
+transform = transforms.Compose([GOSNormalize()])
 
 def load_image(im_path, hypar):
     im = im_reader(im_path)
@@ -55,55 +45,45 @@ def load_image(im_path, hypar):
     return transform(im).unsqueeze(0), shape.unsqueeze(0)
 
 def build_model(hypar, device):
-    net = hypar["model"] # GOSNETINC(3,1)
-
-    # Convert to half precision
+    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)
-
-    if hypar["restore_model"] != "":
+    
+    if hypar["restore_model"]:
         net.load_state_dict(torch.load(
-            hypar["model_path"]+"/"+hypar["restore_model"], 
+            os.path.join(hypar["model_path"], hypar["restore_model"]),
             map_location=device
         ))
-    net.eval()  
+    net.eval()
     return net
 
 def predict(net, inputs_val, shapes_val, hypar, device):
     net.eval()
-
-    if hypar["model_digit"] == "full":
-        inputs_val = inputs_val.type(torch.FloatTensor)
-    else:
-        inputs_val = inputs_val.type(torch.HalfTensor)
-
+    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] # list of 6 results
-    pred_val = ds_val[0][0,:,:,:] # B x 1 x H x W
-
-    # Recover the prediction spatial size to the original image size
-    pred_val = torch.squeeze(F.upsample(
+    ds_val = net(inputs_val_v)[0]
+    pred_val = ds_val[0][0,:,:,:]
+    
+    pred_val = torch.squeeze(F.interpolate(
         torch.unsqueeze(pred_val, 0),
-        (shapes_val[0][0], shapes_val[0][1]),
+        size=(shapes_val[0][0], shapes_val[0][1]),
         mode='bilinear'
     ))
-
-    ma = torch.max(pred_val)
-    mi = torch.min(pred_val)
-    pred_val = (pred_val-mi)/(ma-mi) # max = 1
-
-    if device == 'cuda': 
+    
+    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)
+    return (pred_val.detach().cpu().numpy() * 255).astype(np.uint8)
 
-# Set parameters
+# Model configuration
 hypar = {
-    "model_path": "./saved_models",
+    "model_path": "saved_models",
     "restore_model": "isnet.pth",
     "interm_sup": False,
     "model_digit": "full",
@@ -114,60 +94,70 @@ hypar = {
     "model": ISNetDIS()
 }
 
-# Build model
+# Initialize model
 net = build_model(hypar, device)
 
-def inference(image):
+def process_image(image):
     try:
-        image_path = image.name if hasattr(image, 'name') else image
-        
-        image_tensor, orig_size = load_image(image_path, hypar) 
+        if isinstance(image, str):
+            image_path = image
+        else:
+            image_path = image.name
+            
+        image_tensor, orig_size = load_image(image_path, hypar)
         mask = predict(net, image_tensor, orig_size, hypar, device)
         
-        pil_mask = Image.fromarray(mask).convert('L')
-        im_rgb = Image.open(image_path).convert("RGB")
+        mask_img = Image.fromarray(mask).convert('L')
+        rgb_img = Image.open(image_path).convert("RGB")
         
-        im_rgba = im_rgb.copy()
-        im_rgba.putalpha(pil_mask)
-
-        return [im_rgba, pil_mask]
+        rgba_img = rgb_img.copy()
+        rgba_img.putalpha(mask_img)
+        
+        return rgba_img, mask_img
     except Exception as e:
-        print(f"Error during inference: {str(e)}")
-        raise e
-
-title = "Highly Accurate Dichotomous Image Segmentation"
-description = """
-This is an unofficial demo for DIS, a model that can remove the background from a given image. 
-To use it, simply upload your image, or click one of the examples to load them.
-<br>GitHub: https://github.com/xuebinqin/DIS
-<br>Telegram bot: https://t.me/restoration_photo_bot
-[![](https://img.shields.io/twitter/follow/DoEvent?label=@DoEvent&style=social)](https://twitter.com/DoEvent)
-"""
-article = "<div><center><img src='https://visitor-badge.glitch.me/badge?page_id=max_skobeev_dis_cmp_public' alt='visitor badge'></center></div>"
-
-# Create interface
-interface = gr.Interface(
-    fn=inference,
-    inputs=gr.Image(type="filepath"),
-    outputs=[
-        gr.Image(type="pil", label="Image with Transparency"),
-        gr.Image(type="pil", label="Mask Only")
-    ],
-    examples=[
-        ["robot.png"],
-        ["ship.png"]
-    ],
-    title=title,
-    description=description,
-    article=article,
-    allow_flagging="never"
-)
-
-# Launch with corrected parameters
-interface.launch(
-    server_name="0.0.0.0",
-    server_port=7860,
-    share=False,
-    debug=True,
-    show_error=True
-)
+        raise gr.Error(f"Error processing image: {str(e)}")
+
+# Interface setup
+title = "Image Segmentation Demo"
+description = "Upload an image to extract its foreground"
+
+examples = []
+if os.path.exists("robot.png"):
+    examples.append(["robot.png"])
+if os.path.exists("ship.png"):
+    examples.append(["ship.png"])
+
+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")
+            submit_btn = gr.Button("Process")
+        
+        with gr.Column():
+            output_rgba = gr.Image(label="Transparent Background", type="pil")
+            output_mask = gr.Image(label="Segmentation Mask", type="pil")
+    
+    if examples:
+        gr.Examples(
+            examples=examples,
+            inputs=input_image,
+            outputs=[output_rgba, output_mask],
+            fn=process_image,
+            cache_examples=True
+        )
+    
+    submit_btn.click(
+        fn=process_image,
+        inputs=input_image,
+        outputs=[output_rgba, output_mask]
+    )
+
+if __name__ == "__main__":
+    app.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        show_error=True
+    )