Init

app.py

@@ -0,0 +1,124 @@
+import gradio as gr
+from gradio_image_prompter import ImagePrompter
+from detectron2.config import LazyConfig, instantiate
+from detectron2.checkpoint import DetectionCheckpointer
+import cv2
+import numpy as np
+import torch
+from huggingface_hub import hf_hub_download
+
+DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+model_choice = {
+    'SAM': None, 
+    'HQ-SAM': None, 
+    'SAM2': None
+}
+
+for model_type in model_choice.keys():
+    model_choice[model_type] = hf_hub_download(repo_id="XiaRho/SEMat", filename=f"SEMat_{model_type}.pth", repo_type="model")
+
+def load_model(model_type='SAM2'):
+    assert model_type in model_choice.keys()
+    config_path = './configs/SEMat_{}.py'.format(model_type)
+    cfg = LazyConfig.load(config_path)
+
+    if hasattr(cfg.model.sam_model, 'ckpt_path'):
+        cfg.model.sam_model.ckpt_path = None
+    else:
+        cfg.model.sam_model.checkpoint = None
+    model = instantiate(cfg.model)
+    if model.lora_rank is not None:
+        model.init_lora()
+    model.to(DEVICE)
+    DetectionCheckpointer(model).load(model_choice[model_type])
+    model.eval()
+    return model, model_type
+
+def transform_image_bbox(prompts):
+    if len(prompts["points"]) != 1:
+        raise gr.Error("Please input only one BBox.", duration=5)
+    [[x1, y1, idx_3, x2, y2, idx_6]] = prompts["points"]
+    if idx_3 != 2 or idx_6 != 3:
+        raise gr.Error("Please input BBox instead of point.", duration=5)
+    x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
+
+    img = prompts["image"]
+    ori_H, ori_W, _ = img.shape
+
+    scale = 1024 * 1.0 / max(ori_H, ori_W)
+    new_H, new_W = ori_H * scale, ori_W * scale
+    new_W = int(new_W + 0.5)
+    new_H = int(new_H + 0.5)
+
+    img = cv2.resize(img, (new_W, new_H), interpolation=cv2.INTER_LINEAR)
+    padding = np.zeros([1024, 1024, 3], dtype=img.dtype)
+    padding[: new_H, : new_W, :] = img
+    img = padding
+    # img = img[:, :, ::-1].transpose((2, 0, 1)).astype(np.float32) / 255.0
+    img = img.transpose((2, 0, 1)).astype(np.float32) / 255.0
+
+    [[x1, y1, _, x2, y2, _]] = prompts["points"]
+    x1, y1, x2, y2 = int(x1 * scale + 0.5), int(y1 * scale + 0.5), int(x2 * scale + 0.5), int(y2 * scale + 0.5)
+    bbox = np.clip(np.array([[x1, y1, x2, y2]]) * 1.0, 0, 1023.0)
+
+    return img, bbox, (ori_H, ori_W), (new_H, new_W)
+
+if __name__ == '__main__':
+
+    model, model_type = load_model()
+
+    def inference_image(prompts, input_model_type):
+
+        global model_type
+        global model
+
+        if input_model_type != model_type:
+            gr.Info('Loading SEMat of {} version.'.format(input_model_type), duration=5)
+            _model, _ = load_model(input_model_type)
+            model_type = input_model_type
+            model = _model
+
+        image, bbox, ori_H_W, pad_H_W = transform_image_bbox(prompts)
+        input_data = {
+            'image': torch.from_numpy(image)[None].to(model.device),
+            'bbox': torch.from_numpy(bbox)[None].to(model.device),
+        }
+
+        with torch.no_grad():
+            inputs = model.preprocess_inputs(input_data) 
+            images, bbox, gt_alpha, trimap, condition = inputs['images'], inputs['bbox'], inputs['alpha'], inputs['trimap'], inputs['condition']
+
+            if model.backbone_condition:
+                condition_proj = model.condition_embedding(condition) 
+            elif model.backbone_bbox_prompt is not None or model.bbox_prompt_all_block is not None:
+                condition_proj = bbox
+            else:
+                condition_proj = None
+
+            low_res_masks, pred_alphas, pred_trimap, sam_hq_matting_token = model.forward_samhq_and_matting_decoder(images, bbox, condition_proj)
+
+
+        output_alpha = np.uint8(pred_alphas[0, 0][:pad_H_W[0], :pad_H_W[1], None].repeat(1, 1, 3).cpu().numpy() * 255)
+
+        return output_alpha
+
+    with gr.Blocks() as demo:
+
+        with gr.Row():
+            with gr.Column(scale=45):
+                img_in = ImagePrompter(type='numpy', show_label=False, label="query image")
+                
+            with gr.Column(scale=45):
+                img_out = gr.Image(type='pil', label="output")
+
+        with gr.Row():
+            with gr.Column(scale=45):
+                input_model_type = gr.Dropdown(list(model_choice.keys()), value='SAM2', label="Trained SEMat Version")
+
+            with gr.Column(scale=45):
+                bt = gr.Button()
+
+        bt.click(inference_image, inputs=[img_in, input_model_type], outputs=[img_out]) 
+
+demo.launch()
+

configs/SEMat_HQ-SAM.py

@@ -0,0 +1,48 @@
+from .common.train import train
+from .semantic_enhanced_matting.model import model
+from .common.optimizer import optimizer
+from .common.scheduler import lr_multiplier
+from .semantic_enhanced_matting.dataloader import dataloader
+from modeling.decoder.unet_detail_capture import MattingDetailDecoder
+from detectron2.config import LazyCall as L
+
+model.sam_model.model_type = 'vit_l'
+model.sam_model.checkpoint = None
+model.vis_period = 200
+model.output_dir = '?'
+
+train.max_iter = 60000
+train.eval_period = int(train.max_iter * 1 / 10)
+train.checkpointer.period = int(train.max_iter * 1 / 10)
+train.checkpointer.max_to_keep = 1
+
+optimizer.lr = 5e-5
+
+lr_multiplier.scheduler.values = [1.0, 0.5, 0.2]
+lr_multiplier.scheduler.milestones = [0.5, 0.75]
+lr_multiplier.scheduler.num_updates = train.max_iter
+lr_multiplier.warmup_length = 250 / train.max_iter
+
+train.output_dir = './work_dirs/SEMat_HQ-SAM'
+
+model.lora_rank = 16
+model.lora_alpha = 16
+model.matting_decoder = L(MattingDetailDecoder)(
+    vit_intern_feat_in = 1024,
+    vit_intern_feat_index = [0, 1, 2, 3],
+    norm_type = 'SyncBN',
+    block_num = 2,
+    img_feat_in = 6,
+    norm_mask_logits = 6.5
+)
+model.backbone_bbox_prompt = 'bbox'
+model.backbone_bbox_prompt_loc = [2, 3]
+model.backbone_bbox_prompt_loss_weight = 1.0
+model.matting_token = True
+model.sam_model.matting_token = 3
+model.sam_model.frozen_decoder = True
+model.sam_hq_token_reg = 0.2
+model.reg_w_bce_loss = True
+model.matting_token_sup = 'trimap'
+model.matting_token_sup_loss_weight = 0.05
+model.trimap_loss_type = 'NGHM'

configs/SEMat_SAM.py

@@ -0,0 +1,51 @@
+from .common.train import train
+from .semantic_enhanced_matting.model import model
+from .common.optimizer import optimizer
+from .common.scheduler import lr_multiplier
+from .semantic_enhanced_matting.dataloader import dataloader
+from modeling.decoder.unet_detail_capture import MattingDetailDecoder
+from detectron2.config import LazyCall as L
+
+model.sam_model.model_type = 'vit_l'
+model.sam_model.checkpoint = None
+model.vis_period = 200
+model.output_dir = '?'
+
+train.max_iter = 60000
+train.eval_period = int(train.max_iter * 1 / 10)
+train.checkpointer.period = int(train.max_iter * 1 / 10)
+train.checkpointer.max_to_keep = 1
+
+optimizer.lr = 5e-5
+
+lr_multiplier.scheduler.values = [1.0, 0.5, 0.2]
+lr_multiplier.scheduler.milestones = [0.5, 0.75]
+lr_multiplier.scheduler.num_updates = train.max_iter
+lr_multiplier.warmup_length = 250 / train.max_iter
+
+train.output_dir = './work_dirs/SEMat_SAM'
+
+model.lora_rank = 16
+model.lora_alpha = 16
+model.matting_decoder = L(MattingDetailDecoder)(
+    vit_intern_feat_in = 1024,
+    vit_intern_feat_index = [0, 1, 2, 3],
+    norm_type = 'SyncBN',
+    block_num = 2,
+    img_feat_in = 6,
+    norm_mask_logits = 6.5
+)
+model.backbone_bbox_prompt = 'bbox'
+model.backbone_bbox_prompt_loc = [2, 3]
+model.backbone_bbox_prompt_loss_weight = 1.0
+model.matting_token = True
+model.sam_model.matting_token = 3
+model.sam_model.frozen_decoder = True
+model.sam_hq_token_reg = 0.2
+model.reg_on_sam_logits = True
+model.reg_w_bce_loss = True
+model.matting_token_sup = 'trimap'
+model.matting_token_sup_loss_weight = 0.05
+model.trimap_loss_type = 'NGHM'
+model.sam_model.wo_hq = True
+model.sam_model.mask_matting_res_add = False

configs/SEMat_SAM2.py

@@ -0,0 +1,57 @@
+from .common.train import train
+from .semantic_enhanced_matting.model import model
+from .common.optimizer import optimizer
+from .common.scheduler import lr_multiplier
+from .semantic_enhanced_matting.dataloader import dataloader
+from modeling.decoder.unet_detail_capture import MattingDetailDecoder
+from detectron2.config import LazyCall as L
+from sam2.build_sam import build_sam2
+
+model.sam_model.model_type = 'vit_l'
+model.sam_model.checkpoint = None
+model.vis_period = 200
+model.output_dir = '?'
+
+train.max_iter = 60000
+train.eval_period = int(train.max_iter * 1 / 10)
+train.checkpointer.period = int(train.max_iter * 1 / 10)
+train.checkpointer.max_to_keep = 1
+
+optimizer.lr = 5e-5
+
+lr_multiplier.scheduler.values = [1.0, 0.5, 0.2]
+lr_multiplier.scheduler.milestones = [0.5, 0.75]
+lr_multiplier.scheduler.num_updates = train.max_iter
+lr_multiplier.warmup_length = 250 / train.max_iter
+
+train.output_dir = './work_dirs/SEMat_SAM2'
+
+model.sam2 = True
+model.sam_model = L(build_sam2)(
+    config_file = 'sam2_hiera_l.yaml',
+    ckpt_path = None,
+    device = "cuda",
+    bbox_mask_matting_token = True,
+    mode="train",
+    upscaled_embedding_res_add = False
+)
+model.lora_rank = 16
+model.lora_alpha = 16
+model.matting_decoder = L(MattingDetailDecoder)(
+    vit_intern_feat_in = 1024,
+    vit_intern_feat_index = [0, 1, 2, 3],
+    norm_type = 'SyncBN',
+    block_num = 2,
+    img_feat_in = 6,
+    norm_mask_logits = 6.5,
+    sam2_multi_scale_feates = True
+)
+model.backbone_bbox_prompt = 'bbox'
+model.backbone_bbox_prompt_loc = [2, 3]
+model.backbone_bbox_prompt_loss_weight = 1.0
+model.matting_token = True
+model.sam_hq_token_reg = 0.2
+model.reg_w_bce_loss = True
+model.matting_token_sup = 'trimap'
+model.matting_token_sup_loss_weight = 0.05
+model.trimap_loss_type = 'NGHM'

configs/common/optimizer.py

@@ -0,0 +1,26 @@
+from detectron2 import model_zoo
+from functools import partial
+
+def get_vit_lr_decay_rate(name, lr_decay_rate=1.0, num_layers=12):
+    """
+    Calculate lr decay rate for different ViT blocks.
+    Args:
+        name (string): parameter name.
+        lr_decay_rate (float): base lr decay rate.
+        num_layers (int): number of ViT blocks.
+
+    Returns:
+        lr decay rate for the given parameter.
+    """
+    layer_id = num_layers + 1
+    if name.startswith("backbone"):
+        if ".pos_embed" in name or ".patch_embed" in name:
+            layer_id = 0
+        elif ".blocks." in name and ".residual." not in name:
+            layer_id = int(name[name.find(".blocks.") :].split(".")[2]) + 1
+    return lr_decay_rate ** (num_layers + 1 - layer_id)
+
+# Optimizer
+optimizer = model_zoo.get_config("common/optim.py").AdamW
+optimizer.params.lr_factor_func = partial(get_vit_lr_decay_rate, num_layers=12, lr_decay_rate=0.65)
+optimizer.params.overrides = {"pos_embed": {"weight_decay": 0.0}}

configs/common/scheduler.py

@@ -0,0 +1,13 @@
+from detectron2.config import LazyCall as L
+from detectron2.solver import WarmupParamScheduler
+from fvcore.common.param_scheduler import MultiStepParamScheduler
+
+lr_multiplier = L(WarmupParamScheduler)(
+    scheduler=L(MultiStepParamScheduler)(
+        values=[1.0, 0.1, 0.01],
+        milestones=[96778, 103579],
+        num_updates=100,
+    ),
+    warmup_length=250 / 100,
+    warmup_factor=0.001,
+)

configs/common/train.py

@@ -0,0 +1,17 @@
+train = dict(
+    output_dir="./output",
+    init_checkpoint="",
+    max_iter=90000,
+    amp=dict(enabled=False),  # options for Automatic Mixed Precision
+    ddp=dict(  # options for DistributedDataParallel
+        broadcast_buffers=True,
+        find_unused_parameters=False,
+        fp16_compression=True,
+    ),
+    checkpointer=dict(period=5000, max_to_keep=100),  # options for PeriodicCheckpointer
+    eval_period=5000,
+    log_period=20,
+    device="cuda",
+    seed=42
+    # ...
+)

configs/semantic_enhanced_matting/dataloader.py

@@ -0,0 +1,62 @@
+from omegaconf import OmegaConf
+from torch.utils.data import ConcatDataset
+from detectron2.config import LazyCall as L
+
+from data.dim_dataset import build_d2_test_dataloader, AdobeCompositionEvaluator, adobe_composition_collate_fn, RW100Test, AIM500Test, AM2KTest, P3M500Test, RWP636Test, SIMTest
+
+AIM500_PATH = '/path/to/datasets/AIM-500'
+RW100_PATH = '/path/to/datasets/RefMatte_RW_100'
+AM2K_PATH = '/path/to/datasets/AM-2K'
+P3M500_PATH = '/path/to/datasets/P3M-10k/validation/P3M-500-NP'
+RWP636_PATH = '/path/to/datasets/RealWorldPortrait-636'
+SIM_PATH = '/path/to/datasets/SIMD/generated_testset'
+
+dataloader = OmegaConf.create()
+test_dataset = L(ConcatDataset)(
+    datasets = [
+        L(AIM500Test)(
+            data_dir = AIM500_PATH,
+            target_size = 1024,
+            multi_fg = True,
+        ),
+        L(RW100Test)(
+            data_dir = RW100_PATH,
+            target_size = 1024,
+            multi_fg = True,
+        ),
+        L(AM2KTest)(
+            data_dir = AM2K_PATH,
+            target_size = 1024,
+            multi_fg = True,
+        ),
+        L(P3M500Test)(
+            data_dir = P3M500_PATH,
+            target_size = 1024,
+            multi_fg = True,
+        ),
+        L(RWP636Test)(
+            data_dir = RWP636_PATH,
+            target_size = 1024,
+            multi_fg = True
+        ),
+        L(SIMTest)(
+            data_dir = SIM_PATH,
+            target_size = 1024,
+            multi_fg = True
+        )
+    ]
+)
+
+dataloader.test = L(build_d2_test_dataloader)(
+    dataset = test_dataset,
+    local_batch_size = 1,
+    num_workers = 4,
+    collate_fn = adobe_composition_collate_fn
+)
+
+dataloader.evaluator = L(AdobeCompositionEvaluator)(
+    save_eval_results_step = 10, 
+    output_dir = None,  # modify in EvalHook (do_test)
+    eval_dataset_type = ['RW100', 'AIM500', 'AM2K', 'P3M500', 'RWP636', 'SIM'],
+    distributed = True,
+),

configs/semantic_enhanced_matting/model.py

@@ -0,0 +1,35 @@
+from detectron2.config import LazyCall as L
+
+from modeling import Detail_Capture, MattingCriterion
+from modeling.meta_arch import SamHqMatte
+from modeling.semantic_enhanced_matting.build_sam import sam_model_registry_def
+# from modeling.sam_hq_matting.predictor import SamPredictor
+from modeling.semantic_enhanced_matting import MaskDecoderMatting
+
+mask_token_only = False
+
+model = L(SamHqMatte)(
+
+    # original sam_hq
+    sam_model = L(sam_model_registry_def)(
+        model_type = 'vit_b',
+        checkpoint = None,
+    ),
+    hq_token_only = True,
+    hq_features_type = 'Final',
+    multimask_output = True,
+
+    # loss function
+    criterion=L(MattingCriterion)(
+        losses = ['unknown_l1_loss', 'known_l1_loss', 'loss_pha_laplacian', 'loss_gradient_penalty']
+    ),
+    
+    # other params.
+    pixel_mean = [123.675 / 255., 116.280 / 255., 103.530 / 255.],
+    pixel_std = [58.395 / 255., 57.120 / 255., 57.375 / 255.],
+
+    lora_rank = None,
+    lora_alpha = None,
+    w_dora = False,
+    w_rslora = False,
+)

data/__init__.py

@@ -0,0 +1 @@
+from .dim_dataset import *

data/coconut_dataset.py

@@ -0,0 +1,377 @@
+import os
+import time
+import json
+import torch
+import numpy as np
+import cv2
+from torch.utils.data import Dataset, DistributedSampler, Sampler
+from torchvision import transforms
+from detectron2.utils.logger import setup_logger
+from typing import Optional
+from operator import itemgetter
+from collections import defaultdict
+
+from data.dim_dataset import GenBBox
+
+
+def random_interp():
+    return np.random.choice([cv2.INTER_NEAREST, cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_LANCZOS4])
+
+
+class SplitConcatImage(object):
+
+    def __init__(self, concat_num=4, wo_mask_to_mattes=False):
+        self.concat_num = concat_num
+        self.wo_mask_to_mattes = wo_mask_to_mattes
+        if self.wo_mask_to_mattes:
+            assert self.concat_num == 5
+
+    def __call__(self, concat_image):
+        if isinstance(concat_image, list):
+            concat_image, image_path = concat_image[0], concat_image[1]
+        else:
+            image_path = None
+        H, W, _ = concat_image.shape
+
+        concat_num = self.concat_num
+        if image_path is not None:
+            if '06-14' in image_path:
+                concat_num = 4
+            elif 'ori_mask' in image_path or 'SEMat' in image_path:
+                concat_num = 3
+            else:
+                concat_num = 5
+        
+        assert W % concat_num == 0
+        W = W // concat_num
+
+        image = concat_image[:H, :W]
+        if self.concat_num != 3:
+            trimap = concat_image[:H, (concat_num - 2) * W: (concat_num - 1) * W]
+            if self.wo_mask_to_mattes:
+                alpha = concat_image[:H, 2 * W: 3 * W]
+            else:
+                alpha = concat_image[:H, (concat_num - 1) * W: concat_num * W]
+        else:
+            trimap = concat_image[:H, (concat_num - 1) * W: concat_num * W]
+            alpha = concat_image[:H, (concat_num - 2) * W: (concat_num - 1) * W]
+
+        return {'image': image, 'trimap': trimap, 'alpha': alpha}
+
+
+class RandomHorizontalFlip(object):
+
+    def __init__(self, prob=0.5):
+        self.prob = prob
+
+    def __call__(self, sample):
+        if np.random.uniform(0, 1) < self.prob:
+            for key in sample.keys():
+                sample[key] = cv2.flip(sample[key], 1)
+        return sample
+
+class EmptyAug(object):
+    def __call__(self, sample):
+        return sample
+
+class RandomReszieCrop(object):
+
+    def __init__(self, output_size=1024, aug_scale_min=0.5, aug_scale_max=1.5):
+        self.desired_size = output_size
+        self.aug_scale_min = aug_scale_min
+        self.aug_scale_max = aug_scale_max
+
+    def __call__(self, sample):
+        H, W, _ = sample['image'].shape
+
+        if self.aug_scale_min == 1.0 and self.aug_scale_max == 1.0:
+            crop_H, crop_W = H, W
+            crop_y1, crop_y2 = 0, crop_H
+            crop_x1, crop_x2 = 0, crop_W
+            scale_W, scaled_H = W, H
+        elif self.aug_scale_min == -1.0 and self.aug_scale_max == -1.0:
+            scale = min(self.desired_size / H, self.desired_size / W)
+            scaled_H, scale_W = round(H * scale), round(W * scale)
+            crop_H, crop_W = scaled_H, scale_W
+            crop_y1, crop_y2 = 0, crop_H
+            crop_x1, crop_x2 = 0, crop_W
+        else:
+            # random size
+            random_scale = np.random.uniform(0, 1) * (self.aug_scale_max - self.aug_scale_min) + self.aug_scale_min  # random_val: 0.5 ~ 1.5
+            scaled_size = round(random_scale * self.desired_size)
+
+            scale = min(scaled_size / H, scaled_size / W)
+            scaled_H, scale_W = round(H * scale), round(W * scale)
+
+            # random crop
+            crop_H, crop_W = min(self.desired_size, scaled_H), min(self.desired_size, scale_W)  # crop_size
+            margin_H, margin_W = max(scaled_H - crop_H, 0), max(scale_W - crop_W, 0)
+            offset_H, offset_W = np.random.randint(0, margin_H + 1), np.random.randint(0, margin_W + 1)
+            crop_y1, crop_y2 = offset_H, offset_H + crop_H
+            crop_x1, crop_x2 = offset_W, offset_W + crop_W
+
+        for key in sample.keys():
+            sample[key] = cv2.resize(sample[key], (scale_W, scaled_H), interpolation=random_interp())[crop_y1: crop_y2, crop_x1: crop_x2, :]  # resize and crop
+            padding = np.zeros(shape=(self.desired_size, self.desired_size, 3), dtype=sample[key].dtype)  # pad to desired_size
+            padding[: crop_H, : crop_W, :] = sample[key]
+            sample[key] = padding
+
+        return sample
+
+
+class RandomJitter(object):
+    """
+    Random change the hue of the image
+    """
+
+    def __call__(self, sample):
+
+        image = sample['image']
+
+        # convert to HSV space, convert to float32 image to keep precision during space conversion.
+        image = cv2.cvtColor(image.astype(np.float32)/255.0, cv2.COLOR_BGR2HSV)
+        # Hue noise
+        hue_jitter = np.random.randint(-40, 40)
+        image[:, :, 0] = np.remainder(image[:, :, 0].astype(np.float32) + hue_jitter, 360)
+        # Saturation noise
+        sat_bar = image[:, :, 1].mean()
+
+        sat_jitter = np.random.rand()*(1.1 - sat_bar)/5 - (1.1 - sat_bar) / 10
+        sat = image[:, :, 1]
+        sat = np.abs(sat + sat_jitter)
+        sat[sat>1] = 2 - sat[sat>1]
+        image[:, :, 1] = sat
+        # Value noise
+        val_bar = image[:, :, 2].mean()
+
+        val_jitter = np.random.rand()*(1.1 - val_bar)/5-(1.1 - val_bar) / 10
+        val = image[:, :, 2]
+        val = np.abs(val + val_jitter)
+        val[val>1] = 2 - val[val>1]
+        image[:, :, 2] = val
+        # convert back to BGR space
+        image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR)
+        sample['image'] = image * 255
+
+        return sample
+
+
+class ToTensor(object):
+
+    def __call__(self, sample):
+        image, alpha, trimap = sample['image'][:, :, ::-1], sample['alpha'], sample['trimap']
+
+        # image
+        image = image.transpose((2, 0, 1)) / 255.
+        sample['image'] = torch.from_numpy(image).float()
+
+        # alpha
+        alpha = alpha.transpose((2, 0, 1))[0: 1] / 255.
+        alpha[alpha < 0 ] = 0
+        alpha[alpha > 1] = 1
+        sample['alpha'] = torch.from_numpy(alpha).float()
+
+        # trimap
+        trimap = trimap.transpose((2, 0, 1))[0: 1] / 1.
+        sample['trimap'] = torch.from_numpy(trimap).float()
+        sample['trimap'][sample['trimap'] < 85] = 0
+        sample['trimap'][sample['trimap'] >= 170] = 1
+        sample['trimap'][sample['trimap'] >= 85] = 0.5
+
+        return sample
+    
+
+class COCONutData(Dataset):
+    def __init__(
+        self, 
+        json_path, 
+        data_root_path, 
+        output_size = 512, 
+        aug_scale_min = 0.5, 
+        aug_scale_max = 1.5,
+        with_bbox = False, 
+        bbox_offset_factor = None,
+        phase = "train",
+        min_miou = 95,
+        miou_json = '',
+        remove_coco_transparent = False,
+        coconut_num_ratio = None,
+        return_multi_fg_info = False,
+        wo_accessory_fusion = False,
+        wo_mask_to_mattes = False,
+        return_image_name = False,
+    ):
+        
+        self.data_root_path = data_root_path
+        self.output_size = output_size
+        self.aug_scale_min = aug_scale_min
+        self.aug_scale_max = aug_scale_max
+        self.with_bbox = with_bbox
+        self.bbox_offset_factor = bbox_offset_factor
+        self.phase = phase
+        self.min_miou = min_miou
+        self.miou_json = miou_json
+        self.remove_coco_transparent = remove_coco_transparent
+        self.coconut_num_ratio = coconut_num_ratio
+        self.return_multi_fg_info = return_multi_fg_info
+        self.wo_accessory_fusion = wo_accessory_fusion # TODO
+        self.wo_mask_to_mattes = wo_mask_to_mattes
+        self.return_image_name = return_image_name
+        assert self.wo_accessory_fusion + self.wo_mask_to_mattes <= 1
+        assert self.phase == 'train'
+
+        self.data_path = []
+        with open(json_path, "r") as file:
+            coconut_matting_info = json.load(file)
+        
+        if self.miou_json != '':
+            name_2_miou_dict = defaultdict(int)
+            with open(self.miou_json, "r") as file:
+                coconut_matting_miou = json.load(file)
+            for miou, name in coconut_matting_miou:
+                name_2_miou_dict[name] = miou
+            for i in coconut_matting_info:
+                if 'accessory' in i['save_path']:
+                    self.data_path.append(i['save_path'])
+                elif name_2_miou_dict[i['save_path'].split('/')[-1]] >= self.min_miou:
+                    if not (self.remove_coco_transparent and 'glass' in i['save_path']):
+                        self.data_path.append(i['save_path'])
+        else:
+            for i in coconut_matting_info:
+                self.data_path.append(i['save_path'])
+
+        if 'accessory' in json_path:
+            concat_num = 5
+        elif 'ori_mask' in json_path:
+            concat_num = 3
+        else:
+            concat_num = 4
+
+        train_trans = [
+            SplitConcatImage(concat_num, wo_mask_to_mattes = self.wo_mask_to_mattes),
+            RandomHorizontalFlip(prob=0 if hasattr(self, 'return_image_name') and self.return_image_name else 0.5),
+            RandomReszieCrop(self.output_size, self.aug_scale_min, self.aug_scale_max),
+            EmptyAug() if hasattr(self, 'return_image_name') and self.return_image_name else RandomJitter(),
+            ToTensor(),
+            GenBBox(bbox_offset_factor=self.bbox_offset_factor)
+        ]
+        self.transform = transforms.Compose(train_trans)
+        print('coconut num: ', len(self.data_path) * self.coconut_num_ratio if self.coconut_num_ratio is not None else len(self.data_path))
+
+    def __getitem__(self, idx):
+        if self.coconut_num_ratio is not None:
+            if self.coconut_num_ratio < 1.0 or idx >= len(self.data_path):
+                idx = np.random.randint(0, len(self.data_path))
+        concat_image = cv2.imread(os.path.join(self.data_root_path, self.data_path[idx]))
+        sample = self.transform([concat_image, self.data_path[idx]])
+        sample['dataset_name'] = 'COCONut'
+        if self.return_multi_fg_info:
+            sample['multi_fg'] = False
+        if hasattr(self, 'return_image_name') and self.return_image_name:
+            sample['image_name'] = self.data_path[idx]
+        return sample
+
+    def __len__(self):
+        if self.coconut_num_ratio is not None:
+            return int(len(self.data_path) * self.coconut_num_ratio)
+        else:
+            return len(self.data_path)
+
+
+class DatasetFromSampler(Dataset):
+    """Dataset to create indexes from `Sampler`.
+
+    Args:
+        sampler: PyTorch sampler
+    """
+
+    def __init__(self, sampler: Sampler):
+        """Initialisation for DatasetFromSampler."""
+        self.sampler = sampler
+        self.sampler_list = None
+
+    def __getitem__(self, index: int):
+        """Gets element of the dataset.
+
+        Args:
+            index: index of the element in the dataset
+
+        Returns:
+            Single element by index
+        """
+        if self.sampler_list is None:
+            self.sampler_list = list(self.sampler)
+        return self.sampler_list[index]
+
+    def __len__(self) -> int:
+        """
+        Returns:
+            int: length of the dataset
+        """
+        return len(self.sampler)
+    
+
+class DistributedSamplerWrapper(DistributedSampler):
+    """
+    Wrapper over `Sampler` for distributed training.
+    Allows you to use any sampler in distributed mode.
+    It is especially useful in conjunction with
+    `torch.nn.parallel.DistributedDataParallel`. In such case, each
+    process can pass a DistributedSamplerWrapper instance as a DataLoader
+    sampler, and load a subset of subsampled data of the original dataset
+    that is exclusive to it.
+    .. note::
+        Sampler is assumed to be of constant size.
+    """
+
+    def __init__(
+        self,
+        sampler,
+        num_replicas: Optional[int] = None,
+        rank: Optional[int] = None,
+        shuffle: bool = True,
+    ):
+        """
+        Args:
+            sampler: Sampler used for subsampling
+            num_replicas (int, optional): Number of processes participating in
+              distributed training
+            rank (int, optional): Rank of the current process
+              within ``num_replicas``
+            shuffle (bool, optional): If true (default),
+              sampler will shuffle the indices
+        """
+        super(DistributedSamplerWrapper, self).__init__(
+            DatasetFromSampler(sampler),
+            num_replicas=num_replicas,
+            rank=rank,
+            shuffle=shuffle,
+        )
+        self.sampler = sampler
+
+    def __iter__(self):
+        """@TODO: Docs. Contribution is welcome."""
+        self.dataset = DatasetFromSampler(self.sampler)
+        indexes_of_indexes = super().__iter__()
+        subsampler_indexes = self.dataset
+        return iter(itemgetter(*indexes_of_indexes)(subsampler_indexes))
+    
+
+if __name__ == '__main__':
+
+    
+
+    dataset = COCONutData(
+        json_path = '/root/data/my_path/Matting/DiffMatte-main/24-06-14_coco-nut_matting.json', 
+        data_root_path = '/root/data/my_path/Matting/DiffMatte-main', 
+        output_size = 1024, 
+        aug_scale_min = 0.5, 
+        aug_scale_max = 1.5,
+        with_bbox = True, 
+        bbox_offset_factor = 0.1,
+        phase = "train"
+    )
+    data = dataset[0]
+
+    for key, val in data.items():
+        print(key, val.shape, torch.min(val), torch.max(val))

