Change to tab_batch to take dynamic number of images as the input.

app.py

@@ -36,8 +36,24 @@ model.eval()
 
 
 @spaces.GPU
-def pred_maps(image_1, image_2, image_3, image_4):
-    images = [image_1, image_2, image_3, image_4]
+def pred_maps(images):
+    assert (images is not None), 'AssertionError: images cannot be None.'
+    # For tab_batch
+    save_paths = []
+    save_dir = 'preds-GCoNet_plus'
+    if not os.path.exists(save_dir):
+        os.makedirs(save_dir)
+
+    image_array_lst = []
+    for idx_image, image_src in enumerate(images):
+        save_paths.append(os.path.join(save_dir, "{}.png".format(os.path.splitext(os.path.basename(image_src))[0])))
+        if isinstance(image_src, str):
+            image = np.array(Image.open(image_src))
+        else:
+            image = image_src
+        image_array_lst.append(image)
+    images = image_array_lst
+
     image_shapes = [image.shape[:2] for image in images]
     images = [Image.fromarray(image) for image in images]
 
@@ -50,30 +66,33 @@ def pred_maps(image_1, image_2, image_3, image_4):
     with torch.no_grad():
         scaled_preds_tensor = model(images_proc.to(device))[-1]
     preds = []
-    for image_shape, pred_tensor in zip(image_shapes, scaled_preds_tensor):
+    for image_shape, pred_tensor, save_path in zip(image_shapes, scaled_preds_tensor, save_paths):
         if device == 'cuda':
             pred_tensor = pred_tensor.cpu()
-        preds.append(torch.nn.functional.interpolate(pred_tensor.unsqueeze(0), size=image_shape, mode='bilinear', align_corners=True).squeeze().numpy())
-    image_preds = []
-    for image, pred in zip(images, preds):
-        image_preds.append(
-                np.hstack([np.array(image.convert('RGB')), cv2.cvtColor((pred*255).astype(np.uint8), cv2.COLOR_GRAY2RGB)])
-        )
-    return image_preds[:]
+        pred_tensor = torch.nn.functional.interpolate(pred_tensor.unsqueeze(0), size=image_shape, mode='bilinear', align_corners=True).squeeze().numpy()
+        cv2.imwrite(save_path, pred_tensor)
+
+    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 save_paths, zip_file_path
 
 
-N = 4
+# N = 4
 # examples = [[_] for _ in glob('example_images/butterfly/*')][:N]
 
-ipt = [gr.Image(width=600, height=300) for _ in range(N)]
-opt = [gr.Image(width=600, height=300) for _ in range(N)]
-demo = gr.Interface(
+tab_batch = gr.Interface(
     fn=pred_maps,
-    inputs=ipt,
-    outputs=opt,
-    # examples=examples,
-    # interpretation='default',
-    title='Online demo for `GCoNet+: A Stronger Group Collaborative Co-Salient Object Detector (T-PAMI 2023)`',
-    description='Upload pictures, most of which contain salient objects of the same class. Our demo will give you the binary maps of these co-salient objects :)\n**********Example images need to be dropped into each block, instead of click.**********'
+    inputs=gr.File(label="Upload multiple images in a group", type="filepath", file_count="multiple"),
+    outputs=[gr.Gallery(label="GCoNet+'s predictions"), gr.File(label="Download predicted maps.")],
+    api_name="batch",
+    description='Upload pictures, most of which contain salient objects of the same class. Our demo will give you the binary maps of these co-salient objects :)',
+)
+
+demo = gr.TabbedInterface(
+    [tab_batch],
+    ['batch'],
+    title="Online demo for `GCoNet+: A Stronger Group Collaborative Co-Salient Object Detector (T-PAMI 2023)`",
 )
 demo.launch(debug=True)