From af3e53d3660f0366e2cb8e4fc7db653d28e7e6d6 Mon Sep 17 00:00:00 2001 From: Adrian Bulat Date: Mon, 14 Dec 2020 21:28:11 +0000 Subject: [PATCH] restructure model structure --- Dockerfile | 2 +- README.md | 44 +- face_alignment/api.py | 6 - face_alignment/models/__init__.py | 2 + face_alignment/{models.py => models/fan.py} | 421 ++++++++------------ face_alignment/models/resnet.py | 100 +++++ face_alignment/utils.py | 2 +- 7 files changed, 267 insertions(+), 310 deletions(-) create mode 100644 face_alignment/models/__init__.py rename face_alignment/{models.py => models/fan.py} (59%) create mode 100644 face_alignment/models/resnet.py diff --git a/Dockerfile b/Dockerfile index 5943690..0e34190 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,4 +1,4 @@ -# Based on https://github.com/pytorch/pytorch/blob/master/Dockerfile +# Based on a older version of https://github.com/pytorch/pytorch/blob/master/Dockerfile FROM nvidia/cuda:10.1-cudnn7-devel-ubuntu18.04 RUN apt-get update && apt-get install -y --no-install-recommends \ diff --git a/README.md b/README.md index f406ed1..2e9589d 100644 --- a/README.md +++ b/README.md @@ -85,7 +85,7 @@ Please also see the ``examples`` folder ### Requirements -* Python 3.5+ or Python 2.7 (it may work with other versions too). Support for Python 2.7 is deprecated. +* Python 3.5+ (it may work with other versions too). * Linux, Windows or macOS * pytorch (>=1.0) @@ -104,39 +104,7 @@ Alternatively, bellow, you can find instruction to build it from source. ### From source - Install pytorch and pytorch dependencies. Instructions taken from [pytorch readme](https://github.com/pytorch/pytorch). For a more updated version check the framework github page. - - On Linux -```bash -export CMAKE_PREFIX_PATH="$(dirname $(which conda))/../" # [anaconda root directory] - -# Install basic dependencies -conda install numpy pyyaml mkl setuptools cmake gcc cffi - -# Add LAPACK support for the GPU -conda install -c soumith magma-cuda80 # or magma-cuda75 if CUDA 7.5 -``` - -On OSX -```bash -export CMAKE_PREFIX_PATH=[anaconda root directory] -conda install numpy pyyaml setuptools cmake cffi -``` -#### Get the PyTorch source -```bash -git clone --recursive https://github.com/pytorch/pytorch -``` - -#### Install PyTorch -On Linux -```bash -python setup.py install -``` - -On OSX -```bash -MACOSX_DEPLOYMENT_TARGET=10.9 CC=clang CXX=clang++ python setup.py install -``` + Install pytorch and pytorch dependencies. Please check the [pytorch readme](https://github.com/pytorch/pytorch) for this. #### Get the Face Alignment source code ```bash @@ -150,7 +118,7 @@ python setup.py install ### Docker image -A Dockerfile is provided to build images with cuda support and cudnn v5. For more instructions about running and building a docker image check the orginal Docker documentation. +A Dockerfile is provided to build images with cuda support and cudnn. For more instructions about running and building a docker image check the orginal Docker documentation. ``` docker build -t face-alignment . ``` @@ -175,9 +143,3 @@ All contributions are welcomed. If you encounter any issue (including examples o ``` For citing dlib, pytorch or any other packages used here please check the original page of their respective authors. - -## Acknowledgements - -* To the [pytorch](http://pytorch.org/) team for providing such an awesome deeplearning framework -* To [my supervisor](http://www.cs.nott.ac.uk/~pszyt/) for his patience and suggestions. -* To all other python developers that made available the rest of the packages used in this repository. diff --git a/face_alignment/api.py b/face_alignment/api.py index 7a36488..2cadf68 100644 --- a/face_alignment/api.py +++ b/face_alignment/api.py @@ -1,15 +1,9 @@ -import os import torch from torch.utils.model_zoo import load_url from enum import Enum from skimage import io from skimage import color import numpy as np -import cv2 -try: - import urllib.request as request_file -except BaseException: - import urllib as request_file from .models import FAN, ResNetDepth from .utils import * diff --git a/face_alignment/models/__init__.py b/face_alignment/models/__init__.py new file mode 100644 index 0000000..fb4fa1c --- /dev/null +++ b/face_alignment/models/__init__.py @@ -0,0 +1,2 @@ +from .fan import FAN +from .resnet import ResNetDepth \ No newline at end of file diff --git a/face_alignment/models.py b/face_alignment/models/fan.py similarity index 59% rename from face_alignment/models.py rename to face_alignment/models/fan.py index d8f56bf..6aab129 100644 --- a/face_alignment/models.py +++ b/face_alignment/models/fan.py @@ -1,261 +1,160 @@ -import torch -import torch.nn as nn -import torch.nn.functional as F -import math - - -def conv3x3(in_planes, out_planes, strd=1, padding=1, bias=False): - "3x3 convolution with padding" - return nn.Conv2d(in_planes, out_planes, kernel_size=3, - stride=strd, padding=padding, bias=bias) - - -class ConvBlock(nn.Module): - def __init__(self, in_planes, out_planes): - super(ConvBlock, self).__init__() - self.bn1 = nn.BatchNorm2d(in_planes) - self.conv1 = conv3x3(in_planes, int(out_planes / 2)) - self.bn2 = nn.BatchNorm2d(int(out_planes / 2)) - self.conv2 = conv3x3(int(out_planes / 2), int(out_planes / 4)) - self.bn3 = nn.BatchNorm2d(int(out_planes / 4)) - self.conv3 = conv3x3(int(out_planes / 4), int(out_planes / 4)) - - if in_planes != out_planes: - self.downsample = nn.Sequential( - nn.BatchNorm2d(in_planes), - nn.ReLU(True), - nn.Conv2d(in_planes, out_planes, - kernel_size=1, stride=1, bias=False), - ) - else: - self.downsample = None - - def forward(self, x): - residual = x - - out1 = self.bn1(x) - out1 = F.relu(out1, True) - out1 = self.conv1(out1) - - out2 = self.bn2(out1) - out2 = F.relu(out2, True) - out2 = self.conv2(out2) - - out3 = self.bn3(out2) - out3 = F.relu(out3, True) - out3 = self.conv3(out3) - - out3 = torch.cat((out1, out2, out3), 1) - - if self.downsample is not None: - residual = self.downsample(residual) - - out3 += residual - - return out3 - - -class Bottleneck(nn.Module): - - expansion = 4 - - def __init__(self, inplanes, planes, stride=1, downsample=None): - super(Bottleneck, self).__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 * 4, kernel_size=1, bias=False) - self.bn3 = nn.BatchNorm2d(planes * 4) - self.relu = nn.ReLU(inplace=True) - self.downsample = downsample - self.stride = stride - - def forward(self, x): - residual = x - - out = self.conv1(x) - out = self.bn1(out) - out = self.relu(out) - - out = self.conv2(out) - out = self.bn2(out) - out = self.relu(out) - - out = self.conv3(out) - out = self.bn3(out) - - if self.downsample is not None: - residual = self.downsample(x) - - out += residual - out = self.relu(out) - - return out - - -class HourGlass(nn.Module): - def __init__(self, num_modules, depth, num_features): - super(HourGlass, self).__init__() - self.num_modules = num_modules - self.depth = depth - self.features = num_features - - self._generate_network(self.depth) - - def _generate_network(self, level): - self.add_module('b1_' + str(level), ConvBlock(self.features, self.features)) - - self.add_module('b2_' + str(level), ConvBlock(self.features, self.features)) - - if level > 1: - self._generate_network(level - 1) - else: - self.add_module('b2_plus_' + str(level), ConvBlock(self.features, self.features)) - - self.add_module('b3_' + str(level), ConvBlock(self.features, self.features)) - - def _forward(self, level, inp): - # Upper branch - up1 = inp - up1 = self._modules['b1_' + str(level)](up1) - - # Lower branch - low1 = F.avg_pool2d(inp, 2, stride=2) - low1 = self._modules['b2_' + str(level)](low1) - - if level > 1: - low2 = self._forward(level - 1, low1) - else: - low2 = low1 - low2 = self._modules['b2_plus_' + str(level)](low2) - - low3 = low2 - low3 = self._modules['b3_' + str(level)](low3) - - up2 = F.interpolate(low3, scale_factor=2, mode='nearest') - - return up1 + up2 - - def forward(self, x): - return self._forward(self.depth, x) - - -class FAN(nn.Module): - - def __init__(self, num_modules=1): - super(FAN, self).__init__() - self.num_modules = num_modules - - # Base part - self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3) - self.bn1 = nn.BatchNorm2d(64) - self.conv2 = ConvBlock(64, 128) - self.conv3 = ConvBlock(128, 128) - self.conv4 = ConvBlock(128, 256) - - # Stacking part - for hg_module in range(self.num_modules): - self.add_module('m' + str(hg_module), HourGlass(1, 4, 256)) - self.add_module('top_m_' + str(hg_module), ConvBlock(256, 256)) - self.add_module('conv_last' + str(hg_module), - nn.Conv2d(256, 256, kernel_size=1, stride=1, padding=0)) - self.add_module('bn_end' + str(hg_module), nn.BatchNorm2d(256)) - self.add_module('l' + str(hg_module), nn.Conv2d(256, - 68, kernel_size=1, stride=1, padding=0)) - - if hg_module < self.num_modules - 1: - self.add_module( - 'bl' + str(hg_module), nn.Conv2d(256, 256, kernel_size=1, stride=1, padding=0)) - self.add_module('al' + str(hg_module), nn.Conv2d(68, - 256, kernel_size=1, stride=1, padding=0)) - - def forward(self, x): - x = F.relu(self.bn1(self.conv1(x)), True) - x = F.avg_pool2d(self.conv2(x), 2, stride=2) - x = self.conv3(x) - x = self.conv4(x) - - previous = x - - outputs = [] - for i in range(self.num_modules): - hg = self._modules['m' + str(i)](previous) - - ll = hg - ll = self._modules['top_m_' + str(i)](ll) - - ll = F.relu(self._modules['bn_end' + str(i)] - (self._modules['conv_last' + str(i)](ll)), True) - - # Predict heatmaps - tmp_out = self._modules['l' + str(i)](ll) - outputs.append(tmp_out) - - if i < self.num_modules - 1: - ll = self._modules['bl' + str(i)](ll) - tmp_out_ = self._modules['al' + str(i)](tmp_out) - previous = previous + ll + tmp_out_ - - return outputs - - -class ResNetDepth(nn.Module): - - def __init__(self, block=Bottleneck, layers=[3, 8, 36, 3], num_classes=68): - self.inplanes = 64 - super(ResNetDepth, self).__init__() - self.conv1 = nn.Conv2d(3 + 68, 64, kernel_size=7, stride=2, padding=3, - 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(block, 64, layers[0]) - self.layer2 = self._make_layer(block, 128, layers[1], stride=2) - self.layer3 = self._make_layer(block, 256, layers[2], stride=2) - self.layer4 = self._make_layer(block, 512, layers[3], stride=2) - self.avgpool = nn.AvgPool2d(7) - self.fc = nn.Linear(512 * block.expansion, num_classes) - - for m in self.modules(): - if isinstance(m, nn.Conv2d): - n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels - m.weight.data.normal_(0, math.sqrt(2. / n)) - elif isinstance(m, nn.BatchNorm2d): - m.weight.data.fill_(1) - m.bias.data.zero_() - - def _make_layer(self, block, planes, blocks, stride=1): - downsample = None - if stride != 1 or self.inplanes != planes * block.expansion: - downsample = nn.Sequential( - nn.Conv2d(self.inplanes, planes * block.expansion, - kernel_size=1, stride=stride, bias=False), - nn.BatchNorm2d(planes * block.expansion), - ) - - layers = [] - layers.append(block(self.inplanes, planes, stride, downsample)) - self.inplanes = planes * block.expansion - for i in range(1, blocks): - layers.append(block(self.inplanes, planes)) - - return nn.Sequential(*layers) - - def forward(self, x): - x = self.conv1(x) - x = self.bn1(x) - x = self.relu(x) - x = self.maxpool(x) - - x = self.layer1(x) - x = self.layer2(x) - x = self.layer3(x) - x = self.layer4(x) - - x = self.avgpool(x) - x = x.view(x.size(0), -1) - x = self.fc(x) - - return x +import torch +import torch.nn as nn +import torch.nn.functional as F +import math + +def conv3x3(in_planes, out_planes, strd=1, padding=1, bias=False): + "3x3 convolution with padding" + return nn.Conv2d(in_planes, out_planes, kernel_size=3, + stride=strd, padding=padding, bias=bias) + + +class ConvBlock(nn.Module): + def __init__(self, in_planes, out_planes): + super(ConvBlock, self).__init__() + self.bn1 = nn.BatchNorm2d(in_planes) + self.conv1 = conv3x3(in_planes, int(out_planes / 2)) + self.bn2 = nn.BatchNorm2d(int(out_planes / 2)) + self.conv2 = conv3x3(int(out_planes / 2), int(out_planes / 4)) + self.bn3 = nn.BatchNorm2d(int(out_planes / 4)) + self.conv3 = conv3x3(int(out_planes / 4), int(out_planes / 4)) + + if in_planes != out_planes: + self.downsample = nn.Sequential( + nn.BatchNorm2d(in_planes), + nn.ReLU(True), + nn.Conv2d(in_planes, out_planes, + kernel_size=1, stride=1, bias=False), + ) + else: + self.downsample = None + + def forward(self, x): + residual = x + + out1 = self.bn1(x) + out1 = F.relu(out1, True) + out1 = self.conv1(out1) + + out2 = self.bn2(out1) + out2 = F.relu(out2, True) + out2 = self.conv2(out2) + + out3 = self.bn3(out2) + out3 = F.relu(out3, True) + out3 = self.conv3(out3) + + out3 = torch.cat((out1, out2, out3), 1) + + if self.downsample is not None: + residual = self.downsample(residual) + + out3 += residual + + return out3 + + +class HourGlass(nn.Module): + def __init__(self, num_modules, depth, num_features): + super(HourGlass, self).__init__() + self.num_modules = num_modules + self.depth = depth + self.features = num_features + + self._generate_network(self.depth) + + def _generate_network(self, level): + self.add_module('b1_' + str(level), ConvBlock(self.features, self.features)) + + self.add_module('b2_' + str(level), ConvBlock(self.features, self.features)) + + if level > 1: + self._generate_network(level - 1) + else: + self.add_module('b2_plus_' + str(level), ConvBlock(self.features, self.features)) + + self.add_module('b3_' + str(level), ConvBlock(self.features, self.features)) + + def _forward(self, level, inp): + # Upper branch + up1 = inp + up1 = self._modules['b1_' + str(level)](up1) + + # Lower branch + low1 = F.avg_pool2d(inp, 2, stride=2) + low1 = self._modules['b2_' + str(level)](low1) + + if level > 1: + low2 = self._forward(level - 1, low1) + else: + low2 = low1 + low2 = self._modules['b2_plus_' + str(level)](low2) + + low3 = low2 + low3 = self._modules['b3_' + str(level)](low3) + + up2 = F.interpolate(low3, scale_factor=2, mode='nearest') + + return up1 + up2 + + def forward(self, x): + return self._forward(self.depth, x) + + +class FAN(nn.Module): + + def __init__(self, num_modules=1): + super(FAN, self).__init__() + self.num_modules = num_modules + + # Base part + self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3) + self.bn1 = nn.BatchNorm2d(64) + self.conv2 = ConvBlock(64, 128) + self.conv3 = ConvBlock(128, 128) + self.conv4 = ConvBlock(128, 256) + + # Stacking part + for hg_module in range(self.num_modules): + self.add_module('m' + str(hg_module), HourGlass(1, 4, 256)) + self.add_module('top_m_' + str(hg_module), ConvBlock(256, 256)) + self.add_module('conv_last' + str(hg_module), + nn.Conv2d(256, 256, kernel_size=1, stride=1, padding=0)) + self.add_module('bn_end' + str(hg_module), nn.BatchNorm2d(256)) + self.add_module('l' + str(hg_module), nn.Conv2d(256, + 68, kernel_size=1, stride=1, padding=0)) + + if hg_module < self.num_modules - 1: + self.add_module( + 'bl' + str(hg_module), nn.Conv2d(256, 256, kernel_size=1, stride=1, padding=0)) + self.add_module('al' + str(hg_module), nn.Conv2d(68, + 256, kernel_size=1, stride=1, padding=0)) + + def forward(self, x): + x = F.relu(self.bn1(self.conv1(x)), True) + x = F.avg_pool2d(self.conv2(x), 2, stride=2) + x = self.conv3(x) + x = self.conv4(x) + + previous = x + + outputs = [] + for i in range(self.num_modules): + hg = self._modules['m' + str(i)](previous) + + ll = hg + ll = self._modules['top_m_' + str(i)](ll) + + ll = F.relu(self._modules['bn_end' + str(i)] + (self._modules['conv_last' + str(i)](ll)), True) + + # Predict heatmaps + tmp_out = self._modules['l' + str(i)](ll) + outputs.append(tmp_out) + + if i < self.num_modules - 1: + ll = self._modules['bl' + str(i)](ll) + tmp_out_ = self._modules['al' + str(i)](tmp_out) + previous = previous + ll + tmp_out_ + + return outputs diff --git a/face_alignment/models/resnet.py b/face_alignment/models/resnet.py new file mode 100644 index 0000000..d9b8b7b --- /dev/null +++ b/face_alignment/models/resnet.py @@ -0,0 +1,100 @@ +import math +import torch.nn as nn + +class Bottleneck(nn.Module): + + expansion = 4 + + def __init__(self, inplanes, planes, stride=1, downsample=None): + super(Bottleneck, self).__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 * 4, kernel_size=1, bias=False) + self.bn3 = nn.BatchNorm2d(planes * 4) + self.relu = nn.ReLU(inplace=True) + self.downsample = downsample + self.stride = stride + + def forward(self, x): + residual = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + out = self.relu(out) + + out = self.conv3(out) + out = self.bn3(out) + + if self.downsample is not None: + residual = self.downsample(x) + + out += residual + out = self.relu(out) + + return out + +class ResNetDepth(nn.Module): + + def __init__(self, block=Bottleneck, layers=[3, 8, 36, 3], num_classes=68): + self.inplanes = 64 + super(ResNetDepth, self).__init__() + self.conv1 = nn.Conv2d(3 + 68, 64, kernel_size=7, stride=2, padding=3, + 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(block, 64, layers[0]) + self.layer2 = self._make_layer(block, 128, layers[1], stride=2) + self.layer3 = self._make_layer(block, 256, layers[2], stride=2) + self.layer4 = self._make_layer(block, 512, layers[3], stride=2) + self.avgpool = nn.AvgPool2d(7) + self.fc = nn.Linear(512 * block.expansion, num_classes) + + for m in self.modules(): + if isinstance(m, nn.Conv2d): + n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + m.weight.data.normal_(0, math.sqrt(2. / n)) + elif isinstance(m, nn.BatchNorm2d): + m.weight.data.fill_(1) + m.bias.data.zero_() + + def _make_layer(self, block, planes, blocks, stride=1): + downsample = None + if stride != 1 or self.inplanes != planes * block.expansion: + downsample = nn.Sequential( + nn.Conv2d(self.inplanes, planes * block.expansion, + kernel_size=1, stride=stride, bias=False), + nn.BatchNorm2d(planes * block.expansion), + ) + + layers = [] + layers.append(block(self.inplanes, planes, stride, downsample)) + self.inplanes = planes * block.expansion + for i in range(1, blocks): + layers.append(block(self.inplanes, planes)) + + return nn.Sequential(*layers) + + def forward(self, x): + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + x = self.maxpool(x) + + x = self.layer1(x) + x = self.layer2(x) + x = self.layer3(x) + x = self.layer4(x) + + x = self.avgpool(x) + x = x.view(x.size(0), -1) + x = self.fc(x) + + return x \ No newline at end of file diff --git a/face_alignment/utils.py b/face_alignment/utils.py index f93c529..7ac8eef 100644 --- a/face_alignment/utils.py +++ b/face_alignment/utils.py @@ -249,7 +249,7 @@ def flip(tensor, is_label=False): # Pytorch load supports only pytorch models -def load_file_from_url(url, model_dir=None, map_location=None, progress=True, check_hash=False, file_name=None): +def load_file_from_url(url, model_dir=None, progress=True, check_hash=False, file_name=None): if model_dir is None: hub_dir = get_dir() model_dir = os.path.join(hub_dir, 'checkpoints')