Out-of-Distribution detection for segmentation data, using deep nearest neighbors. Presented at the ICCV23 workshop on Uncertainty Quantification for Computer Vision.
from .backbones.vit_setr import VisionTransformerSETR
from .decode_heads.vit_up_head import VisionTransformerUpHead
models/__init__.py
silviogalesso/dense-ood-knns
# Adapted from https://github.com/fudan-zvg/SETR/blob/main/mmseg/models/backbones/vit.py
# Added resampling of positional embedding

import torch
import torch.nn as nn
from functools import partial
import math
from itertools import repeat
import collections.abc as container_abcs
import warnings

from mmseg.models.builder import BACKBONES
from mmseg.ops import resize


def _cfg(url='', **kwargs):
    return {
        'url': url,
        'num_classes': 1000, 'input_size': (3, 224, 224), 'pool_size': None,
        'crop_pct': .9, 'interpolation': 'bicubic',
        'mean': (0.485, 0.456, 0.406), 'std': (0.229, 0.224, 0.225),
        'first_conv': '', 'classifier': 'head',
        **kwargs
    }


def to_2tuple(x):
    if isinstance(x, container_abcs.Iterable):
        return x
    return tuple(repeat(x, 2))


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):
        if self.drop_prob == 0. or not self.training:
            return x
        keep_prob = 1 - self.drop_prob
        # work with diff dim tensors, not just 2D ConvNets
        shape = (x.shape[0],) + (1,) * (x.ndim - 1)
        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


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)


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
        # 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
        q, k, v = self.qkv(x).reshape(B, N, 3, self.num_heads,
                                      C // self.num_heads).permute(2, 0, 3, 1, 4)

        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):
        x = self.proj(x)
        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


