inital files

app.py

@@ -0,0 +1,97 @@
+import os
+import torch
+import gradio as gr
+import argparse
+import time
+import torch
+import torch.nn as nn
+import torch.optim as optim
+
+from tqdm import tqdm
+from PIL import Image
+from torch.utils.data import DataLoader
+from PIL import Image
+from torchvision import transforms
+
+from pipeline.resnet_csra import ResNet_CSRA
+from pipeline.vit_csra import VIT_B16_224_CSRA, VIT_L16_224_CSRA, VIT_CSRA
+from pipeline.dataset import DataSet
+from torchvision.transforms import transforms
+from utils.evaluation.eval import voc_classes, wider_classes, coco_classes, class_dict
+
+torch.manual_seed(0)
+
+if torch.cuda.is_available():
+    torch.backends.cudnn.deterministic = True
+
+# Device
+DEVICE = "cpu"
+print(DEVICE)
+
+# Make directories
+os.system("mkdir ./models")
+
+# Get model weights
+if not os.path.exists("./models/msl_c_voc.pth"):
+    os.system(
+        "wget -O ./models/msl_c_voc.pth https://github.com/hasibzunair/msl-recognition/releases/download/v1.0-models/msl_c_voc.pth"
+    )
+
+# Load model
+model = ResNet_CSRA(num_heads=1, lam=0.1, num_classes=20)
+normalize = transforms.Normalize(mean=[0, 0, 0], std=[1, 1, 1])
+model.to(DEVICE)
+print("Loading weights from {}".format("./models/msl_c_voc.pth"))
+model.load_state_dict(torch.load("./models/msl_c_voc.pth"))
+
+# Inference!
+def inference(img_path):
+    # read image
+    image = Image.open(img_path).convert("RGB")
+
+    # image pre-process
+    transforms_image = transforms.Compose([
+        transforms.Resize((448, 448)),
+        transforms.ToTensor(),
+        normalize
+    ])
+
+    image = transforms_image(image)
+    image = image.unsqueeze(0)
+
+    # Predict
+    result = []
+    model.eval()
+    with torch.no_grad():
+        image = image.to(DEVICE)
+        logit = model(image).squeeze(0)
+        logit = nn.Sigmoid()(logit)
+
+        pos = torch.where(logit > 0.5)[0].cpu().numpy()
+        for k in pos:
+            result.append(str(class_dict["voc07"][k]))
+    return result
+
+
+# Define ins outs placeholders
+inputs = gr.inputs.Image(type="filepath", label="Input Image")
+
+# Define style
+title = "Learning to Recognize Occluded and Small Objects with Partial Inputs"
+description = "TBA."
+article = "<p style='text-align: center'><a href='https://arxiv.org/abs/1512.03385' target='_blank'>Learning to Recognize Occluded and Small Objects with Partial Inputs</a> | <a href='https://github.com/hasibzunair/msl-recognition' target='_blank'>Github Repo</a></p>"
+
+voc_classes = ("aeroplane", "bicycle", "bird", "boat", "bottle",
+           "bus", "car", "cat", "chair", "cow", "diningtable",
+           "dog", "horse", "motorbike", "person", "pottedplant",
+           "sheep", "sofa", "train", "tvmonitor")
+
+# Run inference
+gr.Interface(inference, 
+            inputs, 
+            outputs="text", 
+            examples=["demo_images/000001.jpg", "demo_images/000006.jpg", "demo_images/000009.jpg"], 
+            title=title, 
+            description=description, 
+            article=article,
+            analytics_enabled=False).launch()

pipeline/csra.py

@@ -0,0 +1,55 @@
+import torch
+import torch.nn as nn
+
+
+
+class CSRA(nn.Module): # one basic block 
+    def __init__(self, input_dim, num_classes, T, lam):
+        super(CSRA, self).__init__()
+        self.T = T      # temperature       
+        self.lam = lam  # Lambda                        
+        self.head = nn.Conv2d(input_dim, num_classes, 1, bias=False)
+        self.softmax = nn.Softmax(dim=2)
+
+    def forward(self, x):
+        # x (B d H W)
+        # normalize classifier
+        # score (B C HxW)
+        score = self.head(x) / torch.norm(self.head.weight, dim=1, keepdim=True).transpose(0,1)
+        score = score.flatten(2)
+        base_logit = torch.mean(score, dim=2)
+
+        if self.T == 99: # max-pooling
+            att_logit = torch.max(score, dim=2)[0]
+        else:
+            score_soft = self.softmax(score * self.T)
+            # https://github.com/Kevinz-code/CSRA/issues/5
+            att_logit = torch.sum(score * score_soft, dim=2)
+
+        return base_logit + self.lam * att_logit
+
+    
+
+
+class MHA(nn.Module):  # multi-head attention
+    temp_settings = {  # softmax temperature settings
+        1: [1],
+        2: [1, 99],
+        4: [1, 2, 4, 99],
+        6: [1, 2, 3, 4, 5, 99],
+        8: [1, 2, 3, 4, 5, 6, 7, 99]
+    }
+
+    def __init__(self, num_heads, lam, input_dim, num_classes):
+        super(MHA, self).__init__()
+        self.temp_list = self.temp_settings[num_heads]
+        self.multi_head = nn.ModuleList([
+            CSRA(input_dim, num_classes, self.temp_list[i], lam)
+            for i in range(num_heads)
+        ])
+
+    def forward(self, x):
+        logit = 0.
+        for head in self.multi_head:
+            logit += head(x)
+        return logit

pipeline/dataset.py

