multiple images

app.py

@@ -1,3 +1,4 @@
+
 import cv2
 import gradio as gr
 import os
@@ -7,11 +8,10 @@ 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")
 import time
+warnings.filterwarnings("ignore")
 
 os.system("git clone https://github.com/xuebinqin/DIS")
 os.system("mv DIS/IS-Net/* .")
@@ -20,14 +20,13 @@ os.system("mv DIS/IS-Net/* .")
 from data_loader_cache import normalize, im_reader, im_preprocess 
 from models import *
 
-#Helpers
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 # Download official weights
 if not os.path.exists("saved_models"):
     os.mkdir("saved_models")
     os.system("mv isnet.pth saved_models/")
-    
+
 class GOSNormalize(object):
     '''
     Normalize the Image using torch.transforms
@@ -45,9 +44,9 @@ 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)
+    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
+    return transform(im).unsqueeze(0), shape.unsqueeze(0)
 
 def build_model(hypar, device):
     net = hypar["model"]
@@ -67,10 +66,7 @@ def build_model(hypar, device):
     net.eval()  
     return net
 
-def predict(net,  inputs_val, shapes_val, hypar, device):
-    '''
-    Given an Image, predict the mask
-    '''
+def predict(net, inputs_val, shapes_val, hypar, device):
     net.eval()
 
     if(hypar["model_digit"]=="full"):
@@ -81,21 +77,21 @@ def predict(net,  inputs_val, shapes_val, hypar, device):
     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.upsample(torch.unsqueeze(pred_val,0),
-                                        (shapes_val[0][0],shapes_val[0][1]),
+    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) # normalize to 0~1
+    pred_val = (pred_val - mi) / (ma - mi + 1e-8)  # normalize to 0~1, +1e-8 to avoid div by zero
 
     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
-hypar = {} 
-hypar["model_path"] ="./saved_models"
+# Parameters
+hypar = {}
+hypar["model_path"] = "./saved_models"
 hypar["restore_model"] = "isnet.pth"
 hypar["interm_sup"] = False
 hypar["model_digit"] = "full"
@@ -108,32 +104,42 @@ hypar["model"] = ISNetDIS()
 # Build Model
 net = build_model(hypar, device)
 
-def inference(image, logs):
+def inference(images, logs):
     start_time = time.time()
-    
-    image_tensor, orig_size = load_image(image, hypar) 
-    mask = predict(net, image_tensor, orig_size, hypar, device)
-  
-    pil_mask = Image.fromarray(mask).convert('L')
-    im_rgb = Image.open(image).convert("RGB")
-
-    im_rgba = im_rgb.copy()
-    im_rgba.putalpha(pil_mask)
+
+    # If user didn't upload images, just return empty
+    if not images:
+        return [], logs, logs
+
+    processed_pairs = []
+    for img_path in images:
+        image_tensor, orig_size = load_image(img_path, hypar)
+        mask = predict(net, image_tensor, orig_size, hypar, device)
+
+        pil_mask = Image.fromarray(mask).convert('L')
+        im_rgb = Image.open(img_path).convert("RGB")
+        im_rgba = im_rgb.copy()
+        im_rgba.putalpha(pil_mask)
+        processed_pairs.append([im_rgba, pil_mask])
     
     end_time = time.time()
     elapsed = round(end_time - start_time, 2)
-    
-    # Update and return logs
+
+    # Flatten the list so that we can display all images in a single Gallery
+    final_images = []
+    for pair in processed_pairs:
+        final_images.extend(pair)
+
+    # Update logs
     logs = logs or ""
-    logs += f"Processed in {elapsed} seconds.\n"
+    logs += f"Processed {len(processed_pairs)} image(s) in {elapsed} seconds.\n"
 
-    # Return (gallery output), the logs state, and the logs display
-    return [im_rgba, pil_mask], logs, logs
+    return final_images, logs, logs
 
 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. "
+    "To use it, simply upload up to 3 images, or click one of the examples to load them. "
     "Read more at the links below.<br>"
     "GitHub: https://github.com/xuebinqin/DIS<br>"
     "Telegram bot: https://t.me/restoration_photo_bot<br>"
@@ -146,13 +152,19 @@ article = (
 
 interface = gr.Interface(
     fn=inference,
-    inputs=[gr.Image(type='filepath'), gr.State()],
+    inputs=[gr.Image(
+                type='filepath', 
+                label='Images (up to 3)', 
+                multiple=True, 
+                max_count=3
+            ), 
+            gr.State()],
     outputs=[
-        gr.Gallery(format="png"),
+        gr.Gallery(label="Output (rgba + mask)"),
         gr.State(),
         gr.Textbox(label="Logs", lines=6)
     ],
-    examples=[['robot.png'], ['ship.png']],
+    examples=[['robot.png'], ['ship.png']],  # for multi-image examples, pass a list like ['robot.png','ship.png']
     title=title,
     description=description,
     article=article,