Upload 10 files

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Prasanna Reddy Pulakurthi
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,14 +1,25 @@
 ---
-title: DR Augmentation
-emoji: 👁
-colorFrom: green
-colorTo: yellow
+title: Dual-Region Augmentation (DRA, Local U²-Net)
+emoji: ✂️
+colorFrom: gray
+colorTo: blue
 sdk: gradio
-sdk_version: 5.44.1
+sdk_version: 4.22.1
 app_file: app.py
 pinned: false
 license: mit
-short_description: 'Dual-Region Foreground-Background Augmentation '
+short_description: An interactive demo for DRA augmentation.
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+### Dual-Region Foreground-Background Augmentation (DRA)
+
+An interactive demo for DRA augmentation.
+
+GitHub repo: https://github.com/PrasannaPulakurthi/Foreground-Background-Augmentation
+
+arxiv.org/abs/2504.13077
+
+- **Grid control**: choose number of patches (2×2, 4×4, 8×8, 16×16). No overlap.
+- **Background**: jigsaw shuffle on the full image, fused only into background via U²‑Net mask.
+- **Foreground**: Gaussian noise applied to **one random rectangular box** inside the foreground.
+- **Weights**: uses local `saved_models/u2net(.pth)` or `saved_models/u2netp(.pth)`.

app.py