@BACKBONES.register_module()
class VisionTransformerSETR(nn.Module):
    """ Vision Transformer with support for patch or hybrid CNN input stage
    """

    def __init__(self, model_name='vit_large_patch16_384', img_size=384, patch_size=16, in_chans=3, embed_dim=1024, depth=24,
                 num_heads=16, num_classes=19, mlp_ratio=4., qkv_bias=True, qk_scale=None, drop_rate=0.1, attn_drop_rate=0.,
                 drop_path_rate=0., hybrid_backbone=None, norm_layer=partial(nn.LayerNorm, eps=1e-6), norm_cfg=None,
                 pos_embed_interp=False, init_cfg=None, align_corners=False, **kwargs):
        super(VisionTransformerSETR, self).__init__(**kwargs)
        self.model_name = model_name
        self.img_size = to_2tuple(img_size)
        self.patch_size = patch_size
        self.in_chans = in_chans
        self.embed_dim = embed_dim
        self.depth = depth
        self.num_heads = num_heads
        self.num_classes = num_classes
        self.mlp_ratio = mlp_ratio
        self.qkv_bias = qkv_bias
        self.qk_scale = qk_scale
        self.drop_rate = drop_rate
        self.attn_drop_rate = attn_drop_rate
        self.drop_path_rate = drop_path_rate
        self.hybrid_backbone = hybrid_backbone
        self.norm_layer = norm_layer
        self.norm_cfg = norm_cfg
        self.pos_embed_interp = pos_embed_interp
        self.init_cfg = init_cfg
        self.align_corners = align_corners

        assert init_cfg is None, "Not implemented!"

        self.num_stages = self.depth
        self.out_indices = tuple(range(self.num_stages))

        self.with_cls_token = True

        if self.hybrid_backbone is not None:
            self.patch_embed = HybridEmbed(
                self.hybrid_backbone, img_size=self.img_size, in_chans=self.in_chans, embed_dim=self.embed_dim)
        else:
            self.patch_embed = PatchEmbed(
                img_size=self.img_size, patch_size=self.patch_size, in_chans=self.in_chans, embed_dim=self.embed_dim)
        self.num_patches = self.patch_embed.num_patches

        self.cls_token = nn.Parameter(torch.zeros(1, 1, self.embed_dim))
        self.pos_embed = nn.Parameter(torch.zeros(
            1, self.num_patches + 1, self.embed_dim))
        self.pos_drop = nn.Dropout(p=self.drop_rate)

        dpr = [x.item() for x in torch.linspace(0, self.drop_path_rate,
                                                self.depth)]  # stochastic depth decay rule
        self.blocks = nn.ModuleList([
            Block(
                dim=self.embed_dim, num_heads=self.num_heads, mlp_ratio=self.mlp_ratio, qkv_bias=self.qkv_bias, qk_scale=self.qk_scale,
                drop=self.drop_rate, attn_drop=self.attn_drop_rate, drop_path=dpr[i], norm_layer=self.norm_layer)
            for i in range(self.depth)])

        # 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)

    def init_weights(self, pretrained=None):
        # nn.init.normal_(self.pos_embed, std=0.02)
        # nn.init.zeros_(self.cls_token)

        for m in self.modules():
            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)

        if not self.random_init:
            raise NotImplementedError
        else:
            print('Initialize weight randomly')

    @property
    def no_weight_decay(self):
        return {'pos_embed', 'cls_token'}

    def _conv_filter(self, 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 to_2D(self, x):
        n, hw, c = x.shape
        h = w = int(math.sqrt(hw))
        x = x.transpose(1, 2).reshape(n, c, h, w)
        return x

    def to_1D(self, x):
        n, c, h, w = x.shape
        x = x.reshape(n, c, -1).transpose(1, 2)
        return x

    # from mmseg
    @staticmethod
    def resize_pos_embed(pos_embed, input_shpae, pos_shape, mode):
        """Resize pos_embed weights.

        Resize pos_embed using bicubic interpolate method.
        Args:
            pos_embed (torch.Tensor): Position embedding weights.
            input_shpae (tuple): Tuple for (downsampled input image height,
                downsampled input image width).
            pos_shape (tuple): The resolution of downsampled origin training
                image.
            mode (str): Algorithm used for upsampling:
                ``'nearest'`` | ``'linear'`` | ``'bilinear'`` | ``'bicubic'`` |
                ``'trilinear'``. Default: ``'nearest'``
        Return:
            torch.Tensor: The resized pos_embed of shape [B, L_new, C]
        """
        assert pos_embed.ndim == 3, 'shape of pos_embed must be [B, L, C]'
        pos_h, pos_w = pos_shape
        # keep dim for easy deployment
        cls_token_weight = pos_embed[:, 0:1]
        pos_embed_weight = pos_embed[:, (-1 * pos_h * pos_w):]
        pos_embed_weight = pos_embed_weight.reshape(
            1, pos_h, pos_w, pos_embed.shape[2]).permute(0, 3, 1, 2)
        pos_embed_weight = resize(
            pos_embed_weight, size=input_shpae, align_corners=False, mode=mode)
        pos_embed_weight = torch.flatten(pos_embed_weight, 2).transpose(1, 2)
        pos_embed = torch.cat((cls_token_weight, pos_embed_weight), dim=1)
        return pos_embed

    def forward(self, x):
        B = x.shape[0]
        x = self.patch_embed(x)
        h, w = x.shape[-2:]
        x = x.flatten(2).transpose(1, 2)

        # stole cls_tokens impl from Phil Wang, thanks
        cls_tokens = self.cls_token.expand(B, -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)
        x = x + self.resize_pos_embed(self.pos_embed, (h, w), (self.img_size[0]//self.patch_size, self.img_size[1]//self.patch_size), 'bicubic')
        x = self.pos_drop(x)

        outs = []
        for i, blk in enumerate(self.blocks):
            x = blk(x)
            if i in self.out_indices:
                outs.append(x[:,1:].transpose(1, 2).reshape(B, -1, h, w))
        return tuple(outs)

models/backbones/vit_setr.py
silviogalesso/dense-ood-knns
import torch.nn as nn
import torch.nn.functional as F
from functools import partial
import math

# from .layers import trunc_normal_

from mmseg.models.builder import HEADS
from mmseg.models.decode_heads.decode_head import BaseDecodeHead

from mmcv.cnn import build_norm_layer

from itertools import repeat
from collections.abc import Iterable


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)


def to_2tuple(x):
    # if isinstance(x, container_abcs.Iterable):
    if isinstance(x, Iterable):
        return x
    return tuple(repeat(x, 2))


@HEADS.register_module()
class VisionTransformerUpHead(BaseDecodeHead):
    """ Vision Transformer with support for patch or hybrid CNN input stage
    """

    def __init__(self, img_size=768, embed_dim=1024,
                 norm_layer=partial(nn.LayerNorm, eps=1e-6), norm_cfg=None,
                 num_conv=1, upsampling_method='bilinear', num_upsampe_layer=1, conv3x3_conv1x1=True, patch_size=16, **kwargs):
        super(VisionTransformerUpHead, self).__init__(**kwargs)
        self.img_size = to_2tuple(img_size)
        self.norm_cfg = norm_cfg
        self.num_conv = num_conv
        self.norm = norm_layer(embed_dim)
        self.upsampling_method = upsampling_method
        self.num_upsampe_layer = num_upsampe_layer
        self.conv3x3_conv1x1 = conv3x3_conv1x1
        self.patch_size=patch_size

        out_channel = self.num_classes

        if self.num_conv == 2:
            if self.conv3x3_conv1x1:
                self.conv_0 = nn.Conv2d(
                    embed_dim, 256, kernel_size=3, stride=1, padding=1)
            else:
                self.conv_0 = nn.Conv2d(embed_dim, 256, 1, 1)
            self.conv_1 = nn.Conv2d(256, out_channel, 1, 1)
            _, self.syncbn_fc_0 = build_norm_layer(self.norm_cfg, 256)

        elif self.num_conv == 4:
            self.conv_0 = nn.Conv2d(
                embed_dim, 256, kernel_size=3, stride=1, padding=1)
            self.conv_1 = nn.Conv2d(
                256, 256, kernel_size=3, stride=1, padding=1)
            self.conv_2 = nn.Conv2d(
                256, 256, kernel_size=3, stride=1, padding=1)
            self.conv_3 = nn.Conv2d(
                256, 256, kernel_size=3, stride=1, padding=1)
            self.conv_4 = nn.Conv2d(256, out_channel, kernel_size=1, stride=1)

            _, self.syncbn_fc_0 = build_norm_layer(self.norm_cfg, 256)
            _, self.syncbn_fc_1 = build_norm_layer(self.norm_cfg, 256)
            _, self.syncbn_fc_2 = build_norm_layer(self.norm_cfg, 256)
            _, self.syncbn_fc_3 = build_norm_layer(self.norm_cfg, 256)

        # Segmentation head

    def init_weights(self):
        for m in self.modules():
            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 forward(self, x):
        x = self._transform_inputs(x)
        if x.dim() == 3:
            if x.shape[1] % 48 != 0:
                x = x[:, 1:]
            x = self.norm(x)

        if self.upsampling_method == 'bilinear':
            if x.dim() == 3:
                n, hw, c = x.shape
                h = self.img_size[0]//self.patch_size
                w = self.img_size[1]//self.patch_size
                x = x.transpose(1, 2).reshape(n, c, h, w)

            if self.num_conv == 2:
                if self.num_upsampe_layer == 2:
                    x = self.conv_0(x)
                    x = self.syncbn_fc_0(x)
                    x = F.relu(x, inplace=True)
                    x = F.interpolate(
                        x, size=x.shape[-1]*4, mode='bilinear', align_corners=self.align_corners)
                    x = self.conv_1(x)
                    x = F.interpolate(
                        x, size=self.img_size, mode='bilinear', align_corners=self.align_corners)
                elif self.num_upsampe_layer == 1:
                    x = self.conv_0(x)
                    x = self.syncbn_fc_0(x)
                    x = F.relu(x, inplace=True)
                    x = self.conv_1(x)
                    x = F.interpolate(
                        x, size=self.img_size, mode='bilinear', align_corners=self.align_corners)
            elif self.num_conv == 4:
                if self.num_upsampe_layer == 4:
                    x = self.conv_0(x)
                    x = self.syncbn_fc_0(x)
                    x = F.relu(x, inplace=True)
                    x = F.interpolate(
                        x, size=x.shape[-1]*2, mode='bilinear', align_corners=self.align_corners)
                    x = self.conv_1(x)
                    x = self.syncbn_fc_1(x)
                    x = F.relu(x, inplace=True)
                    x = F.interpolate(
                        x, size=x.shape[-1]*2, mode='bilinear', align_corners=self.align_corners)
                    x = self.conv_2(x)
                    x = self.syncbn_fc_2(x)
                    x = F.relu(x, inplace=True)
                    x = F.interpolate(
                        x, size=x.shape[-1]*2, mode='bilinear', align_corners=self.align_corners)
                    x = self.conv_3(x)
                    x = self.syncbn_fc_3(x)
                    x = F.relu(x, inplace=True)
                    x = self.conv_4(x)
                    x = F.interpolate(
                        x, size=x.shape[-1]*2, mode='bilinear', align_corners=self.align_corners)

        return x

models/decode_heads/vit_up_head.py
silviogalesso/dense-ood-knns
from .nearest_neighbors import get_reference_features, subsample_features, kNNEngineFAISS, kNNEngineTorch
nearest_neighbors/__init__.py
silviogalesso/dense-ood-knns
from .convnext import FeatureProbeConvNeXt
from .vit import FeatureProbeViT, FeatureProbeViTSETR


nearest_neighbors/feature_probes/__init__.py
silviogalesso/dense-ood-knns
from .utils import FeatureProbe


class FeatureProbeConvNeXt(FeatureProbe):
    @staticmethod
    def _register_hooks(model, stages=(2,)):
        feat_dict = dict()
        for stage in stages:
            FeatureProbe._register_output_activation_hook(model.module.backbone.stages[stage],
                                                   lambda ft: ft[0].permute(1, 2, 0),
                                                   feat_dict,
                                                   f"stage{stage}")
        return feat_dict

    @staticmethod
    def _features_post_process(ft_dict):
        return ft_dict

nearest_neighbors/feature_probes/convnext.py
silviogalesso/dense-ood-knns
import torch

class FeatureProbe:
    def __init__(self, model, **kwargs):
        # hooks, pre-post fns, get features
        self.activation_dict = self._register_hooks(model, **kwargs)

    def get_features(self):
        return self._features_post_process(self.activation_dict)

    @staticmethod
    def _register_input_activation_hook(module, feature_extraction_fn, output_dict, key):
        def activation_hook(self, x, y):
            output_dict[key] = feature_extraction_fn(x)
        module.register_forward_hook(activation_hook)
    
    @staticmethod
    def _register_output_activation_hook(module, feature_extraction_fn, output_dict, key):
        def activation_hook(self, x, y):
            output_dict[key] = feature_extraction_fn(y)
        module.register_forward_hook(activation_hook)

    @staticmethod
    def _register_hooks(model, **kwargs):
        raise NotImplementedError

    @staticmethod
    def _features_post_process(ft_dict):
        raise NotImplementedError


nearest_neighbors/feature_probes/utils.py
silviogalesso/dense-ood-knns
import torch
from .utils import FeatureProbe


class FeatureProbeViT(FeatureProbe):
    @staticmethod
    def _register_hooks(model, layer_index=-1):
        # model.module.backbone.layers[layer_index].attn.attn = MultiHeadAttentionExplicitQKV(
        #     model.module.backbone.layers[layer_index].attn.attn)

        selected_layer = model.module.backbone.layers[layer_index]
        assert selected_layer.attn.attn.in_proj_weight is not None and selected_layer.attn.attn.q_proj_weight is None
        num_extra_tokens = model.module.backbone.num_extra_tokens if hasattr(model.module.backbone, "num_extra_tokens") else 1

        def _extract_keys(x):
            # that's how mmseg vit computes attention...
            q, k, v = torch.nn.functional._in_projection_packed(x[0], x[0], x[0], selected_layer.attn.attn.in_proj_weight, selected_layer.attn.attn.in_proj_bias)
            return k[:,num_extra_tokens:].permute(1, 0, 2)

        feat_dict = dict()
        FeatureProbe._register_output_activation_hook(model.module.backbone.patch_embed,
                                                      lambda x: x[1],  # x[0] is the fwd pass tensor, we only need the shape
                                                      feat_dict,
                                                      "hw_shape")
        
        FeatureProbe._register_input_activation_hook(selected_layer.attn,
                                                     _extract_keys,
                                                     feat_dict,
                                                     "k")
        return feat_dict

    @staticmethod
    def _features_post_process(ft_dict):
        for k, ft in zip("qkv", ft_dict.pop("qkv", [])):
            ft_dict[k] = ft
        hw_shape = ft_dict.pop("hw_shape")
        return {k: ft.reshape(*hw_shape, -1) for k, ft in ft_dict.items()}


class FeatureProbeViTSETR(FeatureProbe):
    @staticmethod
    def _register_hooks(model, layer_index=-1, eval_all_features=False):
        assert not eval_all_features

        feat_dict = dict()

        num_extra_tokens = 0            
        
        FeatureProbe._register_output_activation_hook(model.module.backbone.patch_embed,
                                               lambda x: x.shape,
                                               feat_dict,
                                               "input_shape")
        FeatureProbe._register_output_activation_hook(model.module.backbone.blocks[layer_index].attn.qkv,
                                               lambda x: x, 
                                               feat_dict,
                                               "qkv")
        return feat_dict
    
    @staticmethod
    def _features_post_process(ft_dict):
        B, C, H, W = ft_dict.pop("input_shape")
        qkv = ft_dict["qkv"] # B, N, 3*C
        k = qkv.view(H*W+1, 3, C)[1:,1].view(H, W, C)
        return {"k": k}
nearest_neighbors/feature_probes/vit.py
silviogalesso/dense-ood-knns
import os
from time import time
from tqdm import tqdm
import numpy as np
import random
import logging
from mmseg.datasets import build_dataloader
# from .subsampling import greedy_coreset_subsampling, per_class_greedy_coreset_subsampling
from .utils import *
try:
    import faiss
except ModuleNotFoundError:
    pass
import torch


def get_reference_features(model, dataset, feature_probe, num_samples, ood_scoring_fn_for_scaling):
    ft_list = []
    avg_valids = []
    coordinates = []
    labels_list = []
    max_param_score = -np.inf; min_param_score = np.inf

    indices = random.sample(range(len(dataset)), num_samples)
    subset = torch.utils.data.Subset(dataset, indices)
    loader_cfg = dict(
        dist=False,
        shuffle=False,
        samples_per_gpu=1,
        workers_per_gpu=1
    )
    loader = build_dataloader(subset, **loader_cfg)
    model.eval()
    for smpl_i, (index, sample) in tqdm(enumerate(zip(indices, loader))):
        smpl_ft_list = []
        smpl_coords = []
        smpl_labels = []
        logits = fwd_pass_mmseg(sample, model)
        max_param_score = max(max_param_score, ood_scoring_fn_for_scaling(logits).max().item())
        min_param_score = min(min_param_score, ood_scoring_fn_for_scaling(logits).min().item())
        proc_act_dict = feature_probe.get_features()
        features = [proc_act_dict[k] for k in proc_act_dict]
        segm = torch.tensor(dataset.get_gt_seg_map_by_idx(index))
        for ft in features:
            h, w, c = ft.shape
            downsampled_segm = torch.nn.functional.interpolate(segm.view((1, 1, *segm.shape)).byte(), size=(h, w))
            valid = (downsampled_segm.squeeze() != 255).flatten()
            smpl_ft_list.append(ft.view(h * w, c).cpu()[valid])
            smpl_coords.append(np.stack(np.meshgrid(range(w), range(h))+[(np.ones((h, w))*smpl_i)], -1).reshape(h*w, 3)[valid])
            smpl_labels.append(downsampled_segm.flatten()[valid])
        ft_list.append(smpl_ft_list)
        coordinates.append(smpl_coords)
        avg_valids.append(valid.float().mean())
        labels_list.append(smpl_labels)

    ft_lists = zip(*ft_list)  # from list of feature tuples, to lists of features (n_ft, n_imgs, samples_per_img)
    lb_lists = zip(*labels_list)  # from list of label tuples, to lists of labels (n_ft, n_imgs, samples_per_img)
    ft_dict = {k: torch.cat(fl, dim=0) for k, fl in zip(proc_act_dict.keys(), ft_lists)}  # concatenate for each ft type
    lb_dict = {k: torch.cat(ll, dim=0) for k, ll in zip(proc_act_dict.keys(), lb_lists)}  # concatenate for each ft type

    return ft_dict, lb_dict, ((min_param_score), (max_param_score))


def subsample_features(features_dict, labels_dict, subsample_number, mode="random"):
    output_features = dict()
    output_labels = dict()
    for k, ft in features_dict.items():
        assert len(ft) == len(labels_dict[k])
        assert subsample_number <= len(ft), f"Subsample number ({subsample_number}) is larger than the number of features ({len(ft)})"
        if mode == 'random':
            ids = random.sample(range(len(ft)), subsample_number)
            output_features[k] = ft[ids]
            output_labels[k] = labels_dict[k][ids]
        else:
            raise NotImplementedError("Mode: "+mode)

    return output_features, output_labels



class kNNEngine:
    def __init__(self, reference_features, num_neighbors, dist_type, reduction='mean'):
        self.num_neighbors = num_neighbors
        self.dist_type = dist_type
        self.reduction = reduction
        assert len(reference_features.shape) == 2, "Features should be in (N, c) shape"
        self.reference_features = reference_features
        self.search_runtimes = []
    
    def compute_nn_scores(self, test_features):
        raise NotImplementedError

    def compute_min_max_nndist(self, reference_features, max_samples=30000):
        raise NotImplementedError
    
    def reduce(self, nn_dists):
        if self.reduction == "max":
            return nn_dists.max(-1)[0]
        elif self.reduction == "mean":
            return nn_dists.mean(-1)
        else:
            raise ValueError

    def __repr__(self):
        return f"features: {self.reference_features.shape[0]}x{self.reference_features.shape[1]}, k={self.num_neighbors}, dist={self.dist_type}"


class kNNEngineTorch(kNNEngine):
    def __init__(self, reference_features, num_neighbors, dist_type, reduction='mean'):
        super().__init__(reference_features, num_neighbors, dist_type, reduction)
        self.reference_features = self.reference_features.cuda()

    def _compute_dists(self, ft1, ft2):
        assert len(ft1.shape) == len(ft2.shape) == 2, "Features should be in (N, c) shape"
        if self.dist_type in ["euclidean", "l2"]:
            dists = torch.cdist(ft1.unsqueeze(0), ft2.unsqueeze(0))[0]
        elif self.dist_type == "l1":
            dists = torch.cdist(ft1.unsqueeze(0), ft2.unsqueeze(0), p=1)[0]
        return dists

    def compute_min_max_nndist(self, reference_features, max_samples=30000):
        if len(reference_features) > max_samples:
            reference_features = reference_features[random.sample(range(len(reference_features)), max_samples)]
        reference_features = reference_features.cuda()
        dists = self._compute_dists(reference_features, reference_features)
        nn_dists = torch.topk(dists, self.num_neighbors + 1, dim=-1, largest=False, sorted=True).values[:,1:]
        nn_dists = self.reduce(nn_dists)
        return nn_dists.min().item(), nn_dists.max().item()

    def compute_nn_scores(self, test_features):
        assert len(test_features.shape) == 3, "Test features should be in (h, w, c) shape"
        t0 = time()
        dists = self._compute_dists(test_features.flatten(0, 1).cuda(), self.reference_features)
        nn_dists = torch.topk(dists, self.num_neighbors, dim=-1, largest=False).values
        nn_dists = self.reduce(nn_dists).cpu()
        torch.cuda.synchronize()
        self.search_runtimes.append((time() - t0))
        return nn_dists.view(test_features.shape[:2])

    def __repr__(self):
        return "kNNEngineNaive "+super().__repr__()
    

class kNNEngineFAISS(kNNEngine):
    def __init__(self, reference_features, num_neighbors, dist_type, index_string, reduction='mean'):
        super().__init__(reference_features, num_neighbors, dist_type, reduction)

        res = faiss.StandardGpuResources()
        d = self.reference_features.shape[-1]

        # extract nprobe if in string
        if "/" in index_string:
            nprobe = int(index_string.split("/")[-1])
            index_string = index_string.split("/")[0]
        else:
            nprobe = None

        if dist_type in ["l2", "euclidean"]:
            self.index = faiss.index_factory(d, index_string)
        elif dist_type == "l1":
            assert index_string == "Flat"
            self.index = faiss.IndexFlat(d, faiss.METRIC_L1)
        else:
            raise ValueError
    
        self.index = faiss.index_cpu_to_gpu(res, 0, self.index)
        if index_string != "Flat":
            self.index.train(self.reference_features.numpy())
        self.index.add(self.reference_features.contiguous().numpy())
        assert self.index.is_trained
        if nprobe:
            self.index.nprobe = nprobe

    def compute_min_max_nndist(self, reference_features, max_samples=30000):
        if len(reference_features) > max_samples:
            reference_features = reference_features[random.sample(range(len(reference_features)), max_samples)]
        min_max_index = faiss.index_cpu_to_gpu(faiss.StandardGpuResources(), 0, faiss.IndexFlatL2(reference_features.shape[-1]))
        min_max_index.add(reference_features.numpy())
        nn_dists = np.sqrt(min_max_index.search(reference_features.numpy(), self.num_neighbors+1)[0][:,1:])
        nn_dists = self.reduce(nn_dists)
        return nn_dists.min().item(), nn_dists.max().item()
    
    def compute_nn_scores(self, test_features):
        assert len(test_features.shape) == 3, "Test features should be in (h, w, c) shape"
        test_features = test_features
        t0 = time()
        nn_dists, _ = self.index.search(test_features.flatten(0, 1).cpu().numpy(), self.num_neighbors)
        self.search_runtimes.append((time() - t0))
        nn_dists = self.reduce(torch.tensor(np.sqrt(nn_dists)))
        return nn_dists.view(test_features.shape[:2])
    
    def __repr__(self):
        return "kNNEngineFAISS "+super().__repr__()

nearest_neighbors/nearest_neighbors.py
silviogalesso/dense-ood-knns
import mmcv
import warnings
import torch


def fwd_pass_mmseg(data_sample, model):
    image = data_sample["img"]
    if isinstance(image, mmcv.parallel.data_container.DataContainer):
        image = image.data
    assert isinstance(image, list) and len(image) == 1, "Should be single image in a list... " + str(type(image)) + str(
        len(image))
    image = image[0]
    if len(image.shape) == 3:
        image = image.unsqueeze(0)
    imgh, imgw = image.shape[2:]

    if isinstance(model, torch.nn.DataParallel):
        model = model.module

    with torch.no_grad(), warnings.catch_warnings():
        warnings.filterwarnings("ignore", category=UserWarning)
        logits = model.whole_inference(image.cuda(), data_sample["img_metas"], rescale=False).cpu()

    return logits


def upsample_scores(scores, reference):
    return torch.nn.functional.interpolate(scores[None, None, ...], size=reference.shape).squeeze()


def combine_scores(max_logit, nn_dist, min_ml, max_ml, min_nn, max_nn):
    ml_ = ((max_logit - min_ml) / (max_ml - min_ml))
    nn_ = ((nn_dist-min_nn)/(max_nn-min_nn))
    return ml_+nn_


nearest_neighbors/utils.py
silviogalesso/dense-ood-knns
import os
ds_root = os.environ["DATASETS_ROOT"]
data_root = os.path.join(ds_root, "Cityscapes")

dataset_type = 'CityscapesDataset'

img_norm_cfg = dict(
    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
crop_size = (512, 1024)
train_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='LoadAnnotations'),
    dict(type='Resize', img_scale=(2048, 1024), ratio_range=(0.5, 2.0)),
    dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
    dict(type='RandomFlip', prob=0.5),
    dict(type='PhotoMetricDistortion'),
    dict(type='Normalize', **img_norm_cfg),
    dict(type='Pad', size=crop_size, pad_val=0, seg_pad_val=255),
    dict(type='DefaultFormatBundle'),
    dict(type='Collect', keys=['img', 'gt_semantic_seg']),
]
test_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(
        type='MultiScaleFlipAug',
        img_scale=(2048, 1024),
        # img_ratios=[0.5, 0.75, 1.0, 1.25, 1.5, 1.75],
        flip=False,
        transforms=[
            dict(type='Resize', keep_ratio=True),
            dict(type='RandomFlip'),
            dict(type='Normalize', **img_norm_cfg),
            dict(type='ImageToTensor', keys=['img']),
            dict(type='Collect', keys=['img']),
        ])
]
data = dict(
    samples_per_gpu=2,
    workers_per_gpu=2,
    train=dict(
        type=dataset_type,
        data_root=data_root,
        img_dir='leftImg8bit/train',
        ann_dir='gtFine/train',
        pipeline=train_pipeline),
    val=dict(
        type=dataset_type,
        data_root=data_root,
        img_dir='leftImg8bit/val',
        ann_dir='gtFine/val',
        pipeline=test_pipeline),
    test=dict(
        type=dataset_type,
        data_root=data_root,
        img_dir='leftImg8bit/val',
        ann_dir='gtFine/val',
        pipeline=test_pipeline))

