for gradio app

README.md

@@ -1,49 +1,10 @@
-# Adapt Before Comparision: A New Perspective on Cross-Domain Few-Shot Segmentation
-
-[[`Paper`](https://arxiv.org/abs/2402.17614)] accepted for CVPR'24.
-
-## Preparing Data
-Because we follow the evaluation procedure of PATNet and Remember the Difference (RtD), please refer to their work for prepration of the following datasets:
-- Deepglobe (PAT)
-- ISIC (PAT)
-- Chest X-Ray (Lung) (PAT)
-- FSS-1000 (PAT)
-- SUIM (RtD)
-
-You do not need to get all datasets. Just prepare the one you want to test our method with.
-
-## Python package prerequisites
-1. torch
-2. torchvision
-3. cv2
-4. numpy
-5. for others, follow the console output
-
-## Run it
-Call
-`python main.py --benchmark {} --datapath {} --nshot {}`
-
-for example
-`python main.py --benchmark deepglobe --datapath ./datasets/deepglobe/ --nshot 1`
-
-Available `benchmark` strings: `deepglobe`,`isic`,`lung`,`fss`,`suim`. Easiest to prepare should be `lung` or `fss`.
-
-Default is quick-infer mode.
-To change this, set `config.featext.fit_every_episode=True` in the main file.
-You can change all other parameters likewise, check the available parameters in `core/runner->makeConfig()`.
-
-## Await it
-
-You can experiment with this code. Before opening issues, I suggest awaiting nicer demonstrations and documentation to be added. 
-
-## Cite it
-If you use ABCDFSS in your research, please use the following BibTeX entry.
-```
-@article{herzog2024cdfss,
-      title={Adapt Before Comparison: A New Perspective on Cross-Domain Few-Shot Segmentation}, 
-      author={Jonas Herzog},
-      journal={arXiv:2402.17614},
-      year={2024}
-}
-```
-
+---
+title: Abcdfss
+emoji: 🔥
+colorFrom: yellow
+colorTo: purple
+sdk: gradio
+sdk_version: 4.24.0
+app_file: app.py
+pinned: false
+---

app.py

@@ -0,0 +1,122 @@
+from core import runner
+import torch
+from torch import tensor
+from PIL import Image
+import numpy as np
+import torch.nn.functional as F
+import gradio as gr
+
+#-----------------------global definitions------------------------#
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f'{device.type=}')
+
+description = '<p> Choose an example below; OR <br>\
+    Upload by yourself: <br>\
+    1. Upload any test image (query) with any target object you wish to segment <br>\
+    2. Upload another image (support) with the target object or a variation of it <br>\
+    3. Upload a binary mask that segments the target objet in the support image <br>\
+    </p>'
+    
+# qimg, simg, smask
+example_episodes = [
+            ['./imgs/cake1.png', './imgs/cake2.png', './imgs/cake2_mask.png'],
+            ['./imgs/549870_35.jpg', './imgs/457070_00.jpg', './imgs/457070_00.png'],
+            ['./imgs/ISIC_0000372.jpg', './imgs/ISIC_0013176.jpg', './imgs/ISIC_0013176_segmentation.png'],
+            ['./imgs/d_r_450_.jpg', './imgs/d_r_465_.jpg', './imgs/d_r_465_.bmp'],
+            ['./imgs/CHNCXR_0282_0.png', './imgs/CHNCXR_0324_0.png', './imgs/CHNCXR_0324_0_mask.png'],
+            ['./imgs/1.jpg', './imgs/5.jpg', './imgs/5.png']
+            ]
+blank_img = './imgs/blank.png'
+
+gr_img = lambda name: gr.Image(label=name, sources=['upload', 'webcam'], type="pil")
+inputs = [gr_img('Query Img'), gr_img('Support Img'), gr_img('Support Mask'), gr.Checkbox(label='re-adapt')]
+if device.type=='cpu':
+  inputs.append(gr.Checkbox(label='Confirm CPU run (CHOOSE ONLY WHEN REQUESTED)'))
+
+def prepare_feat_maker():
+    config = runner.makeConfig()
+    class DummyDataset:
+        class_ids = [0]
+    fake_feat_maker = runner.makeFeatureMaker(DummyDataset(), config, device=device)
+    return fake_feat_maker
+
+feat_maker = prepare_feat_maker()
+has_fit = False
+
+#-----------------------------------------------------------------#
+
+def reset_layers():
+    global feat_maker
+    feat_maker = prepare_feat_maker()
+    
+def prepare_batch(q_img_pil, s_img_pil, s_mask_pil):
+    from data.dataset import FSSDataset
+    FSSDataset.initialize(img_size=400,datapath='')
+    q_img_tensor = FSSDataset.transform(q_img_pil)
+    s_img_tensor = FSSDataset.transform(s_img_pil)
+    s_mask_tensor = torch.tensor(np.array(s_mask_pil.convert('L')))
+    s_mask_tensor = F.interpolate(s_mask_tensor.unsqueeze(0).unsqueeze(0).float(), s_img_tensor.size()[-2:], mode='nearest').squeeze()
+    add_batch_dim = lambda t: t.unsqueeze(0)
+    add_kshot_dim = lambda t: t.unsqueeze(1)
+    fake_batch = {'query_img':add_batch_dim(q_img_tensor), 'support_imgs':add_kshot_dim(add_batch_dim(s_img_tensor)), 'support_masks':add_kshot_dim(add_batch_dim(s_mask_tensor)), 'class_id':tensor([0])}
+    return fake_batch
+
+norm = lambda t: (t - t.min()) / (t.max() - t.min())
+def overlay(img, mask):
+    #img h,w,3(float) mask h,w(float)
+    return norm(img)*0.5 + mask[:,:,np.newaxis]*0.5
+    
+def from_model(q_img, s_img, s_mask):
+    batch = prepare_batch(q_img, s_img, s_mask)
+    sseval = runner.SingleSampleEval(batch, feat_maker)
+    pred_logits, pred_mask = sseval.forward()
+    global has_fit
+    has_fit = True
+    # logit mask in range from -1 to 1, and mask-overlaid query image 0 to 1
+    return norm(pred_logits[0].numpy()), overlay(batch['query_img'][0].permute(1,2,0).numpy(), pred_mask[0].numpy())
+
+def predict(q,s,m,re_adapt,confirmed):
+  print(f'predict with {re_adapt=}, {confirmed=}')
+  print(f'{type(q)=}')
+  is_cache_run = re_adapt is None and confirmed is None
+  is_example = any([(np.array_equal(np.array(m),np.array(Image.open(e[2])))) for e in example_episodes]) #[2] pointing to support mask
+  print(f'{is_example=}')
+
+  if is_cache_run:
+    reset_layers()
+    pred = from_model(q,s,m)
+    msg = 'Results ready.'
+    return msg, *pred
+  elif re_adapt:
+    if confirmed:
+        reset_layers()
+        pred = from_model(q,s,m)
+        msg = "Results ready.\nRemember to untick 're-adapt' if you wish to predict more images with the same parameters."
+        return msg, *pred
+    else:
+        msg = "You chose to re-adapt but are on CPU.\nThis may take 1 minute on your local machine or 4 minutes on huggingface space.\nSelect 'Confirm CPU run' to start."
+        return msg, blank_img, blank_img
+  else:
+    if is_example:
+        msg = "Cached results for example have been shown previously already.\nTo view it again, click the example again.\nTo run adaption again from scratch, select 're-adapt'."
+        return msg, blank_img, blank_img
+    else:
+        if has_fit:
+            pred = from_model(q,s,m)
+            msg = "Results predicted based on layers fitted from previous run.\nIf you wish to re-adapt, select 're-adapt'."
+            return msg, *pred
+        else:
+            msg = "This is the first time you predict own images.\nThe attached layers need to be fitted.\nPlease select 're-adapt'."
+            return msg, blank_img, blank_img
+    
+gradio_app = gr.Interface(
+    fn=predict,
+    inputs=inputs,
+    outputs=[gr.Textbox(label="Status"), gr.Image(label="Coarse Query Prediction"), gr.Image(label="Mask Prediction")],
+    description=description,
+    examples=example_episodes,
+    title="abcdfss",
+)
+
+gradio_app.launch()

main.py

@@ -1,37 +0,0 @@
-from core import runner
-import torch
-import argparse
-
-def parse_opts():
-    r"""arguments"""
-    parser = argparse.ArgumentParser(description='Adapt Before Comparison - A New Perspective on Cross-Domain Few-Shot Segmentation')
-
-    # common
-    parser.add_argument('--benchmark', type=str, default='lung', choices=['fss', 'deepglobe', 'lung', 'isic', 'fss', 'lung'])
-    parser.add_argument('--datapath', type=str)
-    parser.add_argument('--nshot', type=int, default=1)
-
-    args = parser.parse_args()
-    return args
-
-
-if __name__ == '__main__':
-    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    args = parse_opts()
-    print(args)
-    runner.args.benchmark = args.benchmark
-    runner.args.datapath = args.datapath
-    runner.args.nshot = args.nshot
-
-    dataloader = runner.makeDataloader()
-    config = runner.makeConfig()
-    feat_maker = runner.makeFeatureMaker(dataloader.dataset, config, device=device)
-    average_meter = runner.AverageMeterWrapper(dataloader, device)
-
-    for idx, batch in enumerate(dataloader):
-        sseval = runner.SingleSampleEval(batch, feat_maker)
-        sseval.forward()
-        sseval.calc_metrics()
-        average_meter.update(sseval)
-        average_meter.write(idx)
-    print('Result m|FB:', average_meter.average_meter.compute_iou())

requirements.txt

@@ -0,0 +1,5 @@
+torch
+torchvision
+opencv-python
+numpy
+tensorboardX