@@ -0,0 +1,255 @@
+import json
+import glob
+import random 
+
+from torch.utils.data import Dataset
+from PIL import Image
+from torchvision.transforms import transforms
+import torch
+import numpy as np
+
+try:
+    from torchvision.transforms import InterpolationMode
+
+    BICUBIC = InterpolationMode.BICUBIC
+except ImportError:
+    BICUBIC = Image.BICUBIC
+
+
+# modify for transformation for vit
+# modfify wider crop-person images
+
+
+###### Base data loader ######
+class DataSet(Dataset):
+    def __init__(
+        self,
+        ann_files,
+        augs,
+        img_size,
+        dataset,
+    ):
+        self.dataset = dataset
+        self.ann_files = ann_files
+        self.augment = self.augs_function(augs, img_size)
+        self.transform = transforms.Compose(
+            [transforms.ToTensor(), transforms.Normalize(mean=[0, 0, 0], std=[1, 1, 1])]
+            # In this paper, we normalize the image data to [0, 1]
+            # You can also use the so called 'ImageNet' Normalization method
+        )
+        self.anns = []
+        self.load_anns()
+        print(self.augment)
+
+        # in wider dataset we use vit models
+        # so transformation has been changed
+        if self.dataset == "wider":
+            self.transform = transforms.Compose(
+                [
+                    transforms.ToTensor(),
+                    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+                ]
+            )
+
+    def augs_function(self, augs, img_size):
+        t = []
+        if "randomflip" in augs:
+            t.append(transforms.RandomHorizontalFlip())
+        if "ColorJitter" in augs:
+            t.append(
+                transforms.ColorJitter(
+                    brightness=0.5, contrast=0.5, saturation=0.5, hue=0
+                )
+            )
+        if "resizedcrop" in augs:
+            t.append(transforms.RandomResizedCrop(img_size, scale=(0.7, 1.0)))
+        if "RandAugment" in augs:
+            t.append(RandAugment())
+
+        t.append(transforms.Resize((img_size, img_size)))
+
+        return transforms.Compose(t)
+
+    def load_anns(self):
+        self.anns = []
+        for ann_file in self.ann_files:
+            json_data = json.load(open(ann_file, "r"))
+            self.anns += json_data
+
+    def __len__(self):
+        return len(self.anns)
+
+    def __getitem__(self, idx):
+        idx = idx % len(self)
+        ann = self.anns[idx]
+        img = Image.open(ann["img_path"]).convert("RGB")
+
+        if self.dataset == "wider":
+            x, y, w, h = ann["bbox"]
+            img_area = img.crop([x, y, x + w, y + h])
+            img_area = self.augment(img_area)
+            img_area = self.transform(img_area)
+            message = {
+                "img_path": ann["img_path"],
+                "target": torch.Tensor(ann["target"]),
+                "img": img_area,
+            }
+        else:  # voc and coco
+            img = self.augment(img)
+            img = self.transform(img)
+            message = {
+                "img_path": ann["img_path"],
+                "target": torch.Tensor(ann["target"]),
+                "img": img,
+            }
+
+        return message
+        # finally, if we use dataloader to get the data, we will get
+        # {
+        #     "img_path": list, # length = batch_size
+        #     "target": Tensor, # shape: batch_size * num_classes
+        #     "img": Tensor, # shape: batch_size * 3 * 224 * 224
+        # }
+
+
+def preprocess_scribble(img, img_size):
+    transform = transforms.Compose(
+        [
+            transforms.Resize(img_size, BICUBIC),
+            transforms.CenterCrop(img_size),
+            #_convert_image_to_rgb,
+            transforms.ToTensor(),
+        ]
+    )
+    return transform(img)
+
+
+class DataSetMaskSup(Dataset):
+    """
+    Data loader with scribbles.
+    """
+    def __init__(
+        self,
+        ann_files,
+        augs,
+        img_size,
+        dataset,
+    ):
+        self.dataset = dataset
+        self.ann_files = ann_files
+        self.img_size = img_size
+        self.augment = self.augs_function(augs, img_size)
+        self.transform = transforms.Compose(
+            [transforms.ToTensor(), transforms.Normalize(mean=[0, 0, 0], std=[1, 1, 1])]
+            # In this paper, we normalize the image data to [0, 1]
+            # You can also use the so called 'ImageNet' Normalization method
+        )
+        self.anns = []
+        self.load_anns()
+        print(self.augment)
+
+        # scribbles
+        self._scribbles_folder = "./datasets/SCRIBBLES"
+        
+        # Type of masks to use, this is hardcoded since we find that high masks
+        # work better in MSL. See paper for details.
+
+        # for low masks
+        # self._scribbles = sorted(glob.glob(self._scribbles_folder + "/*.png"))[
+        #     :1000
+        # ]
+
+        # for high masks
+        self._scribbles = sorted(glob.glob(self._scribbles_folder + "/*.png"))[::-1][
+            :1000
+        ]
+
+        # in wider dataset we use vit models
+        # so transformation has been changed
+        if self.dataset == "wider":
+            self.transform = transforms.Compose(
+                [
+                    transforms.ToTensor(),
+                    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+                ]
+            )
+
+    def augs_function(self, augs, img_size):
+        t = []
+        if "randomflip" in augs:
+            t.append(transforms.RandomHorizontalFlip())
+        if "ColorJitter" in augs:
+            t.append(
+                transforms.ColorJitter(
+                    brightness=0.5, contrast=0.5, saturation=0.5, hue=0
+                )
+            )
+        if "resizedcrop" in augs:
+            t.append(transforms.RandomResizedCrop(img_size, scale=(0.7, 1.0)))
+        if "RandAugment" in augs:
+            t.append(RandAugment())
+
+        t.append(transforms.Resize((img_size, img_size)))
+
+        return transforms.Compose(t)
+
+    def load_anns(self):
+        self.anns = []
+        for ann_file in self.ann_files:
+            json_data = json.load(open(ann_file, "r"))
+            self.anns += json_data
+
+    def __len__(self):
+        return len(self.anns)
+
+    def __getitem__(self, idx):
+        idx = idx % len(self)
+        ann = self.anns[idx]
+        img = Image.open(ann["img_path"]).convert("RGB")
+
+        # get scribble
+        scribble_path = self._scribbles[
+                random.randint(0, 950)
+            ]
+        scribble = Image.open(scribble_path).convert('P')
+        scribble = preprocess_scribble(scribble, self.img_size)
+        
+        scribble_t = (scribble > 0).float() # threshold to [0,1]
+        inv_scribble = (torch.max(scribble_t) - scribble_t) # inverted scribble      
+
+        if self.dataset == "wider":
+            x, y, w, h = ann["bbox"]
+            img_area = img.crop([x, y, x + w, y + h])
+            img_area = self.augment(img_area)
+            img_area = self.transform(img_area)
+
+            # masked image
+            masked_image = img_area * inv_scribble
+            message = {
+                "img_path": ann["img_path"],
+                "target": torch.Tensor(ann["target"]),
+                "img": img_area,
+                "masked_img": masked_image,
+                #"scribble": inv_scribble,
+            }
+        else:  # voc and coco
+            img = self.augment(img)
+            img = self.transform(img)
+            # masked image
+            masked_image = img * inv_scribble
+            message = {
+                "img_path": ann["img_path"],
+                "target": torch.Tensor(ann["target"]),
+                "img": img,
+                "masked_img": masked_image,
+                #"scribble": inv_scribble,
+            }
+
+        return message
+        # finally, if we use dataloader to get the data, we will get
+        # {
+        #     "img_path": list, # length = batch_size
+        #     "target": Tensor, # shape: batch_size * num_classes
+        #     "img": Tensor, # shape: batch_size * 3 * 224 * 224
+        #     "masked_img": Tensor, # shape: batch_size * 3 * 224 * 224
+        # }

pipeline/losses.py

@@ -0,0 +1,56 @@
+import torch
+import torch.nn as nn
+
+"""ASL taken from https://github.com/Alibaba-MIIL/ASL"""
+
+# Usage 
+# global criterion_asl
+# criterion_asl = AsymmetricLoss(gamma_neg=4, gamma_pos=0, clip=0.05, disable_torch_grad_focal_loss=True)
+# loss3 = criterion_asl(pred1, pred2)
+
+class AsymmetricLoss(nn.Module):
+    def __init__(self, gamma_neg=4, gamma_pos=1, clip=0.05, eps=1e-8, disable_torch_grad_focal_loss=True):
+        super(AsymmetricLoss, self).__init__()
+
+        self.gamma_neg = gamma_neg
+        self.gamma_pos = gamma_pos
+        self.clip = clip
+        self.disable_torch_grad_focal_loss = disable_torch_grad_focal_loss
+        self.eps = eps
+
+    def forward(self, x, y):
+        """"
+        Parameters
+        ----------
+        x: input logits
+        y: targets (multi-label binarized vector)
+        """
+
+        # Calculating Probabilities
+        x_sigmoid = torch.sigmoid(x)
+        xs_pos = x_sigmoid
+        xs_neg = 1 - x_sigmoid
+
+        # Asymmetric Clipping
+        if self.clip is not None and self.clip > 0:
+            xs_neg = (xs_neg + self.clip).clamp(max=1)
+
+        # Basic CE calculation
+        los_pos = y * torch.log(xs_pos.clamp(min=self.eps))
+        los_neg = (1 - y) * torch.log(xs_neg.clamp(min=self.eps))
+        loss = los_pos + los_neg
+
+        # Asymmetric Focusing
+        if self.gamma_neg > 0 or self.gamma_pos > 0:
+            if self.disable_torch_grad_focal_loss:
+                torch.set_grad_enabled(False)
+            pt0 = xs_pos * y
+            pt1 = xs_neg * (1 - y)  # pt = p if t > 0 else 1-p
+            pt = pt0 + pt1
+            one_sided_gamma = self.gamma_pos * y + self.gamma_neg * (1 - y)
+            one_sided_w = torch.pow(1 - pt, one_sided_gamma)
+            if self.disable_torch_grad_focal_loss:
+                torch.set_grad_enabled(True)
+            loss *= one_sided_w
+
+        return -loss.sum()

pipeline/models/tresnet/layers/anti_aliasing.py

@@ -0,0 +1,60 @@
+import torch
+import torch.nn.parallel
+import numpy as np
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class AntiAliasDownsampleLayer(nn.Module):
+    def __init__(self, remove_model_jit: bool = False, filt_size: int = 3, stride: int = 2,
+                 channels: int = 0):
+        super(AntiAliasDownsampleLayer, self).__init__()
+        if not remove_model_jit:
+            self.op = DownsampleJIT(filt_size, stride, channels)
+        else:
+            self.op = Downsample(filt_size, stride, channels)
+
+    def forward(self, x):
+        return self.op(x)
+
+
+@torch.jit.script
+class DownsampleJIT(object):
+    def __init__(self, filt_size: int = 3, stride: int = 2, channels: int = 0):
+        self.stride = stride
+        self.filt_size = filt_size
+        self.channels = channels
+
+        assert self.filt_size == 3
+        assert stride == 2
+        a = torch.tensor([1., 2., 1.])
+
+        filt = (a[:, None] * a[None, :]).clone().detach()
+        filt = filt / torch.sum(filt)
+        self.filt = filt[None, None, :, :].repeat((self.channels, 1, 1, 1)).cuda().half()
+
+    def __call__(self, input: torch.Tensor):
+        if input.dtype != self.filt.dtype:
+            self.filt = self.filt.float() 
+        input_pad = F.pad(input, (1, 1, 1, 1), 'reflect')
+        return F.conv2d(input_pad, self.filt, stride=2, padding=0, groups=input.shape[1])
+
+
+class Downsample(nn.Module):
+    def __init__(self, filt_size=3, stride=2, channels=None):
+        super(Downsample, self).__init__()
+        self.filt_size = filt_size
+        self.stride = stride
+        self.channels = channels
+
+
+        assert self.filt_size == 3
+        a = torch.tensor([1., 2., 1.])
+
+        filt = (a[:, None] * a[None, :]).clone().detach()
+        filt = filt / torch.sum(filt)
+        self.filt = filt[None, None, :, :].repeat((self.channels, 1, 1, 1))
+
+    def forward(self, input):
+        input_pad = F.pad(input, (1, 1, 1, 1), 'reflect')
+        return F.conv2d(input_pad, self.filt, stride=self.stride, padding=0, groups=input.shape[1])

pipeline/models/tresnet/layers/avg_pool.py

@@ -0,0 +1,19 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+
+class FastAvgPool2d(nn.Module):
+    def __init__(self, flatten=False):
+        super(FastAvgPool2d, self).__init__()
+        self.flatten = flatten
+
+    def forward(self, x):
+        if self.flatten:
+            in_size = x.size()
+            return x.view((in_size[0], in_size[1], -1)).mean(dim=2)
+        else:
+            return x.view(x.size(0), x.size(1), -1).mean(-1).view(x.size(0), x.size(1), 1, 1)
+
+