nn_ood_configs/datasets/cityscapes.py
silviogalesso/dense-ood-knns
import os
ds_root = os.environ["DATASETS_ROOT"]
data_root = os.path.join(ds_root, 'RoadAnomaly/RoadAnomaly_jpg')

img_norm_cfg = dict(
    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True
)
max_ratio = 2

test_pipeline = [
    dict(type="LoadImageFromFile"),
    dict(
        type="MultiScaleFlipAug",
        img_scale=(1280, 720),
        flip=False,
        transforms=[
            dict(type="Resize", keep_ratio=True),
            # dict(type="RandomFlip"),
            dict(type="Normalize", **img_norm_cfg),
            dict(type='ImageToTensor', keys=['img']),
            dict(type="Collect", keys=["img"]),
        ],
    ),
]
data = dict(
    samples_per_gpu=1,
    workers_per_gpu=1,
    test=dict(
        type='CustomDataset',
        data_root=data_root,
        img_dir="frames",
        ann_dir="frames",
        img_suffix='.jpg',
        seg_map_suffix='.labels/labels_semantic.png',
        pipeline=test_pipeline,
    ),
)

ood_index = 2

nn_ood_configs/datasets/roadanomaly.py
silviogalesso/dense-ood-knns
_base_ = [
    '../datasets/cityscapes.py',
]

