Spaces:
Running
Running
File size: 3,783 Bytes
c9baa67 |
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 |
from collections import OrderedDict
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils import model_zoo
from .features.densenet import RGBDenseNet201
from .modules import FeatureExtractor, Finalizer, DeepGazeIIIMixture
from .layers import FlexibleScanpathHistoryEncoding
from .layers import (
Conv2dMultiInput,
LayerNorm,
LayerNormMultiInput,
Bias,
)
def build_saliency_network(input_channels):
return nn.Sequential(OrderedDict([
('layernorm0', LayerNorm(input_channels)),
('conv0', nn.Conv2d(input_channels, 8, (1, 1), bias=False)),
('bias0', Bias(8)),
('softplus0', nn.Softplus()),
('layernorm1', LayerNorm(8)),
('conv1', nn.Conv2d(8, 16, (1, 1), bias=False)),
('bias1', Bias(16)),
('softplus1', nn.Softplus()),
('layernorm2', LayerNorm(16)),
('conv2', nn.Conv2d(16, 1, (1, 1), bias=False)),
('bias2', Bias(1)),
('softplus2', nn.Softplus()),
]))
def build_scanpath_network():
return nn.Sequential(OrderedDict([
('encoding0', FlexibleScanpathHistoryEncoding(in_fixations=4, channels_per_fixation=3, out_channels=128, kernel_size=[1, 1], bias=True)),
('softplus0', nn.Softplus()),
('layernorm1', LayerNorm(128)),
('conv1', nn.Conv2d(128, 16, (1, 1), bias=False)),
('bias1', Bias(16)),
('softplus1', nn.Softplus()),
]))
def build_fixation_selection_network():
return nn.Sequential(OrderedDict([
('layernorm0', LayerNormMultiInput([1, 16])),
('conv0', Conv2dMultiInput([1, 16], 128, (1, 1), bias=False)),
('bias0', Bias(128)),
('softplus0', nn.Softplus()),
('layernorm1', LayerNorm(128)),
('conv1', nn.Conv2d(128, 16, (1, 1), bias=False)),
('bias1', Bias(16)),
('softplus1', nn.Softplus()),
('conv2', nn.Conv2d(16, 1, (1, 1), bias=False)),
]))
class DeepGazeIII(DeepGazeIIIMixture):
"""DeepGazeIII model
:note
See Kümmerer, M., Bethge, M., & Wallis, T.S.A. (2022). DeepGaze III: Modeling free-viewing human scanpaths with deep learning. Journal of Vision 2022, https://doi.org/10.1167/jov.22.5.7
"""
def __init__(self, pretrained=True):
features = RGBDenseNet201()
feature_extractor = FeatureExtractor(features, [
'1.features.denseblock4.denselayer32.norm1',
'1.features.denseblock4.denselayer32.conv1',
'1.features.denseblock4.denselayer31.conv2',
])
saliency_networks = []
scanpath_networks = []
fixation_selection_networks = []
finalizers = []
for component in range(10):
saliency_network = build_saliency_network(2048)
scanpath_network = build_scanpath_network()
fixation_selection_network = build_fixation_selection_network()
saliency_networks.append(saliency_network)
scanpath_networks.append(scanpath_network)
fixation_selection_networks.append(fixation_selection_network)
finalizers.append(Finalizer(sigma=8.0, learn_sigma=True, saliency_map_factor=4))
super().__init__(
features=feature_extractor,
saliency_networks=saliency_networks,
scanpath_networks=scanpath_networks,
fixation_selection_networks=fixation_selection_networks,
finalizers=finalizers,
downsample=2,
readout_factor=4,
saliency_map_factor=4,
included_fixations=[-1, -2, -3, -4]
)
if pretrained:
self.load_state_dict(model_zoo.load_url('https://github.com/matthias-k/DeepGaze/releases/download/v1.1.0/deepgaze3.pth', map_location=torch.device('cpu'))) |