@@ -0,0 +1,309 @@
+import io, os, random
+from pathlib import Path
+
+import gradio as gr
+import numpy as np
+from PIL import Image
+
+import torch
+from torchvision import transforms
+
+# --- Expect the user's u2net codebase available as a local module folder "u2net"
+try:
+    from u2net.model import U2NET, U2NETP
+except Exception as e:
+    raise RuntimeError(
+        "Could not import 'u2net'. Please place the U^2-Net code folder named 'u2net' "
+        "next to app.py (containing model.py, data_loader.py, ...)."
+    ) from e
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+IMG_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".bmp"}
+
+# Cache for loaded models
+_MODEL_CACHE = {"u2net": None, "u2netp": None}
+
+
+def _pil_to_np(img: Image.Image) -> np.ndarray:
+    return np.array(img.convert("RGB"), dtype=np.float32)
+
+
+def _np_to_pil(arr: np.ndarray) -> Image.Image:
+    return Image.fromarray(np.clip(arr, 0, 255).astype(np.uint8), mode="RGB")
+
+
+def _minmax_norm(t: torch.Tensor, eps: float = 1e-8) -> torch.Tensor:
+    t_min = t.amin(dim=(-2, -1), keepdim=True)
+    t_max = t.amax(dim=(-2, -1), keepdim=True)
+    return (t - t_min) / (t_max - t_min + eps)
+
+
+def _find_weight_file(model_type: str) -> Path:
+    """
+    model_type in {"u2net", "u2netp"}.
+    Looks under: saved_models/{model_type}/{model_type}.pth (preferred)
+    or first *.pth under that subfolder.
+    """
+    base = Path("saved_models").expanduser().resolve()
+    sub = base / model_type
+    preferred = sub / f"{model_type}.pth"
+    if preferred.exists():
+        return preferred
+    # fallback: first .pth in subdir
+    if sub.exists():
+        for p in sorted(sub.glob("*.pth")):
+            return p
+    raise FileNotFoundError(
+        f"Could not find weights for '{model_type}'. Expected at '{preferred}' or any .pth in '{sub}'."
+    )
+
+
+def load_u2net(model_type: str = "u2netp"):
+    assert model_type in {"u2net", "u2netp"}
+    if _MODEL_CACHE.get(model_type) is not None:
+        return _MODEL_CACHE[model_type]
+
+    weights_path = _find_weight_file(model_type)
+    if model_type == "u2net":
+        net = U2NET(3, 1)
+    else:
+        net = U2NETP(3, 1)
+
+    state = torch.load(weights_path, map_location="cpu")
+    net.load_state_dict(state)
+    net.to(DEVICE)
+    net.eval()
+    _MODEL_CACHE[model_type] = net
+    return net
+
+
+def get_u2net_mask_with_local_weights(
+    pil_img: Image.Image,
+    model_type: str = "u2netp",
+    resize_to: int = 320,
+) -> Image.Image:
+    """
+    Single-image inference using user's local U^2-Net/U^2-NetP weights.
+    Returns 8-bit 'L' mask resized back to original W,H.
+    """
+    W, H = pil_img.size
+    net = load_u2net(model_type)
+
+    tr = transforms.Compose([
+        transforms.Resize((resize_to, resize_to), interpolation=Image.BILINEAR),
+        transforms.ToTensor(),  # [0,1], CxHxW
+    ])
+    x = tr(pil_img.convert("RGB")).unsqueeze(0).to(DEVICE)  # 1x3x320x320
+
+    with torch.no_grad():
+        d1, d2, d3, d4, d5, d6, d7 = net(x)
+        pred = d1[:, 0, :, :]  # 1xHxW
+        pred = _minmax_norm(pred)  # min-max normalize per-batch
+
+    pred_np = pred.squeeze(0).detach().cpu().numpy()  # HxW, [0..1]
+    mask_small = Image.fromarray((pred_np * 255).astype(np.uint8), mode="L")
+    mask = mask_small.resize((W, H), resample=Image.BILINEAR)
+    return mask
+
+
+def jigsaw_shuffle_full_image(pil_img: Image.Image, N: int, seed: int) -> Image.Image:
+    """
+    Create a jigsaw-shuffled version of the input image by splitting into an N×N grid
+    with *no overlap*, permuting patches uniformly at random, and reassembling.
+    To keep uniform patch sizes, we center-crop the image to (H2,W2) divisible by N,
+    then paste back to the original canvas.
+    """
+    random.seed(seed)
+    np.random.seed(seed)
+
+    W, H = pil_img.size
+    # compute crop that is divisible by N
+    H2 = (H // N) * N
+    W2 = (W // N) * N
+    pad_canvas = Image.fromarray(np.array(pil_img), mode="RGB")
+    if H2 == 0 or W2 == 0:
+        # too small; just return original
+        return pil_img
+
+    # center crop box
+    y0 = (H - H2) // 2
+    x0 = (W - W2) // 2
+    crop = pil_img.crop((x0, y0, x0 + W2, y0 + H2))
+
+    arr = np.array(crop).copy()
+    out = np.empty_like(arr)
+
+    ph = H2 // N
+    pw = W2 // N
+
+    # build coordinates
+    coords = []
+    for i in range(N):
+        for j in range(N):
+            y1 = i * ph
+            x1 = j * pw
+            coords.append((y1, y1 + ph, x1, x1 + pw))
+
+    perm = np.random.permutation(len(coords))
+    for dst_idx, src_idx in enumerate(perm):
+        yd0, yd1, xd0, xd1 = coords[dst_idx]
+        ys0, ys1, xs0, xs1 = coords[src_idx]
+        out[yd0:yd1, xd0:xd1, :] = arr[ys0:ys1, xs0:xs1, :]
+
+    # paste back into original canvas
+    pad = np.array(pad_canvas)
+    pad[y0:y0 + H2, x0:x0 + W2, :] = out
+    return Image.fromarray(pad.astype(np.uint8), mode="RGB")
+
+
+def add_noise_in_random_fg_box(base: np.ndarray, mask_hard: np.ndarray, sigma: float, seed: int) -> np.ndarray:
+    """
+    Add Gaussian noise only within a randomly selected rectangular region *inside the foreground*.
+    If no foreground is found, no noise is added.
+    """
+    rng = np.random.default_rng(seed)
+    H, W = mask_hard.shape
+    ys, xs = np.where(mask_hard > 0.5)
+    if len(ys) == 0:
+        return base
+
+    y_min, y_max = ys.min(), ys.max()
+    x_min, x_max = xs.min(), xs.max()
+
+    # choose box size as a fraction of the FG bbox (20% ~ 60% of width/height)
+    box_h = max(1, int((y_max - y_min + 1) * float(rng.uniform(0.2, 0.6))))
+    box_w = max(1, int((x_max - x_min + 1) * float(rng.uniform(0.2, 0.6))))
+
+    # random top-left so that box fits within FG bbox
+    y0 = int(rng.integers(y_min, max(y_min, y_max - box_h + 1) + 1))
+    x0 = int(rng.integers(x_min, max(x_min, x_max - box_w + 1) + 1))
+
+    # slice
+    region_mask = mask_hard[y0:y0 + box_h, x0:x0 + box_w]
+    if region_mask.size == 0:
+        return base
+
+    noise = rng.normal(0.0, sigma, size=(box_h, box_w, 3))
+    base[y0:y0 + box_h, x0:x0 + box_w, :] += noise * region_mask[:, :, None]
+    return base
+
+
+def dual_region_augment_dra(
+    img: Image.Image,
+    grid_n: int = 8,
+    fg_noise_std: float = 20.0,
+    seed: int = 0,
+    model_type: str = "u2netp",
+):
+    """
+    DRA:
+    - Background: jigsaw-shuffle full image on an N×N grid (no overlap), then use only on background.
+    - Foreground: add Gaussian noise to a single randomly selected rectangular box (inside FG).
+    - Fusion: FG from noisy image, BG from jigsaw image, using a hard U^2-Net mask.
+    """
+    random.seed(seed)
+    np.random.seed(seed)
+
+    base = _pil_to_np(img)  # (H, W, 3), float32 [0..255]
+    H, W = base.shape[:2]
+
+    # 1) Mask from local weights (no feather; use hard threshold at 0.5)
+    raw_mask_L = get_u2net_mask_with_local_weights(img, model_type=model_type)
+    mask = (np.array(raw_mask_L, dtype=np.float32) / 255.0) >= 0.5
+    mask_hard = mask.astype(np.float32)  # (H,W) in {0,1}
+
+    # 2) Foreground: noise in a random FG rectangle
+    img_fg = base.copy()
+    img_fg = add_noise_in_random_fg_box(img_fg, mask_hard, sigma=fg_noise_std, seed=seed)
+
+    # 3) Background: jigsaw-shuffle full image on N×N grid
+    jig = jigsaw_shuffle_full_image(Image.fromarray(base.astype(np.uint8)), N=grid_n, seed=seed)
+    img_bg = _pil_to_np(jig)
+
+    # 4) Fusion: BG where mask==0, FG where mask==1
+    m3 = np.repeat(mask_hard[:, :, None], 3, axis=2)
+    out = img_bg * (1.0 - m3) + img_fg * m3
+
+    return _np_to_pil(out), raw_mask_L
+
+
+# ---- Gradio UI ----
+GRID_CHOICES = ["2x2", "4x4", "8x8", "16x16"]
+
+def parse_grid_choice(s: str) -> int:
+    try:
+        n = int(s.lower().split('x')[0])
+        return max(2, min(16, n))
+    except Exception:
+        return 8
+
+
+def run_demo(
+    image,
+    grid_choice,
+    fg_noise_std,
+    seed,
+    model_type,
+):
+    if image is None:
+        raise gr.Error("Please upload an image or pick one from the examples.")
+    n = parse_grid_choice(grid_choice)
+    out_img, mask_L = dual_region_augment_dra(
+        image,
+        grid_n=n,
+        fg_noise_std=fg_noise_std,
+        seed=seed,
+        model_type=model_type,
+    )
+    return out_img, mask_L
+
+
+def list_example_images():
+    ex_dir = Path("examples")
+    ex_dir.mkdir(exist_ok=True)
+    files = [
+        str(p) for p in sorted(ex_dir.iterdir())
+        if p.suffix.lower() in IMG_EXTS and p.is_file()
+    ]
+    return files if files else None
+
+
+examples = list_example_images()
+
+with gr.Blocks(title="Dual-Region Augmentation (DRA, Local U²-Net Weights)") as demo:
+    gr.Markdown(
+        "### Dual-Region Augmentation (DRA)\n"
+        "- **Background**: random patch shuffle on an N×N grid (no overlap) in the background region.\n"
+        "- **Foreground**: Gaussian noise in the foreground region.\n"
+        "- **Mask**: U²-Net / U²-NetP"
+    )
+    with gr.Row():
+        with gr.Column():
+            in_img = gr.Image(type="pil", label="Input Image", sources=["upload", "clipboard"])
+            grid_choice = gr.Dropdown(GRID_CHOICES, value="8x8", label="Grid (number of patches)")
+            noise_std = gr.Slider(0, 100, value=50, step=1, label="Foreground Noise σ")
+            seed = gr.Slider(0, 9999, value=69, step=1, label="Seed")
+            model_type = gr.Dropdown(choices=["u2netp", "u2net"], value="u2netp", label="Model Type")
+            btn = gr.Button("Augment")
+        with gr.Column():
+            out_img = gr.Image(type="pil", label="Augmented Output")
+            out_mask = gr.Image(type="pil", label="U²-Net Mask (preview)")
+
+    btn.click(
+        fn=run_demo,
+        inputs=[in_img, grid_choice, noise_std, seed, model_type],
+        outputs=[out_img, out_mask],
+        concurrency_limit=3,
+        api_name="augment",
+    )
+
+    if examples:
+        gr.Examples(
+            examples=examples,
+            inputs=[in_img],
+            examples_per_page=12,
+            label="Examples (loaded from ./examples)"
+        )
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,5 @@
+gradio>=4.22,<5
+numpy
+Pillow
+torch
+torchvision

saved_models/u2net/u2net.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:10025a17f49cd3208afc342b589890e402ee63123d6f2d289a4a0903695cce58
+size 176290937

saved_models/u2netp/u2netp.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e7567cde013fb64813973ce6e1ecc25a80c05c3ca7adbc5a54f3c3d90991b854
+size 4683258

u2net/data_loader.py

@@ -0,0 +1,266 @@
+# data loader
+from __future__ import print_function, division
+import glob
+import torch
+from skimage import io, transform, color
+import numpy as np
+import random
+import math
+import matplotlib.pyplot as plt
+from torch.utils.data import Dataset, DataLoader
+from torchvision import transforms, utils
+from PIL import Image
+
+#==========================dataset load==========================
+class RescaleT(object):
+
+	def __init__(self,output_size):
+		assert isinstance(output_size,(int,tuple))
+		self.output_size = output_size
+
+	def __call__(self,sample):
+		imidx, image, label = sample['imidx'], sample['image'],sample['label']
+
+		h, w = image.shape[:2]
+
+		if isinstance(self.output_size,int):
+			if h > w:
+				new_h, new_w = self.output_size*h/w,self.output_size
+			else:
+				new_h, new_w = self.output_size,self.output_size*w/h
+		else:
+			new_h, new_w = self.output_size
+
+		new_h, new_w = int(new_h), int(new_w)
+
+		# #resize the image to new_h x new_w and convert image from range [0,255] to [0,1]
+		# img = transform.resize(image,(new_h,new_w),mode='constant')
+		# lbl = transform.resize(label,(new_h,new_w),mode='constant', order=0, preserve_range=True)
+
+		img = transform.resize(image,(self.output_size,self.output_size),mode='constant')
+		lbl = transform.resize(label,(self.output_size,self.output_size),mode='constant', order=0, preserve_range=True)
+
+		return {'imidx':imidx, 'image':img,'label':lbl}
+
+class Rescale(object):
+
+	def __init__(self,output_size):
+		assert isinstance(output_size,(int,tuple))
+		self.output_size = output_size
+
+	def __call__(self,sample):
+		imidx, image, label = sample['imidx'], sample['image'],sample['label']
+
+		if random.random() >= 0.5:
+			image = image[::-1]
+			label = label[::-1]
+
+		h, w = image.shape[:2]
+
+		if isinstance(self.output_size,int):
+			if h > w:
+				new_h, new_w = self.output_size*h/w,self.output_size
+			else:
+				new_h, new_w = self.output_size,self.output_size*w/h
+		else:
+			new_h, new_w = self.output_size
+
+		new_h, new_w = int(new_h), int(new_w)
+
+		# #resize the image to new_h x new_w and convert image from range [0,255] to [0,1]
+		img = transform.resize(image,(new_h,new_w),mode='constant')
+		lbl = transform.resize(label,(new_h,new_w),mode='constant', order=0, preserve_range=True)
+
+		return {'imidx':imidx, 'image':img,'label':lbl}
+
+class RandomCrop(object):
+
+	def __init__(self,output_size):
+		assert isinstance(output_size, (int, tuple))
+		if isinstance(output_size, int):
+			self.output_size = (output_size, output_size)
+		else:
+			assert len(output_size) == 2
+			self.output_size = output_size
+	def __call__(self,sample):
+		imidx, image, label = sample['imidx'], sample['image'], sample['label']
+
+		if random.random() >= 0.5:
+			image = image[::-1]
+			label = label[::-1]
+
+		h, w = image.shape[:2]
+		new_h, new_w = self.output_size
+
+		top = np.random.randint(0, h - new_h)
+		left = np.random.randint(0, w - new_w)
+
+		image = image[top: top + new_h, left: left + new_w]
+		label = label[top: top + new_h, left: left + new_w]
+
+		return {'imidx':imidx,'image':image, 'label':label}
+
+class ToTensor(object):
+	"""Convert ndarrays in sample to Tensors."""
+
+	def __call__(self, sample):
+
+		imidx, image, label = sample['imidx'], sample['image'], sample['label']
+
+		tmpImg = np.zeros((image.shape[0],image.shape[1],3))
+		tmpLbl = np.zeros(label.shape)
+
+		image = image/np.max(image)
+		if(np.max(label)<1e-6):
+			label = label
+		else:
+			label = label/np.max(label)
+
+		if image.shape[2]==1:
+			tmpImg[:,:,0] = (image[:,:,0]-0.485)/0.229
+			tmpImg[:,:,1] = (image[:,:,0]-0.485)/0.229
+			tmpImg[:,:,2] = (image[:,:,0]-0.485)/0.229
+		else:
+			tmpImg[:,:,0] = (image[:,:,0]-0.485)/0.229
+			tmpImg[:,:,1] = (image[:,:,1]-0.456)/0.224
+			tmpImg[:,:,2] = (image[:,:,2]-0.406)/0.225
+
+		tmpLbl[:,:,0] = label[:,:,0]
+
+
+		tmpImg = tmpImg.transpose((2, 0, 1))
+		tmpLbl = label.transpose((2, 0, 1))
+
+		return {'imidx':torch.from_numpy(imidx), 'image': torch.from_numpy(tmpImg), 'label': torch.from_numpy(tmpLbl)}
+
+class ToTensorLab(object):
+	"""Convert ndarrays in sample to Tensors."""
+	def __init__(self,flag=0):
+		self.flag = flag
+
+	def __call__(self, sample):
+
+		imidx, image, label =sample['imidx'], sample['image'], sample['label']
+
+		tmpLbl = np.zeros(label.shape)
+
+		if(np.max(label)<1e-6):
+			label = label
+		else:
+			label = label/np.max(label)
+
+		# change the color space
+		if self.flag == 2: # with rgb and Lab colors
+			tmpImg = np.zeros((image.shape[0],image.shape[1],6))
+			tmpImgt = np.zeros((image.shape[0],image.shape[1],3))
+			if image.shape[2]==1:
+				tmpImgt[:,:,0] = image[:,:,0]
+				tmpImgt[:,:,1] = image[:,:,0]
+				tmpImgt[:,:,2] = image[:,:,0]
+			else:
+				tmpImgt = image
+			tmpImgtl = color.rgb2lab(tmpImgt)
+
+			# nomalize image to range [0,1]
+			tmpImg[:,:,0] = (tmpImgt[:,:,0]-np.min(tmpImgt[:,:,0]))/(np.max(tmpImgt[:,:,0])-np.min(tmpImgt[:,:,0]))
+			tmpImg[:,:,1] = (tmpImgt[:,:,1]-np.min(tmpImgt[:,:,1]))/(np.max(tmpImgt[:,:,1])-np.min(tmpImgt[:,:,1]))
+			tmpImg[:,:,2] = (tmpImgt[:,:,2]-np.min(tmpImgt[:,:,2]))/(np.max(tmpImgt[:,:,2])-np.min(tmpImgt[:,:,2]))
+			tmpImg[:,:,3] = (tmpImgtl[:,:,0]-np.min(tmpImgtl[:,:,0]))/(np.max(tmpImgtl[:,:,0])-np.min(tmpImgtl[:,:,0]))
+			tmpImg[:,:,4] = (tmpImgtl[:,:,1]-np.min(tmpImgtl[:,:,1]))/(np.max(tmpImgtl[:,:,1])-np.min(tmpImgtl[:,:,1]))
+			tmpImg[:,:,5] = (tmpImgtl[:,:,2]-np.min(tmpImgtl[:,:,2]))/(np.max(tmpImgtl[:,:,2])-np.min(tmpImgtl[:,:,2]))
+
+			# tmpImg = tmpImg/(np.max(tmpImg)-np.min(tmpImg))
+
+			tmpImg[:,:,0] = (tmpImg[:,:,0]-np.mean(tmpImg[:,:,0]))/np.std(tmpImg[:,:,0])
+			tmpImg[:,:,1] = (tmpImg[:,:,1]-np.mean(tmpImg[:,:,1]))/np.std(tmpImg[:,:,1])
+			tmpImg[:,:,2] = (tmpImg[:,:,2]-np.mean(tmpImg[:,:,2]))/np.std(tmpImg[:,:,2])
+			tmpImg[:,:,3] = (tmpImg[:,:,3]-np.mean(tmpImg[:,:,3]))/np.std(tmpImg[:,:,3])
+			tmpImg[:,:,4] = (tmpImg[:,:,4]-np.mean(tmpImg[:,:,4]))/np.std(tmpImg[:,:,4])
+			tmpImg[:,:,5] = (tmpImg[:,:,5]-np.mean(tmpImg[:,:,5]))/np.std(tmpImg[:,:,5])
+
+		elif self.flag == 1: #with Lab color
+			tmpImg = np.zeros((image.shape[0],image.shape[1],3))
+
+			if image.shape[2]==1:
+				tmpImg[:,:,0] = image[:,:,0]
+				tmpImg[:,:,1] = image[:,:,0]
+				tmpImg[:,:,2] = image[:,:,0]
+			else:
+				tmpImg = image
+
+			tmpImg = color.rgb2lab(tmpImg)
+
+			# tmpImg = tmpImg/(np.max(tmpImg)-np.min(tmpImg))
+
+			tmpImg[:,:,0] = (tmpImg[:,:,0]-np.min(tmpImg[:,:,0]))/(np.max(tmpImg[:,:,0])-np.min(tmpImg[:,:,0]))
+			tmpImg[:,:,1] = (tmpImg[:,:,1]-np.min(tmpImg[:,:,1]))/(np.max(tmpImg[:,:,1])-np.min(tmpImg[:,:,1]))
+			tmpImg[:,:,2] = (tmpImg[:,:,2]-np.min(tmpImg[:,:,2]))/(np.max(tmpImg[:,:,2])-np.min(tmpImg[:,:,2]))
+
+			tmpImg[:,:,0] = (tmpImg[:,:,0]-np.mean(tmpImg[:,:,0]))/np.std(tmpImg[:,:,0])
+			tmpImg[:,:,1] = (tmpImg[:,:,1]-np.mean(tmpImg[:,:,1]))/np.std(tmpImg[:,:,1])
+			tmpImg[:,:,2] = (tmpImg[:,:,2]-np.mean(tmpImg[:,:,2]))/np.std(tmpImg[:,:,2])
+
+		else: # with rgb color
+			tmpImg = np.zeros((image.shape[0],image.shape[1],3))
+			image = image/np.max(image)
+			if image.shape[2]==1:
+				tmpImg[:,:,0] = (image[:,:,0]-0.485)/0.229
+				tmpImg[:,:,1] = (image[:,:,0]-0.485)/0.229
+				tmpImg[:,:,2] = (image[:,:,0]-0.485)/0.229
+			else:
+				tmpImg[:,:,0] = (image[:,:,0]-0.485)/0.229
+				tmpImg[:,:,1] = (image[:,:,1]-0.456)/0.224
+				tmpImg[:,:,2] = (image[:,:,2]-0.406)/0.225
+
+		tmpLbl[:,:,0] = label[:,:,0]
+
+
+		tmpImg = tmpImg.transpose((2, 0, 1))
+		tmpLbl = label.transpose((2, 0, 1))
+
+		return {'imidx':torch.from_numpy(imidx), 'image': torch.from_numpy(tmpImg), 'label': torch.from_numpy(tmpLbl)}
+
+class SalObjDataset(Dataset):
+	def __init__(self,img_name_list,lbl_name_list,transform=None):
+		# self.root_dir = root_dir
+		# self.image_name_list = glob.glob(image_dir+'*.png')
+		# self.label_name_list = glob.glob(label_dir+'*.png')
+		self.image_name_list = img_name_list
+		self.label_name_list = lbl_name_list
+		self.transform = transform
+
+	def __len__(self):
+		return len(self.image_name_list)
+
+	def __getitem__(self,idx):
+
+		# image = Image.open(self.image_name_list[idx])#io.imread(self.image_name_list[idx])
+		# label = Image.open(self.label_name_list[idx])#io.imread(self.label_name_list[idx])
+
+		image = io.imread(self.image_name_list[idx])
+		imname = self.image_name_list[idx]
+		imidx = np.array([idx])
+
+		if(0==len(self.label_name_list)):
+			label_3 = np.zeros(image.shape)
+		else:
+			label_3 = io.imread(self.label_name_list[idx])
+
+		label = np.zeros(label_3.shape[0:2])
+		if(3==len(label_3.shape)):
+			label = label_3[:,:,0]
+		elif(2==len(label_3.shape)):
+			label = label_3
+
+		if(3==len(image.shape) and 2==len(label.shape)):
+			label = label[:,:,np.newaxis]
+		elif(2==len(image.shape) and 2==len(label.shape)):
+			image = image[:,:,np.newaxis]
+			label = label[:,:,np.newaxis]
+
+		sample = {'imidx':imidx, 'image':image, 'label':label}
+
+		if self.transform:
+			sample = self.transform(sample)
+
+		return sample

u2net/model.py

@@ -0,0 +1,525 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class REBNCONV(nn.Module):
+    def __init__(self,in_ch=3,out_ch=3,dirate=1):
+        super(REBNCONV,self).__init__()
+
+        self.conv_s1 = nn.Conv2d(in_ch,out_ch,3,padding=1*dirate,dilation=1*dirate)
+        self.bn_s1 = nn.BatchNorm2d(out_ch)
+        self.relu_s1 = nn.ReLU(inplace=True)
+
+    def forward(self,x):
+
+        hx = x
+        xout = self.relu_s1(self.bn_s1(self.conv_s1(hx)))
+
+        return xout
+
+## upsample tensor 'src' to have the same spatial size with tensor 'tar'
+def _upsample_like(src,tar):
+
+    src = F.upsample(src,size=tar.shape[2:],mode='bilinear')
+
+    return src
+
+
+### RSU-7 ###
+class RSU7(nn.Module):#UNet07DRES(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU7,self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
+        self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool3 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool4 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool5 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv6 = REBNCONV(mid_ch,mid_ch,dirate=1)
+
+        self.rebnconv7 = REBNCONV(mid_ch,mid_ch,dirate=2)
+
+        self.rebnconv6d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv5d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
+
+    def forward(self,x):
+
+        hx = x
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+        hx = self.pool3(hx3)
+
+        hx4 = self.rebnconv4(hx)
+        hx = self.pool4(hx4)
+
+        hx5 = self.rebnconv5(hx)
+        hx = self.pool5(hx5)
+
+        hx6 = self.rebnconv6(hx)
+
+        hx7 = self.rebnconv7(hx6)
+
+        hx6d =  self.rebnconv6d(torch.cat((hx7,hx6),1))
+        hx6dup = _upsample_like(hx6d,hx5)
+
+        hx5d =  self.rebnconv5d(torch.cat((hx6dup,hx5),1))
+        hx5dup = _upsample_like(hx5d,hx4)
+
+        hx4d = self.rebnconv4d(torch.cat((hx5dup,hx4),1))
+        hx4dup = _upsample_like(hx4d,hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4dup,hx3),1))
+        hx3dup = _upsample_like(hx3d,hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
+        hx2dup = _upsample_like(hx2d,hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
+
+        return hx1d + hxin
+
+### RSU-6 ###
+class RSU6(nn.Module):#UNet06DRES(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU6,self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
+        self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool3 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool4 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=1)
+
+        self.rebnconv6 = REBNCONV(mid_ch,mid_ch,dirate=2)
+
+        self.rebnconv5d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
+
+    def forward(self,x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+        hx = self.pool3(hx3)
+
+        hx4 = self.rebnconv4(hx)
+        hx = self.pool4(hx4)
+
+        hx5 = self.rebnconv5(hx)
+
+        hx6 = self.rebnconv6(hx5)
+
+
+        hx5d =  self.rebnconv5d(torch.cat((hx6,hx5),1))
+        hx5dup = _upsample_like(hx5d,hx4)
+
+        hx4d = self.rebnconv4d(torch.cat((hx5dup,hx4),1))
+        hx4dup = _upsample_like(hx4d,hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4dup,hx3),1))
+        hx3dup = _upsample_like(hx3d,hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
+        hx2dup = _upsample_like(hx2d,hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
+
+        return hx1d + hxin
+
+### RSU-5 ###
+class RSU5(nn.Module):#UNet05DRES(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU5,self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
+        self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool3 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
+
+        self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=2)
+
+        self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
+
+    def forward(self,x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+        hx = self.pool3(hx3)
+
+        hx4 = self.rebnconv4(hx)
+
+        hx5 = self.rebnconv5(hx4)
+
+        hx4d = self.rebnconv4d(torch.cat((hx5,hx4),1))
+        hx4dup = _upsample_like(hx4d,hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4dup,hx3),1))
+        hx3dup = _upsample_like(hx3d,hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
+        hx2dup = _upsample_like(hx2d,hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
+
+        return hx1d + hxin
+
+### RSU-4 ###
+class RSU4(nn.Module):#UNet04DRES(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU4,self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
+        self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
+
+        self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=2)
+
+        self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
+
+    def forward(self,x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+
+        hx4 = self.rebnconv4(hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4,hx3),1))
+        hx3dup = _upsample_like(hx3d,hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
+        hx2dup = _upsample_like(hx2d,hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
+
+        return hx1d + hxin
+
+### RSU-4F ###
+class RSU4F(nn.Module):#UNet04FRES(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU4F,self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
+        self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=2)
+        self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=4)
+
+        self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=8)
+
+        self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=4)
+        self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=2)
+        self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
+
+    def forward(self,x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx2 = self.rebnconv2(hx1)
+        hx3 = self.rebnconv3(hx2)
+
+        hx4 = self.rebnconv4(hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4,hx3),1))
+        hx2d = self.rebnconv2d(torch.cat((hx3d,hx2),1))
+        hx1d = self.rebnconv1d(torch.cat((hx2d,hx1),1))
+
+        return hx1d + hxin
+
+
+##### U^2-Net ####
+class U2NET(nn.Module):
+
+    def __init__(self,in_ch=3,out_ch=1):
+        super(U2NET,self).__init__()
+
+        self.stage1 = RSU7(in_ch,32,64)
+        self.pool12 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage2 = RSU6(64,32,128)
+        self.pool23 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage3 = RSU5(128,64,256)
+        self.pool34 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage4 = RSU4(256,128,512)
+        self.pool45 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage5 = RSU4F(512,256,512)
+        self.pool56 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage6 = RSU4F(512,256,512)
+
+        # decoder
+        self.stage5d = RSU4F(1024,256,512)
+        self.stage4d = RSU4(1024,128,256)
+        self.stage3d = RSU5(512,64,128)
+        self.stage2d = RSU6(256,32,64)
+        self.stage1d = RSU7(128,16,64)
+
+        self.side1 = nn.Conv2d(64,out_ch,3,padding=1)
+        self.side2 = nn.Conv2d(64,out_ch,3,padding=1)
+        self.side3 = nn.Conv2d(128,out_ch,3,padding=1)
+        self.side4 = nn.Conv2d(256,out_ch,3,padding=1)
+        self.side5 = nn.Conv2d(512,out_ch,3,padding=1)
+        self.side6 = nn.Conv2d(512,out_ch,3,padding=1)
+
+        self.outconv = nn.Conv2d(6*out_ch,out_ch,1)
+
+    def forward(self,x):
+
+        hx = x
+
+        #stage 1
+        hx1 = self.stage1(hx)
+        hx = self.pool12(hx1)
+
+        #stage 2
+        hx2 = self.stage2(hx)
+        hx = self.pool23(hx2)
+
+        #stage 3
+        hx3 = self.stage3(hx)
+        hx = self.pool34(hx3)
+
+        #stage 4
+        hx4 = self.stage4(hx)
+        hx = self.pool45(hx4)
+
+        #stage 5
+        hx5 = self.stage5(hx)
+        hx = self.pool56(hx5)
+
+        #stage 6
+        hx6 = self.stage6(hx)
+        hx6up = _upsample_like(hx6,hx5)
+
+        #-------------------- decoder --------------------
+        hx5d = self.stage5d(torch.cat((hx6up,hx5),1))
+        hx5dup = _upsample_like(hx5d,hx4)
+
+        hx4d = self.stage4d(torch.cat((hx5dup,hx4),1))
+        hx4dup = _upsample_like(hx4d,hx3)
+
+        hx3d = self.stage3d(torch.cat((hx4dup,hx3),1))
+        hx3dup = _upsample_like(hx3d,hx2)
+
+        hx2d = self.stage2d(torch.cat((hx3dup,hx2),1))
+        hx2dup = _upsample_like(hx2d,hx1)
+
+        hx1d = self.stage1d(torch.cat((hx2dup,hx1),1))
+
+
+        #side output
+        d1 = self.side1(hx1d)
+
+        d2 = self.side2(hx2d)
+        d2 = _upsample_like(d2,d1)
+
+        d3 = self.side3(hx3d)
+        d3 = _upsample_like(d3,d1)
+
+        d4 = self.side4(hx4d)
+        d4 = _upsample_like(d4,d1)
+
+        d5 = self.side5(hx5d)
+        d5 = _upsample_like(d5,d1)
+
+        d6 = self.side6(hx6)
+        d6 = _upsample_like(d6,d1)
+
+        d0 = self.outconv(torch.cat((d1,d2,d3,d4,d5,d6),1))
+
+        return F.sigmoid(d0), F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)
+
+### U^2-Net small ###
+class U2NETP(nn.Module):
+
+    def __init__(self,in_ch=3,out_ch=1):
+        super(U2NETP,self).__init__()
+
+        self.stage1 = RSU7(in_ch,16,64)
+        self.pool12 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage2 = RSU6(64,16,64)
+        self.pool23 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage3 = RSU5(64,16,64)
+        self.pool34 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage4 = RSU4(64,16,64)
+        self.pool45 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage5 = RSU4F(64,16,64)
+        self.pool56 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage6 = RSU4F(64,16,64)
+
+        # decoder
+        self.stage5d = RSU4F(128,16,64)
+        self.stage4d = RSU4(128,16,64)
+        self.stage3d = RSU5(128,16,64)
+        self.stage2d = RSU6(128,16,64)
+        self.stage1d = RSU7(128,16,64)
+
+        self.side1 = nn.Conv2d(64,out_ch,3,padding=1)
+        self.side2 = nn.Conv2d(64,out_ch,3,padding=1)
+        self.side3 = nn.Conv2d(64,out_ch,3,padding=1)
+        self.side4 = nn.Conv2d(64,out_ch,3,padding=1)
+        self.side5 = nn.Conv2d(64,out_ch,3,padding=1)
+        self.side6 = nn.Conv2d(64,out_ch,3,padding=1)
+
+        self.outconv = nn.Conv2d(6*out_ch,out_ch,1)
+
+    def forward(self,x):
+
+        hx = x
+
+        #stage 1
+        hx1 = self.stage1(hx)
+        hx = self.pool12(hx1)
+
+        #stage 2
+        hx2 = self.stage2(hx)
+        hx = self.pool23(hx2)
+
+        #stage 3
+        hx3 = self.stage3(hx)
+        hx = self.pool34(hx3)
+
+        #stage 4
+        hx4 = self.stage4(hx)
+        hx = self.pool45(hx4)
+
+        #stage 5
+        hx5 = self.stage5(hx)
+        hx = self.pool56(hx5)
+
+        #stage 6
+        hx6 = self.stage6(hx)
+        hx6up = _upsample_like(hx6,hx5)
+
+        #decoder
+        hx5d = self.stage5d(torch.cat((hx6up,hx5),1))
+        hx5dup = _upsample_like(hx5d,hx4)
+
+        hx4d = self.stage4d(torch.cat((hx5dup,hx4),1))
+        hx4dup = _upsample_like(hx4d,hx3)
+
+        hx3d = self.stage3d(torch.cat((hx4dup,hx3),1))
+        hx3dup = _upsample_like(hx3d,hx2)
+
+        hx2d = self.stage2d(torch.cat((hx3dup,hx2),1))
+        hx2dup = _upsample_like(hx2d,hx1)
+
+        hx1d = self.stage1d(torch.cat((hx2dup,hx1),1))
+
+
+        #side output
+        d1 = self.side1(hx1d)
+
+        d2 = self.side2(hx2d)
+        d2 = _upsample_like(d2,d1)
+
+        d3 = self.side3(hx3d)
+        d3 = _upsample_like(d3,d1)
+
+        d4 = self.side4(hx4d)
+        d4 = _upsample_like(d4,d1)
+
+        d5 = self.side5(hx5d)
+        d5 = _upsample_like(d5,d1)
+
+        d6 = self.side6(hx6)
+        d6 = _upsample_like(d6,d1)
+
+        d0 = self.outconv(torch.cat((d1,d2,d3,d4,d5,d6),1))
+
+        return F.sigmoid(d0), F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)