norm_cfg = dict(type='SyncBN', requires_grad=True)
model = dict(
    type='EncoderDecoder',
    backbone=dict(
        type='VisionTransformerSETR',
        model_name='vit_base_patch16_384',
        img_size=768,
        patch_size=16,
        in_chans=3,
        embed_dim=768,
        depth=12,
        num_heads=12,
        num_classes=19,
        drop_rate=0.,
        norm_cfg=norm_cfg,
        pos_embed_interp=True,
        align_corners=False,
    ),
    decode_head=dict(
        type='VisionTransformerUpHead',
        in_channels=768,
        channels=512,
        in_index=11,
        img_size=768,
        embed_dim=768,
        num_classes=19,
        conv3x3_conv1x1=False,
        norm_cfg=norm_cfg,
        num_conv=2,
        upsampling_method='bilinear',
        align_corners=False,
        ),
    auxiliary_head=[
        dict(
            type='VisionTransformerUpHead',
            in_channels=1024,
            channels=512,
            in_index=2,
            img_size=768,
            embed_dim=768,
            num_classes=19,
            norm_cfg=norm_cfg,
            num_conv=2,
            upsampling_method='bilinear',
            align_corners=False,
            conv3x3_conv1x1=False,
        ),
        dict(
            type='VisionTransformerUpHead',
            in_channels=1024,
            channels=512,
            in_index=5,
            img_size=768,
            embed_dim=768,
            num_classes=19,
            norm_cfg=norm_cfg,
            num_conv=2,
            upsampling_method='bilinear',
            align_corners=False,
            conv3x3_conv1x1=False,
        ),
        dict(
            type='VisionTransformerUpHead',
            in_channels=1024,
            channels=512,
            in_index=8,
            img_size=768,
            embed_dim=768,
            num_classes=19,
            norm_cfg=norm_cfg,
            num_conv=2,
            upsampling_method='bilinear',
            align_corners=False,
            conv3x3_conv1x1=False,
        ),
    ]
)


