Spaces:
Sleeping
Sleeping
| import numpy as np | |
| import torch.nn as nn | |
| import torch | |
| import torch.nn.functional as F | |
| class ResnetEncoder(nn.Module): | |
| def __init__(self, append_layers=None): | |
| super(ResnetEncoder, self).__init__() | |
| from . import resnet | |
| # feature_size = 2048 | |
| self.feature_dim = 2048 | |
| self.encoder = resnet.load_ResNet50Model() # out: 2048 | |
| # regressor | |
| self.append_layers = append_layers | |
| # for normalize input images | |
| MEAN = [0.485, 0.456, 0.406] | |
| STD = [0.229, 0.224, 0.225] | |
| self.register_buffer('MEAN', torch.tensor(MEAN)[None, :, None, None]) | |
| self.register_buffer('STD', torch.tensor(STD)[None, :, None, None]) | |
| def forward(self, inputs): | |
| ''' inputs: [bz, 3, h, w], range: [0,1] | |
| ''' | |
| inputs = (inputs - self.MEAN) / self.STD | |
| features = self.encoder(inputs) | |
| if self.append_layers: | |
| features = self.last_op(features) | |
| return features | |
| class MLP(nn.Module): | |
| def __init__(self, channels=[2048, 1024, 1], last_op=None): | |
| super(MLP, self).__init__() | |
| layers = [] | |
| for l in range(0, len(channels) - 1): | |
| layers.append(nn.Linear(channels[l], channels[l + 1])) | |
| if l < len(channels) - 2: | |
| layers.append(nn.ReLU()) | |
| if last_op: | |
| layers.append(last_op) | |
| self.layers = nn.Sequential(*layers) | |
| def forward(self, inputs): | |
| outs = self.layers(inputs) | |
| return outs | |
| class HRNEncoder(nn.Module): | |
| def __init__(self, append_layers=None): | |
| super(HRNEncoder, self).__init__() | |
| from . import hrnet | |
| self.feature_dim = 2048 | |
| self.encoder = hrnet.load_HRNet(pretrained=True) # out: 2048 | |
| # regressor | |
| self.append_layers = append_layers | |
| # for normalize input images | |
| MEAN = [0.485, 0.456, 0.406] | |
| STD = [0.229, 0.224, 0.225] | |
| self.register_buffer('MEAN', torch.tensor(MEAN)[None, :, None, None]) | |
| self.register_buffer('STD', torch.tensor(STD)[None, :, None, None]) | |
| def forward(self, inputs): | |
| ''' inputs: [bz, 3, h, w], range: [0,1] | |
| ''' | |
| inputs = (inputs - self.MEAN) / self.STD | |
| features = self.encoder(inputs)['concat'] | |
| if self.append_layers: | |
| features = self.last_op(features) | |
| return features | |