add gradio app

README.md

@@ -74,7 +74,7 @@ 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`.
+detection result in `tmp`, with manipulation probablities appended to the file names.
 
 
 

app.py

@@ -0,0 +1,74 @@
+from pathlib import Path
+
+import gradio as gr
+import numpy as np
+import torch
+from albumentations.pytorch.functional import img_to_tensor
+from huggingface_hub import hf_hub_download
+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 greet(input_image):
+    opt, model = _get_model()
+
+    with torch.no_grad():
+        image = input_image
+        image = np.array(image)
+        dsm_image = torch.from_numpy(image).permute(2, 0, 1)
+        image_size = image.shape[:2]
+        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()
+        pred = outputs["ensemble"]["out_map"].max().item()
+        if pred > opt.mask_threshold:
+            output_string = f"Found manipulation (manipulation probability {pred:.2f})."
+        else:
+            output_string = (
+                f"No manipulation found (manipulation probability {pred:.2f})."
+            )
+
+        overlay = draw_segmentation_masks(
+            dsm_image, masks=out_map[0, ...] > opt.mask_threshold
+        )
+        overlay = overlay.permute(1, 2, 0)
+        overlay = overlay.detach().cpu().numpy()
+        overlay = overlay.astype(np.uint8)
+    return overlay, output_string
+
+
+def _get_model(config_path="configs/final.yaml", ckpt_path="tmp/checkpoint.pt"):
+    ckpt_path = Path(ckpt_path)
+    if not ckpt_path.exists():
+        ckpt_path.parent.mkdir(exist_ok=True, parents=True)
+        hf_hub_download(
+            repo_id="yhzhai/WSCL",
+            filename="checkpoint.pt",
+            local_dir=ckpt_path.parent.as_posix(),
+        )
+
+    opt = get_opt(config_path)
+    opt.resume = ckpt_path.as_posix()
+
+    model = get_ensemble_model(opt).to(opt.device)
+    misc.resume_from(model, opt.resume)
+    return opt, model
+
+
+iface = gr.Interface(
+    fn=greet,
+    title="WSCL: Image Manipulation Detection",
+    inputs=gr.Image(),
+    outputs=["image", "text"],
+    examples=[["demo/au.jpg"], ["demo/tp.jpg"]],
+    cache_examples=True,
+)
+iface.launch()

demo.py

@@ -1,9 +1,10 @@
+from pathlib import Path
+
 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
 
@@ -31,8 +32,8 @@ def demo(folder_path, output_path=Path("tmp")):
             image = image.to(opt.device).unsqueeze(0)
             outputs = model(image, seg_size=image_size)
             out_map = outputs["ensemble"]["out_map"][0, ...].detach().cpu()
-            pred = outputs["ensemble"]["out_map"].max().item() 
-            if  pred > opt.mask_threshold:
+            pred = outputs["ensemble"]["out_map"].max().item()
+            if pred > opt.mask_threshold:
                 print(f"Found manipulation in {image_path.name}")
             else:
                 print(f"No manipulation found in {image_path.name}")