nn_ood_configs/models/SETR_Naive_ViT-B_cityscapes.py
silviogalesso/dense-ood-knns
_base_ = [
    '../datasets/cityscapes.py',
]

norm_cfg = dict(type='SyncBN', requires_grad=True)
model = dict(
    type='EncoderDecoder',
    backbone=dict(
        type='VisionTransformerSETR',
        model_name='vit_large_patch16_384',
        img_size=768,
        patch_size=16,
        in_chans=3,
        embed_dim=1024,
        depth=24,
        num_heads=16,
        num_classes=19,
        drop_rate=0.1,
        norm_cfg=norm_cfg,
        pos_embed_interp=True,
        align_corners=False,
    ),
    decode_head=dict(
        type='VisionTransformerUpHead',
        in_channels=1024,
        channels=512,
        in_index=23,
        img_size=768,
        embed_dim=1024,
        num_classes=19,
        conv3x3_conv1x1=False,
        norm_cfg=norm_cfg,
        num_conv=2,
        upsampling_method='bilinear',
        align_corners=False,
        ),
    auxiliary_head=[
        dict(
            type='VisionTransformerUpHead',
            in_channels=1024,
            channels=512,
            in_index=2,
            img_size=768,
            embed_dim=1024,
            num_classes=19,
            norm_cfg=norm_cfg,
            num_conv=2,
            upsampling_method='bilinear',
            align_corners=False,
            conv3x3_conv1x1=False,
        ),
        dict(
            type='VisionTransformerUpHead',
            in_channels=1024,
            channels=512,
            in_index=5,
            img_size=768,
            embed_dim=1024,
            num_classes=19,
            norm_cfg=norm_cfg,
            num_conv=2,
            upsampling_method='bilinear',
            align_corners=False,
            conv3x3_conv1x1=False,
        ),
        dict(
            type='VisionTransformerUpHead',
            in_channels=1024,
            channels=512,
            in_index=8,
            img_size=768,
            embed_dim=1024,
            num_classes=19,
            norm_cfg=norm_cfg,
            num_conv=2,
            upsampling_method='bilinear',
            align_corners=False,
            conv3x3_conv1x1=False,
        ),
    ]
)


