diff --git a/cv/face/arcface/pytorch/.gitignore b/cv/face/arcface/pytorch/.gitignore new file mode 100755 index 0000000000000000000000000000000000000000..7290d7cc89e8065b72c674f05f466fb1cd71f983 --- /dev/null +++ b/cv/face/arcface/pytorch/.gitignore @@ -0,0 +1,139 @@ +# ignore map, miou, datasets +map_out/ +miou_out/ +VOCdevkit/ +datasets/ +Medical_Datasets/ +lfw/ +logs/ +model_data/ + +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +pip-wheel-metadata/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +.python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ diff --git a/cv/face/arcface/pytorch/LICENSE b/cv/face/arcface/pytorch/LICENSE new file mode 100755 index 0000000000000000000000000000000000000000..0f36b54e8c55acfe13cf2d5a96ce965eed586876 --- /dev/null +++ b/cv/face/arcface/pytorch/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2022 Bubbliiiing + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/cv/face/arcface/pytorch/README.md b/cv/face/arcface/pytorch/README.md new file mode 100644 index 0000000000000000000000000000000000000000..896b1dc3d9feafe77a599bce730f385b5760d05f --- /dev/null +++ b/cv/face/arcface/pytorch/README.md @@ -0,0 +1,33 @@ +# arcface +## Model description +This repo is a pytorch implement of ArcFace, which propose an Additive Angular Margin Loss to obtain highly discriminative features for face recognition. The proposed ArcFace has a clear geometric interpretation due to the exact correspondence to the geodesic distance on the hypersphere. ArcFace consistently outperforms the state-of-the-art and can be easily implemented with negligible computational overhead + +## install +```bash +pip3 install requiresments.txt +wget http://www.zlib.net/fossils/zlib-1.2.9.tar.gz +tar xvf zlib-1.2.9.tar.gz +cd zlib-1.2.9/ +./configure && make install +``` +## download + +```bash +cd datasets +``` +download dataset in this way: +download link: https://pan.baidu.com/s/1qMxFR8H_ih0xmY-rKgRejw password: bcrq + +## preprocess + +```bash +python3 txt_annotation.py +``` + +## train + +```bash +bash run.sh $GPUS +``` + + diff --git a/cv/face/arcface/pytorch/arcface.py b/cv/face/arcface/pytorch/arcface.py new file mode 100755 index 0000000000000000000000000000000000000000..8991d519a0863a301b3d2725ec71e4d5b4ca1e1d --- /dev/null +++ b/cv/face/arcface/pytorch/arcface.py @@ -0,0 +1,141 @@ +import matplotlib.pyplot as plt +import numpy as np +import torch +import torch.backends.cudnn as cudnn + +from nets.arcface import Arcface as arcface +from utils.utils import preprocess_input, resize_image, show_config + + +class Arcface(object): + _defaults = { + #--------------------------------------------------------------------------# + # 使用自己训练好的模型进行预测要修改model_path，指向logs文件夹下的权值文件 + # 训练好后logs文件夹下存在多个权值文件，选择验证集损失较低的即可。 + # 验证集损失较低不代表准确度较高，仅代表该权值在验证集上泛化性能较好。 + #--------------------------------------------------------------------------# + "model_path" : "model_data/arcface_mobilefacenet.pth", + #-------------------------------------------# + # 输入图片的大小。 + #-------------------------------------------# + "input_shape" : [112, 112, 3], + #-------------------------------------------# + # 所使用到的主干特征提取网络，与训练的相同 + # mobilefacenet + # mobilenetv1 + # iresnet18 + # iresnet34 + # iresnet50 + # iresnet100 + # iresnet200 + #-------------------------------------------# + "backbone" : "mobilefacenet", + #-------------------------------------------# + # 是否进行不失真的resize + #-------------------------------------------# + "letterbox_image" : True, + #-------------------------------------------# + # 是否使用Cuda + # 没有GPU可以设置成False + #-------------------------------------------# + "cuda" : True, + } + + @classmethod + def get_defaults(cls, n): + if n in cls._defaults: + return cls._defaults[n] + else: + return "Unrecognized attribute name '" + n + "'" + + #---------------------------------------------------# + # 初始化Arcface + #---------------------------------------------------# + def __init__(self, **kwargs): + self.__dict__.update(self._defaults) + for name, value in kwargs.items(): + setattr(self, name, value) + + self.generate() + + show_config(**self._defaults) + + def generate(self): + #---------------------------------------------------# + # 载入模型与权值 + #---------------------------------------------------# + print('Loading weights into state dict...') + device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + self.net = arcface(backbone=self.backbone, mode="predict").eval() + self.net.load_state_dict(torch.load(self.model_path, map_location=device), strict=False) + print('{} model loaded.'.format(self.model_path)) + + if self.cuda: + self.net = torch.nn.DataParallel(self.net) + cudnn.benchmark = True + self.net = self.net.cuda() + + #---------------------------------------------------# + # 检测图片 + #---------------------------------------------------# + def detect_image(self, image_1, image_2): + #---------------------------------------------------# + # 图片预处理，归一化 + #---------------------------------------------------# + with torch.no_grad(): + image_1 = resize_image(image_1, [self.input_shape[1], self.input_shape[0]], letterbox_image=self.letterbox_image) + image_2 = resize_image(image_2, [self.input_shape[1], self.input_shape[0]], letterbox_image=self.letterbox_image) + + photo_1 = torch.from_numpy(np.expand_dims(np.transpose(preprocess_input(np.array(image_1, np.float32)), (2, 0, 1)), 0)) + photo_2 = torch.from_numpy(np.expand_dims(np.transpose(preprocess_input(np.array(image_2, np.float32)), (2, 0, 1)), 0)) + + if self.cuda: + photo_1 = photo_1.cuda() + photo_2 = photo_2.cuda() + + #---------------------------------------------------# + # 图片传入网络进行预测 + #---------------------------------------------------# + output1 = self.net(photo_1).cpu().numpy() + output2 = self.net(photo_2).cpu().numpy() + + #---------------------------------------------------# + # 计算二者之间的距离 + #---------------------------------------------------# + l1 = np.linalg.norm(output1 - output2, axis=1) + + plt.subplot(1, 2, 1) + plt.imshow(np.array(image_1)) + + plt.subplot(1, 2, 2) + plt.imshow(np.array(image_2)) + plt.text(-12, -12, 'Distance:%.3f' % l1, ha='center', va= 'bottom',fontsize=11) + plt.show() + return l1 + + def get_FPS(self, image, test_interval): + #---------------------------------------------------# + # 对图片进行不失真的resize + #---------------------------------------------------# + image_data = resize_image(image, [self.input_shape[1], self.input_shape[0]], self.letterbox_image) + #---------------------------------------------------------# + # 归一化+添加上batch_size维度 + #---------------------------------------------------------# + image_data = torch.from_numpy(np.expand_dims(np.transpose(preprocess_input(np.array(image_data, np.float32)), (2, 0, 1)), 0)) + with torch.no_grad(): + #---------------------------------------------------# + # 图片传入网络进行预测 + #---------------------------------------------------# + preds = self.net(image_data).cpu().numpy() + + import time + t1 = time.time() + for _ in range(test_interval): + with torch.no_grad(): + #---------------------------------------------------# + # 图片传入网络进行预测 + #---------------------------------------------------# + preds = self.net(image_data).cpu().numpy() + t2 = time.time() + tact_time = (t2 - t1) / test_interval + return tact_time diff --git a/cv/face/arcface/pytorch/eval_LFW.py b/cv/face/arcface/pytorch/eval_LFW.py new file mode 100755 index 0000000000000000000000000000000000000000..53b2a0fc4ec417e43429bdc08c476c242dfa5148 --- /dev/null +++ b/cv/face/arcface/pytorch/eval_LFW.py @@ -0,0 +1,65 @@ +import torch +import torch.backends.cudnn as cudnn + +from nets.arcface import Arcface +from utils.dataloader import LFWDataset +from utils.utils_metrics import test + + +if __name__ == "__main__": + #--------------------------------------# + # 是否使用Cuda + # 没有GPU可以设置成False + #--------------------------------------# + cuda = True + #--------------------------------------# + # 主干特征提取网络的选择 + # mobilefacenet + # mobilenetv1 + # iresnet18 + # iresnet34 + # iresnet50 + # iresnet100 + # iresnet200 + #--------------------------------------# + backbone = "mobilefacenet" + #--------------------------------------# + # 输入图像大小 + #--------------------------------------# + input_shape = [112, 112, 3] + #--------------------------------------# + # 训练好的权值文件 + #--------------------------------------# + model_path = "model_data/arcface_mobilefacenet.pth" + #--------------------------------------# + # LFW评估数据集的文件路径 + # 以及对应的txt文件 + #--------------------------------------# + lfw_dir_path = "lfw" + lfw_pairs_path = "model_data/lfw_pair.txt" + #--------------------------------------# + # 评估的批次大小和记录间隔 + #--------------------------------------# + batch_size = 256 + log_interval = 1 + #--------------------------------------# + # ROC图的保存路径 + #--------------------------------------# + png_save_path = "model_data/roc_test.png" + + test_loader = torch.utils.data.DataLoader( + LFWDataset(dir=lfw_dir_path, pairs_path=lfw_pairs_path, image_size=input_shape), batch_size=batch_size, shuffle=False) + + model = Arcface(backbone=backbone, mode="predict") + + print('Loading weights into state dict...') + device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + model.load_state_dict(torch.load(model_path, map_location=device), strict=False) + model = model.eval() + + if cuda: + model = torch.nn.DataParallel(model) + cudnn.benchmark = True + model = model.cuda() + + test(test_loader, model, png_save_path, log_interval, batch_size, cuda) diff --git a/cv/face/arcface/pytorch/img/1_001.jpg b/cv/face/arcface/pytorch/img/1_001.jpg new file mode 100755 index 0000000000000000000000000000000000000000..66ca17ff6433b9b693f7df24f76f369c3503f2f4 Binary files /dev/null and b/cv/face/arcface/pytorch/img/1_001.jpg differ diff --git a/cv/face/arcface/pytorch/img/1_002.jpg b/cv/face/arcface/pytorch/img/1_002.jpg new file mode 100755 index 0000000000000000000000000000000000000000..b260be6fc0de53d88e86e31b7e41339a6e4fa774 Binary files /dev/null and b/cv/face/arcface/pytorch/img/1_002.jpg differ diff --git a/cv/face/arcface/pytorch/img/2_001.jpg b/cv/face/arcface/pytorch/img/2_001.jpg new file mode 100755 index 0000000000000000000000000000000000000000..5922d291e1e7ee1f27d7c6167a0348408868bf99 Binary files /dev/null and b/cv/face/arcface/pytorch/img/2_001.jpg differ diff --git a/cv/face/arcface/pytorch/nets/__init__.py b/cv/face/arcface/pytorch/nets/__init__.py new file mode 100755 index 0000000000000000000000000000000000000000..4287ca8617970fa8fc025b75cb319c7032706910 --- /dev/null +++ b/cv/face/arcface/pytorch/nets/__init__.py @@ -0,0 +1 @@ +# \ No newline at end of file diff --git a/cv/face/arcface/pytorch/nets/arcface.py b/cv/face/arcface/pytorch/nets/arcface.py new file mode 100755 index 0000000000000000000000000000000000000000..c44cb0dcbb445de3c7154a09e448a846c7d3eb93 --- /dev/null +++ b/cv/face/arcface/pytorch/nets/arcface.py @@ -0,0 +1,90 @@ +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.nn import Module, Parameter + +from nets.iresnet import (iresnet18, iresnet34, iresnet50, iresnet100, + iresnet200) +from nets.mobilefacenet import get_mbf +from nets.mobilenet import get_mobilenet + +class Arcface_Head(Module): + def __init__(self, embedding_size=128, num_classes=10575, s=64., m=0.5): + super(Arcface_Head, self).__init__() + self.s = s + self.m = m + self.weight = Parameter(torch.FloatTensor(num_classes, embedding_size)) + nn.init.xavier_uniform_(self.weight) + + self.cos_m = math.cos(m) + self.sin_m = math.sin(m) + self.th = math.cos(math.pi - m) + self.mm = math.sin(math.pi - m) * m + + def forward(self, input, label): + cosine = F.linear(input, F.normalize(self.weight)) + sine = torch.sqrt((1.0 - torch.pow(cosine, 2)).clamp(0, 1)) + phi = cosine * self.cos_m - sine * self.sin_m + phi = torch.where(cosine.float() > self.th, phi.float(), cosine.float() - self.mm) + + one_hot = torch.zeros(cosine.size()).type_as(phi).long() + one_hot.scatter_(1, label.view(-1, 1).long(), 1) + output = (one_hot * phi) + ((1.0 - one_hot) * cosine) + output *= self.s + return output + +class Arcface(nn.Module): + def __init__(self, num_classes=None, backbone="mobilefacenet", pretrained=False, mode="train"): + super(Arcface, self).__init__() + if backbone=="mobilefacenet": + embedding_size = 128 + s = 32 + self.arcface = get_mbf(embedding_size=embedding_size, pretrained=pretrained) + + elif backbone=="mobilenetv1": + embedding_size = 512 + s = 64 + self.arcface = get_mobilenet(dropout_keep_prob=0.5, embedding_size=embedding_size, pretrained=pretrained) + + elif backbone=="iresnet18": + embedding_size = 512 + s = 64 + self.arcface = iresnet18(dropout_keep_prob=0.5, embedding_size=embedding_size, pretrained=pretrained) + + elif backbone=="iresnet34": + embedding_size = 512 + s = 64 + self.arcface = iresnet34(dropout_keep_prob=0.5, embedding_size=embedding_size, pretrained=pretrained) + + elif backbone=="iresnet50": + embedding_size = 512 + s = 64 + self.arcface = iresnet50(dropout_keep_prob=0.5, embedding_size=embedding_size, pretrained=pretrained) + + elif backbone=="iresnet100": + embedding_size = 512 + s = 64 + self.arcface = iresnet100(dropout_keep_prob=0.5, embedding_size=embedding_size, pretrained=pretrained) + + elif backbone=="iresnet200": + embedding_size = 512 + s = 64 + self.arcface = iresnet200(dropout_keep_prob=0.5, embedding_size=embedding_size, pretrained=pretrained) + else: + raise ValueError('Unsupported backbone - `{}`, Use mobilefacenet, mobilenetv1.'.format(backbone)) + + self.mode = mode + if mode == "train": + self.head = Arcface_Head(embedding_size=embedding_size, num_classes=num_classes, s=s) + + def forward(self, x, y = None, mode = "predict"): + x = self.arcface(x) + x = x.view(x.size()[0], -1) + x = F.normalize(x) + if mode == "predict": + return x + else: + x = self.head(x, y) + return x diff --git a/cv/face/arcface/pytorch/nets/arcface_training.py b/cv/face/arcface/pytorch/nets/arcface_training.py new file mode 100755 index 0000000000000000000000000000000000000000..2c4f688e1169c0aa4888bf27fb3963138b104776 --- /dev/null +++ b/cv/face/arcface/pytorch/nets/arcface_training.py @@ -0,0 +1,46 @@ +import math +from functools import partial + + +def get_lr_scheduler(lr_decay_type, lr, min_lr, total_iters, warmup_iters_ratio = 0.1, warmup_lr_ratio = 0.1, no_aug_iter_ratio = 0.3, step_num = 10): + def yolox_warm_cos_lr(lr, min_lr, total_iters, warmup_total_iters, warmup_lr_start, no_aug_iter, iters): + if iters <= warmup_total_iters: + # lr = (lr - warmup_lr_start) * iters / float(warmup_total_iters) + warmup_lr_start + lr = (lr - warmup_lr_start) * pow(iters / float(warmup_total_iters), 2 + ) + warmup_lr_start + elif iters >= total_iters - no_aug_iter: + lr = min_lr + else: + lr = min_lr + 0.5 * (lr - min_lr) * ( + 1.0 + + math.cos( + math.pi + * (iters - warmup_total_iters) + / (total_iters - warmup_total_iters - no_aug_iter) + ) + ) + return lr + + def step_lr(lr, decay_rate, step_size, iters): + if step_size < 1: + raise ValueError("step_size must above 1.") + n = iters // step_size + out_lr = lr * decay_rate ** n + return out_lr + + if lr_decay_type == "cos": + warmup_total_iters = min(max(warmup_iters_ratio * total_iters, 1), 3) + warmup_lr_start = max(warmup_lr_ratio * lr, 1e-6) + no_aug_iter = min(max(no_aug_iter_ratio * total_iters, 1), 15) + func = partial(yolox_warm_cos_lr ,lr, min_lr, total_iters, warmup_total_iters, warmup_lr_start, no_aug_iter) + else: + decay_rate = (min_lr / lr) ** (1 / (step_num - 1)) + step_size = total_iters / step_num + func = partial(step_lr, lr, decay_rate, step_size) + + return func + +def set_optimizer_lr(optimizer, lr_scheduler_func, epoch): + lr = lr_scheduler_func(epoch) + for param_group in optimizer.param_groups: + param_group['lr'] = lr diff --git a/cv/face/arcface/pytorch/nets/iresnet.py b/cv/face/arcface/pytorch/nets/iresnet.py new file mode 100755 index 0000000000000000000000000000000000000000..4f818b56e1b0b1663432a13c92b994ca5fcfce04 --- /dev/null +++ b/cv/face/arcface/pytorch/nets/iresnet.py @@ -0,0 +1,184 @@ + +import torch +from torch import nn + +__all__ = ['iresnet18', 'iresnet34', 'iresnet50', 'iresnet100', 'iresnet200'] + + +def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1): + return nn.Conv2d(in_planes, + out_planes, + kernel_size=3, + stride=stride, + padding=dilation, + groups=groups, + bias=False, + dilation=dilation) + + +def conv1x1(in_planes, out_planes, stride=1): + return nn.Conv2d(in_planes, + out_planes, + kernel_size=1, + stride=stride, + bias=False) + + +class IBasicBlock(nn.Module): + expansion = 1 + def __init__(self, inplanes, planes, stride=1, downsample=None, + groups=1, base_width=64, dilation=1): + super(IBasicBlock, self).__init__() + if groups != 1 or base_width != 64: + raise ValueError('BasicBlock only supports groups=1 and base_width=64') + if dilation > 1: + raise NotImplementedError("Dilation > 1 not supported in BasicBlock") + self.bn1 = nn.BatchNorm2d(inplanes, eps=1e-05,) + self.conv1 = conv3x3(inplanes, planes) + self.bn2 = nn.BatchNorm2d(planes, eps=1e-05,) + self.prelu = nn.PReLU(planes) + self.conv2 = conv3x3(planes, planes, stride) + self.bn3 = nn.BatchNorm2d(planes, eps=1e-05,) + self.downsample = downsample + self.stride = stride + + def forward(self, x): + identity = x + out = self.bn1(x) + out = self.conv1(out) + out = self.bn2(out) + out = self.prelu(out) + out = self.conv2(out) + out = self.bn3(out) + if self.downsample is not None: + identity = self.downsample(x) + out += identity + return out + + +class IResNet(nn.Module): + fc_scale = 7 * 7 + def __init__(self, + block, layers, dropout_keep_prob=0, embedding_size=512, zero_init_residual=False, + groups=1, width_per_group=64, replace_stride_with_dilation=None, fp16=False): + super(IResNet, self).__init__() + self.fp16 = fp16 + self.inplanes = 64 + self.dilation = 1 + if replace_stride_with_dilation is None: + replace_stride_with_dilation = [False, False, False] + if len(replace_stride_with_dilation) != 3: + raise ValueError("replace_stride_with_dilation should be None " + "or a 3-element tuple, got {}".format(replace_stride_with_dilation)) + self.groups = groups + self.base_width = width_per_group + self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=3, stride=1, padding=1, bias=False) + self.bn1 = nn.BatchNorm2d(self.inplanes, eps=1e-05) + self.prelu = nn.PReLU(self.inplanes) + self.layer1 = self._make_layer(block, 64, layers[0], stride=2) + self.layer2 = self._make_layer(block, + 128, + layers[1], + stride=2, + dilate=replace_stride_with_dilation[0]) + self.layer3 = self._make_layer(block, + 256, + layers[2], + stride=2, + dilate=replace_stride_with_dilation[1]) + self.layer4 = self._make_layer(block, + 512, + layers[3], + stride=2, + dilate=replace_stride_with_dilation[2]) + self.bn2 = nn.BatchNorm2d(512 * block.expansion, eps=1e-05,) + self.dropout = nn.Dropout(p=dropout_keep_prob, inplace=True) + self.fc = nn.Linear(512 * block.expansion * self.fc_scale, embedding_size) + self.features = nn.BatchNorm1d(embedding_size, eps=1e-05) + nn.init.constant_(self.features.weight, 1.0) + self.features.weight.requires_grad = False + + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.normal_(m.weight, 0, 0.1) + elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)): + nn.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + + if zero_init_residual: + for m in self.modules(): + if isinstance(m, IBasicBlock): + nn.init.constant_(m.bn2.weight, 0) + + def _make_layer(self, block, planes, blocks, stride=1, dilate=False): + downsample = None + previous_dilation = self.dilation + if dilate: + self.dilation *= stride + stride = 1 + if stride != 1 or self.inplanes != planes * block.expansion: + downsample = nn.Sequential( + conv1x1(self.inplanes, planes * block.expansion, stride), + nn.BatchNorm2d(planes * block.expansion, eps=1e-05, ), + ) + layers = [] + layers.append( + block(self.inplanes, planes, stride, downsample, self.groups, + self.base_width, previous_dilation)) + self.inplanes = planes * block.expansion + for _ in range(1, blocks): + layers.append( + block(self.inplanes, + planes, + groups=self.groups, + base_width=self.base_width, + dilation=self.dilation)) + + return nn.Sequential(*layers) + + def forward(self, x): + x = self.conv1(x) + x = self.bn1(x) + x = self.prelu(x) + x = self.layer1(x) + x = self.layer2(x) + x = self.layer3(x) + x = self.layer4(x) + x = self.bn2(x) + x = torch.flatten(x, 1) + x = self.dropout(x) + x = self.fc(x) + x = self.features(x) + return x + + +def _iresnet(arch, block, layers, pretrained, progress, **kwargs): + model = IResNet(block, layers, **kwargs) + if pretrained: + raise ValueError("No pretrained model for iresnet") + return model + + +def iresnet18(pretrained=False, progress=True, **kwargs): + return _iresnet('iresnet18', IBasicBlock, [2, 2, 2, 2], pretrained, + progress, **kwargs) + + +def iresnet34(pretrained=False, progress=True, **kwargs): + return _iresnet('iresnet34', IBasicBlock, [3, 4, 6, 3], pretrained, + progress, **kwargs) + + +def iresnet50(pretrained=False, progress=True, **kwargs): + return _iresnet('iresnet50', IBasicBlock, [3, 4, 14, 3], pretrained, + progress, **kwargs) + + +def iresnet100(pretrained=False, progress=True, **kwargs): + return _iresnet('iresnet100', IBasicBlock, [3, 13, 30, 3], pretrained, + progress, **kwargs) + + +def iresnet200(pretrained=False, progress=True, **kwargs): + return _iresnet('iresnet200', IBasicBlock, [6, 26, 60, 6], pretrained, + progress, **kwargs) \ No newline at end of file diff --git a/cv/face/arcface/pytorch/nets/mobilefacenet.py b/cv/face/arcface/pytorch/nets/mobilefacenet.py new file mode 100755 index 0000000000000000000000000000000000000000..c39afb4ebaf50409b161ca93e04341bd40328e5c --- /dev/null +++ b/cv/face/arcface/pytorch/nets/mobilefacenet.py @@ -0,0 +1,131 @@ +from torch import nn +from torch.nn import BatchNorm2d, Conv2d, Module, PReLU, Sequential + +class Flatten(Module): + def forward(self, input): + return input.view(input.size(0), -1) + +class Linear_block(Module): + def __init__(self, in_c, out_c, kernel=(1, 1), stride=(1, 1), padding=(0, 0), groups=1): + super(Linear_block, self).__init__() + self.conv = Conv2d(in_c, out_channels=out_c, kernel_size=kernel, groups=groups, stride=stride, padding=padding, bias=False) + self.bn = BatchNorm2d(out_c) + def forward(self, x): + x = self.conv(x) + x = self.bn(x) + return x + +class Residual_Block(Module): + def __init__(self, in_c, out_c, residual = False, kernel=(3, 3), stride=(2, 2), padding=(1, 1), groups=1): + super(Residual_Block, self).__init__() + self.conv = Conv_block(in_c, out_c=groups, kernel=(1, 1), padding=(0, 0), stride=(1, 1)) + self.conv_dw = Conv_block(groups, groups, groups=groups, kernel=kernel, padding=padding, stride=stride) + self.project = Linear_block(groups, out_c, kernel=(1, 1), padding=(0, 0), stride=(1, 1)) + self.residual = residual + def forward(self, x): + if self.residual: + short_cut = x + x = self.conv(x) + x = self.conv_dw(x) + x = self.project(x) + if self.residual: + output = short_cut + x + else: + output = x + return output + +class Residual(Module): + def __init__(self, c, num_block, groups, kernel=(3, 3), stride=(1, 1), padding=(1, 1)): + super(Residual, self).__init__() + modules = [] + for _ in range(num_block): + modules.append(Residual_Block(c, c, residual=True, kernel=kernel, padding=padding, stride=stride, groups=groups)) + self.model = Sequential(*modules) + def forward(self, x): + return self.model(x) + +class Conv_block(Module): + def __init__(self, in_c, out_c, kernel=(1, 1), stride=(1, 1), padding=(0, 0), groups=1): + super(Conv_block, self).__init__() + self.conv = Conv2d(in_c, out_channels=out_c, kernel_size=kernel, groups=groups, stride=stride, padding=padding, bias=False) + self.bn = BatchNorm2d(out_c) + self.prelu = PReLU(out_c) + def forward(self, x): + x = self.conv(x) + x = self.bn(x) + x = self.prelu(x) + return x + +class MobileFaceNet(Module): + def __init__(self, embedding_size): + super(MobileFaceNet, self).__init__() + # 112,112,3 -> 56,56,64 + self.conv1 = Conv_block(3, 64, kernel=(3, 3), stride=(2, 2), padding=(1, 1)) + + # 56,56,64 -> 56,56,64 + self.conv2_dw = Conv_block(64, 64, kernel=(3, 3), stride=(1, 1), padding=(1, 1), groups=64) + + # 56,56,64 -> 28,28,64 + self.conv_23 = Residual_Block(64, 64, kernel=(3, 3), stride=(2, 2), padding=(1, 1), groups=128) + self.conv_3 = Residual(64, num_block=4, groups=128, kernel=(3, 3), stride=(1, 1), padding=(1, 1)) + + # 28,28,64 -> 14,14,128 + self.conv_34 = Residual_Block(64, 128, kernel=(3, 3), stride=(2, 2), padding=(1, 1), groups=256) + self.conv_4 = Residual(128, num_block=6, groups=256, kernel=(3, 3), stride=(1, 1), padding=(1, 1)) + + # 14,14,128 -> 7,7,128 + self.conv_45 = Residual_Block(128, 128, kernel=(3, 3), stride=(2, 2), padding=(1, 1), groups=512) + self.conv_5 = Residual(128, num_block=2, groups=256, kernel=(3, 3), stride=(1, 1), padding=(1, 1)) + + self.sep = nn.Conv2d(128, 512, kernel_size=1, bias=False) + self.sep_bn = nn.BatchNorm2d(512) + self.prelu = nn.PReLU(512) + + self.GDC_dw = nn.Conv2d(512, 512, kernel_size=7, bias=False, groups=512) + self.GDC_bn = nn.BatchNorm2d(512) + + self.features = nn.Conv2d(512, embedding_size, kernel_size=1, bias=False) + self.last_bn = nn.BatchNorm2d(embedding_size) + + self._initialize_weights() + + def _initialize_weights(self): + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') + if m.bias is not None: + m.bias.data.zero_() + elif isinstance(m, nn.BatchNorm2d): + m.weight.data.fill_(1) + m.bias.data.zero_() + elif isinstance(m, nn.Linear): + nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') + if m.bias is not None: + m.bias.data.zero_() + + def forward(self, x): + x = self.conv1(x) + x = self.conv2_dw(x) + x = self.conv_23(x) + x = self.conv_3(x) + x = self.conv_34(x) + x = self.conv_4(x) + x = self.conv_45(x) + x = self.conv_5(x) + + x = self.sep(x) + x = self.sep_bn(x) + x = self.prelu(x) + + x = self.GDC_dw(x) + x = self.GDC_bn(x) + + x = self.features(x) + x = self.last_bn(x) + return x + + +def get_mbf(embedding_size, pretrained): + if pretrained: + raise ValueError("No pretrained model for mobilefacenet") + return MobileFaceNet(embedding_size) diff --git a/cv/face/arcface/pytorch/nets/mobilenet.py b/cv/face/arcface/pytorch/nets/mobilenet.py new file mode 100755 index 0000000000000000000000000000000000000000..b02dca3540e9142fec994ac70a86595946fe1999 --- /dev/null +++ b/cv/face/arcface/pytorch/nets/mobilenet.py @@ -0,0 +1,86 @@ +import torch +import torch.nn as nn + + +def conv_bn(inp, oup, stride = 1): + return nn.Sequential( + nn.Conv2d(inp, oup, 3, stride, 1, bias=False), + nn.BatchNorm2d(oup), + nn.ReLU6(inplace=True) + ) + +def conv_dw(inp, oup, stride = 1): + return nn.Sequential( + nn.Conv2d(inp, inp, 3, stride, 1, groups=inp, bias=False), + nn.BatchNorm2d(inp), + nn.ReLU6(inplace=True), + + nn.Conv2d(inp, oup, 1, 1, 0, bias=False), + nn.BatchNorm2d(oup), + nn.ReLU6(inplace=True), + ) + +class MobileNetV1(nn.Module): + fc_scale = 7 * 7 + def __init__(self, dropout_keep_prob, embedding_size, pretrained): + super(MobileNetV1, self).__init__() + self.stage1 = nn.Sequential( + conv_bn(3, 32, 1), # 3 + conv_dw(32, 64, 1), # 7 + + conv_dw(64, 128, 2), # 11 + conv_dw(128, 128, 1), # 19 + + conv_dw(128, 256, 2), # 27 + conv_dw(256, 256, 1), # 43 + ) + self.stage2 = nn.Sequential( + conv_dw(256, 512, 2), # 43 + 16 = 59 + conv_dw(512, 512, 1), # 59 + 32 = 91 + conv_dw(512, 512, 1), # 91 + 32 = 123 + conv_dw(512, 512, 1), # 123 + 32 = 155 + conv_dw(512, 512, 1), # 155 + 32 = 187 + conv_dw(512, 512, 1), # 187 + 32 = 219 + ) + self.stage3 = nn.Sequential( + conv_dw(512, 1024, 2), # 219 +3 2 = 241 + conv_dw(1024, 1024, 1), # 241 + 64 = 301 + ) + + self.sep = nn.Conv2d(1024, 512, kernel_size=1, bias=False) + self.sep_bn = nn.BatchNorm2d(512) + self.prelu = nn.PReLU(512) + + self.bn2 = nn.BatchNorm2d(512, eps=1e-05) + self.dropout = nn.Dropout(p=dropout_keep_prob, inplace=True) + self.linear = nn.Linear(512 * self.fc_scale, embedding_size) + self.features = nn.BatchNorm1d(embedding_size, eps=1e-05) + + if pretrained: + self.load_state_dict(torch.load("model_data/mobilenet_v1_backbone_weights.pth"), strict = False) + else: + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.normal_(m.weight, 0, 0.1) + elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)): + nn.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + + def forward(self, x): + x = self.stage1(x) + x = self.stage2(x) + x = self.stage3(x) + + x = self.sep(x) + x = self.sep_bn(x) + x = self.prelu(x) + + x = self.bn2(x) + x = torch.flatten(x, 1) + x = self.dropout(x) + x = self.linear(x) + x = self.features(x) + return x + +def get_mobilenet(dropout_keep_prob, embedding_size, pretrained): + return MobileNetV1(dropout_keep_prob, embedding_size, pretrained) diff --git a/cv/face/arcface/pytorch/predict.py b/cv/face/arcface/pytorch/predict.py new file mode 100755 index 0000000000000000000000000000000000000000..be5a6aa07602c6a6f1de24ccd350f49f2e17d9d7 --- /dev/null +++ b/cv/face/arcface/pytorch/predict.py @@ -0,0 +1,44 @@ +from PIL import Image + +from arcface import Arcface + +if __name__ == "__main__": + model = Arcface() + + #----------------------------------------------------------------------------------------------------------# + # mode用于指定测试的模式： + # 'predict'表示单张图片预测，如果想对预测过程进行修改，如保存图片，截取对象等，可以先看下方详细的注释 + # 'fps'表示测试fps，使用的图片是img里面的street.jpg，详情查看下方注释。 + #----------------------------------------------------------------------------------------------------------# + mode = "predict" + #-------------------------------------------------------------------------# + # test_interval 用于指定测量fps的时候，图片检测的次数 + # 理论上test_interval越大，fps越准确。 + # fps_test_image fps测试图片 + #-------------------------------------------------------------------------# + test_interval = 100 + fps_test_image = 'img/1_001.jpg' + + if mode == "predict": + while True: + image_1 = input('Input image_1 filename:') + try: + image_1 = Image.open(image_1) + except: + print('Image_1 Open Error! Try again!') + continue + + image_2 = input('Input image_2 filename:') + try: + image_2 = Image.open(image_2) + except: + print('Image_2 Open Error! Try again!') + continue + + probability = model.detect_image(image_1,image_2) + print(probability) + + elif mode == "fps": + img = Image.open(fps_test_image) + tact_time = model.get_FPS(img, test_interval) + print(str(tact_time) + ' seconds, ' + str(1/tact_time) + 'FPS, @batch_size 1') \ No newline at end of file diff --git a/cv/face/arcface/pytorch/requirements.txt b/cv/face/arcface/pytorch/requirements.txt new file mode 100755 index 0000000000000000000000000000000000000000..2859a26662d631096adfa189ebf667dffe98fa7d --- /dev/null +++ b/cv/face/arcface/pytorch/requirements.txt @@ -0,0 +1,2 @@ +matplotlib +scikit-learn \ No newline at end of file diff --git a/cv/face/arcface/pytorch/run.sh b/cv/face/arcface/pytorch/run.sh new file mode 100755 index 0000000000000000000000000000000000000000..1837fc7ddc4f74a5200b61175443f30a7885bbf6 --- /dev/null +++ b/cv/face/arcface/pytorch/run.sh @@ -0,0 +1,27 @@ +# Copyright (c) 2023, Shanghai Iluvatar CoreX Semiconductor Co., Ltd. +# All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); you may +# not use this file except in compliance with the License. You may obtain +# a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +# License for the specific language governing permissions and limitations +# under the License. + +GPUS=$1 +python3 -m torch.distributed.launch --nproc_per_node=$GPUS --use_env \ + train.py --distributed True \ + --sync_bn False \ + --fp16 False \ + --backbone mobilefacenet \ + --Epoch 100 \ + --batch_size 128 \ + --save_period 1 \ + --save_dir logs + +# python3 train.py --distributed False \ No newline at end of file diff --git a/cv/face/arcface/pytorch/summary.py b/cv/face/arcface/pytorch/summary.py new file mode 100755 index 0000000000000000000000000000000000000000..3a14dfeb640842c436558e059f71f40e2d2368db --- /dev/null +++ b/cv/face/arcface/pytorch/summary.py @@ -0,0 +1,29 @@ +#--------------------------------------------# +# 该部分代码只用于看网络结构，并非测试代码 +#--------------------------------------------# +import torch +from thop import clever_format, profile +from torchsummary import summary + +from nets.arcface import Arcface + +if __name__ == "__main__": + input_shape = [112, 112] + backbone = 'mobilefacenet' + + device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + model = Arcface(num_classes=10575, backbone=backbone, mode="predict").to(device) + summary(model, (3, input_shape[0], input_shape[1])) + + dummy_input = torch.randn(1, 3, input_shape[0], input_shape[1]).to(device) + flops, params = profile(model.to(device), (dummy_input, ), verbose=False) + #--------------------------------------------------------# + # flops * 2是因为profile没有将卷积作为两个operations + # 有些论文将卷积算乘法、加法两个operations。此时乘2 + # 有些论文只考虑乘法的运算次数，忽略加法。此时不乘2 + # 本代码选择乘2，参考YOLOX。 + #--------------------------------------------------------# + flops = flops * 2 + flops, params = clever_format([flops, params], "%.3f") + print('Total GFLOPS: %s' % (flops)) + print('Total params: %s' % (params)) \ No newline at end of file diff --git a/cv/face/arcface/pytorch/train.py b/cv/face/arcface/pytorch/train.py new file mode 100755 index 0000000000000000000000000000000000000000..10ae6422dfa8a52fcd62d1f88e6c18edcfd1d477 --- /dev/null +++ b/cv/face/arcface/pytorch/train.py @@ -0,0 +1,363 @@ +# Copyright (c) 2022, Shanghai Iluvatar CoreX Semiconductor Co., Ltd. +import os +import argparse + +import numpy as np +import torch +import torch.backends.cudnn as cudnn +import torch.distributed as dist +import torch.optim as optim +from torch.utils.data import DataLoader + +from nets.arcface import Arcface +from nets.arcface_training import get_lr_scheduler, set_optimizer_lr +from utils.callback import LossHistory +from utils.dataloader import FacenetDataset, LFWDataset, dataset_collate +from utils.utils import get_num_classes, show_config +from utils.utils_fit import fit_one_epoch + +parser = argparse.ArgumentParser(add_help=False) +parser.add_argument("--distributed", default=True, type=bool, help="if use ddp") +parser.add_argument("--sync_bn", default=False, type=bool, help="sync_bn for ddp") +parser.add_argument("--fp16", default=False, type=bool, help="whether use fp16") +parser.add_argument("--backbone", default="mobilefacenet", type=str, help="backbone") +parser.add_argument("--pretrained", default=False, type=bool, help="state_dict for backbone") +parser.add_argument("--load_from", default="", type=str, help="state_dict for whole model") + +parser.add_argument("--Epoch", default=100, type=int, help="state_dict for whole model") +parser.add_argument("--batch_size", default=128, type=int, help="batchsize for single gpu") +parser.add_argument("--save_period", default=1, type=int, help="save_period") +parser.add_argument("--save_dir", default="logs", type=str, help="save_dir") + +args = parser.parse_args() + + +if __name__ == "__main__": + #-------------------------------# + # 是否使用Cuda + # 没有GPU可以设置成False + #-------------------------------# + Cuda = True + #---------------------------------------------------------------------# + # distributed 用于指定是否使用单机多卡分布式运行 + # 终端指令仅支持Ubuntu。CUDA_VISIBLE_DEVICES用于在Ubuntu下指定显卡。 + # Windows系统下默认使用DP模式调用所有显卡，不支持DDP。 + # DP模式： + # 设置 distributed = False + # 在终端中输入 CUDA_VISIBLE_DEVICES=0,1 python train.py + # DDP模式： + # 设置 distributed = True + # 在终端中输入 CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 train.py + #---------------------------------------------------------------------# + distributed = args.distributed + #---------------------------------------------------------------------# + # sync_bn 是否使用sync_bn，DDP模式多卡可用 + #---------------------------------------------------------------------# + sync_bn = args.sync_bn + #---------------------------------------------------------------------# + # fp16 是否使用混合精度训练 + # 可减少约一半的显存、需要pytorch1.7.1以上 + #---------------------------------------------------------------------# + fp16 = args.fp16 + #--------------------------------------------------------# + # 指向根目录下的cls_train.txt，读取人脸路径与标签 + #--------------------------------------------------------# + annotation_path = "cls_train.txt" + #--------------------------------------------------------# + # 输入图像大小 + #--------------------------------------------------------# + input_shape = [112, 112, 3] + #--------------------------------------------------------# + # 主干特征提取网络的选择 + # mobilefacenet + # mobilenetv1 + # iresnet18 + # iresnet34 + # iresnet50 + # iresnet100 + # iresnet200 + # + # 除了mobilenetv1外，其它的backbone均可从0开始训练。 + # 这是由于mobilenetv1没有残差边，收敛速度慢，因此建议： + # 如果使用mobilenetv1为主干, 则设置pretrain = True + # 如果使用其它网络为主干， 则设置pretrain = False + #--------------------------------------------------------# + backbone = "mobilefacenet" + #----------------------------------------------------------------------------------------------------------------------------# + # 如果训练过程中存在中断训练的操作，可以将model_path设置成logs文件夹下的权值文件，将已经训练了一部分的权值再次载入。 + # 同时修改下方的训练的参数，来保证模型epoch的连续性。 + # + # 当model_path = ''的时候不加载整个模型的权值。 + # + # 此处使用的是整个模型的权重，因此是在train.py进行加载的，pretrain不影响此处的权值加载。 + # 如果想要让模型从主干的预训练权值开始训练，则设置model_path = ''，pretrain = True，此时仅加载主干。 + # 如果想要让模型从0开始训练，则设置model_path = ''，pretrain = Fasle，此时从0开始训练。 + #----------------------------------------------------------------------------------------------------------------------------# + model_path = args.load_from + #----------------------------------------------------------------------------------------------------------------------------# + # 是否使用主干网络的预训练权重，此处使用的是主干的权重，因此是在模型构建的时候进行加载的。 + # 如果设置了model_path，则主干的权值无需加载，pretrained的值无意义。 + # 如果不设置model_path，pretrained = True，此时仅加载主干开始训练。 + # 如果不设置model_path，pretrained = False，此时从0开始训练。 + # 除了mobilenetv1外，其它的backbone均未提供预训练权重。 + #----------------------------------------------------------------------------------------------------------------------------# + pretrained = args.pretrained + + #----------------------------------------------------------------------------------------------------------------------------# + # 显存不足与数据集大小无关，提示显存不足请调小batch_size。 + # 受到BatchNorm层影响，不能为1。 + # + # 在此提供若干参数设置建议，各位训练者根据自己的需求进行灵活调整： + # （一）从预训练权重开始训练： + # Adam： + # Init_Epoch = 0，Epoch = 100，optimizer_type = 'adam'，Init_lr = 1e-3，weight_decay = 0。 + # SGD： + # Init_Epoch = 0，Epoch = 100，optimizer_type = 'sgd'，Init_lr = 1e-2，weight_decay = 5e-4。 + # 其中：UnFreeze_Epoch可以在100-300之间调整。 + # （二）batch_size的设置： + # 在显卡能够接受的范围内，以大为好。显存不足与数据集大小无关，提示显存不足（OOM或者CUDA out of memory）请调小batch_size。 + # 受到BatchNorm层影响，batch_size最小为2，不能为1。 + # 正常情况下Freeze_batch_size建议为Unfreeze_batch_size的1-2倍。不建议设置的差距过大，因为关系到学习率的自动调整。 + #----------------------------------------------------------------------------------------------------------------------------# + #------------------------------------------------------# + # 训练参数 + # Init_Epoch 模型当前开始的训练世代 + # Epoch 模型总共训练的epoch + # batch_size 每次输入的图片数量 + #------------------------------------------------------# + Init_Epoch = 0 + Epoch = args.Epoch + batch_size = args.batch_size + + #------------------------------------------------------------------# + # 其它训练参数：学习率、优化器、学习率下降有关 + #------------------------------------------------------------------# + #------------------------------------------------------------------# + # Init_lr 模型的最大学习率 + # Min_lr 模型的最小学习率，默认为最大学习率的0.01 + #------------------------------------------------------------------# + Init_lr = 1e-2 + Min_lr = Init_lr * 0.01 + #------------------------------------------------------------------# + # optimizer_type 使用到的优化器种类，可选的有adam、sgd + # 当使用Adam优化器时建议设置 Init_lr=1e-3 + # 当使用SGD优化器时建议设置 Init_lr=1e-2 + # momentum 优化器内部使用到的momentum参数 + # weight_decay 权值衰减，可防止过拟合 + # adam会导致weight_decay错误，使用adam时建议设置为0。 + #------------------------------------------------------------------# + optimizer_type = "sgd" + momentum = 0.9 + weight_decay = 5e-4 + #------------------------------------------------------------------# + # lr_decay_type 使用到的学习率下降方式，可选的有step、cos + #------------------------------------------------------------------# + lr_decay_type = "cos" + #------------------------------------------------------------------# + # save_period 多少个epoch保存一次权值，默认每个世代都保存 + #------------------------------------------------------------------# + save_period = args.save_period + #------------------------------------------------------------------# + # save_dir 权值与日志文件保存的文件夹 + #------------------------------------------------------------------# + save_dir = args.save_dir + #------------------------------------------------------------------# + # 用于设置是否使用多线程读取数据 + # 开启后会加快数据读取速度，但是会占用更多内存 + # 内存较小的电脑可以设置为2或者0 + #------------------------------------------------------------------# + num_workers = 4 + #------------------------------------------------------------------# + # 是否开启LFW评估 + #------------------------------------------------------------------# + lfw_eval_flag = True + #------------------------------------------------------------------# + # LFW评估数据集的文件路径和对应的txt文件 + #------------------------------------------------------------------# + lfw_dir_path = "datasets/lfw" + lfw_pairs_path = "datasets/lfw_pair.txt" + + #------------------------------------------------------# + # 设置用到的显卡 + #------------------------------------------------------# + ngpus_per_node = torch.cuda.device_count() + if distributed: + if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ: + rank = int(os.environ["RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + local_rank = int(os.environ['LOCAL_RANK']) + else: + print('Not using distributed mode') + import sys + sys.exit() + + dist_backend = "nccl" + print('| distributed init (rank {}) (size {})'.format(rank, world_size), flush=True) + torch.distributed.init_process_group(backend=dist_backend, init_method='env://', + world_size=world_size, rank=local_rank) + + torch.cuda.set_device(local_rank) + torch.distributed.barrier() + + num_classes = get_num_classes(annotation_path) + #---------------------------------# + # 载入模型并加载预训练权重 + #---------------------------------# + model = Arcface(num_classes=num_classes, backbone=backbone, pretrained=pretrained) + + if model_path != '': + #------------------------------------------------------# + # 权值文件请看README，百度网盘下载 + #------------------------------------------------------# + if local_rank == 0: + print('Load weights {}.'.format(model_path)) + + #------------------------------------------------------# + # 根据预训练权重的Key和模型的Key进行加载 + #------------------------------------------------------# + model_dict = model.state_dict() + pretrained_dict = torch.load(model_path, map_location = device) + load_key, no_load_key, temp_dict = [], [], {} + for k, v in pretrained_dict.items(): + if k in model_dict.keys() and np.shape(model_dict[k]) == np.shape(v): + temp_dict[k] = v + load_key.append(k) + else: + no_load_key.append(k) + model_dict.update(temp_dict) + model.load_state_dict(model_dict) + #------------------------------------------------------# + # 显示没有匹配上的Key + #------------------------------------------------------# + if local_rank == 0: + print("\nSuccessful Load Key:", str(load_key)[:500], "……\nSuccessful Load Key Num:", len(load_key)) + print("\nFail To Load Key:", str(no_load_key)[:500], "……\nFail To Load Key num:", len(no_load_key)) + print("\n\033[1;33;44m温馨提示，head部分没有载入是正常现象，Backbone部分没有载入是错误的。\033[0m") + + #----------------------# + # 记录Loss + #----------------------# + if local_rank == 0: + loss_history = LossHistory(save_dir, model, input_shape=input_shape) + else: + loss_history = None + + #------------------------------------------------------------------# + # torch 1.2不支持amp，建议使用torch 1.7.1及以上正确使用fp16 + # 因此torch1.2这里显示"could not be resolve" + #------------------------------------------------------------------# + if fp16: + from torch.cuda.amp import GradScaler as GradScaler + scaler = GradScaler() + else: + scaler = None + + model_train = model.train() + #----------------------------# + # 多卡同步Bn + #----------------------------# + if sync_bn and ngpus_per_node > 1 and distributed: + model_train = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model_train) + elif sync_bn: + print("Sync_bn is not support in one gpu or not distributed.") + + if Cuda: + if distributed: + #----------------------------# + # 多卡平行运行 + #----------------------------# + model_train = model_train.cuda(local_rank) + model_train = torch.nn.parallel.DistributedDataParallel(model_train, device_ids=[local_rank], find_unused_parameters=True) + else: + model_train = torch.nn.DataParallel(model) + cudnn.benchmark = True + model_train = model_train.cuda() + + #---------------------------------# + # LFW估计 + #---------------------------------# + LFW_loader = torch.utils.data.DataLoader( + LFWDataset(dir=lfw_dir_path, pairs_path=lfw_pairs_path, image_size=input_shape), batch_size=32, shuffle=False) if lfw_eval_flag else None + + #-------------------------------------------------------# + # 0.01用于验证，0.99用于训练 + #-------------------------------------------------------# + val_split = 0.01 + with open(annotation_path,"r") as f: + lines = f.readlines() + np.random.seed(10101) + np.random.shuffle(lines) + np.random.seed(None) + num_val = int(len(lines)*val_split) + num_train = len(lines) - num_val + + show_config( + num_classes = num_classes, backbone = backbone, model_path = model_path, input_shape = input_shape, \ + Init_Epoch = Init_Epoch, Epoch = Epoch, batch_size = batch_size, \ + Init_lr = Init_lr, Min_lr = Min_lr, optimizer_type = optimizer_type, momentum = momentum, lr_decay_type = lr_decay_type, \ + save_period = save_period, save_dir = save_dir, num_workers = num_workers, num_train = num_train, num_val = num_val + ) + + if True: + #-------------------------------------------------------------------# + # 判断当前batch_size，自适应调整学习率 + #-------------------------------------------------------------------# + nbs = 64 + lr_limit_max = 1e-3 if optimizer_type == 'adam' else 1e-1 + lr_limit_min = 3e-4 if optimizer_type == 'adam' else 5e-4 + Init_lr_fit = min(max(batch_size / nbs * Init_lr, lr_limit_min), lr_limit_max) + Min_lr_fit = min(max(batch_size / nbs * Min_lr, lr_limit_min * 1e-2), lr_limit_max * 1e-2) + + #---------------------------------------# + # 根据optimizer_type选择优化器 + #---------------------------------------# + optimizer = { + 'adam' : optim.Adam(model.parameters(), Init_lr_fit, betas = (momentum, 0.999), weight_decay = weight_decay), + 'sgd' : optim.SGD(model.parameters(), Init_lr_fit, momentum=momentum, nesterov=True, weight_decay = weight_decay) + }[optimizer_type] + + #---------------------------------------# + # 获得学习率下降的公式 + #---------------------------------------# + lr_scheduler_func = get_lr_scheduler(lr_decay_type, Init_lr_fit, Min_lr_fit, Epoch) + + #---------------------------------------# + # 判断每一个世代的长度 + #---------------------------------------# + epoch_step = num_train // batch_size + epoch_step_val = num_val // batch_size + + if epoch_step == 0 or epoch_step_val == 0: + raise ValueError("数据集过小，无法继续进行训练，请扩充数据集。") + + #---------------------------------------# + # 构建数据集加载器。 + #---------------------------------------# + train_dataset = FacenetDataset(input_shape, lines[:num_train], random = True) + val_dataset = FacenetDataset(input_shape, lines[num_train:], random = False) + + if distributed: + train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset, shuffle=True,) + val_sampler = torch.utils.data.distributed.DistributedSampler(val_dataset, shuffle=False,) + batch_size = batch_size // ngpus_per_node + shuffle = False + else: + train_sampler = None + val_sampler = None + shuffle = True + + gen = DataLoader(train_dataset, shuffle=shuffle, batch_size=batch_size, num_workers=num_workers, pin_memory=True, + drop_last=True, collate_fn=dataset_collate, sampler=train_sampler) + gen_val = DataLoader(val_dataset, shuffle=shuffle, batch_size=batch_size, num_workers=num_workers, pin_memory=True, + drop_last=True, collate_fn=dataset_collate, sampler=val_sampler) + + for epoch in range(Init_Epoch, Epoch): + if distributed: + train_sampler.set_epoch(epoch) + + set_optimizer_lr(optimizer, lr_scheduler_func, epoch) + + fit_one_epoch(model_train, model, loss_history, optimizer, epoch, epoch_step, epoch_step_val, gen, gen_val, Epoch, Cuda, LFW_loader, lfw_eval_flag, fp16, scaler, save_period, save_dir, local_rank) + + if local_rank == 0: + loss_history.writer.close() diff --git a/cv/face/arcface/pytorch/txt_annotation.py b/cv/face/arcface/pytorch/txt_annotation.py new file mode 100755 index 0000000000000000000000000000000000000000..c79725a9a04bd48e7bcbea2cd0828d30220b0c2d --- /dev/null +++ b/cv/face/arcface/pytorch/txt_annotation.py @@ -0,0 +1,27 @@ +# Copyright (c) 2022, Shanghai Iluvatar CoreX Semiconductor Co., Ltd. + +#------------------------------------------------# +# 进行训练前需要利用这个文件生成cls_train.txt +#------------------------------------------------# +import os + +if __name__ == "__main__": + #---------------------# + # 训练集所在的路径 + #---------------------# + datasets_path = "datasets/datasets" + + types_name = os.listdir(datasets_path) + types_name = sorted(types_name) + + list_file = open('cls_train.txt', 'w') + for cls_id, type_name in enumerate(types_name): + photos_path = os.path.join(datasets_path, type_name) + if not os.path.isdir(photos_path): + continue + photos_name = os.listdir(photos_path) + + for photo_name in photos_name: + list_file.write(str(cls_id) + ";" + '%s'%(os.path.join(os.path.abspath(datasets_path), type_name, photo_name))) + list_file.write('\n') + list_file.close() diff --git a/cv/face/arcface/pytorch/utils/__init__.py b/cv/face/arcface/pytorch/utils/__init__.py new file mode 100755 index 0000000000000000000000000000000000000000..4287ca8617970fa8fc025b75cb319c7032706910 --- /dev/null +++ b/cv/face/arcface/pytorch/utils/__init__.py @@ -0,0 +1 @@ +# \ No newline at end of file diff --git a/cv/face/arcface/pytorch/utils/callback.py b/cv/face/arcface/pytorch/utils/callback.py new file mode 100755 index 0000000000000000000000000000000000000000..386f88f75cf820d03b38ffd303ce801b19890d78 --- /dev/null +++ b/cv/face/arcface/pytorch/utils/callback.py @@ -0,0 +1,85 @@ +import datetime +import os + +import torch +import matplotlib +matplotlib.use('Agg') +import scipy.signal +from matplotlib import pyplot as plt +from torch.utils.tensorboard import SummaryWriter + +class LossHistory(): + def __init__(self, log_dir, model, input_shape): + time_str = datetime.datetime.strftime(datetime.datetime.now(),'%Y_%m_%d_%H_%M_%S') + self.log_dir = os.path.join(log_dir, "loss_" + str(time_str)) + self.acc = [] + self.losses = [] + self.val_loss = [] + + os.makedirs(self.log_dir) + self.writer = SummaryWriter(self.log_dir) + dummy_input = torch.randn(2, 3, input_shape[0], input_shape[1]) + self.writer.add_graph(model, dummy_input) + + def append_loss(self, epoch, acc, loss, val_loss): + if not os.path.exists(self.log_dir): + os.makedirs(self.log_dir) + + self.acc.append(acc) + self.losses.append(loss) + self.val_loss.append(val_loss) + + with open(os.path.join(self.log_dir, "epoch_acc.txt"), 'a') as f: + f.write(str(acc)) + f.write("\n") + with open(os.path.join(self.log_dir, "epoch_loss.txt"), 'a') as f: + f.write(str(loss)) + f.write("\n") + with open(os.path.join(self.log_dir, "epoch_val_loss.txt"), 'a') as f: + f.write(str(val_loss)) + f.write("\n") + + self.writer.add_scalar('loss', loss, epoch) + self.writer.add_scalar('val_loss', val_loss, epoch) + self.loss_plot() + + def loss_plot(self): + iters = range(len(self.losses)) + + plt.figure() + plt.plot(iters, self.losses, 'red', linewidth = 2, label='train loss') + plt.plot(iters, self.val_loss, 'coral', linewidth = 2, label='val loss') + try: + if len(self.losses) < 25: + num = 5 + else: + num = 15 + plt.plot(iters, scipy.signal.savgol_filter(self.losses, num, 3), 'green', linestyle = '--', linewidth = 2, label='smooth train loss') + plt.plot(iters, scipy.signal.savgol_filter(self.val_loss, num, 3), '#8B4513', linestyle = '--', linewidth = 2, label='smooth val loss') + except: + pass + plt.grid(True) + plt.xlabel('Epoch') + plt.ylabel('Loss') + plt.legend(loc="upper right") + plt.savefig(os.path.join(self.log_dir, "epoch_loss.png")) + plt.cla() + plt.close("all") + + plt.figure() + plt.plot(iters, self.acc, 'red', linewidth = 2, label='lfw acc') + try: + if len(self.losses) < 25: + num = 5 + else: + num = 15 + plt.plot(iters, scipy.signal.savgol_filter(self.acc, num, 3), 'green', linestyle = '--', linewidth = 2, label='smooth lfw acc') + except: + pass + plt.grid(True) + plt.xlabel('Epoch') + plt.ylabel('Lfw Acc') + plt.legend(loc="upper right") + plt.savefig(os.path.join(self.log_dir, "epoch_acc.png")) + plt.cla() + plt.close("all") diff --git a/cv/face/arcface/pytorch/utils/dataloader.py b/cv/face/arcface/pytorch/utils/dataloader.py new file mode 100755 index 0000000000000000000000000000000000000000..9836fda1ad6a8187a7cc103129e137b9d283394d --- /dev/null +++ b/cv/face/arcface/pytorch/utils/dataloader.py @@ -0,0 +1,105 @@ +import os + +import numpy as np +import torch +import torch.utils.data as data +import torchvision.datasets as datasets +from PIL import Image + +from .utils import cvtColor, preprocess_input, resize_image + + +class FacenetDataset(data.Dataset): + def __init__(self, input_shape, lines, random): + self.input_shape = input_shape + self.lines = lines + self.random = random + + def __len__(self): + return len(self.lines) + + def rand(self, a=0, b=1): + return np.random.rand()*(b-a) + a + + def __getitem__(self, index): + annotation_path = self.lines[index].split(';')[1].split()[0] + y = int(self.lines[index].split(';')[0]) + + image = cvtColor(Image.open(annotation_path)) + #------------------------------------------# + # 翻转图像 + #------------------------------------------# + if self.rand()<.5 and self.random: + image = image.transpose(Image.FLIP_LEFT_RIGHT) + image = resize_image(image, [self.input_shape[1], self.input_shape[0]], letterbox_image = True) + + image = np.transpose(preprocess_input(np.array(image, dtype='float32')), (2, 0, 1)) + return image, y + +def dataset_collate(batch): + images = [] + targets = [] + for image, y in batch: + images.append(image) + targets.append(y) + images = torch.from_numpy(np.array(images)).type(torch.FloatTensor) + targets = torch.from_numpy(np.array(targets)).long() + return images, targets + +class LFWDataset(datasets.ImageFolder): + def __init__(self, dir, pairs_path, image_size, transform=None): + super(LFWDataset, self).__init__(dir,transform) + self.image_size = image_size + self.pairs_path = pairs_path + self.validation_images = self.get_lfw_paths(dir) + + def read_lfw_pairs(self,pairs_filename): + pairs = [] + with open(pairs_filename, 'r') as f: + for line in f.readlines()[1:]: + pair = line.strip().split() + pairs.append(pair) + return np.array(pairs) + + def get_lfw_paths(self,lfw_dir,file_ext="jpg"): + + pairs = self.read_lfw_pairs(self.pairs_path) + + nrof_skipped_pairs = 0 + path_list = [] + issame_list = [] + + for i in range(len(pairs)): + #for pair in pairs: + pair = pairs[i] + if len(pair) == 3: + path0 = os.path.join(lfw_dir, pair[0], pair[0] + '_' + '%04d' % int(pair[1])+'.'+file_ext) + path1 = os.path.join(lfw_dir, pair[0], pair[0] + '_' + '%04d' % int(pair[2])+'.'+file_ext) + issame = True + elif len(pair) == 4: + path0 = os.path.join(lfw_dir, pair[0], pair[0] + '_' + '%04d' % int(pair[1])+'.'+file_ext) + path1 = os.path.join(lfw_dir, pair[2], pair[2] + '_' + '%04d' % int(pair[3])+'.'+file_ext) + issame = False + if os.path.exists(path0) and os.path.exists(path1): # Only add the pair if both paths exist + path_list.append((path0,path1,issame)) + issame_list.append(issame) + else: + nrof_skipped_pairs += 1 + if nrof_skipped_pairs>0: + print('Skipped %d image pairs' % nrof_skipped_pairs) + + return path_list + + def __getitem__(self, index): + (path_1, path_2, issame) = self.validation_images[index] + image1, image2 = Image.open(path_1), Image.open(path_2) + + image1 = resize_image(image1, [self.image_size[1], self.image_size[0]], letterbox_image = True) + image2 = resize_image(image2, [self.image_size[1], self.image_size[0]], letterbox_image = True) + + image1, image2 = np.transpose(preprocess_input(np.array(image1, np.float32)),[2, 0, 1]), np.transpose(preprocess_input(np.array(image2, np.float32)),[2, 0, 1]) + + return image1, image2, issame + + def __len__(self): + return len(self.validation_images) diff --git a/cv/face/arcface/pytorch/utils/utils.py b/cv/face/arcface/pytorch/utils/utils.py new file mode 100755 index 0000000000000000000000000000000000000000..c8ddd431cd7ae77bf1fe632894444d1b3318740d --- /dev/null +++ b/cv/face/arcface/pytorch/utils/utils.py @@ -0,0 +1,65 @@ +import numpy as np +import torch +from PIL import Image + +#---------------------------------------------------------# +# 将图像转换成RGB图像，防止灰度图在预测时报错。 +# 代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB +#---------------------------------------------------------# +def cvtColor(image): + if len(np.shape(image)) == 3 and np.shape(image)[2] == 3: + return image + else: + image = image.convert('RGB') + return image + +#---------------------------------------------------# +# 对输入图像进行resize +#---------------------------------------------------# +def resize_image(image, size, letterbox_image): + iw, ih = image.size + w, h = size + if letterbox_image: + scale = min(w/iw, h/ih) + nw = int(iw*scale) + nh = int(ih*scale) + + image = image.resize((nw,nh), Image.BICUBIC) + new_image = Image.new('RGB', size, (128,128,128)) + new_image.paste(image, ((w-nw)//2, (h-nh)//2)) + else: + new_image = image.resize((w, h), Image.BICUBIC) + return new_image + +def get_num_classes(annotation_path): + with open(annotation_path) as f: + dataset_path = f.readlines() + + labels = [] + for path in dataset_path: + path_split = path.split(";") + labels.append(int(path_split[0])) + num_classes = np.max(labels) + 1 + return num_classes + +#---------------------------------------------------# +# 获得学习率 +#---------------------------------------------------# +def get_lr(optimizer): + for param_group in optimizer.param_groups: + return param_group['lr'] + +def preprocess_input(image): + image /= 255.0 + image -= 0.5 + image /= 0.5 + return image + +def show_config(**kwargs): + print('Configurations:') + print('-' * 70) + print('|%25s | %40s|' % ('keys', 'values')) + print('-' * 70) + for key, value in kwargs.items(): + print('|%25s | %40s|' % (str(key), str(value))) + print('-' * 70) \ No newline at end of file diff --git a/cv/face/arcface/pytorch/utils/utils_fit.py b/cv/face/arcface/pytorch/utils/utils_fit.py new file mode 100755 index 0000000000000000000000000000000000000000..121153205c574ab3d4f7522a6f5129ebb4e7d3b8 --- /dev/null +++ b/cv/face/arcface/pytorch/utils/utils_fit.py @@ -0,0 +1,126 @@ +import os + +import numpy as np +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.nn.functional as F +from tqdm import tqdm + +from .utils import get_lr +from .utils_metrics import evaluate + + +def fit_one_epoch(model_train, model, loss_history, optimizer, epoch, epoch_step, epoch_step_val, gen, gen_val, Epoch, cuda, test_loader, lfw_eval_flag, fp16, scaler, save_period, save_dir, local_rank=0): + total_loss = 0 + total_accuracy = 0 + + val_total_loss = 0 + val_total_accuracy = 0 + + if local_rank == 0: + print('Start Train') + pbar = tqdm(total=epoch_step,desc=f'Epoch {epoch + 1}/{Epoch}',postfix=dict,mininterval=0.3) + model_train.train() + for iteration, batch in enumerate(gen): + if iteration >= epoch_step: + break + images, labels = batch + with torch.no_grad(): + if cuda: + images = images.cuda(local_rank) + labels = labels.cuda(local_rank) + + #----------------------# + # 清零梯度 + #----------------------# + optimizer.zero_grad() + if not fp16: + outputs = model_train(images, labels, mode="train") + loss = nn.NLLLoss()(F.log_softmax(outputs, -1), labels) + + loss.backward() + optimizer.step() + else: + from torch.cuda.amp import autocast + with autocast(): + outputs = model_train(images, labels, mode="train") + loss = nn.NLLLoss()(F.log_softmax(outputs, -1), labels) + #----------------------# + # 反向传播 + #----------------------# + scaler.scale(loss).backward() + scaler.step(optimizer) + scaler.update() + + with torch.no_grad(): + accuracy = torch.mean((torch.argmax(F.softmax(outputs, dim=-1), dim=-1) == labels).type(torch.FloatTensor)) + + total_loss += loss.item() + total_accuracy += accuracy.item() + + if local_rank == 0: + pbar.set_postfix(**{'total_loss': total_loss / (iteration + 1), + 'accuracy' : total_accuracy / (iteration + 1), + 'lr' : get_lr(optimizer)}) + pbar.update(1) + + if local_rank == 0: + pbar.close() + print('Finish Train') + print('Start Validation') + pbar = tqdm(total=epoch_step_val, desc=f'Epoch {epoch + 1}/{Epoch}',postfix=dict,mininterval=0.3) + model_train.eval() + for iteration, batch in enumerate(gen_val): + if iteration >= epoch_step_val: + break + images, labels = batch + with torch.no_grad(): + if cuda: + images = images.cuda(local_rank) + labels = labels.cuda(local_rank) + + optimizer.zero_grad() + outputs = model_train(images, labels, mode="train") + loss = nn.NLLLoss()(F.log_softmax(outputs, -1), labels) + + accuracy = torch.mean((torch.argmax(F.softmax(outputs, dim=-1), dim=-1) == labels).type(torch.FloatTensor)) + + val_total_loss += loss.item() + val_total_accuracy += accuracy.item() + + if local_rank == 0: + pbar.set_postfix(**{'total_loss': val_total_loss / (iteration + 1), + 'accuracy' : val_total_accuracy / (iteration + 1), + 'lr' : get_lr(optimizer)}) + pbar.update(1) + + if lfw_eval_flag: + print("开始进行LFW数据集的验证。") + labels, distances = [], [] + for _, (data_a, data_p, label) in enumerate(test_loader): + with torch.no_grad(): + data_a, data_p = data_a.type(torch.FloatTensor), data_p.type(torch.FloatTensor) + if cuda: + data_a, data_p = data_a.cuda(local_rank), data_p.cuda(local_rank) + + out_a, out_p = model_train(data_a), model_train(data_p) + dists = torch.sqrt(torch.sum((out_a - out_p) ** 2, 1)) + distances.append(dists.data.cpu().numpy()) + labels.append(label.data.cpu().numpy()) + + labels = np.array([sublabel for label in labels for sublabel in label]) + distances = np.array([subdist for dist in distances for subdist in dist]) + _, _, accuracy, _, _, _, _ = evaluate(distances,labels) + + if local_rank == 0: + pbar.close() + print('Finish Validation') + + if lfw_eval_flag: + print('LFW_Accuracy: %2.5f+-%2.5f' % (np.mean(accuracy), np.std(accuracy))) + + loss_history.append_loss(epoch, np.mean(accuracy) if lfw_eval_flag else total_accuracy / epoch_step, total_loss / epoch_step, val_total_loss / epoch_step_val) + print('Total Loss: %.4f' % (total_loss / epoch_step)) + if (epoch + 1) % save_period == 0 or epoch + 1 == Epoch: + torch.save(model.state_dict(), os.path.join(save_dir, 'ep%03d-loss%.3f-val_loss%.3f.pth'%((epoch+1), total_loss / epoch_step, val_total_loss / epoch_step_val))) diff --git a/cv/face/arcface/pytorch/utils/utils_metrics.py b/cv/face/arcface/pytorch/utils/utils_metrics.py new file mode 100755 index 0000000000000000000000000000000000000000..110fc2bc16a38cf662abdff124c8c504ef099686 --- /dev/null +++ b/cv/face/arcface/pytorch/utils/utils_metrics.py @@ -0,0 +1,157 @@ +import numpy as np +import torch +from scipy import interpolate +from sklearn.model_selection import KFold +from tqdm import tqdm + +def evaluate(distances, labels, nrof_folds=10): + # Calculate evaluation metrics + thresholds = np.arange(0, 4, 0.01) + tpr, fpr, accuracy, best_thresholds = calculate_roc(thresholds, distances, + labels, nrof_folds=nrof_folds) + thresholds = np.arange(0, 4, 0.001) + val, val_std, far = calculate_val(thresholds, distances, + labels, 1e-3, nrof_folds=nrof_folds) + return tpr, fpr, accuracy, val, val_std, far, best_thresholds + +def calculate_roc(thresholds, distances, labels, nrof_folds=10): + + nrof_pairs = min(len(labels), len(distances)) + nrof_thresholds = len(thresholds) + k_fold = KFold(n_splits=nrof_folds, shuffle=False) + + tprs = np.zeros((nrof_folds,nrof_thresholds)) + fprs = np.zeros((nrof_folds,nrof_thresholds)) + accuracy = np.zeros((nrof_folds)) + + indices = np.arange(nrof_pairs) + + for fold_idx, (train_set, test_set) in enumerate(k_fold.split(indices)): + + # Find the best threshold for the fold + acc_train = np.zeros((nrof_thresholds)) + for threshold_idx, threshold in enumerate(thresholds): + _, _, acc_train[threshold_idx] = calculate_accuracy(threshold, distances[train_set], labels[train_set]) + + best_threshold_index = np.argmax(acc_train) + for threshold_idx, threshold in enumerate(thresholds): + tprs[fold_idx,threshold_idx], fprs[fold_idx,threshold_idx], _ = calculate_accuracy(threshold, distances[test_set], labels[test_set]) + _, _, accuracy[fold_idx] = calculate_accuracy(thresholds[best_threshold_index], distances[test_set], labels[test_set]) + tpr = np.mean(tprs,0) + fpr = np.mean(fprs,0) + return tpr, fpr, accuracy, thresholds[best_threshold_index] + +def calculate_accuracy(threshold, dist, actual_issame): + predict_issame = np.less(dist, threshold) + tp = np.sum(np.logical_and(predict_issame, actual_issame)) + fp = np.sum(np.logical_and(predict_issame, np.logical_not(actual_issame))) + tn = np.sum(np.logical_and(np.logical_not(predict_issame), np.logical_not(actual_issame))) + fn = np.sum(np.logical_and(np.logical_not(predict_issame), actual_issame)) + + tpr = 0 if (tp+fn==0) else float(tp) / float(tp+fn) + fpr = 0 if (fp+tn==0) else float(fp) / float(fp+tn) + acc = float(tp+tn)/dist.size + return tpr, fpr, acc + +def calculate_val(thresholds, distances, labels, far_target=1e-3, nrof_folds=10): + nrof_pairs = min(len(labels), len(distances)) + nrof_thresholds = len(thresholds) + k_fold = KFold(n_splits=nrof_folds, shuffle=False) + + val = np.zeros(nrof_folds) + far = np.zeros(nrof_folds) + + indices = np.arange(nrof_pairs) + + for fold_idx, (train_set, test_set) in enumerate(k_fold.split(indices)): + # Find the threshold that gives FAR = far_target + far_train = np.zeros(nrof_thresholds) + for threshold_idx, threshold in enumerate(thresholds): + _, far_train[threshold_idx] = calculate_val_far(threshold, distances[train_set], labels[train_set]) + if np.max(far_train)>=far_target: + f = interpolate.interp1d(far_train, thresholds, kind='slinear') + threshold = f(far_target) + else: + threshold = 0.0 + + val[fold_idx], far[fold_idx] = calculate_val_far(threshold, distances[test_set], labels[test_set]) + + val_mean = np.mean(val) + far_mean = np.mean(far) + val_std = np.std(val) + return val_mean, val_std, far_mean + +def calculate_val_far(threshold, dist, actual_issame): + predict_issame = np.less(dist, threshold) + true_accept = np.sum(np.logical_and(predict_issame, actual_issame)) + false_accept = np.sum(np.logical_and(predict_issame, np.logical_not(actual_issame))) + n_same = np.sum(actual_issame) + n_diff = np.sum(np.logical_not(actual_issame)) + if n_diff == 0: + n_diff = 1 + if n_same == 0: + return 0,0 + val = float(true_accept) / float(n_same) + far = float(false_accept) / float(n_diff) + return val, far + +def test(test_loader, model, png_save_path, log_interval, batch_size, cuda): + labels, distances = [], [] + pbar = tqdm(enumerate(test_loader)) + for batch_idx, (data_a, data_p, label) in pbar: + with torch.no_grad(): + #--------------------------------------# + # 加载数据，设置成cuda + #--------------------------------------# + data_a, data_p = data_a.type(torch.FloatTensor), data_p.type(torch.FloatTensor) + if cuda: + data_a, data_p = data_a.cuda(), data_p.cuda() + #--------------------------------------# + # 传入模型预测，获得预测结果 + # 获得预测结果的距离 + #--------------------------------------# + out_a, out_p = model(data_a), model(data_p) + dists = torch.sqrt(torch.sum((out_a - out_p) ** 2, 1)) + + #--------------------------------------# + # 将结果添加进列表中 + #--------------------------------------# + distances.append(dists.data.cpu().numpy()) + labels.append(label.data.cpu().numpy()) + + #--------------------------------------# + # 打印 + #--------------------------------------# + if batch_idx % log_interval == 0: + pbar.set_description('Test Epoch: [{}/{} ({:.0f}%)]'.format( + batch_idx * batch_size, len(test_loader.dataset), + 100. * batch_idx / len(test_loader))) + + #--------------------------------------# + # 转换成numpy + #--------------------------------------# + labels = np.array([sublabel for label in labels for sublabel in label]) + distances = np.array([subdist for dist in distances for subdist in dist]) + + tpr, fpr, accuracy, val, val_std, far, best_thresholds = evaluate(distances,labels) + print('Accuracy: %2.5f+-%2.5f' % (np.mean(accuracy), np.std(accuracy))) + print('Best_thresholds: %2.5f' % best_thresholds) + print('Validation rate: %2.5f+-%2.5f @ FAR=%2.5f' % (val, val_std, far)) + plot_roc(fpr, tpr, figure_name = png_save_path) + +def plot_roc(fpr, tpr, figure_name = "roc.png"): + import matplotlib.pyplot as plt + from sklearn.metrics import auc, roc_curve + roc_auc = auc(fpr, tpr) + fig = plt.figure() + lw = 2 + plt.plot(fpr, tpr, color='darkorange', + lw=lw, label='ROC curve (area = %0.2f)' % roc_auc) + plt.plot([0, 1], [0, 1], color='navy', lw=lw, linestyle='--') + plt.xlim([0.0, 1.0]) + plt.ylim([0.0, 1.05]) + plt.xlabel('False Positive Rate') + plt.ylabel('True Positive Rate') + plt.title('Receiver operating characteristic') + plt.legend(loc="lower right") + fig.savefig(figure_name, dpi=fig.dpi)