data/dim_dataset.py

@@ -0,0 +1,1476 @@
+'''
+Dataloader to process Adobe Image Matting Dataset.
+
+From GCA_Matting(https://github.com/Yaoyi-Li/GCA-Matting/tree/master/dataloader)
+'''
+import os
+import glob
+import logging
+import os.path as osp
+import functools
+import numpy as np
+import torch
+import cv2
+import math
+import numbers
+import random
+import pickle
+from   torch.utils.data import Dataset, DataLoader
+from   torch.nn import functional as F
+from   torchvision import transforms
+from easydict import EasyDict
+from detectron2.utils.logger import setup_logger
+from detectron2.utils import comm
+from detectron2.data import build_detection_test_loader
+import torchvision.transforms.functional
+
+import json
+from PIL import Image
+from detectron2.evaluation.evaluator import DatasetEvaluator
+from collections import defaultdict
+
+from data.evaluate import compute_sad_loss, compute_mse_loss, compute_mad_loss, compute_gradient_loss, compute_connectivity_error
+
+# Base default config
+CONFIG = EasyDict({})
+
+# Model config
+CONFIG.model = EasyDict({})
+# one-hot or class, choice: [3, 1]
+CONFIG.model.trimap_channel = 1
+
+# Dataloader config
+CONFIG.data = EasyDict({})
+# feed forward image size (untested)
+CONFIG.data.crop_size = 512
+# composition of two foregrounds, affine transform, crop and HSV jitter
+CONFIG.data.cutmask_prob = 0.25
+CONFIG.data.augmentation = True
+CONFIG.data.random_interp = True
+
+class Prefetcher():
+    """
+    Modified from the data_prefetcher in https://github.com/NVIDIA/apex/blob/master/examples/imagenet/main_amp.py
+    """
+    def __init__(self, loader):
+        self.orig_loader = loader
+        self.stream = torch.cuda.Stream()
+        self.next_sample = None
+
+    def preload(self):
+        try:
+            self.next_sample = next(self.loader)
+        except StopIteration:
+            self.next_sample = None
+            return
+
+        with torch.cuda.stream(self.stream):
+            for key, value in self.next_sample.items():
+                if isinstance(value, torch.Tensor):
+                    self.next_sample[key] = value.cuda(non_blocking=True)
+
+    def __next__(self):
+        torch.cuda.current_stream().wait_stream(self.stream)
+        sample = self.next_sample
+        if sample is not None:
+            for key, value in sample.items():
+                if isinstance(value, torch.Tensor):
+                    sample[key].record_stream(torch.cuda.current_stream())
+            self.preload()
+        else:
+            # throw stop exception if there is no more data to perform as a default dataloader
+            raise StopIteration("No samples in loader. example: `iterator = iter(Prefetcher(loader)); "
+                                "data = next(iterator)`")
+        return sample
+
+    def __iter__(self):
+        self.loader = iter(self.orig_loader)
+        self.preload()
+        return self
+
+
+class ImageFile(object):
+    def __init__(self, phase='train'):
+        self.phase = phase
+        self.rng = np.random.RandomState(0)
+
+    def _get_valid_names(self, *dirs, shuffle=True):
+        name_sets = [self._get_name_set(d) for d in dirs]
+
+        def _join_and(a, b):
+            return a & b
+
+        valid_names = list(functools.reduce(_join_and, name_sets))
+        if shuffle:
+            self.rng.shuffle(valid_names)
+
+        return valid_names
+
+    @staticmethod
+    def _get_name_set(dir_name):
+        path_list = glob.glob(os.path.join(dir_name, '*'))
+        name_set = set()
+        for path in path_list:
+            name = os.path.basename(path)
+            name = os.path.splitext(name)[0]
+            name_set.add(name)
+        return name_set
+
+    @staticmethod
+    def _list_abspath(data_dir, ext, data_list):
+        return [os.path.join(data_dir, name + ext)
+                for name in data_list]
+
+class ImageFileTrain(ImageFile):
+    def __init__(
+        self,
+        alpha_dir="train_alpha",
+        fg_dir="train_fg",
+        bg_dir="train_bg",
+        alpha_ext=".jpg",
+        fg_ext=".jpg",
+        bg_ext=".jpg",
+        fg_have_bg_num=None,
+        alpha_ratio_json = None,
+        alpha_min_ratio = None,
+        key_sample_ratio = None,
+    ):
+        super(ImageFileTrain, self).__init__(phase="train")
+
+        self.alpha_dir  = alpha_dir
+        self.fg_dir     = fg_dir
+        self.bg_dir     = bg_dir
+        self.alpha_ext  = alpha_ext
+        self.fg_ext     = fg_ext
+        self.bg_ext     = bg_ext
+        logger = setup_logger(name=__name__)
+
+        if not isinstance(self.alpha_dir, str):
+            assert len(self.alpha_dir) == len(self.fg_dir) == len(alpha_ext) == len(fg_ext)
+            self.valid_fg_list = []
+            self.alpha = []
+            self.fg = []
+            self.key_alpha = []
+            self.key_fg = []
+            for i in range(len(self.alpha_dir)):
+                valid_fg_list = self._get_valid_names(self.fg_dir[i], self.alpha_dir[i])
+                valid_fg_list.sort()
+                alpha = self._list_abspath(self.alpha_dir[i], self.alpha_ext[i], valid_fg_list)
+                fg = self._list_abspath(self.fg_dir[i], self.fg_ext[i], valid_fg_list)
+                self.valid_fg_list += valid_fg_list
+
+                self.alpha += alpha * fg_have_bg_num[i]
+                self.fg += fg * fg_have_bg_num[i]
+
+                if alpha_ratio_json[i] is not None:
+                    tmp_key_alpha = []
+                    tmp_key_fg = []
+                    name_to_alpha_path = dict()
+                    for name in alpha:
+                        name_to_alpha_path[name.split('/')[-1].split('.')[0]] = name
+                    name_to_fg_path = dict()
+                    for name in fg:
+                        name_to_fg_path[name.split('/')[-1].split('.')[0]] = name
+
+                    with open(alpha_ratio_json[i], 'r') as file:
+                        alpha_ratio_list = json.load(file)
+                    for ratio, name in alpha_ratio_list:
+                        if ratio < alpha_min_ratio[i]:
+                            break
+                        tmp_key_alpha.append(name_to_alpha_path[name.split('.')[0]])
+                        tmp_key_fg.append(name_to_fg_path[name.split('.')[0]])
+
+                    self.key_alpha.extend(tmp_key_alpha * fg_have_bg_num[i])
+                    self.key_fg.extend(tmp_key_fg * fg_have_bg_num[i])
+
+            if len(self.key_alpha) != 0 and key_sample_ratio > 0:
+                repeat_num = key_sample_ratio * (len(self.alpha) - len(self.key_alpha)) / len(self.key_alpha) / (1 - key_sample_ratio) - 1
+                print('key sample num:', len(self.key_alpha), ', repeat num: ', repeat_num)
+                for i in range(math.ceil(repeat_num)):
+                    self.alpha += self.key_alpha
+                    self.fg += self.key_fg
+
+        else:
+            self.valid_fg_list = self._get_valid_names(self.fg_dir, self.alpha_dir)
+            self.valid_fg_list.sort()
+            self.alpha = self._list_abspath(self.alpha_dir, self.alpha_ext, self.valid_fg_list)
+            self.fg = self._list_abspath(self.fg_dir, self.fg_ext, self.valid_fg_list)
+            
+        self.valid_bg_list = [os.path.splitext(name)[0] for name in os.listdir(self.bg_dir)]
+        self.valid_bg_list.sort()
+
+        if fg_have_bg_num is not None:
+            # assert fg_have_bg_num * len(self.valid_fg_list) <= len(self.valid_bg_list)
+            # self.valid_bg_list = self.valid_bg_list[: fg_have_bg_num * len(self.valid_fg_list)]
+            assert len(self.alpha) <= len(self.valid_bg_list)
+            self.valid_bg_list = self.valid_bg_list[: len(self.alpha)]
+
+        self.bg = self._list_abspath(self.bg_dir, self.bg_ext, self.valid_bg_list)
+
+    def __len__(self):  
+        return len(self.alpha)
+
+class ImageFileTest(ImageFile):
+    def __init__(self,
+                 alpha_dir="test_alpha",
+                 merged_dir="test_merged",
+                 trimap_dir="test_trimap",
+                 alpha_ext=".png",
+                 merged_ext=".png",
+                 trimap_ext=".png"):
+        super(ImageFileTest, self).__init__(phase="test")
+
+        self.alpha_dir  = alpha_dir
+        self.merged_dir = merged_dir
+        self.trimap_dir = trimap_dir
+        self.alpha_ext  = alpha_ext
+        self.merged_ext = merged_ext
+        self.trimap_ext = trimap_ext
+
+        self.valid_image_list = self._get_valid_names(self.alpha_dir, self.merged_dir, self.trimap_dir, shuffle=False)
+
+        self.alpha = self._list_abspath(self.alpha_dir, self.alpha_ext, self.valid_image_list)
+        self.merged = self._list_abspath(self.merged_dir, self.merged_ext, self.valid_image_list)
+        self.trimap = self._list_abspath(self.trimap_dir, self.trimap_ext, self.valid_image_list)
+
+    def __len__(self):
+        return len(self.alpha)
+
+interp_list = [cv2.INTER_NEAREST, cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_LANCZOS4]
+
+
+def maybe_random_interp(cv2_interp):
+    if CONFIG.data.random_interp:
+        return np.random.choice(interp_list)
+    else:
+        return cv2_interp
+
+
+class ToTensor(object):
+    """
+    Convert ndarrays in sample to Tensors with normalization.
+    """
+    def __init__(self, phase="test"):
+        self.mean = torch.tensor([0.485, 0.456, 0.406]).view(3,1,1)
+        self.std = torch.tensor([0.229, 0.224, 0.225]).view(3,1,1)
+        self.phase = phase
+
+    def __call__(self, sample):
+        image, alpha, trimap, mask = sample['image'][:,:,::-1], sample['alpha'], sample['trimap'], sample['mask']
+        
+        alpha[alpha < 0 ] = 0
+        alpha[alpha > 1] = 1
+     
+        image = image.transpose((2, 0, 1)).astype(np.float32)
+        alpha = np.expand_dims(alpha.astype(np.float32), axis=0)
+        
+        mask = np.expand_dims(mask.astype(np.float32), axis=0)
+
+        image /= 255.
+
+        if self.phase == "train":
+            fg = sample['fg'][:,:,::-1].transpose((2, 0, 1)).astype(np.float32) / 255.
+            sample['fg'] = torch.from_numpy(fg)
+            bg = sample['bg'][:,:,::-1].transpose((2, 0, 1)).astype(np.float32) / 255.
+            sample['bg'] = torch.from_numpy(bg)
+
+        sample['image'], sample['alpha'], sample['trimap'] = \
+            torch.from_numpy(image), torch.from_numpy(alpha), torch.from_numpy(trimap).to(torch.long)
+        sample['image'] = sample['image']
+
+        if CONFIG.model.trimap_channel == 3:
+            sample['trimap'] = F.one_hot(sample['trimap'], num_classes=3).permute(2,0,1).float()
+        elif CONFIG.model.trimap_channel == 1:
+            sample['trimap'] = sample['trimap'][None,...].float()
+        else:
+            raise NotImplementedError("CONFIG.model.trimap_channel can only be 3 or 1")
+        sample['trimap'][sample['trimap'] < 85] = 0
+        sample['trimap'][sample['trimap'] >= 170] = 1
+        sample['trimap'][sample['trimap'] >= 85] = 0.5
+
+        sample['mask'] = torch.from_numpy(mask).float()
+
+        return sample
+
+
+class RandomAffine(object):
+    """
+    Random affine translation
+    """
+    def __init__(self, degrees, translate=None, scale=None, shear=None, flip=None, resample=False, fillcolor=0):
+        if isinstance(degrees, numbers.Number):
+            if degrees < 0:
+                raise ValueError("If degrees is a single number, it must be positive.")
+            self.degrees = (-degrees, degrees)
+        else:
+            assert isinstance(degrees, (tuple, list)) and len(degrees) == 2, \
+                "degrees should be a list or tuple and it must be of length 2."
+            self.degrees = degrees
+
+        if translate is not None:
+            assert isinstance(translate, (tuple, list)) and len(translate) == 2, \
+                "translate should be a list or tuple and it must be of length 2."
+            for t in translate:
+                if not (0.0 <= t <= 1.0):
+                    raise ValueError("translation values should be between 0 and 1")
+        self.translate = translate
+
+        if scale is not None:
+            assert isinstance(scale, (tuple, list)) and len(scale) == 2, \
+                "scale should be a list or tuple and it must be of length 2."
+            for s in scale:
+                if s <= 0:
+                    raise ValueError("scale values should be positive")
+        self.scale = scale
+
+        if shear is not None:
+            if isinstance(shear, numbers.Number):
+                if shear < 0:
+                    raise ValueError("If shear is a single number, it must be positive.")
+                self.shear = (-shear, shear)
+            else:
+                assert isinstance(shear, (tuple, list)) and len(shear) == 2, \
+                    "shear should be a list or tuple and it must be of length 2."
+                self.shear = shear
+        else:
+            self.shear = shear
+
+        self.resample = resample
+        self.fillcolor = fillcolor
+        self.flip = flip
+
+    @staticmethod
+    def get_params(degrees, translate, scale_ranges, shears, flip, img_size):
+        """Get parameters for affine transformation
+
+        Returns:
+            sequence: params to be passed to the affine transformation
+        """
+        angle = random.uniform(degrees[0], degrees[1])
+        if translate is not None:
+            max_dx = translate[0] * img_size[0]
+            max_dy = translate[1] * img_size[1]
+            translations = (np.round(random.uniform(-max_dx, max_dx)),
+                            np.round(random.uniform(-max_dy, max_dy)))
+        else:
+            translations = (0, 0)
+
+        if scale_ranges is not None:
+            scale = (random.uniform(scale_ranges[0], scale_ranges[1]),
+                     random.uniform(scale_ranges[0], scale_ranges[1]))
+        else:
+            scale = (1.0, 1.0)
+
+        if shears is not None:
+            shear = random.uniform(shears[0], shears[1])
+        else:
+            shear = 0.0
+
+        if flip is not None:
+            flip = (np.random.rand(2) < flip).astype(np.int32) * 2 - 1
+
+        return angle, translations, scale, shear, flip
+
+    def __call__(self, sample):
+        fg, alpha = sample['fg'], sample['alpha']
+        rows, cols, ch = fg.shape
+        if np.maximum(rows, cols) < 1024:
+            params = self.get_params((0, 0), self.translate, self.scale, self.shear, self.flip, fg.size)
+        else:
+            params = self.get_params(self.degrees, self.translate, self.scale, self.shear, self.flip, fg.size)
+
+        center = (cols * 0.5 + 0.5, rows * 0.5 + 0.5)
+        M = self._get_inverse_affine_matrix(center, *params)
+        M = np.array(M).reshape((2, 3))
+
+        fg = cv2.warpAffine(fg, M, (cols, rows),
+                            flags=maybe_random_interp(cv2.INTER_NEAREST) + cv2.WARP_INVERSE_MAP)
+        alpha = cv2.warpAffine(alpha, M, (cols, rows),
+                               flags=maybe_random_interp(cv2.INTER_NEAREST) + cv2.WARP_INVERSE_MAP)
+
+        sample['fg'], sample['alpha'] = fg, alpha
+
+        return sample
+
+
+    @ staticmethod
+    def _get_inverse_affine_matrix(center, angle, translate, scale, shear, flip):
+
+        angle = math.radians(angle)
+        shear = math.radians(shear)
+        scale_x = 1.0 / scale[0] * flip[0]
+        scale_y = 1.0 / scale[1] * flip[1]
+
+        # Inverted rotation matrix with scale and shear
+        d = math.cos(angle + shear) * math.cos(angle) + math.sin(angle + shear) * math.sin(angle)
+        matrix = [
+            math.cos(angle) * scale_x, math.sin(angle + shear) * scale_x, 0,
+            -math.sin(angle) * scale_y, math.cos(angle + shear) * scale_y, 0
+        ]
+        matrix = [m / d for m in matrix]
+
+        # Apply inverse of translation and of center translation: RSS^-1 * C^-1 * T^-1
+        matrix[2] += matrix[0] * (-center[0] - translate[0]) + matrix[1] * (-center[1] - translate[1])
+        matrix[5] += matrix[3] * (-center[0] - translate[0]) + matrix[4] * (-center[1] - translate[1])
+
+        # Apply center translation: C * RSS^-1 * C^-1 * T^-1
+        matrix[2] += center[0]
+        matrix[5] += center[1]
+
+        return matrix
+
+
+class RandomJitter(object):
+    """
+    Random change the hue of the image
+    """
+
+    def __call__(self, sample):
+        sample_ori = sample.copy()
+        fg, alpha = sample['fg'], sample['alpha']
+        # if alpha is all 0 skip
+        if np.all(alpha==0):
+            return sample_ori
+        # convert to HSV space, convert to float32 image to keep precision during space conversion.
+        fg = cv2.cvtColor(fg.astype(np.float32)/255.0, cv2.COLOR_BGR2HSV)
+        # Hue noise
+        hue_jitter = np.random.randint(-40, 40)
+        fg[:, :, 0] = np.remainder(fg[:, :, 0].astype(np.float32) + hue_jitter, 360)
+        # Saturation noise
+        sat_bar = fg[:, :, 1][alpha > 0].mean()
+        if np.isnan(sat_bar):
+            return sample_ori
+        sat_jitter = np.random.rand()*(1.1 - sat_bar)/5 - (1.1 - sat_bar) / 10
+        sat = fg[:, :, 1]
+        sat = np.abs(sat + sat_jitter)
+        sat[sat>1] = 2 - sat[sat>1]
+        fg[:, :, 1] = sat
+        # Value noise
+        val_bar = fg[:, :, 2][alpha > 0].mean()
+        if np.isnan(val_bar):
+            return sample_ori
+        val_jitter = np.random.rand()*(1.1 - val_bar)/5-(1.1 - val_bar) / 10
+        val = fg[:, :, 2]
+        val = np.abs(val + val_jitter)
+        val[val>1] = 2 - val[val>1]
+        fg[:, :, 2] = val
+        # convert back to BGR space
+        fg = cv2.cvtColor(fg, cv2.COLOR_HSV2BGR)
+        sample['fg'] = fg*255
+
+        return sample
+
+
+class RandomHorizontalFlip(object):
+    """
+    Random flip image and label horizontally
+    """
+    def __init__(self, prob=0.5):
+        self.prob = prob
+    def __call__(self, sample):
+        fg, alpha = sample['fg'], sample['alpha']
+        if np.random.uniform(0, 1) < self.prob:
+            fg = cv2.flip(fg, 1)
+            alpha = cv2.flip(alpha, 1)
+        sample['fg'], sample['alpha'] = fg, alpha
+
+        return sample
+
+
+class RandomCrop(object):
+    """
+    Crop randomly the image in a sample, retain the center 1/4 images, and resize to 'output_size'
+
+    :param output_size (tuple or int): Desired output size. If int, square crop
+            is made.
+    """
+
+    def __init__(self, output_size=( CONFIG.data.crop_size, CONFIG.data.crop_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
+        self.margin = output_size[0] // 2
+        self.logger = logging.getLogger("Logger")
+
+    def __call__(self, sample):
+        fg, alpha, trimap, mask, name = sample['fg'],  sample['alpha'], sample['trimap'], sample['mask'], sample['image_name']
+        bg = sample['bg']
+        h, w = trimap.shape
+        bg = cv2.resize(bg, (w, h), interpolation=maybe_random_interp(cv2.INTER_CUBIC))
+        if w < self.output_size[0]+1 or h < self.output_size[1]+1:
+            ratio = 1.1*self.output_size[0]/h if h < w else 1.1*self.output_size[1]/w
+            # self.logger.warning("Size of {} is {}.".format(name, (h, w)))
+            while h < self.output_size[0]+1 or w < self.output_size[1]+1:
+                fg = cv2.resize(fg, (int(w*ratio), int(h*ratio)), interpolation=maybe_random_interp(cv2.INTER_NEAREST))
+                alpha = cv2.resize(alpha, (int(w*ratio), int(h*ratio)),
+                                   interpolation=maybe_random_interp(cv2.INTER_NEAREST))
+                trimap = cv2.resize(trimap, (int(w*ratio), int(h*ratio)), interpolation=cv2.INTER_NEAREST)
+                bg = cv2.resize(bg, (int(w*ratio), int(h*ratio)), interpolation=maybe_random_interp(cv2.INTER_CUBIC))
+                mask = cv2.resize(mask, (int(w*ratio), int(h*ratio)), interpolation=cv2.INTER_NEAREST)
+                h, w = trimap.shape
+        small_trimap = cv2.resize(trimap, (w//4, h//4), interpolation=cv2.INTER_NEAREST)
+        unknown_list = list(zip(*np.where(small_trimap[self.margin//4:(h-self.margin)//4,
+                                                       self.margin//4:(w-self.margin)//4] == 128)))
+        unknown_num = len(unknown_list)
+        if len(unknown_list) < 10:
+            left_top = (np.random.randint(0, h-self.output_size[0]+1), np.random.randint(0, w-self.output_size[1]+1))
+        else:
+            idx = np.random.randint(unknown_num)
+            left_top = (unknown_list[idx][0]*4, unknown_list[idx][1]*4)
+
+        fg_crop = fg[left_top[0]:left_top[0]+self.output_size[0], left_top[1]:left_top[1]+self.output_size[1],:]
+        alpha_crop = alpha[left_top[0]:left_top[0]+self.output_size[0], left_top[1]:left_top[1]+self.output_size[1]]
+        bg_crop = bg[left_top[0]:left_top[0]+self.output_size[0], left_top[1]:left_top[1]+self.output_size[1],:]
+        trimap_crop = trimap[left_top[0]:left_top[0]+self.output_size[0], left_top[1]:left_top[1]+self.output_size[1]]
+        mask_crop = mask[left_top[0]:left_top[0]+self.output_size[0], left_top[1]:left_top[1]+self.output_size[1]]
+
+        if len(np.where(trimap==128)[0]) == 0:
+            self.logger.error("{} does not have enough unknown area for crop. Resized to target size."
+                                "left_top: {}".format(name, left_top))
+            fg_crop = cv2.resize(fg, self.output_size[::-1], interpolation=maybe_random_interp(cv2.INTER_NEAREST))
+            alpha_crop = cv2.resize(alpha, self.output_size[::-1], interpolation=maybe_random_interp(cv2.INTER_NEAREST))
+            trimap_crop = cv2.resize(trimap, self.output_size[::-1], interpolation=cv2.INTER_NEAREST)
+            bg_crop = cv2.resize(bg, self.output_size[::-1], interpolation=maybe_random_interp(cv2.INTER_CUBIC))
+            mask_crop = cv2.resize(mask, self.output_size[::-1], interpolation=cv2.INTER_NEAREST)
+        
+        sample.update({'fg': fg_crop, 'alpha': alpha_crop, 'trimap': trimap_crop, 'mask': mask_crop, 'bg': bg_crop})
+        return sample
+
+
+class OriginScale(object):
+    def __call__(self, sample):
+        h, w = sample["alpha_shape"]
+
+        if h % 32 == 0 and w % 32 == 0:
+            return sample
+
+        target_h = 32 * ((h - 1) // 32 + 1)
+        target_w = 32 * ((w - 1) // 32 + 1)
+        pad_h = target_h - h
+        pad_w = target_w - w
+
+        padded_image = np.pad(sample['image'], ((0,pad_h), (0, pad_w), (0,0)), mode="reflect")
+        padded_trimap = np.pad(sample['trimap'], ((0,pad_h), (0, pad_w)), mode="reflect")
+        padded_mask = np.pad(sample['mask'], ((0,pad_h), (0, pad_w)), mode="reflect")
+
+        sample['image'] = padded_image
+        sample['trimap'] = padded_trimap
+        sample['mask'] = padded_mask
+
+        return sample
+
+
+class GenMask(object):
+    def __init__(self):
+        self.erosion_kernels = [None] + [cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (size, size)) for size in range(1,30)]
+
+    def __call__(self, sample):
+        alpha_ori = sample['alpha']
+        h, w = alpha_ori.shape
+
+        max_kernel_size = 30
+        alpha = cv2.resize(alpha_ori, (640,640), interpolation=maybe_random_interp(cv2.INTER_NEAREST))
+
+        ### generate trimap
+        fg_mask = (alpha + 1e-5).astype(np.int32).astype(np.uint8)
+        bg_mask = (1 - alpha + 1e-5).astype(np.int32).astype(np.uint8)
+        fg_mask = cv2.erode(fg_mask, self.erosion_kernels[np.random.randint(1, max_kernel_size)])
+        bg_mask = cv2.erode(bg_mask, self.erosion_kernels[np.random.randint(1, max_kernel_size)])
+
+        fg_width = np.random.randint(1, 30)
+        bg_width = np.random.randint(1, 30)
+        fg_mask = (alpha + 1e-5).astype(np.int32).astype(np.uint8)
+        bg_mask = (1 - alpha + 1e-5).astype(np.int32).astype(np.uint8)
+        fg_mask = cv2.erode(fg_mask, self.erosion_kernels[fg_width])
+        bg_mask = cv2.erode(bg_mask, self.erosion_kernels[bg_width])
+
+        trimap = np.ones_like(alpha) * 128
+        trimap[fg_mask == 1] = 255
+        trimap[bg_mask == 1] = 0
+
+        trimap = cv2.resize(trimap, (w,h), interpolation=cv2.INTER_NEAREST)
+        sample['trimap'] = trimap
+
+        ### generate mask
+        low = 0.01
+        high = 1.0
+        thres = random.random() * (high - low) + low
+        seg_mask = (alpha >= thres).astype(np.int32).astype(np.uint8)
+        random_num = random.randint(0,3)
+        if random_num == 0:
+            seg_mask = cv2.erode(seg_mask, self.erosion_kernels[np.random.randint(1, max_kernel_size)])
+        elif random_num == 1:
+            seg_mask = cv2.dilate(seg_mask, self.erosion_kernels[np.random.randint(1, max_kernel_size)])
+        elif random_num == 2:
+            seg_mask = cv2.erode(seg_mask, self.erosion_kernels[np.random.randint(1, max_kernel_size)])
+            seg_mask = cv2.dilate(seg_mask, self.erosion_kernels[np.random.randint(1, max_kernel_size)])
+        elif random_num == 3:
+            seg_mask = cv2.dilate(seg_mask, self.erosion_kernels[np.random.randint(1, max_kernel_size)])
+            seg_mask = cv2.erode(seg_mask, self.erosion_kernels[np.random.randint(1, max_kernel_size)])
+        
+        seg_mask = cv2.resize(seg_mask, (w,h), interpolation=cv2.INTER_NEAREST)
+        sample['mask'] = seg_mask
+
+        return sample
+
+
+class Composite(object):
+    def __call__(self, sample):
+        fg, bg, alpha = sample['fg'], sample['bg'], sample['alpha']
+        alpha[alpha < 0 ] = 0
+        alpha[alpha > 1] = 1
+        fg[fg < 0 ] = 0
+        fg[fg > 255] = 255
+        bg[bg < 0 ] = 0
+        bg[bg > 255] = 255
+
+        image = fg * alpha[:, :, None] + bg * (1 - alpha[:, :, None])
+        sample['image'] = image
+        return sample
+
+
+class CutMask(object):
+    def __init__(self, perturb_prob = 0):
+        self.perturb_prob = perturb_prob
+
+    def __call__(self, sample):
+        if np.random.rand() < self.perturb_prob:
+            return sample
+
+        mask = sample['mask'] # H x W, trimap 0--255, segmask 0--1, alpha 0--1
+        h, w = mask.shape
+        perturb_size_h, perturb_size_w = random.randint(h // 4, h // 2), random.randint(w // 4, w // 2)
+        x = random.randint(0, h - perturb_size_h)
+        y = random.randint(0, w - perturb_size_w)
+        x1 = random.randint(0, h - perturb_size_h)
+        y1 = random.randint(0, w - perturb_size_w)
+        
+        mask[x:x+perturb_size_h, y:y+perturb_size_w] = mask[x1:x1+perturb_size_h, y1:y1+perturb_size_w].copy()
+        
+        sample['mask'] = mask
+        return sample
+
+
+class ScaleFg(object):
+    def __init__(self, min_scale_fg_scale=0.5, max_scale_fg_scale=1.0):
+        self.min_scale_fg_scale = min_scale_fg_scale
+        self.max_scale_fg_scale = max_scale_fg_scale
+
+    def __call__(self, sample):
+        scale_factor = np.random.uniform(low=self.min_scale_fg_scale, high=self.max_scale_fg_scale)
+
+        fg, alpha = sample['fg'], sample['alpha']  # np.array(): [H, W, 3] 0 ~ 255 , [H, W] 0.0 ~ 1.0
+        h, w = alpha.shape
+        scale_h, scale_w = int(h * scale_factor), int(w * scale_factor)
+
+        new_fg, new_alpha = np.zeros_like(fg), np.zeros_like(alpha)
+        fg = cv2.resize(fg, (scale_w, scale_h), interpolation=cv2.INTER_LINEAR) 
+        alpha = cv2.resize(alpha, (scale_w, scale_h), interpolation=cv2.INTER_LINEAR) 
+
+        if scale_factor <= 1:
+            offset_h, offset_w = np.random.randint(h - scale_h + 1), np.random.randint(w - scale_w + 1)
+            new_fg[offset_h: offset_h + scale_h, offset_w: offset_w + scale_w, :] = fg
+            new_alpha[offset_h: offset_h + scale_h, offset_w: offset_w + scale_w] = alpha
+        else:
+            offset_h, offset_w = np.random.randint(scale_h - h + 1), np.random.randint(scale_w - w + 1)
+            new_fg = fg[offset_h: offset_h + scale_h, offset_w: offset_w + scale_w, :]
+            new_alpha = alpha[offset_h: offset_h + scale_h, offset_w: offset_w + scale_w]
+
+        sample['fg'], sample['alpha'] = new_fg, new_alpha
+        return sample
+
+class GenBBox(object):
+    def __init__(self, bbox_offset_factor = 0.1, random_crop_bbox = None, train_or_test = 'train', dataset_type = None, random_auto_matting=None):
+        self.bbox_offset_factor = bbox_offset_factor
+        self.random_crop_bbox = random_crop_bbox
+        self.train_or_test = train_or_test
+        self.dataset_type = dataset_type
+        self.random_auto_matting = random_auto_matting
+
+    def __call__(self, sample):
+
+        alpha = sample['alpha']  # [1, H, W] 0.0 ~ 1.0
+        indices = torch.nonzero(alpha[0], as_tuple=True)
+
+        if len(indices[0]) > 0:
+
+            min_x, min_y = torch.min(indices[1]), torch.min(indices[0])
+            max_x, max_y = torch.max(indices[1]), torch.max(indices[0])
+
+            if self.random_crop_bbox is not None and np.random.uniform(0, 1) < self.random_crop_bbox:
+                ori_h_w = (sample['alpha'].shape[-2], sample['alpha'].shape[-1])
+                sample['alpha'] = F.interpolate(sample['alpha'][None, :, min_y: max_y + 1, min_x: max_x + 1], size=ori_h_w, mode='bilinear', align_corners=False)[0]
+                sample['image'] = F.interpolate(sample['image'][None, :, min_y: max_y + 1, min_x: max_x + 1], size=ori_h_w, mode='bilinear', align_corners=False)[0]
+                sample['trimap'] = F.interpolate(sample['trimap'][None, :, min_y: max_y + 1, min_x: max_x + 1], size=ori_h_w, mode='nearest')[0]
+                bbox = torch.tensor([[0, 0, ori_h_w[1] - 1, ori_h_w[0] - 1]])
+
+            elif self.bbox_offset_factor != 0:
+                bbox_w = max(1, max_x - min_x)
+                bbox_h = max(1, max_y - min_y)
+                offset_w = math.ceil(self.bbox_offset_factor * bbox_w)
+                offset_h = math.ceil(self.bbox_offset_factor * bbox_h)
+
+                min_x = max(0, min_x + np.random.randint(-offset_w, offset_w))
+                max_x = min(alpha.shape[2] - 1, max_x + np.random.randint(-offset_w, offset_w))
+                min_y = max(0, min_y + np.random.randint(-offset_h, offset_h))
+                max_y = min(alpha.shape[1] - 1, max_y + np.random.randint(-offset_h, offset_h))
+                bbox = torch.tensor([[min_x, min_y, max_x, max_y]])
+            else:
+                bbox = torch.tensor([[min_x, min_y, max_x, max_y]])
+            
+            if self.random_auto_matting is not None and np.random.uniform(0, 1) < self.random_auto_matting:
+                bbox = torch.tensor([[0, 0, alpha.shape[2] - 1, alpha.shape[1] - 1]])
+
+        else:
+            bbox = torch.zeros(1, 4)
+
+        sample['bbox'] = bbox.float()
+        return sample
+
+class DataGenerator(Dataset):
+    def __init__(
+            self, 
+            data, 
+            phase="train", 
+            crop_size=512, 
+            remove_multi_fg=False, 
+            min_scale_fg_scale=None, 
+            max_scale_fg_scale=None, 
+            with_bbox = False, 
+            bbox_offset_factor = None,
+            return_keys = None,
+            random_crop_bbox = None,
+            dataset_name = None,
+            random_auto_matting = None,
+        ):
+        self.phase = phase
+        # self.crop_size = CONFIG.data.crop_size
+        self.crop_size = crop_size
+        self.remove_multi_fg = remove_multi_fg
+        self.with_bbox = with_bbox
+        self.bbox_offset_factor = bbox_offset_factor
+        self.alpha = data.alpha
+        self.return_keys = return_keys
+        self.random_crop_bbox = random_crop_bbox
+        self.dataset_name = dataset_name
+        self.random_auto_matting = random_auto_matting
+
+        if self.phase == "train":
+            self.fg = data.fg
+            self.bg = data.bg
+            self.merged = []
+            self.trimap = []
+        else:
+            self.fg = []
+            self.bg = []
+            self.merged = data.merged
+            self.trimap = data.trimap
+
+        train_trans = [
+            RandomAffine(degrees=30, scale=[0.8, 1.25], shear=10, flip=0.5),
+            GenMask(),
+            CutMask(perturb_prob=CONFIG.data.cutmask_prob),
+            RandomCrop((self.crop_size, self.crop_size)),
+            RandomJitter(),
+            Composite(),
+            ToTensor(phase="train")
+        ]
+        if min_scale_fg_scale is not None:
+            train_trans.insert(0, ScaleFg(min_scale_fg_scale, max_scale_fg_scale))
+        if self.with_bbox:
+            train_trans.append(GenBBox(bbox_offset_factor=self.bbox_offset_factor, random_crop_bbox=self.random_crop_bbox, random_auto_matting=self.random_auto_matting))
+
+        test_trans = [ OriginScale(), ToTensor() ]
+
+        self.transform = {
+            'train':
+                transforms.Compose(train_trans),
+            'val':
+                transforms.Compose([
+                    OriginScale(),
+                    ToTensor()
+                ]),
+            'test':
+                transforms.Compose(test_trans)
+        }[phase]
+
+        self.fg_num = len(self.fg)
+
+    def select_keys(self, sample):
+        new_sample = {}
+        for key, val in sample.items():
+            if key in self.return_keys:
+                new_sample[key] = val
+        return new_sample
+
+    def __getitem__(self, idx):
+        if self.phase == "train":
+            fg = cv2.imread(self.fg[idx % self.fg_num])
+            alpha = cv2.imread(self.alpha[idx % self.fg_num], 0).astype(np.float32)/255
+            bg = cv2.imread(self.bg[idx], 1)
+
+            if not self.remove_multi_fg:
+                fg, alpha, multi_fg = self._composite_fg(fg, alpha, idx)
+            else:
+                multi_fg = False
+            image_name = os.path.split(self.fg[idx % self.fg_num])[-1]
+            sample = {'fg': fg, 'alpha': alpha, 'bg': bg, 'image_name': image_name, 'multi_fg': multi_fg}
+
+        else:
+            image = cv2.imread(self.merged[idx])
+            alpha = cv2.imread(self.alpha[idx], 0)/255.
+            trimap = cv2.imread(self.trimap[idx], 0)
+            mask = (trimap >= 170).astype(np.float32)
+            image_name = os.path.split(self.merged[idx])[-1]
+
+            sample = {'image': image, 'alpha': alpha, 'trimap': trimap, 'mask': mask, 'image_name': image_name, 'alpha_shape': alpha.shape}
+
+        sample = self.transform(sample)
+
+        if self.return_keys is not None:
+            sample = self.select_keys(sample)
+        if self.dataset_name is not None:
+            sample['dataset_name'] = self.dataset_name
+        return sample
+
+    def _composite_fg(self, fg, alpha, idx):
+        
+        multi_fg = False
+        if np.random.rand() < 0.5:
+            idx2 = np.random.randint(self.fg_num) + idx
+            fg2 = cv2.imread(self.fg[idx2 % self.fg_num])
+            alpha2 = cv2.imread(self.alpha[idx2 % self.fg_num], 0).astype(np.float32)/255.
+            h, w = alpha.shape
+            fg2 = cv2.resize(fg2, (w, h), interpolation=maybe_random_interp(cv2.INTER_NEAREST))
+            alpha2 = cv2.resize(alpha2, (w, h), interpolation=maybe_random_interp(cv2.INTER_NEAREST))
+
+            alpha_tmp = 1 - (1 - alpha) * (1 - alpha2)
+            if  np.any(alpha_tmp < 1):
+                fg = fg.astype(np.float32) * alpha[:,:,None] + fg2.astype(np.float32) * (1 - alpha[:,:,None])
+                # The overlap of two 50% transparency should be 25%
+                alpha = alpha_tmp
+                fg = fg.astype(np.uint8)
+            multi_fg = True
+
+        if np.random.rand() < 0.25:
+            # fg = cv2.resize(fg, (640, 640), interpolation=maybe_random_interp(cv2.INTER_NEAREST))
+            # alpha = cv2.resize(alpha, (640, 640), interpolation=maybe_random_interp(cv2.INTER_NEAREST))
+            fg = cv2.resize(fg, (1280, 1280), interpolation=maybe_random_interp(cv2.INTER_NEAREST))
+            alpha = cv2.resize(alpha, (1280, 1280), interpolation=maybe_random_interp(cv2.INTER_NEAREST))
+
+        return fg, alpha, multi_fg
+
+    def __len__(self):
+        if self.phase == "train":
+            return len(self.bg)
+        else:
+            return len(self.alpha)
+
+
+class ResziePad(object):
+
+    def __init__(self, target_size=1024):
+        self.target_size = target_size
+
+    def __call__(self, sample):
+        _, H, W = sample['image'].shape
+
+        scale = self.target_size * 1.0 / max(H, W)
+        new_H, new_W = H * scale, W * scale
+        new_W = int(new_W + 0.5)
+        new_H = int(new_H + 0.5)
+
+        choice = {'image', 'trimap', 'alpha'} if 'trimap' in sample.keys() else {'image', 'alpha'}
+        for key in choice:
+            if key in {'image', 'trimap'}:
+                sample[key] = F.interpolate(sample[key][None], size=(new_H, new_W), mode='bilinear', align_corners=False)[0]
+            else:
+                # sample[key] = F.interpolate(sample[key][None], size=(new_H, new_W), mode='nearest')[0]
+                sample[key] = F.interpolate(sample[key][None], size=(new_H, new_W), mode='bilinear', align_corners=False)[0]
+            padding = torch.zeros([sample[key].shape[0], self.target_size, self.target_size], dtype=sample[key].dtype, device=sample[key].device)
+            padding[:, : new_H, : new_W] = sample[key]
+            sample[key] = padding
+
+        return sample
+    
+
+class Cv2ResziePad(object):
+
+    def __init__(self, target_size=1024):
+        self.target_size = target_size
+
+    def __call__(self, sample):
+        H, W, _ = sample['image'].shape
+
+        scale = self.target_size * 1.0 / max(H, W)
+        new_H, new_W = H * scale, W * scale
+        new_W = int(new_W + 0.5)
+        new_H = int(new_H + 0.5)
+
+        choice = {'image', 'trimap', 'alpha'} if 'trimap' in sample.keys() and sample['trimap'] is not None else {'image', 'alpha'}
+        for key in choice:
+            sample[key] = cv2.resize(sample[key], (new_W, new_H), interpolation=cv2.INTER_LINEAR)  # cv2.INTER_NEAREST, cv2.INTER_LINEAR, cv2.INTER_CUBIC
+
+            if key == 'image':
+                padding = np.zeros([self.target_size, self.target_size, sample[key].shape[-1]], dtype=sample[key].dtype)
+                padding[: new_H, : new_W, :] = sample[key]
+                sample[key] = padding
+                sample[key] = sample[key][:, :, ::-1].transpose((2, 0, 1)).astype(np.float32) #/ 255.0
+            else:
+                padding = np.zeros([self.target_size, self.target_size], dtype=sample[key].dtype)
+                padding[: new_H, : new_W] = sample[key]
+                sample[key] = padding
+                sample[key] = sample[key][None].astype(np.float32)
+            sample[key] = torch.from_numpy(sample[key])
+
+        return sample
+    
+
+class AdobeCompositionTest(Dataset):
+    def __init__(self, data_dir, target_size=1024, multi_fg=None):
+        self.data_dir = data_dir
+        self.file_names = sorted(os.listdir(os.path.join(self.data_dir, 'merged')))
+        
+        test_trans = [
+            ResziePad(target_size=target_size),
+            GenBBox(bbox_offset_factor=0)
+        ]
+        self.transform = transforms.Compose(test_trans)
+        self.multi_fg = multi_fg
+
+    def __len__(self):  # 1000
+        return len(self.file_names)
+
+    def __getitem__(self, idx):
+        phas = Image.open(os.path.join(self.data_dir, 'alpha_copy', self.file_names[idx])).convert('L')
+        tris = Image.open(os.path.join(self.data_dir, 'trimaps', self.file_names[idx]))
+        imgs = Image.open(os.path.join(self.data_dir, 'merged', self.file_names[idx]))
+        sample = {
+            'ori_h_w': (imgs.size[1], imgs.size[0]),
+            'data_type': 'Adobe'
+        }
+
+        sample['alpha'] = torchvision.transforms.functional.to_tensor(phas)  # [1, H, W] 0.0 ~ 1.0
+        sample['trimap'] = torchvision.transforms.functional.to_tensor(tris) * 255.0
+        sample['image'] = torchvision.transforms.functional.to_tensor(imgs)
+        sample['image_name'] = 'Adobe_' + self.file_names[idx]
+
+        sample = self.transform(sample)
+        sample['trimap'][sample['trimap'] < 85] = 0
+        sample['trimap'][sample['trimap'] >= 170] = 1
+        sample['trimap'][sample['trimap'] >= 85] = 0.5
+
+        if self.multi_fg is not None:
+            sample['multi_fg'] = torch.tensor(self.multi_fg)
+
+        return sample
+
+
+class SIMTest(Dataset):
+    def __init__(self, data_dir, target_size=1024, multi_fg=None):
+        self.data_dir = data_dir
+        self.file_names = sorted(glob.glob(os.path.join(*[data_dir, '*', 'alpha', '*'])))  # [: 10]
+        test_trans = [
+            ResziePad(target_size=target_size),
+            GenBBox(bbox_offset_factor=0)
+        ]
+        self.transform = transforms.Compose(test_trans)
+        self.multi_fg = multi_fg
+
+    def __len__(self):  # 1000
+        return len(self.file_names)
+
+    def __getitem__(self, idx):
+        phas = Image.open(self.file_names[idx]).convert('L')
+        # tris = Image.open(self.file_names[idx].replace('alpha', 'trimap'))
+        imgs = Image.open(self.file_names[idx].replace('alpha', 'merged'))
+        sample = {
+            'ori_h_w': (imgs.size[1], imgs.size[0]),
+            'data_type': 'SIM'
+        }
+
+        sample['alpha'] = torchvision.transforms.functional.to_tensor(phas)  # [1, H, W] 0.0 ~ 1.0
+        # sample['trimap'] = torchvision.transforms.functional.to_tensor(tris) * 255.0
+        sample['image'] = torchvision.transforms.functional.to_tensor(imgs)
+        sample['image_name'] = 'SIM_{}_{}'.format(self.file_names[idx].split('/')[-3], self.file_names[idx].split('/')[-1])
+
+        sample = self.transform(sample)
+        # sample['trimap'][sample['trimap'] < 85] = 0
+        # sample['trimap'][sample['trimap'] >= 170] = 1
+        # sample['trimap'][sample['trimap'] >= 85] = 0.5
+
+        if self.multi_fg is not None:
+            sample['multi_fg'] = torch.tensor(self.multi_fg)
+
+        return sample
+    
+
+class RW100Test(Dataset):
+    def __init__(self, data_dir, target_size=1024, multi_fg=None):
+        self.data_dir = data_dir
+        self.file_names = sorted(glob.glob(os.path.join(*[data_dir, 'mask', '*'])))
+
+        self.name_to_idx = dict()
+        for idx, file_name in enumerate(self.file_names):
+            self.name_to_idx[file_name.split('/')[-1].split('.')[0]] = idx
+            
+        test_trans = [
+            ResziePad(target_size=target_size),
+            GenBBox(bbox_offset_factor=0, train_or_test='test', dataset_type='RW100')
+        ]
+        self.transform = transforms.Compose(test_trans)
+        self.multi_fg = multi_fg
+
+    def __len__(self):  # 1000
+        return len(self.file_names)
+
+    def __getitem__(self, idx):
+        phas = Image.open(self.file_names[idx]).convert('L')
+        imgs = Image.open(self.file_names[idx].replace('mask', 'image')[:-6] + '.jpg')
+        sample = {
+            'ori_h_w': (imgs.size[1], imgs.size[0]),
+            'data_type': 'RW100'
+        }
+
+        sample['alpha'] = torchvision.transforms.functional.to_tensor(phas)  # [1, H, W] 0.0 ~ 1.0
+        sample['image'] = torchvision.transforms.functional.to_tensor(imgs)
+        sample['image_name'] = 'RW100_' + self.file_names[idx].split('/')[-1]
+        
+        sample = self.transform(sample)
+
+        if self.multi_fg is not None:
+            sample['multi_fg'] = torch.tensor(self.multi_fg)
+
+        return sample
+    
+    
+class AIM500Test(Dataset):
+    def __init__(self, data_dir, target_size=1024, multi_fg=None):
+        self.data_dir = data_dir
+        self.file_names = sorted(glob.glob(os.path.join(*[data_dir, 'original', '*'])))
+
+        self.name_to_idx = dict()
+        for idx, file_name in enumerate(self.file_names):
+            self.name_to_idx[file_name.split('/')[-1].split('.')[0]] = idx
+
+        test_trans = [
+            ResziePad(target_size=target_size),
+            GenBBox(bbox_offset_factor=0)
+        ]
+        self.transform = transforms.Compose(test_trans)
+        self.multi_fg = multi_fg
+
+    def __len__(self):  # 1000
+        return len(self.file_names)
+
+    def __getitem__(self, idx):
+        phas = Image.open(self.file_names[idx].replace('original', 'mask').replace('jpg', 'png')).convert('L')
+        # tris = Image.open(self.file_names[idx].replace('original', 'trimap').replace('jpg', 'png')).convert('L')
+        imgs = Image.open(self.file_names[idx])
+        sample = {
+            'ori_h_w': (imgs.size[1], imgs.size[0]),
+            'data_type': 'AIM500'
+        }
+
+        sample['alpha'] = torchvision.transforms.functional.to_tensor(phas)  # [1, H, W] 0.0 ~ 1.0
+        # sample['trimap'] = torchvision.transforms.functional.to_tensor(tris) * 255.0
+        sample['image'] = torchvision.transforms.functional.to_tensor(imgs)
+        sample['image_name'] = 'AIM500_' + self.file_names[idx].split('/')[-1]
+
+        sample = self.transform(sample)
+        # sample['trimap'][sample['trimap'] < 85] = 0
+        # sample['trimap'][sample['trimap'] >= 170] = 1
+        # sample['trimap'][sample['trimap'] >= 85] = 0.5
+
+        if self.multi_fg is not None:
+            sample['multi_fg'] = torch.tensor(self.multi_fg)
+
+        return sample
+
+
+class RWP636Test(Dataset):
+    def __init__(self, data_dir, target_size=1024, multi_fg=None):
+        self.data_dir = data_dir
+        self.file_names = sorted(glob.glob(os.path.join(*[data_dir, 'image', '*'])))
+
+        self.name_to_idx = dict()
+        for idx, file_name in enumerate(self.file_names):
+            self.name_to_idx[file_name.split('/')[-1].split('.')[0]] = idx
+
+        test_trans = [
+            ResziePad(target_size=target_size),
+            GenBBox(bbox_offset_factor=0)
+        ]
+        self.transform = transforms.Compose(test_trans)
+        self.multi_fg = multi_fg
+
+    def __len__(self):  # 1000
+        return len(self.file_names)
+
+    def __getitem__(self, idx):
+        phas = Image.open(self.file_names[idx].replace('image', 'alpha').replace('jpg', 'png')).convert('L')
+        imgs = Image.open(self.file_names[idx])
+        sample = {
+            'ori_h_w': (imgs.size[1], imgs.size[0]),
+            'data_type': 'RWP636'
+        }
+
+        sample['alpha'] = torchvision.transforms.functional.to_tensor(phas)  # [1, H, W] 0.0 ~ 1.0
+        sample['image'] = torchvision.transforms.functional.to_tensor(imgs)
+        sample['image_name'] = 'RWP636_' + self.file_names[idx].split('/')[-1]
+
+        sample = self.transform(sample)
+
+        if self.multi_fg is not None:
+            sample['multi_fg'] = torch.tensor(self.multi_fg)
+
+        return sample
+
+
+class AM2KTest(Dataset):
+    def __init__(self, data_dir, target_size=1024, multi_fg=None):
+        self.data_dir = data_dir
+        self.file_names = sorted(glob.glob(os.path.join(*[data_dir, 'validation/original', '*'])))
+        test_trans = [
+            ResziePad(target_size=target_size),
+            GenBBox(bbox_offset_factor=0)
+        ]
+        self.transform = transforms.Compose(test_trans)
+        self.multi_fg = multi_fg
+
+    def __len__(self):  # 1000
+        return len(self.file_names)
+
+    def __getitem__(self, idx):
+        phas = Image.open(self.file_names[idx].replace('original', 'mask').replace('jpg', 'png')).convert('L')
+        # tris = Image.open(self.file_names[idx].replace('original', 'trimap').replace('jpg', 'png')).convert('L')
+        imgs = Image.open(self.file_names[idx])
+        sample = {
+            'ori_h_w': (imgs.size[1], imgs.size[0]),
+            'data_type': 'AM2K'
+        }
+
+        sample['alpha'] = torchvision.transforms.functional.to_tensor(phas)  # [1, H, W] 0.0 ~ 1.0
+        # sample['trimap'] = torchvision.transforms.functional.to_tensor(tris) * 255.0
+        sample['image'] = torchvision.transforms.functional.to_tensor(imgs)
+        sample['image_name'] = 'AM2K_' + self.file_names[idx].split('/')[-1]
+
+        sample = self.transform(sample)
+        # sample['trimap'][sample['trimap'] < 85] = 0
+        # sample['trimap'][sample['trimap'] >= 170] = 1
+        # sample['trimap'][sample['trimap'] >= 85] = 0.5
+
+        if self.multi_fg is not None:
+            sample['multi_fg'] = torch.tensor(self.multi_fg)
+
+        return sample
+
+
+class P3M500Test(Dataset):
+    def __init__(self, data_dir, target_size=1024, multi_fg=None):
+        self.data_dir = data_dir
+        self.file_names = sorted(glob.glob(os.path.join(*[data_dir, 'original_image', '*'])))
+
+        self.name_to_idx = dict()
+        for idx, file_name in enumerate(self.file_names):
+            self.name_to_idx[file_name.split('/')[-1].split('.')[0]] = idx
+
+        test_trans = [
+            ResziePad(target_size=target_size),
+            GenBBox(bbox_offset_factor=0)
+        ]
+        self.transform = transforms.Compose(test_trans)
+        self.multi_fg = multi_fg
+
+    def __len__(self):  # 1000
+        return len(self.file_names)
+
+    def __getitem__(self, idx):
+        phas = Image.open(self.file_names[idx].replace('original_image', 'mask').replace('jpg', 'png')).convert('L')
+        # tris = Image.open(self.file_names[idx].replace('original_image', 'trimap').replace('jpg', 'png')).convert('L')
+        imgs = Image.open(self.file_names[idx])
+        sample = {
+            'ori_h_w': (imgs.size[1], imgs.size[0]),
+            'data_type': 'P3M500'
+        }
+
+        sample['alpha'] = torchvision.transforms.functional.to_tensor(phas)  # [1, H, W] 0.0 ~ 1.0
+        # sample['trimap'] = torchvision.transforms.functional.to_tensor(tris) * 255.0
+        sample['image'] = torchvision.transforms.functional.to_tensor(imgs)
+        sample['image_name'] = 'P3M500_' + self.file_names[idx].split('/')[-1]
+
+        sample = self.transform(sample)
+        # sample['trimap'][sample['trimap'] < 85] = 0
+        # sample['trimap'][sample['trimap'] >= 170] = 1
+        # sample['trimap'][sample['trimap'] >= 85] = 0.5
+
+        if self.multi_fg is not None:
+            sample['multi_fg'] = torch.tensor(self.multi_fg)
+
+        return sample
+
+
+class MattingTest(Dataset):
+    def __init__(
+        self, 
+        data_type,
+        data_dir,
+        image_sub_path,
+        alpha_sub_path,
+        trimpa_sub_path=None,
+        target_size=1024, 
+        multi_fg=None,
+    ):
+        self.data_type = data_type
+        self.data_dir = data_dir
+
+        self.image_paths = sorted(glob.glob(os.path.join(*[data_dir, image_sub_path])))
+        self.alpha_paths = sorted(glob.glob(os.path.join(*[data_dir, alpha_sub_path])))
+        self.trimpa_paths = sorted(glob.glob(os.path.join(*[data_dir, trimpa_sub_path]))) if trimpa_sub_path is not None else None
+
+        self.name_to_idx = dict()
+        for idx, file_name in enumerate(self.image_paths):
+            self.name_to_idx[file_name.split('/')[-1].split('.')[0]] = idx
+
+        test_trans = [
+            Cv2ResziePad(target_size=target_size),
+            GenBBox(bbox_offset_factor=0)
+        ]
+        self.transform = transforms.Compose(test_trans)
+        self.multi_fg = multi_fg
+
+    def __len__(self):  # 1000
+        return len(self.image_paths)
+
+    def __getitem__(self, idx):
+
+        img = cv2.imread(self.image_paths[idx])
+        sample = {
+            'image': img.astype(np.float32) / 255,
+            'alpha': cv2.imread(self.alpha_paths[idx], 0).astype(np.float32) / 255,
+            'trimap': cv2.imread(self.trimpa_paths[idx], 0) if self.trimpa_paths is not None else None,
+            'ori_h_w': (img.shape[0], img.shape[1]),
+            'data_type': self.data_type,
+            'image_name': self.data_type + '_' + self.image_paths[idx].split('/')[-1]
+        }
+
+        sample = self.transform(sample)
+        if self.trimpa_paths is not None:
+            sample['trimap'][sample['trimap'] < 85] = 0
+            sample['trimap'][sample['trimap'] >= 170] = 1
+            sample['trimap'][sample['trimap'] >= 85] = 0.5
+        else:
+            del sample['trimap']
+
+        if self.multi_fg is not None:
+            sample['multi_fg'] = torch.tensor(self.multi_fg)
+
+        return sample
+
+
+def adobe_composition_collate_fn(batch):
+    new_batch = defaultdict(list)
+    for sub_batch in batch:
+        for key in sub_batch.keys():
+            new_batch[key].append(sub_batch[key])
+    for key in new_batch: 
+        if isinstance(new_batch[key][0], torch.Tensor):
+            new_batch[key] = torch.stack(new_batch[key])
+    return dict(new_batch)
+
+
+def build_d2_test_dataloader(
+    dataset,
+    mapper=None,
+    total_batch_size=None,
+    local_batch_size=None,
+    num_workers=0,
+    collate_fn=None
+):
+
+    assert (total_batch_size is None) != (
+        local_batch_size is None
+    ), "Either total_batch_size or local_batch_size must be specified"
+
+    world_size = comm.get_world_size()
+
+    if total_batch_size is not None:
+        assert (
+            total_batch_size > 0 and total_batch_size % world_size == 0
+        ), "Total batch size ({}) must be divisible by the number of gpus ({}).".format(
+            total_batch_size, world_size
+        )
+        batch_size = total_batch_size // world_size
+
+    if local_batch_size is not None:
+        batch_size = local_batch_size
+
+    logger = logging.getLogger(__name__)
+    if batch_size != 1:
+        logger.warning(
+            "When testing, batch size is set to 1. "
+            "This is the only mode that is supported for d2."
+        )
+
+    return build_detection_test_loader(
+        dataset=dataset,
+        mapper=mapper,
+        sampler=None,
+        num_workers=num_workers,
+        collate_fn=collate_fn,
+    )
+
+
+class AdobeCompositionEvaluator(DatasetEvaluator):
+
+    def __init__(
+        self, 
+        save_eval_results_step=-1, 
+        output_dir=None, 
+        eval_dataset_type=['Adobe'],
+        distributed=True,
+        eval_w_sam_hq_mask = False,
+    ):  
+
+        self.save_eval_results_step = save_eval_results_step
+        self.output_dir = output_dir
+        self.eval_index = 0
+        self.eval_dataset_type = eval_dataset_type
+        self.eval_w_sam_hq_mask = eval_w_sam_hq_mask
+
+        self._distributed = distributed
+        self._logger = logging.getLogger(__name__)
+
+    def reset(self):
+        self.eval_metric = dict()
+        for i in self.eval_dataset_type:
+            self.eval_metric[i + '_MSE'] = []
+            self.eval_metric[i + '_SAD'] = []
+            self.eval_metric[i + '_MAD'] = []
+            self.eval_metric[i + '_Grad'] = []
+            self.eval_metric[i + '_Conn'] = []
+
+        os.makedirs(self.output_dir, exist_ok=True) if self.output_dir is not None else None
+
+    def process(self, inputs, outputs):
+        """
+        Args:
+            inputs: {'alpha', 'trimap', 'image', 'bbox', 'image_name'}
+            outputs: [1, 1, H, W] 0. ~ 1.
+        """
+
+        # crop the black pad area
+        assert inputs['image'].shape[-1] == inputs['image'].shape[-2] == 1024 and len(inputs['ori_h_w']) == 1
+        inputs['ori_h_w'] = inputs['ori_h_w'][0]
+        before_pad_h, before_pad_w = int(1024 / max(inputs['ori_h_w']) * inputs['ori_h_w'][0] + 0.5), int(1024 / max(inputs['ori_h_w']) * inputs['ori_h_w'][1] + 0.5)
+        inputs['image'] = inputs['image'][:, :, :before_pad_h, :before_pad_w]
+        inputs['alpha'] = inputs['alpha'][:, :, :before_pad_h, :before_pad_w]
+
+        if self.eval_w_sam_hq_mask:
+            outputs, samhq_low_res_masks = outputs[0][:, :, :before_pad_h, :before_pad_w], outputs[1][:, :, :before_pad_h, :before_pad_w]
+            pred_alpha, label_alpha, samhq_low_res_masks = outputs.cpu().numpy(), inputs['alpha'].numpy(), (samhq_low_res_masks > 0).float().cpu()
+        else:
+            outputs = outputs[:, :, :before_pad_h, :before_pad_w]
+            pred_alpha, label_alpha = outputs.cpu().numpy(), inputs['alpha'].numpy()
+
+        # if 'trimap' in inputs.keys():
+        #     inputs['trimap'] = inputs['trimap'][:, :, :before_pad_h, :before_pad_w]
+        #     trimap = inputs['trimap'].numpy()
+        #     assert np.max(trimap) <= 1 and len(np.unique(trimap)) <= 3
+        #     sad_loss_unknown = compute_sad_loss(pred_alpha, label_alpha, trimap, area='unknown')
+        #     mse_loss_unknown = compute_mse_loss(pred_alpha, label_alpha, trimap, area='unknown')
+
+        #     self.eval_metric[inputs['data_type'][0] + '_unknown_mse (1e-3)'].append(mse_loss_unknown)
+        #     self.eval_metric[inputs['data_type'][0] + '_unknown_sad (1e3)'].append(sad_loss_unknown)
+
+        # calculate loss
+        assert np.max(pred_alpha) <= 1 and np.max(label_alpha) <= 1
+        eval_pred = np.uint8(pred_alpha[0, 0] * 255.0 + 0.5) * 1.0
+        eval_gt = label_alpha[0, 0] * 255.0
+
+        detailmap = np.zeros_like(eval_gt) + 128
+        mse_loss_ = compute_mse_loss(eval_pred, eval_gt, detailmap)
+        sad_loss_ = compute_sad_loss(eval_pred, eval_gt, detailmap)[0]
+        mad_loss_ = compute_mad_loss(eval_pred, eval_gt, detailmap)
+        grad_loss_ = compute_gradient_loss(eval_pred, eval_gt, detailmap)
+        conn_loss_ = compute_connectivity_error(eval_pred, eval_gt, detailmap)
+
+        self.eval_metric[inputs['data_type'][0] + '_MSE'].append(mse_loss_)
+        self.eval_metric[inputs['data_type'][0] + '_SAD'].append(sad_loss_)
+        self.eval_metric[inputs['data_type'][0] + '_MAD'].append(mad_loss_)
+        self.eval_metric[inputs['data_type'][0] + '_Grad'].append(grad_loss_)
+        self.eval_metric[inputs['data_type'][0] + '_Conn'].append(conn_loss_)
+
+        # vis results
+        if self.save_eval_results_step != -1 and self.eval_index % self.save_eval_results_step == 0:
+            if self.eval_w_sam_hq_mask:
+                self.save_vis_results(inputs, pred_alpha, samhq_low_res_masks)
+            else:
+                self.save_vis_results(inputs, pred_alpha)
+        self.eval_index += 1
+
+    def save_vis_results(self, inputs, pred_alpha, samhq_low_res_masks=None):
+
+        # image
+        image = inputs['image'][0].permute(1, 2, 0) * 255.0
+        l, u, r, d = int(inputs['bbox'][0, 0, 0].item()), int(inputs['bbox'][0, 0, 1].item()), int(inputs['bbox'][0, 0, 2].item()), int(inputs['bbox'][0, 0, 3].item())
+        red_line = torch.tensor([[255., 0., 0.]], device=image.device, dtype=image.dtype)
+        image[u: d, l, :] = red_line
+        image[u: d, r, :] = red_line
+        image[u, l: r, :] = red_line
+        image[d, l: r, :] = red_line
+        image = np.uint8(image.numpy())
+
+        # trimap, pred_alpha, label_alpha
+        save_results = [image]
+
+        choice = [inputs['trimap'], torch.from_numpy(pred_alpha), inputs['alpha']] if 'trimap' in inputs.keys() else [torch.from_numpy(pred_alpha), inputs['alpha']]
+        for val in choice:
+            val = val[0].permute(1, 2, 0).repeat(1, 1, 3) * 255.0 + 0.5  # +0.5 and int() = round()
+            val = np.uint8(val.numpy())
+            save_results.append(val)
+
+        if samhq_low_res_masks is not None:
+            save_results.append(np.uint8(samhq_low_res_masks[0].permute(1, 2, 0).repeat(1, 1, 3).numpy() * 255.0))
+
+        save_results = np.concatenate(save_results, axis=1)
+        save_name = os.path.join(self.output_dir, inputs['image_name'][0])
+        Image.fromarray(save_results).save(save_name.replace('.jpg', '.png'))
+
+    def evaluate(self):
+        
+        if self._distributed:
+            comm.synchronize()
+            eval_metric = comm.gather(self.eval_metric, dst=0)
+
+            if not comm.is_main_process():
+                return {}
+            
+            merges_eval_metric = defaultdict(list)
+            for sub_eval_metric in eval_metric:
+                for key, val in sub_eval_metric.items():
+                    merges_eval_metric[key] += val
+            eval_metric = merges_eval_metric
+
+        else:
+            eval_metric = self.eval_metric
+
+        eval_results = {}
+
+        for key, val in eval_metric.items():
+            if len(val) != 0:
+                # if 'mse' in key:
+                #     eval_results[key] = np.array(val).mean() * 1e3
+                # else:
+                #     assert 'sad' in key
+                #     eval_results[key] = np.array(val).mean() / 1e3
+                eval_results[key] = np.array(val).mean()
+
+        return eval_results
+
+
+if __name__ == '__main__':
+    pass

data/evaluate.py

@@ -0,0 +1,102 @@
+import scipy.ndimage
+import numpy as np
+from skimage.measure import label
+import scipy.ndimage.morphology
+
+
+def gauss(x, sigma):
+    y = np.exp(-x ** 2 / (2 * sigma ** 2)) / (sigma * np.sqrt(2 * np.pi))
+    return y
+
+
+def dgauss(x, sigma):
+    y = -x * gauss(x, sigma) / (sigma ** 2)
+    return y
+
+
+def gaussgradient(im, sigma):
+    epsilon = 1e-2
+    halfsize = np.ceil(sigma * np.sqrt(-2 * np.log(np.sqrt(2 * np.pi) * sigma * epsilon))).astype(np.int32)
+    size = 2 * halfsize + 1
+    hx = np.zeros((size, size))
+    for i in range(0, size):
+        for j in range(0, size):
+            u = [i - halfsize, j - halfsize]
+            hx[i, j] = gauss(u[0], sigma) * dgauss(u[1], sigma)
+
+    hx = hx / np.sqrt(np.sum(np.abs(hx) * np.abs(hx)))
+    hy = hx.transpose()
+
+    gx = scipy.ndimage.convolve(im, hx, mode='nearest')
+    gy = scipy.ndimage.convolve(im, hy, mode='nearest')
+
+    return gx, gy
+
+
+def compute_gradient_loss(pred, target, trimap):
+
+    pred = pred / 255.0
+    target = target / 255.0
+
+    pred_x, pred_y = gaussgradient(pred, 1.4)
+    target_x, target_y = gaussgradient(target, 1.4)
+
+    pred_amp = np.sqrt(pred_x ** 2 + pred_y ** 2)
+    target_amp = np.sqrt(target_x ** 2 + target_y ** 2)
+
+    error_map = (pred_amp - target_amp) ** 2
+    loss = np.sum(error_map[trimap == 128])
+
+    return loss / 1000.
+
+
+def getLargestCC(segmentation):
+    labels = label(segmentation, connectivity=1)
+    largestCC = labels == np.argmax(np.bincount(labels.flat))
+    return largestCC
+
+
+def compute_connectivity_error(pred, target, trimap, step=0.1):
+    pred = pred / 255.0
+    target = target / 255.0
+    h, w = pred.shape
+
+    thresh_steps = list(np.arange(0, 1 + step, step))
+    l_map = np.ones_like(pred, dtype=np.float32) * -1
+    for i in range(1, len(thresh_steps)):
+        pred_alpha_thresh = (pred >= thresh_steps[i]).astype(np.int32)
+        target_alpha_thresh = (target >= thresh_steps[i]).astype(np.int32)
+
+        omega = getLargestCC(pred_alpha_thresh * target_alpha_thresh).astype(np.int32)
+        flag = ((l_map == -1) & (omega == 0)).astype(np.int32)
+        l_map[flag == 1] = thresh_steps[i - 1]
+
+    l_map[l_map == -1] = 1
+
+    pred_d = pred - l_map
+    target_d = target - l_map
+    pred_phi = 1 - pred_d * (pred_d >= 0.15).astype(np.int32)
+    target_phi = 1 - target_d * (target_d >= 0.15).astype(np.int32)
+    loss = np.sum(np.abs(pred_phi - target_phi)[trimap == 128])
+
+    return loss / 1000.
+
+
+def compute_mse_loss(pred, target, trimap):
+    error_map = (pred - target) / 255.0
+    loss = np.sum((error_map ** 2) * (trimap == 128)) / (np.sum(trimap == 128) + 1e-8)
+
+    return loss
+
+
+def compute_sad_loss(pred, target, trimap):
+    error_map = np.abs((pred - target) / 255.0)
+    loss = np.sum(error_map * (trimap == 128))
+
+    return loss / 1000, np.sum(trimap == 128) / 1000
+
+def compute_mad_loss(pred, target, trimap):
+    error_map = np.abs((pred - target) / 255.0)
+    loss = np.sum(error_map * (trimap == 128)) / (np.sum(trimap == 128) + 1e-8)
+
+    return loss

data/p3m10k_dataset.py

@@ -0,0 +1,325 @@
+import os
+import torch
+import numpy as np
+import cv2
+from torch.utils.data import Dataset
+from torchvision import transforms
+import math
+import torch.nn.functional as F
+
+
+class GenBBox(object):
+    def __init__(self, bbox_offset_factor = 0.1, random_crop_bbox = None, train_or_test = 'train', dataset_type = None, random_auto_matting=None):
+        self.bbox_offset_factor = bbox_offset_factor
+        self.random_crop_bbox = random_crop_bbox
+        self.train_or_test = train_or_test
+        self.dataset_type = dataset_type
+        self.random_auto_matting = random_auto_matting
+
+    def __call__(self, sample):
+
+        alpha = sample['alpha']  # [1, H, W] 0.0 ~ 1.0
+        indices = torch.nonzero(alpha[0], as_tuple=True)
+
+        if len(indices[0]) > 0:
+
+            min_x, min_y = torch.min(indices[1]), torch.min(indices[0])
+            max_x, max_y = torch.max(indices[1]), torch.max(indices[0])
+
+            if self.random_crop_bbox is not None and np.random.uniform(0, 1) < self.random_crop_bbox:
+                ori_h_w = (sample['alpha'].shape[-2], sample['alpha'].shape[-1])
+                sample['alpha'] = F.interpolate(sample['alpha'][None, :, min_y: max_y + 1, min_x: max_x + 1], size=ori_h_w, mode='bilinear', align_corners=False)[0]
+                sample['image'] = F.interpolate(sample['image'][None, :, min_y: max_y + 1, min_x: max_x + 1], size=ori_h_w, mode='bilinear', align_corners=False)[0]
+                sample['trimap'] = F.interpolate(sample['trimap'][None, :, min_y: max_y + 1, min_x: max_x + 1], size=ori_h_w, mode='nearest')[0]
+                bbox = torch.tensor([[0, 0, ori_h_w[1] - 1, ori_h_w[0] - 1]])
+
+            elif self.bbox_offset_factor != 0:
+                bbox_w = max(1, max_x - min_x)
+                bbox_h = max(1, max_y - min_y)
+                offset_w = math.ceil(self.bbox_offset_factor * bbox_w)
+                offset_h = math.ceil(self.bbox_offset_factor * bbox_h)
+
+                min_x = max(0, min_x + np.random.randint(-offset_w, offset_w))
+                max_x = min(alpha.shape[2] - 1, max_x + np.random.randint(-offset_w, offset_w))
+                min_y = max(0, min_y + np.random.randint(-offset_h, offset_h))
+                max_y = min(alpha.shape[1] - 1, max_y + np.random.randint(-offset_h, offset_h))
+                bbox = torch.tensor([[min_x, min_y, max_x, max_y]])
+            else:
+                bbox = torch.tensor([[min_x, min_y, max_x, max_y]])
+            
+            if self.random_auto_matting is not None and np.random.uniform(0, 1) < self.random_auto_matting:
+                bbox = torch.tensor([[0, 0, alpha.shape[2] - 1, alpha.shape[1] - 1]])
+
+        else:
+            bbox = torch.zeros(1, 4)
+
+        sample['bbox'] = bbox.float()
+        return sample
+
+def random_interp():
+    return np.random.choice([cv2.INTER_NEAREST, cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_LANCZOS4])
+
+
+class SplitConcatImage(object):
+
+    def __init__(self, concat_num=4, wo_mask_to_mattes=False):
+        self.concat_num = concat_num
+        self.wo_mask_to_mattes = wo_mask_to_mattes
+        if self.wo_mask_to_mattes:
+            assert self.concat_num == 5
+
+    def __call__(self, concat_image):
+        if isinstance(concat_image, list):
+            concat_image, image_path = concat_image[0], concat_image[1]
+        else:
+            image_path = None
+        H, W, _ = concat_image.shape
+
+        concat_num = self.concat_num
+        if image_path is not None:
+            if '06-14' in image_path:
+                concat_num = 4
+            elif 'ori_mask' in image_path or 'SEMat' in image_path:
+                concat_num = 3
+            else:
+                concat_num = 5
+        
+        assert W % concat_num == 0
+        W = W // concat_num
+
+        image = concat_image[:H, :W]
+        if self.concat_num != 3:
+            trimap = concat_image[:H, (concat_num - 2) * W: (concat_num - 1) * W]
+            if self.wo_mask_to_mattes:
+                alpha = concat_image[:H, 2 * W: 3 * W]
+            else:
+                alpha = concat_image[:H, (concat_num - 1) * W: concat_num * W]
+        else:
+            trimap = concat_image[:H, (concat_num - 1) * W: concat_num * W]
+            alpha = concat_image[:H, (concat_num - 2) * W: (concat_num - 1) * W]
+
+        return {'image': image, 'trimap': trimap, 'alpha': alpha}
+
+
+class RandomHorizontalFlip(object):
+
+    def __init__(self, prob=0.5):
+        self.prob = prob
+
+    def __call__(self, sample):
+        if np.random.uniform(0, 1) < self.prob:
+            for key in sample.keys():
+                sample[key] = cv2.flip(sample[key], 1)
+        return sample
+
+class EmptyAug(object):
+    def __call__(self, sample):
+        return sample
+
+class RandomReszieCrop(object):
+
+    def __init__(self, output_size=1024, aug_scale_min=0.5, aug_scale_max=1.5):
+        self.desired_size = output_size
+        self.aug_scale_min = aug_scale_min
+        self.aug_scale_max = aug_scale_max
+
+    def __call__(self, sample):
+        H, W, _ = sample['image'].shape
+        sample['trimap'] = sample['trimap'][:, :, None].repeat(3, axis=-1)
+        sample['alpha'] = sample['alpha'][:, :, None].repeat(3, axis=-1)
+
+        if self.aug_scale_min == 1.0 and self.aug_scale_max == 1.0:
+            crop_H, crop_W = H, W
+            crop_y1, crop_y2 = 0, crop_H
+            crop_x1, crop_x2 = 0, crop_W
+            scale_W, scaled_H = W, H
+        elif self.aug_scale_min == -1.0 and self.aug_scale_max == -1.0:
+            scale = min(self.desired_size / H, self.desired_size / W)
+            scaled_H, scale_W = round(H * scale), round(W * scale)
+            crop_H, crop_W = scaled_H, scale_W
+            crop_y1, crop_y2 = 0, crop_H
+            crop_x1, crop_x2 = 0, crop_W
+        else:
+            # random size
+            random_scale = np.random.uniform(0, 1) * (self.aug_scale_max - self.aug_scale_min) + self.aug_scale_min  # random_val: 0.5 ~ 1.5
+            scaled_size = round(random_scale * self.desired_size)
+
+            scale = min(scaled_size / H, scaled_size / W)
+            scaled_H, scale_W = round(H * scale), round(W * scale)
+
+            # random crop
+            crop_H, crop_W = min(self.desired_size, scaled_H), min(self.desired_size, scale_W)  # crop_size
+            margin_H, margin_W = max(scaled_H - crop_H, 0), max(scale_W - crop_W, 0)
+            offset_H, offset_W = np.random.randint(0, margin_H + 1), np.random.randint(0, margin_W + 1)
+            crop_y1, crop_y2 = offset_H, offset_H + crop_H
+            crop_x1, crop_x2 = offset_W, offset_W + crop_W
+
+        for key in sample.keys():
+            sample[key] = cv2.resize(sample[key], (scale_W, scaled_H), interpolation=random_interp())[crop_y1: crop_y2, crop_x1: crop_x2, :]  # resize and crop
+            padding = np.zeros(shape=(self.desired_size, self.desired_size, 3), dtype=sample[key].dtype)  # pad to desired_size
+            padding[: crop_H, : crop_W, :] = sample[key]
+            sample[key] = padding
+
+        return sample
+
+
+class RandomJitter(object):
+    """
+    Random change the hue of the image
+    """
+
+    def __call__(self, sample):
+
+        image = sample['image']
+
+        # convert to HSV space, convert to float32 image to keep precision during space conversion.
+        image = cv2.cvtColor(image.astype(np.float32)/255.0, cv2.COLOR_BGR2HSV)
+        # Hue noise
+        hue_jitter = np.random.randint(-40, 40)
+        image[:, :, 0] = np.remainder(image[:, :, 0].astype(np.float32) + hue_jitter, 360)
+        # Saturation noise
+        sat_bar = image[:, :, 1].mean()
+
+        sat_jitter = np.random.rand()*(1.1 - sat_bar)/5 - (1.1 - sat_bar) / 10
+        sat = image[:, :, 1]
+        sat = np.abs(sat + sat_jitter)
+        sat[sat>1] = 2 - sat[sat>1]
+        image[:, :, 1] = sat
+        # Value noise
+        val_bar = image[:, :, 2].mean()
+
+        val_jitter = np.random.rand()*(1.1 - val_bar)/5-(1.1 - val_bar) / 10
+        val = image[:, :, 2]
+        val = np.abs(val + val_jitter)
+        val[val>1] = 2 - val[val>1]
+        image[:, :, 2] = val
+        # convert back to BGR space
+        image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR)
+        sample['image'] = image * 255
+
+        return sample
+
+
+class ToTensor(object):
+
+    def __call__(self, sample):
+        image, alpha, trimap = sample['image'][:, :, ::-1], sample['alpha'], sample['trimap']
+
+        # image
+        image = image.transpose((2, 0, 1)) / 255.
+        sample['image'] = torch.from_numpy(image).float()
+
+        # alpha
+        alpha = alpha.transpose((2, 0, 1))[0: 1] / 255.
+        alpha[alpha < 0 ] = 0
+        alpha[alpha > 1] = 1
+        sample['alpha'] = torch.from_numpy(alpha).float()
+
+        # trimap
+        trimap = trimap.transpose((2, 0, 1))[0: 1] / 1.
+        sample['trimap'] = torch.from_numpy(trimap).float()
+        sample['trimap'][sample['trimap'] < 85] = 0
+        sample['trimap'][sample['trimap'] >= 170] = 1
+        sample['trimap'][sample['trimap'] >= 85] = 0.5
+
+        return sample
+    
+
+class GenTrimap(object):
+    def __init__(self):
+        self.erosion_kernels = [None] + [cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (size, size)) for size in range(1,100)]
+
+    def __call__(self, sample):
+        alpha = sample['alpha']
+        h, w = alpha.shape
+
+        max_kernel_size = max(30, int((min(h,w) / 2048) * 30))
+
+        ### generate trimap
+        fg_mask = (alpha / 255.0 + 1e-5).astype(np.int32).astype(np.uint8)
+        bg_mask = (1 - alpha / 255.0 + 1e-5).astype(np.int32).astype(np.uint8)
+        fg_mask = cv2.erode(fg_mask, self.erosion_kernels[np.random.randint(1, max_kernel_size)])
+        bg_mask = cv2.erode(bg_mask, self.erosion_kernels[np.random.randint(1, max_kernel_size)])
+
+        trimap = np.ones_like(alpha) * 128
+        trimap[fg_mask == 1] = 255
+        trimap[bg_mask == 1] = 0
+
+        trimap = cv2.resize(trimap, (w,h), interpolation=cv2.INTER_NEAREST)
+        sample['trimap'] = trimap
+
+        return sample
+    
+
+class P3MData(Dataset):
+    def __init__(
+        self, 
+        data_root_path = '/root/data/my_path_b/public_data/data/matting/P3M-10k/train/blurred_image/', 
+        output_size = 1024, 
+        aug_scale_min = 0.8, 
+        aug_scale_max = 1.5,
+        with_bbox = True, 
+        bbox_offset_factor = 0.05,
+        num_ratio = 4.06,  # 9421 * 4.06 = 38249.26 (38251)
+    ):
+        
+        self.data_root_path = data_root_path
+        self.output_size = output_size
+        self.aug_scale_min = aug_scale_min
+        self.aug_scale_max = aug_scale_max
+        self.with_bbox = with_bbox
+        self.bbox_offset_factor = bbox_offset_factor
+        self.num_ratio = num_ratio
+
+        self.image_names = os.listdir(self.data_root_path)
+        self.image_names = [i for i in self.image_names if 'jpg' in i]
+        self.image_names.sort()
+
+        train_trans = [
+            RandomHorizontalFlip(prob=0 if hasattr(self, 'return_image_name') and self.return_image_name else 0.5),
+            GenTrimap(),
+            RandomReszieCrop(self.output_size, self.aug_scale_min, self.aug_scale_max),
+            RandomJitter(),
+            ToTensor(),
+            GenBBox(bbox_offset_factor=self.bbox_offset_factor)
+        ]
+        self.transform = transforms.Compose(train_trans)
+
+    def __getitem__(self, idx):
+
+        if self.num_ratio is not None:
+            if self.num_ratio < 1.0:
+                idx = np.random.randint(0, len(self.image_names))
+            else:
+                idx = idx % len(self.image_names)
+
+        image_path = os.path.join(self.data_root_path, self.image_names[idx])
+        alpha_path = image_path.replace('jpg', 'png').replace('blurred_image', 'mask')
+
+        sample = self.transform({
+            'image': cv2.imread(image_path),
+            'alpha': cv2.imread(alpha_path, 0),
+        })
+
+        sample['dataset_name'] = 'P3M'
+        sample['multi_fg'] = False
+
+        return sample
+
+    def __len__(self):
+        if self.num_ratio is not None:
+            return int(len(self.image_names) * self.num_ratio)
+        else:
+            return len(self.image_names)
+
+
+if __name__ == '__main__':
+
+    dataset = P3MData()
+    data = dataset[0]
+    print(len(dataset))
+    for key, val in data.items():
+        if isinstance(val, torch.Tensor):
+            print(key, val.shape, torch.min(val), torch.max(val), torch.unique(val))
+        else:
+            print(key, val)

data/rand_augment.py

@@ -0,0 +1,196 @@
+# copyright: https://github.com/ildoonet/pytorch-randaugment
+# code in this file is adpated from rpmcruz/autoaugment
+# https://github.com/rpmcruz/autoaugment/blob/master/transformations.py
+# This code is modified version of one of ildoonet, for randaugmentation of fixmatch.
+
+import random
+
+import PIL, PIL.ImageOps, PIL.ImageEnhance, PIL.ImageDraw
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image
+
+
+def AutoContrast(img, _):
+    return PIL.ImageOps.autocontrast(img)
+
+
+def Brightness(img, v):
+    assert v >= 0.0
+    return PIL.ImageEnhance.Brightness(img).enhance(v)
+
+
+def Color(img, v):
+    assert v >= 0.0
+    return PIL.ImageEnhance.Color(img).enhance(v)
+
+
+def Contrast(img, v):
+    assert v >= 0.0
+    return PIL.ImageEnhance.Contrast(img).enhance(v)
+
+
+def Equalize(img, _):
+    return PIL.ImageOps.equalize(img)
+
+
+def Invert(img, _):
+    return PIL.ImageOps.invert(img)
+
+
+def Identity(img, v):
+    return img
+
+
+def Posterize(img, v):  # [4, 8]
+    v = int(v)
+    v = max(1, v)
+    return PIL.ImageOps.posterize(img, v)
+
+
+def Rotate(img, v):  # [-30, 30]
+    #assert -30 <= v <= 30
+    #if random.random() > 0.5:
+    #    v = -v
+    return img.rotate(v)
+
+
+
+def Sharpness(img, v):  # [0.1,1.9]
+    assert v >= 0.0
+    return PIL.ImageEnhance.Sharpness(img).enhance(v)
+
+
+def ShearX(img, v):  # [-0.3, 0.3]
+    #assert -0.3 <= v <= 0.3
+    #if random.random() > 0.5:
+    #    v = -v
+    return img.transform(img.size, PIL.Image.AFFINE, (1, v, 0, 0, 1, 0))
+
+
+def ShearY(img, v):  # [-0.3, 0.3]
+    #assert -0.3 <= v <= 0.3
+    #if random.random() > 0.5:
+    #    v = -v
+    return img.transform(img.size, PIL.Image.AFFINE, (1, 0, 0, v, 1, 0))
+
+
+def TranslateX(img, v):  # [-150, 150] => percentage: [-0.45, 0.45]
+    #assert -0.3 <= v <= 0.3
+    #if random.random() > 0.5:
+    #    v = -v
+    v = v * img.size[0]
+    return img.transform(img.size, PIL.Image.AFFINE, (1, 0, v, 0, 1, 0))
+
+
+def TranslateXabs(img, v):  # [-150, 150] => percentage: [-0.45, 0.45]
+    #assert v >= 0.0
+    #if random.random() > 0.5:
+    #    v = -v
+    return img.transform(img.size, PIL.Image.AFFINE, (1, 0, v, 0, 1, 0))
+
+
+def TranslateY(img, v):  # [-150, 150] => percentage: [-0.45, 0.45]
+    #assert -0.3 <= v <= 0.3
+    #if random.random() > 0.5:
+    #    v = -v
+    v = v * img.size[1]
+    return img.transform(img.size, PIL.Image.AFFINE, (1, 0, 0, 0, 1, v))
+
+
+def TranslateYabs(img, v):  # [-150, 150] => percentage: [-0.45, 0.45]
+    #assert 0 <= v
+    #if random.random() > 0.5:
+    #    v = -v
+    return img.transform(img.size, PIL.Image.AFFINE, (1, 0, 0, 0, 1, v))
+
+
+def Solarize(img, v):  # [0, 256]
+    assert 0 <= v <= 256
+    return PIL.ImageOps.solarize(img, v)
+
+
+def Cutout(img, v):  #[0, 60] => percentage: [0, 0.2] => change to [0, 0.5]
+    assert 0.0 <= v <= 0.5
+    if v <= 0.:
+        return img
+
+    v = v * img.size[0]
+    return CutoutAbs(img, v)
+
+
+def CutoutAbs(img, v):  # [0, 60] => percentage: [0, 0.2]
+    # assert 0 <= v <= 20
+    if v < 0:
+        return img
+    w, h = img.size
+    x0 = np.random.uniform(w)
+    y0 = np.random.uniform(h)
+
+    x0 = int(max(0, x0 - v / 2.))
+    y0 = int(max(0, y0 - v / 2.))
+    x1 = min(w, x0 + v)
+    y1 = min(h, y0 + v)
+
+    xy = (x0, y0, x1, y1)
+    color = (125, 123, 114)
+    # color = (0, 0, 0)
+    img = img.copy()
+    PIL.ImageDraw.Draw(img).rectangle(xy, color)
+    return img
+
+    
+def augment_list():  
+    l = [
+        (AutoContrast, 0, 1),
+        (Brightness, 0.05, 0.95),
+        (Color, 0.05, 0.95),
+        (Contrast, 0.05, 0.95),
+        (Equalize, 0, 1),
+        (Identity, 0, 1),
+        (Posterize, 4, 8),
+        # (Rotate, -30, 30),
+        (Sharpness, 0.05, 0.95),
+        # (ShearX, -0.3, 0.3),
+        # (ShearY, -0.3, 0.3),
+        (Solarize, 0, 256),
+        # (TranslateX, -0.3, 0.3),
+        # (TranslateY, -0.3, 0.3)
+    ]
+    return l
+
+    
+class RandAugment:
+    def __init__(self, n, m):
+        self.n = n
+        self.m = m      # [0, 30] in fixmatch, deprecated.
+        self.augment_list = augment_list()
+
+        
+    def __call__(self, img, cutout=True):
+        ops = random.choices(self.augment_list, k=self.n)
+        for op, min_val, max_val in ops:
+            val = min_val + float(max_val - min_val)*random.random()
+            img = op(img, val) 
+        if cutout:
+            cutout_val = random.random() * 0.5 
+            img = Cutout(img, cutout_val) #for fixmatch
+        return img
+
+    
+if __name__ == '__main__':
+    # randaug = RandAugment(3,5)
+    # print(randaug)
+    # for item in randaug.augment_list:
+    #     print(item)
+    import os
+
+    os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
+    img = PIL.Image.open('./u.jpg')
+    randaug = RandAugment(3,6)
+    img = randaug(img)
+    import matplotlib
+    from matplotlib import pyplot as plt 
+    plt.imshow(img)
+    plt.show()

data/refmatte_dataset.py

@@ -0,0 +1,418 @@
+import os
+import torch
+import numpy as np
+import cv2
+from torch.utils.data import Dataset
+from torchvision import transforms
+import random
+import imgaug.augmenters as iaa
+import numbers
+import math
+
+
+def random_interp():
+    return np.random.choice([cv2.INTER_NEAREST, cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_LANCZOS4])
+
+class RandomAffine(object):
+    """
+    Random affine translation
+    """
+    def __init__(self, degrees, translate=None, scale=None, shear=None, flip=None, resample=False, fillcolor=0):
+        if isinstance(degrees, numbers.Number):
+            if degrees < 0:
+                raise ValueError("If degrees is a single number, it must be positive.")
+            self.degrees = (-degrees, degrees)
+        else:
+            assert isinstance(degrees, (tuple, list)) and len(degrees) == 2, \
+                "degrees should be a list or tuple and it must be of length 2."
+            self.degrees = degrees
+
+        if translate is not None:
+            assert isinstance(translate, (tuple, list)) and len(translate) == 2, \
+                "translate should be a list or tuple and it must be of length 2."
+            for t in translate:
+                if not (0.0 <= t <= 1.0):
+                    raise ValueError("translation values should be between 0 and 1")
+        self.translate = translate
+
+        if scale is not None:
+            assert isinstance(scale, (tuple, list)) and len(scale) == 2, \
+                "scale should be a list or tuple and it must be of length 2."
+            for s in scale:
+                if s <= 0:
+                    raise ValueError("scale values should be positive")
+        self.scale = scale
+
+        if shear is not None:
+            if isinstance(shear, numbers.Number):
+                if shear < 0:
+                    raise ValueError("If shear is a single number, it must be positive.")
+                self.shear = (-shear, shear)
+            else:
+                assert isinstance(shear, (tuple, list)) and len(shear) == 2, \
+                    "shear should be a list or tuple and it must be of length 2."
+                self.shear = shear
+        else:
+            self.shear = shear
+
+        self.resample = resample
+        self.fillcolor = fillcolor
+        self.flip = flip
+
+    @staticmethod
+    def get_params(degrees, translate, scale_ranges, shears, flip, img_size):
+        """Get parameters for affine transformation
+
+        Returns:
+            sequence: params to be passed to the affine transformation
+        """
+        angle = random.uniform(degrees[0], degrees[1])
+        if translate is not None:
+            max_dx = translate[0] * img_size[0]
+            max_dy = translate[1] * img_size[1]
+            translations = (np.round(random.uniform(-max_dx, max_dx)),
+                            np.round(random.uniform(-max_dy, max_dy)))
+        else:
+            translations = (0, 0)
+
+        if scale_ranges is not None:
+            scale = (random.uniform(scale_ranges[0], scale_ranges[1]),
+                     random.uniform(scale_ranges[0], scale_ranges[1]))
+        else:
+            scale = (1.0, 1.0)
+
+        if shears is not None:
+            shear = random.uniform(shears[0], shears[1])
+        else:
+            shear = 0.0
+
+        if flip is not None:
+            flip = (np.random.rand(2) < flip).astype(np.int32) * 2 - 1
+
+        return angle, translations, scale, shear, flip
+
+    def __call__(self, sample):
+        fg, alpha = sample['fg'], sample['alpha']
+        rows, cols, ch = fg.shape
+        if np.maximum(rows, cols) < 1024:
+            params = self.get_params((0, 0), self.translate, self.scale, self.shear, self.flip, fg.size)
+        else:
+            params = self.get_params(self.degrees, self.translate, self.scale, self.shear, self.flip, fg.size)
+
+        center = (cols * 0.5 + 0.5, rows * 0.5 + 0.5)
+        M = self._get_inverse_affine_matrix(center, *params)
+        M = np.array(M).reshape((2, 3))
+
+        fg = cv2.warpAffine(fg, M, (cols, rows), flags=random_interp() + cv2.WARP_INVERSE_MAP)
+        alpha = cv2.warpAffine(alpha, M, (cols, rows), flags=random_interp() + cv2.WARP_INVERSE_MAP)
+
+        sample['fg'], sample['alpha'] = fg, alpha
+
+        return sample
+
+    @ staticmethod
+    def _get_inverse_affine_matrix(center, angle, translate, scale, shear, flip):
+        # Helper method to compute inverse matrix for affine transformation
+
+        # As it is explained in PIL.Image.rotate
+        # We need compute INVERSE of affine transformation matrix: M = T * C * RSS * C^-1
+        # where T is translation matrix: [1, 0, tx | 0, 1, ty | 0, 0, 1]
+        # C is translation matrix to keep center: [1, 0, cx | 0, 1, cy | 0, 0, 1]
+        # RSS is rotation with scale and shear matrix
+        # It is different from the original function in torchvision
+        # The order are changed to flip -> scale -> rotation -> shear
+        # x and y have different scale factors
+        # RSS(shear, a, scale, f) = [ cos(a + shear)*scale_x*f -sin(a + shear)*scale_y     0]
+        # [ sin(a)*scale_x*f          cos(a)*scale_y             0]
+        # [     0                       0                      1]
+        # Thus, the inverse is M^-1 = C * RSS^-1 * C^-1 * T^-1
+
+        angle = math.radians(angle)
+        shear = math.radians(shear)
+        scale_x = 1.0 / scale[0] * flip[0]
+        scale_y = 1.0 / scale[1] * flip[1]
+
+        # Inverted rotation matrix with scale and shear
+        d = math.cos(angle + shear) * math.cos(angle) + math.sin(angle + shear) * math.sin(angle)
+        matrix = [
+            math.cos(angle) * scale_x, math.sin(angle + shear) * scale_x, 0,
+            -math.sin(angle) * scale_y, math.cos(angle + shear) * scale_y, 0
+        ]
+        matrix = [m / d for m in matrix]
+
+        # Apply inverse of translation and of center translation: RSS^-1 * C^-1 * T^-1
+        matrix[2] += matrix[0] * (-center[0] - translate[0]) + matrix[1] * (-center[1] - translate[1])
+        matrix[5] += matrix[3] * (-center[0] - translate[0]) + matrix[4] * (-center[1] - translate[1])
+
+        # Apply center translation: C * RSS^-1 * C^-1 * T^-1
+        matrix[2] += center[0]
+        matrix[5] += center[1]
+
+        return matrix
+    
+
+class GenTrimap(object):
+    def __init__(self):
+        self.erosion_kernels = [None] + [cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (size, size)) for size in range(1,100)]
+
+    def __call__(self, sample):
+        alpha = sample['alpha']
+        h, w = alpha.shape
+
+        max_kernel_size = max(30, int((min(h,w) / 2048) * 30))
+
+        ### generate trimap
+        fg_mask = (alpha + 1e-5).astype(np.int32).astype(np.uint8)
+        bg_mask = (1 - alpha + 1e-5).astype(np.int32).astype(np.uint8)
+        fg_mask = cv2.erode(fg_mask, self.erosion_kernels[np.random.randint(1, max_kernel_size)])
+        bg_mask = cv2.erode(bg_mask, self.erosion_kernels[np.random.randint(1, max_kernel_size)])
+
+        trimap = np.ones_like(alpha) * 128
+        trimap[fg_mask == 1] = 255
+        trimap[bg_mask == 1] = 0
+
+        trimap = cv2.resize(trimap, (w,h), interpolation=cv2.INTER_NEAREST)
+        sample['trimap'] = trimap
+
+        return sample
+    
+
+class RandomCrop(object):
+    """
+    Crop randomly the image in a sample, retain the center 1/4 images, and resize to 'output_size'
+
+    :param output_size (tuple or int): Desired output size. If int, square crop
+            is made.
+    """
+
+    def __init__(self, output_size=(1024, 1024)):
+        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
+        self.margin = output_size[0] // 2
+
+    def __call__(self, sample):
+        fg, alpha, trimap, name = sample['fg'],  sample['alpha'], sample['trimap'], sample['image_name']
+        bg = sample['bg']
+        h, w = trimap.shape
+        bg = cv2.resize(bg, (w, h), interpolation=random_interp())
+        if w < self.output_size[0]+1 or h < self.output_size[1]+1:
+            ratio = 1.1*self.output_size[0]/h if h < w else 1.1*self.output_size[1]/w
+            # self.logger.warning("Size of {} is {}.".format(name, (h, w)))
+            while h < self.output_size[0]+1 or w < self.output_size[1]+1:
+                fg = cv2.resize(fg, (int(w*ratio), int(h*ratio)), interpolation=random_interp())
+                alpha = cv2.resize(alpha, (int(w*ratio), int(h*ratio)),
+                                   interpolation=random_interp())
+                trimap = cv2.resize(trimap, (int(w*ratio), int(h*ratio)), interpolation=cv2.INTER_NEAREST)
+                bg = cv2.resize(bg, (int(w*ratio), int(h*ratio)), interpolation=random_interp())
+                h, w = trimap.shape
+        small_trimap = cv2.resize(trimap, (w//4, h//4), interpolation=cv2.INTER_NEAREST)
+        unknown_list = list(zip(*np.where(small_trimap[self.margin//4:(h-self.margin)//4,
+                                                       self.margin//4:(w-self.margin)//4] == 128)))
+        unknown_num = len(unknown_list)
+        if len(unknown_list) < 10:
+            left_top = (np.random.randint(0, h-self.output_size[0]+1), np.random.randint(0, w-self.output_size[1]+1))
+        else:
+            idx = np.random.randint(unknown_num)
+            left_top = (unknown_list[idx][0]*4, unknown_list[idx][1]*4)
+
+        fg_crop = fg[left_top[0]:left_top[0]+self.output_size[0], left_top[1]:left_top[1]+self.output_size[1],:]
+        alpha_crop = alpha[left_top[0]:left_top[0]+self.output_size[0], left_top[1]:left_top[1]+self.output_size[1]]
+        bg_crop = bg[left_top[0]:left_top[0]+self.output_size[0], left_top[1]:left_top[1]+self.output_size[1],:]
+        trimap_crop = trimap[left_top[0]:left_top[0]+self.output_size[0], left_top[1]:left_top[1]+self.output_size[1]]
+
+        if len(np.where(trimap==128)[0]) == 0:
+            fg_crop = cv2.resize(fg, self.output_size[::-1], interpolation=random_interp())
+            alpha_crop = cv2.resize(alpha, self.output_size[::-1], interpolation=random_interp())
+            trimap_crop = cv2.resize(trimap, self.output_size[::-1], interpolation=cv2.INTER_NEAREST)
+            bg_crop = cv2.resize(bg, self.output_size[::-1], interpolation=random_interp())
+        
+        sample.update({'fg': fg_crop, 'alpha': alpha_crop, 'trimap': trimap_crop, 'bg': bg_crop})
+        return sample
+    
+
+class Composite_Seg(object):
+    def __call__(self, sample):
+        fg, bg, alpha = sample['fg'], sample['bg'], sample['alpha']
+        fg[fg < 0 ] = 0
+        fg[fg > 255] = 255
+        image = fg
+        sample['image'] = image
+        return sample
+    
+
+class ToTensor(object):
+    """
+    Convert ndarrays in sample to Tensors with normalization.
+    """
+    def __init__(self, phase="test", real_world_aug = False):
+        # self.mean = torch.tensor([0.485, 0.456, 0.406]).view(3,1,1)
+        # self.std = torch.tensor([0.229, 0.224, 0.225]).view(3,1,1)
+        self.mean = torch.tensor([0.0, 0.0, 0.0]).view(3,1,1)
+        self.std = torch.tensor([1.0, 1.0, 1.0]).view(3,1,1)
+        self.phase = phase
+        if real_world_aug:
+            self.RWA = iaa.SomeOf((1, None), [
+                iaa.LinearContrast((0.6, 1.4)),
+                iaa.JpegCompression(compression=(0, 60)),
+                iaa.GaussianBlur(sigma=(0.0, 3.0)),
+                iaa.AdditiveGaussianNoise(scale=(0, 0.1*255))
+            ], random_order=True)
+        else:
+            self.RWA = None
+    
+    def get_box_from_alpha(self, alpha_final):
+        bi_mask = np.zeros_like(alpha_final)
+        bi_mask[alpha_final>0.5] = 1
+        #bi_mask[alpha_final<=0.5] = 0
+        fg_set = np.where(bi_mask != 0)
+        if len(fg_set[1]) == 0 or len(fg_set[0]) == 0:
+            x_min = random.randint(1, 511)
+            x_max = random.randint(1, 511) + x_min
+            y_min = random.randint(1, 511)
+            y_max = random.randint(1, 511) + y_min
+        else:
+            x_min = np.min(fg_set[1])
+            x_max = np.max(fg_set[1])
+            y_min = np.min(fg_set[0])
+            y_max = np.max(fg_set[0])
+        bbox = np.array([x_min, y_min, x_max, y_max])
+        #cv2.rectangle(image, (x_min, y_min), (x_max, y_max), (0,255,0), 2)
+        #cv2.imwrite('../outputs/test.jpg', image)
+        #cv2.imwrite('../outputs/test_gt.jpg', alpha_single)
+        return bbox
+
+    def __call__(self, sample):
+        # convert GBR images to RGB
+        image, alpha, trimap = sample['image'][:,:,::-1], sample['alpha'], sample['trimap']
+        
+        alpha[alpha < 0 ] = 0
+        alpha[alpha > 1] = 1
+        
+        bbox = self.get_box_from_alpha(alpha)
+
+        if self.phase == 'train' and self.RWA is not None and np.random.rand() < 0.5:
+            image[image > 255] = 255
+            image[image < 0] = 0
+            image = np.round(image).astype(np.uint8)
+            image = np.expand_dims(image, axis=0)
+            image = self.RWA(images=image)
+            image = image[0, ...]
+
+        # swap color axis because
+        # numpy image: H x W x C
+        # torch image: C X H X W
+        image = image.transpose((2, 0, 1)).astype(np.float32)
+        alpha = np.expand_dims(alpha.astype(np.float32), axis=0)
+        trimap[trimap < 85] = 0
+        trimap[trimap >= 170] = 2
+        trimap[trimap >= 85] = 1
+        #image = cv2.rectangle(image, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (255,0,0), 3)
+        #cv2.imwrite(os.path.join('outputs', 'img_bbox.png'), image.astype('uint8'))
+        # normalize image
+        image /= 255.
+
+        if self.phase == "train":
+            # convert GBR images to RGB
+            fg = sample['fg'][:,:,::-1].transpose((2, 0, 1)).astype(np.float32) / 255.
+            sample['fg'] = torch.from_numpy(fg).sub_(self.mean).div_(self.std)
+            bg = sample['bg'][:,:,::-1].transpose((2, 0, 1)).astype(np.float32) / 255.
+            sample['bg'] = torch.from_numpy(bg).sub_(self.mean).div_(self.std)
+            del sample['image_name']
+        
+        sample['boxes'] = torch.from_numpy(bbox).to(torch.float)[None,...]
+
+        sample['image'], sample['alpha'], sample['trimap'] = \
+            torch.from_numpy(image), torch.from_numpy(alpha), torch.from_numpy(trimap).to(torch.long)
+        sample['image'] = sample['image'].sub_(self.mean).div_(self.std)
+        sample['trimap'] = sample['trimap'][None,...].float()
+
+        return sample
+
+
+class RefMatteData(Dataset):
+    def __init__(
+        self, 
+        data_root_path,
+        num_ratio = 0.34,
+    ):
+        self.data_root_path = data_root_path
+        self.num_ratio = num_ratio
+
+        self.rim_img = [os.path.join(data_root_path, name) for name in sorted(os.listdir(data_root_path))]
+        self.rim_pha = [os.path.join(data_root_path.replace('img', 'mask'), name) for name in sorted(os.listdir(data_root_path.replace('img', 'mask')))]
+        self.rim_num = len(self.rim_pha)
+
+        self.transform_spd = transforms.Compose([
+            RandomAffine(degrees=30, scale=[0.8, 1.5], shear=10, flip=0.5),
+            GenTrimap(),
+            RandomCrop((1024, 1024)),
+            Composite_Seg(),
+            ToTensor(phase="train", real_world_aug=False)
+        ])
+
+    def __getitem__(self, idx):
+        if self.num_ratio is not None:
+            if self.num_ratio < 1.0 or idx >= self.rim_num:
+                idx = np.random.randint(0, self.rim_num)
+        alpha = cv2.imread(self.rim_pha[idx % self.rim_num], 0).astype(np.float32)/255
+        alpha_img_name = self.rim_pha[idx % self.rim_num].split('/')[-1]
+        fg_img_name = alpha_img_name[:-6] + '.jpg'
+
+        fg = cv2.imread(os.path.join(self.data_root_path, fg_img_name))
+
+        if np.random.rand() < 0.25:
+            fg = cv2.resize(fg, (1280, 1280), interpolation=random_interp())
+            alpha = cv2.resize(alpha, (1280, 1280), interpolation=random_interp())
+
+        image_name = alpha_img_name  # os.path.split(self.rim_img[idx % self.rim_num])[-1]
+        sample = {'fg': fg, 'alpha': alpha, 'bg': fg, 'image_name': image_name}
+        sample = self.transform_spd(sample)
+
+        converted_sample = {
+            'image': sample['image'],
+            'trimap': sample['trimap'] / 2.0,
+            'alpha': sample['alpha'],
+            'bbox': sample['boxes'],
+            'dataset_name': 'RefMatte',
+            'multi_fg': False,
+        }
+        return converted_sample
+
+    def __len__(self):
+        if self.num_ratio is not None:
+            return int(self.rim_num * self.num_ratio)  # 112506 * 0.34 = 38252 (COCONut_num-38251 + 1)
+        else:
+            return self.rim_num  # 112506
+
+
+    
+if __name__ == '__main__':
+    dataset = RefMatteData(
+        data_root_path = '/data/my_path_b/public_data/data/matting/RefMatte/RefMatte/train/img', 
+        num_ratio=0.34,
+    )
+    data = dataset[0]
+    '''
+    fg torch.Size([3, 1024, 1024]) tensor(-2.1179) tensor(2.6400)
+    alpha torch.Size([1, 1024, 1024]) tensor(0.) tensor(1.)
+    bg torch.Size([3, 1024, 1024]) tensor(-2.1179) tensor(2.6400)
+    trimap torch.Size([1, 1024, 1024]) 0.0 or 1.0 or 2.0
+    image torch.Size([3, 1024, 1024]) tensor(-2.1179) tensor(2.6400)
+    boxes torch.Size([1, 4]) tensor(72.) tensor(676.)  0.0~1024.0
+
+    COCONut:
+        image torch.Size([3, 1024, 1024]) tensor(0.0006) tensor(0.9991)
+        trimap torch.Size([1, 1024, 1024]) 0.0 or 0.5 or 1.0
+        alpha torch.Size([1, 1024, 1024]) tensor(0.) tensor(1.)
+        bbox torch.Size([1, 4]) tensor(0.) tensor(590.)
+        dataset_name: 'COCONut'
+    '''
+    for key, val in data.items():
+        if isinstance(val, torch.Tensor):
+            print(key, val.shape, torch.min(val), torch.max(val))
+        else:
+            print(key, val.shape)

engine/__init__.py

@@ -0,0 +1 @@
+from .mattingtrainer import MattingTrainer

engine/hooks.py

@@ -0,0 +1,52 @@
+import inspect
+import detectron2.utils.comm as comm
+from detectron2.engine import EvalHook as _EvalHook
+from detectron2.evaluation.testing import flatten_results_dict
+
+
+class EvalHook(_EvalHook):
+    def __init__(self, eval_period, eval_function):
+        super().__init__(eval_period, eval_function)
+        func_args = inspect.getfullargspec(eval_function).args
+        assert {"final_iter", "next_iter"}.issubset(set(func_args)), (
+            f"Eval function must have either 'final_iter' or 'next_iter' as an argument."
+            f"Got {func_args} instead."
+        )
+
+    def _do_eval(self, final_iter=False, next_iter=0):
+        results = self._func(final_iter=final_iter, next_iter=next_iter)
+
+        if results:
+            assert isinstance(
+                results, dict
+            ), "Eval function must return a dict. Got {} instead.".format(results)
+
+            flattened_results = flatten_results_dict(results)
+            for k, v in flattened_results.items():
+                try:
+                    v = float(v)
+                except Exception as e:
+                    raise ValueError(
+                        "[EvalHook] eval_function should return a nested dict of float. "
+                        "Got '{}: {}' instead.".format(k, v)
+                    ) from e
+            self.trainer.storage.put_scalars(**flattened_results, smoothing_hint=False)
+
+        # Evaluation may take different time among workers.
+        # A barrier make them start the next iteration together.
+        comm.synchronize()
+
+    def after_step(self):
+        next_iter = self.trainer.iter + 1
+        if self._period > 0 and next_iter % self._period == 0:
+            # do the last eval in after_train
+            if next_iter != self.trainer.max_iter:
+                self._do_eval(next_iter=next_iter)
+
+    def after_train(self):
+        # This condition is to prevent the eval from running after a failed training
+        if self.trainer.iter + 1 >= self.trainer.max_iter:
+            self._do_eval(final_iter=True)
+        # func is likely a closure that holds reference to the trainer
+        # therefore we clean it to avoid circular reference in the end
+        del self._func

engine/mattingtrainer.py

@@ -0,0 +1,171 @@
+from detectron2.engine import AMPTrainer
+import torch
+import time
+import logging
+
+logger = logging.getLogger("detectron2")
+
+import typing
+from collections import defaultdict
+import tabulate
+from torch import nn
+
+
+def parameter_count(model: nn.Module, trainable_only: bool = False) -> typing.DefaultDict[str, int]:
+    """
+    Count parameters of a model and its submodules.
+
+    Args:
+        model: a torch module
+
+    Returns:
+        dict (str-> int): the key is either a parameter name or a module name.
+        The value is the number of elements in the parameter, or in all
+        parameters of the module. The key "" corresponds to the total
+        number of parameters of the model.
+    """
+    r = defaultdict(int)
+    for name, prm in model.named_parameters():
+        if trainable_only:
+            if not prm.requires_grad:
+                continue
+        size = prm.numel()
+        name = name.split(".")
+        for k in range(0, len(name) + 1):
+            prefix = ".".join(name[:k])
+            r[prefix] += size
+    return r
+
+
+def parameter_count_table(
+    model: nn.Module, max_depth: int = 3, trainable_only: bool = False
+) -> str:
+    """
+    Format the parameter count of the model (and its submodules or parameters)
+    in a nice table. It looks like this:
+
+    ::
+
+        | name                            | #elements or shape   |
+        |:--------------------------------|:---------------------|
+        | model                           | 37.9M                |
+        |  backbone                       |  31.5M               |
+        |   backbone.fpn_lateral3         |   0.1M               |
+        |    backbone.fpn_lateral3.weight |    (256, 512, 1, 1)  |
+        |    backbone.fpn_lateral3.bias   |    (256,)            |
+        |   backbone.fpn_output3          |   0.6M               |
+        |    backbone.fpn_output3.weight  |    (256, 256, 3, 3)  |
+        |    backbone.fpn_output3.bias    |    (256,)            |
+        |   backbone.fpn_lateral4         |   0.3M               |
+        |    backbone.fpn_lateral4.weight |    (256, 1024, 1, 1) |
+        |    backbone.fpn_lateral4.bias   |    (256,)            |
+        |   backbone.fpn_output4          |   0.6M               |
+        |    backbone.fpn_output4.weight  |    (256, 256, 3, 3)  |
+        |    backbone.fpn_output4.bias    |    (256,)            |
+        |   backbone.fpn_lateral5         |   0.5M               |
+        |    backbone.fpn_lateral5.weight |    (256, 2048, 1, 1) |
+        |    backbone.fpn_lateral5.bias   |    (256,)            |
+        |   backbone.fpn_output5          |   0.6M               |
+        |    backbone.fpn_output5.weight  |    (256, 256, 3, 3)  |
+        |    backbone.fpn_output5.bias    |    (256,)            |
+        |   backbone.top_block            |   5.3M               |
+        |    backbone.top_block.p6        |    4.7M              |
+        |    backbone.top_block.p7        |    0.6M              |
+        |   backbone.bottom_up            |   23.5M              |
+        |    backbone.bottom_up.stem      |    9.4K              |
+        |    backbone.bottom_up.res2      |    0.2M              |
+        |    backbone.bottom_up.res3      |    1.2M              |
+        |    backbone.bottom_up.res4      |    7.1M              |
+        |    backbone.bottom_up.res5      |    14.9M             |
+        |    ......                       |    .....             |
+
+    Args:
+        model: a torch module
+        max_depth (int): maximum depth to recursively print submodules or
+            parameters
+
+    Returns:
+        str: the table to be printed
+    """
+    count: typing.DefaultDict[str, int] = parameter_count(model, trainable_only)
+    # pyre-fixme[24]: Generic type `tuple` expects at least 1 type parameter.
+    param_shape: typing.Dict[str, typing.Tuple] = {
+        k: tuple(v.shape) for k, v in model.named_parameters()
+    }
+
+    # pyre-fixme[24]: Generic type `tuple` expects at least 1 type parameter.
+    table: typing.List[typing.Tuple] = []
+
+    def format_size(x: int) -> str:
+        if x > 1e8:
+            return "{:.1f}G".format(x / 1e9)
+        if x > 1e5:
+            return "{:.1f}M".format(x / 1e6)
+        if x > 1e2:
+            return "{:.1f}K".format(x / 1e3)
+        return str(x)
+
+    def fill(lvl: int, prefix: str) -> None:
+        if lvl >= max_depth:
+            return
+        for name, v in count.items():
+            if name.count(".") == lvl and name.startswith(prefix):
+                indent = " " * (lvl + 1)
+                if name in param_shape:
+                    table.append((indent + name, indent + str(param_shape[name])))
+                else:
+                    table.append((indent + name, indent + format_size(v)))
+                    fill(lvl + 1, name + ".")
+
+    table.append(("model", format_size(count.pop(""))))
+    fill(0, "")
+
+    old_ws = tabulate.PRESERVE_WHITESPACE
+    tabulate.PRESERVE_WHITESPACE = True
+    tab = tabulate.tabulate(table, headers=["name", "#elements or shape"], tablefmt="pipe")
+    tabulate.PRESERVE_WHITESPACE = old_ws
+    return tab
+
+
+def cycle(iterable):
+    while True:
+        for x in iterable:
+            yield x
+
+class MattingTrainer(AMPTrainer):
+    def __init__(self, model, data_loader, optimizer, grad_scaler=None):
+        super().__init__(model, data_loader, optimizer, grad_scaler=None)
+        self.data_loader_iter = iter(cycle(self.data_loader))
+
+        # print model parameters
+        logger.info("All parameters: \n" + parameter_count_table(model))
+        logger.info("Trainable parameters: \n" + parameter_count_table(model, trainable_only=True, max_depth=8))
+
+    def run_step(self):
+        """
+        Implement the AMP training logic.
+        """
+        assert self.model.training, "[AMPTrainer] model was changed to eval mode!"
+        assert torch.cuda.is_available(), "[AMPTrainer] CUDA is required for AMP training!"
+        from torch.cuda.amp import autocast
+
+        #matting pass
+        start = time.perf_counter()        
+        data = next(self.data_loader_iter)
+        data_time = time.perf_counter() - start
+
+        with autocast():
+            loss_dict = self.model(data)
+            if isinstance(loss_dict, torch.Tensor):
+                losses = loss_dict
+                loss_dict = {"total_loss": loss_dict}
+            else:
+                losses = sum(loss_dict.values())
+
+        self.optimizer.zero_grad()
+        self.grad_scaler.scale(losses).backward()
+
+        self._write_metrics(loss_dict, data_time)
+
+        self.grad_scaler.step(self.optimizer)
+        self.grad_scaler.update()

modeling/__init__.py

@@ -0,0 +1,5 @@
+from .backbone import *
+from .criterion import *
+from .decoder import *
+from .meta_arch import *
+from .semantic_enhanced_matting import *

modeling/backbone/__init__.py

@@ -0,0 +1,2 @@
+from .backbone import *
+from .vit import *

modeling/backbone/backbone.py

@@ -0,0 +1,74 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+from abc import ABCMeta, abstractmethod
+from typing import Dict
+import torch.nn as nn
+
+from detectron2.layers import ShapeSpec
+
+__all__ = ["Backbone"]
+
+
+class Backbone(nn.Module, metaclass=ABCMeta):
+    """
+    Abstract base class for network backbones.
+    """
+
+    def __init__(self):
+        """
+        The `__init__` method of any subclass can specify its own set of arguments.
+        """
+        super().__init__()
+
+    @abstractmethod
+    def forward(self):
+        """
+        Subclasses must override this method, but adhere to the same return type.
+
+        Returns:
+            dict[str->Tensor]: mapping from feature name (e.g., "res2") to tensor
+        """
+        pass
+
+    @property
+    def size_divisibility(self) -> int:
+        """
+        Some backbones require the input height and width to be divisible by a
+        specific integer. This is typically true for encoder / decoder type networks
+        with lateral connection (e.g., FPN) for which feature maps need to match
+        dimension in the "bottom up" and "top down" paths. Set to 0 if no specific
+        input size divisibility is required.
+        """
+        return 0
+
+    @property
+    def padding_constraints(self) -> Dict[str, int]:
+        """
+        This property is a generalization of size_divisibility. Some backbones and training
+        recipes require specific padding constraints, such as enforcing divisibility by a specific
+        integer (e.g., FPN) or padding to a square (e.g., ViTDet with large-scale jitter
+        in :paper:vitdet). `padding_constraints` contains these optional items like:
+        {
+            "size_divisibility": int,
+            "square_size": int,
+            # Future options are possible
+        }
+        `size_divisibility` will read from here if presented and `square_size` indicates the
+        square padding size if `square_size` > 0.
+
+        TODO: use type of Dict[str, int] to avoid torchscipt issues. The type of padding_constraints
+        could be generalized as TypedDict (Python 3.8+) to support more types in the future.
+        """
+        return {}
+
+    def output_shape(self):
+        """
+        Returns:
+            dict[str->ShapeSpec]
+        """
+        # this is a backward-compatible default
+        return {
+            name: ShapeSpec(
+                channels=self._out_feature_channels[name], stride=self._out_feature_strides[name]
+            )
+            for name in self._out_features
+        }

modeling/backbone/utils.py

@@ -0,0 +1,186 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+__all__ = [
+    "window_partition",
+    "window_unpartition",
+    "add_decomposed_rel_pos",
+    "get_abs_pos",
+    "PatchEmbed",
+]
+
+
+def window_partition(x, window_size):
+    """
+    Partition into non-overlapping windows with padding if needed.
+    Args:
+        x (tensor): input tokens with [B, H, W, C].
+        window_size (int): window size.
+
+    Returns:
+        windows: windows after partition with [B * num_windows, window_size, window_size, C].
+        (Hp, Wp): padded height and width before partition
+    """
+    B, H, W, C = x.shape
+
+    pad_h = (window_size - H % window_size) % window_size
+    pad_w = (window_size - W % window_size) % window_size
+    if pad_h > 0 or pad_w > 0:
+        x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h))
+    Hp, Wp = H + pad_h, W + pad_w
+
+    x = x.view(B, Hp // window_size, window_size, Wp // window_size, window_size, C)
+    windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
+    return windows, (Hp, Wp)
+
+
+def window_unpartition(windows, window_size, pad_hw, hw):
+    """
+    Window unpartition into original sequences and removing padding.
+    Args:
+        x (tensor): input tokens with [B * num_windows, window_size, window_size, C].
+        window_size (int): window size.
+        pad_hw (Tuple): padded height and width (Hp, Wp).
+        hw (Tuple): original height and width (H, W) before padding.
+
+    Returns:
+        x: unpartitioned sequences with [B, H, W, C].
+    """
+    Hp, Wp = pad_hw
+    H, W = hw
+    B = windows.shape[0] // (Hp * Wp // window_size // window_size)
+    x = windows.view(B, Hp // window_size, Wp // window_size, window_size, window_size, -1)
+    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, Hp, Wp, -1)
+
+    if Hp > H or Wp > W:
+        x = x[:, :H, :W, :].contiguous()
+    return x
+
+
+def get_rel_pos(q_size, k_size, rel_pos):
+    """
+    Get relative positional embeddings according to the relative positions of
+        query and key sizes.
+    Args:
+        q_size (int): size of query q.
+        k_size (int): size of key k.
+        rel_pos (Tensor): relative position embeddings (L, C).
+
+    Returns:
+        Extracted positional embeddings according to relative positions.
+    """
+    max_rel_dist = int(2 * max(q_size, k_size) - 1)
+    # Interpolate rel pos if needed.
+    if rel_pos.shape[0] != max_rel_dist:
+        # Interpolate rel pos.
+        rel_pos_resized = F.interpolate(
+            rel_pos.reshape(1, rel_pos.shape[0], -1).permute(0, 2, 1),
+            size=max_rel_dist,
+            mode="linear",
+        )
+        rel_pos_resized = rel_pos_resized.reshape(-1, max_rel_dist).permute(1, 0)
+    else:
+        rel_pos_resized = rel_pos
+
+    # Scale the coords with short length if shapes for q and k are different.
+    q_coords = torch.arange(q_size)[:, None] * max(k_size / q_size, 1.0)
+    k_coords = torch.arange(k_size)[None, :] * max(q_size / k_size, 1.0)
+    relative_coords = (q_coords - k_coords) + (k_size - 1) * max(q_size / k_size, 1.0)
+
+    return rel_pos_resized[relative_coords.long()]
+
+
+def add_decomposed_rel_pos(attn, q, rel_pos_h, rel_pos_w, q_size, k_size):
+    """
+    Calculate decomposed Relative Positional Embeddings from :paper:`mvitv2`.
+    https://github.com/facebookresearch/mvit/blob/19786631e330df9f3622e5402b4a419a263a2c80/mvit/models/attention.py   # noqa B950
+    Args:
+        attn (Tensor): attention map.
+        q (Tensor): query q in the attention layer with shape (B, q_h * q_w, C).
+        rel_pos_h (Tensor): relative position embeddings (Lh, C) for height axis.
+        rel_pos_w (Tensor): relative position embeddings (Lw, C) for width axis.
+        q_size (Tuple): spatial sequence size of query q with (q_h, q_w).
+        k_size (Tuple): spatial sequence size of key k with (k_h, k_w).
+
+    Returns:
+        attn (Tensor): attention map with added relative positional embeddings.
+    """
+    q_h, q_w = q_size
+    k_h, k_w = k_size
+    Rh = get_rel_pos(q_h, k_h, rel_pos_h)
+    Rw = get_rel_pos(q_w, k_w, rel_pos_w)
+
+    B, _, dim = q.shape
+    r_q = q.reshape(B, q_h, q_w, dim)
+    rel_h = torch.einsum("bhwc,hkc->bhwk", r_q, Rh)
+    rel_w = torch.einsum("bhwc,wkc->bhwk", r_q, Rw)
+
+    attn = (
+        attn.view(B, q_h, q_w, k_h, k_w) + rel_h[:, :, :, :, None] + rel_w[:, :, :, None, :]
+    ).view(B, q_h * q_w, k_h * k_w)
+
+    return attn
+
+
+def get_abs_pos(abs_pos, has_cls_token, hw):
+    """
+    Calculate absolute positional embeddings. If needed, resize embeddings and remove cls_token
+        dimension for the original embeddings.
+    Args:
+        abs_pos (Tensor): absolute positional embeddings with (1, num_position, C).
+        has_cls_token (bool): If true, has 1 embedding in abs_pos for cls token.
+        hw (Tuple): size of input image tokens.
+
+    Returns:
+        Absolute positional embeddings after processing with shape (1, H, W, C)
+    """
+    h, w = hw
+    if has_cls_token:
+        abs_pos = abs_pos[:, 1:]
+    xy_num = abs_pos.shape[1]
+    size = int(math.sqrt(xy_num))
+    assert size * size == xy_num
+
+    if size != h or size != w:
+        new_abs_pos = F.interpolate(
+            abs_pos.reshape(1, size, size, -1).permute(0, 3, 1, 2),
+            size=(h, w),
+            mode="bicubic",
+            align_corners=False,
+        )
+
+        return new_abs_pos.permute(0, 2, 3, 1)
+    else:
+        return abs_pos.reshape(1, h, w, -1)
+
+
+class PatchEmbed(nn.Module):
+    """
+    Image to Patch Embedding.
+    """
+
+    def __init__(
+        self, kernel_size=(16, 16), stride=(16, 16), padding=(0, 0), in_chans=3, embed_dim=768
+    ):
+        """
+        Args:
+            kernel_size (Tuple): kernel size of the projection layer.
+            stride (Tuple): stride of the projection layer.
+            padding (Tuple): padding size of the projection layer.
+            in_chans (int): Number of input image channels.
+            embed_dim (int):  embed_dim (int): Patch embedding dimension.
+        """
+        super().__init__()
+
+        self.proj = nn.Conv2d(
+            in_chans, embed_dim, kernel_size=kernel_size, stride=stride, padding=padding
+        )
+
+    def forward(self, x):
+        x = self.proj(x)
+        # B C H W -> B H W C
+        x = x.permute(0, 2, 3, 1)
+        return x

modeling/backbone/vit.py

@@ -0,0 +1,404 @@
+import logging
+import math
+import fvcore.nn.weight_init as weight_init
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from detectron2.layers import CNNBlockBase, Conv2d, get_norm
+from detectron2.modeling.backbone.fpn import _assert_strides_are_log2_contiguous
+from fairscale.nn.checkpoint import checkpoint_wrapper
+from timm.models.layers import DropPath, Mlp, trunc_normal_
+from .backbone import Backbone
+from .utils import (
+    PatchEmbed,
+    add_decomposed_rel_pos,
+    get_abs_pos,
+    window_partition,
+    window_unpartition,
+)
+
+logger = logging.getLogger(__name__)
+
+
+__all__ = ["ViT"]
+
+
+class Attention(nn.Module):
+    """Multi-head Attention block with relative position embeddings."""
+
+    def __init__(
+        self,
+        dim,
+        num_heads=8,
+        qkv_bias=True,
+        use_rel_pos=False,
+        rel_pos_zero_init=True,
+        input_size=None,
+    ):
+        """
+        Args:
+            dim (int): Number of input channels.
+            num_heads (int): Number of attention heads.
+            qkv_bias (bool:  If True, add a learnable bias to query, key, value.
+            rel_pos (bool): If True, add relative positional embeddings to the attention map.
+            rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.
+            input_size (int or None): Input resolution for calculating the relative positional
+                parameter size.
+        """
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim**-0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.proj = nn.Linear(dim, dim)
+
+        self.use_rel_pos = use_rel_pos
+        if self.use_rel_pos:
+            # initialize relative positional embeddings
+            self.rel_pos_h = nn.Parameter(torch.zeros(2 * input_size[0] - 1, head_dim))
+            self.rel_pos_w = nn.Parameter(torch.zeros(2 * input_size[1] - 1, head_dim))
+
+            if not rel_pos_zero_init:
+                trunc_normal_(self.rel_pos_h, std=0.02)
+                trunc_normal_(self.rel_pos_w, std=0.02)
+
+    def forward(self, x):
+        B, H, W, _ = x.shape
+        # qkv with shape (3, B, nHead, H * W, C)
+        qkv = self.qkv(x).reshape(B, H * W, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
+        # q, k, v with shape (B * nHead, H * W, C)
+        q, k, v = qkv.reshape(3, B * self.num_heads, H * W, -1).unbind(0)
+
+        attn = (q * self.scale) @ k.transpose(-2, -1)
+
+        if self.use_rel_pos:
+            attn = add_decomposed_rel_pos(attn, q, self.rel_pos_h, self.rel_pos_w, (H, W), (H, W))
+
+        attn = attn.softmax(dim=-1)
+        x = (attn @ v).view(B, self.num_heads, H, W, -1).permute(0, 2, 3, 1, 4).reshape(B, H, W, -1)
+        x = self.proj(x)
+
+        return x
+
+class LayerNorm(nn.Module):
+    r""" LayerNorm that supports two data formats: channels_last (default) or channels_first. 
+    The ordering of the dimensions in the inputs. channels_last corresponds to inputs with 
+    shape (batch_size, height, width, channels) while channels_first corresponds to inputs 
+    with shape (batch_size, channels, height, width).
+    """
+    def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(normalized_shape))
+        self.bias = nn.Parameter(torch.zeros(normalized_shape))
+        self.eps = eps
+        self.data_format = data_format
+        if self.data_format not in ["channels_last", "channels_first"]:
+            raise NotImplementedError 
+        self.normalized_shape = (normalized_shape, )
+    
+    def forward(self, x):
+        if self.data_format == "channels_last":
+            return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        elif self.data_format == "channels_first":
+            u = x.mean(1, keepdim=True)
+            s = (x - u).pow(2).mean(1, keepdim=True)
+            x = (x - u) / torch.sqrt(s + self.eps)
+            x = self.weight[:, None, None] * x + self.bias[:, None, None]
+            return x
+
+class ResBottleneckBlock(CNNBlockBase):
+    """
+    The standard bottleneck residual block without the last activation layer.
+    It contains 3 conv layers with kernels 1x1, 3x3, 1x1.
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        bottleneck_channels,
+        norm="LN",
+        act_layer=nn.GELU,
+        conv_kernels=3,
+        conv_paddings=1,
+    ):
+        """
+        Args:
+            in_channels (int): Number of input channels.
+            out_channels (int): Number of output channels.
+            bottleneck_channels (int): number of output channels for the 3x3
+                "bottleneck" conv layers.
+            norm (str or callable): normalization for all conv layers.
+                See :func:`layers.get_norm` for supported format.
+            act_layer (callable): activation for all conv layers.
+        """
+        super().__init__(in_channels, out_channels, 1)
+
+        self.conv1 = Conv2d(in_channels, bottleneck_channels, 1, bias=False)
+        self.norm1 = get_norm(norm, bottleneck_channels)
+        self.act1 = act_layer()
+
+        self.conv2 = Conv2d(
+            bottleneck_channels,
+            bottleneck_channels,
+            conv_kernels,
+            padding=conv_paddings,
+            bias=False,
+        )
+        self.norm2 = get_norm(norm, bottleneck_channels)
+        self.act2 = act_layer()
+
+        self.conv3 = Conv2d(bottleneck_channels, out_channels, 1, bias=False)
+        self.norm3 = get_norm(norm, out_channels)
+
+        for layer in [self.conv1, self.conv2, self.conv3]:
+            weight_init.c2_msra_fill(layer)
+        for layer in [self.norm1, self.norm2]:
+            layer.weight.data.fill_(1.0)
+            layer.bias.data.zero_()
+        # zero init last norm layer.
+        self.norm3.weight.data.zero_()
+        self.norm3.bias.data.zero_()
+
+    def forward(self, x):
+        out = x
+        for layer in self.children():
+            out = layer(out)
+
+        out = x + out
+        return out
+
+
+class Block(nn.Module):
+    """Transformer blocks with support of window attention and residual propagation blocks"""
+
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        mlp_ratio=4.0,
+        qkv_bias=True,
+        drop_path=0.0,
+        norm_layer=nn.LayerNorm,
+        act_layer=nn.GELU,
+        use_rel_pos=False,
+        rel_pos_zero_init=True,
+        window_size=0,
+        use_cc_attn = False,
+        use_residual_block=False,
+        use_convnext_block=False,
+        input_size=None,
+        res_conv_kernel_size=3,
+        res_conv_padding=1,
+    ):
+        """
+        Args:
+            dim (int): Number of input channels.
+            num_heads (int): Number of attention heads in each ViT block.
+            mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+            qkv_bias (bool): If True, add a learnable bias to query, key, value.
+            drop_path (float): Stochastic depth rate.
+            norm_layer (nn.Module): Normalization layer.
+            act_layer (nn.Module): Activation layer.
+            use_rel_pos (bool): If True, add relative positional embeddings to the attention map.
+            rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.
+            window_size (int): Window size for window attention blocks. If it equals 0, then not
+                use window attention.
+            use_residual_block (bool): If True, use a residual block after the MLP block.
+            input_size (int or None): Input resolution for calculating the relative positional
+                parameter size.
+        """
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            use_rel_pos=use_rel_pos,
+            rel_pos_zero_init=rel_pos_zero_init,
+            input_size=input_size if window_size == 0 else (window_size, window_size),
+        )
+
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        self.mlp = Mlp(in_features=dim, hidden_features=int(dim * mlp_ratio), act_layer=act_layer)
+
+        self.window_size = window_size
+
+        self.use_residual_block = use_residual_block
+        if use_residual_block:
+            # Use a residual block with bottleneck channel as dim // 2
+            self.residual = ResBottleneckBlock(
+                in_channels=dim,
+                out_channels=dim,
+                bottleneck_channels=dim // 2,
+                norm="LN",
+                act_layer=act_layer,
+                conv_kernels=res_conv_kernel_size,
+                conv_paddings=res_conv_padding,
+            )
+        self.use_convnext_block = use_convnext_block
+        if use_convnext_block:
+            self.convnext = ConvNextBlock(dim = dim)
+
+        if use_cc_attn:
+            self.attn = CrissCrossAttention(dim)
+
+
+    def forward(self, x):
+        shortcut = x
+        x = self.norm1(x)
+        # Window partition
+        if self.window_size > 0:
+            H, W = x.shape[1], x.shape[2]
+            x, pad_hw = window_partition(x, self.window_size)
+
+        x = self.attn(x)
+
+        # Reverse window partition
+        if self.window_size > 0:
+            x = window_unpartition(x, self.window_size, pad_hw, (H, W))
+
+        x = shortcut + self.drop_path(x)
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+
+        if self.use_residual_block:
+            x = self.residual(x.permute(0, 3, 1, 2)).permute(0, 2, 3, 1)
+        if self.use_convnext_block:
+            x = self.convnext(x.permute(0, 3, 1, 2)).permute(0, 2, 3, 1)
+
+        return x
+
+
+class ViT(Backbone):
+    """
+    This module implements Vision Transformer (ViT) backbone in :paper:`vitdet`.
+    "Exploring Plain Vision Transformer Backbones for Object Detection",
+    https://arxiv.org/abs/2203.16527
+    """
+
+    def __init__(
+        self,
+        img_size=1024,
+        patch_size=16,
+        in_chans=3,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4.0,
+        qkv_bias=True,
+        drop_path_rate=0.0,
+        norm_layer=nn.LayerNorm,
+        act_layer=nn.GELU,
+        use_abs_pos=True,
+        use_rel_pos=False,
+        rel_pos_zero_init=True,
+        window_size=0,
+        window_block_indexes=(),
+        residual_block_indexes=(),
+        use_act_checkpoint=False,
+        pretrain_img_size=224,
+        pretrain_use_cls_token=True,
+        out_feature="last_feat",
+        res_conv_kernel_size=3, 
+        res_conv_padding=1,
+    ):
+        """
+        Args:
+            img_size (int): Input image size.
+            patch_size (int): Patch size.
+            in_chans (int): Number of input image channels.
+            embed_dim (int): Patch embedding dimension.
+            depth (int): Depth of ViT.
+            num_heads (int): Number of attention heads in each ViT block.
+            mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+            qkv_bias (bool): If True, add a learnable bias to query, key, value.
+            drop_path_rate (float): Stochastic depth rate.
+            norm_layer (nn.Module): Normalization layer.
+            act_layer (nn.Module): Activation layer.
+            use_abs_pos (bool): If True, use absolute positional embeddings.
+            use_rel_pos (bool): If True, add relative positional embeddings to the attention map.
+            rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.
+            window_size (int): Window size for window attention blocks.
+            window_block_indexes (list): Indexes for blocks using window attention.
+            residual_block_indexes (list): Indexes for blocks using conv propagation.
+            use_act_checkpoint (bool): If True, use activation checkpointing.
+            pretrain_img_size (int): input image size for pretraining models.
+            pretrain_use_cls_token (bool): If True, pretrainig models use class token.
+            out_feature (str): name of the feature from the last block.
+        """
+        super().__init__()
+        self.pretrain_use_cls_token = pretrain_use_cls_token
+
+        self.patch_embed = PatchEmbed(
+            kernel_size=(patch_size, patch_size),
+            stride=(patch_size, patch_size),
+            in_chans=in_chans,
+            embed_dim=embed_dim,
+        )
+
+        if use_abs_pos:
+            # Initialize absolute positional embedding with pretrain image size.
+            num_patches = (pretrain_img_size // patch_size) * (pretrain_img_size // patch_size)
+            num_positions = (num_patches + 1) if pretrain_use_cls_token else num_patches
+            self.pos_embed = nn.Parameter(torch.zeros(1, num_positions, embed_dim))
+        else:
+            self.pos_embed = None
+
+        # stochastic depth decay rule
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]
+
+        self.blocks = nn.ModuleList()
+        for i in range(depth):
+            block = Block(
+                dim=embed_dim,
+                num_heads=num_heads,
+                mlp_ratio=mlp_ratio,
+                qkv_bias=qkv_bias,
+                drop_path=dpr[i],
+                norm_layer=norm_layer,
+                act_layer=act_layer,
+                use_rel_pos=use_rel_pos,
+                rel_pos_zero_init=rel_pos_zero_init,
+                window_size=window_size if i in window_block_indexes else 0,
+                use_residual_block=i in residual_block_indexes,
+                input_size=(img_size // patch_size, img_size // patch_size),
+                res_conv_kernel_size=res_conv_kernel_size,
+                res_conv_padding=res_conv_padding,
+            )
+            if use_act_checkpoint:
+                block = checkpoint_wrapper(block)
+            self.blocks.append(block)
+
+        self._out_feature_channels = {out_feature: embed_dim}
+        self._out_feature_strides = {out_feature: patch_size}
+        self._out_features = [out_feature]
+
+        if self.pos_embed is not None:
+            trunc_normal_(self.pos_embed, std=0.02)
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=0.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def forward(self, x):
+        x = self.patch_embed(x)
+        if self.pos_embed is not None:
+            x = x + get_abs_pos(
+                self.pos_embed, self.pretrain_use_cls_token, (x.shape[1], x.shape[2])
+            )
+
+        for blk in self.blocks:
+            x = blk(x)
+
+        outputs = {self._out_features[0]: x.permute(0, 3, 1, 2)}
+
+        return outputs['last_feat']

modeling/criterion/__init__.py

@@ -0,0 +1 @@
+from .matting_criterion import MattingCriterion

modeling/criterion/matting_criterion.py

@@ -0,0 +1,271 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from collections import defaultdict
+
+
+class MattingCriterion(nn.Module):
+    def __init__(
+        self,
+        *,
+        losses,
+        image_size = 1024,
+    ):
+        super(MattingCriterion, self).__init__()
+        self.losses = losses
+        self.image_size = image_size
+
+    def loss_gradient_penalty(self, sample_map, preds, targets):
+
+        #sample_map for unknown area
+        if torch.sum(sample_map) == 0:
+            scale = 0
+        else:
+            scale = sample_map.shape[0] * (self.image_size ** 2) / torch.sum(sample_map)
+
+        #gradient in x
+        sobel_x_kernel = torch.tensor([[[[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]]]]).type(dtype=preds.type())
+        delta_pred_x = F.conv2d(preds, weight=sobel_x_kernel, padding=1)
+        delta_gt_x = F.conv2d(targets, weight=sobel_x_kernel, padding=1)
+
+        #gradient in y 
+        sobel_y_kernel = torch.tensor([[[[-1, -2, -1], [0, 0, 0], [1, 2, 1]]]]).type(dtype=preds.type())
+        delta_pred_y = F.conv2d(preds, weight=sobel_y_kernel, padding=1)
+        delta_gt_y = F.conv2d(targets, weight=sobel_y_kernel, padding=1)
+
+        #loss
+        loss = (F.l1_loss(delta_pred_x * sample_map, delta_gt_x * sample_map) * scale + \
+            F.l1_loss(delta_pred_y * sample_map, delta_gt_y * sample_map) * scale + \
+            0.01 * torch.mean(torch.abs(delta_pred_x * sample_map)) * scale +  \
+            0.01 * torch.mean(torch.abs(delta_pred_y * sample_map)) * scale)
+
+        return dict(loss_gradient_penalty=loss)
+
+    def loss_pha_laplacian(self, preds, targets):
+        loss = laplacian_loss(preds, targets)
+        return dict(loss_pha_laplacian=loss)
+
+    def unknown_l1_loss(self, sample_map, preds, targets):
+        
+        if torch.sum(sample_map) == 0:
+            scale = 0
+        else:
+            scale = sample_map.shape[0] * (self.image_size ** 2) / torch.sum(sample_map)
+        # scale = 1
+
+        loss = F.l1_loss(preds * sample_map, targets * sample_map) * scale
+
+        return dict(unknown_l1_loss=loss)
+
+    def known_l1_loss(self, sample_map, preds, targets):
+        new_sample_map = torch.zeros_like(sample_map)
+        new_sample_map[sample_map==0] = 1
+        
+        if torch.sum(new_sample_map) == 0:
+            scale = 0
+        else:
+            scale = new_sample_map.shape[0] * (self.image_size ** 2) / torch.sum(new_sample_map)
+        # scale = 1
+        
+        loss = F.l1_loss(preds * new_sample_map, targets * new_sample_map) * scale
+
+        return dict(known_l1_loss=loss)
+
+    def get_loss(self, k, sample_map, preds, targets):
+        if k=='unknown_l1_loss' or k=='known_l1_loss' or k=='loss_gradient_penalty':
+            losses = getattr(self, k)(sample_map, preds, targets)
+        else:
+            losses = getattr(self, k)(preds, targets)
+        assert len(list(losses.keys())) == 1
+        return losses[list(losses.keys())[0]]
+
+    def forward(self, sample_map, preds, targets, batch_weight=None):
+        losses = {i: torch.tensor(0.0, device=sample_map.device) for i in self.losses}
+        for k in self.losses:
+            if batch_weight is None:
+                losses[k] += self.get_loss(k, sample_map, preds, targets)
+            else:
+                for i, loss_weight in enumerate(batch_weight):
+                    if loss_weight == -1.0 and k != 'known_l1_loss':
+                        continue
+                    else:
+                        losses[k] += self.get_loss(k, sample_map[i: i + 1], preds[i: i + 1], targets[i: i + 1]) * abs(loss_weight)
+        return losses
+
+
+#-----------------Laplacian Loss-------------------------#
+def laplacian_loss(pred, true, max_levels=5):
+    kernel = gauss_kernel(device=pred.device, dtype=pred.dtype)
+    pred_pyramid = laplacian_pyramid(pred, kernel, max_levels)
+    true_pyramid = laplacian_pyramid(true, kernel, max_levels)
+    loss = 0
+    for level in range(max_levels):
+        loss += (2 ** level) * F.l1_loss(pred_pyramid[level], true_pyramid[level])
+    return loss / max_levels
+
+def laplacian_pyramid(img, kernel, max_levels):
+    current = img
+    pyramid = []
+    for _ in range(max_levels):
+        current = crop_to_even_size(current)
+        down = downsample(current, kernel)
+        up = upsample(down, kernel)
+        diff = current - up
+        pyramid.append(diff)
+        current = down
+    return pyramid
+
+def gauss_kernel(device='cpu', dtype=torch.float32):
+    kernel = torch.tensor([[1,  4,  6,  4, 1],
+                        [4, 16, 24, 16, 4],
+                        [6, 24, 36, 24, 6],
+                        [4, 16, 24, 16, 4],
+                        [1,  4,  6,  4, 1]], device=device, dtype=dtype)
+    kernel /= 256
+    kernel = kernel[None, None, :, :]
+    return kernel
+
+def gauss_convolution(img, kernel):
+    B, C, H, W = img.shape
+    img = img.reshape(B * C, 1, H, W)
+    img = F.pad(img, (2, 2, 2, 2), mode='reflect')
+    img = F.conv2d(img, kernel)
+    img = img.reshape(B, C, H, W)
+    return img
+
+def downsample(img, kernel):
+    img = gauss_convolution(img, kernel)
+    img = img[:, :, ::2, ::2]
+    return img
+
+def upsample(img, kernel):
+    B, C, H, W = img.shape
+    out = torch.zeros((B, C, H * 2, W * 2), device=img.device, dtype=img.dtype)
+    out[:, :, ::2, ::2] = img * 4
+    out = gauss_convolution(out, kernel)
+    return out
+
+def crop_to_even_size(img):
+    H, W = img.shape[2:]
+    H = H - H % 2
+    W = W - W % 2
+    return img[:, :, :H, :W]
+
+def normalized_focal_loss(pred, gt, gamma=2, class_num=3, norm=True, beta_detach=False, beta_sum_detach=False):
+    pred_logits = F.softmax(pred, dim=1)  # [B, 3, H, W]
+    gt_one_hot = F.one_hot(gt, class_num).permute(0, 3, 1, 2)  # [B, 3, H, W]
+    p = (pred_logits * gt_one_hot).sum(dim=1)  # [B, H, W]
+    beta = (1 - p) ** gamma  # [B, H, W]
+    beta_sum = torch.sum(beta, dim=(-2, -1), keepdim=True) / (pred.shape[-1] * pred.shape[-2])  # [B, 1, 1]
+
+    if beta_detach:
+        beta = beta.detach()
+    if beta_sum_detach:
+        beta_sum = beta_sum.detach()
+
+    if norm:
+        loss = 1 / beta_sum * beta * (-torch.log(p))
+        return torch.mean(loss)
+    else:
+        loss = beta * (-torch.log(p))
+        return torch.mean(loss)
+
+class GHMC(nn.Module):
+    def __init__(self, bins=10, momentum=0.75, loss_weight=1.0, device='cuda', norm=False):
+        super(GHMC, self).__init__()
+        self.bins = bins
+        self.momentum = momentum
+        self.edges = torch.arange(bins + 1).float().cuda() / bins
+        self.edges[-1] += 1e-6
+        if momentum > 0:
+            self.acc_sum = torch.zeros(bins).cuda()
+        self.loss_weight = loss_weight
+        self.device = device
+        self.norm = norm
+
+    def forward(self, pred, target, *args, **kwargs):
+        """Calculate the GHM-C loss.
+        Args:
+            pred (float tensor of size [batch_num, class_num]):
+                The direct prediction of classification fc layer.
+            target (float tensor of size [batch_num, class_num]):
+                Binary class target for each sample.
+            label_weight (float tensor of size [batch_num, class_num]):
+                the value is 1 if the sample is valid and 0 if ignored.
+        Returns:
+            The gradient harmonized loss.
+        """
+
+        # the target should be binary class label
+        # if pred.dim() != target.dim():
+        #     target, label_weight = _expand_binary_labels(
+        #                             target, label_weight, pred.size(-1))
+        # target, label_weight = target.float(), label_weight.float()
+        # pdb.set_trace()
+
+        # pred: [B, C, H, W], target: [B, H, W]
+        pred = pred.permute(0, 2, 3, 1).reshape(-1, 3)  # [B x H x W, C]
+        target = target.reshape(-1)  # [B x H x W]
+        # self.acc_sum = self.acc_sum.type(pred.dtype)
+
+        edges = self.edges
+        mmt = self.momentum
+        weights = torch.zeros((target.shape),dtype=pred.dtype).to(self.device)
+
+        # gradient length
+        #g = 1 - torch.index_select(F.softmax(pred,dim=1).detach(), dim=0, index=target)
+        g = 1 - torch.gather(F.softmax(pred,dim=1).detach(),dim=1,index=target.unsqueeze(1))
+        #g = torch.abs(pred.softmax(2).detach() - target)
+
+        tot = 1.0
+        n = 0  # n valid bins
+        for i in range(self.bins):
+            inds = (g >= edges[i]) & (g < edges[i+1])
+            num_in_bin = inds.sum().item()
+            if num_in_bin > 0:
+                idx = torch.nonzero(inds)[:, 0]
+                if mmt > 0:
+                    self.acc_sum[i] = mmt * self.acc_sum[i] \
+                        + (1 - mmt) * num_in_bin
+                    # pdb.set_trace()#scatter_ index_put_
+                    #BB=torch.nonzero(inds)
+                    _weight_idx = tot / self.acc_sum[i]
+                    weights = weights.to(dtype=_weight_idx.dtype)
+                    weights[idx] = _weight_idx
+                    # weights.scatter_(0, torch.nonzero(inds)[:,0], tot / self.acc_sum[i])
+                    # # weights.index_put_(inds, tot / self.acc_sum[i])
+                    # weights[inds] = tot / self.acc_sum[i] # * torch.ones((len(inds)))
+                else:
+                    weights[idx] = tot / num_in_bin
+                n += 1
+        if n > 0:
+            weights = weights / n
+
+            # pdb.set_trace()
+            # loss = (weights * F.cross_entropy(pred, target, reduction='none')).sum() / tot / pred.shape[0]
+        if self.norm:
+            weights = weights / torch.sum(weights).detach()
+
+        loss = - ((weights.unsqueeze(1) * torch.gather(F.log_softmax(pred, dim=1), dim=1, index=target.unsqueeze(1))).sum() )  # / pred.shape[0]
+
+        # loss3= F.cross_entropy(pred, target, reduction='mean')
+        # loss4 = - ((torch.gather(F.log_softmax(pred, dim=1), dim=1, index=target.unsqueeze(1))).sum() / pred.shape[0])
+
+        # pro = F.softmax(logits, dim=1)
+        #
+        # label_onehot = torch.zeros_like(logits).scatter_(1, labels.unsqueeze(1), 1)
+        # with torch.no_grad():
+        #     weight_matrix = (1 - pro) ** self.gamma
+        # # pdb.set_trace()
+        # fl = - (weight_matrix * (label_onehot * (pro + self.eps).log())).sum() / pro.shape[0]
+
+        return loss
+
+if __name__ == '__main__':
+    pred = torch.randn(2, 3, 1024, 1024)
+    gt =torch.argmax(torch.randn(2, 3, 1024, 1024), dim=1)
+    loss = normalized_focal_loss(pred, gt)
+    print(loss)
+    
+
+

modeling/decoder/__init__.py

@@ -0,0 +1 @@
+from .detail_capture import Detail_Capture, Ori_Detail_Capture

modeling/decoder/detail_capture.py

@@ -0,0 +1,185 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+class Basic_Conv3x3(nn.Module):
+    """
+    Basic convolution layers including: Conv3x3, BatchNorm2d, ReLU layers.
+    """
+    def __init__(
+        self,
+        in_chans,
+        out_chans,
+        stride=2,
+        padding=1,
+    ):
+        super().__init__()
+        self.conv = nn.Conv2d(in_chans, out_chans, 3, stride, padding, bias=False)
+        self.bn = nn.BatchNorm2d(out_chans)
+        self.relu = nn.ReLU(True)
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        x = self.relu(x)
+
+        return x
+
+class ConvStream(nn.Module):
+    """
+    Simple ConvStream containing a series of basic conv3x3 layers to extract detail features.
+    """
+    def __init__(
+        self,
+        in_chans = 4,
+        out_chans = [48, 96, 192],
+    ):
+        super().__init__()
+        self.convs = nn.ModuleList()
+        
+        self.conv_chans = out_chans.copy()
+        self.conv_chans.insert(0, in_chans)
+        
+        for i in range(len(self.conv_chans)-1):
+            in_chan_ = self.conv_chans[i]
+            out_chan_ = self.conv_chans[i+1]
+            self.convs.append(
+                Basic_Conv3x3(in_chan_, out_chan_)
+            )
+    
+    def forward(self, x):
+        out_dict = {'D0': x}
+        for i in range(len(self.convs)):
+            x = self.convs[i](x)
+            name_ = 'D'+str(i+1)
+            out_dict[name_] = x
+        
+        return out_dict
+
+class Fusion_Block(nn.Module):
+    """
+    Simple fusion block to fuse feature from ConvStream and Plain Vision Transformer.
+    """
+    def __init__(
+        self,
+        in_chans,
+        out_chans,
+    ):
+        super().__init__()
+        self.conv = Basic_Conv3x3(in_chans, out_chans, stride=1, padding=1)
+
+    def forward(self, x, D):
+        F_up = F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=False)
+        out = torch.cat([D, F_up], dim=1)
+        out = self.conv(out)
+
+        return out    
+
+class Matting_Head(nn.Module):
+    """
+    Simple Matting Head, containing only conv3x3 and conv1x1 layers.
+    """
+    def __init__(
+        self,
+        in_chans = 32,
+        mid_chans = 16,
+    ):
+        super().__init__()
+        self.matting_convs = nn.Sequential(
+            nn.Conv2d(in_chans, mid_chans, 3, 1, 1),
+            nn.BatchNorm2d(mid_chans),
+            nn.ReLU(True),
+            nn.Conv2d(mid_chans, 1, 1, 1, 0)
+            )
+
+    def forward(self, x):
+        x = self.matting_convs(x)
+
+        return x
+
+class Detail_Capture(nn.Module):
+    """
+    Simple and Lightweight Detail Capture Module for ViT Matting.
+    """
+    def __init__(
+        self,
+        in_chans = [384, 1],
+        img_chans=4,
+        convstream_out = [48, 96, 192],
+        fusion_out = [256, 128, 64, 32],
+    ):
+        super().__init__()
+        assert len(fusion_out) == len(convstream_out) + 1
+
+        self.convstream = ConvStream(in_chans=img_chans, out_chans=convstream_out)
+        self.conv_chans = self.convstream.conv_chans  # [4, 48, 96, 192]
+
+        self.fusion_blks = nn.ModuleList()
+        self.fus_channs = fusion_out.copy()
+        self.fus_channs.insert(0, in_chans[0])  # [384, 256, 128, 64, 32]
+        for i in range(len(self.fus_channs)-1):
+            in_channels = self.fus_channs[i] + self.conv_chans[-(i+1)] if i != 2 else in_chans[1] + self.conv_chans[-(i+1)]  # [256 + 192 = 448, 256 + 96 = 352, 128 + 48 = 176, 64 + 4 = 68]
+            out_channels = self.fus_channs[i+1]  # [256, 128, 64, 32]
+            self.fusion_blks.append(
+                Fusion_Block(
+                    in_chans = in_channels,
+                    out_chans = out_channels,
+                )
+            )
+
+        self.matting_head = Matting_Head(  # 32 --> 1
+            in_chans = fusion_out[-1],  
+        )
+
+    def forward(self, features, images):
+        detail_features = self.convstream(images)  # [1, 4, 672, 992] --> D0: [1, 4, 672, 992], D1: [1, 48, 336, 496], D2: [1, 96, 168, 248], D3: [1, 192, 84, 124]
+        for i in range(len(self.fusion_blks)):  # D3 
+            d_name_ = 'D'+str(len(self.fusion_blks)-i-1)
+            features = self.fusion_blks[i](features, detail_features[d_name_])
+        
+        phas = torch.sigmoid(self.matting_head(features))
+
+        return {'phas': phas}
+
+
+class Ori_Detail_Capture(nn.Module):
+    """
+    Simple and Lightweight Detail Capture Module for ViT Matting.
+    """
+    def __init__(
+        self,
+        in_chans = 384,
+        img_chans=4,
+        convstream_out = [48, 96, 192],
+        fusion_out = [256, 128, 64, 32],
+    ):
+        super().__init__()
+        assert len(fusion_out) == len(convstream_out) + 1
+
+        self.convstream = ConvStream(in_chans = img_chans)
+        self.conv_chans = self.convstream.conv_chans
+
+        self.fusion_blks = nn.ModuleList()
+        self.fus_channs = fusion_out.copy()
+        self.fus_channs.insert(0, in_chans)
+        for i in range(len(self.fus_channs)-1):
+            self.fusion_blks.append(
+                Fusion_Block(
+                    in_chans = self.fus_channs[i] + self.conv_chans[-(i+1)],
+                    out_chans = self.fus_channs[i+1],
+                )
+            )
+
+        self.matting_head = Matting_Head(
+            in_chans = fusion_out[-1],
+        )
+
+    def forward(self, features, images):
+        detail_features = self.convstream(images)
+        for i in range(len(self.fusion_blks)):
+            d_name_ = 'D'+str(len(self.fusion_blks)-i-1)
+            features = self.fusion_blks[i](features, detail_features[d_name_])
+        
+        phas = torch.sigmoid(self.matting_head(features))
+
+        return {'phas': phas}

modeling/decoder/unet_detail_capture.py

@@ -0,0 +1,429 @@
+import cv2
+import torch
+from torch import nn
+from torch.nn import functional as F
+# from nnMorpho.binary_operators import erosion
+from detectron2.layers.batch_norm import NaiveSyncBatchNorm
+
+
+class GenTrimapTorch(object):
+    def __init__(self, max_kernal=200):
+        self.max_kernal = max_kernal
+        self.erosion_kernels = [None] + [torch.from_numpy(cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (size, size))).float().cuda() for size in range(1, self.max_kernal)]
+
+    def __call__(self, mask, kernel_size):
+        
+        fg_width = kernel_size
+        bg_width = kernel_size
+
+        fg_mask = mask
+        bg_mask = 1 - mask
+
+        fg_mask = erosion(fg_mask, self.erosion_kernels[fg_width], border='a')
+        bg_mask = erosion(bg_mask, self.erosion_kernels[bg_width], border='a')
+
+        trimap = torch.ones_like(mask) * 0.5
+        trimap[fg_mask == 1] = 1.0
+        trimap[bg_mask == 1] = 0.0
+
+        return trimap
+
+
+class LayerNorm2d(nn.Module):
+    def __init__(self, num_channels: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(num_channels))
+        self.bias = nn.Parameter(torch.zeros(num_channels))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        u = x.mean(1, keepdim=True)
+        s = (x - u).pow(2).mean(1, keepdim=True)
+        x = (x - u) / torch.sqrt(s + self.eps)
+        x = self.weight[:, None, None] * x + self.bias[:, None, None]
+        return x
+
+
+class BasicDownBlock(nn.Module):
+    def __init__(self, in_channel, out_channel, res = True, norm=LayerNorm2d, block_num=1, kernel_size=3):
+        super().__init__()
+
+        self.res = res
+        self.basic_layer = nn.ModuleList()
+        for i in range(block_num):
+            if i == 0:
+                basic_layer_in_ch = in_channel
+                stride = 2
+            else:
+                basic_layer_in_ch = out_channel
+                stride = 1
+                self.basic_layer.append(nn.GELU())
+            self.basic_layer.append(nn.Sequential(
+                nn.Conv2d(basic_layer_in_ch, out_channel, kernel_size, stride, kernel_size // 2), 
+                norm(out_channel),
+                nn.GELU(),
+                nn.Conv2d(out_channel, out_channel, kernel_size, 1, kernel_size // 2), 
+                norm(out_channel),
+            ))
+        self.act = nn.GELU()
+
+        if self.res:
+            self.res_layer = nn.Conv2d(in_channel, out_channel, kernel_size, 2, kernel_size // 2)
+
+    def forward(self, x):
+
+        if self.res:
+            identity = self.res_layer(x)
+        else:
+            identity = F.interpolate(x, size=(out.shape[-2], out.shape[-1]), mode='bilinear', align_corners=False)
+
+        out = x
+        for layer in self.basic_layer:
+            out = layer(out)
+        
+        out = out + identity
+        out = self.act(out)
+
+        return out
+
+
+class BasicUpBlock(nn.Module):
+
+    def __init__( self, in_channel, out_channel, res = True, skip_connect = 'concat', norm=LayerNorm2d, block_num=1, kernel_size=3):
+        super().__init__()
+        assert skip_connect in {'sum', 'concat'}
+
+        self.res = res
+        self.skip_connect = skip_connect
+        self.basic_layer = nn.ModuleList()
+        for i in range(block_num):
+            if i == 0:
+                basic_layer_in_ch = in_channel
+                first_conv = nn.ConvTranspose2d(basic_layer_in_ch, out_channel, 2, 2)
+            else:
+                basic_layer_in_ch = out_channel
+                first_conv = nn.Conv2d(out_channel, out_channel, kernel_size, 1, kernel_size // 2)
+                self.basic_layer.append(nn.GELU())
+            self.basic_layer.append(nn.Sequential(
+                first_conv, 
+                norm(out_channel),
+                nn.GELU(),
+                nn.Conv2d(out_channel, out_channel, kernel_size, 1, kernel_size // 2), 
+                norm(out_channel),
+            ))
+        self.act = nn.GELU()
+
+        if self.res:
+            self.res_layer = nn.Conv2d(in_channel, out_channel, kernel_size, 1, kernel_size // 2)
+
+
+    def forward(self, x, skip_feat, concat_feat=None):
+
+        if self.skip_connect == 'sum':
+            x = x + skip_feat
+        else:
+            x = torch.concat((x, skip_feat), dim=1)
+
+        if concat_feat is not None:
+            x = torch.concat((x, concat_feat), dim=1)
+
+        out = x
+        for layer in self.basic_layer:
+            out = layer(out)
+        # out = self.basic_layer(x)
+        
+        identity = F.interpolate(x, size=(out.shape[-2], out.shape[-1]), mode='bilinear', align_corners=False)
+        if self.res:
+            identity = self.res_layer(identity)
+
+        out = out + identity
+        out = self.act(out)
+
+        return out
+    
+
+
+class DetailUNet(nn.Module):
+    def __init__(
+        self,
+        img_feat_in = 4,
+        vit_early_feat_in = 768,
+        matting_feat_in = 5,
+        downsample_in_out = [(4, 32), (32, 64), (64, 128), (128, 256)],
+        upsample_in_out = [(256, 128), (128, 64), (64, 32), (32, 16)],
+        matting_head_in = 16,
+        skip_connect = 'sum',
+        norm_type = 'LN',
+    ):
+        super().__init__()
+
+        assert len(downsample_in_out) == len(upsample_in_out)
+        downsample_in_out[0] = (img_feat_in, downsample_in_out[0][1])
+
+        assert norm_type in {'BN', 'LN', 'SyncBN'}
+        if norm_type == 'BN':
+            self.norm = torch.nn.BatchNorm2d
+        elif norm_type == 'SyncBN':
+            self.norm = NaiveSyncBatchNorm
+        else:
+            self.norm = LayerNorm2d
+
+        self.down_blks = nn.ModuleList()
+        for in_ch, out_ch in downsample_in_out:
+            self.down_blks.append(
+                BasicDownBlock(in_ch, out_ch, norm=self.norm)
+            )
+        
+        self.mid_layer = nn.Sequential(
+            nn.Conv2d(vit_early_feat_in, downsample_in_out[-1][1], 1, 1), 
+            self.norm(downsample_in_out[-1][1]),
+            nn.GELU(),
+        )
+
+        self.up_blks = nn.ModuleList()
+        for i, (in_ch, out_ch) in enumerate(upsample_in_out):
+            if i == 2:
+                in_ch += matting_feat_in
+            self.up_blks.append(
+                BasicUpBlock(in_ch, out_ch, skip_connect=skip_connect, norm=self.norm)
+            )
+
+        self.matting_head = nn.Conv2d(matting_head_in, 1, 3, 1, 1)
+
+
+    def forward(self, x, vit_early_feat, matting_feat, return_alpha_logits=False):
+        details = []
+        dfeatures = x
+
+        for i in range(len(self.down_blks)):
+            dfeatures = self.down_blks[i](dfeatures)
+            details.append(dfeatures)
+
+        out = self.mid_layer(vit_early_feat)
+        for i in range(len(self.up_blks)):
+            if i == 2:
+                out = self.up_blks[i](out, details[-i - 1], matting_feat)
+            else:
+                out = self.up_blks[i](out, details[-i - 1])
+        alpha = self.matting_head(out)
+        if return_alpha_logits:
+            return alpha, out
+        else:
+            return alpha
+    
+
+class MattingDetailDecoder(nn.Module):
+    def __init__(
+        self,
+        img_feat_in = 4,
+        vit_intern_feat_in = 1024,
+        vit_intern_feat_index = [0, 1, 2, 3],
+        downsample_in_out = [(4, 32), (32, 64), (64, 128), (128, 256)],
+        upsample_in_out = [(256, 128), (128, 64), (64, 32), (32, 16)],
+        matting_head_in = 16,
+        skip_connect = 'sum',
+        norm_type = 'BN',
+        norm_mask_logits = 6.5,
+        with_trimap = False,
+        min_kernel_size = 20,
+        kernel_div = 10,
+        concat_gen_trimap = False,
+        wo_hq_features = False,
+        block_num = 1,
+        wo_big_kernel = False,
+        sam2_multi_scale_feates = False,
+    ):
+        super().__init__()
+
+        assert len(downsample_in_out) == len(upsample_in_out)
+        assert skip_connect in {'sum', 'concat'}
+        downsample_in_out[0] = (img_feat_in, downsample_in_out[0][1])
+        
+        self.vit_intern_feat_in = vit_intern_feat_in
+        self.vit_intern_feat_index = vit_intern_feat_index
+        self.norm_mask_logits = norm_mask_logits
+        self.with_trimap = with_trimap
+        self.min_kernel_size = min_kernel_size
+        self.kernel_div = kernel_div
+        self.concat_gen_trimap = concat_gen_trimap
+        self.wo_hq_features = wo_hq_features
+        self.block_num = block_num
+        self.wo_big_kernel = wo_big_kernel
+        self.sam2_multi_scale_feates = sam2_multi_scale_feates
+        if self.sam2_multi_scale_feates:
+            assert downsample_in_out[0][0] == 6
+            downsample_in_out = [(4, 32), (32, 64), (64 + 32, 128), (128 + 64, 256)]
+            upsample_in_out = [(256, 128), (128, 64), (64, 32), (32, 16)]
+
+        if self.with_trimap and not self.concat_gen_trimap:
+            self.gen_trimap = GenTrimapTorch()
+        assert norm_type in {'BN', 'LN', 'SyncBN'}
+        if norm_type == 'BN':
+            self.norm = torch.nn.BatchNorm2d
+        elif norm_type == 'SyncBN':
+            self.norm = NaiveSyncBatchNorm
+        else:
+            self.norm = LayerNorm2d
+
+        if self.block_num >= 2 and not self.wo_big_kernel:
+            self.big_kernel_process = nn.Sequential(
+                nn.Conv2d(img_feat_in, 16, kernel_size=13, stride=1, padding=6), 
+                self.norm(16),
+                nn.GELU(),
+                nn.Conv2d(16, 32, kernel_size=13, stride=1, padding=6), 
+                self.norm(32),
+                nn.GELU(),
+            )
+            downsample_in_out[0] = (32, downsample_in_out[0][1])
+
+        if not self.sam2_multi_scale_feates:
+            self.vit_feat_proj = nn.ModuleDict()
+            for idx in self.vit_intern_feat_index:
+                self.vit_feat_proj[str(idx)] = nn.Conv2d(self.vit_intern_feat_in, self.vit_intern_feat_in // len(self.vit_intern_feat_index), 1, 1)
+        self.vit_feat_aggregation = nn.Sequential(
+            nn.Conv2d(self.vit_intern_feat_in // len(self.vit_intern_feat_index) * len(self.vit_intern_feat_index), downsample_in_out[-1][1], 3, 1, 1), 
+            self.norm(downsample_in_out[-1][1]),
+            nn.GELU(),
+        )
+
+        self.down_blks = nn.ModuleList()
+        for in_ch, out_ch in downsample_in_out:
+            self.down_blks.append(
+                BasicDownBlock(in_ch, out_ch, norm=self.norm, block_num=self.block_num, kernel_size=5 if self.block_num >= 2 else 3)
+            )
+        
+        if self.sam2_multi_scale_feates:
+            self.mid_layer = nn.ModuleList([
+                nn.Sequential(
+                    nn.Conv2d(32, 32, 1, 1), 
+                    self.norm(32),
+                    nn.GELU(),
+                ),
+                nn.Sequential(
+                    nn.Conv2d(64, 64, 1, 1), 
+                    self.norm(64),
+                    nn.GELU(),
+                ),
+                nn.Sequential(
+                    nn.Conv2d(256, 256, 1, 1), 
+                    self.norm(256),
+                    nn.GELU(),
+                ),
+                nn.Sequential(
+                    nn.Conv2d(512, 256, 3, 1, 1), 
+                    self.norm(256),
+                    nn.GELU(),
+                ),
+            ])
+        else:
+            self.mid_layer = nn.Sequential(
+                nn.Conv2d(downsample_in_out[-1][1] * 2, downsample_in_out[-1][1], 1, 1), 
+                self.norm(downsample_in_out[-1][1]),
+                nn.GELU(),
+            )
+
+        self.up_blks = nn.ModuleList()
+        for _, (in_ch, out_ch) in enumerate(upsample_in_out):
+            if skip_connect == 'concat':
+                self.up_blks.append(BasicUpBlock(in_ch * 2, out_ch, skip_connect=skip_connect, norm=self.norm, block_num=self.block_num))
+            else:
+                self.up_blks.append(BasicUpBlock(in_ch, out_ch, skip_connect=skip_connect, norm=self.norm, block_num=self.block_num))
+
+        self.matting_head = nn.Conv2d(matting_head_in, 1, 3, 1, 1)
+
+        if self.norm_mask_logits == 'BN':
+            self.logits_norm = self.norm(1)
+
+
+    def preprocess_inputs(self, images, hq_features, pred_trimap):
+
+        if self.wo_hq_features:
+            return images
+
+        if isinstance(self.norm_mask_logits, float):
+            norm_hq_features = hq_features / self.norm_mask_logits
+        elif self.norm_mask_logits == 'BN':
+            norm_hq_features = self.logits_norm(hq_features)
+        elif self.norm_mask_logits == 'Sigmoid':
+            if hq_features.shape[1] == 1:
+                norm_hq_features = torch.sigmoid(hq_features)
+            else:
+                norm_hq_features = torch.softmax(hq_features, dim=1)
+        elif self.norm_mask_logits:
+            norm_hq_features = hq_features / torch.std(hq_features, dim=(1, 2, 3), keepdim=True)
+        else:
+            norm_hq_features = hq_features
+
+        if self.concat_gen_trimap:
+            pred_trimap = F.interpolate(pred_trimap, size=(images.shape[-2], images.shape[-1]), mode='bilinear', align_corners=False)
+            pred_trimap = torch.argmax(pred_trimap, dim=1, keepdim=True).float() / 2.0
+            norm_hq_features = torch.concat((norm_hq_features, pred_trimap.detach()), dim=1)
+        elif self.with_trimap:
+            mask = (norm_hq_features > 0).float()
+            for i_batch in range(images.shape[0]):
+                mask_area = torch.sum(mask[i_batch])
+                kernel_size = max(self.min_kernel_size, int((mask_area ** 0.5) / self.kernel_div))
+                kernel_size = min(kernel_size, self.gen_trimap.max_kernal - 1)
+                mask[i_batch, 0] = self.gen_trimap(mask[i_batch, 0], kernel_size=kernel_size)
+            trimaps = mask
+            norm_hq_features = torch.concat((norm_hq_features, trimaps), dim=1)
+
+        conditional_images = torch.concatenate((images, norm_hq_features), dim=1)
+        return conditional_images
+
+    def forward(self, images, hq_features, vit_intern_feat, return_alpha_logits=False, pred_trimap=None):
+        
+        condition_input = self.preprocess_inputs(images, hq_features, pred_trimap)
+
+        if not self.sam2_multi_scale_feates:
+            # aggregate 4 vit_intern_feat
+            # assert len(vit_intern_feat) == self.vit_intern_feat_num
+            vit_feats = []
+            for idx in self.vit_intern_feat_index:
+                vit_feats.append(self.vit_feat_proj[str(idx)](vit_intern_feat[idx].permute(0, 3, 1, 2)))
+            vit_feats = torch.concat(vit_feats, dim=1)
+            vit_aggregation_feats = self.vit_feat_aggregation(vit_feats)
+
+        details = []
+        dfeatures = condition_input
+
+        if hasattr(self, 'big_kernel_process'):
+            dfeatures = self.big_kernel_process(dfeatures)
+
+        for i in range(len(self.down_blks)):
+            if self.sam2_multi_scale_feates:
+                if i == 2:
+                    dfeatures = torch.concat((dfeatures, self.mid_layer[0](vit_intern_feat['high_res_feats'][0])), dim=1)
+                elif i == 3:
+                    dfeatures = torch.concat((dfeatures, self.mid_layer[1](vit_intern_feat['high_res_feats'][1])), dim=1)
+            dfeatures = self.down_blks[i](dfeatures)
+            details.append(dfeatures)
+
+        if self.sam2_multi_scale_feates:
+            out = torch.concat((details[-1], self.mid_layer[2](vit_intern_feat['image_embed'])), dim=1)
+            out = self.mid_layer[3](out)
+        else:
+            out = self.mid_layer(torch.concat((details[-1], vit_aggregation_feats), dim=1))
+        for i in range(len(self.up_blks)):
+            out = self.up_blks[i](out, details[-i - 1])
+        alpha = torch.sigmoid(self.matting_head(out))
+        if return_alpha_logits:
+            return alpha, out
+        else:
+            return alpha
+
+
+
+if __name__ == '__main__':
+
+    from engine.mattingtrainer import parameter_count_table
+
+    model = MattingDetailDecoder(img_feat_in = 5, vit_intern_feat_index=[0])
+    x = torch.randn((2, 5, 1024, 1024))
+    hq_features = torch.randn((2, 1, 1024, 1024))
+    vit_feat = [torch.randn((2, 64, 64, 1024)) for _ in range(4)]
+
+    out = model(x, hq_features, vit_feat)
+    print(out.shape)
+
+    print("Trainable parameters: \n" + parameter_count_table(model, trainable_only=True, max_depth=5))

modeling/meta_arch/__init__.py

@@ -0,0 +1 @@
+from .sam_hq_matting import SamHqMatte

modeling/meta_arch/sam_hq_matting.py

@@ -0,0 +1,671 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision
+import os
+import numpy as np
+from PIL import Image
+from copy import deepcopy
+from collections import defaultdict
+
+from detectron2.structures import ImageList
+from detectron2.utils.comm import get_local_rank
+from modeling.semantic_enhanced_matting.predictor import SamPredictor
+from modeling.semantic_enhanced_matting.condition_conv import ConditionConv, ConditionEmbedding, ConditionAdd, BBoxEmbedInteract, BBoxInteract, BBoxInteractInOut
+from modeling.semantic_enhanced_matting.modeling.image_encoder import PatchEmbed
+from modeling.semantic_enhanced_matting.modeling.common import LayerNorm2d
+from modeling.decoder.unet_detail_capture import MattingDetailDecoder
+from modeling.semantic_enhanced_matting.feature_fusion import FeatureFusion
+from sam2.sam2_image_predictor import SAM2ImagePredictor
+
+from modeling.semantic_enhanced_matting.modeling.mask_decoder_hq_matting import MaskDecoderHQMatting
+from modeling.semantic_enhanced_matting.modeling import TwoWayTransformer
+
+from peft import LoraConfig, get_peft_model
+from peft.tuners.lora.layer import LoraLayer
+from peft.tuners.tuners_utils import BaseTunerLayer
+
+from data.rand_augment import RandAugment
+import random
+import kornia.filters as kf
+
+
+class SamHqMatte(nn.Module):
+
+    target_length = 1024
+
+    def __init__(
+        self,
+        *,
+        sam_model,
+        hq_token_only,
+        hq_features_type,
+        matting_decoder,
+        criterion,
+        pixel_mean,
+        pixel_std,
+        multimask_output=False,
+        vis_period=None,
+        output_dir=None,
+        lora_rank = None,
+        lora_alpha = None,
+        lora_target_modules = ["qkv", "proj"],
+        lora_dropout = 0.1,
+        w_dora = False,
+        w_rslora = False,
+        lora_on_mask_decoder = False,
+        frozen_sam_hq_reg = None,
+        reg_margin = 0.85,
+        w_attention_mask = False,
+        alpha_reg_range = None,
+        alpha_reg_weight = 1.0,
+        coconut_pl = False,
+        coconut_pl_alpha = 1.0,
+        coconut_self_training = False,
+        eval_w_sam_hq_mask = False,
+        backbone_condition = False,
+        condition_wo_conv = False,
+        w_only_bbox_cond = False,
+        coconut_only_known_l1 = False,
+        backbone_bbox_prompt = None,
+        backbone_bbox_prompt_loc = [2, 3], 
+        backbone_bbox_prompt_loss_weight = 1.0,
+        concat_gen_trimap = False,
+        multi_matting_decoder = None,
+        w_all_logits = False,
+        bbox_prompt_all_block = None,
+        matting_token = False,
+        test_w_hq_token = False,
+        sam_hq_token_reg = None,
+        feat_cross_attn_fusion = False,
+        trimap_loss_type = None,
+        reg_on_sam_logits = False,
+        reg_w_bce_loss = False,
+        complex_trimap_pred_layer = False,
+        matting_token_sup = None,
+        matting_token_sup_loss_weight = None,
+        sam2 = False,
+    ):
+        super(SamHqMatte, self).__init__()
+
+        self.sam_model = sam_model
+        self.sam_predictor = SamPredictor(self.sam_model) if not sam2 else SAM2ImagePredictor(self.sam_model)  # already in eval mode and no_grad
+        self.hq_token_only = hq_token_only
+        self.multimask_output = multimask_output
+        self.hq_features_type = hq_features_type
+
+        self.matting_decoder = matting_decoder
+
+        self.criterion = criterion
+
+        self.register_buffer(
+            "pixel_mean", torch.tensor(pixel_mean).view(-1, 1, 1), False
+        )
+        self.register_buffer("pixel_std", torch.tensor(pixel_std).view(-1, 1, 1), False)
+        assert (
+            self.pixel_mean.shape == self.pixel_std.shape
+        ), f"{self.pixel_mean} and {self.pixel_std} have different shapes!"
+
+        self.vis_period = vis_period
+        if output_dir is not None and output_dir != '?':
+            self.output_dir = os.path.join(output_dir, 'vis_results')
+            os.makedirs(self.output_dir, exist_ok=True)
+        self.train_iter_index = 0
+
+        self.lora_rank = lora_rank
+        self.lora_alpha = lora_alpha
+        self.lora_target_modules = lora_target_modules
+        self.lora_dropout = lora_dropout
+        self.w_dora = w_dora
+        self.w_rslora = w_rslora
+        self.lora_on_mask_decoder = lora_on_mask_decoder
+        self.frozen_sam_hq_reg = frozen_sam_hq_reg
+        self.reg_margin = reg_margin
+        self.w_attention_mask = w_attention_mask
+        self.alpha_reg_range = alpha_reg_range
+        self.alpha_reg_weight = alpha_reg_weight
+        self.coconut_pl = coconut_pl
+        self.coconut_pl_alpha = coconut_pl_alpha
+        self.coconut_self_training = coconut_self_training
+        self.eval_w_sam_hq_mask = eval_w_sam_hq_mask
+        self.backbone_condition = backbone_condition
+        self.condition_wo_conv = condition_wo_conv
+        self.w_only_bbox_cond = w_only_bbox_cond
+        self.coconut_only_known_l1 = coconut_only_known_l1
+        self.backbone_bbox_prompt = backbone_bbox_prompt
+        self.backbone_bbox_prompt_loc = backbone_bbox_prompt_loc
+        self.backbone_bbox_prompt_loss_weight = backbone_bbox_prompt_loss_weight
+        self.concat_gen_trimap = concat_gen_trimap
+        self.multi_matting_decoder = multi_matting_decoder
+        self.w_all_logits = w_all_logits
+        self.bbox_prompt_all_block = bbox_prompt_all_block
+        self.matting_token = matting_token
+        self.test_w_hq_token = test_w_hq_token
+        self.sam_hq_token_reg = sam_hq_token_reg
+        self.feat_cross_attn_fusion = feat_cross_attn_fusion
+        self.trimap_loss_type = trimap_loss_type
+        self.reg_on_sam_logits = reg_on_sam_logits
+        self.reg_w_bce_loss = reg_w_bce_loss
+        self.complex_trimap_pred_layer = complex_trimap_pred_layer
+        self.matting_token_sup = matting_token_sup
+        self.sam2 = sam2
+        assert self.matting_token_sup in {'alpha', 'trimap', None}
+        self.matting_token_sup_loss_weight = matting_token_sup_loss_weight
+        if self.matting_token_sup is not None:
+            assert self.backbone_bbox_prompt in {'bbox', None}
+        if self.frozen_sam_hq_reg is not None:
+            assert self.lora_rank is not None
+        if self.w_attention_mask:
+            self.attention_head = deepcopy(self.matting_decoder)
+        if self.coconut_self_training:
+            self.rand_aug = RandAugment(3,6)
+            self.warm_iter_coconut_self_training = 5000
+        if self.backbone_condition:
+            assert self.lora_rank is not None
+        if self.backbone_bbox_prompt is not None:
+            assert self.lora_rank is not None
+        if self.w_all_logits:
+            self.sam_predictor.model.mask_decoder.w_all_logits = True
+        if self.bbox_prompt_all_block:
+            assert self.lora_rank is not None
+        if self.matting_token and not self.sam2:
+            self.sam_predictor.model.mask_decoder.hq_token_only = self.hq_token_only
+
+    @property
+    def device(self):
+        return self.pixel_mean.device
+
+    def init_lora(self, model=None):
+        if model is not None and self.lora_rank >= 1:
+            if self.lora_on_mask_decoder:
+                self.lora_target_modules += ["q_proj", "k_proj", "v_proj", "out_proj"]
+                modules_to_save = None
+            else:
+                modules_to_save = ['matting_decoder']
+
+            lora_config = LoraConfig(
+                r=self.lora_rank,
+                lora_alpha=self.lora_alpha,
+                use_rslora=self.w_rslora,
+                use_dora=self.w_dora,
+                init_lora_weights="gaussian",
+                target_modules=self.lora_target_modules,
+                lora_dropout=self.lora_dropout,
+                modules_to_save=modules_to_save
+            )
+            model = get_peft_model(model, lora_config)
+            if self.lora_on_mask_decoder:
+                for n, p in model.matting_decoder.named_parameters():
+                    if n.split('modules_to_save.default.')[-1] in model.matting_decoder.trainable_params_str:
+                        p.requires_grad = True
+            else:
+                for n, p in model.matting_decoder.named_parameters():
+                    if n.split('modules_to_save.default.')[-1] in model.matting_decoder.frozen_params_str:
+                        p.requires_grad = False
+            return model
+        elif self.lora_rank >= 1:
+            lora_config = LoraConfig(
+                r=self.lora_rank,
+                lora_alpha=self.lora_alpha,
+                use_rslora=self.w_rslora,
+                use_dora=self.w_dora,
+                init_lora_weights="gaussian",
+                target_modules=self.lora_target_modules,
+                lora_dropout=self.lora_dropout,
+            )
+            self.sam_predictor.model.image_encoder = get_peft_model(self.sam_predictor.model.image_encoder, lora_config)
+
+            if self.sam2:
+                for n, p in self.sam_predictor.model.image_encoder.named_parameters():
+                    if 'bbox_mask' in n:
+                        p.requires_grad = True
+
+        if self.backbone_condition:
+            if self.w_only_bbox_cond:
+                self.condition_embedding = ConditionEmbedding(condition_num = 4, pos_embedding_dim = 160)
+            else:
+                self.condition_embedding = ConditionEmbedding(condition_num = 5, pos_embedding_dim = 128)
+
+            if self.condition_wo_conv:
+                self.condition_conv = nn.ModuleList([ConditionAdd() for _ in range(4)])
+            else:
+                self.condition_conv = nn.ModuleList([ConditionConv(
+                    in_channels = self.sam_predictor.model.image_encoder.embed_dim, 
+                    out_channels = self.sam_predictor.model.image_encoder.embed_dim,
+                    bottleneck_channels = 512
+                ) for _ in range(4)])
+        
+        if self.backbone_bbox_prompt is not None and not self.sam2:
+            self.condition_layer = nn.ModuleDict()
+            self.condition_layer['patch_embed'] =  PatchEmbed(
+                kernel_size=(self.sam_predictor.model.image_encoder.patch_size, self.sam_predictor.model.image_encoder.patch_size),
+                stride=(self.sam_predictor.model.image_encoder.patch_size, self.sam_predictor.model.image_encoder.patch_size),
+                in_chans=4,
+                embed_dim=self.sam_predictor.model.image_encoder.embed_dim,
+            )
+            if self.multi_matting_decoder is None:
+                if self.backbone_bbox_prompt in {'trimap', 'alpha_trimap'}:
+                    transformer_dim = self.sam_predictor.model.image_encoder.embed_dim
+                    for i in self.backbone_bbox_prompt_loc:
+                        if self.complex_trimap_pred_layer:
+                            self.condition_layer['{}_pred_layer'.format(i)] = nn.Sequential(
+                                nn.ConvTranspose2d(transformer_dim, transformer_dim // 2, kernel_size=2, stride=2),
+                                LayerNorm2d(transformer_dim // 2),  # 512
+                                nn.GELU(),
+                                nn.Conv2d(transformer_dim // 2, transformer_dim // 4, kernel_size=3, stride=1, padding=1),
+                                LayerNorm2d(transformer_dim // 4),  # 256
+                                nn.GELU(),
+                                nn.ConvTranspose2d(transformer_dim // 4, transformer_dim // 8, kernel_size=2, stride=2),
+                                LayerNorm2d(transformer_dim // 8),  # 128
+                                nn.GELU(),
+                                nn.Conv2d(transformer_dim // 8, transformer_dim // 16, kernel_size=3, stride=1, padding=1),
+                                LayerNorm2d(transformer_dim // 16),  # 64
+                                nn.GELU(),
+                                nn.Conv2d(transformer_dim // 16, 3, kernel_size=3, stride=1, padding=1),
+                            )
+                        else:
+                            self.condition_layer['{}_pred_layer'.format(i)] = nn.Sequential(
+                                nn.ConvTranspose2d(transformer_dim, transformer_dim // 4, kernel_size=2, stride=2),
+                                LayerNorm2d(transformer_dim // 4),
+                                nn.GELU(),
+                                nn.ConvTranspose2d(transformer_dim // 4, transformer_dim // 8, kernel_size=2, stride=2),
+                                nn.GELU(),
+                                nn.Conv2d(transformer_dim // 8, 3, kernel_size=1, stride=1),
+                            )
+                elif self.backbone_bbox_prompt == 'alpha':
+                    transformer_dim = self.sam_predictor.model.image_encoder.embed_dim
+                    for i in self.backbone_bbox_prompt_loc:
+                        if self.complex_trimap_pred_layer:
+                            self.condition_layer['{}_pred_layer'.format(i)] = nn.Sequential(
+                                nn.ConvTranspose2d(transformer_dim, transformer_dim // 2, kernel_size=2, stride=2),
+                                LayerNorm2d(transformer_dim // 2),  # 512
+                                nn.GELU(),
+                                nn.Conv2d(transformer_dim // 2, transformer_dim // 4, kernel_size=3, stride=1, padding=1),
+                                LayerNorm2d(transformer_dim // 4),  # 256
+                                nn.GELU(),
+                                nn.ConvTranspose2d(transformer_dim // 4, transformer_dim // 8, kernel_size=2, stride=2),
+                                LayerNorm2d(transformer_dim // 8),  # 128
+                                nn.GELU(),
+                                nn.Conv2d(transformer_dim // 8, transformer_dim // 16, kernel_size=3, stride=1, padding=1),
+                                LayerNorm2d(transformer_dim // 16),  # 64
+                                nn.GELU(),
+                                nn.Conv2d(transformer_dim // 16, 1, kernel_size=3, stride=1, padding=1),
+                                nn.Sigmoid()
+                            )
+                        else:
+                            self.condition_layer['{}_pred_layer'.format(i)] = nn.Sequential(
+                                nn.ConvTranspose2d(transformer_dim, transformer_dim // 4, kernel_size=2, stride=2),
+                                LayerNorm2d(transformer_dim // 4),
+                                nn.GELU(),
+                                nn.ConvTranspose2d(transformer_dim // 4, transformer_dim // 8, kernel_size=2, stride=2),
+                                nn.GELU(),
+                                nn.Conv2d(transformer_dim // 8, 1, kernel_size=1, stride=1),
+                                nn.Sigmoid()
+                            )
+            if self.bbox_prompt_all_block is not None:
+                if self.bbox_prompt_all_block == 'reuse_cross-self-attn':
+                    self.condition_layer['prompt_layer'] = BBoxInteract(
+                        position_point_embedding = deepcopy(self.sam_predictor.model.prompt_encoder.pe_layer), 
+                        point_weight = deepcopy(self.sam_predictor.model.prompt_encoder.point_embeddings)
+                    )
+                elif self.bbox_prompt_all_block == 'in-out-bbox_cross-self-attn':
+                    self.condition_layer['prompt_layer'] = BBoxInteractInOut(downsample_rate = 2)
+                else:
+                    embed_type, interact_type = self.bbox_prompt_all_block.split('_')
+                    self.condition_layer['prompt_layer'] = BBoxEmbedInteract(embed_type, interact_type)
+
+        if self.feat_cross_attn_fusion:
+            self.condition_layer['feature_fusion'] = FeatureFusion(in_channels=self.sam_predictor.model.image_encoder.embed_dim, attn_compression_ratio=8)
+
+    def condition_bbox_and_instance_num(self):
+        self.sam_predictor.model.image_encoder.conv_necks = None
+
+    def forward_samhq_and_matting_decoder(self, images, bbox, condition_proj=None, return_hq_token=False):
+        # get features from SAM image encoder
+        if self.sam2:
+            interm_features, sam2_logits, matting_logits, pred_trimap = self.forward_samhq(images, bbox, condition_proj)
+            sam2_logits = F.interpolate(sam2_logits, size=images.shape[-2:], mode='bilinear', align_corners=False)
+            matting_logits = F.interpolate(matting_logits, size=images.shape[-2:], mode='bilinear', align_corners=False)
+            sam_hq_matting_token = {
+                'masks_hq': sam2_logits,
+                'masks_matting': matting_logits
+            }
+            hq_features = matting_logits
+            low_res_masks = matting_logits
+        else:
+            if self.matting_token:
+                features, image_pe, sparse_embeddings, dense_embeddings, interm_features, sam_hq_matting_token, pred_trimap = self.forward_samhq(images, bbox, condition_proj)
+                if return_hq_token:
+                    return sam_hq_matting_token['masks_hq']
+                else:
+                    if not self.training and self.test_w_hq_token:
+                        low_res_masks, hq_features = sam_hq_matting_token['masks_hq'], sam_hq_matting_token['masks_hq']
+                    else:
+                        low_res_masks, hq_features = sam_hq_matting_token['masks_matting'], sam_hq_matting_token['masks_matting']
+            else:
+                features, image_pe, sparse_embeddings, dense_embeddings, interm_features, hq_features, sam_logits, low_res_masks, pred_trimap = self.forward_samhq(images, bbox, condition_proj)
+                if return_hq_token:
+                    return hq_features
+                sam_hq_matting_token = {'masks_hq': hq_features, 'masks_sam': sam_logits}
+
+        # get alpha from our proposed matting_decoder
+        if isinstance(self.matting_decoder, MattingDetailDecoder):
+            pred_alpha = self.matting_decoder(
+                images = images,
+                hq_features = hq_features,
+                vit_intern_feat = interm_features,
+                return_alpha_logits = (self.alpha_reg_range is not None),
+                pred_trimap = pred_trimap
+            )
+        else:
+            pred_alpha = self.matting_decoder(
+                image_embeddings = features,  # [B, 256, 64, 64]
+                image_pe = image_pe,
+                sparse_prompt_embeddings = sparse_embeddings,
+                dense_prompt_embeddings = dense_embeddings,
+                multimask_output = False,
+                interm_embeddings = interm_features,  # [B, 256, 64, 64]
+                hq_features = hq_features,
+                images = images,
+                return_alpha_logits = (self.alpha_reg_range is not None),
+                pred_trimap = pred_trimap
+            )
+        return low_res_masks, pred_alpha, pred_trimap, sam_hq_matting_token
+
+    def forward(self, batched_inputs):  # image: [1, 3, 643, 960]: 0.0~1.0, trimap: [1, 1, 643, 960]: 0.0~1.0
+
+        inputs = self.preprocess_inputs(batched_inputs) 
+        images, bbox, gt_alpha, trimap, condition = inputs['images'], inputs['bbox'], inputs['alpha'], inputs['trimap'], inputs['condition']
+
+        if self.backbone_condition:
+            condition_proj = self.condition_embedding(condition) 
+        elif self.backbone_bbox_prompt is not None or self.bbox_prompt_all_block is not None:
+            condition_proj = bbox
+        else:
+            condition_proj = None
+
+        low_res_masks, pred_alpha, pred_trimap, sam_hq_matting_token = self.forward_samhq_and_matting_decoder(images, bbox, condition_proj)
+        
+        assert not self.training
+        if self.eval_w_sam_hq_mask:
+            self.sam_predictor.model.image_encoder.disable_adapter_layers()
+            with torch.no_grad():
+                ori_features, ori_interm_features = self.sam_predictor.model.image_encoder(images)
+                samhq_low_res_masks = self.forward_samhq_others(images, bbox, ori_features, ori_interm_features)[-1]
+                samhq_low_res_masks = F.interpolate(samhq_low_res_masks, size=(images.shape[-2], images.shape[-1]), mode='bilinear', align_corners=False)
+            self.sam_predictor.model.image_encoder.enable_adapter_layers()
+
+            return pred_alpha, samhq_low_res_masks
+        else:
+            return pred_alpha
+        
+    def forward_samhq_image_encoder(self, images, condition_proj=None):
+        if self.sam2:
+            backbone_out = self.sam_predictor.model.forward_image([images, condition_proj])
+            _, vision_feats, _, _ = self.sam_predictor.model._prepare_backbone_features(backbone_out)
+            # Add no_mem_embed, which is added to the lowest rest feat. map during training on videos
+            if self.sam_predictor.model.directly_add_no_mem_embed:
+                vision_feats[-1] = vision_feats[-1] + self.sam_predictor.model.no_mem_embed
+            feats = [
+                feat.permute(1, 2, 0).view(feat.shape[1], -1, *feat_size)
+                for feat, feat_size in zip(vision_feats[::-1], self.sam_predictor._bb_feat_sizes[::-1])
+            ][::-1]
+            return {"image_embed": feats[-1], "high_res_feats": feats[:-1]}, None, None
+        else:
+            if self.backbone_condition:
+                condition_layer = self.condition_conv 
+            elif self.backbone_bbox_prompt:
+                condition_layer = self.condition_layer
+            else:
+                condition_layer = None
+            # [B, 3, 1024, 1024]: -2. ~ 2. --> [B, 256, 64, 64], 4 x [B, 64, 64, 768]
+            features, interm_features, pred_trimap = self.sam_predictor.model.image_encoder(images, condition_proj, condition_layer)
+            return features, interm_features, pred_trimap
+    
+    # @torch.no_grad()
+    def forward_samhq_others(self, images, bbox, features, interm_features):
+        if self.sam2:
+            sam2_logits, matting_logits = self.sam_predictor.predict_batch_boxes_and_features(bbox, features)
+            return features, sam2_logits, matting_logits
+        
+        image_pe = self.sam_predictor.model.prompt_encoder.get_dense_pe()
+
+        cat_sparse_embeddings = []
+        cat_dense_prompt_embeddings = []
+        cat_hq_features = []
+        cat_sam_logits = []
+        cat_low_res_masks = []
+        cat_sam_hq_matting_token = defaultdict(list)
+
+        for idx in range(images.shape[0]):
+            # get hq_features from SAM_HQ mask decoder
+
+                # Embed prompts
+            sparse_embeddings, dense_embeddings = self.sam_predictor.model.prompt_encoder(
+                points=None,
+                # boxes=bbox[idx: idx + 1],
+                boxes=bbox[idx],  # [N, 4]
+                masks=None,
+            )  # [B, 2, 256], [B, 256, 64, 64]
+
+                # Predict masks
+            if isinstance(self.sam_predictor.model.mask_decoder, MaskDecoderHQMatting):
+                sam_hq_matting_token = self.sam_predictor.model.mask_decoder(
+                    image_embeddings = features[idx: idx + 1],
+                    image_pe = image_pe,
+                    sparse_prompt_embeddings = sparse_embeddings,
+                    dense_prompt_embeddings = dense_embeddings,
+                    multimask_output = self.multimask_output,
+                    interm_embeddings = [interm_feature[idx: idx + 1] for interm_feature in interm_features],
+                )
+                for key in sam_hq_matting_token.keys():
+                    cat_sam_hq_matting_token[key].append(sam_hq_matting_token[key])
+            else:
+                low_res_masks, masks_sam, hq_features = self.sam_predictor.model.mask_decoder(
+                    image_embeddings = features[idx: idx + 1],
+                    image_pe = image_pe,
+                    sparse_prompt_embeddings = sparse_embeddings,
+                    dense_prompt_embeddings = dense_embeddings,
+                    multimask_output = self.multimask_output,
+                    hq_token_only = self.hq_token_only,
+                    interm_embeddings = [interm_feature[idx: idx + 1] for interm_feature in interm_features],
+                    return_hq_features_type = self.hq_features_type
+                )
+                cat_hq_features.append(hq_features)
+                cat_sam_logits.append(masks_sam)
+                cat_low_res_masks.append(low_res_masks)
+
+            cat_sparse_embeddings.append(sparse_embeddings)
+            cat_dense_prompt_embeddings.append(dense_embeddings)
+            
+        sparse_embeddings = torch.stack(cat_sparse_embeddings, dim=0)  # [B, 1, 2, 256]
+        dense_embeddings = torch.stack(cat_dense_prompt_embeddings, dim=0)  # [B, 1, 256, 64, 64]
+        
+        if self.matting_token:
+            for key in cat_sam_hq_matting_token.keys():
+                cat_sam_hq_matting_token[key] = torch.cat(cat_sam_hq_matting_token[key], dim=0)
+                cat_sam_hq_matting_token[key] = F.interpolate(cat_sam_hq_matting_token[key], size=images.shape[-2:], mode='bilinear', align_corners=False)
+            sam_hq_matting_token = cat_sam_hq_matting_token
+            return features, image_pe, sparse_embeddings, dense_embeddings, interm_features, sam_hq_matting_token
+        else:
+            hq_features = torch.cat(cat_hq_features, dim=0)  # [B, 1, 256, 256]
+            low_res_masks = torch.cat(cat_low_res_masks, dim=0)  # [B, 1, 256, 256]
+            hq_features = F.interpolate(hq_features, size=images.shape[-2:], mode='bilinear', align_corners=False)  # [B, 1, 256, 256] --> [B, 1, 1024, 1024]
+            sam_logits = torch.cat(cat_sam_logits, dim=0)
+            sam_logits = F.interpolate(sam_logits, size=images.shape[-2:], mode='bilinear', align_corners=False)  # [B, 1, 256, 256] --> [B, 1, 1024, 1024]
+            return features, image_pe, sparse_embeddings, dense_embeddings, interm_features, hq_features, sam_logits, low_res_masks
+
+    def forward_samhq(self, images, bbox, condition_proj=None):
+        if self.lora_rank is None:
+            with torch.no_grad():
+                features, interm_features, pred_trimap = self.forward_samhq_image_encoder(images, condition_proj)
+        else:
+            features, interm_features, pred_trimap = self.forward_samhq_image_encoder(images, condition_proj)
+
+        return self.forward_samhq_others(images, bbox, features, interm_features) + (pred_trimap, )
+
+    def get_frozen_sam_logits(self, images, bbox, mask_type='hq'):
+        
+        if self.sam2:
+            features, _, _ = self.forward_samhq_image_encoder(images)
+            sam2_logits = self.sam_predictor.predict_batch_boxes_and_features(bbox, features, wo_matting_token=True)
+            sam2_logits = F.interpolate(sam2_logits, size=images.shape[-2:], mode='bilinear', align_corners=False)
+            return sam2_logits
+
+        assert mask_type in {'hq', 'sam'} 
+        features, interm_features, _ = self.forward_samhq_image_encoder(images)
+        image_pe = self.sam_predictor.model.prompt_encoder.get_dense_pe()
+
+        cat_logits = []
+        for idx in range(images.shape[0]):
+            sparse_embeddings, dense_embeddings = self.sam_predictor.model.prompt_encoder(points=None, boxes=bbox[idx], masks=None)
+
+            low_res_masks, masks_sam, hq_features = self.sam_predictor.model.frozen_mask_decoder(
+                image_embeddings = features[idx: idx + 1],
+                image_pe = image_pe,
+                sparse_prompt_embeddings = sparse_embeddings,
+                dense_prompt_embeddings = dense_embeddings,
+                multimask_output = self.multimask_output,
+                hq_token_only = self.hq_token_only,
+                interm_embeddings = [interm_feature[idx: idx + 1] for interm_feature in interm_features],
+                return_hq_features_type = self.hq_features_type
+            )
+            if mask_type == 'hq':
+                cat_logits.append(hq_features)  
+            else:
+                cat_logits.append(masks_sam)  
+        
+        logits = torch.cat(cat_logits, dim=0)  # [B, 1, 256, 256]
+        logits = F.interpolate(logits, size=images.shape[-2:], mode='bilinear', align_corners=False)  # [B, 1, 256, 256] --> [B, 1, 1024, 1024]
+        return logits
+
+    def vis_training_results(self, **kwargs):
+        # images, bbox, trimap, low_res_masks, pred_alpha, alpha
+        self.train_iter_index += 1
+        if self.train_iter_index % self.vis_period == 0:
+            batch_save_results = []
+            save_path = os.path.join(self.output_dir, '{:06d}_rank{}.jpg'.format(self.train_iter_index, get_local_rank()))
+            
+            # [('images', (4, 3, 1024, 1024), -2.117904, 2.64), ('bbox', (4, 1, 4), 0.0, 1023.0), ('trimap', (4, 1, 1024, 1024), 0.0, 1.0), ('low_res_masks', (4, 1, 256, 256), -20.38, 10.15), ('pred_alpha', (4, 1, 1024, 1024), 0.1547, 0.791), ('alpha', (4, 1, 1024, 1024), 0.0, 1.0)]
+            for key in kwargs.keys():
+                if key == 'bbox':
+                    continue
+                # turn all tensor to [B, H, W, 3]: 0~255 np.int8
+                if key == 'images':
+                    kwargs[key] = kwargs[key] * self.pixel_std + self.pixel_mean
+                    kwargs[key] = kwargs[key].permute(0, 2, 3, 1) * 255.0
+                    for i in range(kwargs['images'].shape[0]):
+                        l, u, r, d = int(kwargs['bbox'][i, 0, 0].item()), int(kwargs['bbox'][i, 0, 1].item()), int(kwargs['bbox'][i, 0, 2].item()), int(kwargs['bbox'][i, 0, 3].item())
+                        red_line = torch.tensor([[255., 0., 0.]], device=kwargs[key].device, dtype=kwargs[key].dtype)
+                        kwargs[key][i, u: d, l, :] = red_line
+                        kwargs[key][i, u: d, r, :] = red_line
+                        kwargs[key][i, u, l: r, :] = red_line
+                        kwargs[key][i, d, l: r, :] = red_line
+                elif key in {'low_res_masks', 'frozen_hq_token'}:
+                    if torch.max(kwargs[key]) <= 1:  # coconut ori alpha
+                        kwargs[key] = kwargs[key].permute(0, 2, 3, 1).repeat(1, 1, 1, 3) * 255.0
+                    else:
+                        kwargs[key] = F.interpolate(kwargs[key], size=(kwargs['images'].shape[-3], kwargs['images'].shape[-2]), mode='bilinear', align_corners=False)
+                        kwargs[key] = (kwargs[key] > self.sam_predictor.model.mask_threshold).float().permute(0, 2, 3, 1).repeat(1, 1, 1, 3) * 255.0
+                else:
+                    kwargs[key] = kwargs[key].permute(0, 2, 3, 1).repeat(1, 1, 1, 3) * 255.0
+
+                kwargs[key] = np.uint8(kwargs[key].detach().cpu().numpy())
+
+            for i in range(kwargs['images'].shape[0]):
+                save_results = []
+                for key in kwargs.keys():
+                    if key != 'bbox':
+                        save_results.append(kwargs[key][i])
+                batch_save_results.append(np.concatenate(save_results, axis=1))
+            
+            Image.fromarray(np.concatenate(batch_save_results, axis=0)).save(save_path)
+
+    def preprocess_inputs(self, batched_inputs):
+        """
+        Normalize, pad and batch the input images.
+        """
+        output = dict()
+
+        if "alpha" in batched_inputs:
+            alpha = batched_inputs["alpha"].to(self.device)
+        else:
+            alpha = None
+
+        bbox = batched_inputs["bbox"].to(self.device)
+
+        if self.training and self.coconut_self_training and sum([i == 'COCONut' for i in batched_inputs['dataset_name']]) >= 1:
+            output['coconut_ori_img'] = []
+            output['coconut_trimap'] = []
+            output['coconut_bbox'] = []
+            output['coconut_idx'] = []
+            for i, dataset_name in enumerate(batched_inputs['dataset_name']):
+                if dataset_name == 'COCONut':
+                    # generate coconut_aug_img
+                    img_np = np.uint8(batched_inputs["image"][i].permute(1, 2, 0).cpu().numpy() * 255.)
+                    strong_aug_img = self.rand_aug(Image.fromarray(img_np), cutout = False)
+                    strong_aug_img_tensor = torch.from_numpy(np.array(strong_aug_img)).to(self.device).permute(2, 0, 1)[None] / 255.
+                    blur_kernel_sigma = 1.0 + random.random()  # random from 1.0 ~ 2.0
+                    blur_filter = kf.GaussianBlur2d((101, 101), (blur_kernel_sigma, blur_kernel_sigma))
+                    blur_strong_aug_img_tensor = blur_filter(strong_aug_img_tensor)[0]
+
+                    output['coconut_ori_img'].append(batched_inputs["image"][i])
+                    batched_inputs["image"][i] = blur_strong_aug_img_tensor
+
+                    # generate coconut_trimap
+                    coconut_mask = (alpha[i] != 0).float()
+                    mask_area = torch.sum(coconut_mask)
+                    kernel_size = max(self.matting_decoder.min_kernel_size, int((mask_area ** 0.5) / 7))  # self.matting_decoder.kernel_div
+                    kernel_size = min(kernel_size, self.matting_decoder.gen_trimap.max_kernal - 1)
+                    output['coconut_trimap'].append(self.matting_decoder.gen_trimap(coconut_mask[0], kernel_size=kernel_size)[None])
+
+                    output['coconut_bbox'].append(bbox[i])
+                    output['coconut_idx'].append(i)
+
+            output['coconut_ori_img'] = torch.stack(output['coconut_ori_img']).to(self.device)
+            output['coconut_ori_img'] = (output['coconut_ori_img'] - self.pixel_mean) / self.pixel_std
+            output['coconut_trimap'] = torch.stack(output['coconut_trimap']).to(self.device)
+            output['coconut_bbox'] = torch.stack(output['coconut_bbox']).to(self.device)
+
+        images = batched_inputs["image"].to(self.device)
+        images = (images - self.pixel_mean) / self.pixel_std
+        assert images.shape[-2] == images.shape[-1] == 1024
+
+        if 'trimap' in batched_inputs.keys():
+            trimap = batched_inputs["trimap"].to(self.device)
+            assert len(torch.unique(trimap)) <= 3
+        else:
+            trimap = None
+
+        output['images'] = images
+        output['bbox'] = bbox
+        output['alpha'] = alpha
+        output['trimap'] = trimap
+
+        if 'hr_images' in batched_inputs.keys():
+            hr_images = batched_inputs["hr_images"].to(self.device)
+            hr_images = (hr_images - self.pixel_mean) / self.pixel_std
+            _, _, H, W = hr_images.shape
+            if hr_images.shape[-1] % 16 != 0 or hr_images.shape[-2] % 16 != 0:
+                new_H = (16 - hr_images.shape[-2] % 16) + H if hr_images.shape[-2] % 16 != 0 else H
+                new_W = (16 - hr_images.shape[-1] % 16) + W if hr_images.shape[-1] % 16 != 0 else W
+                new_hr_images = torch.zeros((hr_images.shape[0], hr_images.shape[1], new_H, new_W)).to(self.device)
+                new_hr_images[:,:,:H,:W] = hr_images[:,:,:,:]
+                del hr_images
+                hr_images = new_hr_images
+            output['hr_images'] = hr_images
+            output['hr_images_ori_h_w'] = (H, W)
+
+        if 'dataset_name' in batched_inputs.keys():
+            output['dataset_name'] = batched_inputs["dataset_name"]
+
+        if self.backbone_condition:
+            if self.w_only_bbox_cond:
+                output['condition'] = output['bbox'][:, 0, :]
+            else:
+                multi_fg_float = batched_inputs["multi_fg"].to(bbox.device).float()[:, None] * 512
+                output['condition'] = torch.concat((output['bbox'][:, 0, :], multi_fg_float), dim=-1)
+        else:
+            output['condition'] = None
+
+        return output

modeling/semantic_enhanced_matting/__init__.py

@@ -0,0 +1,17 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from .build_sam import (
+    build_sam,
+    build_sam_vit_h,
+    build_sam_vit_l,
+    build_sam_vit_b,
+    sam_model_registry,
+)
+from .build_sam_baseline import sam_model_registry_baseline
+from .predictor import SamPredictor
+from .automatic_mask_generator import SamAutomaticMaskGenerator
+from .mask_decoder_matting import MaskDecoderMatting