utils.py

@@ -0,0 +1,515 @@
+import logging
+import math
+import os
+from PIL import Image
+import yaml
+
+from sklearn.metrics import confusion_matrix
+import torch
+import torch.nn.functional as F
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel
+from torchvision import transforms
+
+from moco.loader import GaussianBlur
+import numpy as np
+from augmentations import JigsawPuzzle, JigsawPuzzle_l, JigsawPuzzle_all, RandomErasing, RandomPatchNoise, RandomPatchErase
+
+
+LOG_FORMAT = "[%(levelname)s] %(asctime)s %(filename)s:%(lineno)s %(message)s"
+LOG_DATEFMT = "%Y-%m-%d %H:%M:%S"
+
+NUM_CLASSES = {"domainnet-126": 126, "VISDA-C": 12, "PACS": 7}
+
+import torch
+import numpy as np
+from PIL import Image
+
+
+def configure_logger(rank, log_path=None):
+    if log_path:
+        log_dir = os.path.dirname(log_path)
+        os.makedirs(log_dir, exist_ok=True)
+
+    # only master process will print & write
+    level = logging.INFO if rank in {-1, 0} else logging.WARNING
+    handlers = [logging.StreamHandler()]
+    if rank in {0, -1} and log_path:
+        handlers.append(logging.FileHandler(log_path, "w"))
+
+    logging.basicConfig(
+        level=level,
+        format=LOG_FORMAT,
+        datefmt=LOG_DATEFMT,
+        handlers=handlers,
+    )
+
+
+class UnevenBatchLoader:
+    """Loader that loads data from multiple datasets with different length."""
+
+    def __init__(self, data_loaders, is_ddp=False):
+        # register N data loaders with epoch counters.
+        self.data_loaders = data_loaders
+        self.epoch_counters = [0 for _ in range(len(data_loaders))]
+
+        # set_epoch() needs to be called before creating the iterator
+        self.is_ddp = is_ddp
+        if is_ddp:
+            for data_loader in data_loaders:
+                data_loader.sampler.set_epoch(0)
+        self.iterators = [iter(data_loader) for data_loader in data_loaders]
+
+    def next_batch(self):
+        """Load the next batch by collecting from N data loaders.
+        Args:
+            None
+        Returns:
+            data: a list of N items from N data loaders. each item has the format
+                output by a single data loader.
+        """
+        data = []
+        for i, iterator in enumerate(self.iterators):
+            try:
+                batch_i = next(iterator)
+            except StopIteration:
+                self.epoch_counters[i] += 1
+                # create a new iterator
+                if self.is_ddp:
+                    self.data_loaders[i].sampler.set_epoch(self.epoch_counters[i])
+                new_iterator = iter(self.data_loaders[i])
+                self.iterators[i] = new_iterator
+                batch_i = next(new_iterator)
+            data.append(batch_i)
+
+        return data
+
+    def update_loader(self, idx, loader, epoch=None):
+        if self.is_ddp and isinstance(epoch, int):
+            loader.sampler.set_epoch(epoch)
+        self.iterators[idx] = iter(loader)
+
+
+class CustomDistributedDataParallel(DistributedDataParallel):
+    """A wrapper class over DDP that relay "module" attribute."""
+
+    def __init__(self, model, **kwargs):
+        super(CustomDistributedDataParallel, self).__init__(model, **kwargs)
+
+    def __getattr__(self, name):
+        try:
+            return super(CustomDistributedDataParallel, self).__getattr__(name)
+        except AttributeError:
+            return getattr(self.module, name)
+
+
+@torch.no_grad()
+def concat_all_gather(tensor):
+    """
+    Performs all_gather operation on the provided tensors.
+    *** Warning ***: torch.distributed.all_gather has no gradient.
+    """
+    tensors_gather = [torch.ones_like(tensor) for _ in range(dist.get_world_size())]
+    dist.all_gather(tensors_gather, tensor, async_op=False)
+
+    output = torch.cat(tensors_gather, dim=0)
+    return output
+
+
+@torch.no_grad()
+def remove_wrap_arounds(tensor, ranks):
+    if ranks == 0:
+        return tensor
+
+    world_size = dist.get_world_size()
+    single_length = len(tensor) // world_size
+    output = []
+    for rank in range(world_size):
+        sub_tensor = tensor[rank * single_length : (rank + 1) * single_length]
+        if rank >= ranks:
+            output.append(sub_tensor[:-1])
+        else:
+            output.append(sub_tensor)
+    output = torch.cat(output)
+
+    return output
+
+
+def get_categories(category_file):
+    """Return a list of categories ordered by corresponding label.
+
+    Args:
+        category_file: str, path to the category file. can be .yaml or .txt
+
+    Returns:
+        categories: List[str], a list of categories ordered by label.
+    """
+    if category_file.endswith(".yaml"):
+        with open(category_file, "r") as fd:
+            cat_mapping = yaml.load(fd, Loader=yaml.SafeLoader)
+        categories = list(cat_mapping.keys())
+        categories.sort(key=lambda x: cat_mapping[x])
+    elif category_file.endswith(".txt"):
+        with open(category_file, "r") as fd:
+            categories = fd.readlines()
+        categories = [cat.strip() for cat in categories if cat]
+    else:
+        raise NotImplementedError()
+
+    categories = [cat.replace("_", " ") for cat in categories]
+    return categories
+
+
+def get_augmentation(aug_type, patch_height=28, mix_prob=0.8, normalize=None):
+    if not normalize:
+        normalize = transforms.Normalize(
+            mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
+        )
+    if aug_type == "moco-v2":
+        image_aug = transforms.Compose(
+            [
+                transforms.RandomResizedCrop(224, scale=(0.2, 1.0)),
+                transforms.RandomApply(
+                    [transforms.ColorJitter(0.4, 0.4, 0.4, 0.1)],
+                    p=0.8,  # not strengthened
+                ),
+                transforms.RandomGrayscale(p=0.2),
+                transforms.RandomApply([GaussianBlur([0.1, 2.0])], p=0.5),
+                transforms.RandomHorizontalFlip(),
+                transforms.ToTensor(),
+                normalize,
+            ]
+        )
+    elif aug_type == "moco-v1":
+        image_aug = transforms.Compose(
+            [
+                transforms.RandomResizedCrop(224, scale=(0.2, 1.0)),
+                transforms.RandomGrayscale(p=0.2),
+                transforms.ColorJitter(0.4, 0.4, 0.4, 0.4),
+                transforms.RandomHorizontalFlip(),
+                transforms.ToTensor(),
+                normalize,
+            ]
+        )
+    elif aug_type == "plain":
+        image_aug = transforms.Compose(
+            [
+                transforms.Resize((256, 256)),
+                transforms.RandomCrop(224),
+                transforms.RandomHorizontalFlip(),
+                transforms.ToTensor(),
+                normalize,
+            ]
+        )
+    elif aug_type == "clip_inference":
+        image_aug = transforms.Compose(
+            [
+                transforms.Resize(224, interpolation=Image.BICUBIC),
+                transforms.CenterCrop(224),
+                transforms.ToTensor(),
+                normalize,
+            ]
+        )
+    elif aug_type == "test":
+        image_aug = transforms.Compose(
+            [
+                transforms.Resize((256, 256)),
+                transforms.CenterCrop(224),
+                transforms.ToTensor(),
+                normalize,
+            ]
+        )
+    elif aug_type == "jigsaw":
+        image_aug =  transforms.Compose(
+            [
+                transforms.Resize((256, 256)),
+                transforms.CenterCrop(224),
+                # transforms.RandomHorizontalFlip(),
+                JigsawPuzzle(patch_height=patch_height, patch_width=patch_height, mix_prob=1),
+                transforms.ToTensor(),
+                normalize,
+            ]
+        )
+    elif aug_type == "jigsaw_all":
+        image_aug =  transforms.Compose(
+            [
+                transforms.Resize((256, 256)),
+                transforms.CenterCrop(224),
+                # transforms.RandomHorizontalFlip(),
+                JigsawPuzzle_all(patch_height=patch_height, patch_width=patch_height, mix_prob=1),
+                transforms.ToTensor(),
+                normalize,
+            ]
+        )
+    elif aug_type == "jigsaw_l":
+        image_aug = transforms.Compose(
+            [
+                transforms.Resize((256, 256)),
+                transforms.CenterCrop(224),
+                # transforms.RandomHorizontalFlip(),
+                JigsawPuzzle_l(patch_height=patch_height, patch_width=patch_height, mix_prob=1),
+                transforms.ToTensor(),
+                normalize,
+            ]
+        )
+    elif aug_type == "rpe":
+        image_aug = transforms.Compose(
+            [
+                transforms.Resize((256, 256)),
+                transforms.CenterCrop(224),
+                # transforms.RandomHorizontalFlip(),
+                RandomPatchErase(patch_height=patch_height, patch_width=patch_height, mix_prob=1),
+                transforms.ToTensor(),
+                normalize,
+            ]
+        )
+    elif aug_type == "rpn":
+        image_aug = transforms.Compose(
+            [
+                transforms.Resize((256, 256)),
+                transforms.CenterCrop(224),
+                # transforms.RandomHorizontalFlip(),
+                RandomPatchNoise(patch_height=patch_height, patch_width=patch_height, mix_prob=1),
+                transforms.ToTensor(),
+                normalize,
+            ]
+        )
+    elif aug_type in ["ours", "ours_1"]:
+        image_aug = transforms.Compose(
+            [
+                transforms.Resize((256, 256)),
+                transforms.CenterCrop(224),
+                JigsawPuzzle_all(patch_height=patch_height, patch_width=patch_height, mix_prob=mix_prob),
+                transforms.ToTensor(),
+            ]
+        )
+    else:
+        image_aug = None
+
+    return DualTransform(
+        aug_type=aug_type,
+        image_transform=image_aug,
+        patch_height=patch_height, 
+        patch_width=patch_height,
+        mix_prob=mix_prob,
+    )
+
+def fuse_foreground_background(img1, img2, mask):
+    """
+    Given a (C,H,W) image tensor and a (possibly 2D) mask,
+    multiply img by mask to black out the background.
+    Expects 0 as background in the mask.
+    """
+    mask = (mask > 0.5)
+    output = img1 * mask + img2 * (~mask)
+
+    return output
+
+def normalize(tensor):
+    T = transforms.Normalize(
+            mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
+        )
+    return T(tensor)
+
+class DualTransform:
+    """
+    A wrapper that can apply image-only transforms or image+mask transforms.
+    """
+    def __init__(self, aug_type, image_transform=None, patch_height=28, patch_width=28,mix_prob=1.0):
+        self.image_transform = image_transform
+        self.aug_type = aug_type
+        self.base_transform = transforms.Compose(
+            [
+                transforms.Resize((256, 256)),
+                transforms.CenterCrop(224),
+                transforms.ToTensor(),
+            ]
+        )
+
+        self.moco_transform = transforms.Compose(
+            [   
+                transforms.RandomResizedCrop(224, scale=(0.2, 1.0)),
+                transforms.RandomApply(
+                    [transforms.ColorJitter(0.4, 0.4, 0.4, 0.1)],
+                    p=0.8,  # not strengthened
+                ),
+                transforms.RandomGrayscale(p=0.2),
+                transforms.RandomApply([GaussianBlur([0.1, 2.0])], p=0.5),
+                transforms.RandomHorizontalFlip(),
+                # RandomErasing(mode='soft_pixel'),
+                transforms.ToTensor(),
+                normalize,
+            ]
+        )
+        self.fpn = RandomPatchNoise(patch_height=28, patch_width=28, mix_prob=mix_prob)
+        self.to_pil = transforms.ToPILImage()
+        self.to_tensor = transforms.ToTensor()
+        self.jigsaw = JigsawPuzzle(patch_height=28, patch_width=28, mix_prob=mix_prob)
+        self.jigsaw_all = JigsawPuzzle_all(mix_prob=mix_prob)
+
+    def __call__(self, img, mask=None):
+        if self.aug_type == "mask":
+            mask = self.base_transform(mask)
+            return normalize(mask)
+        elif self.aug_type == "foreground":
+            mask = self.base_transform(mask)
+            img = self.base_transform(img)
+            return normalize(img * (mask>0.5).float())
+        elif self.aug_type == "fpn":
+            mask = self.base_transform(mask)
+            img = self.base_transform(img)
+            img_n = self.to_tensor(self.fpn(self.to_pil(img)))
+            return normalize(img_n * (mask>0.5).float())
+        elif self.aug_type == "bps":
+            mask = self.base_transform(mask)
+            img = self.base_transform(img)
+            img_jigsaw = self.to_tensor(self.jigsaw_all(self.to_pil(img)))
+            return normalize(img_jigsaw * (mask<0.5).float())
+        elif self.aug_type == "ours_raw":
+            mask = self.base_transform(mask)
+            img = self.base_transform(img)
+            img_n = self.to_tensor(self.fpn(self.to_pil(img)))
+            img_jigsaw = self.to_tensor(self.jigsaw_all(self.to_pil(img)))
+            img_out = fuse_foreground_background(img_n, img_jigsaw, mask)
+            return normalize(img_out)
+        elif self.aug_type == "ours":
+            mask = self.base_transform(mask)
+            img = self.base_transform(img)
+            img_n = self.to_tensor(self.fpn(self.to_pil(img)))
+            img_jigsaw = self.to_tensor(self.jigsaw_all(self.to_pil(img)))
+            img_out = fuse_foreground_background(img_n, img_jigsaw, mask)
+            return self.moco_transform(self.to_pil(img_out))
+        elif self.aug_type == "ours_fpn":
+            mask = self.base_transform(mask)
+            img = self.base_transform(img)
+            img_n = self.to_tensor(self.fpn(self.to_pil(img)))
+            img_out = fuse_foreground_background(img_n, img, mask)
+            return self.moco_transform(self.to_pil(img_out))
+        elif self.aug_type == "ours_bps":
+            mask = self.base_transform(mask)
+            img = self.base_transform(img)
+            # img_n = self.to_tensor(self.fpn(self.to_pil(img)))
+            img_jigsaw = self.to_tensor(self.jigsaw_all(self.to_pil(img)))
+            img_out = fuse_foreground_background(img, img_jigsaw, mask)
+            return self.moco_transform(self.to_pil(img_out))
+        # Always transform the image if we have an image_transform
+        else:
+            return self.image_transform(img)
+        '''
+        elif self.aug_type == "ours_old":
+            img_t = self.image_transform(img)
+            img = self.base_transform(img)
+            mask = self.base_transform(mask)
+            img_t1 = fuse_foreground_background(img, img_t, mask)
+            img_t1_pil = self.to_pil(img_t1)
+            output = self.moco_transform(img_t1_pil)
+            return output
+        elif self.aug_type == "ours_1":
+            img_t = self.image_transform(img)
+            img = self.base_transform(img)
+            mask = self.base_transform(mask)
+            img_t1 = fuse_foreground_background(img, img_t, mask)
+            return normalize(img_t1)'
+        '''
+    
+class AverageMeter(object):
+    """Computes and stores the average and current value"""
+
+    def __init__(self, name, fmt=":f"):
+        self.name = name
+        self.fmt = fmt
+        self.reset()
+
+    def reset(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, n=1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
+
+    def __str__(self):
+        fmtstr = "{name} {val" + self.fmt + "} ({avg" + self.fmt + "})"
+        return fmtstr.format(**self.__dict__)
+
+
+class ProgressMeter(object):
+    def __init__(self, num_batches, meters, prefix=""):
+        self.batch_fmtstr = self._get_batch_fmtstr(num_batches)
+        self.meters = meters
+        self.prefix = prefix
+
+    def display(self, batch):
+        entries = [self.prefix + self.batch_fmtstr.format(batch)]
+        entries += [str(meter) for meter in self.meters]
+        logging.info("\t".join(entries))
+
+    def _get_batch_fmtstr(self, num_batches):
+        num_digits = len(str(num_batches // 1))
+        fmt = "{:" + str(num_digits) + "d}"
+        return "[" + fmt + "/" + fmt.format(num_batches) + "]"
+
+
+def save_checkpoint(model, optimizer, epoch, save_path="checkpoint.pth.tar"):
+    state = {
+        "state_dict": model.state_dict(),
+        "optimizer": optimizer.state_dict(),
+        "epoch": epoch,
+    }
+    torch.save(state, save_path)
+
+
+def adjust_learning_rate(optimizer, progress, args):
+    """
+    Decay the learning rate based on epoch or iteration.
+    """
+    if args.optim.cos:
+        decay = 0.5 * (1.0 + math.cos(math.pi * progress / args.learn.full_progress))
+    elif args.optim.exp:
+        decay = (1 + 10 * progress / args.learn.full_progress) ** -0.75
+    else:
+        decay = 1.0
+        for milestone in args.optim.schedule:
+            decay *= args.optim.gamma if progress >= milestone else 1.0
+    for param_group in optimizer.param_groups:
+        param_group["lr"] = param_group["lr0"] * decay
+
+    return decay
+
+
+def per_class_accuracy(y_true, y_pred):
+    matrix = confusion_matrix(y_true, y_pred)
+    acc_per_class = (matrix.diagonal() / matrix.sum(axis=1) * 100.0).round(2)
+    logging.info(
+        f"Accuracy per class: {acc_per_class}, mean: {acc_per_class.mean().round(2)}"
+    )
+
+    return acc_per_class
+
+
+def get_distances(X, Y, dist_type="euclidean"):
+    """
+    Args:
+        X: (N, D) tensor
+        Y: (M, D) tensor
+    """
+    if dist_type == "euclidean":
+        distances = torch.cdist(X, Y)
+    elif dist_type == "cosine":
+        distances = 1 - torch.matmul(F.normalize(X, dim=1), F.normalize(Y, dim=1).T)
+    else:
+        raise NotImplementedError(f"{dist_type} distance not implemented.")
+
+    return distances
+
+
+def is_master(args):
+    return args.rank % args.ngpus_per_node == 0
+
+
+def use_wandb(args):
+    return is_master(args) and args.use_wandb