pipeline/models/tresnet/layers/general_layers.py

@@ -0,0 +1,93 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from pipeline.models.tresnet.layers.avg_pool import FastAvgPool2d
+
+
+class Flatten(nn.Module):
+    def forward(self, x):
+        return x.view(x.size(0), -1)
+
+
+class DepthToSpace(nn.Module):
+
+    def __init__(self, block_size):
+        super().__init__()
+        self.bs = block_size
+
+    def forward(self, x):
+        N, C, H, W = x.size()
+        x = x.view(N, self.bs, self.bs, C // (self.bs ** 2), H, W)  # (N, bs, bs, C//bs^2, H, W)
+        x = x.permute(0, 3, 4, 1, 5, 2).contiguous()  # (N, C//bs^2, H, bs, W, bs)
+        x = x.view(N, C // (self.bs ** 2), H * self.bs, W * self.bs)  # (N, C//bs^2, H * bs, W * bs)
+        return x
+
+
+class SpaceToDepthModule(nn.Module):
+    def __init__(self, remove_model_jit=False):
+        super().__init__()
+        if not remove_model_jit:
+            self.op = SpaceToDepthJit()
+        else:
+            self.op = SpaceToDepth()
+
+    def forward(self, x):
+        return self.op(x)
+
+
+class SpaceToDepth(nn.Module):
+    def __init__(self, block_size=4):
+        super().__init__()
+        assert block_size == 4
+        self.bs = block_size
+
+    def forward(self, x):
+        N, C, H, W = x.size()
+        x = x.view(N, C, H // self.bs, self.bs, W // self.bs, self.bs)  # (N, C, H//bs, bs, W//bs, bs)
+        x = x.permute(0, 3, 5, 1, 2, 4).contiguous()  # (N, bs, bs, C, H//bs, W//bs)
+        x = x.view(N, C * (self.bs ** 2), H // self.bs, W // self.bs)  # (N, C*bs^2, H//bs, W//bs)
+        return x
+
+
+@torch.jit.script
+class SpaceToDepthJit(object):
+    def __call__(self, x: torch.Tensor):
+        # assuming hard-coded that block_size==4 for acceleration
+        N, C, H, W = x.size()
+        x = x.view(N, C, H // 4, 4, W // 4, 4)  # (N, C, H//bs, bs, W//bs, bs)
+        x = x.permute(0, 3, 5, 1, 2, 4).contiguous()  # (N, bs, bs, C, H//bs, W//bs)
+        x = x.view(N, C * 16, H // 4, W // 4)  # (N, C*bs^2, H//bs, W//bs)
+        return x
+
+
+class hard_sigmoid(nn.Module):
+    def __init__(self, inplace=True):
+        super(hard_sigmoid, self).__init__()
+        self.inplace = inplace
+
+    def forward(self, x):
+        if self.inplace:
+            return x.add_(3.).clamp_(0., 6.).div_(6.)
+        else:
+            return F.relu6(x + 3.) / 6.
+
+
+class SEModule(nn.Module):
+
+    def __init__(self, channels, reduction_channels, inplace=True):
+        super(SEModule, self).__init__()
+        self.avg_pool = FastAvgPool2d()
+        self.fc1 = nn.Conv2d(channels, reduction_channels, kernel_size=1, padding=0, bias=True)
+        self.relu = nn.ReLU(inplace=inplace)
+        self.fc2 = nn.Conv2d(reduction_channels, channels, kernel_size=1, padding=0, bias=True)
+        # self.activation = hard_sigmoid(inplace=inplace)
+        self.activation = nn.Sigmoid()
+
+    def forward(self, x):
+        x_se = self.avg_pool(x)
+        x_se2 = self.fc1(x_se)
+        x_se2 = self.relu(x_se2)
+        x_se = self.fc2(x_se2)
+        x_se = self.activation(x_se)
+        return x * x_se

pipeline/models/tresnet/tresnet.py

@@ -0,0 +1,268 @@
+import torch
+import torch.nn as nn
+from torch.nn import Module as Module
+from collections import OrderedDict
+from pipeline.models.tresnet.layers.anti_aliasing import AntiAliasDownsampleLayer
+from .layers.avg_pool import FastAvgPool2d
+from .layers.general_layers import SEModule, SpaceToDepthModule
+from inplace_abn import InPlaceABN, ABN
+import torch.nn.functional as F
+
+def InplacABN_to_ABN(module: nn.Module) -> nn.Module:
+    # convert all InplaceABN layer to bit-accurate ABN layers.
+    if isinstance(module, InPlaceABN):
+        module_new = ABN(module.num_features, activation=module.activation,
+                         activation_param=module.activation_param)
+        for key in module.state_dict():
+            module_new.state_dict()[key].copy_(module.state_dict()[key])
+        module_new.training = module.training
+        module_new.weight.data = module_new.weight.abs() + module_new.eps
+        return module_new
+    for name, child in reversed(module._modules.items()):
+        new_child = InplacABN_to_ABN(child)
+        if new_child != child:
+            module._modules[name] = new_child
+    return module
+
+class bottleneck_head(nn.Module):
+    def __init__(self, num_features, num_classes, bottleneck_features=200):
+        super(bottleneck_head, self).__init__()
+        self.embedding_generator = nn.ModuleList()
+        self.embedding_generator.append(nn.Linear(num_features, bottleneck_features))
+        self.embedding_generator = nn.Sequential(*self.embedding_generator)
+        self.FC = nn.Linear(bottleneck_features, num_classes)
+
+    def forward(self, x):
+        self.embedding = self.embedding_generator(x)
+        logits = self.FC(self.embedding)
+        return logits
+
+
+def conv2d(ni, nf, stride):
+    return nn.Sequential(
+        nn.Conv2d(ni, nf, kernel_size=3, stride=stride, padding=1, bias=False),
+        nn.BatchNorm2d(nf),
+        nn.ReLU(inplace=True)
+    )
+
+
+def conv2d_ABN(ni, nf, stride, activation="leaky_relu", kernel_size=3, activation_param=1e-2, groups=1):
+    return nn.Sequential(
+        nn.Conv2d(ni, nf, kernel_size=kernel_size, stride=stride, padding=kernel_size // 2, groups=groups,
+                  bias=False),
+        InPlaceABN(num_features=nf, activation=activation, activation_param=activation_param)
+    )
+
+
+class BasicBlock(Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, use_se=True, anti_alias_layer=None):
+        super(BasicBlock, self).__init__()
+        if stride == 1:
+            self.conv1 = conv2d_ABN(inplanes, planes, stride=1, activation_param=1e-3)
+        else:
+            if anti_alias_layer is None:
+                self.conv1 = conv2d_ABN(inplanes, planes, stride=2, activation_param=1e-3)
+            else:
+                self.conv1 = nn.Sequential(conv2d_ABN(inplanes, planes, stride=1, activation_param=1e-3),
+                                           anti_alias_layer(channels=planes, filt_size=3, stride=2))
+
+        self.conv2 = conv2d_ABN(planes, planes, stride=1, activation="identity")
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+        reduce_layer_planes = max(planes * self.expansion // 4, 64)
+        self.se = SEModule(planes * self.expansion, reduce_layer_planes) if use_se else None
+
+    def forward(self, x):
+        if self.downsample is not None:
+            residual = self.downsample(x)
+        else:
+            residual = x
+
+        out = self.conv1(x)
+        out = self.conv2(out)
+
+        if self.se is not None: out = self.se(out)
+
+        out += residual
+
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, use_se=True, anti_alias_layer=None):
+        super(Bottleneck, self).__init__()
+        self.conv1 = conv2d_ABN(inplanes, planes, kernel_size=1, stride=1, activation="leaky_relu",
+                                activation_param=1e-3)
+        if stride == 1:
+            self.conv2 = conv2d_ABN(planes, planes, kernel_size=3, stride=1, activation="leaky_relu",
+                                    activation_param=1e-3)
+        else:
+            if anti_alias_layer is None:
+                self.conv2 = conv2d_ABN(planes, planes, kernel_size=3, stride=2, activation="leaky_relu",
+                                        activation_param=1e-3)
+            else:
+                self.conv2 = nn.Sequential(conv2d_ABN(planes, planes, kernel_size=3, stride=1,
+                                                      activation="leaky_relu", activation_param=1e-3),
+                                           anti_alias_layer(channels=planes, filt_size=3, stride=2))
+
+        self.conv3 = conv2d_ABN(planes, planes * self.expansion, kernel_size=1, stride=1,
+                                activation="identity")
+
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+        reduce_layer_planes = max(planes * self.expansion // 8, 64)
+        self.se = SEModule(planes, reduce_layer_planes) if use_se else None
+
+    def forward(self, x):
+        if self.downsample is not None:
+            residual = self.downsample(x)
+        else:
+            residual = x
+
+        out = self.conv1(x)
+        out = self.conv2(out)
+        if self.se is not None: out = self.se(out)
+
+        out = self.conv3(out)
+        out = out + residual  # no inplace
+        out = self.relu(out)
+
+        return out
+
+
+class TResNet(Module):
+
+    def __init__(self, layers, in_chans=3, num_classes=1000, width_factor=1.0,
+                 do_bottleneck_head=False,bottleneck_features=512):
+        super(TResNet, self).__init__()
+
+        # Loss function
+        self.loss_func = F.binary_cross_entropy_with_logits
+
+        # JIT layers
+        space_to_depth = SpaceToDepthModule()
+        anti_alias_layer = AntiAliasDownsampleLayer
+        global_pool_layer = FastAvgPool2d(flatten=True)
+
+        # TResnet stages
+        self.inplanes = int(64 * width_factor)
+        self.planes = int(64 * width_factor)
+        conv1 = conv2d_ABN(in_chans * 16, self.planes, stride=1, kernel_size=3)
+        layer1 = self._make_layer(BasicBlock, self.planes, layers[0], stride=1, use_se=True,
+                                  anti_alias_layer=anti_alias_layer)  # 56x56
+        layer2 = self._make_layer(BasicBlock, self.planes * 2, layers[1], stride=2, use_se=True,
+                                  anti_alias_layer=anti_alias_layer)  # 28x28
+        layer3 = self._make_layer(Bottleneck, self.planes * 4, layers[2], stride=2, use_se=True,
+                                  anti_alias_layer=anti_alias_layer)  # 14x14
+        layer4 = self._make_layer(Bottleneck, self.planes * 8, layers[3], stride=2, use_se=False,
+                                  anti_alias_layer=anti_alias_layer)  # 7x7
+
+        # body
+        self.body = nn.Sequential(OrderedDict([
+            ('SpaceToDepth', space_to_depth),
+            ('conv1', conv1),
+            ('layer1', layer1),
+            ('layer2', layer2),
+            ('layer3', layer3),
+            ('layer4', layer4)]))
+
+        # head
+        self.embeddings = []
+        self.global_pool = nn.Sequential(OrderedDict([('global_pool_layer', global_pool_layer)]))
+        self.num_features = (self.planes * 8) * Bottleneck.expansion
+        if do_bottleneck_head:
+            fc = bottleneck_head(self.num_features, num_classes,
+                                 bottleneck_features=bottleneck_features)
+        else:
+            fc = nn.Linear(self.num_features , num_classes)
+
+        self.head = nn.Sequential(OrderedDict([('fc', fc)]))
+
+        # model initilization
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='leaky_relu')
+            elif isinstance(m, nn.BatchNorm2d) or isinstance(m, InPlaceABN):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # residual connections special initialization
+        for m in self.modules():
+            if isinstance(m, BasicBlock):
+                m.conv2[1].weight = nn.Parameter(torch.zeros_like(m.conv2[1].weight))  # BN to zero
+            if isinstance(m, Bottleneck):
+                m.conv3[1].weight = nn.Parameter(torch.zeros_like(m.conv3[1].weight))  # BN to zero
+            if isinstance(m, nn.Linear): m.weight.data.normal_(0, 0.01)
+
+    def _make_layer(self, block, planes, blocks, stride=1, use_se=True, anti_alias_layer=None):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            layers = []
+            if stride == 2:
+                # avg pooling before 1x1 conv
+                layers.append(nn.AvgPool2d(kernel_size=2, stride=2, ceil_mode=True, count_include_pad=False))
+            layers += [conv2d_ABN(self.inplanes, planes * block.expansion, kernel_size=1, stride=1,
+                                  activation="identity")]
+            downsample = nn.Sequential(*layers)
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample, use_se=use_se,
+                            anti_alias_layer=anti_alias_layer))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks): layers.append(
+            block(self.inplanes, planes, use_se=use_se, anti_alias_layer=anti_alias_layer))
+        return nn.Sequential(*layers)
+
+    def forward_train(self, x, target):
+        x = self.body(x)
+        self.embeddings = self.global_pool(x)
+        logits = self.head(self.embeddings)
+        loss = self.loss_func(logits, target, reduction="mean")
+        return logits, loss
+
+    def forward_test(self, x):
+        x = self.body(x)
+        self.embeddings = self.global_pool(x)
+        logits = self.head(self.embeddings)
+        return logits
+
+    def forward(self, x, target=None):
+        if target is not None:
+            return self.forward_train(x, target)
+        else:
+            return self.forward_test(x)
+
+
+def TResnetM(num_classes):
+    """Constructs a medium TResnet model.
+    """
+    in_chans = 3
+    model = TResNet(layers=[3, 4, 11, 3], num_classes=num_classes, in_chans=in_chans)
+    return model
+
+
+def TResnetL(num_classes):
+    """Constructs a large TResnet model.
+    """
+    in_chans = 3
+    do_bottleneck_head = False
+    model = TResNet(layers=[4, 5, 18, 3], num_classes=num_classes, in_chans=in_chans, width_factor=1.2,
+                    do_bottleneck_head=do_bottleneck_head)
+    return model
+
+
+def TResnetXL(num_classes):
+    """Constructs a xlarge TResnet model.
+    """
+    in_chans = 3
+    model = TResNet(layers=[4, 5, 24, 3], num_classes=num_classes, in_chans=in_chans, width_factor=1.3)
+    return model

pipeline/models/utils/__init__.py

@@ -0,0 +1,2 @@
+from .factory import create_model
+__all__ = ['create_model']

pipeline/models/utils/factory.py

@@ -0,0 +1,25 @@
+import logging
+
+logger = logging.getLogger(__name__)
+
+from ..tresnet import TResnetM, TResnetL, TResnetXL
+
+
+def create_model(args):
+    """Create a model
+    """
+    model_params = {'args': args, 'num_classes': args.num_classes}
+    args = model_params['args']
+    args.model_name = args.model_name.lower()
+
+    if args.model_name=='tresnet_m':
+        model = TResnetM(model_params)
+    elif args.model_name=='tresnet_l':
+        model = TResnetL(model_params)
+    elif args.model_name=='tresnet_xl':
+        model = TResnetXL(model_params)
+    else:
+        print("model: {} not found !!".format(args.model_name))
+        exit(-1)
+
+    return model

pipeline/resnet_csra.py

@@ -0,0 +1,94 @@
+from torchvision.models import ResNet
+from torchvision.models.resnet import Bottleneck, BasicBlock
+from .csra import CSRA, MHA
+import torch.utils.model_zoo as model_zoo
+import logging
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+model_urls = {
+    "resnet18": "https://download.pytorch.org/models/resnet18-5c106cde.pth",
+    "resnet34": "https://download.pytorch.org/models/resnet34-333f7ec4.pth",
+    "resnet50": "https://download.pytorch.org/models/resnet50-19c8e357.pth",
+    "resnet101": "https://download.pytorch.org/models/resnet101-5d3b4d8f.pth",
+    "resnet152": "https://download.pytorch.org/models/resnet152-b121ed2d.pth",
+}
+
+
+class ResNet_CSRA(ResNet):
+    arch_settings = {
+        18: (BasicBlock, (2, 2, 2, 2)),
+        34: (BasicBlock, (3, 4, 6, 3)),
+        50: (Bottleneck, (3, 4, 6, 3)),
+        101: (Bottleneck, (3, 4, 23, 3)),
+        152: (Bottleneck, (3, 8, 36, 3)),
+    }
+
+    def __init__(
+        self, num_heads, lam, num_classes, depth=101, input_dim=2048, cutmix=None
+    ):
+        self.block, self.layers = self.arch_settings[depth]
+        self.depth = depth
+        super(ResNet_CSRA, self).__init__(self.block, self.layers)
+        self.init_weights(pretrained=True, cutmix=cutmix)
+
+        self.classifier = MHA(num_heads, lam, input_dim, num_classes)
+        self.loss_func = F.binary_cross_entropy_with_logits
+        # todo
+        # criterion = nn.BCEWithLogitsLoss()  # loss combines a Sigmoid layer and the BCELoss in one single class
+
+    def backbone(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        return x
+
+    def forward_train(self, x, target):
+        x = self.backbone(x)
+        logit = self.classifier(x)
+        loss = self.loss_func(logit, target, reduction="mean")
+        return logit, loss
+
+    def forward_test(self, x):
+        x = self.backbone(x)
+        x = self.classifier(x)
+        return x
+
+    def forward(self, x, target=None):
+        if target is not None:
+            return self.forward_train(x, target)
+        else:
+            return self.forward_test(x)
+
+    def init_weights(self, pretrained=True, cutmix=None):
+        if cutmix is not None:
+            print("backbone params inited by CutMix pretrained model")
+            state_dict = torch.load(cutmix)
+        elif pretrained:
+            print("backbone params inited by Pytorch official model")
+            model_url = model_urls["resnet{}".format(self.depth)]
+            state_dict = model_zoo.load_url(model_url)
+
+        model_dict = self.state_dict()
+        try:
+            pretrained_dict = {k: v for k, v in state_dict.items() if k in model_dict}
+            self.load_state_dict(pretrained_dict)
+        except:
+            logger = logging.getLogger()
+            logger.info(
+                "the keys in pretrained model is not equal to the keys in the ResNet you choose, trying to fix..."
+            )
+            state_dict = self._keysFix(model_dict, state_dict)
+            self.load_state_dict(state_dict)
+
+        # remove the original 1000-class fc
+        self.fc = nn.Sequential()

pipeline/timm_utils/__init__.py

@@ -0,0 +1,4 @@
+from .tuple import to_ntuple, to_2tuple, to_3tuple, to_4tuple
+from .drop import DropBlock2d, DropPath, drop_block_2d, drop_path
+from .weight_init import trunc_normal_
+

pipeline/timm_utils/drop.py

@@ -0,0 +1,168 @@
+""" DropBlock, DropPath
+
+PyTorch implementations of DropBlock and DropPath (Stochastic Depth) regularization layers.
+
+Papers:
+DropBlock: A regularization method for convolutional networks (https://arxiv.org/abs/1810.12890)
+
+Deep Networks with Stochastic Depth (https://arxiv.org/abs/1603.09382)
+
+Code:
+DropBlock impl inspired by two Tensorflow impl that I liked:
+ - https://github.com/tensorflow/tpu/blob/master/models/official/resnet/resnet_model.py#L74
+ - https://github.com/clovaai/assembled-cnn/blob/master/nets/blocks.py
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def drop_block_2d(
+        x, drop_prob: float = 0.1, block_size: int = 7,  gamma_scale: float = 1.0,
+        with_noise: bool = False, inplace: bool = False, batchwise: bool = False):
+    """ DropBlock. See https://arxiv.org/pdf/1810.12890.pdf
+
+    DropBlock with an experimental gaussian noise option. This layer has been tested on a few training
+    runs with success, but needs further validation and possibly optimization for lower runtime impact.
+    """
+    B, C, H, W = x.shape
+    total_size = W * H
+    clipped_block_size = min(block_size, min(W, H))
+    # seed_drop_rate, the gamma parameter
+    gamma = gamma_scale * drop_prob * total_size / clipped_block_size ** 2 / (
+        (W - block_size + 1) * (H - block_size + 1))
+
+    # Forces the block to be inside the feature map.
+    w_i, h_i = torch.meshgrid(torch.arange(W).to(x.device), torch.arange(H).to(x.device))
+    valid_block = ((w_i >= clipped_block_size // 2) & (w_i < W - (clipped_block_size - 1) // 2)) & \
+                  ((h_i >= clipped_block_size // 2) & (h_i < H - (clipped_block_size - 1) // 2))
+    valid_block = torch.reshape(valid_block, (1, 1, H, W)).to(dtype=x.dtype)
+
+    if batchwise:
+        # one mask for whole batch, quite a bit faster
+        uniform_noise = torch.rand((1, C, H, W), dtype=x.dtype, device=x.device)
+    else:
+        uniform_noise = torch.rand_like(x)
+    block_mask = ((2 - gamma - valid_block + uniform_noise) >= 1).to(dtype=x.dtype)
+    block_mask = -F.max_pool2d(
+        -block_mask,
+        kernel_size=clipped_block_size,  # block_size,
+        stride=1,
+        padding=clipped_block_size // 2)
+
+    if with_noise:
+        normal_noise = torch.randn((1, C, H, W), dtype=x.dtype, device=x.device) if batchwise else torch.randn_like(x)
+        if inplace:
+            x.mul_(block_mask).add_(normal_noise * (1 - block_mask))
+        else:
+            x = x * block_mask + normal_noise * (1 - block_mask)
+    else:
+        normalize_scale = (block_mask.numel() / block_mask.to(dtype=torch.float32).sum().add(1e-7)).to(x.dtype)
+        if inplace:
+            x.mul_(block_mask * normalize_scale)
+        else:
+            x = x * block_mask * normalize_scale
+    return x
+
+
+def drop_block_fast_2d(
+        x: torch.Tensor, drop_prob: float = 0.1, block_size: int = 7,
+        gamma_scale: float = 1.0, with_noise: bool = False, inplace: bool = False, batchwise: bool = False):
+    """ DropBlock. See https://arxiv.org/pdf/1810.12890.pdf
+
+    DropBlock with an experimental gaussian noise option. Simplied from above without concern for valid
+    block mask at edges.
+    """
+    B, C, H, W = x.shape
+    total_size = W * H
+    clipped_block_size = min(block_size, min(W, H))
+    gamma = gamma_scale * drop_prob * total_size / clipped_block_size ** 2 / (
+            (W - block_size + 1) * (H - block_size + 1))
+
+    if batchwise:
+        # one mask for whole batch, quite a bit faster
+        block_mask = torch.rand((1, C, H, W), dtype=x.dtype, device=x.device) < gamma
+    else:
+        # mask per batch element
+        block_mask = torch.rand_like(x) < gamma
+    block_mask = F.max_pool2d(
+        block_mask.to(x.dtype), kernel_size=clipped_block_size, stride=1, padding=clipped_block_size // 2)
+
+    if with_noise:
+        normal_noise = torch.randn((1, C, H, W), dtype=x.dtype, device=x.device) if batchwise else torch.randn_like(x)
+        if inplace:
+            x.mul_(1. - block_mask).add_(normal_noise * block_mask)
+        else:
+            x = x * (1. - block_mask) + normal_noise * block_mask
+    else:
+        block_mask = 1 - block_mask
+        normalize_scale = (block_mask.numel() / block_mask.to(dtype=torch.float32).sum().add(1e-7)).to(dtype=x.dtype)
+        if inplace:
+            x.mul_(block_mask * normalize_scale)
+        else:
+            x = x * block_mask * normalize_scale
+    return x
+
+
+class DropBlock2d(nn.Module):
+    """ DropBlock. See https://arxiv.org/pdf/1810.12890.pdf
+    """
+    def __init__(self,
+                 drop_prob=0.1,
+                 block_size=7,
+                 gamma_scale=1.0,
+                 with_noise=False,
+                 inplace=False,
+                 batchwise=False,
+                 fast=True):
+        super(DropBlock2d, self).__init__()
+        self.drop_prob = drop_prob
+        self.gamma_scale = gamma_scale
+        self.block_size = block_size
+        self.with_noise = with_noise
+        self.inplace = inplace
+        self.batchwise = batchwise
+        self.fast = fast  # FIXME finish comparisons of fast vs not
+
+    def forward(self, x):
+        if not self.training or not self.drop_prob:
+            return x
+        if self.fast:
+            return drop_block_fast_2d(
+                x, self.drop_prob, self.block_size, self.gamma_scale, self.with_noise, self.inplace, self.batchwise)
+        else:
+            return drop_block_2d(
+                x, self.drop_prob, self.block_size, self.gamma_scale, self.with_noise, self.inplace, self.batchwise)
+
+
+def drop_path(x, drop_prob: float = 0., training: bool = False):
+    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
+
+    This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
+    the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
+    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
+    changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
+    'survival rate' as the argument.
+
+    """
+    if drop_prob == 0. or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
+    random_tensor.floor_()  # binarize
+    output = x.div(keep_prob) * random_tensor
+    return output
+
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
+    """
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)

pipeline/timm_utils/tuple.py

@@ -0,0 +1,27 @@
+""" Layer/Module Helpers
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+from itertools import repeat
+from torch._six import container_abcs
+
+
+# From PyTorch internals
+def _ntuple(n):
+    def parse(x):
+        if isinstance(x, container_abcs.Iterable):
+            return x
+        return tuple(repeat(x, n))
+    return parse
+
+
+to_1tuple = _ntuple(1)
+to_2tuple = _ntuple(2)
+to_3tuple = _ntuple(3)
+to_4tuple = _ntuple(4)
+to_ntuple = _ntuple
+
+
+
+
+

pipeline/timm_utils/weight_init.py

@@ -0,0 +1,60 @@
+import torch
+import math
+import warnings
+
+
+def _no_grad_trunc_normal_(tensor, mean, std, a, b):
+    # Cut & paste from PyTorch official master until it's in a few official releases - RW
+    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
+    def norm_cdf(x):
+        # Computes standard normal cumulative distribution function
+        return (1. + math.erf(x / math.sqrt(2.))) / 2.
+
+    if (mean < a - 2 * std) or (mean > b + 2 * std):
+        warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
+                      "The distribution of values may be incorrect.",
+                      stacklevel=2)
+
+    with torch.no_grad():
+        # Values are generated by using a truncated uniform distribution and
+        # then using the inverse CDF for the normal distribution.
+        # Get upper and lower cdf values
+        l = norm_cdf((a - mean) / std)
+        u = norm_cdf((b - mean) / std)
+
+        # Uniformly fill tensor with values from [l, u], then translate to
+        # [2l-1, 2u-1].
+        tensor.uniform_(2 * l - 1, 2 * u - 1)
+
+        # Use inverse cdf transform for normal distribution to get truncated
+        # standard normal
+        tensor.erfinv_()
+
+        # Transform to proper mean, std
+        tensor.mul_(std * math.sqrt(2.))
+        tensor.add_(mean)
+
+        # Clamp to ensure it's in the proper range
+        tensor.clamp_(min=a, max=b)
+        return tensor
+
+
+def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
+    # type: (Tensor, float, float, float, float) -> Tensor
+    r"""Fills the input Tensor with values drawn from a truncated
+    normal distribution. The values are effectively drawn from the
+    normal distribution :math:`\mathcal{N}(\text{mean}, \text{std}^2)`
+    with values outside :math:`[a, b]` redrawn until they are within
+    the bounds. The method used for generating the random values works
+    best when :math:`a \leq \text{mean} \leq b`.
+    Args:
+        tensor: an n-dimensional `torch.Tensor`
+        mean: the mean of the normal distribution
+        std: the standard deviation of the normal distribution
+        a: the minimum cutoff value
+        b: the maximum cutoff value
+    Examples:
+        >>> w = torch.empty(3, 5)
+        >>> nn.init.trunc_normal_(w)
+    """
+    return _no_grad_trunc_normal_(tensor, mean, std, a, b)

pipeline/vit_csra.py

@@ -0,0 +1,303 @@
+""" Vision Transformer (ViT) in PyTorch
+
+A PyTorch implement of Vision Transformers as described in
+'An Image Is Worth 16 x 16 Words: Transformers for Image Recognition at Scale' - https://arxiv.org/abs/2010.11929
+
+The official jax code is released and available at https://github.com/google-research/vision_transformer
+
+Status/TODO:
+* Models updated to be compatible with official impl. Args added to support backward compat for old PyTorch weights.
+* Weights ported from official jax impl for 384x384 base and small models, 16x16 and 32x32 patches.
+* Trained (supervised on ImageNet-1k) my custom 'small' patch model to 77.9, 'base' to 79.4 top-1 with this code.
+* Hopefully find time and GPUs for SSL or unsupervised pretraining on OpenImages w/ ImageNet fine-tune in future.
+
+Acknowledgments:
+* The paper authors for releasing code and weights, thanks!
+* I fixed my class token impl based on Phil Wang's https://github.com/lucidrains/vit-pytorch ... check it out
+for some einops/einsum fun
+* Simple transformer style inspired by Andrej Karpathy's https://github.com/karpathy/minGPT
+* Bert reference code checks against Huggingface Transformers and Tensorflow Bert
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.model_zoo as model_zoo
+from functools import partial
+from .timm_utils import DropPath, to_2tuple, trunc_normal_
+from .csra import MHA, CSRA
+
+
+default_cfgs = {
+    'vit_base_patch16_224':  'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_p16_224-80ecf9dd.pth',
+    'vit_large_patch16_224':'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_large_p16_224-4ee7a4dc.pth'
+    }
+
+
+
+class Mlp(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads # 64
+        # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights
+        self.scale = qk_scale or head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        # qkv (3, B, 12, N, C/12)
+        # q (B, 12, N, C/12)
+        # k (B, 12, N, C/12)
+        # v (B, 12, N, C/12)
+        # attn (B, 12, N, N)
+        # x (B, 12, N, C/12)
+        q, k, v = qkv[0], qkv[1], qkv[2]   # make torchscript happy (cannot use tensor as tuple)
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+
+        x = self.proj(x)
+        x = self.proj_drop(x)
+
+        return x
+
+
+class Block(nn.Module):
+
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+    def forward(self, x):
+        x = x + self.drop_path(self.attn(self.norm1(x)))
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x
+
+
+class PatchEmbed(nn.Module):
+    """ Image to Patch Embedding
+    """
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_patches = num_patches
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        # FIXME look at relaxing size constraints
+        assert H == self.img_size[0] and W == self.img_size[1], \
+            f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+        x = self.proj(x).flatten(2).transpose(1, 2)
+        return x
+
+
+class HybridEmbed(nn.Module):
+    """ CNN Feature Map Embedding
+    Extract feature map from CNN, flatten, project to embedding dim.
+    """
+    def __init__(self, backbone, img_size=224, feature_size=None, in_chans=3, embed_dim=768):
+        super().__init__()
+        assert isinstance(backbone, nn.Module)
+        img_size = to_2tuple(img_size)
+        self.img_size = img_size
+        self.backbone = backbone
+        if feature_size is None:
+            with torch.no_grad():
+                # FIXME this is hacky, but most reliable way of determining the exact dim of the output feature
+                # map for all networks, the feature metadata has reliable channel and stride info, but using
+                # stride to calc feature dim requires info about padding of each stage that isn't captured.
+                training = backbone.training
+                if training:
+                    backbone.eval()
+                o = self.backbone(torch.zeros(1, in_chans, img_size[0], img_size[1]))[-1]
+                feature_size = o.shape[-2:]
+                feature_dim = o.shape[1]
+                backbone.train(training)
+        else:
+            feature_size = to_2tuple(feature_size)
+            feature_dim = self.backbone.feature_info.channels()[-1]
+        self.num_patches = feature_size[0] * feature_size[1]
+        self.proj = nn.Linear(feature_dim, embed_dim)
+
+    def forward(self, x):
+        x = self.backbone(x)[-1]
+        x = x.flatten(2).transpose(1, 2)
+        x = self.proj(x)
+        return x
+
+
+class VIT_CSRA(nn.Module):
+    """ Vision Transformer with support for patch or hybrid CNN input stage
+    """
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12,
+                 num_heads=12, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0.,
+                 drop_path_rate=0., hybrid_backbone=None, norm_layer=nn.LayerNorm, cls_num_heads=1, cls_num_cls=80, lam=0.3):
+        super().__init__()
+        self.add_w = 0.
+        self.normalize = False
+        self.num_classes = num_classes
+        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+
+        if hybrid_backbone is not None:
+            self.patch_embed = HybridEmbed(
+                hybrid_backbone, img_size=img_size, in_chans=in_chans, embed_dim=embed_dim)
+        else:
+            self.patch_embed = PatchEmbed(
+                img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
+        num_patches = self.patch_embed.num_patches
+        self.HW = int(math.sqrt(num_patches))
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+        self.blocks = nn.ModuleList([
+            Block(
+                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
+                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer)
+            for i in range(depth)])
+        self.norm = norm_layer(embed_dim)
+
+        # NOTE as per official impl, we could have a pre-logits representation dense layer + tanh here
+        #self.repr = nn.Linear(embed_dim, representation_size)
+        #self.repr_act = nn.Tanh()
+
+        trunc_normal_(self.pos_embed, std=.02)
+        trunc_normal_(self.cls_token, std=.02)
+        self.apply(self._init_weights)
+
+        # We add our MHA (CSRA) beside the orginal VIT structure below
+        self.head = nn.Sequential() # delete original classifier
+        self.classifier = MHA(input_dim=embed_dim, num_heads=cls_num_heads, num_classes=cls_num_cls, lam=lam)
+
+        self.loss_func = F.binary_cross_entropy_with_logits
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.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 backbone(self, x):
+        B = x.shape[0]
+        x = self.patch_embed(x)
+
+        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
+        x = torch.cat((cls_tokens, x), dim=1)
+        x = x + self.pos_embed
+        x = self.pos_drop(x)
+
+        for blk in self.blocks:
+            x = blk(x)
+        x = self.norm(x)
+
+        # (B, 1+HW, C)
+        # we use all the feature to form the tensor like B C H W
+        x = x[:, 1:]
+        b, hw, c = x.shape
+        x = x.transpose(1, 2)
+        x = x.reshape(b, c, self.HW, self.HW)
+
+        return x
+
+    def forward_train(self, x, target):
+        x = self.backbone(x)
+        logit = self.classifier(x)
+        loss = self.loss_func(logit, target, reduction="mean")
+        return logit, loss
+
+    def forward_test(self, x):
+        x = self.backbone(x)
+        x = self.classifier(x)
+        return x
+
+    def forward(self, x, target=None):
+        if target is not None:
+            return self.forward_train(x, target)
+        else:
+            return self.forward_test(x)
+    
+
+        
+
+def _conv_filter(state_dict, patch_size=16):
+    """ convert patch embedding weight from manual patchify + linear proj to conv"""
+    out_dict = {}
+    for k, v in state_dict.items():
+        if 'patch_embed.proj.weight' in k:
+            v = v.reshape((v.shape[0], 3, patch_size, patch_size))
+        out_dict[k] = v
+    return out_dict
+
+
+def VIT_B16_224_CSRA(pretrained=True, cls_num_heads=1, cls_num_cls=80, lam=0.3):
+    model = VIT_CSRA(
+        patch_size=16, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6), cls_num_heads=cls_num_heads, cls_num_cls=cls_num_cls, lam=lam)
+
+    model_url = default_cfgs['vit_base_patch16_224']
+    if pretrained:
+        state_dict = model_zoo.load_url(model_url)
+        model.load_state_dict(state_dict, strict=False)
+    return model
+
+
+def VIT_L16_224_CSRA(pretrained=True, cls_num_heads=1, cls_num_cls=80, lam=0.3):
+    model = VIT_CSRA(
+        patch_size=16, embed_dim=1024, depth=24, num_heads=16, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6), cls_num_heads=cls_num_heads, cls_num_cls=cls_num_cls, lam=lam)
+
+    model_url = default_cfgs['vit_large_patch16_224']
+    if pretrained:
+        state_dict = model_zoo.load_url(model_url)
+        model.load_state_dict(state_dict, strict=False)
+        # load_pretrained(model, num_classes=model.num_classes, in_chans=kwargs.get('in_chans', 3))
+    return model

requirements.txt

@@ -0,0 +1,3 @@
+torch
+torchvision
+Pillow

utils/evaluation/cal_PR.py

@@ -0,0 +1,79 @@
+import json
+import numpy as np
+
+
+
+def json_metric(score_json, target_json, num_classes, types):
+    assert len(score_json) == len(target_json)
+    scores = np.zeros((len(score_json), num_classes))
+    targets = np.zeros((len(target_json), num_classes))
+    for index in range(len(score_json)):
+        scores[index] = score_json[index]["scores"]
+        targets[index] = target_json[index]["target"]
+
+
+    return metric(scores, targets, types)
+
+def json_metric_top3(score_json, target_json, num_classes, types):
+    assert len(score_json) == len(target_json)
+    scores = np.zeros((len(score_json), num_classes))
+    targets = np.zeros((len(target_json), num_classes))
+    for index in range(len(score_json)):
+        tmp = np.array(score_json[index]['scores'])
+        idx = np.argsort(-tmp)
+        idx_after_3 = idx[3:]
+        tmp[idx_after_3] = 0.
+
+        scores[index] = tmp
+        # scores[index] = score_json[index]["scores"]
+        targets[index] = target_json[index]["target"]
+
+    return metric(scores, targets, types)
+
+
+def metric(scores, targets, types):
+    """
+    :param scores: the output the model predict
+    :param targets: the gt label
+    :return: OP, OR, OF1, CP, CR, CF1
+    calculate the Precision of every class by: TP/TP+FP i.e. TP/total predict
+    calculate the Recall by: TP/total GT
+    """
+    num, num_class = scores.shape
+    gt_num = np.zeros(num_class)
+    tp_num = np.zeros(num_class)
+    predict_num = np.zeros(num_class)
+
+
+    for index in range(num_class):
+        score = scores[:, index]
+        target = targets[:, index]
+        if types == 'wider':
+            tmp = np.where(target == 99)[0]
+            # score[tmp] = 0
+            target[tmp] = 0
+        
+        if types == 'voc07':
+            tmp = np.where(target != 0)[0]
+            score = score[tmp]
+            target = target[tmp]
+            neg_id = np.where(target == -1)[0]
+            target[neg_id] = 0
+
+
+        gt_num[index] = np.sum(target == 1)
+        predict_num[index] = np.sum(score >= 0.5)
+        tp_num[index] = np.sum(target * (score >= 0.5))
+
+    predict_num[predict_num == 0] = 1  # avoid dividing 0
+    OP = np.sum(tp_num) / np.sum(predict_num)
+    OR = np.sum(tp_num) / np.sum(gt_num)
+    OF1 = (2 * OP * OR) / (OP + OR)
+
+    #print(tp_num / predict_num)
+    #print(tp_num / gt_num)
+    CP = np.sum(tp_num / predict_num) / num_class
+    CR = np.sum(tp_num / gt_num) / num_class
+    CF1 = (2 * CP * CR) / (CP + CR)
+
+    return OP, OR, OF1, CP, CR, CF1

utils/evaluation/cal_mAP.py

@@ -0,0 +1,56 @@
+import os
+import numpy as np
+import torch
+import json
+
+
+def json_map(cls_id, pred_json, ann_json, types):
+    assert len(ann_json) == len(pred_json)
+    num = len(ann_json)
+    predict = np.zeros((num), dtype=np.float64)
+    target = np.zeros((num), dtype=np.float64)
+
+    for i in range(num):
+        predict[i] = pred_json[i]["scores"][cls_id]
+        target[i] = ann_json[i]["target"][cls_id]
+
+    if types == 'wider':
+        tmp = np.where(target != 99)[0]
+        predict = predict[tmp]
+        target = target[tmp]
+        num = len(tmp)
+
+    if types == 'voc07':
+        tmp = np.where(target != 0)[0]
+        predict = predict[tmp]
+        target = target[tmp]
+        neg_id = np.where(target == -1)[0]
+        target[neg_id] = 0
+        num = len(tmp)
+
+
+    tmp = np.argsort(-predict)
+    target = target[tmp]
+    predict = predict[tmp]
+
+
+    pre, obj = 0, 0
+    for i in range(num):
+        if target[i] == 1:
+            obj += 1.0
+            pre += obj / (i+1)
+    pre /= obj
+    return pre
+
+
+
+
+
+
+
+
+
+
+
+
+

utils/evaluation/eval.py

@@ -0,0 +1,64 @@
+import argparse
+import torch
+import numpy as np
+import json
+from tqdm import tqdm
+from .cal_mAP import json_map
+from .cal_PR import json_metric, metric, json_metric_top3
+
+
+voc_classes = ("aeroplane", "bicycle", "bird", "boat", "bottle",
+           "bus", "car", "cat", "chair", "cow", "diningtable",
+           "dog", "horse", "motorbike", "person", "pottedplant",
+           "sheep", "sofa", "train", "tvmonitor")
+coco_classes = ('person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
+               'train', 'truck', 'boat', 'traffic_light', 'fire_hydrant',
+               'stop_sign', 'parking_meter', 'bench', 'bird', 'cat', 'dog',
+               'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe',
+               'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
+               'skis', 'snowboard', 'sports_ball', 'kite', 'baseball_bat',
+               'baseball_glove', 'skateboard', 'surfboard', 'tennis_racket',
+               'bottle', 'wine_glass', 'cup', 'fork', 'knife', 'spoon', 'bowl',
+               'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot',
+               'hot_dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
+               'potted_plant', 'bed', 'dining_table', 'toilet', 'tv', 'laptop',
+               'mouse', 'remote', 'keyboard', 'cell_phone', 'microwave',
+               'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock',
+               'vase', 'scissors', 'teddy_bear', 'hair_drier', 'toothbrush')
+wider_classes = (
+                "Male","longHair","sunglass","Hat","Tshiirt","longSleeve","formal",
+                "shorts","jeans","longPants","skirt","faceMask", "logo","stripe")
+
+class_dict = {
+    "voc07": voc_classes,
+    "coco": coco_classes,
+    "wider": wider_classes,
+}
+
+
+
+def evaluation(result, types, ann_path):
+    print("Evaluation")
+    classes = class_dict[types]
+    aps = np.zeros(len(classes), dtype=np.float64)
+
+    ann_json = json.load(open(ann_path, "r"))
+    pred_json = result
+
+    for i, _ in enumerate(tqdm(classes)):
+        ap = json_map(i, pred_json, ann_json, types)
+        aps[i] = ap
+    OP, OR, OF1, CP, CR, CF1 = json_metric(pred_json, ann_json, len(classes), types)
+    print("mAP: {:4f}".format(np.mean(aps)))
+    print("CP: {:4f}, CR: {:4f}, CF1 :{:4F}".format(CP, CR, CF1))
+    print("OP: {:4f}, OR: {:4f}, OF1 {:4F}".format(OP, OR, OF1))
+
+# I added it here
+class WarmUpLR(torch.optim.lr_scheduler._LRScheduler):
+    def __init__(self, optimizer, total_iters, last_epoch=-1):
+        self.total_iters = total_iters
+        super().__init__(optimizer, last_epoch=last_epoch)
+
+    def get_lr(self):
+        return [base_lr * self.last_epoch / (self.total_iters + 1e-8) for base_lr in self.base_lrs]
+

utils/evaluation/warmUpLR.py

@@ -0,0 +1,11 @@
+import torch
+
+
+class WarmUpLR(torch.optim.lr_scheduler._LRScheduler):
+    def __init__(self, optimizer, total_iters, last_epoch=-1):
+        self.total_iters = total_iters
+        super().__init__(optimizer, last_epoch=last_epoch)
+
+    def get_lr(self):
+        return [base_lr * self.last_epoch / (self.total_iters + 1e-8) for base_lr in self.base_lrs]
+

utils/prepare/prepare_coco.py

@@ -0,0 +1,67 @@
+import os
+import json
+import argparse
+import numpy as np
+from pycocotools.coco import COCO
+
+
+
+def make_data(data_path=None, tag="train"):
+    annFile = os.path.join(data_path, "annotations/instances_{}2014.json".format(tag))
+    coco = COCO(annFile)
+
+    img_id = coco.getImgIds()
+    cat_id = coco.getCatIds()
+    img_id = list(sorted(img_id))
+    cat_trans = {}
+    for i in range(len(cat_id)):
+        cat_trans[cat_id[i]] = i
+
+    message = []
+
+
+    for i in img_id:
+        data = {}
+        target = [0] * 80
+        path = ""
+        img_info = coco.loadImgs(i)[0]
+        ann_ids = coco.getAnnIds(imgIds = i)
+        anns = coco.loadAnns(ann_ids)
+        if len(anns) == 0:
+            continue
+        else:
+            for i in range(len(anns)):
+                cls = anns[i]['category_id']
+                cls = cat_trans[cls]
+                target[cls] = 1
+            path = img_info['file_name']
+        data['target'] = target
+        data['img_path'] = os.path.join(os.path.join(data_path, "images/{}2014/".format(tag)), path)
+        message.append(data)
+
+    with open('data/coco/{}_coco2014.json'.format(tag), 'w') as f:
+        json.dump(message, f)
+            
+
+
+# The final json file include: train_coco2014.json & val_coco2014.json
+# which is the following format:
+# [item1, item2, item3, ......,]
+# item1 = {
+#      "target": 
+#      "img_path":      
+# }
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    # Usage: --data_path /your/dataset/path/COCO2014
+    parser.add_argument("--data_path", default="Dataset/COCO2014/", type=str, help="The absolute path of COCO2014")
+    args = parser.parse_args()
+
+    if not os.path.exists("data/coco"):
+        os.makedirs("data/coco")
+    
+    make_data(data_path=args.data_path, tag="train")
+    make_data(data_path=args.data_path, tag="val")
+
+    print("COCO data ready!")
+    print("data/coco/train_coco2014.json, data/coco/val_coco2014.json")

utils/prepare/prepare_voc.py

@@ -0,0 +1,149 @@
+import os
+import json
+import argparse
+import numpy as np
+import xml.dom.minidom as XML
+
+
+
+voc_cls_id = {"aeroplane":0, "bicycle":1, "bird":2, "boat":3, "bottle":4,
+               "bus":5, "car":6, "cat":7, "chair":8, "cow":9,
+               "diningtable":10, "dog":11, "horse":12, "motorbike":13, "person":14,
+               "pottedplant":15, "sheep":16, "sofa":17, "train":18, "tvmonitor":19}
+
+
+def get_label(data_path):
+    print("generating labels for VOC07 dataset")
+    xml_paths = os.path.join(data_path, "VOC2007/Annotations/")
+    save_dir = "data/voc07/labels"
+
+    if not os.path.exists(save_dir):
+        os.makedirs(save_dir)
+
+    for i in os.listdir(xml_paths):
+        if not i.endswith(".xml"):
+            continue
+        s_name = i.split('.')[0] + ".txt"
+        s_dir = os.path.join(save_dir, s_name)
+        xml_path = os.path.join(xml_paths, i)
+        DomTree = XML.parse(xml_path)
+        Root = DomTree.documentElement
+
+        obj_all = Root.getElementsByTagName("object")
+        leng = len(obj_all)
+        cls = []
+        difi_tag = []
+        for obj in obj_all:
+            # get the classes
+            obj_name = obj.getElementsByTagName('name')[0]
+            one_class = obj_name.childNodes[0].data
+            cls.append(voc_cls_id[one_class])
+
+            difficult = obj.getElementsByTagName('difficult')[0]
+            difi_tag.append(difficult.childNodes[0].data)
+
+        for i, c in enumerate(cls):
+            with open(s_dir, "a") as f:
+                f.writelines("%s,%s\n" % (c, difi_tag[i]))
+
+
+def transdifi(data_path):
+    print("generating final json file for VOC07 dataset")
+    label_dir = "data/voc07/labels/"
+    img_dir = os.path.join(data_path, "VOC2007/JPEGImages/")
+
+    # get trainval test id
+    id_dirs = os.path.join(data_path, "VOC2007/ImageSets/Main/")
+    f_train = open(os.path.join(id_dirs, "train.txt"), "r").readlines()
+    f_val = open(os.path.join(id_dirs, "val.txt"), "r").readlines()
+    f_trainval = f_train + f_val
+    f_test = open(os.path.join(id_dirs, "test.txt"), "r")
+
+    trainval_id =  np.sort([int(line.strip()) for line in f_trainval]).tolist()
+    test_id = [int(line.strip()) for line in f_test]
+    trainval_data = []
+    test_data = []
+
+    # ternary label
+    # -1 means negative
+    # 0 means difficult
+    # +1 means positive
+
+    # binary label
+    # 0 means negative
+    # +1 means positive 
+
+    # we use binary labels in our implementation
+
+    for item in sorted(os.listdir(label_dir)):
+        with open(os.path.join(label_dir, item), "r") as f:
+
+            target = np.array([-1] * 20)
+            classes = []
+            diffi_tag = []
+
+            for line in f.readlines():
+                cls, tag = map(int, line.strip().split(','))
+                classes.append(cls)
+                diffi_tag.append(tag)
+
+            classes = np.array(classes)
+            diffi_tag = np.array(diffi_tag)
+            for i in range(20):
+                if i in classes:
+                    i_index = np.where(classes == i)[0]
+                    if len(i_index) == 1:
+                        target[i] = 1 - diffi_tag[i_index]
+                    else:
+                        if len(i_index) == sum(diffi_tag[i_index]):
+                            target[i] = 0
+                        else:
+                            target[i] = 1
+                else:
+                    continue
+            img_path = os.path.join(img_dir, item.split('.')[0]+".jpg")
+
+            if int(item.split('.')[0]) in trainval_id:
+                target[target == -1] = 0  # from ternary to binary by treating difficult as negatives
+                data = {"target": target.tolist(), "img_path": img_path}      
+                trainval_data.append(data)
+            if int(item.split('.')[0]) in test_id:
+                data = {"target": target.tolist(), "img_path": img_path}      
+                test_data.append(data)
+
+    json.dump(trainval_data, open("data/voc07/trainval_voc07.json", "w"))
+    json.dump(test_data, open("data/voc07/test_voc07.json", "w"))
+    print("VOC07 data preparing finished!")
+    print("data/voc07/trainval_voc07.json data/voc07/test_voc07.json")
+    
+    # remove label cash
+    for item in os.listdir(label_dir):
+        os.remove(os.path.join(label_dir, item))
+    os.rmdir(label_dir)
+    
+
+# We treat difficult classes in trainval_data as negtive while ignore them in test_data
+# The ignoring operation can be automatically done during evaluation (testing).
+# The final json file include: trainval_voc07.json & test_voc07.json
+# which is the following format:
+# [item1, item2, item3, ......,]
+# item1 = {
+#      "target": 
+#      "img_path":      
+# }
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    # Usage: --data_path /your/dataset/path/VOCdevkit
+    parser.add_argument("--data_path", default="Dataset/VOCdevkit/", type=str, help="The absolute path of VOCdevkit")
+    args = parser.parse_args()
+
+    if not os.path.exists("data/voc07"):
+        os.makedirs("data/voc07")
+    
+    if 'VOCdevkit' not in args.data_path:
+        print("WARNING: please include \'VOCdevkit\' str in your args.data_path")
+        # exit()
+
+    get_label(args.data_path)
+    transdifi(args.data_path)

utils/prepare/prepare_wider.py

@@ -0,0 +1,58 @@
+import os
+import json
+import random
+import argparse
+
+
+def make_wider(tag, value, data_path):
+    img_path = os.path.join(data_path, "Image")
+    ann_path = os.path.join(data_path, "Annotations")
+    ann_file = os.path.join(ann_path, "wider_attribute_{}.json".format(tag))
+
+    data = json.load(open(ann_file, "r"))
+
+    final = []
+    image_list = data['images']
+    for image in image_list:
+        for person in image["targets"]: # iterate over each person
+            tmp = {}
+            tmp['img_path'] = os.path.join(img_path, image['file_name'])
+            tmp['bbox'] = person['bbox']
+            attr = person["attribute"]
+            for i, item in enumerate(attr):
+                if item == -1:
+                    attr[i] = 0
+                if item == 0:
+                    attr[i] = value  # pad un-specified samples
+                if item == 1:
+                    attr[i] = 1
+            tmp["target"] = attr
+            final.append(tmp)
+
+    json.dump(final, open("data/wider/{}_wider.json".format(tag), "w"))
+    print("data/wider/{}_wider.json".format(tag))
+
+
+
+# which is the following format:
+# [item1, item2, item3, ......,]
+# item1 = {
+#      "target": 
+#      "img_path":      
+# }
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--data_path", default="Dataset/WIDER_ATTRIBUTE", type=str)
+    args = parser.parse_args()
+
+    if not os.path.exists("data/wider"):
+        os.makedirs("data/wider")
+
+    # 0 (zero) means negative, we treat un-specified attribute as negative in the trainval set
+    make_wider(tag='trainval', value=0, data_path=args.data_path) 
+
+    # 99 means we ignore un-specified attribute in the test set, following previous work
+    # the number 99 can be properly identified when evaluating mAP
+    make_wider(tag='test', value=99, data_path=args.data_path)

utils/visualize.py

@@ -0,0 +1,42 @@
+from PIL import Image
+import json
+import torch
+from torchvision import transforms
+import cv2
+import numpy as np
+import os
+import torch.nn as nn
+
+def show_cam_on_img(img, mask, img_path_save):
+    heat_map = cv2.applyColorMap(np.uint8(255*mask), cv2.COLORMAP_JET)
+    heat_map = np.float32(heat_map) / 255
+
+    cam = heat_map + np.float32(img)
+    cam = cam / np.max(cam)
+    cv2.imwrite(img_path_save, np.uint8(255 * cam))
+
+    
+img_path_read = ""
+img_path_save = ""
+
+
+
+
+def main():
+  img = cv2.imread(img_path_read, flags=1)
+
+  img = np.float32(cv2.resize(img, (224, 224))) / 255
+
+  # cam_all is the score tensor of shape (B, C, H, W), similar to y_raw in out Figure 1
+  # cls_idx specifying the i-th class out of C class
+  # visualize the 0's class heatmap
+  cls_idx = 0
+  cam = cam_all[cls_idx]
+
+
+  # cam = nn.ReLU()(cam)
+  cam = cam / torch.max(cam)
+  
+  cam = cv2.resize(np.array(cam), (224, 224))
+  show_cam_on_img(img, cam, img_path_save)
+