nn_ood_configs/models/SETR_Naive_ViT-L_cityscapes.py
silviogalesso/dense-ood-knns
_base_ = '../datasets/cityscapes.py'

backbone_norm_cfg = dict(type='LN', eps=1e-6, requires_grad=True)
model = dict(
    type='EncoderDecoder',
    backbone=dict(
        type='VisionTransformer',
        img_size=(768, 768),
        patch_size=16,
        in_channels=3,
        embed_dims=768,
        num_layers=12,
        num_heads=12,
        drop_path_rate=0.1,
        attn_drop_rate=0.0,
        drop_rate=0.0,
        final_norm=True,
        norm_cfg=backbone_norm_cfg,
        with_cls_token=True,
        interpolate_mode='bicubic',
    ),
    decode_head=dict(
        type='FCNHead',
        in_channels=768,
        channels=768,
        num_convs=0,
        dropout_ratio=0.0,
        concat_input=False,
        num_classes=19,
        loss_decode=dict(
            type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0)
    ),
    test_cfg=dict(mode='slide', crop_size=(512, 512), stride=(480, 480)),
)

nn_ood_configs/models/segmenter_deit-b_linear_cityscapes.py
silviogalesso/dense-ood-knns
_base_ = '../datasets/cityscapes.py'

norm_cfg = dict(type='SyncBN', requires_grad=True)
custom_imports = dict(imports='mmcls.models', allow_failed_imports=False)
model = dict(
    type='EncoderDecoder',
    pretrained=None,
    backbone=dict(
        type='mmcls.ConvNeXt',
        arch='small',
        out_indices=[0, 1, 2, 3],
        drop_path_rate=0.3,
        layer_scale_init_value=1.0,
        gap_before_final_norm=False),
    decode_head=dict(
        type='UPerHead',
        in_channels=[96, 192, 384, 768],
        in_index=[0, 1, 2, 3],
        pool_scales=(1, 2, 3, 6),
        channels=512,
        dropout_ratio=0.1,
        num_classes=19,
        norm_cfg=norm_cfg,
        align_corners=False,
        loss_decode=dict(
            type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0)),
    auxiliary_head=dict(
        type='FCNHead',
        in_channels=384,
        in_index=2,
        channels=256,
        num_convs=1,
        concat_input=False,
        dropout_ratio=0.1,
        num_classes=19,
        norm_cfg=norm_cfg,
        align_corners=False,
        loss_decode=dict(
            type='CrossEntropyLoss', use_sigmoid=False, loss_weight=0.4)),
    # model training and testing settings
    train_cfg=dict(),
    test_cfg=dict(mode='whole'))

