Spaces:
Sleeping
Sleeping
File size: 3,164 Bytes
a972d65 717222a a972d65 717222a a972d65 717222a a972d65 717222a a972d65 717222a a972d65 717222a a972d65 717222a a972d65 | 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 | import torch
from torch import nn
from torch.nn import functional as F
class Bottleneck(nn.Module):
expansion = 4
def __init__(self, inplanes, planes, stride=1, downsample=None):
super().__init__()
self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
self.bn1 = nn.BatchNorm2d(planes)
self.conv2 = nn.Conv2d(
planes, planes, kernel_size=3, stride=stride, padding=1, bias=False
)
self.bn2 = nn.BatchNorm2d(planes)
self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1, bias=False)
self.bn3 = nn.BatchNorm2d(planes * self.expansion)
self.relu = nn.ReLU(inplace=True)
self.downsample = downsample
def forward(self, x):
identity = x
out = self.relu(self.bn1(self.conv1(x)))
out = self.relu(self.bn2(self.conv2(out)))
out = self.bn3(self.conv3(out))
if self.downsample is not None:
identity = self.downsample(x)
out += identity
return self.relu(out)
class NPRModel(nn.Module):
def __init__(self):
super().__init__()
self.unfold_size = 2
self.unfold_index = 0
self.inplanes = 64
self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1, bias=False)
self.bn1 = nn.BatchNorm2d(64)
self.relu = nn.ReLU(inplace=True)
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.layer1 = self._make_layer(64, 3)
self.layer2 = self._make_layer(128, 4, stride=2)
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
self.fc1 = nn.Linear(512, 1)
def _make_layer(self, planes, blocks, stride=1):
downsample = None
if stride != 1 or self.inplanes != planes * Bottleneck.expansion:
downsample = nn.Sequential(
nn.Conv2d(
self.inplanes,
planes * Bottleneck.expansion,
kernel_size=1,
stride=stride,
bias=False,
),
nn.BatchNorm2d(planes * Bottleneck.expansion),
)
layers = [Bottleneck(self.inplanes, planes, stride, downsample)]
self.inplanes = planes * Bottleneck.expansion
for _ in range(1, blocks):
layers.append(Bottleneck(self.inplanes, planes))
return nn.Sequential(*layers)
@staticmethod
def interpolate(image, factor):
return F.interpolate(
F.interpolate(image, scale_factor=factor, mode="nearest", recompute_scale_factor=True),
scale_factor=1 / factor,
mode="nearest",
recompute_scale_factor=True,
)
def forward(self, x):
_, _, height, width = x.shape
if height % 2 == 1:
x = x[:, :, :-1, :]
if width % 2 == 1:
x = x[:, :, :, :-1]
x = (x - self.interpolate(x, 0.5)) * (2.0 / 3.0)
x = self.relu(self.bn1(self.conv1(x)))
x = self.maxpool(x)
x = self.layer1(x)
x = self.layer2(x)
x = torch.flatten(self.avgpool(x), 1)
return self.fc1(x)
|