Add tab of batch inference with saving function.

app.py

@@ -57,15 +57,37 @@ birefnet = AutoModelForImageSegmentation.from_pretrained('/'.join(('zhengpeng7',
 birefnet.to(device)
 birefnet.eval()
 
+    # for idx, image_path in enumerate(images):
+    #     im = load_img(image_path, output_type="pil")
+    #     if im is None:
+    #         continue
+        
+    #     im = im.convert("RGB")
+    #     image_size = im.size
+    #     input_images = transform_image(im).unsqueeze(0).to("cpu")
+        
+    #     with torch.no_grad():
+    #         preds = birefnet(input_images)[-1].sigmoid().cpu()
+    #     pred = preds[0].squeeze()
+    #     pred_pil = transforms.ToPILImage()(pred)
+    #     mask = pred_pil.resize(image_size)
+        
+    #     im.putalpha(mask)
+    #     output_file_path = os.path.join(save_dir, f"output_image_batch_{idx + 1}.png")
+    #     im.save(output_file_path)
+    #     output_paths.append(output_file_path)
+
+    # zip_file_path = os.path.join(save_dir, "processed_images.zip")
+    # with zipfile.ZipFile(zip_file_path, 'w') as zipf:
+    #     for file in output_paths:
+    #         zipf.write(file, os.path.basename(file))
+
+    # return output_paths, zip_file_path
 
 @spaces.GPU
-def predict(image, resolution, weights_file):
-    assert (image is not None), 'AssertionError: image cannot be None.'
-        
-    if isinstance(image, str):
-        response = requests.get(image)
-        image_data = BytesIO(response.content)
-        image = np.array(Image.open(image_data))
+def predict(images, resolution, weights_file):
+    assert (images is not None), 'AssertionError: images cannot be None.'
+
     global birefnet
     # Load BiRefNet with chosen weights
     _weights_file = '/'.join(('zhengpeng7', usage_to_weights_file[weights_file] if weights_file is not None else usage_to_weights_file['General']))
@@ -74,33 +96,63 @@ def predict(image, resolution, weights_file):
     birefnet.to(device)
     birefnet.eval()
 
-    resolution = f"{image.shape[1]}x{image.shape[0]}" if resolution == '' else resolution
-    resolution = [int(int(reso)//32*32) for reso in resolution.strip().split('x')]
+    try:
+        resolution = [int(int(reso)//32*32) for reso in resolution.strip().split('x')]
+    except:
+        resolution = [1024, 1024]
+        print('Invalid resolution input. Automatically changed to 1024x1024.')
+
+    if isinstance(images, list):
+        save_dir = 'preds-BiRefNet'
+        if not os.path.exists(save_dir):
+            os.makedirs(save_dir)
+    else:
+        # For tab_batch
+        save_paths = []
+        images = [images]
+
+    for idx_image, image_src in enumerate(images):
+        if isinstance(image_src, str):
+            response = requests.get(image_src)
+            image_data = BytesIO(response.content)
+            image = np.array(Image.open(image_data))
+        else:
+            image = image_src
     
-    image_shape = image.shape[:2]
-    image_pil = array_to_pil_image(image, tuple(resolution))
-
-    # Preprocess the image
-    image_preprocessor = ImagePreprocessor(resolution=tuple(resolution))
-    image_proc = image_preprocessor.proc(image_pil)
-    image_proc = image_proc.unsqueeze(0)
-
-    # Perform the prediction
-    with torch.no_grad():
-        scaled_pred_tensor = birefnet(image_proc.to(device))[-1].sigmoid()
-
-    if device == 'cuda':
-        scaled_pred_tensor = scaled_pred_tensor.cpu()
+        image_shape = image.shape[:2]
+        image_pil = array_to_pil_image(image, tuple(resolution))
     
-    # Resize the prediction to match the original image shape
-    pred = torch.nn.functional.interpolate(scaled_pred_tensor, size=image_shape, mode='bilinear', align_corners=True).squeeze().numpy()
-
-    # Apply the prediction mask to the original image
-    image_pil = image_pil.resize(pred.shape[::-1])
-    pred = np.repeat(np.expand_dims(pred, axis=-1), 3, axis=-1)
-    image_pred = (pred * np.array(image_pil)).astype(np.uint8)
+        # Preprocess the image
+        image_preprocessor = ImagePreprocessor(resolution=tuple(resolution))
+        image_proc = image_preprocessor.proc(image_pil)
+        image_proc = image_proc.unsqueeze(0)
+    
+        # Perform the prediction
+        with torch.no_grad():
+            scaled_pred_tensor = birefnet(image_proc.to(device))[-1].sigmoid()
+    
+        if device == 'cuda':
+            scaled_pred_tensor = scaled_pred_tensor.cpu()
+        
+        # Resize the prediction to match the original image shape
+        pred = torch.nn.functional.interpolate(scaled_pred_tensor, size=image_shape, mode='bilinear', align_corners=True).squeeze().numpy()
+    
+        # Apply the prediction mask to the original image
+        image_pil = image_pil.resize(pred.shape[::-1])
+        pred = np.repeat(np.expand_dims(pred, axis=-1), 3, axis=-1)
+        image_pred = (pred * np.array(image_pil)).astype(np.uint8)
+    
+        torch.cuda.empty_cache()
+        
+        save_file_path = os.path.join(save_dir, "{}.png".format(os.path.splitext(os.path.basename(image_src))[0]))
+        cv2.imwrite(save_file_path)
+        save_paths.append(save_file_path)
 
-    torch.cuda.empty_cache()
+    if len(images) > 1:
+        zip_file_path = os.path.join(save_dir, "{}.zip".format(save_dir))
+        with zipfile.ZipFile(zip_file_path, 'w') as zipf:
+            for file in save_paths:
+                zipf.write(file, os.path.basename(file))
 
     return image, image_pred
 
@@ -118,6 +170,11 @@ examples_url = [
 for idx_example_url, example_url in enumerate(examples_url):
     examples_url[idx_example_url].append('1024x1024')
 
+descriptions = ('Upload a picture, our model will extract a highly accurate segmentation of the subject in it.\n)'
+                 ' The resolution used in our training was `1024x1024`, thus the suggested resolution to obtain good results!\n'
+                 ' Our codes can be found at https://github.com/ZhengPeng7/BiRefNet.\n'
+                 ' We also maintain the HF model of BiRefNet at https://huggingface.co/ZhengPeng7/BiRefNet for easier access.')
+
 tab_image = gr.Interface(
     fn=predict,
     inputs=[
@@ -128,10 +185,7 @@ tab_image = gr.Interface(
     outputs=ImageSlider(label="BiRefNet's prediction", type="pil"),
     examples=examples,
     api_name="image",
-    description=('Upload a picture, our model will extract a highly accurate segmentation of the subject in it.\n)'
-                 ' The resolution used in our training was `1024x1024`, thus the suggested resolution to obtain good results!\n'
-                 ' Our codes can be found at https://github.com/ZhengPeng7/BiRefNet.\n'
-                 ' We also maintain the HF model of BiRefNet at https://huggingface.co/ZhengPeng7/BiRefNet for easier access.'),
+    description=descriptions,
 )
 
 tab_text = gr.Interface(
@@ -144,15 +198,20 @@ tab_text = gr.Interface(
     outputs=ImageSlider(label="BiRefNet's prediction", type="pil"),
     examples=examples_url,
     api_name="text",
-    description=('Upload a URL, our model will extract a highly accurate segmentation of the subject in it.\n)'
-                 ' The resolution used in our training was `1024x1024`, thus the suggested resolution to obtain good results!\n'
-                 ' Our codes can be found at https://github.com/ZhengPeng7/BiRefNet.\n'
-                 ' We also maintain the HF model of BiRefNet at https://huggingface.co/ZhengPeng7/BiRefNet for easier access.'),
+    description=descriptions+'\nTab-URL is partially modified from https://huggingface.co/spaces/not-lain/background-removal, thanks to this great work!',
+)
+
+tab_batch = gr.Interface(
+    fn=predict,
+    inputs=gr.File(label="Upload multiple images", type="filepath", file_count="multiple"),
+    outputs=[gr.Gallery(label="BiRefNet's predictions"), gr.File(label="Download masked images.")],
+    api_name="batch",
+    description=descriptions+'\nTab-batch is partially modified from https://huggingface.co/spaces/NegiTurkey/Multi_Birefnetfor_Background_Removal, thanks to this great work!',
 )
 
 demo = gr.TabbedInterface(
-    [tab_image, tab_text],
-    ["image", "text"],
+    [tab_image, tab_text, tab_batch],
+    ['image', 'text', 'batch'],
     title="BiRefNet demo for subject extraction (general / salient / camouflaged / portrait).",
 )