nn_ood_configs/models/upernet_convnext_small_cityscapes.py
silviogalesso/dense-ood-knns
# import faiss
import os
import argparse
import random
import torch
import numpy as np
from tqdm import tqdm
import warnings
from pprint import pprint

import mmcv
from mmcv import DictAction
from mmcv.cnn.utils import revert_sync_batchnorm
from mmseg.models import build_segmentor
from mmseg.datasets import build_dataset
from mmseg.utils import build_dp
import models

# from nearest_neighbors import get_reference_features, compute_nn_scores, subsample_features, compute_min_max_nndist
from nearest_neighbors import get_reference_features, subsample_features, kNNEngineTorch  #, kNNEngineFAISS
from nearest_neighbors.utils import upsample_scores, combine_scores
from nearest_neighbors.feature_probes import FeatureProbeConvNeXt, FeatureProbeViT, FeatureProbeViTSETR

from utils import StreamingEval, ood_score_functions, load_checkpoint

def parse_args():
    parser = argparse.ArgumentParser(
        description='Eval a model for out-of-distribution segmentation')

    parser.add_argument('config', help='test config file path, for the model and for the reference data if --reference-data-config is not set')
    parser.add_argument('checkpoint', help='checkpoint file')
    parser.add_argument('test_data_config', help="config file for the test dataset")

    parser.add_argument('--reference-data-config', default=None, help='Config file for the reference dataset. If not provided, the main/model config is used.')
    parser.add_argument('--test-data-cfg-split', type=str, default='test')
    parser.add_argument('--nn-k', type=int, default=3, help='number of nearest neighbors')
    parser.add_argument('--nn-reference-samples', type=int, default=500, help='number of NN reference samples')
    parser.add_argument('--nn-subsample-mode', type=str, default="random", choices=["random"])
    parser.add_argument('--nn-subsample-number', type=int, default=100000)
    parser.add_argument('--nn-dist-type', default="euclidean")
    parser.add_argument('--parametric-score-fn', default="logsumexp")
    # parser.add_argument('--faiss-index', default="Flat", type=str, help="the string defining the faiss index type")
    parser.add_argument('--random-seed', default=0, type=int)
    parser.add_argument('--cfg-options', nargs='+', action=DictAction, help='override config setting, see mmcv/mmseg')
    args = parser.parse_args()

    random.seed(args.random_seed)
    np.random.seed(args.random_seed)

    return args


def main():
    args = parse_args()

    print("checkpoint:", args.checkpoint)
    print("config:", args.config)

    cfg = mmcv.Config.fromfile(args.config)
    if args.cfg_options is not None:
        cfg.merge_from_dict(args.cfg_options)

    model = build_segmentor(cfg.model).eval().cuda()
    load_checkpoint(model, args.checkpoint)
    model = revert_sync_batchnorm(model)
    model = build_dp(model, 'cuda')

    ## DATASETS
    # test data
    test_data_config = mmcv.Config.fromfile(args.test_data_config)
    ood_idx = test_data_config.ood_index
    test_data_config = getattr(test_data_config.data, args.test_data_cfg_split)
    test_dataset = build_dataset(test_data_config)
    test_dataset.ood_idx = ood_idx

    # reference data
    reference_data_config = cfg.data.train if not args.reference_data_config else mmcv.Config.fromfile(args.reference_data_config).data.train
    if isinstance(reference_data_config, list): reference_data_config = reference_data_config[0]
    reference_data_config.pipeline = test_data_config.pipeline
    reference_dataset = build_dataset(reference_data_config)

    ## NNs
    if cfg.model.backbone.type.startswith("ResNet"):
        raise NotImplementedError
    elif "ConvNeXt" in cfg.model.backbone.type:
        feature_probe = FeatureProbeConvNeXt(model)
    elif cfg.model.backbone.type in ["MixVisionTransformer"]:
        raise NotImplementedError
    elif cfg.model.backbone.type in ["VisionTransformer", "DistilledVisionTransformer", "SETR"]:
        feature_probe = FeatureProbeViT(model)
    elif cfg.model.backbone.type == "VisionTransformerSETR":
        feature_probe = FeatureProbeViTSETR(model)
    else:
        raise ValueError(cfg.model.backbone.type)
    print('Registered activation hook(s)')

    # gather reference features dataset
    print(f"Getting reference features")
    reference_features, reference_labels, (min_score, max_score) = get_reference_features(
        model, reference_dataset, feature_probe,
        num_samples=args.nn_reference_samples,
        ood_scoring_fn_for_scaling=ood_score_functions[args.parametric_score_fn]
    )
    print(f"Extracted features from {args.nn_reference_samples} images")
    
    # subsample features
    if args.nn_subsample_number > 0:
        subsampled_features, subsampled_labels = subsample_features(
            reference_features, reference_labels,
            args.nn_subsample_number,
            args.nn_subsample_mode,
        )
    else:
        subsampled_features = reference_features
    [print(f"Subsampled '{k}' features to {len(ft)}") for k, ft in subsampled_features.items()]

    # create kNN engines, get nn dist stats
    knnengines = {k: kNNEngineTorch(ref_ft, args.nn_k, args.nn_dist_type, reduction="mean") for k, ref_ft in subsampled_features.items()}
    min_max_nndist = {k: knnengines[k].compute_min_max_nndist(ft) for k, ft in reference_features.items()}

    for k, eng in knnengines.items():
        print(f"features '{k}', engine: {eng}")

    eval_ood(model, test_dataset, feature_probe,
             knn_engines=knnengines,
             score_fn=args.parametric_score_fn,
             scores_min_max=(min_score, max_score),
             nn_min_max=min_max_nndist,
             )


