Update app.py

code minimized

app.py

@@ -1,196 +1,130 @@
+#######################################################################################
+#
+# MIT License
+#
+# Copyright (c) [2025] [leonelhs@gmail.com]
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#######################################################################################
+
+# This file implements an API endpoint for DIS background image removal system.
+#
+# Source code is based on or inspired by several projects.
+# For more details and proper attribution, please refer to the following resources:
+#
+# - [DIS] - [https://github.com/xuebinqin/DIS]
+
 import gradio as gr
-import gradio as gr
-from pathlib import Path
-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 matplotlib.pyplot as plt
-import warnings
-from zipfile import ZipFile
-
-warnings.filterwarnings("ignore")
-
+from PIL import Image
+from huggingface_hub import hf_hub_download
+from torch.autograd import Variable
+from torchvision.transforms.functional import normalize
 
 # project imports
-from data_loader_cache import normalize, im_reader, im_preprocess 
-from models import *
+from models.isnet import ISNetDIS
+
+REPO_ID = "leonelhs/removators"
 
-#Helpers
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
-    
-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]):
-        self.mean = mean
-        self.std = std
-
-    def __call__(self,image):
-        image = normalize(image,self.mean,self.std)
-        return image
-
-
-transform =  transforms.Compose([GOSNormalize([0.5,0.5,0.5],[1.0,1.0,1.0])])
-
-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) # make a batch of image, shape
-
-
-def build_model(hypar,device):
-    net = hypar["model"]#GOSNETINC(3,1)
-
-    # convert to half precision
-    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"]!=""):
-        net.load_state_dict(torch.load(hypar["model_path"]+"/"+hypar["restore_model"], map_location=device))
-        net.to(device)
-    net.eval()  
-    return net
-
-    
-def predict(net,  inputs_val, shapes_val, hypar, device):
-    '''
-    Given an Image, predict the mask
-    '''
-    net.eval()
-
-    if(hypar["model_digit"]=="full"):
-        inputs_val = inputs_val.type(torch.FloatTensor)
+net = ISNetDIS()
+
+model_path = hf_hub_download(repo_id=REPO_ID, filename='isnet.pth')
+net.load_state_dict(torch.load(model_path, map_location=device))
+net.to(device)
+net.eval()
+
+def im_preprocess(im,size):
+    if len(im.shape) < 3:
+        im = im[:, :, np.newaxis]
+    if im.shape[2] == 1:
+        im = np.repeat(im, 3, axis=2)
+    im_tensor = torch.tensor(im.copy(), dtype=torch.float32)
+    im_tensor = torch.transpose(torch.transpose(im_tensor,1,2),0,1)
+    if len(size)<2:
+        return im_tensor, im.shape[0:2]
     else:
