├── deepvac ├── experimental │ ├── __init__.py │ ├── feature │ │ └── __init__.py │ └── core │ │ ├── __init__.py │ │ └── distill.py ├── version.py ├── loss │ ├── __init__.py │ └── loss.py ├── utils │ ├── ddp.py │ ├── time.py │ ├── __init__.py │ ├── meter.py │ ├── annotation.py │ ├── user_config.py │ ├── pallete.py │ ├── crypto.py │ ├── log.py │ ├── face_utils.py │ └── face_align.py ├── __init__.py ├── aug │ ├── __init__.py │ ├── emboss_helper.py │ ├── text_aug.py │ ├── remaper_helper.py │ ├── composer.py │ ├── factory.py │ ├── perspective_helper.py │ ├── seg_audit.py │ ├── warp_mls_helper.py │ ├── haishoku_helper.py │ ├── yolo_aug.py │ ├── seg_aug.py │ └── line_helper.py ├── datasets │ ├── __init__.py │ ├── os_walk.py │ ├── base_dataset.py │ ├── file_line.py │ └── coco.py ├── core │ ├── qat.py │ ├── __init__.py │ └── config.py ├── backbones │ ├── __init__.py │ ├── op_layer.py │ ├── activate_layer.py │ ├── se_layer.py │ ├── spp_layer.py │ ├── weights_init.py │ ├── ssh_layer.py │ ├── fpn_layer.py │ ├── module_helper.py │ ├── bottleneck_layer.py │ ├── res_layer.py │ ├── mobilenet.py │ ├── conv_layer.py │ └── regnet.py └── cast │ ├── __init__.py │ ├── onnx.py │ ├── mnn.py │ ├── trace.py │ ├── script.py │ ├── tensorrt.py │ ├── tnn.py │ ├── ncnn.py │ ├── cast_helper.py │ ├── coreml.py │ └── base.py ├── .github └── ISSUE_TEMPLATE │ └── bug--.md ├── docs ├── lib.md ├── deepvac_checklist.md ├── arch.md ├── code_standard.md └── deepvac_standard.md ├── setup.py ├── .gitignore └── README.md /deepvac/experimental/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /deepvac/experimental/feature/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /deepvac/version.py: -------------------------------------------------------------------------------- 1 | __version__ = '0.6.0' 2 | -------------------------------------------------------------------------------- /deepvac/experimental/core/__init__.py: -------------------------------------------------------------------------------- 1 | from .distill import * -------------------------------------------------------------------------------- /deepvac/loss/__init__.py: -------------------------------------------------------------------------------- 1 | from .loss import * 2 | from .face_loss import * -------------------------------------------------------------------------------- /deepvac/utils/ddp.py: -------------------------------------------------------------------------------- 1 | import sys 2 | is_ddp = False 3 | if '--rank' in sys.argv: 4 | is_ddp = True -------------------------------------------------------------------------------- /deepvac/__init__.py: -------------------------------------------------------------------------------- 1 | from .version import __version__ 2 | from .core import * 3 | from .utils import LOG, is_ddp 4 | 5 | 6 | 7 | -------------------------------------------------------------------------------- /deepvac/aug/__init__.py: -------------------------------------------------------------------------------- 1 | from .factory import * 2 | from .base_aug import * 3 | from .composer import * 4 | from .seg_audit import * 5 | -------------------------------------------------------------------------------- /deepvac/datasets/__init__.py: -------------------------------------------------------------------------------- 1 | from .base_dataset import * 2 | from .file_line import * 3 | from .os_walk import * 4 | from .coco import * 5 | -------------------------------------------------------------------------------- /deepvac/utils/time.py: -------------------------------------------------------------------------------- 1 | from datetime import datetime 2 | 3 | def getPrintTime(): 4 | return (str(datetime.now())[:-10]).replace(' ','-').replace(':','-') -------------------------------------------------------------------------------- /deepvac/core/qat.py: -------------------------------------------------------------------------------- 1 | from ..utils.ddp import is_ddp 2 | if is_ddp: 3 | from .deepvac import DeepvacDDP as DeepvacTrain 4 | else: 5 | from .deepvac import DeepvacTrain 6 | 7 | class DeepvacQAT(DeepvacTrain): 8 | pass -------------------------------------------------------------------------------- /deepvac/utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .log import LOG, assertAndGetGitBranch 2 | from .annotation import * 3 | from .time import * 4 | from .meter import * 5 | from .ddp import is_ddp 6 | from .user_config import * 7 | from .pallete import * -------------------------------------------------------------------------------- /deepvac/aug/emboss_helper.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | emboss_kernal = np.array([ 4 | [-2, -1, 0], 5 | [-1, 1, 1], 6 | [0, 1, 2] 7 | ]) 8 | 9 | def apply_emboss(word_img): 10 | return cv2.filter2D(word_img, -1, emboss_kernal) -------------------------------------------------------------------------------- /deepvac/core/__init__.py: -------------------------------------------------------------------------------- 1 | from ..utils import is_ddp 2 | 3 | if is_ddp: 4 | from .deepvac import DeepvacDDP as DeepvacTrain, Deepvac 5 | else: 6 | from .deepvac import DeepvacTrain, Deepvac 7 | 8 | from .config import AttrDict, interpret, newDict, new, fork 9 | from .report import * -------------------------------------------------------------------------------- /deepvac/backbones/__init__.py: -------------------------------------------------------------------------------- 1 | from .activate_layer import * 2 | from .bottleneck_layer import * 3 | from .conv_layer import * 4 | from .fpn_layer import * 5 | from .op_layer import * 6 | from .res_layer import * 7 | from .se_layer import * 8 | from .spp_layer import * 9 | from .ssh_layer import * 10 | from .weights_init import * 11 | -------------------------------------------------------------------------------- /deepvac/utils/meter.py: -------------------------------------------------------------------------------- 1 | class AverageMeter: 2 | def __init__(self): 3 | self.reset() 4 | 5 | def reset(self): 6 | self.val = 0 7 | self.avg = 0 8 | self.sum = 0 9 | self.count = 0 10 | 11 | def update(self, val, n=1): 12 | self.val = val 13 | self.sum += val * n 14 | self.count += n 15 | self.avg = self.sum / self.count -------------------------------------------------------------------------------- /deepvac/backbones/op_layer.py: -------------------------------------------------------------------------------- 1 | from typing import List 2 | import torch 3 | import torch.nn as nn 4 | 5 | class Concat(nn.Module): 6 | # Concatenate a list of tensors along dimension 7 | def __init__(self, dimension=1): 8 | super(Concat, self).__init__() 9 | self.d = dimension 10 | 11 | def forward(self, x: List[torch.Tensor]) -> torch.Tensor: 12 | return torch.cat(x, self.d) 13 | -------------------------------------------------------------------------------- /deepvac/utils/annotation.py: -------------------------------------------------------------------------------- 1 | #function only be executed once. 2 | def syszux_once(f): 3 | def wrapper(*args, **kwargs): 4 | if not wrapper.has_run: 5 | wrapper.has_run = True 6 | return f(*args, **kwargs) 7 | wrapper.has_run = False 8 | return wrapper 9 | 10 | if __name__ == "__main__": 11 | @syszux_once 12 | def test(x,y): 13 | print(x**y) 14 | 15 | test(2,3) 16 | test(3,4) 17 | -------------------------------------------------------------------------------- /deepvac/backbones/activate_layer.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | class hswish(nn.Module): 4 | def __init__(self, inplace=True): 5 | super(hswish, self).__init__() 6 | self.relu = nn.ReLU6(inplace=inplace) 7 | 8 | def forward(self, x): 9 | out = x * self.relu(x+3) / 6 10 | return out 11 | 12 | class hsigmoid(nn.Module): 13 | def __init__(self, inplace=True): 14 | super(hsigmoid, self).__init__() 15 | self.relu = nn.ReLU6(inplace=inplace) 16 | 17 | def forward(self, x): 18 | out = self.relu(x+3) / 6 19 | return out -------------------------------------------------------------------------------- /deepvac/cast/__init__.py: -------------------------------------------------------------------------------- 1 | from .trace import TraceCast 2 | from .script import ScriptCast 3 | from .onnx import OnnxCast 4 | from .coreml import CoremlCast 5 | from .ncnn import NcnnCast 6 | from .tensorrt import TensorrtCast 7 | from .mnn import MnnCast 8 | from .tnn import TnnCast 9 | 10 | caster_list = ['Trace', 'Script', 'Onnx', 'Coreml', 'Ncnn', 'Tensorrt', 'Mnn', 'Tnn'] 11 | def export3rd(trainer_config, cast_config, output_file=None): 12 | for caster_name in caster_list: 13 | caster = eval('{}Cast(trainer_config, cast_config)'.format(caster_name)) 14 | caster(output_file) 15 | 16 | -------------------------------------------------------------------------------- /deepvac/cast/onnx.py: -------------------------------------------------------------------------------- 1 | from ..utils import LOG 2 | from .base import DeepvacCast 3 | 4 | class OnnxCast(DeepvacCast): 5 | def auditConfig(self): 6 | if not self.config.onnx_model_dir: 7 | return False 8 | return True 9 | 10 | def process(self, cast_output_file=None): 11 | output_onnx_file = self.config.onnx_model_dir 12 | if cast_output_file: 13 | output_onnx_file = '{}/onnx__{}.onnx'.format(self.trainer_config.output_dir, cast_output_file) 14 | self.config.onnx_model_dir = output_onnx_file 15 | LOG.logI("config.cast.OnnxCast.onnx_model_dir found, save onnx model to {}...".format(output_onnx_file)) 16 | self.exportOnnx() -------------------------------------------------------------------------------- /deepvac/backbones/se_layer.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | from .module_helper import makeDivisible 3 | from .activate_layer import hsigmoid 4 | 5 | class SELayer(nn.Module): 6 | def __init__(self, channel, reduction=4): 7 | super(SELayer, self).__init__() 8 | self.avg_pool = nn.AdaptiveAvgPool2d(1) 9 | self.fc = nn.Sequential( 10 | nn.Linear(channel, makeDivisible(channel // reduction, 8)), 11 | nn.ReLU(inplace=True), 12 | nn.Linear(makeDivisible(channel // reduction, 8), channel), 13 | hsigmoid(inplace=True) 14 | ) 15 | 16 | def forward(self, x): 17 | b, c, _, _ = x.size() 18 | y = self.avg_pool(x).view(b, c) 19 | y = self.fc(y).view(b, c, 1, 1) 20 | return x * y -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/bug--.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: BUG报告 3 | about: 按照模板来报告有助于我们的理解 4 | title: 标题请以占位符[BUG]开头 5 | labels: '' 6 | assignees: '' 7 | 8 | --- 9 | 10 | **bug描述** 11 | 请描述该bug。 12 | **如何复现** 13 | 复现步骤: 14 | 1. 步骤1xxxx 15 | 2. 步骤2xxxx 16 | 3. 步骤3xxxx 17 | 4. 看，错误就是这个...... 18 | 19 | **预期结果** 20 | 请描述你原本期望的结果是什么。 21 | 22 | **截图** 23 | 如果有必要的话，请添加截图。 24 | 25 | **如果使用的是MLab HomePod，请填写** 26 | - 宿主机 cpu/ram/cuda设备： [比如. intel i9-9820X/32GB/RTX2080ti ] 27 | - 宿主机操作系统/内核版本/GPU驱动： [比如. ubuntu 20.04/5.4.0-74-generic/460.80 ] 28 | - MLab HomePod版本 [比如： 2.0-pro] 29 | 30 | **如果使用的不是MLab HomePod，请填写** 31 | - 宿主机 cpu/ram/cuda设备： [比如. intel i9-9820X/32GB/RTX2080ti ] 32 | - 宿主机操作系统/内核版本/GPU驱动： [比如. ubuntu 20.04/5.4.0-74-generic/460.80 33 | - 所有出错时刻调用栈中用到的三方包及其版本。 34 | 35 | **上下文备注** 36 | 任何你觉得有必要说明的上下文. 37 | -------------------------------------------------------------------------------- /deepvac/backbones/spp_layer.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from .conv_layer import Conv2dBNHardswish 4 | 5 | class SPP(nn.Module): 6 | # Spatial pyramid pooling layer used in YOLOv3-SPP 7 | def __init__(self, in_planes, out_planes, pool_kernel_size=(5, 9, 13)): 8 | super(SPP, self).__init__() 9 | hidden_planes = in_planes // 2 # hidden channels 10 | self.conv1 = Conv2dBNHardswish(in_planes, hidden_planes, 1, 1) 11 | self.conv2 = Conv2dBNHardswish(hidden_planes * (len(pool_kernel_size) + 1), out_planes, 1, 1) 12 | self.pool_list = nn.ModuleList([nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2) for k in pool_kernel_size]) 13 | 14 | def forward(self, x): 15 | x = self.conv1(x) 16 | return self.conv2(torch.cat([x] + [pool(x) for pool in self.pool_list], 1)) 17 | 18 | 19 | 20 | -------------------------------------------------------------------------------- /docs/lib.md: -------------------------------------------------------------------------------- 1 | # deepvac库 2 | deepvac库对云上炼丹师提供以下模块： 3 | 4 | | 模块 | 使用 | 说明 | 5 | |--------------------|------------------------------|---------| 6 | |Deepvac |from deepvac import Deepvac |测试代码的基类| 7 | |DeepvacTrain |from deepvac import DeepvacTrain |训练、验证代码的基类| 8 | |LOG | from deepvac import LOG |日志模块| 9 | |config | from deepvac import AttrDict, new, fork, interpret |配置模块 | 10 | |*Aug | from deepvac.aug import * | 各种数据增强的类实现| 11 | |{,*Aug}Composer | from deepvac.aug import * |动态数据增强的抽象封装| 12 | |*Dataset | from deepvac.datasets import * | Dataset的扩展实现，torch.utils.data.Dataset的子类们，用于个性化各种dataloader| 13 | |{,Face,Ocr,Classifier}Report | from deepvac import * | Report类体系，用于打印各种测试报告| 14 | |各种经典网络模块 | from deepvac.backbones import * | 神经网络中经典block的实现 | 15 | |各种loss函数 | from deepvac.loss import * | 神经网络中各种损失评价函数的实现 | 16 | 17 | -------------------------------------------------------------------------------- /deepvac/utils/user_config.py: -------------------------------------------------------------------------------- 1 | from .log import LOG 2 | 3 | def addUserConfig(module_name, config_name, user_give=None, developer_give=None, is_user_mandatory=False): 4 | if user_give is None and is_user_mandatory: 5 | LOG.logE("You must set mandatory {}.{} in config.py. Developer advised value:{}".format(module_name, config_name, developer_give),exit=True) 6 | if user_give is not None: 7 | return user_give 8 | if developer_give is not None: 9 | return developer_give 10 | LOG.logE("value missing for configuration: {}.{} in config.py".format(module_name, config_name), exit=True) 11 | 12 | def anyFieldsInConfig(name, c, c_strict=True): 13 | if not c: 14 | return False 15 | 16 | if not isinstance(c, list): 17 | return False 18 | 19 | fields = name.split('.') if c_strict is False else [name] 20 | for field in fields: 21 | if field in c: 22 | return True 23 | return False 24 | -------------------------------------------------------------------------------- /deepvac/backbones/weights_init.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | def initWeightsKaiming(civilnet): 4 | for m in civilnet.modules(): 5 | if isinstance(m, nn.Conv2d): 6 | nn.init.kaiming_normal_(m.weight, mode='fan_out') 7 | if m.bias is not None: 8 | nn.init.zeros_(m.bias) 9 | elif isinstance(m, nn.BatchNorm2d): 10 | if m.affine: 11 | nn.init.ones_(m.weight) 12 | nn.init.zeros_(m.bias) 13 | elif isinstance(m, nn.Linear): 14 | nn.init.normal_(m.weight, 0, 0.01) 15 | if m.bias is not None: 16 | nn.init.zeros_(m.bias) 17 | 18 | def initWeightsNormal(civilnet): 19 | for m in civilnet.modules(): 20 | if isinstance(m, nn.Conv2d): 21 | nn.init.normal_(m.weight, 0, 0.02) 22 | if m.bias is not None: 23 | nn.init.zeros_(m.bias) 24 | elif isinstance(m, nn.BatchNorm2d): 25 | if m.affine: 26 | nn.init.normal_(m.weight, 1.0, 0.02) 27 | nn.init.zeros_(m.bias) -------------------------------------------------------------------------------- /deepvac/datasets/os_walk.py: -------------------------------------------------------------------------------- 1 | import os 2 | import cv2 3 | from .base_dataset import DatasetBase 4 | 5 | class OsWalkBaseDataset(DatasetBase): 6 | def __init__(self, deepvac_config, sample_path): 7 | super(OsWalkBaseDataset, self).__init__(deepvac_config) 8 | self.files = [] 9 | for subdir, dirs, fns in os.walk(sample_path): 10 | for fn in fns: 11 | self.files.append(os.path.join(subdir, fn)) 12 | 13 | def __len__(self): 14 | return len(self.files) 15 | 16 | def __getitem__(self, index): 17 | filepath = self.files[index] 18 | sample = cv2.imread(filepath) 19 | sample = self.compose(sample) 20 | return sample, filepath 21 | 22 | class OsWalkDataset(OsWalkBaseDataset): 23 | def __getitem__(self, index): 24 | sample, filepath = super(OsWalkDataset, self).__getitem__(index) 25 | return sample 26 | 27 | class OsWalkContinueDataset(OsWalkBaseDataset): 28 | def __getitem__(self, index): 29 | try: 30 | sample, filepath = super(OsWalkDataset, self).__getitem__(index) 31 | except: 32 | sample, filepath=None,None 33 | return sample 34 | 35 | -------------------------------------------------------------------------------- /deepvac/backbones/ssh_layer.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from .conv_layer import Conv2dBN, Conv2dBNLeakyReLU 4 | 5 | class SSH(nn.Module): 6 | def __init__(self, in_planes: int, out_planes: int): 7 | super(SSH, self).__init__() 8 | assert out_planes % 4 == 0 9 | leaky = 0.1 if out_planes <= 64 else 0 10 | self.conv3X3 = Conv2dBN(in_planes, out_planes//2, padding=1) 11 | 12 | self.conv5X5_1 = Conv2dBNLeakyReLU(in_planes, out_planes//4, padding=1, leaky=leaky) 13 | self.conv5X5_2 = Conv2dBN(out_planes//4, out_planes//4, padding=1) 14 | 15 | self.conv7X7_2 = Conv2dBNLeakyReLU(out_planes//4, out_planes//4, padding=1, leaky=leaky) 16 | self.conv7x7_3 = Conv2dBN(out_planes//4, out_planes//4, padding=1) 17 | 18 | self.relu = nn.ReLU(inplace=True) 19 | 20 | def forward(self, input): 21 | conv3X3 = self.conv3X3(input) 22 | 23 | conv5X5_1 = self.conv5X5_1(input) 24 | conv5X5 = self.conv5X5_2(conv5X5_1) 25 | 26 | conv7X7_2 = self.conv7X7_2(conv5X5_1) 27 | conv7X7 = self.conv7x7_3(conv7X7_2) 28 | 29 | out = torch.cat([conv3X3, conv5X5, conv7X7], dim=1) 30 | out = self.relu(out) 31 | return out -------------------------------------------------------------------------------- /deepvac/utils/pallete.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import cv2 3 | 4 | pallete20 = [[255, 255, 255], 5 | [128, 64, 128], 6 | [244, 35, 232], 7 | [70, 70, 70], 8 | [102, 102, 156], 9 | [190, 153, 153], 10 | [153, 153, 153], 11 | [250, 170, 30], 12 | [220, 220, 0], 13 | [107, 142, 35], 14 | [152, 251, 152], 15 | [70, 130, 180], 16 | [220, 20, 60], 17 | [255, 0, 0], 18 | [0, 0, 142], 19 | [0, 0, 70], 20 | [0, 60, 100], 21 | [0, 80, 100], 22 | [0, 0, 230], 23 | [119, 11, 32]] 24 | 25 | def getOverlayFromSegMask(img, mask): 26 | assert len(mask.shape) == 2, 'only support 2-D mask, got {}'.format(len(mask.shape)) 27 | assert img.shape[:2] == mask.shape[:2], 'img shape not equal mask shape: {} vs {}'.format(img.shape[:2], mask.shape[:2]) 28 | class_map_np_color = np.zeros((mask.shape[0], mask.shape[1], 3), dtype=np.uint8) 29 | for idx in np.unique(mask): 30 | [r, g, b] = pallete20[idx] 31 | class_map_np_color[mask == idx] = [b, g, r] 32 | overlayed = cv2.addWeighted(img, 0.5, class_map_np_color, 0.5, 0) 33 | return overlayed -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | import setuptools 2 | import sys 3 | 4 | def readme(): 5 | with open('README.md') as f: 6 | return f.read() 7 | 8 | def version(): 9 | with open('deepvac/version.py') as f: 10 | version_cmd = f.read() 11 | loc = {} 12 | exec(version_cmd.strip(),globals(), loc) 13 | return loc['__version__'] 14 | 15 | version=version() 16 | print("You are building deepvac with version: {}".format(version)) 17 | setuptools.setup(name='deepvac', 18 | version=version, 19 | description='PyTorch python project standard', 20 | long_description='PyTorch python project standard', 21 | keywords='PyTorch python project standard', 22 | url='https://github.com/DeepVAC/deepvac', 23 | download_url="", 24 | author='Gemfield', 25 | author_email='gemfield@civilnet.cn', 26 | packages=setuptools.find_packages(), 27 | include_package_data=True, 28 | package_dir={ 29 | 'deepvac': 'deepvac', 30 | }, 31 | 32 | install_requires=[ 33 | 'opencv_python', 34 | 'numpy', 35 | 'Pillow', 36 | 'scipy', 37 | 'tensorboard' 38 | ], 39 | classifiers = [ 40 | "Programming Language :: Python :: 3", 41 | "Topic :: Software Development :: Libraries" 42 | ], 43 | zip_safe=False) 44 | -------------------------------------------------------------------------------- /deepvac/backbones/fpn_layer.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import torch.nn.functional as F 3 | from .conv_layer import Conv2dBNLeakyReLU 4 | 5 | class FPN(nn.Module): 6 | def __init__(self, in_planes: list, out_planes: int): 7 | super(FPN,self).__init__() 8 | 9 | leaky = 0.1 if out_planes <= 64 else 0 10 | self.conv1 = Conv2dBNLeakyReLU(in_planes[0], out_planes, kernel_size=1, padding=0, leaky=leaky) 11 | self.conv2 = Conv2dBNLeakyReLU(in_planes[1], out_planes, kernel_size=1, padding=0, leaky=leaky) 12 | self.conv3 = Conv2dBNLeakyReLU(in_planes[2], out_planes, kernel_size=1, padding=0, leaky=leaky) 13 | 14 | self.conv4 = Conv2dBNLeakyReLU(out_planes, out_planes, padding=1, leaky=leaky) 15 | self.conv5 = Conv2dBNLeakyReLU(out_planes, out_planes, padding=1, leaky=leaky) 16 | 17 | def forward(self, input): 18 | output1 = self.conv1(input[0]) 19 | output2 = self.conv2(input[1]) 20 | output3 = self.conv3(input[2]) 21 | 22 | up3 = F.interpolate(output3, size=[output2.size(2), output2.size(3)], mode="nearest") 23 | output2 = output2 + up3 24 | output2 = self.conv5(output2) 25 | 26 | up2 = F.interpolate(output2, size=[output1.size(2), output1.size(3)], mode="nearest") 27 | output1 = output1 + up2 28 | output1 = self.conv4(output1) 29 | out = [output1, output2, output3] 30 | return out -------------------------------------------------------------------------------- /deepvac/backbones/module_helper.py: -------------------------------------------------------------------------------- 1 | from ..utils import LOG 2 | 3 | def makeDivisible(v, divisor=8, min_value=None): 4 | if min_value is None: 5 | min_value = divisor 6 | new_v = max(min_value, int(v + divisor / 2) // divisor * divisor) 7 | # Make sure that round down does not go down by more than 10%. 8 | if new_v < 0.9 * v: 9 | new_v += divisor 10 | return new_v 11 | 12 | def separateBN4OptimizerPG(modules): 13 | paras_only_bn = [] 14 | paras_wo_bn = [] 15 | memo = set() 16 | gemfield_set = set() 17 | gemfield_set.update(set(modules.parameters())) 18 | LOG.logI("separateBN4OptimizerPG set len: {}".format(len(gemfield_set))) 19 | named_modules = modules.named_modules(prefix='') 20 | for module_prefix, module in named_modules: 21 | if "module" not in module_prefix: 22 | LOG.logI("separateBN4OptimizerPG skip {}".format(module_prefix)) 23 | continue 24 | 25 | members = module._parameters.items() 26 | for k, v in members: 27 | name = module_prefix + ('.' if module_prefix else '') + k 28 | if v is None: 29 | continue 30 | if v in memo: 31 | continue 32 | memo.add(v) 33 | if "batchnorm" in str(module.__class__): 34 | paras_only_bn.append(v) 35 | else: 36 | paras_wo_bn.append(v) 37 | 38 | LOG.logI("separateBN4OptimizerPG param len: {} - {}".format(len(paras_wo_bn),len(paras_only_bn))) 39 | return paras_only_bn, paras_wo_bn 40 | 41 | 42 | 43 | 44 | -------------------------------------------------------------------------------- /deepvac/aug/text_aug.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import random 3 | from .base_aug import CvAugBase 4 | from .perspective_helper import apply_perspective_transform 5 | from .remaper_helper import Remaper 6 | from .line_helper import Liner 7 | from .emboss_helper import apply_emboss 8 | 9 | class TextRendererPerspectiveAug(CvAugBase): 10 | def auditConfig(self): 11 | self.config.max_x = self.addUserConfig('max_x', self.config.max_x, 10) 12 | self.config.max_y = self.addUserConfig('max_y', self.config.max_y, 10) 13 | self.config.max_z = self.addUserConfig('max_z', self.config.max_z, 5) 14 | 15 | def forward(self, img): 16 | return apply_perspective_transform(img, self.config.max_x, self.config.max_y, self.config.max_z) 17 | 18 | class TextRendererCurveAug(CvAugBase): 19 | def forward(self, img): 20 | h, w = img.shape[:2] 21 | re_img, text_box_pnts = Remaper().apply(img, [[0,0],[w,0],[w,h],[0,h]]) 22 | return re_img 23 | 24 | class TextRendererLineAug(CvAugBase): 25 | def auditConfig(self): 26 | self.config.offset = self.addUserConfig('offset', self.config.offset, 5) 27 | 28 | def forward(self, img): 29 | h, w = img.shape[:2] 30 | offset = self.config.offset 31 | pos = [[offset,offset],[w-offset,offset],[w-offset,h-offset],[offset,h-offset]] 32 | re_img, text_box_pnts = Liner().apply(img, pos) 33 | return re_img 34 | 35 | class TextRendererEmbossAug(CvAugBase): 36 | def forward(self, img): 37 | return apply_emboss(img) 38 | 39 | class TextRendererReverseAug(CvAugBase): 40 | def forward(self, img): 41 | offset = np.random.randint(-10, 10) 42 | return 255 + offset - img -------------------------------------------------------------------------------- /deepvac/datasets/base_dataset.py: -------------------------------------------------------------------------------- 1 | from torch.utils.data import Dataset 2 | from ..utils import LOG, addUserConfig 3 | from ..core import AttrDict 4 | 5 | class DatasetBase(Dataset): 6 | def __init__(self, deepvac_config): 7 | super(DatasetBase, self).__init__() 8 | self.deepvac_datasets_config = deepvac_config.datasets 9 | self.initConfig() 10 | self.auditConfig() 11 | 12 | def initConfig(self): 13 | if self.name() not in self.deepvac_datasets_config.keys(): 14 | self.deepvac_datasets_config[self.name()] = AttrDict() 15 | self.config = self.deepvac_datasets_config[self.name()] 16 | 17 | def auditConfig(self): 18 | pass 19 | 20 | def addUserConfig(self, config_name, user_give=None, developer_give=None, is_user_mandatory=False): 21 | module_name = 'config.datasets.{}'.format(self.name()) 22 | return addUserConfig(module_name, config_name, user_give=user_give, developer_give=developer_give, is_user_mandatory=is_user_mandatory) 23 | 24 | def compose(self, img): 25 | if isinstance(self.config.transform, (list,tuple)): 26 | for t in self.config.transform: 27 | img = t(img) 28 | elif self.config.transform is not None: 29 | img = self.config.transform(img) 30 | 31 | if isinstance(self.config.composer, (list,tuple)): 32 | for c in self.config.composer: 33 | img = c(img) 34 | elif self.config.composer is not None: 35 | img = self.config.composer(img) 36 | 37 | return img 38 | 39 | def name(self): 40 | return self.__class__.__name__ 41 | 42 | def setAttr(self, k, v): 43 | self.config[k] = v 44 | 45 | def getAttr(self,k): 46 | return self.config[k] 47 | -------------------------------------------------------------------------------- /deepvac/aug/remaper_helper.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | 3 | class Remaper(object): 4 | def __init__(self): 5 | self.period = 360 6 | self.min = 1 7 | self.max = 5 8 | 9 | def apply(self, word_img, text_box_pnts): 10 | """ 11 | :param word_img: word image with big background 12 | :param text_box_pnts: left-top, right-top, right-bottom, left-bottom of text word 13 | :return: 14 | """ 15 | max_val = np.random.uniform(self.min, self.max) 16 | 17 | h = word_img.shape[0] 18 | w = word_img.shape[1] 19 | 20 | img_x = np.zeros((h, w), np.float32) 21 | img_y = np.zeros((h, w), np.float32) 22 | 23 | xmin = text_box_pnts[0][0] 24 | xmax = text_box_pnts[1][0] 25 | ymin = text_box_pnts[0][1] 26 | ymax = text_box_pnts[2][1] 27 | 28 | remap_y_min = ymin 29 | remap_y_max = ymax 30 | 31 | for y in range(h): 32 | for x in range(w): 33 | remaped_y = y + self._remap_y(x, max_val) 34 | 35 | if y == ymin: 36 | if remaped_y < remap_y_min: 37 | remap_y_min = remaped_y 38 | 39 | if y == ymax: 40 | if remaped_y > remap_y_max: 41 | remap_y_max = remaped_y 42 | 43 | # 某一个位置的 y 值应该为哪个位置的 y 值 44 | img_y[y, x] = remaped_y 45 | # 某一个位置的 x 值应该为哪个位置的 x 值 46 | img_x[y, x] = x 47 | 48 | remaped_text_box_pnts = [ 49 | [xmin, remap_y_min], 50 | [xmax, remap_y_min], 51 | [xmax, remap_y_max], 52 | [xmin, remap_y_max] 53 | ] 54 | 55 | # TODO: use cuda::remap 56 | dst = cv2.remap(word_img, img_x, img_y, cv2.INTER_CUBIC) 57 | return dst, remaped_text_box_pnts 58 | 59 | def _remap_y(self, x, max_val): 60 | return int(max_val * np.math.sin(2 * 3.14 * x / self.period)) 61 | -------------------------------------------------------------------------------- /deepvac/cast/mnn.py: -------------------------------------------------------------------------------- 1 | import os 2 | from ..utils import LOG 3 | from .base import DeepvacCast 4 | 5 | class MnnCast(DeepvacCast): 6 | def auditConfig(self): 7 | if not self.config.model_dir: 8 | return False 9 | 10 | self.config.onnx2mnn = self.addUserConfig("onnx2mnn", self.config.onnx2mnn, "/bin/MNNConvert") 11 | exist = os.path.exists(self.config.onnx2mnn) 12 | if not exist: 13 | LOG.logE("The config.cast.MnnCast.onnx2mnn is invalid. We try to use default setting of config.cast.MnnCast.onnx2mnn is failed too. \ 14 | If you want to compile onnx2mnn tools, reference https://www.yuque.com/mnn/cn/cvrt_linux_mac ", exit=True) 15 | return False 16 | 17 | return True 18 | 19 | def process(self, cast_output_file=None): 20 | output_mnn_file = self.config.model_dir 21 | if cast_output_file: 22 | output_mnn_file = '{}/mnn__{}.mnn'.format(self.trainer_config.output_dir, cast_output_file) 23 | self.config.model_dir = output_mnn_file 24 | 25 | LOG.logI("config.cast.MnnCast.model_dir found, save mnn model to {}...".format(self.config.model_dir)) 26 | 27 | #to onnx, also set self.config.onnx_model_dir, self.config.onnx_input_names and self.config.onnx_output_names 28 | self.exportOnnx() 29 | 30 | cmd = "{} -f ONNX --modelFile {} --MNNModel {}".format(self.config.onnx2mnn, self.config.onnx_model_dir, output_mnn_file) 31 | if self.config.save_static_model: 32 | cmd += " --saveStaticModel " 33 | rc, out_text, err_text = self.runCmd(cmd) 34 | 35 | if "" != err_text: 36 | LOG.logE(err_text + ". Error occured when export mnn model. ", exit=True) 37 | if "Converted Success" in out_text: 38 | LOG.logI("Pytorch model convert to Mnn model succeed, save mnn model file in {}.".format(output_mnn_file)) 39 | return 40 | LOG.logE(out_text + ". Error occured when export mnn model. ", exit=True) 41 | -------------------------------------------------------------------------------- /docs/deepvac_checklist.md: -------------------------------------------------------------------------------- 1 | ## DeepVAC checklist 2 | **AI模型：** 3 | **研究人员：** 4 | **研发周期：** 5 | 6 | | 序号 | 检查项 | AI1 |AI2 |AI3 | 7 | |----|-------------|---------|----|----| 8 | | 1 | 训练集是否维护在MLab RookPod存储上 | ||| 9 | | 2 | 测试集是否维护在MLab RookPod存储上 | ||| 10 | | 3 | 验收集是否维护在MLab RookPod存储上 | ||| 11 | | 4 | 模型定义是否复用了deepvac.backbones | ||| 12 | | 5 | 模型转换是否复用了deepvac.cast | ||| 13 | | 6 | loss定义是否复用了deepvac.loss | ||| 14 | | 7 | 训练代码是否基于deepvac库 | ||| 15 | | 8 | 测试代码是否基于deepvac库 | ||| 16 | | 9 | 日志代码是否基于deepvac库 | ||| 17 | |10 | 配置代码是否基于deepvac库 | ||| 18 | |11 | 数据合成代码是否基于deepvac库 | ||| 19 | |12 | 数据增强和动态增强代码是否基于deepvac.aug | ||| 20 | |13 | dataloader代码是否基于deepvac.datasets | ||| 21 | |14 | 模型性能报告代码是否基于deepvac库 | ||| 22 | |15 | 项目代码是否符合DeepVAC代码规范 | ||| 23 | |16 | 项目目录和文件名、git分支、配置项是否符合DeepVAC软件工程规范 | ||| 24 | |17 | 项目代码、dockerfile、README.md是否维护在了MLab代码服务系统上 | ||| 25 | |18 | 项目所依赖的预训练模型是否维护在了MLab RookPod存储上 | ||| 26 | |19 | 项目是否完成了部署目标为x86+CUDA Linux的训练 | ||| 27 | |20 | 上述训练的验收集指标是否符合要求 | ||| 28 | |21 | 上述训练输出的SOTA模型是否维护在了MLab RookPod存储上 | ||| 29 | |22 | 上述训练的性能报告是否记录到deepvac-product | ||| 30 | |23 | 项目是否完成了部署目标为x86 Linux、Arm Linux、ARM Android/iOS的训练 | ||| 31 | |24 | 上述训练的验收集指标是否符合要求 | ||| 32 | |25 | 上述训练输出的SOTA模型是否维护在了MLab RookPod存储上 | ||| 33 | |26 | 上述训练的性能报告是否记录到deepvac-product | ||| 34 | |27 | 项目是否成功进行了x86 + CUDA Linux的部署测试 | ||| 35 | |28 | 上述测试成功的镜像是否维护在了MLab docker registry | ||| 36 | |29 | 项目是否成功进行了x86 Linux的部署测试 | ||| 37 | |30 | 上述测试成功的镜像是否维护在了MLab docker registry | ||| 38 | |31 | 项目是否成功进行了Arm Android的部署测试 | ||| 39 | |32 | 上述测试成功的libdeepvac库是否维护在了MLab RookPod存储上 | ||| 40 | |33 | 项目或者产品release中诞生的contribute是否记录在了contribute清单中 | ||| 41 | |34 | 项目或者产品release中诞生的里程碑是否记录在了里程碑清单中 | ||| 42 | -------------------------------------------------------------------------------- /deepvac/cast/trace.py: -------------------------------------------------------------------------------- 1 | from ..utils import LOG 2 | from .base import DeepvacCast, FXQuantizeAndTrace 3 | 4 | class TraceCast(DeepvacCast): 5 | def auditConfig(self): 6 | if not self.config.model_dir: 7 | if self.config.dynamic_quantize_dir: 8 | LOG.logE("You must set config.cast.TraceCast.model_dir when you trying to enable config.cast.TraceCast.dynamic_quantize_dir.", exit=True) 9 | if self.config.static_quantize_dir: 10 | LOG.logE("You must set config.cast.TraceCast.model_dir when you trying to enable config.cast.TraceCast.static_quantize_dir.", exit=True) 11 | return False 12 | 13 | if self.config.static_quantize_dir: 14 | if self.trainer_config.is_forward_only and self.trainer_config.test_loader is None: 15 | LOG.logE("You must set config.core.test_loader in config.py when static_quantize_dir is enabled in TEST.", exit=True) 16 | if not self.trainer_config.is_forward_only and self.trainer_config.val_loader is None: 17 | LOG.logE("You must set config.core.val_loader in config.py when static_quantize_dir is enabled in TRAIN.", exit=True) 18 | 19 | return True 20 | 21 | def process(self, cast_output_file=None): 22 | output_trace_file = self.config.model_dir 23 | if cast_output_file: 24 | output_trace_file = '{}/trace_{}.pt'.format(self.trainer_config.output_dir, cast_output_file) 25 | 26 | LOG.logI("config.trace_model_dir found, save trace model to {}...".format(output_trace_file)) 27 | trace_model = FXQuantizeAndTrace(self.trainer_config, output_trace_file) 28 | trace_model.save() 29 | 30 | if self.config.dynamic_quantize_dir: 31 | LOG.logI("You have enabled config.cast.TraceCast.dynamic_quantize_dir, will dynamic quantize the model...") 32 | trace_model.saveDQ() 33 | 34 | if self.config.static_quantize_dir: 35 | LOG.logI("You have enabled config.cast.TraceCast.static_quantize_dir, will static quantize the model...") 36 | trace_model.saveSQ() -------------------------------------------------------------------------------- /deepvac/cast/script.py: -------------------------------------------------------------------------------- 1 | from ..utils import LOG 2 | from .base import DeepvacCast, FXQuantizeAndScript 3 | 4 | class ScriptCast(DeepvacCast): 5 | def auditConfig(self): 6 | if not self.config.model_dir: 7 | if self.config.dynamic_quantize_dir: 8 | LOG.logE("You must set config.cast.ScriptCast.model_dir when you trying to enable config.cast.ScriptCast.dynamic_quantize_dir.", exit=True) 9 | if self.config.static_quantize_dir: 10 | LOG.logE("You must set config.cast.ScriptCast.model_dir when you trying to enable config.cast.ScriptCast.static_quantize_dir.", exit=True) 11 | return False 12 | 13 | if self.config.static_quantize_dir: 14 | if self.trainer_config.is_forward_only and self.trainer_config.test_loader is None: 15 | LOG.logE("You must set config.core.test_loader in config.py when static_quantize_dir is enabled in TEST.", exit=True) 16 | if not self.trainer_config.is_forward_only and self.trainer_config.val_loader is None: 17 | LOG.logE("You must set config.core.val_loader in config.py when static_quantize_dir is enabled in TRAIN.", exit=True) 18 | return True 19 | 20 | def process(self, cast_output_file=None): 21 | output_script_file = self.config.model_dir 22 | if cast_output_file: 23 | output_script_file = '{}/script__{}.pt'.format(self.trainer_config.output_dir, cast_output_file) 24 | 25 | LOG.logI("config.script_model_dir found, save script model to {}...".format(output_script_file)) 26 | 27 | script_model = FXQuantizeAndScript(self.trainer_config, output_script_file) 28 | script_model.save() 29 | 30 | if self.config.dynamic_quantize_dir: 31 | LOG.logI("You have enabled config.cast.ScriptCast.dynamic_quantize_dir, will dynamic quantize the model...") 32 | script_model.saveDQ() 33 | 34 | if self.config.static_quantize_dir: 35 | LOG.logI("You have enabled config.cast.ScriptCast.static_quantize_dir, will static quantize the model...") 36 | script_model.saveSQ() 37 | -------------------------------------------------------------------------------- /docs/arch.md: -------------------------------------------------------------------------------- 1 | # 软件工程规范 2 | 定义PyTorch项目的软件工程规范，包含： 3 | - 训练测试的代码目录/文件规范； 4 | - git分支规范； 5 | - 配置规范； 6 | 7 | ## 训练测试的代码目录/文件规范 8 | 每一个符合deepvac规范的PyTorch模型训练项目，都包含下面这些目录和文件： 9 | 10 | | 目录/文件 | 说明 | 是否必须 | 11 | |--------------|---------|---------| 12 | |README.md |项目的说明、git分支数量及其介绍、原始数据的存放路径说明 | 是 | 13 | |train.py |模型训练和验证文件,继承DeepvacTrain类的扩展实现| 是 | 14 | |test.py |模型测试文件, 继承Deepvac类的扩展实现| 是 | 15 | |config.py |项目整体配置文件| 是 | 16 | |modules/model.py | 模型定义文件，PyTorch Module类的扩展实现|单个模型的情况下，是 | 17 | |modules/model_{name}.py | 同上，有多个model的时候，使用suffix区分|多个模型的情况下，是 | 18 | |modules/loss.py | loss实现。如果该实现比较轻量的话，可以直接放在modules/model.py中|否 | 19 | |modules/utils.py | 工具类/方法的定义文件|否 | 20 | |modules/utils_{name}.py | 同上，有多个工具类/函数文件的时候，使用suffix区分|否 | 21 | |synthesis/synthesis.py| 数据合成或清洗代码|否 | 22 | |synthesis/config.py|synthesis.py的配置文件|否 | 23 | |data/dataloader.py | dataset类的自定义实现|否 | 24 | |data/train.txt | 训练集清单文件|否 | 25 | |data/val.txt | 验证集清单文件|否 | 26 | |aug/aug.py|数据增强的代码。如果该实现比较轻量的话，可以直接放在dataset类的文件中|否 | 27 | |aug/config.py|aug.py的配置文件|否 | 28 | |log/*.log |日志输出目录 |是 | 29 | |output/model__*.pth | 输出或输入的模型文件 |默认Deepvac输出 | 30 | |output/checkpoint__*.pth | 输出或输入的checkpoint文件 |默认Deepvac输出 | 31 | 32 | 这些文件覆盖了一个PyTorch模型训练的整个生命周期： 33 | - 原始数据，在README.md中描述； 34 | - 数据清洗/合成，在synthesis/synthesis.py中定义； 35 | - 数据增强，在aug/aug.py中定义（轻量实现的话在dataset类中定义）； 36 | - 数据输入，在data/dataloader.py中定义； 37 | - 模型训练，在train.py中定义； 38 | - 模型验证，在train.py中定义； 39 | - 模型测试，在test.py中定义； 40 | - 模型输出，在output目录中存放； 41 | - 日志输出，在log目录中存放； 42 | 43 | ## git分支规范 44 | 每一个符合DeepVAC规范的PyTorch项目，都会面临一个问题：如何并行的进行多种实验？ 45 | DeepVAC采用的是git branch的解决方案。DeepVAC规定： 46 | - 每一种实验都有对应的git branch； 47 | - DeepVAC规定了两种git branch，长期运行分支和临时分支； 48 | - 长期运行分支的名字为master、main或者以LTS_开头，临时分支的名字以PROTO_开头； 49 | - 分支的名字要准确表达当前训练的上下文，使用变量名的代码规范，小写字母，以下划线分割；如：LTS_ocr_train_on_5synthesis_9aug_100w； 50 | - 因此，DeepVAC在检测不到当前代码所在的合法git branch时，将会终止运行（除非在config.py中显式配置）； 51 | 52 | 53 | ## 配置规范 54 | - config.py是DeepVAC规范中的一等公民； 55 | - 用户接口层面的配置均在config.py中； 56 | - 内部开发定义的配置均在类的auditConfig方法中，且可以被config.py中的值覆盖； 57 | - 开启分布式训练时，由于--rank和--gpu参数为进程级别，由argparse.ArgumentParser模块来传递，用户需要在命令行指定； 58 | - 类的构造函数的入参一般为config实例； 59 | - 再重复一遍，config.py在DeepVAC规范中为一等公民。 60 | -------------------------------------------------------------------------------- /docs/code_standard.md: -------------------------------------------------------------------------------- 1 | # Code Standard 2 | DeepVAC项目中的代码规范。使用的是C++举例，Java、Python等其它语言也同样遵守。 3 | 4 | ## 代码理念 5 | 代码理念和Python之禅一致，如下所示： 6 | ```python 7 | >>> import this 8 | The Zen of DeepVAC, by Gemfield 9 | 10 | 1.漂亮胜于一切； 11 | 2.可读性很重要； 12 | 3.变量名即注释； 13 | 4.明了胜于晦涩； 14 | 5.简洁胜于复杂； 15 | 6.复杂胜于凌乱； 16 | 7.扁平胜于嵌套； 17 | 8.超过三层嵌套应坐立不安； 18 | 9.尽量复用代码； 19 | 10.应为没有模块化而寝食难安； 20 | 11.多行短代码胜于一行长代码； 21 | 12.不容易实现的想法就不是好想法； 22 | 13.不容易解释的实现就不是好实现； 23 | 14.容易解释的代码才可能是好代码； 24 | 15.善用容器和巧妙算法来重构冗长逻辑； 25 | 16.命名空间是个绝妙的设计； 26 | 17.条件分支尽量重构为表驱动； 27 | 18.写if前要三思而后行； 28 | 19.善用列表推导式； 29 | 20.不要以特殊理由逃避上述规则。 30 | ``` 31 | 当在实践中遇到冲突的理念时，必须： 32 | - 首先确保代码漂亮； 33 | - 其次确保代码简洁； 34 | - 其次确保代码扁平； 35 | - 其次确保代码明了； 36 | - 其次确保代码可读性； 37 | - 其次确保短代码； 38 | - 其次确保代码可解释； 39 | - 其次确保代码巧妙； 40 | - 以上无一确保时，确保离职手续已办妥。 41 | 42 | ## 代码规范 43 | 44 | - file name (文件名) 45 | ```bash 46 | syszux_msg.h 47 | syszux.h 48 | syszux_msg.cpp 49 | ``` 50 | - variable name (变量名) 51 | ```c++ 52 | int gemfield; 53 | int gemfield_flag; 54 | //global 55 | int g_gemfield_flag; 56 | //const 57 | const int kGemfieldFlag; 58 | ``` 59 | - class name (类名) 60 | ```c++ 61 | class Syszux; 62 | class SyszuxMsg; 63 | ``` 64 | - class data member name (类的数据成员) 65 | ```c++ 66 | int gemfield_; 67 | int syszux_msg_; 68 | ``` 69 | 70 | - function name (函数名) 71 | ```c++ 72 | void get(); 73 | void getGemfieldConfig(); 74 | ``` 75 | - namespace name (命名空间) 76 | ```c++ 77 | namespace gemfield; 78 | namespace gemfield_log; 79 | ``` 80 | 81 | - ENUM & MACRO (枚举&宏) 82 | ```c++ 83 | //enum 84 | enum GEMFIELD; 85 | //macro 86 | define GEMFIELD gemfield; 87 | define GEM_FIELD 7030; 88 | ``` 89 | 90 | - flow control(流控制) 91 | ```c++ 92 | //if 93 | if(true){ 94 | int gemfield = 719; 95 | } 96 | 97 | //for 98 | for(...){ 99 | int gemfield = 719; 100 | } 101 | 102 | //switch 103 | switch( level ){ 104 | case gemfield::None: 105 | resip::Log::setLevel( resip::Log::None ); 106 | break; 107 | default: 108 | resip::Log::setLevel( resip::Log::None); 109 | break; 110 | } 111 | ``` 112 | 113 | ## python的import顺序 114 | 按照如下优先级顺序来import模块: 115 | - import标准库 116 | - import三方库 117 | - import cv系列 118 | - import torch系列 119 | - import torchvision系列 120 | - import deepvac系列 121 | - import 本项目 122 | -------------------------------------------------------------------------------- /deepvac/backbones/bottleneck_layer.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from .conv_layer import * 4 | 5 | 6 | class Bottleneck(nn.Module): 7 | expansion: int = 4 8 | def __init__(self, inplanes: int, outplanes: int, stride: int = 1): 9 | super(Bottleneck, self).__init__() 10 | self.conv1 = Conv2dBNReLU(in_planes=inplanes, out_planes=outplanes, kernel_size=1) 11 | self.conv2 = Conv2dBNReLU(in_planes=outplanes, out_planes=outplanes, kernel_size=3, stride=stride) 12 | 13 | outplanes_after_expansion = outplanes * self.expansion 14 | self.conv3 = nn.Conv2d(outplanes, outplanes_after_expansion, kernel_size=1) 15 | self.bn3 = nn.BatchNorm2d(outplanes_after_expansion) 16 | self.relu = nn.ReLU(inplace=True) 17 | 18 | self.stride = stride 19 | self.downsample = None 20 | 21 | if stride != 1 or inplanes != outplanes_after_expansion: 22 | self.downsample = nn.Sequential(nn.Conv2d(inplanes, outplanes_after_expansion, kernel_size=1, stride=stride, bias=False), nn.BatchNorm2d(outplanes_after_expansion)) 23 | 24 | def forward(self, x): 25 | identity = x 26 | out = self.conv1(x) 27 | out = self.conv2(out) 28 | out = self.conv3(out) 29 | out = self.bn3(out) 30 | 31 | if self.downsample is not None: 32 | identity = self.downsample(x) 33 | 34 | out += identity 35 | return self.relu(out) 36 | 37 | class BottleneckIR(nn.Module): 38 | def __init__(self, inplanes: int, outplanes: int, stride: int): 39 | super(BottleneckIR, self).__init__() 40 | 41 | self.shortcut_layer = nn.MaxPool2d(kernel_size=1, stride=stride) if inplanes == outplanes else nn.Sequential( 42 | nn.Conv2d(inplanes, outplanes, kernel_size=1, stride=stride ,bias=False), nn.BatchNorm2d(outplanes)) 43 | 44 | self.res_layer = nn.Sequential( 45 | nn.BatchNorm2d(inplanes), 46 | nn.Conv2d(inplanes, outplanes, kernel_size=3, stride=1, padding=1 ,bias=False), 47 | nn.PReLU(outplanes), 48 | nn.Conv2d(outplanes, outplanes, kernel_size=3, stride=stride, padding=1 ,bias=False), 49 | nn.BatchNorm2d(outplanes) 50 | ) 51 | 52 | def forward(self, x): 53 | shortcut = self.shortcut_layer(x) 54 | res = self.res_layer(x) 55 | return res + shortcut 56 | 57 | -------------------------------------------------------------------------------- /deepvac/cast/tensorrt.py: -------------------------------------------------------------------------------- 1 | from ..utils import LOG, syszux_once 2 | from .base import DeepvacCast 3 | 4 | TRT_LOGGER = None 5 | @syszux_once 6 | def initLog(trt): 7 | global TRT_LOGGER 8 | TRT_LOGGER = trt.Logger(trt.Logger.WARNING) 9 | 10 | class TensorrtCast(DeepvacCast): 11 | def auditConfig(self): 12 | if not self.config.model_dir: 13 | return False 14 | 15 | return True 16 | 17 | def process(self, cast_output_file=None): 18 | try: 19 | import tensorrt as trt 20 | initLog(trt) 21 | except: 22 | LOG.logE("You must install tensorrt package if you want to convert pytorch to onnx. 1. Download Tensorrt7.2.3(for CUDA11.0) from https://developer.nvidia.com/tensorrt \ 23 | 2. unpack Tensorrt*.tar.gz 3. pip install tensorrt-x-cpx-none-linux_x86_64.whl in Tensorrt*(your_tensorrt_path)/python", exit=True) 24 | 25 | output_trt_file = self.config.model_dir 26 | if cast_output_file: 27 | output_trt_file = '{}/trt__{}.trt'.format(self.trainer_config.output_dir, cast_output_file) 28 | self.config.model_dir = output_trt_file 29 | 30 | LOG.logI("config.trt_model_dir found, save tensorrt model to {}...".format(self.config.model_dir)) 31 | 32 | #to onnx, also set self.config.onnx_model_dir, self.config.onnx_input_names and self.config.onnx_output_names 33 | self.exportOnnx() 34 | 35 | with trt.Builder(TRT_LOGGER) as builder, builder.create_network(1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) as network, trt.OnnxParser(network, TRT_LOGGER) as parser: 36 | with open(self.config.onnx_model_dir, 'rb') as model: 37 | parser.parse(model.read()) 38 | config = builder.create_builder_config() 39 | config.max_workspace_size = 1 << 30 40 | 41 | if self.config.onnx_dynamic_ax: 42 | profile = builder.create_optimization_profile() 43 | profile.set_shape(self.config.onnx_input_names[0], self.config.input_min_dims, self.config.input_opt_dims, self.config.input_max_dims) 44 | config.add_optimization_profile(profile) 45 | engine = builder.build_engine(network, config) 46 | with open(output_trt_file, "wb") as f: 47 | f.write(engine.serialize()) 48 | LOG.logI("Pytorch model convert to TensorRT model succeed, save model in {}".format(output_trt_file)) 49 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib64/ 18 | parts/ 19 | sdist/ 20 | var/ 21 | wheels/ 22 | pip-wheel-metadata/ 23 | share/python-wheels/ 24 | *.egg-info/ 25 | .installed.cfg 26 | *.egg 27 | MANIFEST 28 | 29 | # PyInstaller 30 | # Usually these files are written by a python script from a template 31 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 32 | *.manifest 33 | *.spec 34 | 35 | # Installer logs 36 | pip-log.txt 37 | pip-delete-this-directory.txt 38 | 39 | # Unit test / coverage reports 40 | htmlcov/ 41 | .tox/ 42 | .nox/ 43 | .coverage 44 | .coverage.* 45 | .cache 46 | nosetests.xml 47 | coverage.xml 48 | *.cover 49 | *.py,cover 50 | .hypothesis/ 51 | .pytest_cache/ 52 | 53 | # Translations 54 | *.mo 55 | *.pot 56 | 57 | # Django stuff: 58 | *.log 59 | local_settings.py 60 | db.sqlite3 61 | db.sqlite3-journal 62 | 63 | # Flask stuff: 64 | instance/ 65 | .webassets-cache 66 | 67 | # Scrapy stuff: 68 | .scrapy 69 | 70 | # Sphinx documentation 71 | docs/_build/ 72 | 73 | # PyBuilder 74 | target/ 75 | 76 | # Jupyter Notebook 77 | .ipynb_checkpoints 78 | 79 | # IPython 80 | profile_default/ 81 | ipython_config.py 82 | 83 | # pyenv 84 | .python-version 85 | 86 | # pipenv 87 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 88 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 89 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 90 | # install all needed dependencies. 91 | #Pipfile.lock 92 | 93 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 94 | __pypackages__/ 95 | 96 | # Celery stuff 97 | celerybeat-schedule 98 | celerybeat.pid 99 | 100 | # SageMath parsed files 101 | *.sage.py 102 | 103 | # Environments 104 | .env 105 | .venv 106 | env/ 107 | venv/ 108 | ENV/ 109 | env.bak/ 110 | venv.bak/ 111 | 112 | # Spyder project settings 113 | .spyderproject 114 | .spyproject 115 | 116 | # Rope project settings 117 | .ropeproject 118 | 119 | # mkdocs documentation 120 | /site 121 | 122 | # mypy 123 | .mypy_cache/ 124 | .dmypy.json 125 | dmypy.json 126 | 127 | # Pyre type checker 128 | .pyre/ 129 | -------------------------------------------------------------------------------- /deepvac/cast/tnn.py: -------------------------------------------------------------------------------- 1 | import os 2 | import time 3 | from ..utils import LOG 4 | from .base import DeepvacCast 5 | 6 | class TnnCast(DeepvacCast): 7 | def auditConfig(self): 8 | if not self.config.model_dir: 9 | return False 10 | 11 | return True 12 | 13 | def process(self, cast_output_file=None): 14 | try: 15 | import onnx2tnn 16 | except: 17 | LOG.logE("You must install onnx2tnn package if you want to convert pytorch to tnn. \ 18 | 1. compile onnx2tnn, reference https://github.com/Tencent/TNN/blob/master/doc/en/user/onnx2tnn_en.md#compile. \ 19 | 2. copy /tools/onnx2tnn/onnx-converter/*.so to your python site-packages locate(you can use pip3 -V show where your packages are located).", exit=True) 20 | return 21 | 22 | output_tnn_file = self.config.model_dir 23 | if cast_output_file: 24 | output_tnn_file = '{}/tnn__{}.tnnmodel'.format(self.trainer_config.output_dir, cast_output_file) 25 | self.config.model_dir = output_tnn_file 26 | self.config.arch_dir = '{}.tnnproto'.format(output_tnn_file) 27 | 28 | LOG.logI("config.cast.TnnCast.model_dir found, save tnn model to {}...".format(self.config.model_dir)) 29 | 30 | #to onnx, also set self.config.onnx_model_dir, self.config.onnx_input_names and self.config.onnx_output_names 31 | self.exportOnnx() 32 | if self.config.optimize: 33 | self.onnxSimplify(input_shapes={self.config.onnx_input_names[0]: self.trainer_config.sample.shape} if self.config.onnx_dynamic_ax else None, perform_optimization=False, dynamic_input_shape=True if self.config.onnx_dynamic_ax else False) 34 | 35 | output_dir = os.path.dirname(self.config.model_dir) 36 | command = "python3 -c \"import onnx2tnn;onnx2tnn.convert('{}','{}','v1.0','{}',0,{},'')\"". \ 37 | format(self.config.onnx_model_dir, output_dir, time.strftime("%Y%m%d %H:%M:%S", time.localtime()), str(1 if self.config.optimize else 0)) 38 | rc, out_text, err_text = self.runCmd(command) 39 | 40 | if rc != 0: 41 | LOG.logE(err_text + ". Error occured when export tnn model.", exit=True) 42 | 43 | #rename tnnmodel and tnnproto, cuz we can't control tnn model name 44 | prefix = os.path.splitext(os.path.basename(self.config.onnx_model_dir))[0] 45 | os.rename(os.path.join(output_dir, prefix+".tnnmodel"), self.config.model_dir) 46 | os.rename(os.path.join(output_dir, prefix+".tnnproto"), self.config.arch_dir) 47 | LOG.logI("Pytorch model convert to Tnn model succeed, save tnn model file in {}, save tnn proto file in {}".format(self.config.model_dir, self.config.arch_dir)) -------------------------------------------------------------------------------- /deepvac/utils/crypto.py: -------------------------------------------------------------------------------- 1 | import base64 2 | import hashlib 3 | import os 4 | import sys 5 | import string 6 | import random 7 | from .log import LOG 8 | 9 | try: 10 | from Crypto import Random 11 | from Crypto.Cipher import AES 12 | except: 13 | LOG.logE("Crypto package missing, you must install it with [pip3 install pycrypto]", exit=True) 14 | 15 | class SyszuxCipher(object): 16 | def __init__(self, key): 17 | self.aes_bs = AES.block_size 18 | LOG.logI('Using block size = {}'.format(self.aes_bs)) 19 | self.key = hashlib.sha256(key.encode()).digest() 20 | LOG.logI('Hash of key={} is {}'.format(key, self.key)) 21 | 22 | def encrypt(self, pt_byte): 23 | pt_byte = self._pad(pt_byte) 24 | iv = Random.new().read(AES.block_size) 25 | LOG.logI('iv: {}'.format(iv)) 26 | cipher = AES.new(self.key, AES.MODE_CBC, iv) 27 | return iv + cipher.encrypt(pt_byte) 28 | 29 | def decrypt(self, pt_byte): 30 | iv = pt_byte[:self.aes_bs] 31 | cipher = AES.new(self.key, AES.MODE_CBC, iv) 32 | return self._unpad(cipher.decrypt(pt_byte[self.aes_bs:]))#.decode('utf-8') 33 | 34 | def _pad(self, s): 35 | return s + str.encode((self.aes_bs - len(s) % self.aes_bs) * chr(self.aes_bs - len(s) % self.aes_bs)) 36 | 37 | def __call__(self, path, mode='encrypt'): 38 | filename, file_extension = os.path.splitext(path) 39 | 40 | if mode == 'encrypt': 41 | output_filename = "{}_syszux.so".format(filename) 42 | else: 43 | output_filename = "{}.{}".format(filename, 'pt') 44 | 45 | if not os.path.exists(path): 46 | LOG.logE('File does not exist at path: %s'.format(path), exit=True) 47 | 48 | pt_byte = None 49 | with open(path, 'rb') as f: 50 | pt_byte = f.read() 51 | 52 | if mode == 'encrypt': 53 | pt_byte = self.encrypt(pt_byte) 54 | else: 55 | pt_byte = self.decrypt(pt_byte) 56 | 57 | with open(output_filename, 'wb') as f: 58 | f.write(pt_byte) 59 | 60 | LOG.logI('Saved with key={} to {}'.format(self.key, output_filename)) 61 | 62 | @staticmethod 63 | def _unpad(s): 64 | return s[:-ord(s[-1:])] 65 | 66 | if __name__ == "__main__": 67 | modes = ['encrypt', 'decrypt'] 68 | if len(sys.argv) < 3: 69 | LOG.logE('Usage: {} '.format(sys.argv[0]), exit=True) 70 | 71 | file = sys.argv[1] 72 | mode = sys.argv[2] 73 | if mode not in modes: 74 | LOG.logE("Unsupported mode: {}".format(mode), exit=True) 75 | 76 | key = sys.argv[3] 77 | sc = SyszuxCipher(key) 78 | 79 | LOG.logI("You will use {} to {} {}".format(key, mode, file)) 80 | sc(file, mode) 81 | -------------------------------------------------------------------------------- /deepvac/cast/ncnn.py: -------------------------------------------------------------------------------- 1 | import os 2 | from ..utils import LOG 3 | from .base import DeepvacCast 4 | 5 | class NcnnCast(DeepvacCast): 6 | def auditConfig(self): 7 | if not self.config.model_dir: 8 | return False 9 | 10 | self.config.onnx2ncnn = self.addUserConfig("onnx2ncnn", self.config.onnx2ncnn, "/bin/onnx2ncnn") 11 | exist = os.path.exists(self.config.onnx2ncnn) 12 | if not exist: 13 | LOG.logE("The config.cast.NcnnCast.onnx2ncnn is invalid. We try to use default setting of config.cast.NcnnCast.onnx2ncnn is failed too. \ 14 | If you want to compile onnx2ncnn tools, reference https://github.com/Tencent/ncnn/wiki/how-to-build#build-for-linux-x86 ", exit=True) 15 | return False 16 | 17 | return True 18 | 19 | def process(self, cast_output_file=None): 20 | output_ncnn_file = self.config.model_dir 21 | 22 | if cast_output_file: 23 | output_ncnn_file = '{}/ncnn__{}.bin'.format(self.trainer_config.output_dir, cast_output_file) 24 | self.config.model_dir = output_ncnn_file 25 | self.config.arch_dir = '{}.param'.format(output_ncnn_file) 26 | LOG.logI("config.cast.NcnnCast.model_dir found, save ncnn model to {}...".format(self.config.model_dir)) 27 | 28 | #to onnx, also set self.config.onnx_model_dir, self.config.onnx_input_names and self.config.onnx_output_names 29 | self.exportOnnx() 30 | 31 | cmd = "{} {} {} {}".format(self.config.onnx2ncnn, self.config.onnx_model_dir, self.config.arch_dir, output_ncnn_file) 32 | rc, out_text, err_text = self.runCmd(cmd) 33 | if err_text == "": 34 | LOG.logI("Pytorch model convert to NCNN model succeed, save ncnn param file in {}, save ncnn bin file in {}.".format(self.config.arch_dir, output_ncnn_file)) 35 | return 36 | 37 | LOG.logE(err_text + ". Error occured when export ncnn model. Deepvac try to simplify the model first.") 38 | 39 | try: 40 | import onnx 41 | from onnxsim import simplify 42 | except: 43 | LOG.logE("You must install onnx and onnxsim package first if you want to convert pytorch to ncnn.", exit=True) 44 | 45 | dynamic_input_shape = False 46 | input_shapes = None 47 | if self.config.onnx_dynamic_ax: 48 | dynamic_input_shape = True 49 | input_shapes = {self.config.onnx_input_names[0]: self.trainer_config.sample.shape} 50 | 51 | self.onnxSimplify(check_n=3, perform_optimization=True, skip_fuse_bn=True, skip_shape_inference=False, dynamic_input_shape=dynamic_input_shape, input_shapes=input_shapes) 52 | 53 | rc, out_text, err_text = self.runCmd(cmd) 54 | if err_text != "": 55 | LOG.logE(err_text + ". we can't guarantee generate model is right.", exit=True) 56 | 57 | if not os.path.isfile(output_ncnn_file): 58 | LOG.logE("Error: ncnn model not generated due to internal error!", exit=True) 59 | 60 | LOG.logI("Pytorch model convert to NCNN model succeed, save ncnn param file in {}, save ncnn bin file in {}".format(self.config.arch_dir, output_ncnn_file)) -------------------------------------------------------------------------------- /deepvac/core/config.py: -------------------------------------------------------------------------------- 1 | import copy 2 | primitive = (int, float, str, bool, list, dict, tuple, set) 3 | 4 | class AttrDict(dict): 5 | def __getattr__(self, key): 6 | return self.get(key) 7 | 8 | def __setattr__(self, key, value): 9 | if key in self.__dict__: 10 | self.__dict__[key] = value 11 | else: 12 | self[key] = value 13 | 14 | def __deepcopy__(self, memo=None): 15 | try: 16 | ad = AttrDict(copy.deepcopy(dict(self), memo=memo)) 17 | except Exception as e: 18 | print("Warning, issue happened in clone() API. Error: {}".format(str(e))) 19 | ad = AttrDict() 20 | return ad 21 | 22 | def clone(self): 23 | return copy.deepcopy(self) 24 | 25 | def interpret(name): 26 | fields = name.split('.') 27 | c = AttrDict() 28 | if len(fields) < 2: 29 | return c 30 | cd = c 31 | for f in fields[1:]: 32 | cd[f] = AttrDict() 33 | cd=cd[f] 34 | return c 35 | 36 | def newDict(): 37 | return AttrDict() 38 | 39 | def new(trainer=None): 40 | config = AttrDict() 41 | config.core = AttrDict() 42 | config.feature = AttrDict() 43 | config.aug = AttrDict() 44 | config.cast = AttrDict() 45 | config.backbones = AttrDict() 46 | config.loss = AttrDict() 47 | config.datasets = AttrDict() 48 | if trainer is None: 49 | return config 50 | 51 | config.core[trainer] = AttrDict() 52 | ## ------------- common ------------------ 53 | config.core[trainer].device = "cuda" 54 | config.core[trainer].output_dir = "output" 55 | config.core[trainer].log_dir = "log" 56 | config.core[trainer].log_every = 10 57 | config.core[trainer].disable_git = False 58 | config.core[trainer].cast2cpu = True 59 | config.core[trainer].model_reinterpret_cast=False 60 | config.core[trainer].cast_state_dict_strict=True 61 | config.core[trainer].model_path_omit_keys=[] 62 | config.core[trainer].net_omit_keys_strict=[] 63 | 64 | ## ----------------- ddp -------------------- 65 | config.core[trainer].dist_url = "tcp://localhost:27030" 66 | config.core[trainer].world_size = 2 67 | config.core[trainer].shuffle = False 68 | 69 | ## ------------------- loader ------------------ 70 | config.core[trainer].num_workers = 3 71 | config.core[trainer].nominal_batch_factor = 1 72 | 73 | ## ------------------ train ------------------ 74 | config.core[trainer].train_batch_size = 128 75 | config.core[trainer].epoch_num = 30 76 | #model save number during an epoch 77 | config.core[trainer].save_num = 5 78 | config.core[trainer].checkpoint_suffix = '' 79 | 80 | ## ----------------- val ------------------ 81 | config.core[trainer].val_batch_size = None 82 | config.core[trainer].acc = 0 83 | 84 | return config 85 | 86 | def fork(deepvac_config, field=['aug','datasets']): 87 | if not isinstance(field, list): 88 | field = [field] 89 | c = new() 90 | for f in field: 91 | if f not in deepvac_config.keys(): 92 | print("ERROR: deepvac fork found unsupport field: {}".format(f)) 93 | return None 94 | c[f] = deepvac_config[f].clone() 95 | return c -------------------------------------------------------------------------------- /deepvac/datasets/file_line.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | import cv2 4 | from PIL import Image 5 | from ..utils import LOG 6 | from .base_dataset import DatasetBase 7 | 8 | class FileLineDataset(DatasetBase): 9 | def __init__(self, deepvac_config, fileline_path, delimiter=' ', sample_path_prefix=''): 10 | super(FileLineDataset, self).__init__(deepvac_config) 11 | self.sample_path_prefix = sample_path_prefix 12 | self.fileline_path = fileline_path 13 | self.delimiter = delimiter 14 | self.samples = [] 15 | mark = [] 16 | 17 | with open(self.fileline_path) as f: 18 | for line in f: 19 | label = self._buildLabelFromLine(line) 20 | self.samples.append(label) 21 | mark.append(label[1]) 22 | 23 | self.len = len(self.samples) 24 | self.class_num = len(np.unique(mark)) 25 | LOG.logI('FileLineDataset size: {} / {}'.format(self.len, self.class_num)) 26 | 27 | def _buildLabelFromLine(self, line): 28 | line = line.strip().split(self.delimiter) 29 | return [line[0], int(line[1])] 30 | 31 | def __getitem__(self, index): 32 | path, target = self.samples[index] 33 | abs_path = os.path.join(self.sample_path_prefix, path) 34 | sample = self._buildSampleFromPath(abs_path) 35 | sample = self.compose(sample) 36 | return sample, target 37 | 38 | def _buildSampleFromPath(self, abs_path): 39 | #we just set default loader with Pillow Image 40 | sample = Image.open(abs_path).convert('RGB') 41 | return sample 42 | 43 | def __len__(self): 44 | return self.len 45 | 46 | class FileLineCvStrDataset(FileLineDataset): 47 | def _buildLabelFromLine(self, line): 48 | line = line.strip().split(self.delimiter, 1) 49 | return [line[0], line[1]] 50 | 51 | def _buildSampleFromPath(self, abs_path): 52 | #we just set default loader with Pillow Image 53 | sample = cv2.imread(abs_path) 54 | return sample 55 | 56 | class FileLineCvSegDataset(FileLineCvStrDataset): 57 | def _buildLabelFromPath(self, abs_path): 58 | sample = cv2.imread(abs_path, cv2.IMREAD_GRAYSCALE) 59 | return sample 60 | 61 | def __getitem__(self, index): 62 | image_path, label_path = self.samples[index] 63 | sample = self._buildSampleFromPath(os.path.join(self.sample_path_prefix, image_path.strip())) 64 | label = self._buildLabelFromPath(os.path.join(self.sample_path_prefix, label_path.strip())) 65 | 66 | sample, label = self.compose([sample, label]) 67 | return sample, label 68 | 69 | class FileLineCvSegAuditDataset(FileLineCvSegDataset): 70 | def __getitem__(self, index): 71 | image_path, label_path = self.samples[index] 72 | image_path = os.path.join(self.sample_path_prefix, image_path.strip() ) 73 | label_path = os.path.join(self.sample_path_prefix, label_path.strip() ) 74 | sample = self._buildSampleFromPath(image_path) 75 | label = self._buildLabelFromPath(label_path) 76 | 77 | cls_masks = {} 78 | for cls_idx in np.unique(label): 79 | cls_masks[cls_idx] = label == cls_idx 80 | sample, label, cls_masks, label_path = self.compose([sample, label, cls_masks, label_path]) 81 | return sample, label, cls_masks, label_path -------------------------------------------------------------------------------- /deepvac/experimental/core/distill.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | from ...core import DeepvacTrain 3 | from ...utils import LOG 4 | 5 | class DeepvacDistill(DeepvacTrain): 6 | def auditConfig(self): 7 | super(DeepvacDistill, self).auditConfig() 8 | if self.config.teacher.scheduler is None: 9 | LOG.logE("You must set config.train.teacher.scheduler in config.py.", exit=True) 10 | LOG.logI("You set config.train.teacher.scheduler to {}".format(self.config.teacher.scheduler)) 11 | 12 | if self.config.teacher.optimizer is None: 13 | LOG.logE("You must set config.train.teacher.optimizer in config.py.",exit=True) 14 | LOG.logI("You set config.train.teacher.optimizer to {} in config.py".format(self.config.teacher.optimizer)) 15 | 16 | def initNetWithCode(self): 17 | super(DeepvacDistill, self).initNetWithCode() 18 | if self.config.teacher.net is None: 19 | LOG.logE("You must implement and set config.train.teacher.net to a torch.nn.Module instance in config.py.", exit=True) 20 | if not isinstance(self.config.teacher.net, nn.Module): 21 | LOG.logE("You must set config.train.teacher.net to a torch.nn.Module instance in config.py.", exit=True) 22 | 23 | def initStateDict(self): 24 | super(DeepvacDistill, self).initStateDict() 25 | LOG.logI('Loading State Dict from {}'.format(self.config.teacher.model_path)) 26 | self.config.teacher.state_dict = self.auditStateDict(self.config.teacher) 27 | 28 | def loadStateDict(self): 29 | super(DeepvacDistill, self).loadStateDict() 30 | self.config.teacher.net = self.config.teacher.net.to(self.config.device) 31 | if not self.config.teacher.state_dict: 32 | LOG.logI("self.config.teacher.state_dict not initialized, omit loadStateDict()") 33 | return 34 | 35 | if self.config.model_reinterpret_cast: 36 | self.config.teacher.state_dict = self.castStateDict(self.config.teacher) 37 | 38 | self.config.teacher.net.load_state_dict(self.config.teacher.state_dict, strict=False) 39 | 40 | def doForward(self): 41 | super(DeepvacDistill, self).doForward() 42 | self.config.teacher.output = self.config.teacher.net(self.config.sample) 43 | 44 | def initCriterion(self): 45 | super(DeepvacDistill, self).initCriterion() 46 | if self.config.teacher.criterion is None: 47 | LOG.logE("You must set config.train.criterion in config.py, e.g. config.train.teacher.criterion=torch.nn.CrossEntropyLoss()",exit=True) 48 | LOG.logI("You set config.train.teacher.criterion to {}".format(self.config.teacher.criterion)) 49 | 50 | def doLoss(self): 51 | super(DeepvacDistill, self).doLoss() 52 | LOG.logE("You have to reimplement doLoss() in DeepvacDistill subclass {} and set config.train.teacher.loss".format(self.name()), exit=True) 53 | 54 | def doBackward(self): 55 | super(DeepvacDistill, self).doBackward() 56 | if self.config.teacher.use_backword is True: 57 | self.config.teacher.loss.backward() 58 | 59 | def doOptimize(self): 60 | super(DeepvacDistill, self).doOptimize() 61 | if self.config.iter % self.config.nominal_batch_factor != 0: 62 | return 63 | self.config.teacher.optimizer.step() 64 | self.config.teacher.optimizer.zero_grad() 65 | -------------------------------------------------------------------------------- /docs/deepvac_standard.md: -------------------------------------------------------------------------------- 1 | # DeepVAC标准 2 | DeepVAC标准是由MLab团队设立，用来定义和约束AI模型的训练、测试、验收、部署。覆盖以下内容： 3 | - 数据集 4 | - 代码 5 | - 训练 6 | - 部署 7 | 8 | 基于此的[DeepVAC checklist (检查单)](./deepvac_checklist.md)是项目管理中必不可少的文件。 9 | 10 | ## 数据集 11 | 数据集划分为如下四种： 12 | | 名称 | 说明 | 和训练集分布一致？ |来源 | 准备人员 | 维护在 | 13 | |------------|---------|---------|---------|---------|---------| 14 | | 训练集| 基础标注数据中划分80% ~ 90%| 是 |公开数据集/自行标注| 算法人员 | 维护在MLab RookPod 存储上| 15 | | 验证集| 基础标注数据中划分10% ~ 20%| 是 |公开数据集/自行标注| 算法人员 |维护在MLab RookPod 存储上| 16 | | 测试集| 用来对模型做初步测试 | 否 |公开数据集/自行标注| 算法人员 |维护在MLab RookPod 存储上| 17 | | 验收集| 用来对模型做终极测试| 否|业务场景中真实数据的标注，其分布对算法人员不可见| 验收人员| 维护在MLab RookPod 存储上| 18 | 19 | ## 代码 20 | 包含如下三个方面： 21 | #### 使用[deepvac库](./lib.md) 22 | 代码必须基于deepvac库，且维护在MLab代码服务系统上。特别的： 23 | - 训练、验证代码基于DeepvacTrain类； 24 | - 测试代码基于Deepvac类； 25 | - 日志基于LOG模块； 26 | - 配置基于config模块; 27 | - 数据合成基于本项目的synthesis目录，或者通用的Synthesis项目; 28 | - 数据增强基于deepvac.aug模块及该模块中的Composer类体系; 29 | - 自定义dataloader基于deepvac.datasets; 30 | - 模型性能报告基于Report模块; 31 | - 模型定义基于deepvac.backbones; 32 | - 损失函数基于deepvac.loss； 33 | - 模型到各平台的转换基于deepvac.cast； 34 | 35 | #### 使用DeepVAC代码规范 36 | 访问：[代码规范](./code_standard.md) 37 | 38 | #### 使用DeepVAC软件工程规范 39 | 访问：[软件工程规范](./arch.md) 40 | 41 | 42 | ## 训练 43 | 模型训练默认包含两种： 44 | - 部署目标为x86+CUDA Linux的训练; 45 | - 部署目标为x86 Linux、Arm Linux、ARM Android/iOS的训练； 46 | 47 | 并且最终需要满足： 48 | - 在验收集上的指标必须符合要求； 49 | - 各种SOTA模型要维护在MLab RookPod 存储上； 50 | - 报告要记录在deepvac-product项目页上。报告格式如下： 51 | 52 | dataset: <验收集的名称> 53 | tester: <测试人员的名称> 54 | date: <测试日期> 55 | 56 | |dataset|total|duration|accuracy|precision|recall|miss|error| 57 | |--|--|--|--|--|--|--|--| 58 | |gemfield|144|75.579|0.631944444444|1.0|0.631944444444|0.118055555556|0.368055555556| 59 | |self|2047|1065.102|0.829018075232|1.0|0.829018075232|0.0229604298974|0.170981924768| 60 | 61 | #### 部署目标为x86+CUDA Linux的训练 62 | 开启如下开关： 63 | - config.cast.ScriptCast.model_dir（必须） 64 | - config.cast.TraceCast.model_dir（可选） 65 | - config.cast.ScriptCast.static_quantize_dir（必须） 66 | - config.cast.ScriptCast.dynamic_quantize_dir（可选） 67 | - config.cast.TraceCast.static_quantize_dir（必须） 68 | - config.cast.TraceCast.dynamic_quantize_dir（可选） 69 | - config.core..ema（可选） 70 | - config.core..tensorboard_*（可选） 71 | - config.core..amp（可选） 72 | - config.core..dist_url（可选） 73 | 74 | #### 部署目标为x86 Linux、Arm Linux、ARM Android/iOS的训练 75 | 在部署目标为x86+CUDA Linux的训练基础上，开启如下开关： 76 | - config.cast.OnnxCast.onnx_model_dir（可选，需要ONNX时开启） 77 | - config.cast.NcnnCast.model_dir, config.cast.NcnnCast.onnx2ncnn（可选，需要NCNN时开启） 78 | - config.cast.CoremlCast.model_dir（可选，需要CoreML时开启） 79 | - config.cast.MnnCast.model_dir, config.cast.MnnCast.onnx2mnn（可选，需要MNN时开启） 80 | - config.cast.TnnCast.model_dir（可选，需要TNN时开启） 81 | - config.cast.TensorrtCast.model_dir（可选，需要TensorRT时开启） 82 | ## 部署方式 83 | 所有的AI产品默认进行3种部署测试： 84 | - x86 + CUDA Linux 85 | - x86 Linux 86 | - Arm Android 87 | 88 | 可选2种部署测试： 89 | - Arm iOS 90 | - Arm Linux。 91 | 92 | 整体部署测试说明如下： 93 | | 部署目标 | 部署方式 | 大小 | 是否必须 | 94 | |------------|---------|---------|---------| 95 | |x86 + CUDA Linux| 基于 libdeepvac cuda docker image | docker image大小 | 是| 96 | |x86 Linux | 基于 libdeepvac x86 docker image | docker image大小 | 是| 97 | |ARM Android | App + libdeepvac.so | libdeepvac.so大小 | 是| 98 | |Arm iOS | App + libdeepvac.dylib | libdeepvac.dylib大小 | 否| 99 | |Arm Linux | App + libdeepvac.so | libdeepvac.so大小 | 否| 100 | 101 | 测试成功后： 102 | - Docker镜像维护在MLab docker registry上； 103 | - 库维护在MLab RookPod 存储上； 104 | 105 | ## DeepVAC checklist 106 | 请访问：[DeepVAC checklist (检查单)](./deepvac_checklist.md) 107 | 108 | -------------------------------------------------------------------------------- /deepvac/cast/cast_helper.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.quantization.fuser_method_mappings import DEFAULT_OP_LIST_TO_FUSER_METHOD 3 | from torch.quantization import QuantStub, DeQuantStub 4 | 5 | def get_module(model, submodule_key): 6 | tokens = submodule_key.split('.') 7 | cur_mod = model 8 | for s in tokens: 9 | cur_mod = getattr(cur_mod, s) 10 | return cur_mod 11 | 12 | def set_module(model, submodule_key, module): 13 | tokens = submodule_key.split('.') 14 | sub_tokens = tokens[:-1] 15 | cur_mod = model 16 | for s in sub_tokens: 17 | cur_mod = getattr(cur_mod, s) 18 | setattr(cur_mod, tokens[-1], module) 19 | 20 | def get_fuser_module_index(mod_list): 21 | rc = [] 22 | if len(mod_list) < 2: 23 | return rc 24 | keys = sorted(list(DEFAULT_OP_LIST_TO_FUSER_METHOD.keys()), key=lambda x: len(x), reverse=True) 25 | mod2fused_list = [list(x) for x in keys] 26 | 27 | for mod2fused in mod2fused_list: 28 | if len(mod2fused) > len(mod_list): 29 | continue 30 | mod2fused_idx = [(i, i+len(mod2fused)) for i in range(len(mod_list) - len(mod2fused) + 1) if mod_list[i:i+len(mod2fused)] == mod2fused] 31 | if not mod2fused_idx: 32 | continue 33 | 34 | for idx in mod2fused_idx: 35 | start,end = idx 36 | mod_list[start: end] = [None] * len(mod2fused) 37 | 38 | rc.extend(mod2fused_idx) 39 | 40 | return rc 41 | 42 | def auto_fuse_model(model): 43 | module_names = [] 44 | module_types = [] 45 | for name, m in model.named_modules(): 46 | module_names.append(name) 47 | module_types.append(type(m)) 48 | 49 | if len(module_types) < 2: 50 | return model 51 | 52 | module_idxs = get_fuser_module_index(module_types) 53 | modules_to_fuse = [module_names[mi[0]:mi[1]] for mi in module_idxs] 54 | new_model = torch.quantization.fuse_modules(model, modules_to_fuse) 55 | return new_model 56 | 57 | class DeQuantStub(nn.Module): 58 | def __init__(self): 59 | super(DeQuantStub, self).__init__() 60 | 61 | def forward(self, x: Any) -> Any: 62 | return x 63 | 64 | def calibrate(model, data_loader): 65 | model.eval() 66 | with torch.no_grad(): 67 | for sample, target in data_loader: 68 | model(sample) 69 | 70 | def prepareQAT(deepvac_core_config): 71 | if not deepvac_core_config.qat_dir: 72 | return 73 | 74 | if deepvac_core_config.is_forward_only: 75 | LOG.logI("You are in forward_only mode, omit the parepareQAT()") 76 | return 77 | 78 | LOG.logI("You have enabled QAT, this step is only for prepare.") 79 | 80 | if deepvac_core_config.qat_net_prepared: 81 | LOG.logE("Error: You have already prepared the model for QAT.", exit=True) 82 | 83 | backend = 'fbgemm' 84 | if deepvac_core_config.quantize_backend: 85 | backend = deepvac_core_config.quantize_backend 86 | deepvac_core_config.qat_net_prepared = DeepvacQAT(deepvac_core_config.net).to(deepvac_core_config.device) 87 | deepvac_core_config.qat_net_prepared.qconfig = torch.quantization.get_default_qat_qconfig(backend) 88 | torch.quantization.prepare_qat(deepvac_core_config.qat_net_prepared, inplace=True) 89 | #after this, train net will be transfered to QAT ! 90 | deepvac_core_config.net = deepvac_core_config.qat_net_prepared 91 | 92 | class DeepvacQAT(torch.nn.Module): 93 | def __init__(self, net2qat): 94 | super(DeepvacQAT, self).__init__() 95 | self.quant = QuantStub() 96 | self.net2qat = auto_fuse_model(net2qat) 97 | self.dequant = DeQuantStub() 98 | 99 | def forward(self, x): 100 | x = self.quant(x) 101 | x = self.net2qat(x) 102 | x = self.dequant(x) 103 | return x -------------------------------------------------------------------------------- /deepvac/backbones/res_layer.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | from .module_helper import makeDivisible 3 | from .conv_layer import Conv2dBNReLU, Conv2dBNHswish 4 | from .se_layer import SELayer 5 | from .activate_layer import hswish 6 | 7 | class ResBnBlock(nn.Module): 8 | def __init__(self, in_channel, expand_channel, out_channel, kernel_size, stride, shortcut=True): 9 | super(ResBnBlock, self).__init__() 10 | self.layer = nn.Sequential( 11 | Conv2dBNReLU(in_channel, expand_channel, 1, 1), 12 | Conv2dBNReLU(expand_channel, expand_channel, kernel_size, stride, kernel_size//2, expand_channel), 13 | Conv2dBN(expand_channel, out_channel, 1, 1), 14 | ) 15 | self.shortcut = shortcut and (in_channel == out_channel) 16 | 17 | def forward(self, x): 18 | return self.layer(x) + x if self.shortcut else self.layer(x) 19 | 20 | class InvertedResidual(nn.Module): 21 | def __init__(self, inp, oup, kernel_size, stride, expand_ratio, use_se, use_hs, padding=None): 22 | super(InvertedResidual, self).__init__() 23 | if padding is None: 24 | padding = (kernel_size - 1) // 2 25 | hidden_dim = makeDivisible(inp * expand_ratio, 8) 26 | assert stride in [1, 2, (2, 1)] 27 | assert kernel_size in [3,5] 28 | 29 | self.use_res_connect = stride == 1 and inp == oup 30 | layers = [] 31 | if expand_ratio != 1: 32 | layers.append(Conv2dBNHswish(inp, hidden_dim, kernel_size=1)) # pw 33 | 34 | layers.extend([ 35 | # dw 36 | nn.Conv2d(hidden_dim, hidden_dim, kernel_size, stride, padding, groups=hidden_dim, bias=False), 37 | nn.BatchNorm2d(hidden_dim), 38 | hswish() if use_hs else nn.ReLU(inplace=True), 39 | # Squeeze-and-Excite 40 | SELayer(hidden_dim) if use_se else nn.Identity(), 41 | # pw-linear 42 | nn.Conv2d(hidden_dim, oup, kernel_size=1), 43 | nn.BatchNorm2d(oup), 44 | ]) 45 | self.conv = nn.Sequential(*layers) 46 | 47 | def forward(self, x): 48 | if self.use_res_connect: 49 | return x + self.conv(x) 50 | else: 51 | return self.conv(x) 52 | 53 | class InvertedResidualFacialKpBlock(nn.Module): 54 | def __init__(self, in_channels: int, out_channels: int, expansion_factor: int=6, kernel_size: int=3, stride: int=2, padding: int=1, 55 | is_residual: bool=True): 56 | super(InvertedResidualFacialKpBlock, self).__init__() 57 | assert stride in [1, 2] 58 | 59 | self.block = nn.Sequential( 60 | nn.Conv2d(in_channels, in_channels * expansion_factor, 1, bias=False), 61 | nn.BatchNorm2d(in_channels * expansion_factor), 62 | nn.ReLU(inplace=True), 63 | nn.Conv2d(in_channels * expansion_factor, in_channels * expansion_factor, 64 | kernel_size, stride, padding, 1, 65 | groups=in_channels * expansion_factor, bias=False), 66 | nn.BatchNorm2d(in_channels * expansion_factor), 67 | nn.ReLU(inplace=True), 68 | nn.Conv2d(in_channels * expansion_factor, out_channels, 1, 69 | bias=False), 70 | nn.BatchNorm2d(out_channels)) 71 | 72 | self.is_residual = is_residual if stride == 1 else False 73 | self.is_conv_res = False if in_channels == out_channels else True 74 | #todo 75 | if stride == 1 and self.is_conv_res: 76 | self.conv_res = nn.Sequential(nn.Conv2d(in_channels, out_channels, 1, bias=False), 77 | nn.BatchNorm2d(out_channels)) 78 | 79 | def forward(self, x): 80 | block = self.block(x) 81 | if self.is_residual: 82 | if self.is_conv_res: 83 | return self.conv_res(x) + block 84 | return x + block 85 | return block -------------------------------------------------------------------------------- /deepvac/utils/log.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import os 3 | from datetime import datetime 4 | from enum import Enum 5 | import logging 6 | import subprocess 7 | 8 | def getCurrentGitBranch(): 9 | try: 10 | branch = subprocess.check_output(['git', 'rev-parse', '--abbrev-ref', 'HEAD']).strip().decode() 11 | except: 12 | branch = None 13 | return branch 14 | 15 | def getTime(): 16 | return (str(datetime.now())[:-10]).replace(' ','-').replace(':','-') 17 | 18 | def getArgv(): 19 | argv = ''.join(sys.argv) 20 | argv = argv.replace(' ','').replace('/','').replace('-','_').replace('.py','') 21 | return argv[:64] 22 | 23 | deepvac_branch_name = getCurrentGitBranch() 24 | if deepvac_branch_name is None: 25 | deepvac_branch_name = 'not_in_git' 26 | 27 | deepvac_pid = os.getpid() 28 | deepvac_time = getTime() 29 | 30 | #according deepvac standard, log should in log directory. 31 | if not os.path.exists('log'): 32 | os.makedirs('log') 33 | 34 | logger=logging.getLogger("DEEPVAC") 35 | logger.setLevel(logging.DEBUG) 36 | 37 | deepvac_log_format = '%(asctime)s %(levelname)-8s %(message)s' 38 | fh = logging.FileHandler('log/{}:{}:{}:{}.log'.format(deepvac_pid, getArgv(),deepvac_time, deepvac_branch_name)) 39 | fh.setLevel(logging.INFO) 40 | fh.setFormatter( logging.Formatter(deepvac_log_format) ) 41 | logger.addHandler(fh) 42 | 43 | # add console output 44 | console_log_format = '%(asctime)s {}:{} %(levelname)-8s %(message)s'.format(deepvac_pid, deepvac_branch_name) 45 | console = logging.StreamHandler() 46 | console.setLevel(logging.INFO) 47 | console.setFormatter( logging.Formatter(console_log_format) ) 48 | logger.addHandler(console) 49 | 50 | logger.info("deepvac log imported...") 51 | 52 | class LOG(object): 53 | class S(Enum): 54 | I = 'Info' 55 | W = 'Warning' 56 | E = 'Error' 57 | 58 | LOG_LEVEL = S.I 59 | logfunc = {S.I: lambda x : logger.info(x), 60 | S.W: lambda x : logger.warning(x), 61 | S.E: lambda x : logger.error(x) 62 | } 63 | 64 | intfunc = {S.I: 0, S.W: 1, S.E: 2} 65 | @staticmethod 66 | def log(level, str): 67 | if level not in LOG.logfunc.keys(): 68 | LOG.logfunc[LOG.S.E]("incorrect value of parameter level when call log function.") 69 | 70 | if LOG.intfunc[level] < LOG.intfunc[LOG.LOG_LEVEL]: 71 | return 72 | 73 | LOG.logfunc[level](str) 74 | 75 | @staticmethod 76 | def logI(str): 77 | LOG.log(LOG.S.I, str) 78 | 79 | @staticmethod 80 | def logW(str): 81 | LOG.log(LOG.S.W, str) 82 | 83 | @staticmethod 84 | def logE(str, exit=False): 85 | LOG.log(LOG.S.E, str) 86 | if exit: 87 | sys.exit(1) 88 | 89 | # according to DeepVAC Standard, we must run in git repo 90 | def assertAndGetGitBranch(is_disable_git): 91 | if os.environ.get("disable_git"): 92 | branch = "disable_git" 93 | return branch 94 | 95 | if is_disable_git: 96 | branch = "disable_git" 97 | return branch 98 | 99 | branch = getCurrentGitBranch() 100 | if branch is None: 101 | LOG.logE('According to deepvac standard, you must working in a git repo.', exit=True) 102 | 103 | if len(branch) < 4: 104 | LOG.logE('According to deepvac standard, your git branch name is too short: {}'.format(branch), exit=True) 105 | 106 | LOG.logI('You are running on git branch: {}'.format(branch)) 107 | 108 | if branch.startswith('LTS_'): 109 | return branch 110 | 111 | if branch.startswith('PROTO_'): 112 | return branch 113 | 114 | if branch in ['master','main']: 115 | LOG.logW('However, DeepVAC does not suggest running on the master/main branch.') 116 | return branch 117 | 118 | LOG.logE('According to DeepVAC standard, git branch name should start from LTS_ or PROTO_: {}'.format(branch), exit=True) 119 | -------------------------------------------------------------------------------- /deepvac/cast/coreml.py: -------------------------------------------------------------------------------- 1 | from ..utils import LOG 2 | from .base import DeepvacCast 3 | 4 | class CoremlCast(DeepvacCast): 5 | def auditConfig(self): 6 | if not self.config.model_dir: 7 | return False 8 | 9 | if not self.deepvac_cast_config.TraceCast.model_dir and not self.deepvac_cast_config.ScriptCast.model_dir: 10 | LOG.logE("CoreML converter now has dependency on TorchScript model, you need to enable config.cast.TraceCast.model_dir or config.cast.ScriptCast.model_dir", exit=True) 11 | 12 | if self.config.input_type not in [None, 'image','tensor']: 13 | LOG.logE("coreml input type must be {}".format([None, 'image','tensor']), exit=True) 14 | 15 | if self.config.input_type == 'image': 16 | LOG.logI("You are in coreml_input_type=image mod") 17 | if self.config.scale is None: 18 | LOG.logE("You need to set config.cast.CoremlCast.scale in config.py, e.g. config.cast.CoremlCast.scale = 1.0 / (0.226 * 255.0)", exit=True) 19 | 20 | if self.config.color_layout is None: 21 | LOG.logE("You need to set config.cast.CoremlCast.color_layout in config.py, e.g. config.cast.CoremlCast.color_layout = 'BGR' ", exit=True) 22 | 23 | if self.config.blue_bias is None: 24 | LOG.logE("You need to set config.cast.CoremlCast.blue_bias in config.py, e.g. config.cast.CoremlCast.blue_bias = -0.406 / 0.226 ", exit=True) 25 | 26 | if self.config.green_bias is None: 27 | LOG.logE("You need to set config.cast.CoremlCast.green_bias in config.py, e.g. config.cast.CoremlCast.green_bias = -0.456 / 0.226", exit=True) 28 | 29 | if self.config.red_bias is None: 30 | LOG.logE("You need to set config.cast.CoremlCast.red_bias in config.py, e.g. config.cast.CoremlCast.red_bias = -0.485 / 0.226 ", exit=True) 31 | 32 | return True 33 | 34 | def process(self, cast_output_file=None): 35 | try: 36 | import coremltools 37 | except: 38 | LOG.logE("You need to install coremltools package if you want to convert PyTorch to CoreML model. E.g. pip install --upgrade coremltools") 39 | return 40 | 41 | output_coreml_file = self.config.model_dir 42 | if cast_output_file: 43 | output_coreml_file = '{}/coreml__{}.mlmodel'.format(self.trainer_config.output_dir, cast_output_file) 44 | self.config.model_dir = output_coreml_file 45 | 46 | LOG.logI("config.cast.CoremlCast.model_dir found, save coreml model to {}...".format(self.config.model_dir)) 47 | model = self.deepvac_cast_config.ScriptCast.model_dir 48 | if self.deepvac_cast_config.ScriptCast.model_dir is None: 49 | model = self.deepvac_cast_config.TraceCast.model_dir 50 | #input mode 51 | if self.config.input_type == 'image': 52 | input = coremltools.ImageType(name="input", 53 | shape=tuple(self.trainer_config.sample.shape), 54 | scale=self.config.scale, 55 | color_layout = self.config.color_layout, 56 | bias = [self.config.blue_bias, self.config.green_bias, self.config.red_bias]) 57 | else: 58 | input = coremltools.TensorType(name='input', shape=tuple(self.trainer_config.sample.shape)) 59 | #classifier 60 | if self.config.classifier_config is not None: 61 | self.config.classifier_config = coremltools.ClassifierConfig(self.config.classifier_config) 62 | #convert 63 | coreml_model = coremltools.convert(model=model, inputs=[input], 64 | classifier_config=self.config.classifier_config, minimum_deployment_target=self.config.minimum_deployment_target) 65 | 66 | # Set feature descriptions (these show up as comments in XCode) 67 | coreml_model.input_description["input"] = "Deepvac Model Input" 68 | 69 | # Set model author name 70 | coreml_model.author = '"DeepVAC' 71 | 72 | # Set the license of the model 73 | coreml_model.license = "Deepvac Lincense" 74 | coreml_model.short_description = "Powered by DeepVAC" 75 | 76 | # Set a version for the model 77 | coreml_model.version = self.config.version if self.config.version else "1.0" 78 | # Save the CoreML model 79 | coreml_model.save(output_coreml_file) 80 | -------------------------------------------------------------------------------- /deepvac/backbones/mobilenet.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | from .module_helper import makeDivisible 3 | from .activate_layer import hswish 4 | from .conv_layer import Conv2dBNHswish 5 | from .res_layer import InvertedResidual 6 | from .weights_init import initWeightsKaiming 7 | 8 | class MobileNetV3(nn.Module): 9 | def __init__(self, class_num=1000, width_mult=1.): 10 | super(MobileNetV3, self).__init__() 11 | self.width_mult = width_mult 12 | self.class_num = class_num 13 | # setting of inverted residual blocks 14 | self.auditConfig() 15 | # building first layer 16 | input_channel = makeDivisible(16 * self.width_mult, 8) 17 | layers = [Conv2dBNHswish(3, input_channel, stride=2)] 18 | 19 | # building inverted residual blocks 20 | for k, t, c, use_se, use_hs, s in self.cfgs: 21 | exp_size = makeDivisible(input_channel * t, 8) 22 | output_channel = makeDivisible(c * self.width_mult, 8) 23 | layers.append(InvertedResidual(input_channel, output_channel, k, s, t, use_se, use_hs)) 24 | input_channel = output_channel 25 | self.features = nn.Sequential(*layers) 26 | # building last several layers 27 | self.conv = Conv2dBNHswish(input_channel, exp_size, kernel_size=1) 28 | self.fc_inp = exp_size 29 | self.initFc() 30 | initWeightsKaiming(self) 31 | 32 | def forward(self, x): 33 | x = self.features(x) 34 | x = self.conv(x) 35 | x = self.pool(x) 36 | x = x.view(x.size(0), -1) 37 | x = self.classifier(x) 38 | return x 39 | 40 | def initFc(self): 41 | self.pool = nn.AdaptiveAvgPool2d((1, 1)) 42 | output_channel = makeDivisible(self.last_output_channel * self.width_mult, 8) if self.width_mult > 1.0 else self.last_output_channel 43 | self.classifier = nn.Sequential( 44 | nn.Linear(self.fc_inp, output_channel), 45 | hswish(), 46 | nn.Dropout(0.2), 47 | nn.Linear(output_channel, self.class_num), 48 | ) 49 | 50 | def auditConfig(self): 51 | self.last_output_channel = 1024 52 | self.cfgs = [ 53 | # k, t, c, SE, HS, s 54 | [3, 1, 16, True, False, 2], 55 | [3, 4.5, 24, False, False, 2], 56 | [3, 3.67, 24, False, False, 1], 57 | [5, 4, 40, True, True, 2], 58 | [5, 6, 40, True, True, 1], 59 | [5, 6, 40, True, True, 1], 60 | [5, 3, 48, True, True, 1], 61 | [5, 3, 48, True, True, 1], 62 | [5, 6, 96, True, True, 2], 63 | [5, 6, 96, True, True, 1], 64 | [5, 6, 96, True, True, 1], 65 | ] 66 | 67 | 68 | class MobileNetV3Large(MobileNetV3): 69 | def auditConfig(self): 70 | self.last_output_channel = 1280 71 | self.cfgs = [ 72 | # k, t, c, SE, HS, s 73 | [3, 1, 16, False, False, 1], 74 | [3, 4, 24, False, False, 2], 75 | [3, 3, 24, False, False, 1], 76 | [5, 3, 40, True, False, 2], 77 | [5, 3, 40, True, False, 1], 78 | [5, 3, 40, True, False, 1], 79 | [3, 6, 80, False, True, 2], 80 | [3, 2.5, 80, False, True, 1], 81 | [3, 2.3, 80, False, True, 1], 82 | [3, 2.3, 80, False, True, 1], 83 | [3, 6, 112, True, True, 1], 84 | [3, 6, 112, True, True, 1], 85 | [5, 6, 160, True, True, 2], 86 | [5, 6, 160, True, True, 1], 87 | [5, 6, 160, True, True, 1] 88 | ] 89 | 90 | class MobileNetV3OCR(MobileNetV3): 91 | def auditConfig(self): 92 | self.last_output_channel = 1024 93 | self.cfgs = [ 94 | # k, t, c, SE, HS, s 95 | [3, 1, 16, True, False, 1], 96 | [3, 4.5, 24, False, False, (2,1)], 97 | [3, 3.67, 24, False, False, 1], 98 | [5, 4, 40, True, True, (2,1)], 99 | [5, 6, 40, True, True, 1], 100 | [5, 6, 40, True, True, 1], 101 | [5, 3, 48, True, True, 1], 102 | [5, 3, 48, True, True, 1], 103 | [5, 6, 96, True, True, (2,1)], 104 | [5, 6, 96, True, True, 1], 105 | [5, 6, 96, True, True, 1], 106 | ] 107 | def initFc(self): 108 | self.pool = nn.MaxPool2d(2) 109 | 110 | def forward(self,x): 111 | x = self.features(x) 112 | x = self.conv(x) 113 | x = self.pool(x) 114 | 115 | b, c, h, w = x.size() 116 | assert h == 1, "the height of conv must be 1" 117 | 118 | x = x.squeeze(2) # b *c * width 119 | x = x.permute(2, 0, 1) # [w, b, c] 120 | return x 121 | -------------------------------------------------------------------------------- /deepvac/aug/composer.py: -------------------------------------------------------------------------------- 1 | import random 2 | from collections import OrderedDict 3 | from ..core import AttrDict 4 | from .factory import AugFactory 5 | 6 | class Composer(object): 7 | def __init__(self,deepvac_config): 8 | self._graph = OrderedDict() 9 | self._graph_p = OrderedDict() 10 | self.deepvac_aug_config = deepvac_config.aug 11 | self.initConfig() 12 | self.auditConfig() 13 | 14 | def initConfig(self): 15 | if self.name() not in self.deepvac_aug_config.keys(): 16 | self.deepvac_aug_config[self.name()] = AttrDict() 17 | self.config = self.deepvac_aug_config[self.name()] 18 | 19 | def name(self): 20 | return self.__class__.__name__ 21 | 22 | def setAttr(self, k, v): 23 | self.config[k] = v 24 | 25 | def getAttr(self,k): 26 | return self.config[k] 27 | 28 | def addAugFactory(self, name, chain, p=1.0): 29 | self._graph[name] = chain 30 | self._graph_p[name] = p 31 | 32 | def auditConfig(self): 33 | pass 34 | 35 | def addOp(self, name, op, p=1.0): 36 | self._graph[name].addOp(op,p) 37 | 38 | def remove(self, name): 39 | try: 40 | del self._graph[name] 41 | except KeyError: 42 | pass 43 | 44 | def __call__(self, img): 45 | for k in self._graph: 46 | if random.random() < self._graph_p[k]: 47 | img = self._graph[k](img) 48 | return img 49 | 50 | class DiceComposer(Composer): 51 | def __call__(self, img): 52 | i = random.randrange(len(self._graph_p)) 53 | return list(self._graph.values())[i](img) 54 | 55 | class PickOneComposer(Composer): 56 | def __call__(self, img): 57 | for k in self._graph: 58 | if random.random() >= self._graph_p[k]: 59 | continue 60 | return self._graph[k](img) 61 | return self._graph[k](img) 62 | 63 | class MultiInputCompose(object): 64 | def __init__(self, transforms): 65 | self.transforms = transforms 66 | 67 | def __call__(self, img): 68 | for t in self.transforms: 69 | img = t(img) 70 | return img 71 | 72 | class FaceAugComposer(Composer): 73 | def __init__(self, deepvac_config): 74 | super(FaceAugComposer, self).__init__(deepvac_config) 75 | ac1 = AugFactory('RandomColorJitterAug@0.5 => MosaicAug@0.5',deepvac_config) 76 | ac2 = AugFactory('IPCFaceAug || MotionAug',deepvac_config) 77 | self.addAugFactory('ac1', ac1, 1) 78 | self.addAugFactory('ac2', ac2, 0.5) 79 | 80 | class OcrAugComposer(Composer): 81 | def __init__(self, deepvac_config): 82 | super(OcrAugComposer, self).__init__(deepvac_config) 83 | 84 | ac1 = AugFactory('SpeckleAug@0.2 => HorlineAug@0.2 => NoisyAug@0.2',deepvac_config) 85 | ac2 = AugFactory('MotionAug || AffineAug || PerspectAug || GaussianAug || VerlineAug || LRmotionAug || UDmotionAug || PerspectiveAug || StretchAug',deepvac_config) 86 | 87 | self.addAugFactory('ac1', ac1, 0.2) 88 | self.addAugFactory('ac2', ac2, 0.9) 89 | 90 | class YoloAugComposer(Composer): 91 | def __init__(self, deepvac_config): 92 | super(YoloAugComposer, self).__init__(deepvac_config) 93 | if deepvac_config.hflip is None: 94 | deepvac_config.hflip = 0.5 95 | ac = AugFactory("YoloPerspectiveAug => HSVAug => YoloNormalizeAug => YoloHFlipAug@{}".format(deepvac_config.hflip), deepvac_config) 96 | self.addAugFactory("ac", ac) 97 | 98 | class RetinaAugComposer(PickOneComposer): 99 | def __init__(self, deepvac_config): 100 | super(RetinaAugComposer, self).__init__(deepvac_config) 101 | ac1 = AugFactory('CropFacialWithBoxesAndLmksAug => BrightDistortFacialAug@0.5 => ContrastDistortFacialAug@0.5 => SaturationDistortFacialAug@0.5 \ 102 | => HueDistortFacialAug@0.5 => Pad2SquareFacialAug => MirrorFacialAug@0.5 => ResizeSubtractMeanFacialAug', deepvac_config) 103 | ac2 = AugFactory('CropFacialWithBoxesAndLmksAug => BrightDistortFacialAug@0.5 => SaturationDistortFacialAug@0.5 => HueDistortFacialAug@0.5 \ 104 | => ContrastDistortFacialAug@0.5 => Pad2SquareFacialAug => MirrorFacialAug@0.5 => ResizeSubtractMeanFacialAug', deepvac_config) 105 | self.addAugFactory("ac1", ac1, 0.5) 106 | self.addAugFactory("ac2", ac2, 0.5) 107 | 108 | class OcrDetectAugComposer(Composer): 109 | def __init__(self, deepvac_config): 110 | super(OcrDetectAugComposer, self).__init__(deepvac_config) 111 | ac1 = AugFactory('ImageWithMasksRandomHorizontalFlipAug@0.5 => ImageWithMasksRandomRotateAug => ImageWithMasksRandomCropAug',deepvac_config) 112 | self.addAugFactory('ac1', ac1, 1) 113 | 114 | class SegImageAugComposer(Composer): 115 | def __init__(self, deepvac_config): 116 | super(SegImageAugComposer, self).__init__(deepvac_config) 117 | ac1 = AugFactory('GaussianAug || RandomColorJitterAug || BrightnessJitterAug || ContrastJitterAug',deepvac_config) 118 | ac2 = AugFactory('MotionAug',deepvac_config) 119 | self.addAugFactory('ac1', ac1, 1) 120 | self.addAugFactory('ac2', ac1, 0.2) 121 | 122 | class SegImageWithMaskAugComposer(Composer): 123 | def __init__(self, deepvac_config): 124 | super(SegImageWithMaskAugComposer, self).__init__(deepvac_config) 125 | ac1 = AugFactory('ImageWithMasksRandomCropAug',deepvac_config) 126 | ac2 = AugFactory('ImageWithMasksRandomHorizontalFlipAug || ImageWithMasksRandomRotateAug',deepvac_config) 127 | self.addAugFactory('ac1', ac1, 0.4) 128 | self.addAugFactory('ac2', ac2, 0.6) 129 | -------------------------------------------------------------------------------- /deepvac/backbones/conv_layer.py: -------------------------------------------------------------------------------- 1 | from collections import OrderedDict 2 | import torch.nn as nn 3 | from .activate_layer import hswish 4 | 5 | class Conv2dBN(nn.Sequential): 6 | def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, padding=None, groups=1, bias=False): 7 | if padding is None: 8 | padding = (kernel_size - 1) // 2 9 | super(Conv2dBN, self).__init__( 10 | nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=bias), 11 | nn.BatchNorm2d(out_planes) 12 | ) 13 | 14 | class Conv2dBNWithName(nn.Sequential): 15 | def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, padding=None, groups=1, bias=False): 16 | if padding is None: 17 | padding = (kernel_size - 1) // 2 18 | super(Conv2dBNWithName, self).__init__(OrderedDict([ 19 | ('conv', nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=bias)), 20 | ('bn', nn.BatchNorm2d(out_planes)) 21 | ])) 22 | 23 | class Conv2dBNReLU(nn.Sequential): 24 | def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, padding=None, groups=1, bias=False): 25 | if padding is None: 26 | padding = (kernel_size - 1) // 2 27 | super(Conv2dBNReLU, self).__init__( 28 | nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=bias), 29 | nn.BatchNorm2d(out_planes, momentum=0.1), 30 | nn.ReLU(inplace=True) 31 | ) 32 | 33 | class Conv2dBNPReLU(nn.Sequential): 34 | def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, padding=None, groups=1, bias=False): 35 | if padding is None: 36 | padding = (kernel_size - 1) // 2 37 | super(Conv2dBNPReLU, self).__init__( 38 | nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=bias), 39 | nn.BatchNorm2d(out_planes, momentum=0.1), 40 | nn.PReLU(out_planes) 41 | ) 42 | 43 | class Conv2dBnAct(nn.Sequential): 44 | def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, padding=None, groups=1, bias=False, act=True): 45 | if padding is None: 46 | padding = (kernel_size - 1) // 2 47 | super(Conv2dBnAct, self).__init__( 48 | nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=bias), 49 | nn.BatchNorm2d(out_planes), 50 | nn.SiLU() if (act is True) else (act if isinstance(act, nn.Module) else nn.Identity()) 51 | ) 52 | 53 | class Conv2dBNHardswish(nn.Sequential): 54 | def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, padding=None, groups=1, bias=False): 55 | if padding is None: 56 | padding = (kernel_size - 1) // 2 57 | super(Conv2dBNHardswish, self).__init__( 58 | nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=bias), 59 | nn.BatchNorm2d(out_planes, momentum=0.1), 60 | nn.Hardswish() 61 | ) 62 | 63 | class Conv2dBNHswish(nn.Sequential): 64 | def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, padding=None, groups=1, bias=False): 65 | if padding is None: 66 | padding = (kernel_size - 1) // 2 67 | super(Conv2dBNHswish, self).__init__( 68 | nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=bias), 69 | nn.BatchNorm2d(out_planes, momentum=0.1), 70 | hswish(inplace=True) 71 | ) 72 | 73 | class DepthWiseConv2d(nn.Module): 74 | def __init__(self, inplanes: int, outplanes: int, kernel_size: int, stride: int, padding: int, groups: int, residual: bool=False, bias: bool=False): 75 | super(DepthWiseConv2d, self).__init__() 76 | self.conv1 = Conv2dBNPReLU(inplanes, groups, kernel_size=1, stride=1, padding=0) 77 | self.conv2 = Conv2dBNPReLU(groups, groups, kernel_size, stride, padding, groups) 78 | self.conv3 = nn.Conv2d(groups, outplanes, kernel_size=1, stride=1, padding=0, bias=bias) 79 | self.bn = nn.BatchNorm2d(outplanes) 80 | self.residual = residual 81 | 82 | def forward(self, x): 83 | identity = x 84 | x = self.conv1(x) 85 | x = self.conv2(x) 86 | x = self.conv3(x) 87 | x = self.bn(x) 88 | if self.residual: 89 | x = identity + x 90 | return x 91 | 92 | class Conv2dBNLeakyReLU(nn.Sequential): 93 | def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, padding=None, groups=1, bias=False, leaky=0): 94 | if padding is None: 95 | padding = (kernel_size - 1) // 2 96 | #leaky = 0.1 if out_planes <= 64 else 0 97 | super(Conv2dBNLeakyReLU, self).__init__( 98 | nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=bias), 99 | nn.BatchNorm2d(out_planes), 100 | nn.LeakyReLU(negative_slope=leaky, inplace=True) 101 | ) 102 | 103 | class BNPReLU(nn.Sequential): 104 | def __init__(self, out_planes): 105 | super(BNPReLU, self).__init__( 106 | nn.BatchNorm2d(out_planes), 107 | nn.PReLU(out_planes) 108 | ) 109 | 110 | class Conv2dDilatedBN(nn.Sequential): 111 | def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, dilation=1, padding=None, groups=1, bias=False): 112 | padding = ((kernel_size - 1) // 2) * dilation 113 | super(Conv2dDilatedBN, self).__init__( 114 | nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, dilation=dilation, groups=groups, bias=bias), 115 | nn.BatchNorm2d(out_planes) 116 | ) 117 | 118 | -------------------------------------------------------------------------------- /deepvac/aug/factory.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import re 3 | import random 4 | import cv2 5 | from torchvision import transforms as trans 6 | from ..utils import LOG 7 | from ..core import AttrDict 8 | from . import base_aug, face_aug, seg_aug, text_aug, yolo_aug, seg_audit 9 | 10 | class SyszuxFactory(object): 11 | def __init__(self, syntax, deepvac_config): 12 | self.deepvac_config = deepvac_config 13 | self.deepvac_aug_config = deepvac_config.aug 14 | self.factory_dict = dict() 15 | self.initConfig() 16 | self.auditConfig() 17 | self.initSyntax(syntax) 18 | self.initProducts() 19 | 20 | def initConfig(self): 21 | if self.name() not in self.deepvac_aug_config.keys(): 22 | self.deepvac_aug_config[self.name()] = AttrDict() 23 | self.config = self.deepvac_aug_config[self.name()] 24 | 25 | def name(self): 26 | return self.__class__.__name__ 27 | 28 | def setAttr(self, k, v): 29 | self.config[k] = v 30 | 31 | def getAttr(self,k): 32 | return self.config[k] 33 | 34 | def auditConfig(self): 35 | pass 36 | 37 | def initSyntax(self, syntax): 38 | self.func_dict = dict() 39 | self.initChainKind() 40 | self.auditChainKind(syntax) 41 | flow = syntax.split(self.chain_kind) 42 | self.op_sym_list = [x.strip() for x in flow] 43 | self.p_list = [ float(x.split('@')[1].strip()) if '@' in x else 1 for x in self.op_sym_list ] 44 | self.op_sym_list = [x.split('@')[0].strip() for x in self.op_sym_list if x] 45 | self.op_list = [] 46 | 47 | def __call__(self, img): 48 | return self.func_dict[self.chain_kind](img) 49 | 50 | def addOp(self, op, p=1): 51 | self.op_list.append(op) 52 | self.p_list.append(p) 53 | 54 | def initChainKind(self): 55 | self.func_dict['=>'] = self.process 56 | 57 | def process(self, img): 58 | for t in self.op_list: 59 | img = t(img) 60 | return img 61 | 62 | def auditChainKind(self,flow_list): 63 | self.chain_kind = '=>' 64 | tokens = re.sub('[a-zA-Z0-9.@]'," ", flow_list).split() 65 | 66 | if len(tokens) == 0: 67 | return 68 | 69 | tokens = set(tokens) 70 | 71 | if len(tokens) > 1: 72 | raise Exception('Multi token found in flow list: ', tokens) 73 | 74 | self.chain_kind = tokens.pop() 75 | 76 | if self.chain_kind not in self.func_dict.keys(): 77 | raise Exception("token not supported: ", self.chain_kind) 78 | 79 | def initProducts(self): 80 | LOG.logE("You must reimplement initProducts() function in sublcass {}.".format(self.name()), exit=True) 81 | 82 | def addProduct(self, name, ins): 83 | LOG.logE("You must reimplement addProduct() function in subclass {}.".format(self.name()), exit=True) 84 | 85 | def initSym(self, module): 86 | return dir(module) 87 | 88 | def omitUnderScoreSym(self, sym): 89 | return [x for x in sym if not x.startswith('_')] 90 | 91 | def selectStartPrefix(self, sym, start_prefix): 92 | return [x for x in sym if x.startswith(start_prefix)] 93 | 94 | def selectEndSuffix(self, sym, end_suffix): 95 | return [x for x in sym if x.endswith(end_suffix)] 96 | 97 | def omitOmitList(self, sym, omit_list): 98 | return [x for x in sym if x not in omit_list] 99 | 100 | def getSymsFromProduct(self, module, start_prefix='', end_suffix='', omit_list=[]): 101 | sym = self.initSym(module) 102 | sym = self.omitUnderScoreSym(sym) 103 | sym = self.selectStartPrefix(sym, start_prefix) 104 | sym = self.selectEndSuffix(sym, end_suffix) 105 | return self.omitOmitList(sym, omit_list) 106 | 107 | def addProduct(self, sym, ins): 108 | self.factory_dict[sym] = ins 109 | 110 | def addProducts(self, module, module_name, start_prefix='', end_suffix='', sym_prefix='', omit_list=[]): 111 | for sym in self.getSymsFromProduct(module, start_prefix='', end_suffix='', omit_list=[]): 112 | self.addProduct(sym_prefix + sym, eval('{}.{}'.format(module_name, sym)) ) 113 | 114 | def get(self, ins_name): 115 | if ins_name in self.factory_dict: 116 | return self.factory_dict[ins_name](self.deepvac_config) 117 | 118 | if ins_name in self.config: 119 | return self.config[ins_name] 120 | 121 | LOG.logE("ERROR! {} not found in factory {}.".format(ins_name, self.name()), exit=True) 122 | 123 | 124 | class AugFactory(SyszuxFactory): 125 | def __init__(self, syntax, deepvac_config): 126 | super(AugFactory, self).__init__(syntax, deepvac_config) 127 | self.op_list = [self.get(x) for x in self.op_sym_list] 128 | 129 | def processRandom(self, img): 130 | for t,p in zip(self.op_list, self.p_list): 131 | if random.random() < p: 132 | img = t(img) 133 | return img 134 | 135 | def processDice(self, img): 136 | i = random.randrange(len(self.op_list)) 137 | return self.op_list[i](img) 138 | 139 | def initChainKind(self): 140 | self.func_dict['=>'] = self.processRandom 141 | self.func_dict['||'] = self.processDice 142 | 143 | def initProducts(self): 144 | self.addProducts(base_aug, module_name='base_aug', end_suffix='Aug') 145 | self.addProducts(face_aug, module_name='face_aug', end_suffix='Aug') 146 | self.addProducts(seg_aug, module_name='seg_aug', end_suffix='Aug') 147 | self.addProducts(text_aug, module_name='text_aug', end_suffix='Aug') 148 | self.addProducts(yolo_aug, module_name='yolo_aug', end_suffix='Aug') 149 | 150 | 151 | class ImageWithMaskAuditFactory(AugFactory): 152 | def initProducts(self): 153 | super(ImageWithMaskAuditFactory, self).initProducts() 154 | self.addProducts(seg_audit, module_name='seg_audit', end_suffix='Audit') -------------------------------------------------------------------------------- /deepvac/aug/perspective_helper.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | 3 | # http://planning.cs.uiuc.edu/node102.html 4 | def get_rotate_matrix(x, y, z): 5 | """ 6 | 按照 zyx 的顺序旋转，输入角度单位为 degrees, 均为顺时针旋转 7 | :param x: X-axis 8 | :param y: Y-axis 9 | :param z: Z-axis 10 | :return: 11 | """ 12 | x = math.radians(x) 13 | y = math.radians(y) 14 | z = math.radians(z) 15 | 16 | c, s = math.cos(y), math.sin(y) 17 | M_y = np.matrix([[c, 0., s, 0.], 18 | [0., 1., 0., 0.], 19 | [-s, 0., c, 0.], 20 | [0., 0., 0., 1.]]) 21 | 22 | c, s = math.cos(x), math.sin(x) 23 | M_x = np.matrix([[1., 0., 0., 0.], 24 | [0., c, -s, 0.], 25 | [0., s, c, 0.], 26 | [0., 0., 0., 1.]]) 27 | 28 | c, s = math.cos(z), math.sin(z) 29 | M_z = np.matrix([[c, -s, 0., 0.], 30 | [s, c, 0., 0.], 31 | [0., 0., 1., 0.], 32 | [0., 0., 0., 1.]]) 33 | 34 | return M_x * M_y * M_z 35 | 36 | 37 | def cliped_rand_norm(mu=0, sigma3=1): 38 | """ 39 | :param mu: 均值 40 | :param sigma3: 3 倍标准差， 99% 的数据落在 (mu-3*sigma, mu+3*sigma) 41 | :return: 42 | """ 43 | # 标准差 44 | sigma = sigma3 / 3 45 | dst = sigma * np.random.randn() + mu 46 | dst = np.clip(dst, 0 - sigma3, sigma3) 47 | return dst 48 | 49 | 50 | def warpPerspective(src, M33, sl, gpu): 51 | if gpu: 52 | from libs.gpu.GpuWrapper import cudaWarpPerspectiveWrapper 53 | dst = cudaWarpPerspectiveWrapper(src.astype(np.uint8), M33, (sl, sl), cv2.INTER_CUBIC) 54 | else: 55 | dst = cv2.warpPerspective(src, M33, (sl, sl), flags=cv2.INTER_CUBIC) 56 | return dst 57 | 58 | 59 | # https://stackoverflow.com/questions/17087446/how-to-calculate-perspective-transform-for-opencv-from-rotation-angles 60 | # https://nbviewer.jupyter.org/github/manisoftwartist/perspectiveproj/blob/master/perspective.ipynb 61 | # http://planning.cs.uiuc.edu/node102.html 62 | class PerspectiveTransform(object): 63 | def __init__(self, x, y, z, scale, fovy): 64 | self.x = x 65 | self.y = y 66 | self.z = z 67 | self.scale = scale 68 | self.fovy = fovy 69 | 70 | def transform_image(self, src, gpu=False): 71 | if len(src.shape) > 2: 72 | H, W, C = src.shape 73 | else: 74 | H, W = src.shape 75 | 76 | M33, sl, _, ptsOut = self.get_warp_matrix(W, H, self.x, self.y, self.z, self.scale, self.fovy) 77 | sl = int(sl) 78 | 79 | dst = warpPerspective(src, M33, sl, gpu) 80 | 81 | return dst, M33, ptsOut 82 | 83 | def transform_pnts(self, pnts, M33): 84 | """ 85 | :param pnts: 2D pnts, left-top, right-top, right-bottom, left-bottom 86 | :param M33: output from transform_image() 87 | :return: 2D pnts apply perspective transform 88 | """ 89 | pnts = np.asarray(pnts, dtype=np.float32) 90 | pnts = np.array([pnts]) 91 | dst_pnts = cv2.perspectiveTransform(pnts, M33)[0] 92 | 93 | return dst_pnts 94 | 95 | def get_warped_pnts(self, ptsIn, ptsOut, W, H, sidelength): 96 | ptsIn2D = ptsIn[0, :] 97 | ptsOut2D = ptsOut[0, :] 98 | ptsOut2Dlist = [] 99 | ptsIn2Dlist = [] 100 | 101 | for i in range(0, 4): 102 | ptsOut2Dlist.append([ptsOut2D[i, 0], ptsOut2D[i, 1]]) 103 | ptsIn2Dlist.append([ptsIn2D[i, 0], ptsIn2D[i, 1]]) 104 | 105 | pin = np.array(ptsIn2Dlist) + [W / 2., H / 2.] 106 | pout = (np.array(ptsOut2Dlist) + [1., 1.]) * (0.5 * sidelength) 107 | pin = pin.astype(np.float32) 108 | pout = pout.astype(np.float32) 109 | 110 | return pin, pout 111 | 112 | def get_warp_matrix(self, W, H, x, y, z, scale, fV): 113 | fVhalf = np.deg2rad(fV / 2.) 114 | d = np.sqrt(W * W + H * H) 115 | sideLength = scale * d / np.cos(fVhalf) 116 | h = d / (2.0 * np.sin(fVhalf)) 117 | n = h - (d / 2.0) 118 | f = h + (d / 2.0) 119 | 120 | # Translation along Z-axis by -h 121 | T = np.eye(4, 4) 122 | T[2, 3] = -h 123 | 124 | # Rotation matrices around x,y,z 125 | R = get_rotate_matrix(x, y, z) 126 | 127 | # Projection Matrix 128 | P = np.eye(4, 4) 129 | P[0, 0] = 1.0 / np.tan(fVhalf) 130 | P[1, 1] = P[0, 0] 131 | P[2, 2] = -(f + n) / (f - n) 132 | P[2, 3] = -(2.0 * f * n) / (f - n) 133 | P[3, 2] = -1.0 134 | 135 | # pythonic matrix multiplication 136 | M44 = reduce(lambda x, y: np.matmul(x, y), [P, T, R]) 137 | 138 | # shape should be 1,4,3 for ptsIn and ptsOut since perspectiveTransform() expects data in this way. 139 | # In C++, this can be achieved by Mat ptsIn(1,4,CV_64FC3); 140 | ptsIn = np.array([[ 141 | [-W / 2., H / 2., 0.], 142 | [W / 2., H / 2., 0.], 143 | [W / 2., -H / 2., 0.], 144 | [-W / 2., -H / 2., 0.] 145 | ]]) 146 | ptsOut = cv2.perspectiveTransform(ptsIn, M44) 147 | 148 | ptsInPt2f, ptsOutPt2f = self.get_warped_pnts(ptsIn, ptsOut, W, H, sideLength) 149 | 150 | # check float32 otherwise OpenCV throws an error 151 | assert (ptsInPt2f.dtype == np.float32) 152 | assert (ptsOutPt2f.dtype == np.float32) 153 | M33 = cv2.getPerspectiveTransform(ptsInPt2f, ptsOutPt2f).astype(np.float32) 154 | 155 | return M33, sideLength, ptsInPt2f, ptsOutPt2f 156 | 157 | def apply_perspective_transform(img, max_x, max_y, max_z): 158 | x = cliped_rand_norm(0, max_x) 159 | y = cliped_rand_norm(0, max_y) 160 | z = cliped_rand_norm(0, max_z) 161 | 162 | transformer = PerspectiveTransform(x, y, z, scale=1.0, fovy=50) 163 | dst_img, M33, dst_img_pnts = transformer.transform_image(img, False) 164 | dst_img_pnts = np.array(dst_img_pnts,dtype = np.int32) 165 | 166 | x_min, y_min = np.min(dst_img_pnts,axis=0) 167 | x_max, y_max = np.max(dst_img_pnts,axis=0) 168 | return dst_img[y_min:y_max, x_min:x_max] -------------------------------------------------------------------------------- /deepvac/aug/seg_audit.py: -------------------------------------------------------------------------------- 1 | import os 2 | import cv2 3 | import numpy as np 4 | from ..utils import LOG, addUserConfig 5 | from .base_aug import CvAugBase 6 | 7 | class CvAugSegAuditBase4(CvAugBase): 8 | def __init__(self, deepvac_config): 9 | super(CvAugSegAuditBase4, self).__init__(deepvac_config) 10 | self.input_len = self.addUserConfig('input_len', self.config.input_len, 4) 11 | self.cls_num = self.addUserConfig('cls_num', self.config.cls_num, 4) 12 | self.pallete = [[255, 255, 255], 13 | [255, 0, 0], 14 | [0, 255, 0], 15 | [0, 0, 255], 16 | [102, 102, 156], 17 | [190, 153, 153], 18 | [153, 153, 153], 19 | [250, 170, 30], 20 | [220, 220, 0], 21 | [107, 142, 35], 22 | [152, 251, 152], 23 | [70, 130, 180], 24 | [220, 20, 60], 25 | [255, 0, 0], 26 | [0, 0, 142], 27 | [0, 0, 70], 28 | [0, 60, 100], 29 | [0, 80, 100], 30 | [0, 0, 230], 31 | [119, 11, 32]] 32 | 33 | def putMask(self, img, mask): 34 | h,w = img.shape[:2] 35 | classMap_numpy_color = np.zeros((h, w, 3), dtype=np.uint8) 36 | for idx in range(self.cls_num): 37 | [r, g, b] = self.pallete[idx] 38 | classMap_numpy_color[mask == idx] = [b, g, r] 39 | overlayed = cv2.addWeighted(img, 0.5, classMap_numpy_color, 0.5, 0) 40 | return overlayed 41 | 42 | def getBaseName(self, fn): 43 | return os.path.basename(fn) 44 | 45 | def write(self, img, label, fn): 46 | overlayed = self.putMask(img, label) 47 | basename = self.getBaseName(fn) 48 | cv2.imwrite(os.path.join(self.save_dir, os.path.splitext(basename)[0]+'.png'), overlayed) 49 | 50 | class ImageWithMaskIntersectAudit(CvAugSegAuditBase4): 51 | def auditConfig(self): 52 | self.intersect_ratio = self.addUserConfig('intersect_ratio', self.config.intersect_ratio, 0.10) 53 | self.remask = self.addUserConfig('remask', self.config.remask, True) 54 | self.save_dir = self.addUserConfig('intersect_dir', self.config.intersect_dir, "intersect_dir") 55 | self.is_consider_bg = self.addUserConfig('is_consider_bg', self.config.is_consider_bg, False) 56 | os.makedirs(self.save_dir, exist_ok=True) 57 | 58 | def remaskIfNeeded(self, i, j): 59 | mask_i = self.cls_masks[self.catid[i]] 60 | mask_j = self.cls_masks[self.catid[j]] 61 | mask_i_num = np.sum(mask_i==1) 62 | mask_j_num = np.sum(mask_j==1) 63 | min_area = min(mask_i_num, mask_j_num) 64 | 65 | intersect_area = mask_i*mask_j==1 66 | intersect_ratio = np.sum(intersect_area) / min_area 67 | 68 | if intersect_ratio < self.intersect_ratio: 69 | return 70 | 71 | is_write = self.catid[i] * self.catid[j] != 0 72 | if is_write or self.is_consider_bg: 73 | LOG.logI("Image {} has risk on rule {} with value {} ({} / {})".format(self.fn, self.name(), intersect_ratio, self.catid[i], self.catid[j])) 74 | 75 | if self.remask: 76 | min_idx, max_idx = (i,j) if mask_i_num < mask_j_num else (j,i) 77 | LOG.logI("You have enabled remask, do remask...") 78 | self.label[intersect_area] = self.catid[min_idx] 79 | self.cls_masks[self.catid[max_idx]][intersect_area] = self.catid[min_idx] 80 | 81 | #if bg 82 | if is_write or self.is_consider_bg: 83 | self.write(self.img, self.label, self.fn) 84 | 85 | def forward(self, imgs): 86 | self.img, self.label, self.cls_masks, self.fn = imgs 87 | len_mask = len(self.cls_masks) 88 | if len_mask < 2: 89 | return imgs 90 | self.catid = list(self.cls_masks.keys()) 91 | for i in range(len_mask - 1): 92 | for j in range(i+1, len_mask): 93 | self.remaskIfNeeded(i, j) 94 | 95 | return self.img, self.label, self.cls_masks, self.fn 96 | 97 | class ImageWithMaskSideRatioAudit(CvAugSegAuditBase4): 98 | def auditConfig(self): 99 | self.save_dir = self.addUserConfig('side_ratio_dir', self.config.side_ratio_dir, "side_ratio_dir") 100 | os.makedirs(self.save_dir, exist_ok=True) 101 | 102 | def forward(self, imgs): 103 | img, label, _, fn = imgs 104 | if img.shape[0] == img.shape[1]: 105 | LOG.logI("Image {} has risk on rule {} with size {}".format(fn, self.name(), img.shape)) 106 | self.write(img, label, fn) 107 | return imgs 108 | 109 | class ImageWithMaskTargetSizeAudit(CvAugSegAuditBase4): 110 | def auditConfig(self): 111 | self.min_ratio = self.addUserConfig('min_ratio', self.config.min_ratio, 900) 112 | self.save_dir = self.addUserConfig('target_size_dir', self.config.target_size_dir, "target_size_dir") 113 | os.makedirs(self.save_dir, exist_ok=True) 114 | 115 | def forward(self, imgs): 116 | img, label, cls_masks, fn = imgs 117 | min_pixel_num = int(label.shape[0] * label.shape[1] / self.min_ratio) 118 | for cls_idx in cls_masks.keys(): 119 | cls_pixel_num = np.sum(label==cls_idx) 120 | if cls_pixel_num == 0: 121 | continue 122 | if cls_pixel_num >= min_pixel_num: 123 | continue 124 | LOG.logI("Image {} has risk on rule {} with {}_num = {} (min_pixel_num = {})".format(fn, self.name(), cls_idx, cls_pixel_num, min_pixel_num)) 125 | self.write(img, label, fn) 126 | return imgs 127 | 128 | class ImageWithMaskVisionAudit(CvAugSegAuditBase4): 129 | def auditConfig(self): 130 | self.save_dir = self.addUserConfig('vision_dir', self.config.vision_dir, "vision_dir") 131 | self.mask_dir = self.addUserConfig('mask_dir', self.config.mask_dir, "mask_dir") 132 | os.makedirs(self.save_dir, exist_ok=True) 133 | os.makedirs(self.mask_dir, exist_ok=True) 134 | 135 | def forward(self, imgs): 136 | img, label, _, fn = imgs 137 | self.write(img, label, fn) 138 | 139 | sub_path = self.getBaseName(os.path.dirname(fn)) 140 | new_label_path = os.path.join(self.mask_dir, sub_path) 141 | os.makedirs(new_label_path, exist_ok=True) 142 | basename = self.getBaseName(fn) 143 | cv2.imwrite(os.path.join(new_label_path, os.path.splitext(basename)[0]+'.png'), label) 144 | return imgs 145 | 146 | 147 | -------------------------------------------------------------------------------- /deepvac/aug/warp_mls_helper.py: -------------------------------------------------------------------------------- 1 | # -*- coding:utf-8 -*- 2 | # Author: RubanSeven 3 | import numpy as np 4 | 5 | class WarpMLS: 6 | def __init__(self, src, src_pts, dst_pts, dst_w, dst_h, trans_ratio=1.): 7 | self.src = src 8 | self.src_pts = src_pts 9 | self.dst_pts = dst_pts 10 | self.pt_count = len(self.dst_pts) 11 | self.dst_w = dst_w 12 | self.dst_h = dst_h 13 | self.trans_ratio = trans_ratio 14 | self.grid_size = 100 15 | self.rdx = np.zeros((self.dst_h, self.dst_w)) 16 | self.rdy = np.zeros((self.dst_h, self.dst_w)) 17 | 18 | @staticmethod 19 | def __bilinear_interp(x, y, v11, v12, v21, v22): 20 | return (v11 * (1 - y) + v12 * y) * (1 - x) + (v21 * (1 - y) + v22 * y) * x 21 | 22 | def generate(self): 23 | self.calc_delta() 24 | return self.gen_img() 25 | 26 | def calc_delta(self): 27 | w = np.zeros(self.pt_count, dtype=np.float32) 28 | 29 | if self.pt_count < 2: 30 | return 31 | 32 | i = 0 33 | while 1: 34 | if self.dst_w <= i < self.dst_w + self.grid_size - 1: 35 | i = self.dst_w - 1 36 | elif i >= self.dst_w: 37 | break 38 | 39 | j = 0 40 | while 1: 41 | if self.dst_h <= j < self.dst_h + self.grid_size - 1: 42 | j = self.dst_h - 1 43 | elif j >= self.dst_h: 44 | break 45 | 46 | sw = 0 47 | swp = np.zeros(2, dtype=np.float32) 48 | swq = np.zeros(2, dtype=np.float32) 49 | new_pt = np.zeros(2, dtype=np.float32) 50 | cur_pt = np.array([i, j], dtype=np.float32) 51 | 52 | k = 0 53 | for k in range(self.pt_count): 54 | if i == self.dst_pts[k][0] and j == self.dst_pts[k][1]: 55 | break 56 | 57 | w[k] = 1. / ((i - self.dst_pts[k][0]) * (i - self.dst_pts[k][0]) + 58 | (j - self.dst_pts[k][1]) * (j - self.dst_pts[k][1])) 59 | 60 | sw += w[k] 61 | swp = swp + w[k] * np.array(self.dst_pts[k]) 62 | swq = swq + w[k] * np.array(self.src_pts[k]) 63 | 64 | if k == self.pt_count - 1: 65 | pstar = 1 / sw * swp 66 | qstar = 1 / sw * swq 67 | 68 | miu_s = 0 69 | for k in range(self.pt_count): 70 | if i == self.dst_pts[k][0] and j == self.dst_pts[k][1]: 71 | continue 72 | pt_i = self.dst_pts[k] - pstar 73 | miu_s += w[k] * np.sum(pt_i * pt_i) 74 | 75 | cur_pt -= pstar 76 | cur_pt_j = np.array([-cur_pt[1], cur_pt[0]]) 77 | 78 | for k in range(self.pt_count): 79 | if i == self.dst_pts[k][0] and j == self.dst_pts[k][1]: 80 | continue 81 | 82 | pt_i = self.dst_pts[k] - pstar 83 | pt_j = np.array([-pt_i[1], pt_i[0]]) 84 | 85 | tmp_pt = np.zeros(2, dtype=np.float32) 86 | tmp_pt[0] = np.sum(pt_i * cur_pt) * self.src_pts[k][0] - \ 87 | np.sum(pt_j * cur_pt) * self.src_pts[k][1] 88 | tmp_pt[1] = -np.sum(pt_i * cur_pt_j) * self.src_pts[k][0] + \ 89 | np.sum(pt_j * cur_pt_j) * self.src_pts[k][1] 90 | tmp_pt *= (w[k] / miu_s) 91 | new_pt += tmp_pt 92 | 93 | new_pt += qstar 94 | else: 95 | new_pt = self.src_pts[k] 96 | 97 | self.rdx[j, i] = new_pt[0] - i 98 | self.rdy[j, i] = new_pt[1] - j 99 | 100 | j += self.grid_size 101 | i += self.grid_size 102 | 103 | def gen_img(self): 104 | src_h, src_w = self.src.shape[:2] 105 | dst = np.zeros_like(self.src, dtype=np.float32) 106 | 107 | for i in np.arange(0, self.dst_h, self.grid_size): 108 | for j in np.arange(0, self.dst_w, self.grid_size): 109 | ni = i + self.grid_size 110 | nj = j + self.grid_size 111 | w = h = self.grid_size 112 | if ni >= self.dst_h: 113 | ni = self.dst_h - 1 114 | h = ni - i + 1 115 | if nj >= self.dst_w: 116 | nj = self.dst_w - 1 117 | w = nj - j + 1 118 | 119 | di = np.reshape(np.arange(h), (-1, 1)) 120 | dj = np.reshape(np.arange(w), (1, -1)) 121 | delta_x = self.__bilinear_interp(di / h, dj / w, 122 | self.rdx[i, j], self.rdx[i, nj], 123 | self.rdx[ni, j], self.rdx[ni, nj]) 124 | delta_y = self.__bilinear_interp(di / h, dj / w, 125 | self.rdy[i, j], self.rdy[i, nj], 126 | self.rdy[ni, j], self.rdy[ni, nj]) 127 | nx = j + dj + delta_x * self.trans_ratio 128 | ny = i + di + delta_y * self.trans_ratio 129 | nx = np.clip(nx, 0, src_w - 1) 130 | ny = np.clip(ny, 0, src_h - 1) 131 | nxi = np.array(np.floor(nx), dtype=np.int32) 132 | nyi = np.array(np.floor(ny), dtype=np.int32) 133 | nxi1 = np.array(np.ceil(nx), dtype=np.int32) 134 | nyi1 = np.array(np.ceil(ny), dtype=np.int32) 135 | 136 | if len(self.src.shape) == 3: 137 | x = np.tile(np.expand_dims(ny - nyi, axis=-1), (1, 1, 3)) 138 | y = np.tile(np.expand_dims(nx - nxi, axis=-1), (1, 1, 3)) 139 | else: 140 | x = ny - nyi 141 | y = nx - nxi 142 | dst[i:i + h, j:j + w] = self.__bilinear_interp(x, 143 | y, 144 | self.src[nyi, nxi], 145 | self.src[nyi, nxi1], 146 | self.src[nyi1, nxi], 147 | self.src[nyi1, nxi1] 148 | ) 149 | 150 | dst = np.clip(dst, 0, 255) 151 | dst = np.array(dst, dtype=np.uint8) 152 | 153 | return dst -------------------------------------------------------------------------------- /deepvac/aug/haishoku_helper.py: -------------------------------------------------------------------------------- 1 | from PIL import Image 2 | 3 | def get_image(image_path): 4 | # if the image_path is a remote url, read the image at first 5 | if isinstance(image_path, str): 6 | image = Image.open(image_path) 7 | else: 8 | image = image_path 9 | 10 | # convert non-RGB mode to RGB mode 11 | if image.mode != "RGB": 12 | image = image.convert("RGB") 13 | return image 14 | 15 | def get_thumbnail(image): 16 | image.thumbnail((256, 256)) 17 | return image 18 | 19 | def get_colors(image_path): 20 | """ image instance 21 | """ 22 | image = get_image(image_path) 23 | 24 | """ image thumbnail 25 | size: 256 * 256 26 | reduce the calculate time 27 | """ 28 | thumbnail = get_thumbnail(image) 29 | 30 | 31 | """ calculate the max colors the image cound have 32 | if the color is different in every pixel, the color counts may be the max. 33 | so : 34 | max_colors = image.height * image.width 35 | """ 36 | image_height = thumbnail.height 37 | image_width = thumbnail.width 38 | max_colors = image_height * image_width 39 | 40 | image_colors = image.getcolors(max_colors) 41 | return image_colors 42 | 43 | 44 | def sort_by_rgb(colors_tuple): 45 | """ colors_tuple contains color count and color RGB 46 | we want to sort the tuple by RGB 47 | tuple[1] 48 | """ 49 | sorted_tuple = sorted(colors_tuple, key=lambda x:x[1]) 50 | return sorted_tuple 51 | 52 | def rgb_maximum(colors_tuple): 53 | """ 54 | colors_r max min 55 | colors_g max min 56 | colors_b max min 57 | 58 | """ 59 | r_sorted_tuple = sorted(colors_tuple, key=lambda x:x[1][0]) 60 | g_sorted_tuple = sorted(colors_tuple, key=lambda x:x[1][1]) 61 | b_sorted_tuple = sorted(colors_tuple, key=lambda x:x[1][2]) 62 | 63 | r_min = r_sorted_tuple[0][1][0] 64 | g_min = g_sorted_tuple[0][1][1] 65 | b_min = b_sorted_tuple[0][1][2] 66 | 67 | r_max = r_sorted_tuple[len(colors_tuple)-1][1][0] 68 | g_max = g_sorted_tuple[len(colors_tuple)-1][1][1] 69 | b_max = b_sorted_tuple[len(colors_tuple)-1][1][2] 70 | 71 | return { 72 | "r_max":r_max, 73 | "r_min":r_min, 74 | "g_max":g_max, 75 | "g_min":g_min, 76 | "b_max":b_max, 77 | "b_min":b_min, 78 | "r_dvalue":(r_max-r_min)/3, 79 | "g_dvalue":(g_max-g_min)/3, 80 | "b_dvalue":(b_max-b_min)/3 81 | } 82 | 83 | def group_by_accuracy(sorted_tuple, accuracy=3): 84 | """ group the colors by the accuaracy was given 85 | the R G B colors will be depart to accuracy parts 86 | default accuracy = 3 87 | d_value = (max-min)/3 88 | [min, min+d_value), [min+d_value, min+d_value*2), [min+d_value*2, max) 89 | """ 90 | rgb_maximum_json = rgb_maximum(sorted_tuple) 91 | r_min = rgb_maximum_json["r_min"] 92 | g_min = rgb_maximum_json["g_min"] 93 | b_min = rgb_maximum_json["b_min"] 94 | r_dvalue = rgb_maximum_json["r_dvalue"] 95 | g_dvalue = rgb_maximum_json["g_dvalue"] 96 | b_dvalue = rgb_maximum_json["b_dvalue"] 97 | 98 | rgb = [ 99 | [[[], [], []], [[], [], []], [[], [], []]], 100 | [[[], [], []], [[], [], []], [[], [], []]], 101 | [[[], [], []], [[], [], []], [[], [], []]] 102 | ] 103 | 104 | for color_tuple in sorted_tuple: 105 | r_tmp_i = color_tuple[1][0] 106 | g_tmp_i = color_tuple[1][1] 107 | b_tmp_i = color_tuple[1][2] 108 | r_idx = 0 if r_tmp_i < (r_min+r_dvalue) else 1 if r_tmp_i < (r_min+r_dvalue*2) else 2 109 | g_idx = 0 if g_tmp_i < (g_min+g_dvalue) else 1 if g_tmp_i < (g_min+g_dvalue*2) else 2 110 | b_idx = 0 if b_tmp_i < (b_min+b_dvalue) else 1 if b_tmp_i < (b_min+b_dvalue*2) else 2 111 | rgb[r_idx][g_idx][b_idx].append(color_tuple) 112 | 113 | return rgb 114 | 115 | 116 | def get_weighted_mean(grouped_image_color): 117 | """ calculate every group's weighted mean 118 | 119 | r_weighted_mean = sigma(r * count) / sigma(count) 120 | g_weighted_mean = sigma(g * count) / sigma(count) 121 | b_weighted_mean = sigma(b * count) / sigma(count) 122 | """ 123 | sigma_count = 0 124 | sigma_r = 0 125 | sigma_g = 0 126 | sigma_b = 0 127 | 128 | for item in grouped_image_color: 129 | sigma_count += item[0] 130 | sigma_r += item[1][0] * item[0] 131 | sigma_g += item[1][1] * item[0] 132 | sigma_b += item[1][2] * item[0] 133 | 134 | r_weighted_mean = int(sigma_r / sigma_count) 135 | g_weighted_mean = int(sigma_g / sigma_count) 136 | b_weighted_mean = int(sigma_b / sigma_count) 137 | 138 | weighted_mean = (sigma_count, (r_weighted_mean, g_weighted_mean, b_weighted_mean)) 139 | return weighted_mean 140 | 141 | class Haishoku(object): 142 | 143 | """ init Haishoku obj 144 | """ 145 | def __init__(self): 146 | self.dominant = None 147 | self.palette = None 148 | 149 | """ immediate api 150 | 151 | 1. showPalette 152 | 2. showDominant 153 | 3. getDominant 154 | 4. getPalette 155 | """ 156 | def getColorsMean(image): 157 | # get colors tuple 158 | image_colors = get_colors(image) 159 | 160 | # sort the image colors tuple 161 | sorted_image_colors = sort_by_rgb(image_colors) 162 | 163 | # group the colors by the accuaracy 164 | grouped_image_colors = group_by_accuracy(sorted_image_colors) 165 | 166 | # get the weighted mean of all colors 167 | colors_mean = [] 168 | for i in range(3): 169 | for j in range(3): 170 | for k in range(3): 171 | grouped_image_color = grouped_image_colors[i][j][k] 172 | if 0 != len(grouped_image_color): 173 | color_mean = get_weighted_mean(grouped_image_color) 174 | colors_mean.append(color_mean) 175 | 176 | # return the most 8 colors 177 | temp_sorted_colors_mean = sorted(colors_mean) 178 | if 8 < len(temp_sorted_colors_mean): 179 | colors_mean = temp_sorted_colors_mean[len(temp_sorted_colors_mean)-8 : len(temp_sorted_colors_mean)] 180 | else: 181 | colors_mean = temp_sorted_colors_mean 182 | 183 | # sort the colors_mean 184 | colors_mean = sorted(colors_mean, reverse=True) 185 | 186 | return colors_mean 187 | 188 | def getDominant(image=None): 189 | # get the colors_mean 190 | colors_mean = Haishoku.getColorsMean(image) 191 | colors_mean = sorted(colors_mean, reverse=True) 192 | 193 | # get the dominant color 194 | dominant_tuple = colors_mean[0] 195 | dominant = dominant_tuple[1] 196 | return dominant -------------------------------------------------------------------------------- /deepvac/utils/face_utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | import torch 4 | from itertools import product as product 5 | import numpy as np 6 | from math import ceil 7 | 8 | 9 | # Adapted from https://github.com/Hakuyume/chainer-ssd 10 | def decode(loc, priors, variances): 11 | """Decode locations from predictions using priors to undo 12 | the encoding we did for offset regression at train time. 13 | Args: 14 | loc (tensor): location predictions for loc layers, 15 | Shape: [num_priors,4] 16 | priors (tensor): Prior boxes in center-offset form. 17 | Shape: [num_priors,4]. 18 | variances: (list[float]) Variances of priorboxes 19 | Return: 20 | decoded bounding box predictions 21 | """ 22 | 23 | boxes = torch.cat(( 24 | priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:], 25 | priors[:, 2:] * torch.exp(loc[:, 2:] * variances[1])), 1) 26 | boxes[:, :2] -= boxes[:, 2:] / 2 27 | boxes[:, 2:] += boxes[:, :2] 28 | return boxes 29 | 30 | def decode_landm(pre, priors, variances): 31 | """Decode landm from predictions using priors to undo 32 | the encoding we did for offset regression at train time. 33 | Args: 34 | pre (tensor): landm predictions for loc layers, 35 | Shape: [num_priors,10] 36 | priors (tensor): Prior boxes in center-offset form. 37 | Shape: [num_priors,4]. 38 | variances: (list[float]) Variances of priorboxes 39 | Return: 40 | decoded landm predictions 41 | """ 42 | landms = torch.cat((priors[:, :2] + pre[:, :2] * variances[0] * priors[:, 2:], 43 | priors[:, :2] + pre[:, 2:4] * variances[0] * priors[:, 2:], 44 | priors[:, :2] + pre[:, 4:6] * variances[0] * priors[:, 2:], 45 | priors[:, :2] + pre[:, 6:8] * variances[0] * priors[:, 2:], 46 | priors[:, :2] + pre[:, 8:10] * variances[0] * priors[:, 2:], 47 | ), dim=1) 48 | return landms 49 | 50 | # Original author: Francisco Massa: 51 | # https://github.com/fmassa/object-detection.torch 52 | # Ported to PyTorch by Max deGroot (02/01/2017) 53 | def nms(boxes, scores, overlap=0.5, top_k=200): 54 | """Apply non-maximum suppression at test time to avoid detecting too many 55 | overlapping bounding boxes for a given object. 56 | Args: 57 | boxes: (tensor) The location preds for the img, Shape: [num_priors,4]. 58 | scores: (tensor) The class predscores for the img, Shape:[num_priors]. 59 | overlap: (float) The overlap thresh for suppressing unnecessary boxes. 60 | top_k: (int) The Maximum number of box preds to consider. 61 | Return: 62 | The indices of the kept boxes with respect to num_priors. 63 | """ 64 | 65 | keep = torch.Tensor(scores.size(0)).fill_(0).long() 66 | if boxes.numel() == 0: 67 | return keep 68 | x1 = boxes[:, 0] 69 | y1 = boxes[:, 1] 70 | x2 = boxes[:, 2] 71 | y2 = boxes[:, 3] 72 | area = torch.mul(x2 - x1, y2 - y1) 73 | v, idx = scores.sort(0) # sort in ascending order 74 | # I = I[v >= 0.01] 75 | idx = idx[-top_k:] # indices of the top-k largest vals 76 | xx1 = boxes.new() 77 | yy1 = boxes.new() 78 | xx2 = boxes.new() 79 | yy2 = boxes.new() 80 | w = boxes.new() 81 | h = boxes.new() 82 | 83 | # keep = torch.Tensor() 84 | count = 0 85 | while idx.numel() > 0: 86 | i = idx[-1] # index of current largest val 87 | # keep.append(i) 88 | keep[count] = i 89 | count += 1 90 | if idx.size(0) == 1: 91 | break 92 | idx = idx[:-1] # remove kept element from view 93 | # load bboxes of next highest vals 94 | torch.index_select(x1, 0, idx, out=xx1) 95 | torch.index_select(y1, 0, idx, out=yy1) 96 | torch.index_select(x2, 0, idx, out=xx2) 97 | torch.index_select(y2, 0, idx, out=yy2) 98 | # store element-wise max with next highest score 99 | xx1 = torch.clamp(xx1, min=x1[i]) 100 | yy1 = torch.clamp(yy1, min=y1[i]) 101 | xx2 = torch.clamp(xx2, max=x2[i]) 102 | yy2 = torch.clamp(yy2, max=y2[i]) 103 | w.resize_as_(xx2) 104 | h.resize_as_(yy2) 105 | w = xx2 - xx1 106 | h = yy2 - yy1 107 | # check sizes of xx1 and xx2.. after each iteration 108 | w = torch.clamp(w, min=0.0) 109 | h = torch.clamp(h, min=0.0) 110 | inter = w*h 111 | # IoU = i / (area(a) + area(b) - i) 112 | rem_areas = torch.index_select(area, 0, idx) # load remaining areas) 113 | union = (rem_areas - inter) + area[i] 114 | IoU = inter/union # store result in iou 115 | # keep only elements with an IoU <= overlap 116 | idx = idx[IoU.le(overlap)] 117 | return keep, count 118 | 119 | def py_cpu_nms(dets, thresh): 120 | """Pure Python NMS baseline.""" 121 | x1 = dets[:, 0] 122 | y1 = dets[:, 1] 123 | x2 = dets[:, 2] 124 | y2 = dets[:, 3] 125 | scores = dets[:, 4] 126 | 127 | areas = (x2 - x1 + 1) * (y2 - y1 + 1) 128 | order = scores.argsort()[::-1] 129 | 130 | keep = [] 131 | while order.size > 0: 132 | i = order[0] 133 | keep.append(i) 134 | xx1 = np.maximum(x1[i], x1[order[1:]]) 135 | yy1 = np.maximum(y1[i], y1[order[1:]]) 136 | xx2 = np.minimum(x2[i], x2[order[1:]]) 137 | yy2 = np.minimum(y2[i], y2[order[1:]]) 138 | 139 | w = np.maximum(0.0, xx2 - xx1 + 1) 140 | h = np.maximum(0.0, yy2 - yy1 + 1) 141 | inter = w * h 142 | ovr = inter / (areas[i] + areas[order[1:]] - inter) 143 | 144 | inds = np.where(ovr <= thresh)[0] 145 | order = order[inds + 1] 146 | 147 | return keep 148 | 149 | class PriorBox(object): 150 | def __init__(self, cfg, image_size=None, phase='train'): 151 | super(PriorBox, self).__init__() 152 | self.min_sizes = cfg['min_sizes'] 153 | self.steps = cfg['steps'] 154 | self.clip = cfg['clip'] 155 | self.image_size = image_size 156 | self.feature_maps = [[ceil(self.image_size[0]/step), ceil(self.image_size[1]/step)] for step in self.steps] 157 | self.name = "s" 158 | 159 | def forward(self): 160 | anchors = [] 161 | for k, f in enumerate(self.feature_maps): 162 | min_sizes = self.min_sizes[k] 163 | for i, j in product(range(f[0]), range(f[1])): 164 | for min_size in min_sizes: 165 | s_kx = min_size / self.image_size[1] 166 | s_ky = min_size / self.image_size[0] 167 | dense_cx = [x * self.steps[k] / self.image_size[1] for x in [j + 0.5]] 168 | dense_cy = [y * self.steps[k] / self.image_size[0] for y in [i + 0.5]] 169 | for cy, cx in product(dense_cy, dense_cx): 170 | anchors += [cx, cy, s_kx, s_ky] 171 | 172 | # back to torch land 173 | output = torch.Tensor(anchors).view(-1, 4) 174 | if self.clip: 175 | output.clamp_(max=1, min=0) 176 | return output 177 | -------------------------------------------------------------------------------- /deepvac/loss/loss.py: -------------------------------------------------------------------------------- 1 | # -*- coding:utf-8 -*- 2 | import math 3 | import torch 4 | import torch.nn as nn 5 | from ..utils import LOG, addUserConfig 6 | from ..core import AttrDict 7 | 8 | class LossBase(nn.Module): 9 | def __init__(self, deepvac_config): 10 | super(LossBase, self).__init__() 11 | self.deepvac_loss_config = deepvac_config.loss 12 | self.initConfig() 13 | self.auditConfig() 14 | 15 | def initConfig(self): 16 | if self.name() not in self.deepvac_loss_config.keys(): 17 | self.deepvac_loss_config[self.name()] = AttrDict() 18 | self.config = self.deepvac_loss_config[self.name()] 19 | 20 | def addUserConfig(self, config_name, user_give=None, developer_give=None, is_user_mandatory=False): 21 | module_name = 'config.loss.{}'.format(self.name()) 22 | return addUserConfig(module_name, config_name, user_give, developer_give, is_user_mandatory) 23 | 24 | def name(self): 25 | return self.__class__.__name__ 26 | 27 | def auditConfig(self): 28 | raise Exception("Not implemented!") 29 | 30 | class MaskL1Loss(LossBase): 31 | def __init__(self, deepvac_config): 32 | super(MaskL1Loss, self).__init__(deepvac_config) 33 | 34 | def auditConfig(self): 35 | self.eps = self.addUserConfig('eps', self.config.eps, 1e-6) 36 | 37 | def __call__(self, pred: torch.Tensor, gt, mask): 38 | loss = (torch.abs(pred - gt) * mask).sum() / (mask.sum() + self.eps) 39 | return loss 40 | 41 | class DiceLoss(LossBase): 42 | def __init__(self, deepvac_config): 43 | super(DiceLoss, self).__init__(deepvac_config) 44 | 45 | def auditConfig(self): 46 | self.eps = self.addUserConfig('eps', self.config.eps, 1e-6) 47 | 48 | def __call__(self, pred: torch.Tensor, gt, mask, weights=None): 49 | return self._compute(pred, gt, mask, weights) 50 | 51 | def _compute(self, pred, gt, mask, weights): 52 | if pred.dim() == 4: 53 | pred = pred[:, 0, :, :] 54 | gt = gt[:, 0, :, :] 55 | assert pred.shape == gt.shape 56 | assert pred.shape == mask.shape 57 | if weights is not None: 58 | assert weights.shape == mask.shape 59 | mask = weights * mask 60 | intersection = (pred * gt * mask).sum() 61 | union = (pred * mask).sum() + (gt * mask).sum() + self.eps 62 | loss = 1 - 2.0 * intersection / union 63 | assert loss <= 1 64 | return loss 65 | 66 | class BalanceCrossEntropyLoss(LossBase): 67 | def __init__(self, deepvac_config): 68 | super(BalanceCrossEntropyLoss, self).__init__(deepvac_config) 69 | 70 | def auditConfig(self): 71 | self.negative_ratio = self.addUserConfig('negative_ratio', self.config.negative_ratio, 3.0) 72 | self.eps = self.addUserConfig('eps', self.config.eps, 1e-6) 73 | 74 | def __call__(self, 75 | pred: torch.Tensor, 76 | gt: torch.Tensor, 77 | mask: torch.Tensor, 78 | return_origin=False): 79 | positive = (gt * mask).byte() 80 | negative = ((1 - gt) * mask).byte() 81 | positive_count = int(positive.float().sum()) 82 | negative_count = min(int(negative.float().sum()), int(positive_count * self.negative_ratio)) 83 | loss = nn.functional.binary_cross_entropy(pred, gt, reduction='none') 84 | positive_loss = loss * positive.float() 85 | negative_loss = loss * negative.float() 86 | negative_loss, _ = negative_loss.view(-1).topk(negative_count) 87 | balance_loss = (positive_loss.sum() + negative_loss.sum()) / (positive_count + negative_count + self.eps) 88 | if return_origin: 89 | return balance_loss, loss 90 | return balance_loss 91 | 92 | class BCEBlurWithLogitsLoss(LossBase): 93 | # BCEwithLogitLoss() with reduced missing label effects. 94 | def __init__(self, deepvac_config): 95 | super(BCEBlurWithLogitsLoss, self).__init__(deepvac_config) 96 | 97 | def auditConfig(self): 98 | self.alpha = self.addUserConfig('alpha', self.config.alpha, 0.05) 99 | self.reduction = self.addUserConfig('reduction', self.config.reduction, 'none') 100 | self.loss_fcn = nn.BCEWithLogitsLoss(reduction=self.reduction) 101 | 102 | def __call__(self, pred, true): 103 | loss = self.loss_fcn(pred, true) 104 | pred = torch.sigmoid(pred) 105 | # reduce only missing label effects 106 | dx = pred - true 107 | alpha_factor = 1 - torch.exp((dx - 1) / (self.alpha + 1e-4)) 108 | loss *= alpha_factor 109 | return loss.mean() 110 | 111 | class FocalLoss(LossBase): 112 | # Wraps focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5) 113 | def __init__(self, deepvac_config): 114 | super(FocalLoss, self).__init__(deepvac_config) 115 | 116 | def auditConfig(self): 117 | self.alpha = self.addUserConfig('alpha', self.config.alpha, 0.25) 118 | self.gamma = self.addUserConfig('gamma', self.config.gamma, 1.5) 119 | self.reduction = self.addUserConfig('reduction', self.config.reduction, 'none') 120 | self.loss_fcn = nn.BCEWithLogitsLoss(reduction=self.reduction) 121 | 122 | def __call__(self, pred, true): 123 | loss = self.loss_fcn(pred, true) 124 | # TF implementation https://github.com/tensorflow/addons/blob/v0.7.1/tensorflow_addons/losses/focal_loss.py 125 | pred_prob = torch.sigmoid(pred) 126 | p_t = true * pred_prob + (1 - true) * (1 - pred_prob) 127 | alpha_factor = true * self.alpha + (1 - true) * (1 - self.alpha) 128 | modulating_factor = (1.0 - p_t) ** self.gamma 129 | loss *= alpha_factor * modulating_factor 130 | 131 | if self.reduction == 'mean': 132 | return loss.mean() 133 | elif self.reduction == 'sum': 134 | return loss.sum() 135 | else: 136 | return loss 137 | 138 | class QFocalLoss(LossBase): 139 | # Wraps Quality focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5) 140 | def __init__(self, deepvac_config): 141 | super(QFocalLoss, self).__init__(deepvac_config) 142 | 143 | def auditConfig(self): 144 | self.alpha = self.addUserConfig('alpha', self.config.alpha, 0.25) 145 | self.gamma = self.addUserConfig('gamma', self.config.gamma, 1.5) 146 | self.reduction = self.addUserConfig('reduction', self.config.reduction, 'none') 147 | self.loss_fcn = nn.BCEWithLogitsLoss() 148 | 149 | def __call__(self, pred, true): 150 | loss = self.loss_fcn(pred, true) 151 | pred_prob = torch.sigmoid(pred) 152 | alpha_factor = true * self.alpha + (1 - true) * (1 - self.alpha) 153 | modulating_factor = torch.abs(true - pred_prob) ** self.gamma 154 | loss *= alpha_factor * modulating_factor 155 | 156 | if self.reduction == 'mean': 157 | return loss.mean() 158 | elif self.reduction == 'sum': 159 | return loss.sum() 160 | else: 161 | return loss 162 | 163 | class WingLoss(nn.Module): 164 | def __init__(self): 165 | super(WingLoss, self).__init__() 166 | 167 | def forward(self, pred, truth, w=10.0, epsilon=2.0): 168 | x = truth - pred 169 | c = w * (1.0 - math.log(1.0 + w / epsilon)) 170 | absolute_x = torch.abs(x) 171 | losses = torch.where(w > absolute_x, w * torch.log(1.0 + absolute_x / epsilon), absolute_x - c) 172 | return torch.sum(losses) / (len(losses) * 1.0) 173 | -------------------------------------------------------------------------------- /deepvac/utils/face_align.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | import os 3 | import sys 4 | import numpy as np 5 | from numpy.linalg import inv, norm, lstsq 6 | from numpy.linalg import matrix_rank as rank 7 | import time 8 | import cv2 9 | 10 | class AlignFace(object): 11 | def __init__(self): 12 | # reference facial points, a list of coordinates (x,y) 13 | self.REFERENCE_FACIAL_POINTS_96x112 = [ 14 | [30.29459953, 51.69630051], 15 | [65.53179932, 51.50139999], 16 | [48.02519989, 71.73660278], 17 | [33.54930115, 92.3655014], 18 | [62.72990036, 92.20410156] 19 | ] 20 | # made by gemfield 21 | self.REFERENCE_FACIAL_POINTS_112x112 = [ 22 | [38.29459953, 51.69630051], 23 | [73.53179932, 51.50139999], 24 | [56.02519989, 71.73660278], 25 | [41.54930115, 92.3655014 ], 26 | [70.72990036, 92.20410156] 27 | ] 28 | 29 | def __call__(self, frame, facial_5pts): 30 | # shape from (10,) to (2, 5) 31 | # x1,x2...y1,y2... or x1,y1,x2,y2... 32 | facial = [] 33 | x = facial_5pts[::2] 34 | y = facial_5pts[1::2] 35 | facial.append(x) 36 | facial.append(y) 37 | dst_img = self.warpAndCrop(frame, facial, (112, 112)) 38 | return dst_img 39 | 40 | def warpAndCrop(self, src_img, facial_pts, crop_size): 41 | reference_pts = self.REFERENCE_FACIAL_POINTS_112x112 42 | ref_pts = np.float32(reference_pts) 43 | ref_pts_shp = ref_pts.shape 44 | 45 | if ref_pts_shp[0] == 2: 46 | ref_pts = ref_pts.T 47 | 48 | src_pts = np.float32(facial_pts) 49 | src_pts_shp = src_pts.shape 50 | if max(src_pts_shp) < 3 or min(src_pts_shp) != 2: 51 | raise Exception('facial_pts.shape must be (K,2) or (2,K) and K>2') 52 | # 2*5 to 5*2 53 | if src_pts_shp[0] == 2: 54 | src_pts = src_pts.T 55 | 56 | if src_pts.shape != ref_pts.shape: 57 | raise Exception('facial_pts and reference_pts must have the same shape: {} vs {}'.format(src_pts.shape, ref_pts.shape) ) 58 | 59 | tfm = self.getAffineTransform(src_pts, ref_pts) 60 | 61 | face_img = cv2.warpAffine(src_img, tfm, (crop_size[0], crop_size[1])) 62 | return face_img 63 | 64 | def getAffineTransform(self, uv, xy): 65 | options = {'K': 2} 66 | # Solve for trans1 67 | trans1, trans1_inv = self.findNonreflectiveSimilarity(uv, xy, options) 68 | # manually reflect the xy data across the Y-axis 69 | xyR = xy 70 | xyR[:, 0] = -1 * xyR[:, 0] 71 | 72 | trans2r, trans2r_inv = self.findNonreflectiveSimilarity(uv, xyR, options) 73 | 74 | # manually reflect the tform to undo the reflection done on xyR 75 | TreflectY = np.array([ 76 | [-1, 0, 0], 77 | [0, 1, 0], 78 | [0, 0, 1] 79 | ]) 80 | 81 | trans2 = np.dot(trans2r, TreflectY) 82 | 83 | # Figure out if trans1 or trans2 is better 84 | xy1 = self.tformfwd(trans1, uv) 85 | norm1 = norm(xy1 - xy) 86 | 87 | xy2 = self.tformfwd(trans2, uv) 88 | norm2 = norm(xy2 - xy) 89 | 90 | if norm1 <= norm2: 91 | trans = trans1 92 | else: 93 | trans2_inv = inv(trans2) 94 | trans = trans2 95 | 96 | cv2_trans = trans[:, 0:2].T 97 | return cv2_trans 98 | 99 | def findNonreflectiveSimilarity(self, uv, xy, options=None): 100 | options = {'K': 2} 101 | 102 | K = options['K'] 103 | M = xy.shape[0] 104 | x = xy[:, 0].reshape((-1, 1)) 105 | y = xy[:, 1].reshape((-1, 1)) 106 | 107 | tmp1 = np.hstack((x, y, np.ones((M, 1)), np.zeros((M, 1)))) 108 | tmp2 = np.hstack((y, -x, np.zeros((M, 1)), np.ones((M, 1)))) 109 | X = np.vstack((tmp1, tmp2)) 110 | 111 | u = uv[:, 0].reshape((-1, 1)) 112 | v = uv[:, 1].reshape((-1, 1)) 113 | U = np.vstack((u, v)) 114 | 115 | if rank(X) >= 2 * K: 116 | r, _, _, _ = lstsq(X, U, rcond=-1) 117 | r = np.squeeze(r) 118 | else: 119 | raise Exception('cp2tform:twoUniquePointsReq') 120 | sc = r[0] 121 | ss = r[1] 122 | tx = r[2] 123 | ty = r[3] 124 | 125 | Tinv = np.array([ 126 | [sc, -ss, 0], 127 | [ss, sc, 0], 128 | [tx, ty, 1] 129 | ]) 130 | 131 | T = inv(Tinv) 132 | T[:, 2] = np.array([0, 0, 1]) 133 | return T, Tinv 134 | 135 | def tformfwd(self, trans, uv): 136 | uv = np.hstack(( 137 | uv, np.ones((uv.shape[0], 1)) 138 | )) 139 | xy = np.dot(uv, trans) 140 | xy = xy[:, 0:-1] 141 | return xy 142 | 143 | def drawBoxes(self, img, imgname): 144 | write_path = './results/{}/bbox/'.format( time.strftime("%Y%m%d", time.localtime()) ) 145 | if not os.path.exists(write_path): 146 | os.makedirs(write_path) 147 | cv2.imwrite('{}{}_aligned.jpg'.format(write_path, imgname), img) 148 | 149 | class AlignFaceWith51Points(AlignFace): 150 | def __init__(self): 151 | super(AlignFaceWith51Points, self).__init__() 152 | self.REFERENCE_FACIAL_POINTS_112x112 = self.getPosition() 153 | 154 | def getPosition(self): 155 | 156 | x = [0.000213256, 0.0752622, 0.18113, 0.29077, 0.393397, 0.586856, 0.689483, 0.799124, 157 | 0.904991, 0.98004, 0.490127, 0.490127, 0.490127, 0.490127, 0.36688, 0.426036, 158 | 0.490127, 0.554217, 0.613373, 0.121737, 0.187122, 0.265825, 0.334606, 0.260918, 159 | 0.182743, 0.645647, 0.714428, 0.793132, 0.858516, 0.79751, 0.719335, 0.254149, 160 | 0.340985, 0.428858, 0.490127, 0.551395, 0.639268, 0.726104, 0.642159, 0.556721, 161 | 0.490127, 0.423532, 0.338094, 0.290379, 0.428096, 0.490127, 0.552157, 0.689874, 162 | 0.553364, 0.490127, 0.42689] 163 | 164 | y = [0.106454, 0.038915, 0.0187482, 0.0344891, 0.0773906, 0.0773906, 0.0344891, 165 | 0.0187482, 0.038915, 0.106454, 0.203352, 0.307009, 0.409805, 0.515625, 0.587326, 166 | 0.609345, 0.628106, 0.609345, 0.587326, 0.216423, 0.178758, 0.179852, 0.231733, 167 | 0.245099, 0.244077, 0.231733, 0.179852, 0.178758, 0.216423, 0.244077, 0.245099, 168 | 0.780233, 0.745405, 0.727388, 0.742578, 0.727388, 0.745405, 0.780233, 0.864805, 169 | 0.902192, 0.909281, 0.902192, 0.864805, 0.784792, 0.778746, 0.785343, 0.778746, 170 | 0.784792, 0.824182, 0.831803, 0.824182] 171 | 172 | x, y = np.array(x), np.array(y) 173 | x = x * 112 174 | y = y * 112 175 | return np.array(list(zip(x, y))) 176 | 177 | def getAffineTransform(self, src_pts, ref_pts): 178 | src_pts = np.matrix(src_pts) 179 | ref_pts = np.matrix(ref_pts) 180 | mean1 = np.mean(src_pts, axis=0) 181 | mean2 = np.mean(ref_pts, axis=0) 182 | src_pts -= mean1 183 | ref_pts -= mean2 184 | std1 = np.std(src_pts) 185 | std2 = np.std(ref_pts) 186 | src_pts /= std1 187 | ref_pts /= std2 188 | 189 | U, S, Vt = np.linalg.svd(src_pts.T * ref_pts) 190 | R = (U * Vt).T 191 | M = np.vstack([ 192 | np.hstack(((std2 / std1) * R, 193 | mean2.T - (std2 / std1) * R * mean1.T)), 194 | np.matrix([0., 0., 1.]) 195 | ]) 196 | 197 | return M[:2] 198 | -------------------------------------------------------------------------------- /deepvac/aug/yolo_aug.py: -------------------------------------------------------------------------------- 1 | import math 2 | import cv2 3 | import numpy as np 4 | import random 5 | from .base_aug import CvAugBase2 6 | 7 | ### yolov5 dataset aug 8 | class HSVAug(CvAugBase2): 9 | def auditConfig(self): 10 | self.config.hgain = self.addUserConfig('hgain', self.config.hgain, 0.015) 11 | self.config.sgain = self.addUserConfig('sgain', self.config.sgain, 0.7) 12 | self.config.vgain = self.addUserConfig('vgain', self.config.vgain, 0.4) 13 | 14 | def forward(self, img): 15 | img, label = img 16 | assert isinstance(label, np.ndarray) and label.ndim == 2, "label must be numpy.ndarray, and shape should be (n, 5)" 17 | hue, sat, val = cv2.split(cv2.cvtColor(img, cv2.COLOR_BGR2HSV)) 18 | dtype = img.dtype 19 | x = np.arange(0, 256, dtype=np.int16) 20 | # 随机增幅 21 | r = np.random.uniform(-1, 1, 3) * [self.config.hgain, self.config.sgain, self.config.vgain] + 1 22 | lut_hue = ((x * r[0]) % 180).astype(dtype) 23 | lut_sat = np.clip(x * r[1], 0, 255).astype(dtype) 24 | lut_val = np.clip(x * r[2], 0, 255).astype(dtype) 25 | img_hsv = cv2.merge((cv2.LUT(hue, lut_hue), cv2.LUT(sat, lut_sat), cv2.LUT(val, lut_val))).astype(dtype) 26 | cv2.cvtColor(img_hsv, cv2.COLOR_HSV2BGR, dst=img) 27 | return img, label 28 | 29 | class YoloHFlipAug(CvAugBase2): 30 | def forward(self, img): 31 | img, label = img 32 | assert isinstance(label, np.ndarray) and label.ndim == 2, "label must be numpy.ndarray, and shape should be (n, 5)" 33 | img = np.fliplr(img) 34 | if label.size: 35 | label[:, 1] = 1 - label[:, 1] 36 | return img, label 37 | 38 | class YoloVFlipAug(CvAugBase2): 39 | def forward(self, img): 40 | img, label = img 41 | assert isinstance(label, np.ndarray) and label.ndim == 2, "label must be numpy.ndarray, and shape should be (n, 5)" 42 | img = np.flipud(img) 43 | if label.size: 44 | label[:, 2] = 1 - label[:, 2] 45 | return img, label 46 | 47 | class YoloPerspectiveAug(CvAugBase2): 48 | def auditConfig(self): 49 | self.config.scale = self.addUserConfig('scale', self.config.scale, 0.5) 50 | self.config.shear = self.addUserConfig('shear', self.config.shear, 0.0) 51 | self.config.degrees = self.addUserConfig('degrees', self.config.degrees, 0.0) 52 | self.config.translate = self.addUserConfig('translate', self.config.translate, 0.1) 53 | self.config.perspective = self.addUserConfig('perspective', self.config.perspective, 0.0) 54 | self.config.border = self.addUserConfig('border', self.config.border, -9999, True) 55 | 56 | def _box_candidates(self, box1, box2, wh_thr=2, ar_thr=20, area_thr=0.1): 57 | w1, h1 = box1[2] - box1[0], box1[3] - box1[1] 58 | w2, h2 = box2[2] - box2[0], box2[3] - box2[1] 59 | ar = np.maximum(w2 / (h2 + 1e-16), h2 / (w2 + 1e-16)) 60 | return (w2 > wh_thr) & (h2 > wh_thr) & (w2 * h2 / (w1 * h1 + 1e-16) > area_thr) & (ar < ar_thr) 61 | 62 | def forward(self, img): 63 | img, label = img 64 | assert isinstance(label, np.ndarray) and label.ndim == 2, "label must be numpy.ndarray, and shape should be (n, 5)" 65 | 66 | border = self.config.border 67 | h, w, c = img.shape 68 | 69 | width = int(w + border[1] * 2) 70 | height = int(h + border[0] * 2) 71 | # Center 72 | ''' 73 | [[1, 0, -w/2], 74 | [0, 1, -h/2], 75 | [0, 0, 1. ]] 76 | ''' 77 | C = np.eye(3) 78 | C[0, 2] = -w / 2 79 | C[1, 2] = -h / 2 80 | # Perspective 81 | ''' 82 | [[1, 0, 0], 83 | [0, 1, 0], 84 | [p, p, 1]] 85 | ''' 86 | P = np.eye(3) 87 | P[2, 0] = random.uniform(-self.config.perspective, self.config.perspective) 88 | P[2, 1] = random.uniform(-self.config.perspective, self.config.perspective) 89 | # Rotation and Scale 90 | ''' 91 | [[r, r, r], 92 | [r, r, r], 93 | [0, 0, 1]] 94 | ''' 95 | R = np.eye(3) 96 | a = random.uniform(-self.config.degrees, self.config.degrees) 97 | s = random.uniform(1 - self.config.scale, 1 + self.config.scale) 98 | R[:2] = cv2.getRotationMatrix2D(angle=a, center=(0, 0), scale=s) 99 | # Shear 100 | ''' 101 | [[1, s, 0], 102 | [s, 1, 0], 103 | [0, 0, 1]] 104 | ''' 105 | S = np.eye(3) 106 | S[0, 1] = math.tan(random.uniform(-self.config.shear, self.config.shear) * math.pi / 180) 107 | S[1, 0] = math.tan(random.uniform(-self.config.shear, self.config.shear) * math.pi / 180) 108 | # Translation 109 | ''' 110 | [[1, 0, t], 111 | [0, 1, t], 112 | [0, 0, 1]] 113 | ''' 114 | T = np.eye(3) 115 | T[0, 2] = random.uniform(0.5 - self.config.translate, 0.5 + self.config.translate) * width 116 | T[1, 2] = random.uniform(0.5 - self.config.translate, 0.5 + self.config.translate) * height 117 | # Combined rotation matrix 118 | M = T @ S @ R @ P @ C 119 | # img augment and resize to img_size 120 | if (border[0] != 0) or (border[1] != 0) or (M != np.eye(3)).any(): 121 | if self.config.perspective: 122 | img = cv2.warpPerspective(img, M, dsize=(width, height), borderValue=(114, 114, 114)) 123 | else: 124 | img = cv2.warpAffine(img, M[:2], dsize=(width, height), borderValue=(114, 114, 114)) 125 | 126 | n = len(label) 127 | if n: 128 | # warp points 129 | xy = np.ones((n * 4, 3)) 130 | ''' 131 | [[x1, y1], 132 | [x2, y2], 133 | [x1, y2], 134 | [x2, y1]] 135 | ''' 136 | xy[:, :2] = label[:, [1, 2, 3, 4, 1, 4, 3, 2]].reshape(n * 4, 2) 137 | xy = xy @ M.T 138 | if self.config.perspective: 139 | xy = (xy[:, :2] / xy[:, 2:3]).reshape(n, 8) 140 | else: 141 | xy = xy[:, :2].reshape(n, 8) 142 | # x: [[x1, x2, x1, x2], ...n...] 143 | x = xy[:, [0, 2, 4, 6]] 144 | # y: [[y1, y2, y2, y1], ...n...] 145 | y = xy[:, [1, 3, 5, 7]] 146 | # xy: [[xmin, ymin, xmax, ymax], ...n...] 147 | xy = np.concatenate((x.min(1), y.min(1), x.max(1), y.max(1))).reshape(4, n).T 148 | # clip boxes 149 | xy[:, [0, 2]] = xy[:, [0, 2]].clip(0, width) 150 | xy[:, [1, 3]] = xy[:, [1, 3]].clip(0, height) 151 | # filter candidates 152 | i = self._box_candidates(box1=label[:, 1:5].T * s, box2=xy.T) 153 | label = label[i] 154 | label[:, 1:5] = xy[i] 155 | return img, label 156 | 157 | 158 | class YoloNormalizeAug(CvAugBase2): 159 | def forward(self, img): 160 | ''' 161 | 1. [cls, x1, y1, x2, y2] -> [cls, cx, cy, w, h] 162 | 2. cx, w normalized by img width, cy, h normalized by img height 163 | ''' 164 | img, label = img 165 | assert isinstance(label, np.ndarray) and label.ndim == 2, "label must be numpy.ndarray, and shape should be (n, 5)" 166 | if not label.size: 167 | return img, label 168 | label[:, [3, 4]] -= label[:, [1, 2]] 169 | label[:, [1, 2]] += label[:, [3, 4]] / 2 170 | label[:, [1, 3]] /= img.shape[1] 171 | label[:, [2, 4]] /= img.shape[0] 172 | return img, label 173 | -------------------------------------------------------------------------------- /deepvac/aug/seg_aug.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | import random 4 | import torch 5 | from .base_aug import CvAugBase2 6 | 7 | class ImageWithMasksRandomHorizontalFlipAug(CvAugBase2): 8 | def forward(self, imgs): 9 | for i in range(len(imgs)): 10 | imgs[i] = np.flip(imgs[i], axis=1) 11 | return imgs 12 | 13 | class ImageWithMasksRandomRotateAug(CvAugBase2): 14 | def auditConfig(self): 15 | self.config.max_angle = self.addUserConfig('max_angle', self.config.max_angle, 10) 16 | self.config.input_len = self.addUserConfig('input_len', self.config.input_len, 2) 17 | self.config.label_bg_color = self.addUserConfig('label_bg_color', self.config.label_bg_color, (0, 0, 0)) 18 | 19 | def forward(self, imgs): 20 | angle = random.random() * 2 * self.config.max_angle - self.config.max_angle 21 | # fill color 22 | if isinstance(self.config.fill_color, (tuple, list)): 23 | fill_color = self.config.fill_color 24 | try: 25 | fill_color = [min(max(int(fill_color[i]), 0), 255) for i in range(3)] 26 | except: 27 | raise Exception("fill_color = (int, int, int) while fill_color type = list or tuple") 28 | elif isinstance(self.config.fill_color, bool): 29 | fill_color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)) if self.config.fill_color else (0, 0, 0) 30 | 31 | elif self.config.fill_color is None: 32 | fill_color = (0, 0, 0) 33 | else: 34 | raise TypeError("fill_color type must be (None, bool, tuple, list)") 35 | # ops 36 | w, h = imgs[0].shape[:2] 37 | rotation_matrix = cv2.getRotationMatrix2D((h / 2, w / 2), angle, 1) 38 | imgs[0] = cv2.warpAffine(imgs[0], rotation_matrix, (h, w), borderValue=fill_color) 39 | for i in range(1, len(imgs)): 40 | imgs[i] = cv2.warpAffine(imgs[i], rotation_matrix, (h, w), borderValue=self.config.label_bg_color) 41 | return imgs 42 | 43 | class ImageWithMasksRandom4TextCropAug(CvAugBase2): 44 | def auditConfig(self): 45 | self.config.p = self.addUserConfig('p', self.config.p, 3.0 / 8.0) 46 | self.config.img_size = self.addUserConfig('img_size', self.config.img_size, 224) 47 | self.config.input_len = self.addUserConfig('input_len', self.config.input_len, 2) 48 | 49 | def forward(self, imgs): 50 | h, w, _ = imgs[0].shape 51 | if max(h, w) <= self.config.img_size: 52 | return imgs 53 | 54 | img_size = (self.config.img_size,) * 2 55 | th = min(h, img_size[0]) 56 | tw = min(w, img_size[0]) 57 | 58 | if random.random() > self.config.p and np.max(imgs[1]) > 0: 59 | tl = np.min(np.where(imgs[1] > 0), axis = 1) - img_size 60 | tl[tl < 0] = 0 61 | br = np.max(np.where(imgs[1] > 0), axis = 1) - img_size 62 | br[br < 0] = 0 63 | br[0] = min(br[0], h - th) 64 | br[1] = min(br[1], w - tw) 65 | 66 | i = random.randint(tl[0], br[0]) 67 | j = random.randint(tl[1], br[1]) 68 | else: 69 | i = random.randint(0, h - th) 70 | j = random.randint(0, w - tw) 71 | 72 | # return i, j, th, tw 73 | for idx in range(len(imgs)): 74 | if len(imgs[idx].shape) == 3: 75 | imgs[idx] = imgs[idx][i:i + th, j:j + tw, :] 76 | else: 77 | imgs[idx] = imgs[idx][i:i + th, j:j + tw] 78 | return imgs 79 | 80 | class ImageWithMasksScaleAug(CvAugBase2): 81 | def auditConfig(self): 82 | self.config.w = self.addUserConfig('w', self.config.w, 384, True) 83 | self.config.h = self.addUserConfig('h', self.config.h, 384, True) 84 | 85 | def forward(self, imgs): 86 | img, label = imgs 87 | img = cv2.resize(img, (self.config.w, self.config.h)) 88 | label = cv2.resize(label, (self.config.w, self.config.h), interpolation=cv2.INTER_NEAREST) 89 | return [img, label] 90 | 91 | class ImageWithMasksSafeCropAug(CvAugBase2): 92 | def forward(self, imgs): 93 | img, label = imgs 94 | h, w = img.shape[:2] 95 | 96 | xmin, ymin, xmax, ymax = self.getXY(label) 97 | x1 = random.randint(0, xmin) 98 | y1 = random.randint(0, ymin) 99 | x2 = random.randint(xmax, w) 100 | y2 = random.randint(ymax, h) 101 | 102 | img_crop = img[y1:y2, x1:x2, :] 103 | label_crop = label[y1:y2, x1:x2] 104 | return img_crop, label_crop 105 | 106 | def getXY(self, label): 107 | coord = label.nonzero() 108 | ymin, xmin = coord[0].min(), coord[1].min() 109 | ymax, xmax = coord[0].max(), coord[1].max() 110 | return xmin, ymin, xmax, ymax 111 | 112 | class ImageWithMasksCenterCropAug(CvAugBase2): 113 | def auditConfig(self): 114 | self.config.max_crop_ratio = self.addUserConfig('max_crop_ratio', self.config.max_crop_ratio, 0.1) 115 | 116 | def forward(self, imgs): 117 | img, label = imgs 118 | h, w = img.shape[:2] 119 | x1 = random.randint(0, int(w*self.config.max_crop_ratio)) # 25% to 10% 120 | y1 = random.randint(0, int(h*self.config.max_crop_ratio)) 121 | 122 | img_crop = img[y1:h-y1, x1:w-x1] 123 | label_crop = label[y1:h-y1, x1:w-x1] 124 | 125 | return [img_crop, label_crop] 126 | 127 | class ImageWithMasksHFlipAug(CvAugBase2): 128 | def forward(self, imgs): 129 | img, label = imgs 130 | img = cv2.flip(img, 1) # horizontal flip 131 | label = cv2.flip(label, 1) # horizontal flip 132 | return [img, label] 133 | 134 | class ImageWithMasksVFlipAug(CvAugBase2): 135 | def forward(self, imgs): 136 | img, label = imgs 137 | img = cv2.flip(img, 0) # veritcal flip 138 | label = cv2.flip(label, 0) # veritcal flip 139 | return [img, label] 140 | 141 | class ImageWithMasksNormalizeAug(CvAugBase2): 142 | def auditConfig(self): 143 | self.config.mean = self.addUserConfig('mean', self.config.mean, [137.98341,142.35637,150.78705], True) 144 | self.config.std = self.addUserConfig('std', self.config.std, [62.98702,63.34315,62.743645], True) 145 | 146 | def forward(self, imgs): 147 | img, label = imgs 148 | img = img.astype(np.float32) 149 | for i in range(3): 150 | img[:,:,i] -= self.config.mean[i] 151 | for i in range(3): 152 | img[:,:, i] /= self.config.std[i] 153 | return [img, label] 154 | 155 | class ImageWithMasksToTensorAug(CvAugBase2): 156 | def auditConfig(self): 157 | self.config.scale = self.addUserConfig('scale', self.config.scale, 1) 158 | self.config.force_div255 = self.addUserConfig('force_div255', self.config.force_div255, True) 159 | 160 | def forward(self, imgs): 161 | img, label = imgs 162 | 163 | if self.config.scale != 1: 164 | h, w = label.shape[:2] 165 | img = cv2.resize(img, (int(w), int(h))) 166 | label = cv2.resize(label, (int(w/self.config.scale), int(h/self.config.scale)), interpolation=cv2.INTER_NEAREST) 167 | 168 | default_float_dtype = torch.get_default_dtype() 169 | image_tensor = torch.from_numpy(img.transpose((2, 0, 1))).contiguous() 170 | # backward compatibility 171 | if isinstance(image_tensor, torch.ByteTensor) or self.config.force_div255: 172 | image_tensor = image_tensor.to(dtype=default_float_dtype).div(255) 173 | label_tensor = torch.LongTensor(np.array(label, dtype=np.int)) #torch.from_numpy(label) 174 | 175 | return [image_tensor, label_tensor] 176 | -------------------------------------------------------------------------------- /deepvac/aug/line_helper.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | 3 | class LineState(object): 4 | tableline_x_offsets = range(1, 5) 5 | tableline_y_offsets = range(1, 5) 6 | tableline_thickness = [2, 3] 7 | 8 | # 0/1/2/3: 仅单边（左上右下） 9 | # 4/5/6/7: 两边都有线（左上，右上，右下，左下） 10 | tableline_options = range(0, 8) 11 | 12 | middleline_thickness = [1, 2, 3] 13 | middleline_thickness_p = [0.2, 0.7, 0.1] 14 | 15 | class Liner(object): 16 | def __init__(self): 17 | self.linestate = LineState() 18 | self.cfg = {} 19 | self.cfg['line'] = {} 20 | 21 | self.cfg['line']['under_line'] = {} 22 | self.cfg['line']['under_line']['enable'] = True 23 | self.cfg['line']['under_line']['fraction'] = 0.5 24 | 25 | self.cfg['line']['table_line'] = {} 26 | self.cfg['line']['table_line']['enable'] = True 27 | self.cfg['line']['table_line']['fraction'] = 0.5 28 | 29 | self.cfg['line']['middle_line'] = {} 30 | self.cfg['line']['middle_line']['enable'] = True 31 | self.cfg['line']['middle_line']['fraction'] = 0 32 | 33 | 34 | self.cfg['line_color'] = {} 35 | self.cfg['line_color']['enable'] = True 36 | 37 | self.cfg['line_color']['black'] = {} 38 | self.cfg['line_color']['black']['fraction'] = 0.5 39 | self.cfg['line_color']['black']['l_boundary'] = [0,0,0] 40 | self.cfg['line_color']['black']['h_boundary'] = [64,64,64] 41 | 42 | self.cfg['line_color']['blue'] = {} 43 | self.cfg['line_color']['blue']['fraction'] = 0.5 44 | self.cfg['line_color']['blue']['l_boundary'] = [0,0,150] 45 | self.cfg['line_color']['blue']['h_boundary'] = [60,64,255] 46 | 47 | 48 | def get_line_color(self): 49 | p = [] 50 | colors = [] 51 | for k, v in self.cfg['line_color'].items(): 52 | if k == 'enable': 53 | continue 54 | p.append(v['fraction']) 55 | colors.append(k) 56 | 57 | # pick color by fraction 58 | color_name = np.random.choice(colors, p=p) 59 | l_boundary = self.cfg['line_color'][color_name]['l_boundary'] 60 | h_boundary = self.cfg['line_color'][color_name]['h_boundary'] 61 | # random color by low and high RGB boundary 62 | r = np.random.randint(l_boundary[0], h_boundary[0]) 63 | g = np.random.randint(l_boundary[1], h_boundary[1]) 64 | b = np.random.randint(l_boundary[2], h_boundary[2]) 65 | return b, g, r 66 | 67 | def apply(self, word_img, text_box_pnts): 68 | """ 69 | :param word_img: word image with big background 70 | :param text_box_pnts: left-top, right-top, right-bottom, left-bottom of text word 71 | :return: 72 | """ 73 | line_p = [] 74 | funcs = [] 75 | 76 | if self.cfg['line']['under_line']['enable']: 77 | line_p.append(self.cfg['line']['under_line']['fraction']) 78 | funcs.append(self.apply_under_line) 79 | 80 | if self.cfg['line']['table_line']['enable']: 81 | line_p.append(self.cfg['line']['table_line']['fraction']) 82 | funcs.append(self.apply_table_line) 83 | 84 | if self.cfg['line']['middle_line']['enable']: 85 | line_p.append(self.cfg['line']['middle_line']['fraction']) 86 | funcs.append(self.apply_middle_line) 87 | 88 | 89 | if len(line_p) == 0: 90 | return word_img, text_box_pnts 91 | 92 | line_effect_func = np.random.choice(funcs, p=line_p) 93 | 94 | if self.cfg['line_color']['enable'] or self.cfg['font_color']['enable']: 95 | line_color = self.get_line_color() 96 | else: 97 | line_color = word_color + random.randint(0, 10) 98 | 99 | return line_effect_func(word_img, text_box_pnts, line_color) 100 | 101 | def apply_under_line(self, word_img, text_box_pnts, line_color): 102 | y_offset = random.choice([0, 1]) 103 | 104 | text_box_pnts[2][1] += y_offset 105 | text_box_pnts[3][1] += y_offset 106 | 107 | dst = cv2.line(word_img, 108 | (text_box_pnts[2][0], text_box_pnts[2][1]), 109 | (text_box_pnts[3][0], text_box_pnts[3][1]), 110 | color=line_color, 111 | thickness=3, 112 | lineType=cv2.LINE_AA) 113 | 114 | return dst, text_box_pnts 115 | 116 | def apply_table_line(self, word_img, text_box_pnts, line_color): 117 | """ 118 | 共有 8 种可能的画法，横线横穿整张 word_img 119 | 0/1/2/3: 仅单边（左上右下） 120 | 4/5/6/7: 两边都有线（左上，右上，右下，左下） 121 | """ 122 | dst = word_img 123 | option = random.choice(self.linestate.tableline_options) 124 | thickness = random.choice(self.linestate.tableline_thickness) 125 | 126 | top_y_offset = random.choice(self.linestate.tableline_y_offsets) 127 | bottom_y_offset = random.choice(self.linestate.tableline_y_offsets) 128 | left_x_offset = random.choice(self.linestate.tableline_x_offsets) 129 | right_x_offset = random.choice(self.linestate.tableline_x_offsets) 130 | 131 | def is_top(): 132 | return option in [1, 4, 5] 133 | 134 | def is_bottom(): 135 | return option in [3, 6, 7] 136 | 137 | def is_left(): 138 | return option in [0, 4, 7] 139 | 140 | def is_right(): 141 | return option in [2, 5, 6] 142 | 143 | if is_top(): 144 | text_box_pnts[0][1] -= top_y_offset 145 | text_box_pnts[1][1] -= top_y_offset 146 | 147 | if is_bottom(): 148 | text_box_pnts[2][1] += bottom_y_offset 149 | text_box_pnts[3][1] += bottom_y_offset 150 | 151 | if is_left(): 152 | text_box_pnts[0][0] -= left_x_offset 153 | text_box_pnts[3][0] -= left_x_offset 154 | 155 | if is_right(): 156 | text_box_pnts[1][0] += right_x_offset 157 | text_box_pnts[2][0] += right_x_offset 158 | 159 | if is_bottom(): 160 | dst = cv2.line(dst, 161 | (0, text_box_pnts[2][1]), 162 | (word_img.shape[1], text_box_pnts[3][1]), 163 | color=line_color, 164 | thickness=thickness, 165 | lineType=cv2.LINE_AA) 166 | 167 | if is_top(): 168 | dst = cv2.line(dst, 169 | (0, text_box_pnts[0][1]), 170 | (word_img.shape[1], text_box_pnts[1][1]), 171 | color=line_color, 172 | thickness=thickness, 173 | lineType=cv2.LINE_AA) 174 | 175 | if is_left(): 176 | dst = cv2.line(dst, 177 | (text_box_pnts[0][0], 0), 178 | (text_box_pnts[3][0], word_img.shape[0]), 179 | color=line_color, 180 | thickness=thickness, 181 | lineType=cv2.LINE_AA) 182 | 183 | if is_right(): 184 | dst = cv2.line(dst, 185 | (text_box_pnts[1][0], 0), 186 | (text_box_pnts[2][0], word_img.shape[0]), 187 | color=line_color, 188 | thickness=thickness, 189 | lineType=cv2.LINE_AA) 190 | 191 | return dst, text_box_pnts 192 | 193 | def apply_middle_line(self, word_img, text_box_pnts, line_color): 194 | y_center = int((text_box_pnts[0][1] + text_box_pnts[3][1]) / 2) 195 | 196 | thickness = np.random.choice(self.linestate.middleline_thickness, p=self.linestate.middleline_thickness_p) 197 | 198 | dst = cv2.line(word_img, 199 | (text_box_pnts[0][0], y_center), 200 | (text_box_pnts[1][0], y_center), 201 | color=line_color, 202 | thickness=thickness, 203 | lineType=cv2.LINE_AA) 204 | 205 | return dst, text_box_pnts -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # DeepVAC 2 | DeepVAC提供了基于PyTorch的AI项目的工程化规范。为了达到这一目标，DeepVAC包含了： 3 | - 软件工程规范：[软件工程规范](./docs/arch.md)； 4 | - 代码规范：[代码规范](./docs/code_standard.md)； 5 | - deepvac库：[deepvac库](./docs/lib.md)。 6 | 7 | 诸多PyTorch AI项目的内在逻辑都大同小异，因此DeepVAC致力于把更通用的逻辑剥离出来，从而使得工程代码的准确性、易读性、可维护性上更具优势。 8 | 9 | 如果想使得AI项目符合DeepVAC规范，需要仔细阅读[DeepVAC标准](./docs/deepvac_standard.md)。 10 | 如果想了解deepvac库的设计，请阅读[deepvac库的设计](./docs/design.md)。 11 | 12 | 13 | # 如何基于DeepVAC构建自己的PyTorch AI项目 14 | 15 | ## 1. 阅读[DeepVAC标准](./docs/deepvac_standard.md) 16 | 可以粗略阅读，建立起第一印象。 17 | 18 | ## 2. 环境准备 19 | DeepVAC的依赖有： 20 | - Python3。不支持Python2，其已被废弃； 21 | - 依赖包：torch, torchvision, tensorboard, scipy, numpy, cv2, Pillow； 22 | 23 | 这些依赖使用pip命令（或者git clone）自行安装，不再赘述。 24 | 25 | 对于普通用户来说，最方便高效的方式还是使用[MLab HomePod](https://github.com/DeepVAC/MLab#mlab-homepod)作为DeepVAC的使用环境，这是一个预构建的Docker image，可以帮助用户省掉不必要的环境配置时间。 26 | 同时在MLab组织内部，我们也使用[MLab HomePod](https://github.com/DeepVAC/MLab#mlab-homepod)进行日常的模型的训练任务。 27 | 28 | 29 | ## 3. 安装deepvac库 30 | 可以使用pip来进行安装： 31 | ```pip3 install deepvac``` 32 | 或者 33 | ```python3 -m pip install deepvac``` 34 | 35 | 36 | 如果你需要使用deepvac在github上的最新代码，就需要使用如下的开发者模式： 37 | #### 开发者模式 38 | - 克隆该项目到本地：```git clone https://github.com/DeepVAC/deepvac ``` 39 | - 在你的入口文件中添加： 40 | ```python 41 | import sys 42 | #replace with your local deepvac directory 43 | sys.path.insert(0,'/home/gemfield/github/deepvac') 44 | ``` 45 | 或者设置PYTHONPATH环境变量： 46 | ```bash 47 | export PYTHONPATH=/home/gemfield/github/deepvac 48 | ``` 49 | 50 | ## 4. 创建自己的PyTorch项目 51 | - 初始化自己项目的git仓库； 52 | - 在仓库中创建第一个研究分支，比如分支名为 LTS_b1_aug9_movie_video_plate_130w； 53 | - 切换到上述的LTS_b1分支中，开始工作； 54 | 55 | ## 5. 编写配置文件 56 | 配置文件的文件名均为 config.py，位于你项目的根目录。在代码开始处添加```from deepvac import new, AttrDict```； 57 | 所有用户的配置都存放在这个文件里。config模块提供了6个预定义的作用域：config.core,config.aug,config.cast,config.datasets,config.backbones,config.loss。使用方法如下： 58 | - 所有和trainer相关（包括train、val、test）的配置都定义在config.core.中； 59 | - 所有和deepvac.aug中增强模块相关的配置都定义在config.aug.中； 60 | - 所有和模型转换相关的配置都定义在config.cast.中； 61 | - 所有和Datasets相关的配置都定义在config.datasets.中； 62 | - 所有和loss相关的配置都定义在config.loss.中； 63 | - 用户可以开辟自己的作用域，比如config.my_stuff = AttrDict()，然后config.my_stuff.name = 'gemfield'； 64 | - 用户可以使用new()来初始化config实例，使用clone()来深拷贝config配置项。 65 | 66 | 更多配置： 67 | - 预训练模型加载； 68 | - checkpoint加载； 69 | - tensorboard使用； 70 | - TorchScript使用； 71 | - 转换ONNX； 72 | - 转换NCNN； 73 | - 转换CoreML； 74 | - 转换TensorRT； 75 | - 转换TNN； 76 | - 转换MNN； 77 | - 开启量化； 78 | - 开启EMA； 79 | - 开启自动混合精度训练(AMP)； 80 | 81 | 以及关于配置文件的更详细解释，请阅读[config说明](./docs/config.md). 82 | 83 | 项目根目录下的train.py中用如下方式引用config.py文件: 84 | 85 | ```python 86 | from config import config as deepvac_config 87 | from deepvac import DeepvacTrain 88 | 89 | class MyTrain(DeepvacTrain): 90 | ...... 91 | 92 | my_train = MyTrain(deepvac_config) 93 | my_train() 94 | ``` 95 | 96 | 项目根目录下的test.py中用如下方式引用config.py文件: 97 | ```python 98 | from config import config as deepvac_config 99 | from deepvac import Deepvac 100 | 101 | class MyTest(Deepvac) 102 | ...... 103 | 104 | my_test = MyTest(deepvac_config) 105 | my_test() 106 | ``` 107 | 108 | 之后，train.py/test.py代码中通过如下方式来读写config.core中的配置项 109 | ```python 110 | print(self.config.log_dir) 111 | print(self.config.batch_size) 112 | ...... 113 | ``` 114 | 此外，鉴于config的核心作用，deepvac还设计了如下的API来方便对config模块的使用： 115 | - AttrDict 116 | - new 117 | - interpret 118 | - fork 119 | - clone 120 | ```python 121 | from deepvac import AttrDict, new, interpret, fork 122 | ``` 123 | 关于这些API的使用方法，请访问[config API 说明](./docs/design.md#config-module). 124 | ## 6. 编写synthesis/synthesis.py（可选） 125 | 编写该文件，用于产生本项目的数据集，用于对本项目的数据集进行自动化检查和清洗。 126 | 这一步为可选，如果有需要的话，可以参考Deepvac组织下Synthesis2D项目的实现。 127 | 128 | ## 7. 编写aug/aug.py（可选） 129 | 编写该文件，用于实现数据增强策略。 130 | deepvac.aug模块为数据增强设计了特有的语法，在两个层面实现了复用：aug 和 composer。比如说，我想复用添加随机斑点的SpeckleAug： 131 | ```python 132 | from deepvac.aug.base_aug import SpeckleAug 133 | ``` 134 | 135 | 这是对底层aug算子的复用。我们还可以直接复用别人写好的composer，并且是以直截了当的方式。比如deepvac.aug提供了一个用于人脸检测数据增强的RetinaAugComposer： 136 | 137 | ```python 138 | from deepvac.aug import RetinaAugComposer 139 | ``` 140 | 141 | 以上说的是直接复用，但项目中更多的是自定义扩展，而且大部分情况下也需要复用torchvision的transform的compose，又该怎么办呢？这里解释下，composer是deepvac.aug模块的概念，compose是torchvision transform模块的概念，之所以这么相似纯粹是因为巧合。 142 | 143 | 要扩展自己的composer也是很简单的，比如我可以自定义一个composer（我把它命名为GemfieldComposer），这个composer可以使用/复用以下增强逻辑： 144 | - torchvision transform定义的compose； 145 | - deepvac内置的aug算子； 146 | - 我自己写的aug算子。 147 | 148 | 更详细的步骤请访问：[deepvac.aug模块使用](./docs/design.md#aug) 149 | 150 | ## 8. 编写Dataset类 151 | 代码编写在data/dataloader.py文件中。继承deepvac.datasets类体系，比如FileLineDataset类提供了对如下train.txt这种格式的封装： 152 | ```bash 153 | #train.txt，第一列为图片路径，第二列为label 154 | img0/1.jpg 0 155 | img0/2.jpg 0 156 | ... 157 | img1/0.jpg 1 158 | ... 159 | img2/0.jpg 2 160 | ... 161 | ``` 162 | 有时第二列是字符串，并且想把FileLineDataset中使用Image读取图片对方式替换为cv2，那么可以通过如下的继承方式来重新实现： 163 | ```python 164 | from deepvac.datasets import FileLineDataset 165 | 166 | class FileLineCvStrDataset(FileLineDataset): 167 | def _buildLabelFromLine(self, line): 168 | line = line.strip().split(" ") 169 | return [line[0], line[1]] 170 | 171 | def _buildSampleFromPath(self, abs_path): 172 | #we just set default loader with Pillow Image 173 | sample = cv2.imread(abs_path) 174 | sample = self.compose(sample) 175 | return sample 176 | ``` 177 | 哦，FileLineCvStrDataset也已经是deepvac.datasets中提供的类了。 178 | 179 | 180 | ## 9. 编写训练和验证脚本 181 | 在Deepvac规范中，train.py就代表了训练范式。模型训练的代码写在train.py文件中，继承DeepvacTrain类： 182 | ```python 183 | from deepvac import DeepvacTrain 184 | 185 | class MyTrain(DeepvacTrain): 186 | pass 187 | ``` 188 | 189 | 继承DeepvacTrain的子类可能需要重新实现以下方法才能够开始训练： 190 | | 类的方法（*号表示用户一般要重新实现） | 功能 | 备注 | 191 | | ---- | ---- | ---- | 192 | | preEpoch | 每轮Epoch之前的用户操作，DeepvacTrain啥也不做 | 用户可以重新定义（如果需要的话） | 193 | | preIter | 每个batch迭代之前的用户操作，DeepvacTrain啥也不做 | 用户可以重新定义（如果需要的话） | 194 | | postIter | 每个batch迭代之后的用户操作，DeepvacTrain啥也不做 | 用户可以重新定义（如果需要的话） | 195 | | postEpoch | 每轮Epoch之后的用户操作，DeepvacTrain啥也不做 | 用户可以重新定义（如果需要的话） | 196 | | doFeedData2Device | DeepvacTrain把来自dataloader的sample和target(标签)移动到device设备上 | 用户可以重新定义（如果需要的话） | 197 | | doForward | DeepvacTrain会进行网络推理，推理结果赋值给self.config.output成员 |用户可以重新定义（如果需要的话） | 198 | | doLoss | DeepvacTrain会使用self.config.output和self.config.target进行计算得到此次迭代的loss| 用户可以重新定义（如果需要的话）| 199 | | doBackward | 网络反向传播过程，DeepvacTrain会调用self.config.loss.backward()进行BP| 用户可以重新定义（如果需要的话）| 200 | | doOptimize | 网络权重更新的过程，DeepvacTrain会调用self.config.optimizer.step() | 用户可以重新定义（如果需要的话） | 201 | | doSchedule | 更新学习率的过程，DeepvacTrain会调用self.config.scheduler.step() | 用户可以重新定义（如果需要的话） | 202 | | * doValAcc | 在val模式下计算模型的acc，DeepvacTrain啥也不做 | 用户一般要重新定义，写tensorboard的时候依赖于此 | 203 | 204 | 典型的写法如下： 205 | ```python 206 | class MyTrain(DeepvacTrain): 207 | ... 208 | #因为基类不能处理list类型的标签，重写该方法 209 | def doFeedData2Device(self): 210 | self.config.target = [anno.to(self.config.device) for anno in self.config.target] 211 | self.config.sample = self.config.sample.to(self.config.device) 212 | 213 | #初始化config.core.acc 214 | def doValAcc(self): 215 | self.config.acc = your_acc 216 | LOG.logI('Test accuray: {:.4f}'.format(self.config.acc)) 217 | 218 | 219 | train = MyTrain(deepvac_config) 220 | train() 221 | ``` 222 | ## 10. 编写测试脚本 223 | 在Deepvac规范中，test.py就代表测试范式。测试代码写在test.py文件中，继承Deepvac类。 224 | 225 | 和train.py中的train/val的本质不同在于： 226 | - 舍弃train/val上下文； 227 | - 网络不再使用autograd上下文； 228 | - 不再进行loss、反向、优化等计算； 229 | - 使用Deepvac的*Report模块来进行准确度、速度方面的衡量； 230 | 231 | 继承Deepvac类的子类必须（重新）实现以下方法才能够开始测试： 232 | 233 | | 类的方法（*号表示必需重新实现） | 功能 | 备注 | 234 | | ---- | ---- | ---- | 235 | | preIter | 每个batch迭代之前的用户操作，Deepvac啥也不做 | 用户可以重新定义（如果需要的话） | 236 | | postIter| 每个batch迭代之后的用户操作，Deepvac啥也不做 | 用户可以重新定义（如果需要的话） | 237 | | doFeedData2Device | Deepvac把来自dataloader的sample和target(标签)移动到device设备上 | 用户可以重新定义（如果需要的话） | 238 | | doForward | Deepvac会进行网络推理，推理结果赋值给self.config.output成员 |用户可以重新定义（如果需要的话） | 239 | | doTest | 用户完全自定义的test逻辑，可以通过report.add(gt, pred)添加测试结果，生成报告 | 看下面的测试逻辑 | 240 | 241 | 典型的写法如下： 242 | ```python 243 | class MyTest(Deepvac): 244 | ... 245 | def doTest(self): 246 | ... 247 | 248 | test = MyTest(deepvac_config) 249 | test() 250 | #test(input_tensor) 251 | ``` 252 | 253 | 当执行test()的时候，DeepVAC框架会按照如下的优先级进行测试： 254 | - 如果用户传递了参数，比如test(input_tensor)，则将针对该input_tensor进行doFeedData2Device + doForward，然后测试结束； 255 | - 如果用户重写了doTest()函数，则将执行doTest()，然后测试结束； 256 | - 如果用户配置了config.my_test.test_loader，则将迭代该loader，对每个sample进行doFeedData2Device + doForward，然后测试结束； 257 | - 以上都不符合，报错退出。 258 | 259 | 260 | # DeepVAC的社区产品 261 | | 产品名称 | 简介 |当前版本 | 获取方式/部署形式 | 262 | | ---- | ---- | ---- |---- | 263 | |[DeepVAC](https://github.com/deepvac/deepvac)|独树一帜的PyTorch工程规范 | 0.6.0 | pip install deepvac | 264 | |[libdeepvac](https://github.com/deepvac/libdeepvac) | 独树一帜的PyTorch模型部署框架 | 1.9.0 | SDK，下载 & 解压| 265 | |[MLab HomePod](https://github.com/DeepVAC/MLab#mlab-homepod)| 迄今为止最先进的容器化PyTorch模型训练环境 | 2.0 | docker run / k8s| 266 | |MLab RookPod| 迄今为止最先进的成本10万人民币以下的存储解决方案 | NA | 硬件规范 + k8s yaml| 267 | |[pyRBAC](https://github.com/DeepVAC/pyRBAC) | 基于Keycloak的RBAC python实现 | NA | pip install(敬请期待) | 268 | |DeepVAC版PyTorch | 为MLab HomePod pro版本定制的PyTorch包 |1.9.0 | conda install -c gemfield pytorch | 269 | |[DeepVAC版LibTorch](https://github.com/deepvac/libdeepvac)| 为libdeepvac定制的LibTorch库 | 1.9.0 | 压缩包，下载 & 解压| 270 | -------------------------------------------------------------------------------- /deepvac/backbones/regnet.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn as nn 4 | from .weights_init import initWeightsKaiming 5 | #function begin 6 | 7 | def adjust_block_compatibility(width_per_stage, bot_muls_per_stage, group_w_per_stage): 8 | """Adjusts the compatibility of widths, bottlenecks, and groups.""" 9 | assert len(width_per_stage) == len(bot_muls_per_stage) == len(group_w_per_stage) 10 | assert all(w > 0 and b > 0 and g > 0 for w, b, g in zip(width_per_stage, bot_muls_per_stage, group_w_per_stage)) 11 | assert all(b < 1 or b % 1 == 0 for b in bot_muls_per_stage) 12 | vs = [int(max(1, w * b)) for w, b in zip(width_per_stage, bot_muls_per_stage)] 13 | group_w_per_stage = [int(min(g, v)) for g, v in zip(group_w_per_stage, vs)] 14 | ms = [np.lcm(g, int(b)) if b > 1 else g for g, b in zip(group_w_per_stage, bot_muls_per_stage)] 15 | vs = [max(m, int(round(v / m) * m)) for v, m in zip(vs, ms)] 16 | width_per_stage = [int(v / b) for v, b in zip(vs, bot_muls_per_stage)] 17 | assert all(w * b % g == 0 for w, b, g in zip(width_per_stage, bot_muls_per_stage, group_w_per_stage)) 18 | return width_per_stage, bot_muls_per_stage, group_w_per_stage 19 | 20 | def generate_regnet(anynet_slope, anynet_initial_width, channel_controller, depth, divisor=8): 21 | """Generates per stage widths and depths from RegNet parameters.""" 22 | assert anynet_slope >= 0 and anynet_initial_width > 0 and channel_controller > 1 and anynet_initial_width % divisor == 0 23 | # Generate continuous per-block ws 24 | ws_cont = np.arange(depth) * anynet_slope + anynet_initial_width 25 | # Generate quantized per-block ws 26 | ks = np.round(np.log(ws_cont / anynet_initial_width) / np.log(channel_controller)) 27 | ws_all = anynet_initial_width * np.power(channel_controller, ks) 28 | ws_all = np.round(np.divide(ws_all, divisor)).astype(int) * divisor 29 | # Generate per stage ws and ds (assumes ws_all are sorted) 30 | ws, ds = np.unique(ws_all, return_counts=True) 31 | # Compute number of actual stages and total possible stages 32 | num_stages, total_stages = len(ws), ks.max() + 1 33 | # Convert numpy arrays to lists and return 34 | ws, ds, ws_all, ws_cont = (x.tolist() for x in (ws, ds, ws_all, ws_cont)) 35 | return ws, ds, num_stages, total_stages, ws_all, ws_cont 36 | 37 | class AnyHead(nn.Module): 38 | """AnyNet head: AvgPool, 1x1.""" 39 | 40 | def __init__(self, in_planes, class_num, head_w): 41 | super(AnyHead, self).__init__() 42 | self.head_w = head_w 43 | if head_w>0: 44 | self.conv = nn.Conv2d(in_planes, head_w, 1, stride=1, padding=0, groups=1, bias=False) 45 | self.bn = nn.BatchNorm2d(head_w, eps=1e-5, momentum=0.1) 46 | self.relu = nn.ReLU(inplace=True) 47 | in_planes = head_w 48 | self.avg_pool = nn.AdaptiveAvgPool2d((1, 1)) 49 | self.fc = nn.Linear(in_planes, class_num, bias=True) 50 | 51 | def forward(self, x): 52 | x = self.relu(self.bn(self.conv(x))) if self.head_w > 0 else x 53 | x = self.avg_pool(x) 54 | x = x.view(x.size(0), -1) 55 | x = self.fc(x) 56 | return x 57 | 58 | class SE(nn.Module): 59 | """Squeeze-and-Excitation (SE) block: AvgPool, FC, ReLU, FC, Sigmoid.""" 60 | 61 | def __init__(self, in_planes, se_planes): 62 | super(SE, self).__init__() 63 | self.avg_pool = nn.AdaptiveAvgPool2d((1, 1)) 64 | self.f_ex = nn.Sequential( 65 | nn.Conv2d(in_planes, se_planes, 1, bias=True), 66 | nn.ReLU(inplace=True), 67 | nn.Conv2d(se_planes, in_planes, 1, bias=True), 68 | nn.Sigmoid(), 69 | ) 70 | 71 | def forward(self, x): 72 | return x * self.f_ex(self.avg_pool(x)) 73 | 74 | class BottleneckTransform(nn.Module): 75 | """Bottleneck transformation: 1x1, 3x3 [+SE], 1x1.""" 76 | 77 | def __init__(self, in_planes, out_planes, stride, bottom_mul, group_width, se_ratio): 78 | super(BottleneckTransform, self).__init__() 79 | w_b = int(round(out_planes * bottom_mul)) 80 | g = w_b // group_width 81 | self.a = nn.Conv2d(in_planes, w_b, 1, stride=1, padding=0, bias=False) 82 | self.a_bn = nn.BatchNorm2d(w_b, eps=1e-5, momentum=0.1) 83 | self.a_relu = nn.ReLU(inplace=True) 84 | self.b = nn.Conv2d(w_b, w_b, 3, stride=stride, padding=1, groups=g, bias=False) 85 | self.b_bn = nn.BatchNorm2d(w_b, eps=1e-5, momentum=0.1) 86 | self.b_relu = nn.ReLU(inplace=True) 87 | if se_ratio: 88 | se_planes = int(round(in_planes * se_ratio)) 89 | self.se = SE(w_b, se_planes) 90 | self.c = nn.Conv2d(w_b, out_planes, 1, stride=1, padding=0, bias=False) 91 | self.c_bn = nn.BatchNorm2d(out_planes, eps=1e-5, momentum=0.1) 92 | self.c_bn.final_bn = True 93 | 94 | def forward(self, x): 95 | for layer in self.children(): 96 | x = layer(x) 97 | return x 98 | 99 | class ResBottleneckBlock(nn.Module): 100 | """Residual bottleneck block: x + F(x), F = bottleneck transform.""" 101 | 102 | def __init__(self, in_planes, out_planes, stride, bottom_mul=1.0, group_width=1, se_ratio=None): 103 | super(ResBottleneckBlock, self).__init__() 104 | # Use skip connection with projection if shape changes 105 | self.proj_block = (in_planes != out_planes) or (stride != 1) 106 | if self.proj_block: 107 | self.proj = nn.Conv2d(in_planes, out_planes, 1, stride=stride, padding=0, bias=False) 108 | self.bn = nn.BatchNorm2d(out_planes, eps=1e-5, momentum=0.1) 109 | self.f = BottleneckTransform(in_planes, out_planes, stride, bottom_mul, group_width, se_ratio) 110 | self.relu = nn.ReLU(True) 111 | 112 | def forward(self, x): 113 | if self.proj_block: 114 | x = self.bn(self.proj(x)) + self.f(x) 115 | else: 116 | x = x + self.f(x) 117 | x = self.relu(x) 118 | return x 119 | 120 | class SimpleStemIN(nn.Module): 121 | """Simple stem for ImageNet: 3x3, BN, ReLU.""" 122 | 123 | def __init__(self, in_planes, out_planes): 124 | super(SimpleStemIN, self).__init__() 125 | self.conv = nn.Conv2d(in_planes, out_planes, 3, stride=2, padding=1, bias=False) 126 | self.bn = nn.BatchNorm2d(out_planes, eps=1e-5, momentum=0.1) 127 | self.relu = nn.ReLU(True) 128 | 129 | def forward(self, x): 130 | for layer in self.children(): 131 | x = layer(x) 132 | return x 133 | 134 | class AnyStage(nn.Module): 135 | """AnyNet stage (sequence of blocks w/ the same output shape).""" 136 | 137 | def __init__(self, in_planes, out_planes, stride, depth, block_fun, bottom_mul, group_width, se_ratio): 138 | super(AnyStage, self).__init__() 139 | for i in range(depth): 140 | b_stride = stride if i == 0 else 1 141 | b_w_in = in_planes if i == 0 else out_planes 142 | name = "b{}".format(i + 1) 143 | self.add_module(name, block_fun(b_w_in, out_planes, b_stride, bottom_mul, group_width, se_ratio)) 144 | 145 | def forward(self, x): 146 | for block in self.children(): 147 | x = block(x) 148 | return x 149 | 150 | class RegNet(nn.Module): 151 | def __init__(self, anynet_slope, anynet_initial_width, channel_controller, depth, group_w, class_num=0, head_w=0): 152 | super(RegNet, self).__init__() 153 | stem_w = 32 154 | se_ratio = 0.25 155 | width_per_stage, depth_per_stage = generate_regnet(anynet_slope, anynet_initial_width, channel_controller, depth)[0:2] 156 | group_w_per_stage = [group_w for _ in width_per_stage] 157 | bot_muls_per_stage = [1.0 for _ in width_per_stage] 158 | stride_per_stage = self.auditConfig(width_per_stage) 159 | width_per_stage, bot_muls_per_stage, group_w_per_stage = adjust_block_compatibility(width_per_stage, bot_muls_per_stage, group_w_per_stage) 160 | stage_params = list(zip(depth_per_stage, width_per_stage, stride_per_stage, bot_muls_per_stage, group_w_per_stage)) 161 | 162 | self.stem = SimpleStemIN(3, stem_w) 163 | prev_w = stem_w 164 | for i, (depth, w, s, bottom_mul, group_width) in enumerate(stage_params): 165 | name = "s{}".format(i + 1) 166 | self.add_module(name, AnyStage(prev_w, w, s, depth, ResBottleneckBlock, bottom_mul, group_width, se_ratio)) 167 | prev_w = w 168 | 169 | if class_num > 0: 170 | self.head = AnyHead(in_planes=prev_w, class_num=class_num, head_w=head_w) 171 | 172 | initWeightsKaiming(self) 173 | 174 | def auditConfig(self, width_per_stage): 175 | stride_per_stage = [2 for _ in width_per_stage] 176 | return stride_per_stage 177 | 178 | def forward(self, x): 179 | for module in self.children(): 180 | x = module(x) 181 | return x 182 | 183 | class RegNetSmall(RegNet): 184 | def __init__(self, class_num=0, head_w=0): 185 | super(RegNetSmall, self).__init__(27.89, 48, 2.09, 16, 8, class_num, head_w) 186 | 187 | class RegNetMedium(RegNet): 188 | def __init__(self, class_num=0, head_w=0): 189 | super(RegNetMedium, self).__init__(20.71, 48, 2.65, 27, 24, class_num, head_w) 190 | 191 | class RegNetLarge(RegNet): 192 | def __init__(self, class_num=0, head_w=0): 193 | super(RegNetLarge, self).__init__(31.41, 96, 2.24, 22, 64, class_num, head_w) 194 | 195 | class RegNetOCR(RegNet): 196 | def __init__(self, class_num=0, head_w=0): 197 | super(RegNetOCR, self).__init__(27.89, 48, 2.09, 16, 8, class_num, head_w) 198 | 199 | def auditConfig(self, width_per_stage): 200 | stride_per_stage = [(2,1),(2,1),(2,1),2] 201 | return stride_per_stage 202 | 203 | if __name__ == '__main__': 204 | model = RegNetOCR() 205 | LOG.logI(model) 206 | x = torch.randn((1,3,32,320)) 207 | LOG.logI(model(x).shape) 208 | -------------------------------------------------------------------------------- /deepvac/datasets/coco.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | import cv2 4 | from ..utils import LOG 5 | from .base_dataset import DatasetBase 6 | 7 | class CocoCVDataset(DatasetBase): 8 | def __init__(self, deepvac_config, sample_path, target_path, cat2idx=None): 9 | super(CocoCVDataset, self).__init__(deepvac_config) 10 | try: 11 | from pycocotools.coco import COCO 12 | except: 13 | raise Exception("pycocotools module not found, you should try 'pip3 install pycocotools' first!") 14 | self.sample_path = sample_path 15 | self.coco = COCO(target_path) 16 | self.ids = list(sorted(self.coco.imgs.keys())) 17 | self.cats = list(sorted(self.coco.cats.keys())) 18 | self.cat2idx = cat2idx 19 | 20 | def __len__(self): 21 | return len(self.ids) 22 | 23 | def __getitem__(self, index): 24 | sample, label = self._getSample(index) 25 | sample = self.compose(sample) 26 | # post process 27 | img, label = self._buildSample(sample, label) 28 | return img, label 29 | 30 | def _getSample(self, index): 31 | img = self.loadImgs(index) 32 | category_ids, boxes, masks = self.loadAnns(index) 33 | # return target you want 34 | return img, category_ids 35 | 36 | def _buildSample(self, img, target): 37 | # BGR -> RGB && (H, W, C) -> (C, H, W) 38 | img = img[:, :, ::-1].transpose(2, 0, 1) 39 | img = np.ascontiguousarray(img) 40 | # numpy -> tensor 41 | img = torch.from_numpy(img) 42 | target = torch.from_numpy(target) 43 | return img, target 44 | 45 | def loadImgs(self, index): 46 | file_name = self.coco.loadImgs(self.ids[index])[0]['file_name'] 47 | img = cv2.imread(os.path.join(self.sample_path, file_name), 1) 48 | assert img is not None, "Image {} not found!".format(file_name) 49 | return img 50 | 51 | def loadAnns(self, index): 52 | ann_ids = self.coco.getAnnIds(imgIds=self.ids[index]) 53 | anns = self.coco.loadAnns(ann_ids) 54 | # use coco target idx or define your self target idx 55 | if self.cat2idx is None: 56 | category_ids = np.array([self.cats.index(int(i["category_id"])) for i in anns], dtype=np.float) 57 | else: 58 | category_ids = np.array([self.cat2idx[i["category_id"]] for i in anns], dtype=np.float) 59 | 60 | boxes = np.array([i["bbox"] for i in anns], dtype=np.float) 61 | 62 | if "segmentation" in anns[0].keys(): 63 | masks = np.array([self.coco.annToMask(i) for i in anns], dtype=np.float) 64 | else: 65 | masks = None 66 | return category_ids, boxes, masks 67 | 68 | 69 | class CocoCVSegDataset(DatasetBase): 70 | def __init__(self, deepvac_config, sample_path_prefix, target_path, cat2idx): 71 | super(CocoCVSegDataset, self).__init__(deepvac_config) 72 | try: 73 | from pycocotools.coco import COCO 74 | except: 75 | raise Exception("pycocotools module not found, you should try 'pip3 install pycocotools' first!") 76 | self.sample_path_prefix = sample_path_prefix 77 | self.coco = COCO(target_path) 78 | self.ids = list(sorted(self.coco.imgs.keys())) 79 | self.cats = list(sorted(self.coco.cats.keys())) 80 | self.cat2idx = cat2idx 81 | LOG.logI("Notice: 0 will be treated as background in {}!!!".format(self.name())) 82 | 83 | def auditConfig(self): 84 | self.auto_detect_subdir_with_basenum = self.addUserConfig('auto_detect_subdir_with_basenum', self.config.auto_detect_subdir_with_basenum, 0) 85 | LOG.logI("You set auto_detect_subdir_with_basenum to {}".format(self.auto_detect_subdir_with_basenum)) 86 | 87 | def __len__(self): 88 | return len(self.ids) 89 | 90 | def __getitem__(self, index): 91 | id = self.ids[index] 92 | sample, mask, cls_masks, file_path = self._getSample(id) 93 | sample, mask, cls_masks, file_path = self.compose((sample, mask, cls_masks, os.path.join(self.sample_path_prefix, file_path))) 94 | return sample, mask, cls_masks, file_path 95 | 96 | def updatePath(self, id, file_path): 97 | full_file_path = self.coco.loadImgs(id)[0]["path"] 98 | path_list = full_file_path.split('/') 99 | path_list_num = len(path_list) 100 | 101 | if path_list_num == self.auto_detect_subdir_with_basenum: 102 | withsub_file_path = file_path 103 | elif path_list_num == self.auto_detect_subdir_with_basenum + 1: 104 | withsub_file_path = path_list[-2] + '/' + file_path 105 | file_path = path_list[-2] + '_' + file_path 106 | else: 107 | LOG.logE("path list has {} fields, which should be {} or {}".format(path_list_num, self.auto_detect_subdir_with_basenum, self.auto_detect_subdir_with_basenum+1), exit=True) 108 | 109 | return withsub_file_path, file_path 110 | 111 | def _getSample(self, id: int): 112 | # img 113 | file_path = self.coco.loadImgs(id)[0]["file_name"] 114 | withsub_file_path = file_path 115 | if self.auto_detect_subdir_with_basenum > 0: 116 | withsub_file_path, file_path = self.updatePath(id, file_path) 117 | 118 | img = cv2.imread(os.path.join(self.sample_path_prefix, withsub_file_path), 1) 119 | assert img is not None, "Image {} not found in {} !".format(withsub_file_path, self.sample_path_prefix) 120 | # anno 121 | anns = self.coco.loadAnns(self.coco.getAnnIds(id)) 122 | cls_masks = {} 123 | h, w = img.shape[:2] 124 | mask = np.zeros((h, w)) 125 | bg_mask = None 126 | for idx, ann in enumerate(anns): 127 | # current category as 1, BG and others as 0. 128 | cls_mask = self.coco.annToMask(ann) 129 | cls_idx = self.cat2idx[ann["category_id"]] 130 | cls_masks[cls_idx] = cls_mask 131 | mask[cls_mask==1] = cls_idx 132 | 133 | # return target you want 134 | return img, mask, cls_masks, file_path 135 | 136 | class CocoCVContoursDataset(DatasetBase): 137 | def __init__(self, deepvac_config, sample_path_prefix, target_path): 138 | super(CocoCVContoursDataset, self).__init__(deepvac_config) 139 | try: 140 | from pycocotools.coco import COCO 141 | except: 142 | raise Exception("pycocotools module not found, you should try 'pip3 install pycocotools' first!") 143 | self.sample_path_prefix = sample_path_prefix 144 | self.coco = COCO(target_path) 145 | self.ids = list(sorted(self.coco.imgs.keys())) 146 | self.cats = list(sorted(self.coco.cats.keys())) 147 | LOG.logI("Notice: 0 will be treated as background in {}!!!".format(self.name())) 148 | 149 | def auditConfig(self): 150 | self.auto_detect_subdir_with_basenum = self.addUserConfig('auto_detect_subdir_with_basenum', self.config.auto_detect_subdir_with_basenum, 0) 151 | LOG.logI("You set auto_detect_subdir_with_basenum to {}".format(self.auto_detect_subdir_with_basenum)) 152 | 153 | def __len__(self): 154 | return len(self.ids) 155 | 156 | def __getitem__(self, index): 157 | id = self.ids[index] 158 | sample, bbox_info, file_path, show_img = self._getSample(id) 159 | sample, bbox_info, file_path, show_img = self.compose((sample, bbox_info, os.path.join(self.sample_path_prefix, file_path), show_img)) 160 | return sample, bbox_info, file_path, show_img 161 | 162 | def updatePath(self, id, file_path): 163 | full_file_path = self.coco.loadImgs(id)[0]["path"] 164 | path_list = full_file_path.split('/') 165 | path_list_num = len(path_list) 166 | 167 | if path_list_num == self.auto_detect_subdir_with_basenum: 168 | withsub_file_path = file_path 169 | elif path_list_num == self.auto_detect_subdir_with_basenum + 1: 170 | withsub_file_path = path_list[-2] + '/' + file_path 171 | file_path = path_list[-2] + '_' + file_path 172 | else: 173 | LOG.logE("path list has {} fields, which should be {} or {}".format(path_list_num, self.auto_detect_subdir_with_basenum, self.auto_detect_subdir_with_basenum+1), exit=True) 174 | 175 | return withsub_file_path, file_path 176 | 177 | def _bbox_generator(self, img, anns_list): 178 | bbox_info = [] 179 | show_img = img.copy() 180 | for idx, single in enumerate(anns_list): 181 | s = '' 182 | if single['isbbox']: 183 | x, y, w, h = single['bbox'] 184 | x, y, w, h = int(x), int(y), int(w), int(h) 185 | cv2.rectangle(show_img, (x, y), (x + w, y + h), (0, 0, 255), 2) 186 | s = '{},{},{},{},{},{},{},{},ok'.format(x, y, x+w, y, x+w, y+h, x, y+h) 187 | else: 188 | seg = single['segmentation'] 189 | pts = np.array(seg[0]).reshape((-1,2)).astype(np.int32) 190 | cv2.polylines(show_img, [pts], True, (0,0,255), 2) 191 | for pt in pts.reshape((-1)): 192 | s += '{},'.format(pt) 193 | s += 'ok' 194 | bbox_info.append(s) 195 | return bbox_info, show_img 196 | 197 | def _getSample(self, id: int): 198 | # img 199 | file_path = self.coco.loadImgs(id)[0]["file_name"] 200 | withsub_file_path = file_path 201 | if self.auto_detect_subdir_with_basenum > 0: 202 | withsub_file_path, file_path = self.updatePath(id, file_path) 203 | 204 | img = cv2.imread(os.path.join(self.sample_path_prefix, file_path), 1) 205 | assert img is not None, "Image {} not found in {} !".format(file_path, self.sample_path_prefix) 206 | # anno 207 | anns = self.coco.loadAnns(self.coco.getAnnIds(id)) 208 | 209 | bbox_info, show_img = self._bbox_generator(img, anns) 210 | 211 | # return target you want 212 | return img, bbox_info, file_path, show_img 213 | -------------------------------------------------------------------------------- /deepvac/cast/base.py: -------------------------------------------------------------------------------- 1 | import subprocess 2 | import tempfile 3 | import copy 4 | import torch 5 | from torch.quantization import quantize_dynamic_jit, per_channel_dynamic_qconfig 6 | from torch.quantization import get_default_qconfig, quantize_jit 7 | from torch.quantization import default_dynamic_qconfig, float_qparams_weight_only_qconfig 8 | from torch.quantization.quantize_fx import prepare_fx, convert_fx, prepare_qat_fx 9 | import torch.nn as nn 10 | from ..core.config import AttrDict 11 | from ..utils import LOG, addUserConfig 12 | 13 | class DeepvacCast(object): 14 | def __init__(self, trainer_config, cast_config): 15 | self.trainer_config = trainer_config 16 | self.deepvac_cast_config = cast_config 17 | self.initConfig() 18 | self.proceed = False 19 | self.auditFinalConfig() 20 | 21 | def initConfig(self): 22 | if self.name() not in self.deepvac_cast_config.keys(): 23 | self.deepvac_cast_config[self.name()] = AttrDict() 24 | self.config = self.deepvac_cast_config[self.name()] 25 | 26 | def addUserConfig(self, config_name, user_give=None, developer_give=None, is_user_mandatory=False): 27 | module_name = 'config.cast.{}'.format(self.name()) 28 | return addUserConfig(module_name, config_name, user_give=user_give, developer_give=developer_give, is_user_mandatory=is_user_mandatory) 29 | 30 | def setAttr(self, k, v): 31 | self.config[k] = v 32 | 33 | def getAttr(self,k): 34 | return self.config[k] 35 | 36 | def auditConfig(self): 37 | LOG.logW("You should reimplement auditConfig in your subclass {}.".format(self.name())) 38 | 39 | def name(self): 40 | return self.__class__.__name__ 41 | 42 | def auditFinalConfig(self): 43 | if not self.auditConfig(): 44 | return 45 | if self.trainer_config.net is None: 46 | LOG.logE("You must set config.core.net in config.py", exit=True) 47 | 48 | if self.trainer_config.sample is None: 49 | LOG.logE("You must set config.core.sample, in general, this is done by Deepvac Framework.", exit=True) 50 | 51 | self.net = copy.deepcopy(self.trainer_config.ema if self.trainer_config.ema else self.trainer_config.net) 52 | if self.trainer_config.cast2cpu: 53 | self.net.cpu() 54 | self.trainer_config.sample = self.trainer_config.sample.to('cpu') 55 | 56 | self.proceed = True 57 | 58 | def process(self, cast_output_file=None): 59 | LOG.logE("You must reimplement process() in your subclass {}.".format(self.name()), exit=True) 60 | 61 | def exportOnnx(self): 62 | if self.config.onnx_input_names is None: 63 | self.config.onnx_input_names = ["input"] 64 | 65 | if self.config.onnx_output_names is None: 66 | self.config.onnx_output_names = ["output"] 67 | 68 | if not self.config.onnx_model_dir: 69 | f = tempfile.NamedTemporaryFile(suffix=".onnx", delete=False) 70 | self.config.onnx_model_dir = f.name 71 | 72 | torch.onnx.export(self.net, self.trainer_config.sample, self.config.onnx_model_dir, 73 | input_names=self.config.onnx_input_names, output_names=self.config.onnx_output_names, 74 | dynamic_axes=self.config.onnx_dynamic_ax, opset_version=self.config.onnx_version, export_params=True) 75 | LOG.logI("Pytorch model convert to ONNX model succeed, save model in {}".format(self.config.onnx_model_dir)) 76 | 77 | def onnxSimplify(self, check_n=0, perform_optimization=True, skip_fuse_bn=False, input_shapes=None, 78 | skipped_optimizers=None, skip_shape_inference=False, input_data=None, dynamic_input_shape=False): 79 | try: 80 | import onnx 81 | from onnxsim import simplify 82 | except: 83 | LOG.logE("You must install onnx and onnxsim package first if you want to convert pytorch to ncnn or tnn.", exit=True) 84 | 85 | try: 86 | model_op, check_ok = simplify(self.config.onnx_model_dir, check_n=check_n, perform_optimization=perform_optimization, skip_fuse_bn=skip_fuse_bn, skip_shape_inference=skip_shape_inference, \ 87 | input_shapes=input_shapes, skipped_optimizers=skipped_optimizers, input_data=input_data, dynamic_input_shape=dynamic_input_shape) 88 | except RuntimeError as e: 89 | LOG.logE("Simplify model faild. {}".format(e), exit=True) 90 | 91 | onnx.save(model_op, self.config.onnx_model_dir) 92 | if not check_ok: 93 | LOG.logE("Maybe something wrong when simplify the model, we can't guarantee generate model is right.") 94 | else: 95 | LOG.logI("Simplify model succeed") 96 | 97 | def runCmd(self, cmd): 98 | sub = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True) 99 | output, err = sub.communicate() 100 | return sub.poll(), output.decode("utf-8"), err.decode("utf-8") 101 | 102 | def __call__(self, cast_output_file=None): 103 | if not self.proceed: 104 | return 105 | self.process(cast_output_file) 106 | 107 | def calibrate(model, data_loader): 108 | model.eval() 109 | with torch.no_grad(): 110 | for sample, target in data_loader: 111 | model(sample) 112 | 113 | class ScriptModel(object): 114 | def __init__(self, deepvac_core_config, output_file, backend = 'fbgemm'): 115 | self.trainer_config = deepvac_core_config 116 | self.input_net = copy.deepcopy(self.trainer_config.ema if self.trainer_config.ema else self.trainer_config.net) 117 | self.input_net.to(self.trainer_config.sample.device) 118 | self.input_net.eval() 119 | self.output_file = output_file 120 | self.backend = backend 121 | self.dq_output_file = '{}.dq'.format(output_file) 122 | self.sq_output_file = '{}.sq'.format(output_file) 123 | self.d_qconfig_dict = {'': per_channel_dynamic_qconfig} 124 | self.s_qconfig_dict = {'': get_default_qconfig(self.backend) } 125 | 126 | def _freeze_jit(self, model): 127 | LOG.logE("You must reimplement _freeze_jit() in ScriptModel subclass.", exit=True) 128 | 129 | def _jit(self, model): 130 | LOG.logE("You must reimplement _jit() in ScriptModel subclass.", exit=True) 131 | 132 | def getCalibrateLoader(self): 133 | loader = self.trainer_config.test_loader if self.trainer_config.is_forward_only else self.trainer_config.val_loader 134 | if loader is None: 135 | LOG.logE("You enabled config.static_quantize_dir, but didn't provide config.test_loader in forward-only mode, or self.val_loader in train mode.", exit=True) 136 | return loader 137 | 138 | def save(self): 139 | with torch.no_grad(): 140 | torch.jit.save(self._freeze_jit(self.input_net), self.output_file) 141 | 142 | def saveDQ(self): 143 | LOG.logI("Pytorch model dynamic quantize starting, will save model in {}".format(self.dq_output_file)) 144 | with torch.no_grad(): 145 | quantized_model = quantize_dynamic_jit(self._jit(), self.d_qconfig_dict) 146 | torch.jit.save(quantized_model, self.dq_output_file) 147 | LOG.logI("Pytorch model dynamic quantize succeeded, saved model in {}".format(self.dq_output_file)) 148 | 149 | def saveSQ(self): 150 | loader = self.getCalibrateLoader() 151 | LOG.logI("Pytorch model static quantize starting, will save model in {}".format(self.sq_output_file)) 152 | with torch.no_grad(): 153 | quantized_model = quantize_jit(self._jit(), self.s_qconfig_dict, calibrate, [loader], inplace=False, debug=False) 154 | torch.jit.save(quantized_model, self.sq_output_file) 155 | LOG.logI("Pytorch model static quantize succeeded, saved model in {}".format(self.sq_output_file)) 156 | 157 | class SaveModelByTrace(ScriptModel): 158 | def _freeze_jit(self, model): 159 | return torch.jit.freeze(torch.jit.trace(model, self.trainer_config.sample).eval() ) 160 | 161 | def _jit(self, model): 162 | return torch.jit.trace(model, self.trainer_config.sample).eval() 163 | 164 | class SaveModelByScript(ScriptModel): 165 | def _freeze_jit(self, model): 166 | return torch.jit.freeze(torch.jit.script(model).eval() ) 167 | 168 | def _jit(self, model): 169 | return torch.jit.script(model).eval() 170 | 171 | class FXQuantize(ScriptModel): 172 | def saveDQ(self): 173 | qconfig_dict = { 174 | "object_type": [ 175 | (nn.Embedding, float_qparams_weight_only_qconfig), 176 | (nn.LSTM, default_dynamic_qconfig), 177 | (nn.Linear, default_dynamic_qconfig) 178 | ] 179 | } 180 | prepared_model = prepare_fx(self.input_net, qconfig_dict) 181 | quantized_model = convert_fx(prepared_model) 182 | torch.jit.save(self._freeze_jit(quantized_model), self.dq_output_file) 183 | 184 | def saveSQ(self): 185 | loader = self.getCalibrateLoader() 186 | prepared_model = prepare_fx(self.input_net, self.s_qconfig_dict) 187 | #print(prepared_model.graph) 188 | calibrate(prepared_model, loader) 189 | quantized_model = convert_fx(prepared_model) 190 | torch.jit.save(self._freeze_jit(quantized_model), self.sq_output_file) 191 | 192 | class EagerQuantize(ScriptModel): 193 | def saveDQ(self): 194 | quantized_model = torch.quantization.quantize_dynamic(self.input_net, inplace=False) 195 | torch.jit.save(self._freeze_jit(quantized_model), self.dq_output_file) 196 | 197 | def saveSQ(self): 198 | loader = self.getCalibrateLoader() 199 | fused_model = auto_fuse_model(self.input_net) 200 | fused_model.qconfig = torch.quantization.get_default_qconfig(self.backend) 201 | prepared_model = torch.quantization.prepare(fused_model, inplace=False) 202 | calibrate(prepared_model, loader) 203 | quantized_model = torch.quantization.convert(prepared_model, inplace=False) 204 | torch.jit.save(self._freeze_jit(quantized_model), self.sq_output_file) 205 | 206 | class FXQuantizeAndScript(FXQuantize, SaveModelByScript): 207 | pass 208 | 209 | class FXQuantizeAndTrace(FXQuantize, SaveModelByTrace): 210 | pass 211 | 212 | class EagerQuantizeAndScript(EagerQuantize, SaveModelByScript): 213 | pass 214 | 215 | class EagerQuantizeAndTrace(EagerQuantize, SaveModelByTrace): 216 | pass 217 | --------------------------------------------------------------------------------