def eval_ood(model, test_dataset, feature_probe, knn_engines, score_fn,
             scores_min_max, nn_min_max):
    model.eval()

    ood_eval_param = StreamingEval(test_dataset.ood_idx, ignore_ids=255)
    ood_eval_dnp = {k: StreamingEval(test_dataset.ood_idx, ignore_ids=255) for k in knn_engines.keys()}
    ood_eval_cdnp = {k: StreamingEval(test_dataset.ood_idx, ignore_ids=255) for k in knn_engines.keys()}

    progress = tqdm(total=len(test_dataset))
    for sample_index, sample in enumerate(test_dataset):
        segm = test_dataset.get_gt_seg_map_by_idx(sample_index)
        with torch.no_grad(), warnings.catch_warnings():
            warnings.filterwarnings("ignore", category=UserWarning)
            logits = model.module.whole_inference(sample["img"][0].unsqueeze(0).cuda(), sample["img_metas"],
                                                  rescale=False).cpu()
        param_scores = {k: ood_score_functions[k](logits) for k in ood_score_functions}
        ood_eval_param.add(param_scores[score_fn], segm)
        
        # DNP, cDNP
        test_features = feature_probe.get_features()
        for ft_k, ft in test_features.items():
            nn_scores = knn_engines[ft_k].compute_nn_scores(ft)
            nn_scores = upsample_scores(nn_scores, segm)
            ood_eval_dnp[ft_k].add(nn_scores, segm)

            comb_scores = combine_scores(param_scores[score_fn], nn_scores, *scores_min_max, 0, nn_min_max[ft_k][1])  # min dist is 0
            ood_eval_cdnp[ft_k].add(comb_scores, segm)

        progress.update()
    progress.close()

    print(f"Max GPU memory allocated: {torch.cuda.max_memory_allocated()/1e9:.2f}GB")
    for k, e in knn_engines.items():
        if e.search_runtimes:
            print(f"Avg search time {k} = {np.mean(e.search_runtimes):03f}")

    print("Results:")
    print("parametric:  AP = {:.02f}, FPR95 = {:.02f}".format(*ood_eval_param.get_results()[1:]))
    for k in ood_eval_dnp.keys():
        print(f"ft {k}")
        print("       DNP:  AP = {:.02f}, FPR95 = {:.02f}".format(*ood_eval_dnp[k].get_results()[1:]))
        print("      cDNP:  AP = {:.02f}, FPR95 = {:.02f}".format(*ood_eval_cdnp[k].get_results()[1:]))


if __name__ == '__main__':
    main()

test_ood.py
silviogalesso/dense-ood-knns
from .loading import load_checkpoint
from .evaluation import StreamingEval, ood_score_functions
utils/__init__.py
silviogalesso/dense-ood-knns
from sklearn.metrics import roc_curve, precision_recall_curve, roc_auc_score, average_precision_score
import torch
import numpy as np


def fpr_at_tpr(scores, ood_gt, tpr=0.95):
    fprs, tprs, _ = roc_curve(ood_gt, scores)
    idx = np.argmin(np.abs(np.array(tprs) - tpr))
    return fprs[idx]


def f_max_score(scores, ood_gt):
    precision, recall, thresholds = precision_recall_curve(ood_gt, scores)
    f1_scores = 2 * recall * precision / (recall + precision + 1e-6)
    f_max = np.max(f1_scores)
    return f_max


def ap(scores, ood_gt):
    return average_precision_score(ood_gt, scores)


def auroc(scores, target):
    return roc_auc_score(target, scores)


class StreamingEval:
    def __init__(self, ood_id, ignore_ids):
        self.ood_id = ood_id
        self.collected_scores = []
        self.collected_gts = []
        if isinstance(ignore_ids, int):
            ignore_ids = [ignore_ids]
        self.ignore_ids = ignore_ids

    def add(self, scores, segm_gt):
        valid = np.logical_not(np.in1d(segm_gt, self.ignore_ids))
        ood_gt = (segm_gt == self.ood_id)
        if not isinstance(scores, torch.Tensor):
            scores = torch.tensor(scores)
        if not isinstance(ood_gt, torch.Tensor):
            ood_gt = torch.tensor(ood_gt)
        self.collected_scores.append(scores.cpu().flatten()[valid])
        self.collected_gts.append(ood_gt.cpu().flatten()[valid])

    def get_scores_and_labels(self):
        all_scores = torch.cat(self.collected_scores, 0)
        all_gts = torch.cat(self.collected_gts, 0)
        return all_scores, all_gts

    def get_results(self):
        scores, labels = self.get_scores_and_labels()
        return auroc(scores, labels.int()) * 100, ap(scores, labels.int()) * 100, fpr_at_tpr(scores, labels.int()) * 100


def prediction_entropy(probs, eps=1e-6):
    assert len(probs.shape) == 4, "Required: probs (b, c, h, w), got shape {} instead".format(probs.shape)
    h = -(probs * (probs + eps).log()).sum(1)
    return h


ood_score_functions = {
    "max_softmax": lambda l: 1 - l.softmax(dim=1).max(dim=1)[0].detach(),
    "entropy": lambda l: prediction_entropy(l.softmax(1)).detach(),
    "logsumexp": lambda l: 1-l.logsumexp(dim=1).detach(),
    "max_logit": lambda l: - l.max(dim=1)[0].detach()
}

utils/evaluation.py
silviogalesso/dense-ood-knns
import torch

def load_checkpoint(model, path):
    ckpt = torch.load(path)
    model_state_dict = ckpt["state_dict"] if "state_dict" in ckpt else ckpt
    if "backbone.dist_token" in model_state_dict:
        # Distilled ViT (DEIT), discard distillation token
        del model_state_dict["backbone.dist_token"]
        if model_state_dict["backbone.pos_embed"].shape[1] == model.backbone.pos_embed.shape[1]+1:
            # discard positional embedding for distillation token
            model_state_dict["backbone.pos_embed"] = torch.cat([model_state_dict["backbone.pos_embed"][:,:1], 
                                                                model_state_dict["backbone.pos_embed"][:,2:]], dim=1)
    model.load_state_dict(model_state_dict)
utils/loading.py
silviogalesso/dense-ood-knns