diff --git a/modelscope/metainfo.py b/modelscope/metainfo.py index 95a381b9..5956e596 100644 --- a/modelscope/metainfo.py +++ b/modelscope/metainfo.py @@ -92,6 +92,7 @@ class Models(object): msrresnet_lite = 'msrresnet-lite' object_detection_3d = 'object_detection_3d' ddpm = 'ddpm' + ocr_recognition = 'OCRRecognition' image_quality_assessment_mos = 'image-quality-assessment-mos' nerf_recon_acc = 'nerf-recon-acc' bts_depth_estimation = 'bts-depth-estimation' @@ -837,6 +838,7 @@ class Preprocessors(object): object_detection_scrfd = 'object-detection-scrfd' image_sky_change_preprocessor = 'image-sky-change-preprocessor' image_demoire_preprocessor = 'image-demoire-preprocessor' + ocr_recognition = 'ocr-recognition' nerf_recon_acc_preprocessor = 'nerf-recon-acc-preprocessor' # nlp preprocessor diff --git a/modelscope/models/cv/ocr_recognition/__init__.py b/modelscope/models/cv/ocr_recognition/__init__.py new file mode 100644 index 00000000..39c09384 --- /dev/null +++ b/modelscope/models/cv/ocr_recognition/__init__.py @@ -0,0 +1,22 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +from typing import TYPE_CHECKING + +from modelscope.utils.import_utils import LazyImportModule + +if TYPE_CHECKING: + from .model import OCRRecognition + +else: + _import_structure = { + 'model': ['OCRRecognition'], + } + + import sys + + sys.modules[__name__] = LazyImportModule( + __name__, + globals()['__file__'], + _import_structure, + module_spec=__spec__, + extra_objects={}, + ) diff --git a/modelscope/models/cv/ocr_recognition/model.py b/modelscope/models/cv/ocr_recognition/model.py new file mode 100644 index 00000000..7d76f8e8 --- /dev/null +++ b/modelscope/models/cv/ocr_recognition/model.py @@ -0,0 +1,109 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import os + +import torch +import torch.nn.functional as F + +from modelscope.metainfo import Models +from modelscope.models.base.base_torch_model import TorchModel +from modelscope.models.builder import MODELS +from modelscope.utils.config import Config +from modelscope.utils.constant import ModelFile, Tasks +from modelscope.utils.logger import get_logger +from .modules.convnextvit import ConvNextViT +from .modules.crnn import CRNN + +LOGGER = get_logger() + + +@MODELS.register_module( + Tasks.ocr_recognition, module_name=Models.ocr_recognition) +class OCRRecognition(TorchModel): + + def __init__(self, model_dir: str, **kwargs): + """initialize the ocr recognition model from the `model_dir` path. + + Args: + model_dir (str): the model path. + """ + super().__init__(model_dir, **kwargs) + + model_path = os.path.join(model_dir, ModelFile.TORCH_MODEL_FILE) + cfgs = Config.from_file( + os.path.join(model_dir, ModelFile.CONFIGURATION)) + self.do_chunking = cfgs.model.inference_kwargs.do_chunking + self.recognizer = None + if cfgs.model.recognizer == 'ConvNextViT': + self.recognizer = ConvNextViT() + elif cfgs.model.recognizer == 'CRNN': + self.recognizer = CRNN() + else: + raise TypeError( + f'recognizer should be either ConvNextViT, CRNN, but got {cfgs.model.recognizer}' + ) + if model_path != '': + self.recognizer.load_state_dict( + torch.load(model_path, map_location='cpu')) + + dict_path = os.path.join(model_dir, ModelFile.VOCAB_FILE) + self.labelMapping = dict() + with open(dict_path, 'r', encoding='utf-8') as f: + lines = f.readlines() + cnt = 1 + for line in lines: + line = line.strip('\n') + self.labelMapping[cnt] = line + cnt += 1 + + def forward(self, inputs): + """ + Args: + img (`torch.Tensor`): batched image tensor, + shape of each tensor is [N, 1, H, W]. + + Return: + `probs [T, N, Classes] of the sequence feature` + """ + return self.recognizer(inputs) + + def postprocess(self, inputs): + # naive decoder + if self.do_chunking: + preds = inputs + batchSize, length = preds.shape + PRED_LENTH = 75 + PRED_PAD = 6 + pred_idx = [] + if batchSize == 1: + pred_idx = preds[0].cpu().data.tolist() + else: + for idx in range(batchSize): + if idx == 0: + pred_idx.extend( + preds[idx].cpu().data[:PRED_LENTH + - PRED_PAD].tolist()) + elif idx == batchSize - 1: + pred_idx.extend( + preds[idx].cpu().data[PRED_PAD:].tolist()) + else: + pred_idx.extend( + preds[idx].cpu().data[PRED_PAD:PRED_LENTH + - PRED_PAD].tolist()) + pred_idx = [its - 1 for its in pred_idx if its > 0] + else: + outprobs = inputs + outprobs = F.softmax(outprobs, dim=-1) + preds = torch.argmax(outprobs, -1) + length, batchSize = preds.shape + assert batchSize == 1, 'only support onesample inference' + pred_idx = preds[:, 0].cpu().data.tolist() + + pred_idx = pred_idx + last_p = 0 + str_pred = [] + for p in pred_idx: + if p != last_p and p != 0: + str_pred.append(self.labelMapping[p]) + last_p = p + final_str = ''.join(str_pred) + return final_str diff --git a/modelscope/models/cv/ocr_recognition/modules/__init__.py b/modelscope/models/cv/ocr_recognition/modules/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/modelscope/models/cv/ocr_recognition/modules/convnext.py b/modelscope/models/cv/ocr_recognition/modules/convnext.py new file mode 100644 index 00000000..c0e30616 --- /dev/null +++ b/modelscope/models/cv/ocr_recognition/modules/convnext.py @@ -0,0 +1,163 @@ +# Part of the implementation is borrowed and modified from ConvNext, +# publicly available at https://github.com/facebookresearch/ConvNeXt +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .timm_tinyc import DropPath + + +class Block(nn.Module): + r""" ConvNeXt Block. There are two equivalent implementations: + (1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W) + (2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back + We use (2) as we find it slightly faster in PyTorch + + Args: + dim (int): Number of input channels. + drop_path (float): Stochastic depth rate. Default: 0.0 + layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6. + """ + + def __init__(self, dim, drop_path=0., layer_scale_init_value=1e-6): + super().__init__() + self.dwconv = nn.Conv2d( + dim, dim, kernel_size=7, padding=3, groups=dim) # depthwise conv + self.norm = LayerNorm(dim, eps=1e-6) + self.pwconv1 = nn.Linear( + dim, + 4 * dim) # pointwise/1x1 convs, implemented with linear layers + self.act = nn.GELU() + self.pwconv2 = nn.Linear(4 * dim, dim) + self.gamma = nn.Parameter( + layer_scale_init_value * torch.ones((dim)), + requires_grad=True) if layer_scale_init_value > 0 else None + self.drop_path = DropPath( + drop_path) if drop_path > 0. else nn.Identity() + + def forward(self, x): + input = x + x = self.dwconv(x) + x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C) + x = self.norm(x) + x = self.pwconv1(x) + x = self.act(x) + x = self.pwconv2(x) + if self.gamma is not None: + x = self.gamma * x + x = x.permute(0, 3, 1, 2) # (N, H, W, C) -> (N, C, H, W) + + x = input + self.drop_path(x) + return x + + +class ConvNeXt(nn.Module): + r""" ConvNeXt + A PyTorch impl of : `A ConvNet for the 2020s` - + https://arxiv.org/pdf/2201.03545.pdf + + Args: + in_chans (int): Number of input image channels. Default: 3 + num_classes (int): Number of classes for classification head. Default: 1000 + depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3] + dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768] + drop_path_rate (float): Stochastic depth rate. Default: 0. + layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6. + head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1. + """ + + def __init__( + self, + in_chans=1, + num_classes=1000, + depths=[3, 3, 9, 3], + dims=[96, 192, 384, 768], + drop_path_rate=0., + layer_scale_init_value=1e-6, + head_init_scale=1., + ): + super().__init__() + + self.downsample_layers = nn.ModuleList( + ) # stem and 3 intermediate downsampling conv layers + stem = nn.Sequential( + nn.Conv2d(in_chans, dims[0], kernel_size=4, stride=4), + LayerNorm(dims[0], eps=1e-6, data_format='channels_first')) + self.downsample_layers.append(stem) + for i in range(3): + downsample_layer = nn.Sequential( + LayerNorm(dims[i], eps=1e-6, data_format='channels_first'), + nn.Conv2d( + dims[i], dims[i + 1], kernel_size=(2, 1), stride=(2, 1)), + ) + self.downsample_layers.append(downsample_layer) + + self.stages = nn.ModuleList( + ) # 4 feature resolution stages, each consisting of multiple residual blocks + dp_rates = [ + x.item() for x in torch.linspace(0, drop_path_rate, sum(depths)) + ] + cur = 0 + for i in range(4): + stage = nn.Sequential(*[ + Block( + dim=dims[i], + drop_path=dp_rates[cur + j], + layer_scale_init_value=layer_scale_init_value) + for j in range(depths[i]) + ]) + self.stages.append(stage) + cur += depths[i] + + def _init_weights(self, m): + if isinstance(m, (nn.Conv2d, nn.Linear)): + trunc_normal_(m.weight, std=.02) + nn.init.constant_(m.bias, 0) + + def forward_features(self, x): + for i in range(4): + x = self.downsample_layers[i](x.contiguous()) + x = self.stages[i](x.contiguous()) + return x # global average pooling, (N, C, H, W) -> (N, C) + + def forward(self, x): + x = self.forward_features(x.contiguous()) + + return x.contiguous() + + +class LayerNorm(nn.Module): + r""" LayerNorm that supports two data formats: channels_last (default) or channels_first. + The ordering of the dimensions in the inputs. channels_last corresponds to inputs with + shape (batch_size, height, width, channels) while channels_first corresponds to inputs + with shape (batch_size, channels, height, width). + """ + + def __init__(self, + normalized_shape, + eps=1e-6, + data_format='channels_last'): + super().__init__() + self.weight = nn.Parameter(torch.ones(normalized_shape)) + self.bias = nn.Parameter(torch.zeros(normalized_shape)) + self.eps = eps + self.data_format = data_format + if self.data_format not in ['channels_last', 'channels_first']: + raise NotImplementedError + self.normalized_shape = (normalized_shape, ) + + def forward(self, x): + if self.data_format == 'channels_last': + return F.layer_norm(x, self.normalized_shape, self.weight, + self.bias, self.eps) + elif self.data_format == 'channels_first': + u = x.mean(1, keepdim=True) + s = (x - u).pow(2).mean(1, keepdim=True) + x = (x - u) / torch.sqrt(s + self.eps) + x = self.weight[:, None, None] * x + self.bias[:, None, None] + return x + + +def convnext_tiny(): + model = ConvNeXt(depths=[3, 3, 8, 3], dims=[96, 192, 256, 512]) + return model diff --git a/modelscope/models/cv/ocr_recognition/modules/convnextvit.py b/modelscope/models/cv/ocr_recognition/modules/convnextvit.py new file mode 100644 index 00000000..aaedb697 --- /dev/null +++ b/modelscope/models/cv/ocr_recognition/modules/convnextvit.py @@ -0,0 +1,23 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import torch +import torch.nn as nn + +from .convnext import convnext_tiny +from .vitstr import vitstr_tiny + + +class ConvNextViT(nn.Module): + + def __init__(self): + super(ConvNextViT, self).__init__() + self.cnn_model = convnext_tiny() + self.vitstr = vitstr_tiny(num_tokens=7644) + + def forward(self, input): + """ Transformation stage """ + features = self.cnn_model(input) + prediction = self.vitstr(features) + prediction = torch.nn.functional.softmax(prediction, dim=-1) + + output = torch.argmax(prediction, -1) + return output diff --git a/modelscope/models/cv/ocr_recognition/modules/crnn.py b/modelscope/models/cv/ocr_recognition/modules/crnn.py new file mode 100644 index 00000000..e0e489e9 --- /dev/null +++ b/modelscope/models/cv/ocr_recognition/modules/crnn.py @@ -0,0 +1,99 @@ +# Part of the implementation is borrowed and modified from CRNN, +# publicly available at https://github.com/meijieru/crnn.pytorch +# paper linking at https://arxiv.org/pdf/1507.05717.pdf +import torch +import torch.nn as nn + + +class BidirectionalLSTM(nn.Module): + + def __init__(self, nIn, nHidden, nOut): + super(BidirectionalLSTM, self).__init__() + self.rnn = nn.LSTM(nIn, nHidden, bidirectional=True) + self.embedding = nn.Linear(nHidden * 2, nOut) + + def forward(self, input): + recurrent, _ = self.rnn(input) + T, b, h = recurrent.size() + t_rec = recurrent.view(T * b, h) + + output = self.embedding(t_rec) # [T * b, nOut] + output = output.view(T, b, -1) + return output + + +class CRNN(nn.Module): + + def __init__(self): + super(CRNN, self).__init__() + self.conv0 = nn.Sequential( + nn.Conv2d( + 1, 64, kernel_size=(3, 3), padding=(1, 1), stride=(1, 1)), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True), + ) + self.p0 = nn.MaxPool2d( + kernel_size=(2, 2), stride=(2, 2), padding=(0, 0)) + self.conv1 = nn.Sequential( + nn.Conv2d( + 64, 128, kernel_size=(3, 3), padding=(1, 1), stride=(1, 1)), + nn.BatchNorm2d(128), + nn.ReLU(inplace=True), + ) + self.p1 = nn.MaxPool2d( + kernel_size=(2, 2), stride=(2, 2), padding=(0, 0)) + self.conv2 = nn.Sequential( + nn.Conv2d( + 128, 256, kernel_size=(3, 3), padding=(1, 1), stride=(1, 1)), + nn.BatchNorm2d(256), + nn.ReLU(inplace=True), + nn.Conv2d( + 256, 256, kernel_size=(3, 3), padding=(1, 1), stride=(1, 1)), + nn.BatchNorm2d(256), + nn.ReLU(inplace=True), + ) + self.p2 = nn.MaxPool2d( + kernel_size=(2, 1), stride=(2, 1), padding=(0, 0)) + self.conv3 = nn.Sequential( + nn.Conv2d( + 256, 512, kernel_size=(3, 3), padding=(1, 1), stride=(1, 1)), + nn.BatchNorm2d(512), + nn.ReLU(inplace=True), + nn.Conv2d( + 512, 512, kernel_size=(3, 3), padding=(1, 1), stride=(1, 1)), + nn.BatchNorm2d(512), + nn.ReLU(inplace=True), + ) + self.p3 = nn.MaxPool2d( + kernel_size=(2, 1), stride=(2, 1), padding=(0, 0)) + self.conv4 = nn.Sequential( + nn.Conv2d( + 512, 512, kernel_size=(2, 1), padding=(0, 0), stride=(2, 1)), + nn.BatchNorm2d(512), + nn.ReLU(inplace=True), + ) + + self.rnn = nn.Sequential( + BidirectionalLSTM(512, 256, 256), BidirectionalLSTM(256, 256, 512)) + + self.cls = nn.Linear(512, 7644, bias=False) + + def forward(self, input): + feats = self.conv0(input) + feats = self.p0(feats) + feats = self.conv1(feats) + feats = self.p1(feats) + feats = self.conv2(feats) + feats = self.p2(feats) + feats = self.conv3(feats) + feats = self.p3(feats) + convfeats = self.conv4(feats) + + b, c, h, w = convfeats.size() + assert h == 1, 'the height of conv must be 1' + convfeats = convfeats.squeeze(2) + convfeats = convfeats.permute(2, 0, 1) # [w, b, c] + + rnnfeats = self.rnn(convfeats) + output = self.cls(rnnfeats) + return output diff --git a/modelscope/models/cv/ocr_recognition/modules/timm_tinyc.py b/modelscope/models/cv/ocr_recognition/modules/timm_tinyc.py new file mode 100644 index 00000000..555b1e42 --- /dev/null +++ b/modelscope/models/cv/ocr_recognition/modules/timm_tinyc.py @@ -0,0 +1,332 @@ +# Part of the implementation is borrowed and modified from timm, +# publicly available at https://github.com/rwightman/pytorch-image-models +import collections.abc +import logging +import math +from collections import OrderedDict +from copy import deepcopy +from functools import partial +from itertools import repeat + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +def _ntuple(n): + + def parse(x): + if isinstance(x, collections.abc.Iterable): + return x + return tuple(repeat(x, n)) + + return parse + + +class PatchEmbed(nn.Module): + """ 2D Image to Patch Embedding + """ + + def __init__(self, + img_size=224, + patch_size=16, + in_chans=3, + embed_dim=768, + norm_layer=None, + flatten=True): + super().__init__() + img_size = (1, 75) + to_2tuple = _ntuple(2) + patch_size = to_2tuple(patch_size) + self.img_size = img_size + self.patch_size = patch_size + self.grid_size = (img_size[0] // patch_size[0], + img_size[1] // patch_size[1]) + self.num_patches = self.grid_size[0] * self.grid_size[1] + self.flatten = flatten + + self.proj = nn.Conv2d( + in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) + self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() + + def forward(self, x): + B, C, H, W = x.shape + assert H == self.img_size[0] and W == self.img_size[1], \ + f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})." + x = self.proj(x) + x = x.permute(0, 1, 3, 2) + if self.flatten: + x = x.flatten(2).transpose(1, 2) # BCHW -> BNC + x = self.norm(x) + return x + + +class Mlp(nn.Module): + """ MLP as used in Vision Transformer, MLP-Mixer and related networks + """ + + def __init__(self, + in_features, + hidden_features=None, + out_features=None, + act_layer=nn.GELU, + drop=0.): + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features) + self.act = act_layer() + self.fc2 = nn.Linear(hidden_features, out_features) + self.drop = nn.Dropout(drop) + + def forward(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.drop(x) + x = self.fc2(x) + x = self.drop(x) + return x + + +def drop_path(x, drop_prob: float = 0., training: bool = False): + """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). + + This is the same as the DropConnect impl I created for EfficientNet, etc networks, however, + the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper... + See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for + changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use + 'survival rate' as the argument. + + """ + if drop_prob == 0. or not training: + return x + keep_prob = 1 - drop_prob + shape = (x.shape[0], ) + (1, ) * ( + x.ndim - 1) # work with diff dim tensors, not just 2D ConvNets + random_tensor = keep_prob + torch.rand( + shape, dtype=x.dtype, device=x.device) + random_tensor.floor_() # binarize + output = x.div(keep_prob) * random_tensor + return output + + +class DropPath(nn.Module): + """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). + """ + + def __init__(self, drop_prob=None): + super(DropPath, self).__init__() + self.drop_prob = drop_prob + + def forward(self, x): + return drop_path(x, self.drop_prob, self.training) + + +class Attention(nn.Module): + + def __init__(self, + dim, + num_heads=8, + qkv_bias=False, + attn_drop=0.1, + proj_drop=0.1): + super().__init__() + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = head_dim**-0.5 + + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + + def forward(self, x): + B, N, C = x.shape + qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, + C // self.num_heads).permute(2, 0, 3, 1, 4) + q, k, v = qkv[0], qkv[1], qkv[ + 2] # make torchscript happy (cannot use tensor as tuple) + + attn = (q @ k.transpose(-2, -1)) * self.scale + attn = attn.softmax(dim=-1) + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x + + +class Block(nn.Module): + + def __init__(self, + dim, + num_heads, + mlp_ratio=4., + qkv_bias=False, + drop=0., + attn_drop=0., + drop_path=0., + act_layer=nn.GELU, + norm_layer=nn.LayerNorm): + super().__init__() + self.norm1 = norm_layer(dim) + self.attn = Attention( + dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + attn_drop=attn_drop, + proj_drop=drop) + # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here + self.drop_path = DropPath( + drop_path) if drop_path > 0. else nn.Identity() + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = Mlp( + in_features=dim, + hidden_features=mlp_hidden_dim, + act_layer=act_layer, + drop=drop) + + def forward(self, x): + x = x + self.drop_path(self.attn(self.norm1(x))) + x = x + self.drop_path(self.mlp(self.norm2(x))) + return x + + +class VisionTransformer(nn.Module): + + def __init__(self, + img_size=224, + patch_size=16, + in_chans=3, + num_classes=1000, + embed_dim=768, + depth=12, + num_heads=12, + mlp_ratio=4., + qkv_bias=True, + representation_size=None, + distilled=False, + drop_rate=0.1, + attn_drop_rate=0.1, + drop_path_rate=0., + embed_layer=PatchEmbed, + norm_layer=None, + act_layer=None, + weight_init=''): + """ + Args: + img_size (int, tuple): input image size + patch_size (int, tuple): patch size + in_chans (int): number of input channels + num_classes (int): number of classes for classification head + embed_dim (int): embedding dimension + depth (int): depth of transformer + num_heads (int): number of attention heads + mlp_ratio (int): ratio of mlp hidden dim to embedding dim + qkv_bias (bool): enable bias for qkv if True + representation_size (Optional[int]): enable and set representation layer (pre-logits) to this value if set + distilled (bool): model includes a distillation token and head as in DeiT models + drop_rate (float): dropout rate + attn_drop_rate (float): attention dropout rate + drop_path_rate (float): stochastic depth rate + embed_layer (nn.Module): patch embedding layer + norm_layer: (nn.Module): normalization layer + weight_init: (str): weight init scheme + """ + super().__init__() + self.num_classes = num_classes + self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models + self.num_tokens = 2 if distilled else 1 + norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6) + act_layer = act_layer or nn.GELU + + self.patch_embed = embed_layer( + img_size=img_size, + patch_size=patch_size, + in_chans=in_chans, + embed_dim=embed_dim) + num_patches = self.patch_embed.num_patches + + self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) + self.dist_token = nn.Parameter(torch.zeros( + 1, 1, embed_dim)) if distilled else None + self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim)) + self.pos_drop = nn.Dropout(p=drop_rate) + + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth) + ] # stochastic depth decay rule + self.blocks = nn.Sequential(*[ + Block( + dim=embed_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[i], + norm_layer=norm_layer, + act_layer=act_layer) for i in range(depth) + ]) + self.norm = norm_layer(embed_dim) + + # Representation layer + if representation_size and not distilled: + self.num_features = representation_size + self.pre_logits = nn.Sequential( + OrderedDict([('fc', nn.Linear(embed_dim, representation_size)), + ('act', nn.Tanh())])) + else: + self.pre_logits = nn.Identity() + + # Classifier head(s) + self.head = nn.Linear( + self.num_features, + num_classes) if num_classes > 0 else nn.Identity() + self.head_dist = None + if distilled: + self.head_dist = nn.Linear( + self.embed_dim, + self.num_classes) if num_classes > 0 else nn.Identity() + + def reset_classifier(self, num_classes, global_pool=''): + self.num_classes = num_classes + self.head = nn.Linear( + self.embed_dim, num_classes) if num_classes > 0 else nn.Identity() + if self.num_tokens == 2: + self.head_dist = nn.Linear( + self.embed_dim, + self.num_classes) if num_classes > 0 else nn.Identity() + + def forward_features(self, x): + x = self.patch_embed(x) + cls_token = self.cls_token.expand( + x.shape[0], -1, -1) # stole cls_tokens impl from Phil Wang, thanks + if self.dist_token is None: + x = torch.cat((cls_token, x), dim=1) + else: + x = torch.cat( + (cls_token, self.dist_token.expand(x.shape[0], -1, -1), x), + dim=1) + x = self.pos_drop(x + self.pos_embed) + x = self.blocks(x) + x = self.norm(x) + if self.dist_token is None: + return self.pre_logits(x[:, 0]) + else: + return x[:, 0], x[:, 1] + + def forward(self, x): + x = self.forward_features(x) + if self.head_dist is not None: + x, x_dist = self.head(x[0]), self.head_dist( + x[1]) # x must be a tuple + if self.training and not torch.jit.is_scripting(): + # during inference, return the average of both classifier predictions + return x, x_dist + else: + return (x + x_dist) / 2 + else: + x = self.head(x) + return x diff --git a/modelscope/models/cv/ocr_recognition/modules/vitstr.py b/modelscope/models/cv/ocr_recognition/modules/vitstr.py new file mode 100644 index 00000000..5ce3aeca --- /dev/null +++ b/modelscope/models/cv/ocr_recognition/modules/vitstr.py @@ -0,0 +1,58 @@ +# Part of the implementation is borrowed and modified from ViTSTR, +# publicly available at https://github.com/roatienza/deep-text-recognition-benchmark +from __future__ import absolute_import, division, print_function +import logging +from copy import deepcopy +from functools import partial + +import torch +import torch.nn as nn +import torch.utils.model_zoo as model_zoo + +from .timm_tinyc import VisionTransformer + + +class ViTSTR(VisionTransformer): + ''' + ViTSTR is basically a ViT that uses DeiT weights. + Modified head to support a sequence of characters prediction for STR. + ''' + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def reset_classifier(self, num_classes): + self.num_classes = num_classes + self.head = nn.Linear( + self.embed_dim, num_classes) if num_classes > 0 else nn.Identity() + + def forward_features(self, x): + x = self.patch_embed(x) + + x = x + self.pos_embed + x = self.pos_drop(x) + for blk in self.blocks: + x = blk(x) + + x = self.norm(x) + return x + + def forward(self, x): + x = self.forward_features(x) + b, s, e = x.size() + x = x.reshape(b * s, e) + x = self.head(x).view(b, s, self.num_classes) + return x + + +def vitstr_tiny(num_tokens): + vitstr = ViTSTR( + patch_size=1, + in_chans=512, + embed_dim=192, + depth=12, + num_heads=3, + mlp_ratio=4, + qkv_bias=True) + vitstr.reset_classifier(num_classes=num_tokens) + return vitstr diff --git a/modelscope/models/cv/ocr_recognition/preprocessor.py b/modelscope/models/cv/ocr_recognition/preprocessor.py new file mode 100644 index 00000000..f47dcaef --- /dev/null +++ b/modelscope/models/cv/ocr_recognition/preprocessor.py @@ -0,0 +1,104 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. + +import os + +import cv2 +import numpy as np +import PIL +import torch + +from modelscope.metainfo import Preprocessors +from modelscope.preprocessors import Preprocessor, load_image +from modelscope.preprocessors.builder import PREPROCESSORS +from modelscope.utils.config import Config +from modelscope.utils.constant import Fields, ModeKeys, ModelFile + + +@PREPROCESSORS.register_module( + Fields.cv, module_name=Preprocessors.ocr_recognition) +class OCRRecognitionPreprocessor(Preprocessor): + + def __init__(self, model_dir: str, mode: str = ModeKeys.INFERENCE): + """The base constructor for all ocr recognition preprocessors. + + Args: + model_dir (str): model directory to initialize some resource + mode: The mode for the preprocessor. + """ + super().__init__(mode) + cfgs = Config.from_file( + os.path.join(model_dir, ModelFile.CONFIGURATION)) + self.do_chunking = cfgs.model.inference_kwargs.do_chunking + self.target_height = cfgs.model.inference_kwargs.img_height + self.target_width = cfgs.model.inference_kwargs.img_width + + def keepratio_resize(self, img): + cur_ratio = img.shape[1] / float(img.shape[0]) + mask_height = self.target_height + mask_width = self.target_width + if cur_ratio > float(self.target_width) / self.target_height: + cur_target_height = self.target_height + cur_target_width = self.target_width + else: + cur_target_height = self.target_height + cur_target_width = int(self.target_height * cur_ratio) + img = cv2.resize(img, (cur_target_width, cur_target_height)) + mask = np.zeros([mask_height, mask_width]).astype(np.uint8) + mask[:img.shape[0], :img.shape[1]] = img + img = mask + return img + + def __call__(self, inputs): + """process the raw input data + Args: + inputs: + - A string containing an HTTP link pointing to an image + - A string containing a local path to an image + - An image loaded in PIL or opencv directly + Returns: + outputs: the preprocessed image + """ + if isinstance(inputs, str): + img = np.array(load_image(inputs).convert('L')) + elif isinstance(inputs, PIL.Image.Image): + img = np.array(inputs.convert('L')) + elif isinstance(inputs, np.ndarray): + if len(inputs.shape) == 3: + img = cv2.cvtColor(inputs, cv2.COLOR_RGB2GRAY) + else: + raise TypeError( + f'inputs should be either str, PIL.Image, np.array, but got {type(inputs)}' + ) + + if self.do_chunking: + data = [] + img_h, img_w = img.shape + wh_ratio = img_w / img_h + true_w = int(self.target_height * wh_ratio) + split_batch_cnt = 1 + if true_w < self.target_width * 1.2: + img = cv2.resize( + img, (min(true_w, self.target_width), self.target_height)) + else: + split_batch_cnt = math.ceil((true_w - 48) * 1.0 / 252) + img = cv2.resize(img, (true_w, self.target_height)) + + if split_batch_cnt == 1: + mask = np.zeros((self.target_height, self.target_width)) + mask[:, :img.shape[1]] = img + data.append(mask) + else: + for idx in range(split_batch_cnt): + mask = np.zeros((self.target_height, self.target_width)) + left = (PRED_LENTH * 4 - PRED_PAD * 4) * idx + trunk_img = img[:, left:min(left + PRED_LENTH * 4, true_w)] + mask[:, :trunk_img.shape[1]] = trunk_img + data.append(mask) + + data = torch.FloatTensor(data).view( + len(data), 1, self.target_height, self.target_width) / 255. + else: + data = self.keepratio_resize(img) + data = torch.FloatTensor(data).view(1, 1, self.target_height, + self.target_width) / 255. + return data diff --git a/modelscope/pipelines/cv/ocr_recognition_pipeline.py b/modelscope/pipelines/cv/ocr_recognition_pipeline.py index d90f8db6..f5b2f667 100644 --- a/modelscope/pipelines/cv/ocr_recognition_pipeline.py +++ b/modelscope/pipelines/cv/ocr_recognition_pipeline.py @@ -1,133 +1,74 @@ # Copyright (c) Alibaba, Inc. and its affiliates. -import math -import os.path as osp -from typing import Any, Dict - -import cv2 -import numpy as np -import PIL -import torch - from modelscope.metainfo import Pipelines +from modelscope.models.cv.ocr_recognition import OCRRecognition from modelscope.outputs import OutputKeys from modelscope.pipelines.base import Input, Pipeline from modelscope.pipelines.builder import PIPELINES -from modelscope.pipelines.cv.ocr_utils.model_convnext_transformer import \ - OCRRecModel -from modelscope.preprocessors import load_image -from modelscope.utils.constant import ModelFile, Tasks +from modelscope.utils.constant import Tasks from modelscope.utils.logger import get_logger logger = get_logger() -# constant -NUM_CLASSES = 7644 -IMG_HEIGHT = 32 -IMG_WIDTH = 300 -PRED_LENTH = 75 -PRED_PAD = 6 - @PIPELINES.register_module( Tasks.ocr_recognition, module_name=Pipelines.ocr_recognition) class OCRRecognitionPipeline(Pipeline): + """ OCR Recognition Pipeline. + + Example: + + ```python + >>> from modelscope.pipelines import pipeline + + >>> ocr_recognition = pipeline('ocr-recognition', 'damo/cv_crnn_ocr-recognition-general_damo') + >>> ocr_recognition("http://duguang-labelling.oss-cn-shanghai.aliyuncs.com" + "/mass_img_tmp_20220922/ocr_recognition_handwritten.jpg") + + {'text': '电子元器件提供BOM配单'} + ``` + """ def __init__(self, model: str, **kwargs): """ + use `model` to create a ocr recognition pipeline for prediction Args: - model: model id on modelscope hub. + model: model id on modelscope hub or `OCRRecognition` Model. + preprocessor: `OCRRecognitionPreprocessor`. """ + assert isinstance(model, str), 'model must be a single str' super().__init__(model=model, **kwargs) - model_path = osp.join(self.model, ModelFile.TORCH_MODEL_FILE) - label_path = osp.join(self.model, 'label_dict.txt') - logger.info(f'loading model from {model_path}') + logger.info(f'loading model from dir {model}') + self.ocr_recognizer = self.model.to(self.device) + self.ocr_recognizer.eval() + logger.info('loading model done') - self.device = torch.device( - 'cuda' if torch.cuda.is_available() else 'cpu') - self.infer_model = OCRRecModel(NUM_CLASSES).to(self.device) - self.infer_model.eval() - self.infer_model.load_state_dict( - torch.load(model_path, map_location=self.device)) - self.labelMapping = dict() - with open(label_path, 'r', encoding='utf-8') as f: - lines = f.readlines() - cnt = 2 - for line in lines: - line = line.strip('\n') - self.labelMapping[cnt] = line - cnt += 1 + def __call__(self, input, **kwargs): + """ + Recognize text sequence in the text image. - def preprocess(self, input: Input) -> Dict[str, Any]: - if isinstance(input, str): - img = np.array(load_image(input).convert('L')) - elif isinstance(input, PIL.Image.Image): - img = np.array(input.convert('L')) - elif isinstance(input, np.ndarray): - if len(input.shape) == 3: - img = cv2.cvtColor(input, cv2.COLOR_RGB2GRAY) - else: - raise TypeError(f'input should be either str, PIL.Image,' - f' np.array, but got {type(input)}') - data = [] - img_h, img_w = img.shape - wh_ratio = img_w / img_h - true_w = int(IMG_HEIGHT * wh_ratio) - split_batch_cnt = 1 - if true_w < IMG_WIDTH * 1.2: - img = cv2.resize(img, (min(true_w, IMG_WIDTH), IMG_HEIGHT)) - else: - split_batch_cnt = math.ceil((true_w - 48) * 1.0 / 252) - img = cv2.resize(img, (true_w, IMG_HEIGHT)) + Args: + input (`Image`): + The pipeline handles three types of images: - if split_batch_cnt == 1: - mask = np.zeros((IMG_HEIGHT, IMG_WIDTH)) - mask[:, :img.shape[1]] = img - data.append(mask) - else: - for idx in range(split_batch_cnt): - mask = np.zeros((IMG_HEIGHT, IMG_WIDTH)) - left = (PRED_LENTH * 4 - PRED_PAD * 4) * idx - trunk_img = img[:, left:min(left + PRED_LENTH * 4, true_w)] - mask[:, :trunk_img.shape[1]] = trunk_img - data.append(mask) + - A string containing an HTTP link pointing to an image + - A string containing a local path to an image + - An image loaded in PIL or opencv directly - data = torch.FloatTensor(data).view( - len(data), 1, IMG_HEIGHT, IMG_WIDTH) / 255. - data = data.to(self.device) + The pipeline currently supports single image input. - result = {'img': data} + Return: + A text sequence (string) of the input text image. + """ + return super().__call__(input, **kwargs) - return result + def preprocess(self, inputs): + outputs = self.preprocessor(inputs) + return outputs - def forward(self, input: Dict[str, Any]) -> Dict[str, Any]: - pred = self.infer_model(input['img']) - return {'results': pred} + def forward(self, inputs): + outputs = self.ocr_recognizer(inputs) + return outputs - def postprocess(self, inputs: Dict[str, Any]) -> Dict[str, Any]: - preds = inputs['results'] - batchSize, length = preds.shape - pred_idx = [] - if batchSize == 1: - pred_idx = preds[0].cpu().data.tolist() - else: - for idx in range(batchSize): - if idx == 0: - pred_idx.extend(preds[idx].cpu().data[:PRED_LENTH - - PRED_PAD].tolist()) - elif idx == batchSize - 1: - pred_idx.extend(preds[idx].cpu().data[PRED_PAD:].tolist()) - else: - pred_idx.extend(preds[idx].cpu().data[PRED_PAD:PRED_LENTH - - PRED_PAD].tolist()) - - # ctc decoder - last_p = 0 - str_pred = [] - for p in pred_idx: - if p != last_p and p != 0: - str_pred.append(self.labelMapping[p]) - last_p = p - - final_str = ''.join(str_pred) - result = {OutputKeys.TEXT: final_str} - return result + def postprocess(self, inputs): + outputs = {OutputKeys.TEXT: inputs} + return outputs diff --git a/tests/pipelines/test_ocr_recognition.py b/tests/pipelines/test_ocr_recognition.py index 8d48dd7a..372a4bc4 100644 --- a/tests/pipelines/test_ocr_recognition.py +++ b/tests/pipelines/test_ocr_recognition.py @@ -13,7 +13,7 @@ from modelscope.utils.test_utils import test_level class OCRRecognitionTest(unittest.TestCase, DemoCompatibilityCheck): def setUp(self) -> None: - self.model_id = 'damo/cv_convnextTiny_ocr-recognition-general_damo' + self.model_id = 'damo/cv_crnn_ocr-recognition-general_damo' self.test_image = 'data/test/images/ocr_recognition.jpg' self.task = Tasks.ocr_recognition @@ -23,18 +23,25 @@ class OCRRecognitionTest(unittest.TestCase, DemoCompatibilityCheck): @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') def test_run_with_model_from_modelhub(self): - ocr_recognition = pipeline(Tasks.ocr_recognition, model=self.model_id) + ocr_recognition = pipeline( + Tasks.ocr_recognition, + model=self.model_id, + model_revision='v1.0.0') self.pipeline_inference(ocr_recognition, self.test_image) @unittest.skipUnless(test_level() >= 1, 'skip test in current test level') def test_run_with_model_from_modelhub_PILinput(self): - ocr_recognition = pipeline(Tasks.ocr_recognition, model=self.model_id) + ocr_recognition = pipeline( + Tasks.ocr_recognition, + model=self.model_id, + model_revision='v1.0.0') imagePIL = PIL.Image.open(self.test_image) self.pipeline_inference(ocr_recognition, imagePIL) @unittest.skipUnless(test_level() >= 2, 'skip test in current test level') def test_run_modelhub_default_model(self): - ocr_recognition = pipeline(Tasks.ocr_recognition) + ocr_recognition = pipeline( + Tasks.ocr_recognition, model_revision='v2.0.0') self.pipeline_inference(ocr_recognition, self.test_image) @unittest.skip('demo compatibility test is only enabled on a needed-basis')