model/model.py at main · ASAIAA-Project/model · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
import torch
import torch.nn as nn
import torchvision.models as models
import random

from einops import repeat
from torchvision.models import feature_extraction
from kornia import filters
from collections import OrderedDict


def create_ASAIAANet(model_config):
    regressor = Regressor(model_config.backbone_type, model_config.pretrained,
                          model_config.weight_path)
    extractor_backbone = Regressor(model_config.backbone_type,
                                   model_config.pretrained,
                                   model_config.weight_path).backbone
    extractor_backbone.requires_grad_(False)
    feature_extractor = feature_extraction.create_feature_extractor(
        extractor_backbone, model_config.feature_target_layer)
    distractor = Distractor(
        feature_extractor, Finializer(model_config.center_bias_weight),
        ReadoutNet(model_config.feature_channels_num, model_config.feature_h,
                   model_config.feature_w))

    return ASAIAANet(regressor, distractor, model_config.distracting_block)


class ReadoutNet(nn.Module):
    def __init__(self, feature_channels_num, feature_h, feature_w):
        super(ReadoutNet, self).__init__()
        self.net = nn.Sequential(
            OrderedDict([
                ('layernorm0',
                 nn.LayerNorm([feature_channels_num, feature_h, feature_w])),
                ('conv0', nn.Conv2d(feature_channels_num,
                                    512, (1, 1),
                                    bias=True)),
                ('softplus0', nn.Softplus()),
                ('layernorm1', nn.LayerNorm([512, feature_h, feature_w])),
                ('conv1', nn.Conv2d(512, 256, (1, 1), bias=True)),
                #('softplus1', nn.Softplus()),
                #('conv2', nn.Conv2d(256, 256, (1, 1), bias=True)),
                ('sigmoid', nn.Sigmoid())
            ]))

    def forward(self, x):
        return self.net(x)


class Finializer(nn.Module):
    def __init__(self, center_bias_weight=1, kernel_size=3, sigma=1):
        super(Finializer, self).__init__()
        # self.center_bias_weight = nn.Parameter(
        #     torch.Tensor([center_bias_weight]))
        self.kernel_size = kernel_size
        self.sigma = sigma

    def forward(self, x):
        x = filters.gaussian_blur2d(
            x,
            kernel_size=[self.kernel_size, self.kernel_size],
            sigma=[self.sigma, self.sigma],
            border_type='constant')
        #x = reduce(x, 'b c (h1 h) (w1 w) -> b c h w', 'max', h1=2, w1=2)
        return x


class Distractor(nn.Module):
    def __init__(self, feature_extracter, finilializer, readout_net):
        super(Distractor, self).__init__()
        self.feature_extracter = feature_extracter
        self.readout_net = readout_net
        self.finilializer = finilializer

        self.feature_extracter.eval()
        for param in self.feature_extracter.parameters():
            param.requires_grad = False

    def forward(self, x):
        with torch.no_grad():
            features = list(self.feature_extracter(x).values())
            for idx in range(len(features)):
                if features[idx].shape[2] == 7:
                    features[idx] = repeat(features[idx],
                                           'b c h w -> b c (h2 h) (w2 w)',
                                           h2=2,
                                           w2=2)
            feature = torch.cat(features, dim=1)
        x = self.readout_net(feature)
        x = self.finilializer(x)
        return x


class Regressor(nn.Module):
    def __init__(self,
                 backbone_type='resnet18',
                 pretrained=True,
                 weights_path=None):
        super(Regressor, self).__init__()

        self.backbone = models.__dict__[backbone_type](pretrained=pretrained)
        self.backbone.fc = nn.Sequential(
            nn.Linear(self.backbone.fc.in_features, 10), nn.Softmax(dim=1))

        if weights_path:
            self.backbone.load_state_dict(torch.load(weights_path))

    def forward(self, x):
        x = self.backbone(x)
        return x


class ASAIAANet(nn.Module):
    def __init__(self, regressor, distractor, target_block):
        super(ASAIAANet, self).__init__()
        self.distractor = distractor
        self.target_block = target_block
        self.regressor = regressor

    def forward(self, x, min_max):
        mask = self.distractor(x)
        for name, block in self.regressor.backbone.named_children():
            if name == 'fc':
                x = x.view(-1, 512)
                x = block(x)
            else:
                x = block(x)
                if name == self.target_block:
                    ill_features = x * mask
                    if min_max:
                        x = -random.random() * ill_features + x
                    else:
                        x = -ill_features + x
        return x, mask, ill_features