-        inputs_val = inputs_val.type(torch.HalfTensor)
-
-  
-    inputs_val_v = Variable(inputs_val, requires_grad=False).to(device) # wrap inputs in Variable
-   
-    ds_val = net(inputs_val_v)[0] # list of 6 results
-
-    pred_val = ds_val[0][0,:,:,:] # B x 1 x H x W    # we want the first one which is the most accurate prediction
-
-    ## recover the prediction spatial size to the orignal image size
-    pred_val = torch.squeeze(F.upsample(torch.unsqueeze(pred_val,0),(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': torch.cuda.empty_cache()
-    return (pred_val.detach().cpu().numpy()*255).astype(np.uint8) # it is the mask we need
-    
-# Set Parameters
-hypar = {} # paramters for inferencing
-
-
-hypar["model_path"] ="./saved_models" ## load trained weights from this path
-hypar["restore_model"] = "isnet.pth" ## name of the to-be-loaded weights
-hypar["interm_sup"] = False ## indicate if activate intermediate feature supervision
-
-##  choose floating point accuracy --
-hypar["model_digit"] = "full" ## indicates "half" or "full" accuracy of float number
-hypar["seed"] = 0
-
-hypar["cache_size"] = [1024, 1024] ## cached input spatial resolution, can be configured into different size
-
-## data augmentation parameters ---
-hypar["input_size"] = [1024, 1024] ## mdoel input spatial size, usually use the same value hypar["cache_size"], which means we don't further resize the images
-hypar["crop_size"] = [1024, 1024] ## random crop size from the input, it is usually set as smaller than hypar["cache_size"], e.g., [920,920] for data augmentation
-
-hypar["model"] = ISNetDIS()
-
- # Build Model
-net = build_model(hypar, device)
-
-
-def inference(image_path):
-  
-  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")
-  
-  im_rgba = im_rgb.copy()
-  im_rgba.putalpha(pil_mask)
-  file_name = Path(image_path).stem+"_nobg.png"
-  file_path = Path(Path(image_path).parent,file_name)
-  im_rgba.save(file_path)
-  return str(file_path.resolve())
-
-def bw(image_files):
-    print(image_files)
-    output = []
-    for idx, file in enumerate(image_files):
-        print(file.name)
-        img = Image.open(file.name)
-        img = img.convert("L")
-        output.append(img)
-    print(output)
-    return output
-
-def bw_single(image_file):
-    img = Image.open(image_file)
-    img = img.convert("L")
-    return img
-
-def batch(image_files):
-    output = []
-    for idx, file in enumerate(image_files):
-        file = inference(file.name)
-        output.append(file)
-    
-    with ZipFile("tmp.zip", "w") as zipObj:
-        for idx, file in enumerate(output):
-            zipObj.write(file, file.split("/")[-1])
-    return output,"tmp.zip"
-
-with gr.Blocks() as iface:
-    gr.Markdown("# Remove Background")
-    gr.HTML("Uses <a href='https://github.com/xuebinqin/DIS'>DIS</a> to remove background")
-    with gr.Tab("Single Image"):
-        with gr.Row():
-            with gr.Column():
-                image = gr.Image(type='filepath')
-            with gr.Column():
-                image_output = gr.Image(interactive=False)
-        with gr.Row():
-            with gr.Column():
-                single_removebg = gr.Button("Remove Bg")
-            with gr.Column():
-                single_clear = gr.Button("Clear")
-
-    
-    with gr.Tab("Batch"):
-        with gr.Row():
-            with gr.Column():
-                images = gr.File(file_count="multiple", file_types=["image"])
-            with gr.Column():
-                gallery = gr.Gallery()
-                file_list = gr.Files(interactive=False)
-
-        with gr.Row():
-            with gr.Column():
-                batch_removebg = gr.Button("Batch Process")
-            with gr.Column():
-                batch_clear = gr.Button("Clear")
-    #Events
-    single_removebg.click(inference, inputs=image, outputs=image_output)
-    batch_removebg.click(batch, inputs=images, outputs=[gallery,file_list])
-    single_clear.click(lambda: None, None, image, queue=False)
-    batch_clear.click(lambda: None, None, images, queue=False)
-
-iface.launch()
+        im_tensor = torch.unsqueeze(im_tensor,0)
+        im_tensor = F.interpolate(im_tensor, size, mode="bilinear")
+        im_tensor = torch.squeeze(im_tensor,0)
+
+    return im_tensor.type(torch.uint8), im.shape[0:2]
+
+
+def predict(image):
+    """
+        Remove the background from an image.
+        The function extracts the foreground and generates both a background-removed
+        image and a binary mask.
+
+        Parameters:
+            image (string): File path to the input image.
+        Returns:
+            paths (tuple): paths for background-removed image and cutting mask.
+    """
+
+    im_tensor, shapes = im_preprocess(image, [1024, 1024])
+    shapes = torch.from_numpy(np.array(shapes)).unsqueeze(0)
+
+    im_tensor = torch.divide(im_tensor, 255.0)
+    im_tensor = normalize(im_tensor, mean=[0.5, 0.5, 0.5], std=[1.0, 1.0, 1.0]).unsqueeze(0)
+    im_tensor_v = Variable(im_tensor, requires_grad=False)  # wrap inputs in Variable
+    ds_val = net(im_tensor_v)[0]  # list of 6 results
+    prediction = ds_val[0][0, :, :, :]  # B x 1 x H x W    # we want the first one which is the most accurate prediction
+    ## recover the prediction spatial size to the original image size
+    size = (shapes[0][0], shapes[0][1])
+    prediction = F.interpolate(torch.unsqueeze(prediction, 0), size, mode='bilinear')
+    prediction = torch.squeeze(prediction)
+
+    ma = torch.max(prediction)
+    mi = torch.min(prediction)
+    prediction = (prediction - mi) / (ma - mi)  # max = 1
+
+    torch.cuda.empty_cache()
+    mask = (prediction.detach().cpu().numpy() * 255).astype(np.uint8)  # it is the mask we need
+
+    mask = Image.fromarray(mask).convert('L')
+    image_rgb = Image.fromarray(image).convert("RGB")
+    image_rgb.putalpha(mask)
+    return image_rgb, mask
+
+article = "<div><center>Unofficial demo from:<a href='https://github.com/xuebinqin/DIS'>DIS</<></center></div>"
+
+with gr.Blocks(title="DIS") as app:
+    gr.Markdown("## Dichotomous Image Segmentation")
+    with gr.Row():
+        with gr.Column(scale=1):
+            inp = gr.Image(type="numpy", label="Upload Image")
+            btn_predict = gr.Button("Remove background")
+        with gr.Column(scale=2):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    out = gr.Image(type="filepath", label="Output image")
+                    with gr.Accordion("See intermediates", open=False):
+                        out_mask = gr.Image(type="filepath", label="Mask")
+
+    btn_predict.click(predict, inputs=inp, outputs=[out, out_mask])
+    gr.HTML(article)
+
+app.launch(share=False, debug=True, show_error=True, mcp_server=True, pwa=True)
+app.queue()