add demo

README.md

@@ -11,6 +11,10 @@
 
 This repo contains the MIL-FCN version of our WSCL implementation.
 
+## 🚨News
+
+**03/2024**: add demo script! Check [here](https://github.com/yhZhai/WSCL?tab=readme-ov-file#4-demo) for more details!
+
 ## 1. Setup
 Clone this repo
 
@@ -46,6 +50,21 @@ python main.py --load configs/final.yaml --eval --resume checkpoint-path
 
 We provide our pre-trained checkpoint [here](https://buffalo.box.com/s/2t3eqvwp7ua2ircpdx12sfq04sne4x50).
 
+
+## 4. Demo
+
+Running our manipulation model on your custom data!
+Before running, please configure your desired input and output path in the `demo.py` file.
+
+```shell
+python demo.py --load configs/final.yaml --resume checkpoint-path
+```
+
+By default, it evaluates all  `.jpg` files in the `demo` folder, and saves the
+detection result in `tmp`.
+
+
+
 ## Citation
 If you feel this project is helpful, please consider citing our paper
 ```bibtex

demo.py

@@ -0,0 +1,53 @@
+import albumentations as A
+import cv2
+import torch
+import tqdm
+from albumentations.pytorch.functional import img_to_tensor
+from pathlib import Path
+from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
+from torchvision.utils import draw_segmentation_masks, make_grid, save_image
+
+import utils.misc as misc
+from models import get_ensemble_model
+from opt import get_opt
+
+
+def demo(folder_path, output_path=Path("tmp")):
+    opt = get_opt()
+    model = get_ensemble_model(opt).to(opt.device)
+    misc.resume_from(model, opt.resume)
+
+    with torch.no_grad():
+        for image_path in tqdm.tqdm(folder_path.glob("*.jpg")):
+            image = cv2.imread(image_path.as_posix())
+            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+            dsm_image = torch.from_numpy(image).permute(2, 0, 1)
+            image_size = image.shape[:2]
+            raw_image = img_to_tensor(image)
+            image = img_to_tensor(
+                image,
+                normalize={"mean": IMAGENET_DEFAULT_MEAN, "std": IMAGENET_DEFAULT_STD},
+            )
+            image = image.to(opt.device).unsqueeze(0)
+            outputs = model(image, seg_size=image_size)
+            out_map = outputs["ensemble"]["out_map"][0, ...].detach().cpu()
+
+            overlay = draw_segmentation_masks(
+                dsm_image, masks=out_map[0, ...] > opt.mask_threshold
+            )
+            grid_image = make_grid(
+                [
+                    raw_image,
+                    (out_map.repeat(3, 1, 1) > opt.mask_threshold).float() * 255,
+                    overlay / 255.0,
+                ],
+                padding=5,
+            )
+            save_image(grid_image, (output_path / image_path.name).as_posix())
+
+
+if __name__ == "__main__":
+    folder_path = Path("demo")
+    output_path = Path("tmp")
+    output_path.mkdir(exist_ok=True, parents=True)
+    demo(folder_path)

engine.py

@@ -10,7 +10,7 @@ import prettytable as pt
 import torch
 import torch.nn as nn
 from fast_pytorch_kmeans import KMeans
-from pathlib2 import Path
+from pathlib import Path
 from scipy.stats import hmean
 from sklearn import metrics
 from termcolor import cprint

requirements.txt

@@ -10,7 +10,6 @@ opencv_contrib_python==4.5.3.56
 opencv_python==4.4.0.46
 opencv_python_headless==4.5.3.56
 pandas==1.3.5
-pathlib2==2.3.5
 Pillow==9.4.0
 prettytable==2.2.1
 pydensecrf==1.0rc2