diff --git a/data/test/images/audrey_hepburn.jpg b/data/test/images/audrey_hepburn.jpg new file mode 100644 index 00000000..a5afb854 --- /dev/null +++ b/data/test/images/audrey_hepburn.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:da5cf2f3318e61cd38193af374b21a2dec0e90f2aa0e25b3b1825488eadbdc9d +size 97191 diff --git a/modelscope/metainfo.py b/modelscope/metainfo.py index 78d8d973..0a335cd0 100644 --- a/modelscope/metainfo.py +++ b/modelscope/metainfo.py @@ -67,6 +67,7 @@ class Models(object): real_basicvsr = 'real-basicvsr' rcp_sceneflow_estimation = 'rcp-sceneflow-estimation' image_casmvs_depth_estimation = 'image-casmvs-depth-estimation' + ddcolor = 'ddcolor' # EasyCV models yolox = 'YOLOX' @@ -267,6 +268,7 @@ class Pipelines(object): video_super_resolution = 'realbasicvsr-video-super-resolution' pointcloud_sceneflow_estimation = 'pointcloud-sceneflow-estimation' image_multi_view_depth_estimation = 'image-multi-view-depth-estimation' + ddcolor_image_colorization = 'ddcolor-image-colorization' # nlp tasks automatic_post_editing = 'automatic-post-editing' diff --git a/modelscope/models/cv/image_colorization/__init__.py b/modelscope/models/cv/image_colorization/__init__.py index 9dbb07a5..a9facc8a 100644 --- a/modelscope/models/cv/image_colorization/__init__.py +++ b/modelscope/models/cv/image_colorization/__init__.py @@ -4,13 +4,13 @@ from typing import TYPE_CHECKING from modelscope.utils.import_utils import LazyImportModule if TYPE_CHECKING: - from .unet import DynamicUnetWide, DynamicUnetDeep - from .utils import NormType + from .unet import DynamicUnetWide, DynamicUnetDeep, NormType + from .ddcolor import DDColorForImageColorization else: _import_structure = { - 'unet': ['DynamicUnetWide', 'DynamicUnetDeep'], - 'utils': ['NormType'] + 'unet': ['DynamicUnetWide', 'DynamicUnetDeep', 'NormType'], + 'ddcolor': ['DDColorForImageColorization'], } import sys diff --git a/modelscope/models/cv/image_colorization/ddcolor/__init__.py b/modelscope/models/cv/image_colorization/ddcolor/__init__.py new file mode 100644 index 00000000..a185e294 --- /dev/null +++ b/modelscope/models/cv/image_colorization/ddcolor/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +from .ddcolor_for_image_colorization import DDColorForImageColorization diff --git a/modelscope/models/cv/image_colorization/ddcolor/ddcolor.py b/modelscope/models/cv/image_colorization/ddcolor/ddcolor.py new file mode 100644 index 00000000..75ae44f2 --- /dev/null +++ b/modelscope/models/cv/image_colorization/ddcolor/ddcolor.py @@ -0,0 +1,283 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. + +import torch +import torch.nn as nn + +from .utils.convnext import ConvNeXt +from .utils.position_encoding import PositionEmbeddingSine +from .utils.transformer_utils import (MLP, CrossAttentionLayer, FFNLayer, + SelfAttentionLayer) +from .utils.unet import (CustomPixelShuffle_ICNR, Hook, NormType, + UnetBlockWide, custom_conv_layer) + + +class DDColor(nn.Module): + + def __init__(self, + encoder_name='convnext-l', + input_size=(256, 256), + num_queries=100): + + super().__init__() + + self.encoder = Encoder(encoder_name, + ['norm0', 'norm1', 'norm2', 'norm3']) + self.encoder.eval() + test_input = torch.randn(1, 3, *input_size) + self.encoder(test_input) + + self.decoder = Decoder( + self.encoder.hooks, + nf=512, + last_norm='Spectral', + num_queries=num_queries, + num_scales=3, + dec_layers=9, + ) + self.refine_net = nn.Sequential( + custom_conv_layer( + num_queries + 3, + 2, + ks=1, + use_activ=False, + norm_type=NormType.Spectral)) + + self.register_buffer( + 'mean', + torch.Tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)) + self.register_buffer( + 'std', + torch.Tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)) + + def normalize(self, img): + return (img - self.mean) / self.std + + def forward(self, img): + if img.shape[1] == 3: + img = self.normalize(img) + + self.encoder(img) + out_feat = self.decoder() + coarse_input = torch.cat([out_feat, img], dim=1) + out = self.refine_net(coarse_input) + + return out + + +class Decoder(nn.Module): + + def __init__(self, + hooks, + nf=512, + blur=True, + last_norm='Spectral', + num_queries=100, + num_scales=3, + dec_layers=9): + super().__init__() + self.hooks = hooks + self.nf = nf + self.blur = blur + self.last_norm = getattr(NormType, last_norm) + + self.layers = self.make_layers() + embed_dim = nf // 2 + + self.last_shuf = CustomPixelShuffle_ICNR( + embed_dim, + embed_dim, + blur=self.blur, + norm_type=self.last_norm, + scale=4) + + self.color_decoder = MultiScaleColorDecoder( + in_channels=[512, 512, 256], + num_queries=num_queries, + num_scales=num_scales, + dec_layers=dec_layers, + ) + + def forward(self): + encode_feat = self.hooks[-1].feature + out0 = self.layers[0](encode_feat) + out1 = self.layers[1](out0) + out2 = self.layers[2](out1) + out3 = self.last_shuf(out2) + out = self.color_decoder([out0, out1, out2], out3) + + return out + + def make_layers(self): + decoder_layers = [] + + e_in_c = self.hooks[-1].feature.shape[1] + in_c = e_in_c + + out_c = self.nf + setup_hooks = self.hooks[-2::-1] + for layer_index, hook in enumerate(setup_hooks): + feature_c = hook.feature.shape[1] + if layer_index == len(setup_hooks) - 1: + out_c = out_c // 2 + decoder_layers.append( + UnetBlockWide( + in_c, + feature_c, + out_c, + hook, + blur=self.blur, + self_attention=False, + norm_type=NormType.Spectral)) + in_c = out_c + return nn.Sequential(*decoder_layers) + + +class Encoder(nn.Module): + + def __init__(self, encoder_name, hook_names, **kwargs): + super().__init__() + if encoder_name == 'convnext-t' or encoder_name == 'convnext': + self.arch = ConvNeXt() + elif encoder_name == 'convnext-s': + self.arch = ConvNeXt( + depths=[3, 3, 27, 3], dims=[96, 192, 384, 768]) + elif encoder_name == 'convnext-b': + self.arch = ConvNeXt( + depths=[3, 3, 27, 3], dims=[128, 256, 512, 1024]) + elif encoder_name == 'convnext-l': + self.arch = ConvNeXt( + depths=[3, 3, 27, 3], dims=[192, 384, 768, 1536]) + else: + raise NotImplementedError + + self.hook_names = hook_names + self.hooks = self.setup_hooks() + + def setup_hooks(self): + hooks = [Hook(self.arch._modules[name]) for name in self.hook_names] + return hooks + + def forward(self, img): + return self.arch(img) + + +class MultiScaleColorDecoder(nn.Module): + + def __init__(self, + in_channels, + hidden_dim=256, + num_queries=100, + nheads=8, + dim_feedforward=2048, + dec_layers=9, + pre_norm=False, + color_embed_dim=256, + enforce_input_project=True, + num_scales=3): + super().__init__() + + # positional encoding + N_steps = hidden_dim // 2 + self.pe_layer = PositionEmbeddingSine(N_steps, normalize=True) + + # define Transformer decoder + self.num_heads = nheads + self.num_layers = dec_layers + self.transformer_self_attention_layers = nn.ModuleList() + self.transformer_cross_attention_layers = nn.ModuleList() + self.transformer_ffn_layers = nn.ModuleList() + + for _ in range(self.num_layers): + self.transformer_self_attention_layers.append( + SelfAttentionLayer( + d_model=hidden_dim, + nhead=nheads, + dropout=0.0, + normalize_before=pre_norm, + )) + self.transformer_cross_attention_layers.append( + CrossAttentionLayer( + d_model=hidden_dim, + nhead=nheads, + dropout=0.0, + normalize_before=pre_norm, + )) + self.transformer_ffn_layers.append( + FFNLayer( + d_model=hidden_dim, + dim_feedforward=dim_feedforward, + dropout=0.0, + normalize_before=pre_norm, + )) + + self.decoder_norm = nn.LayerNorm(hidden_dim) + + self.num_queries = num_queries + # learnable color query features + self.query_feat = nn.Embedding(num_queries, hidden_dim) + # learnable color query p.e. + self.query_embed = nn.Embedding(num_queries, hidden_dim) + + # level embedding + self.num_feature_levels = num_scales + self.level_embed = nn.Embedding(self.num_feature_levels, hidden_dim) + + # input projections + self.input_proj = nn.ModuleList() + for i in range(self.num_feature_levels): + if in_channels[i] != hidden_dim or enforce_input_project: + self.input_proj.append( + nn.Conv2d(in_channels[i], hidden_dim, kernel_size=1)) + nn.init.kaiming_uniform_(self.input_proj[-1].weight, a=1) + if self.input_proj[-1].bias is not None: + nn.init.constant_(self.input_proj[-1].bias, 0) + else: + self.input_proj.append(nn.Sequential()) + + # output FFNs + self.color_embed = MLP(hidden_dim, hidden_dim, color_embed_dim, 3) + + def forward(self, feature_pyramid, last_img_feature): + assert len(feature_pyramid) == self.num_feature_levels + src, pos = [], [] + + for i in range(self.num_feature_levels): + pos.append(self.pe_layer(feature_pyramid[i], None).flatten(2)) + src.append(self.input_proj[i](feature_pyramid[i]).flatten(2) + + self.level_embed.weight[i][None, :, None]) + + # flatten NxCxHxW to HWxNxC + pos[-1] = pos[-1].permute(2, 0, 1) + src[-1] = src[-1].permute(2, 0, 1) + + _, bs, _ = src[0].shape + + # QxNxC + query_embed = self.query_embed.weight.unsqueeze(1).repeat(1, bs, 1) + output = self.query_feat.weight.unsqueeze(1).repeat(1, bs, 1) + + for i in range(self.num_layers): + level_index = i % self.num_feature_levels + # attention: cross-attention first + output = self.transformer_cross_attention_layers[i]( + output, + src[level_index], + memory_mask=None, + memory_key_padding_mask=None, + pos=pos[level_index], + query_pos=query_embed) + output = self.transformer_self_attention_layers[i]( + output, + tgt_mask=None, + tgt_key_padding_mask=None, + query_pos=query_embed) + # FFN + output = self.transformer_ffn_layers[i](output) + + decoder_output = self.decoder_norm(output) + decoder_output = decoder_output.transpose( + 0, 1) # [N, bs, C] -> [bs, N, C] + color_embed = self.color_embed(decoder_output) + out = torch.einsum('bqc,bchw->bqhw', color_embed, last_img_feature) + + return out diff --git a/modelscope/models/cv/image_colorization/ddcolor/ddcolor_for_image_colorization.py b/modelscope/models/cv/image_colorization/ddcolor/ddcolor_for_image_colorization.py new file mode 100644 index 00000000..0d2acbd2 --- /dev/null +++ b/modelscope/models/cv/image_colorization/ddcolor/ddcolor_for_image_colorization.py @@ -0,0 +1,52 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import os.path as osp +from typing import Dict, Union + +from modelscope.metainfo import Models +from modelscope.models.base import Tensor, TorchModel +from modelscope.models.builder import MODELS +from modelscope.utils.constant import ModelFile, Tasks +from modelscope.utils.logger import get_logger +from .ddcolor import DDColor + +logger = get_logger() + +__all__ = ['DDColorForImageColorization'] + + +@MODELS.register_module(Tasks.image_colorization, module_name=Models.ddcolor) +class DDColorForImageColorization(TorchModel): + + def __init__(self, + model_dir, + encoder_name='convnext-l', + input_size=(512, 512), + num_queries=100, + *args, + **kwargs): + """initialize the image colorization model from the `model_dir` path. + + Args: + model_dir (str): the model path. + encoder_name (str): the encoder name. + input_size (tuple): size of the model input image. + num_queries (int): number of decoder queries + """ + super().__init__(model_dir, *args, **kwargs) + + self.model = DDColor(encoder_name, input_size, num_queries) + + model_path = osp.join(model_dir, ModelFile.TORCH_MODEL_FILE) + self.model = self._load_pretrained(self.model, model_path) + + def forward(self, input: Dict[str, + Tensor]) -> Dict[str, Union[list, Tensor]]: + """return the result of the model + + Args: + inputs (Tensor): the preprocessed data + + Returns: + Dict[str, Tensor]: results + """ + return self.model(**input) diff --git a/modelscope/models/cv/image_colorization/ddcolor/utils/__init__.py b/modelscope/models/cv/image_colorization/ddcolor/utils/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/modelscope/models/cv/image_colorization/ddcolor/utils/convnext.py b/modelscope/models/cv/image_colorization/ddcolor/utils/convnext.py new file mode 100644 index 00000000..3da14c48 --- /dev/null +++ b/modelscope/models/cv/image_colorization/ddcolor/utils/convnext.py @@ -0,0 +1,177 @@ +# The implementation here is modified based on ConvNeXt, originally MIT license +# and publicly available at https://github.com/facebookresearch/ConvNeXt/blob/main/models/convnext.py + +import torch +import torch.nn as nn +import torch.nn.functional as F +from timm.models.layers import DropPath, trunc_normal_ + + +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=3, + 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, stride=2), + ) + 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] + + # add norm layers for each output + out_indices = (0, 1, 2, 3) + for i in out_indices: + layer = LayerNorm(dims[i], eps=1e-6, data_format='channels_first') + # layer = nn.Identity() + layer_name = f'norm{i}' + self.add_module(layer_name, layer) + + self.norm = nn.LayerNorm(dims[-1], eps=1e-6) # final norm layer + self.head_cls = nn.Linear(dims[-1], 4) + + self.apply(self._init_weights) + self.head_cls.weight.data.mul_(head_init_scale) + self.head_cls.bias.data.mul_(head_init_scale) + + 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) + x = self.stages[i](x) + + # add extra norm + norm_layer = getattr(self, f'norm{i}') + norm_layer(x) + + return self.norm(x.mean( + [-2, -1])) # global average pooling, (N, C, H, W) -> (N, C) + + def forward(self, x): + x = self.forward_features(x) + x = self.head_cls(x) + return x + + +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': # B H W C + return F.layer_norm(x, self.normalized_shape, self.weight, + self.bias, self.eps) + elif self.data_format == 'channels_first': # B C H W + 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 diff --git a/modelscope/models/cv/image_colorization/ddcolor/utils/position_encoding.py b/modelscope/models/cv/image_colorization/ddcolor/utils/position_encoding.py new file mode 100644 index 00000000..e613d15b --- /dev/null +++ b/modelscope/models/cv/image_colorization/ddcolor/utils/position_encoding.py @@ -0,0 +1,57 @@ +# The implementation here is modified based on Mask2Former, originally MIT license and publicly available at +# https://github.com/facebookresearch/Mask2Former/blob/main/mask2former/modeling/transformer_decoder/position_encoding.py + +import math + +import torch +from torch import nn + + +class PositionEmbeddingSine(nn.Module): + """ + This is a more standard version of the position embedding, very similar to the one + used by the Attention is all you need paper, generalized to work on images. + """ + + def __init__(self, + num_pos_feats=64, + temperature=10000, + normalize=False, + scale=None): + super().__init__() + self.num_pos_feats = num_pos_feats + self.temperature = temperature + self.normalize = normalize + if scale is not None and normalize is False: + raise ValueError('normalize should be True if scale is passed') + if scale is None: + scale = 2 * math.pi + self.scale = scale + + def forward(self, x, mask=None): + if mask is None: + mask = torch.zeros((x.size(0), x.size(2), x.size(3)), + device=x.device, + dtype=torch.bool) + not_mask = ~mask + y_embed = not_mask.cumsum(1, dtype=torch.float32) + x_embed = not_mask.cumsum(2, dtype=torch.float32) + if self.normalize: + eps = 1e-6 + y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale + x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale + + dim_t = torch.arange( + self.num_pos_feats, dtype=torch.float32, device=x.device) + dim_t = self.temperature**(2 * (dim_t // 2) / self.num_pos_feats) + + pos_x = x_embed[:, :, :, None] / dim_t + pos_y = y_embed[:, :, :, None] / dim_t + pos_x = torch.stack( + (pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), + dim=4).flatten(3) + pos_y = torch.stack( + (pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), + dim=4).flatten(3) + pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2) + return pos diff --git a/modelscope/models/cv/image_colorization/ddcolor/utils/transformer_utils.py b/modelscope/models/cv/image_colorization/ddcolor/utils/transformer_utils.py new file mode 100644 index 00000000..66f2d821 --- /dev/null +++ b/modelscope/models/cv/image_colorization/ddcolor/utils/transformer_utils.py @@ -0,0 +1,232 @@ +# The implementation here is modified based on Mask2Former, originally MIT license and publicly available at +# https://github.com/facebookresearch/Mask2Former/blob/main/mask2former/modeling/transformer_decoder/mask2former_transformer_decoder.py + +from typing import Optional + +from torch import Tensor, nn +from torch.nn import functional as F + + +class SelfAttentionLayer(nn.Module): + + def __init__(self, + d_model, + nhead, + dropout=0.0, + activation='relu', + normalize_before=False): + super().__init__() + self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) + + self.norm = nn.LayerNorm(d_model) + self.dropout = nn.Dropout(dropout) + + self.activation = _get_activation_fn(activation) + self.normalize_before = normalize_before + + self._reset_parameters() + + def _reset_parameters(self): + for p in self.parameters(): + if p.dim() > 1: + nn.init.xavier_uniform_(p) + + def with_pos_embed(self, tensor, pos: Optional[Tensor]): + return tensor if pos is None else tensor + pos + + def forward_post(self, + tgt, + tgt_mask: Optional[Tensor] = None, + tgt_key_padding_mask: Optional[Tensor] = None, + query_pos: Optional[Tensor] = None): + q = k = self.with_pos_embed(tgt, query_pos) + tgt2 = self.self_attn( + q, + k, + value=tgt, + attn_mask=tgt_mask, + key_padding_mask=tgt_key_padding_mask)[0] + tgt = tgt + self.dropout(tgt2) + tgt = self.norm(tgt) + + return tgt + + def forward_pre(self, + tgt, + tgt_mask: Optional[Tensor] = None, + tgt_key_padding_mask: Optional[Tensor] = None, + query_pos: Optional[Tensor] = None): + tgt2 = self.norm(tgt) + q = k = self.with_pos_embed(tgt2, query_pos) + tgt2 = self.self_attn( + q, + k, + value=tgt2, + attn_mask=tgt_mask, + key_padding_mask=tgt_key_padding_mask)[0] + tgt = tgt + self.dropout(tgt2) + + return tgt + + def forward(self, + tgt, + tgt_mask: Optional[Tensor] = None, + tgt_key_padding_mask: Optional[Tensor] = None, + query_pos: Optional[Tensor] = None): + if self.normalize_before: + return self.forward_pre(tgt, tgt_mask, tgt_key_padding_mask, + query_pos) + return self.forward_post(tgt, tgt_mask, tgt_key_padding_mask, + query_pos) + + +class CrossAttentionLayer(nn.Module): + + def __init__(self, + d_model, + nhead, + dropout=0.0, + activation='relu', + normalize_before=False): + super().__init__() + self.multihead_attn = nn.MultiheadAttention( + d_model, nhead, dropout=dropout) + + self.norm = nn.LayerNorm(d_model) + self.dropout = nn.Dropout(dropout) + + self.activation = _get_activation_fn(activation) + self.normalize_before = normalize_before + + self._reset_parameters() + + def _reset_parameters(self): + for p in self.parameters(): + if p.dim() > 1: + nn.init.xavier_uniform_(p) + + def with_pos_embed(self, tensor, pos: Optional[Tensor]): + return tensor if pos is None else tensor + pos + + def forward_post(self, + tgt, + memory, + memory_mask: Optional[Tensor] = None, + memory_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None, + query_pos: Optional[Tensor] = None): + tgt2 = self.multihead_attn( + query=self.with_pos_embed(tgt, query_pos), + key=self.with_pos_embed(memory, pos), + value=memory, + attn_mask=memory_mask, + key_padding_mask=memory_key_padding_mask)[0] + tgt = tgt + self.dropout(tgt2) + tgt = self.norm(tgt) + + return tgt + + def forward_pre(self, + tgt, + memory, + memory_mask: Optional[Tensor] = None, + memory_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None, + query_pos: Optional[Tensor] = None): + tgt2 = self.norm(tgt) + tgt2 = self.multihead_attn( + query=self.with_pos_embed(tgt2, query_pos), + key=self.with_pos_embed(memory, pos), + value=memory, + attn_mask=memory_mask, + key_padding_mask=memory_key_padding_mask)[0] + tgt = tgt + self.dropout(tgt2) + + return tgt + + def forward(self, + tgt, + memory, + memory_mask: Optional[Tensor] = None, + memory_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None, + query_pos: Optional[Tensor] = None): + if self.normalize_before: + return self.forward_pre(tgt, memory, memory_mask, + memory_key_padding_mask, pos, query_pos) + return self.forward_post(tgt, memory, memory_mask, + memory_key_padding_mask, pos, query_pos) + + +class FFNLayer(nn.Module): + + def __init__(self, + d_model, + dim_feedforward=2048, + dropout=0.0, + activation='relu', + normalize_before=False): + super().__init__() + # Implementation of Feedforward model + self.linear1 = nn.Linear(d_model, dim_feedforward) + self.dropout = nn.Dropout(dropout) + self.linear2 = nn.Linear(dim_feedforward, d_model) + + self.norm = nn.LayerNorm(d_model) + + self.activation = _get_activation_fn(activation) + self.normalize_before = normalize_before + + self._reset_parameters() + + def _reset_parameters(self): + for p in self.parameters(): + if p.dim() > 1: + nn.init.xavier_uniform_(p) + + def with_pos_embed(self, tensor, pos: Optional[Tensor]): + return tensor if pos is None else tensor + pos + + def forward_post(self, tgt): + tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt)))) + tgt = tgt + self.dropout(tgt2) + tgt = self.norm(tgt) + return tgt + + def forward_pre(self, tgt): + tgt2 = self.norm(tgt) + tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2)))) + tgt = tgt + self.dropout(tgt2) + return tgt + + def forward(self, tgt): + if self.normalize_before: + return self.forward_pre(tgt) + return self.forward_post(tgt) + + +def _get_activation_fn(activation): + """Return an activation function given a string""" + if activation == 'relu': + return F.relu + if activation == 'gelu': + return F.gelu + if activation == 'glu': + return F.glu + raise RuntimeError(F'activation should be relu/gelu, not {activation}.') + + +class MLP(nn.Module): + """ Very simple multi-layer perceptron (also called FFN)""" + + def __init__(self, input_dim, hidden_dim, output_dim, num_layers): + super().__init__() + self.num_layers = num_layers + h = [hidden_dim] * (num_layers - 1) + self.layers = nn.ModuleList( + nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim])) + + def forward(self, x): + for i, layer in enumerate(self.layers): + x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x) + return x diff --git a/modelscope/models/cv/image_colorization/ddcolor/utils/unet.py b/modelscope/models/cv/image_colorization/ddcolor/utils/unet.py new file mode 100644 index 00000000..2b8b9c14 --- /dev/null +++ b/modelscope/models/cv/image_colorization/ddcolor/utils/unet.py @@ -0,0 +1,203 @@ +# The implementation here is modified based on DeOldify, originally MIT License +# and publicly available at https://github.com/jantic/DeOldify/blob/master/deoldify/unet.py + +import collections +from enum import Enum + +import torch +import torch.nn as nn +from torch.nn import functional as F + +NormType = Enum('NormType', 'Batch BatchZero Weight Spectral') + + +class Hook: + feature = None + + def __init__(self, module): + self.hook = module.register_forward_hook(self.hook_fn) + + def hook_fn(self, module, input, output): + if isinstance(output, torch.Tensor): + self.feature = output + elif isinstance(output, collections.OrderedDict): + self.feature = output['out'] + + def remove(self): + self.hook.remove() + + +class SelfAttention(nn.Module): + 'Self attention layer for nd.' + + def __init__(self, n_channels: int): + super().__init__() + self.query = conv1d(n_channels, n_channels // 8) + self.key = conv1d(n_channels, n_channels // 8) + self.value = conv1d(n_channels, n_channels) + self.gamma = nn.Parameter(torch.tensor([0.])) + + def forward(self, x): + # Notation from https://arxiv.org/pdf/1805.08318.pdf + size = x.size() + x = x.view(*size[:2], -1) + f, g, h = self.query(x), self.key(x), self.value(x) + beta = F.softmax(torch.bmm(f.permute(0, 2, 1).contiguous(), g), dim=1) + o = self.gamma * torch.bmm(h, beta) + x + return o.view(*size).contiguous() + + +def batchnorm_2d(nf: int, norm_type: NormType = NormType.Batch): + 'A batchnorm2d layer with `nf` features initialized depending on `norm_type`.' + bn = nn.BatchNorm2d(nf) + with torch.no_grad(): + bn.bias.fill_(1e-3) + bn.weight.fill_(0. if norm_type == NormType.BatchZero else 1.) + return bn + + +def init_default(m: nn.Module, func=nn.init.kaiming_normal_) -> None: + 'Initialize `m` weights with `func` and set `bias` to 0.' + if func: + if hasattr(m, 'weight'): + func(m.weight) + if hasattr(m, 'bias') and hasattr(m.bias, 'data'): + m.bias.data.fill_(0.) + return m + + +def icnr(x, scale=2, init=nn.init.kaiming_normal_): + 'ICNR init of `x`, with `scale` and `init` function.' + ni, nf, h, w = x.shape + ni2 = int(ni / (scale**2)) + k = init(torch.zeros([ni2, nf, h, w])).transpose(0, 1) + k = k.contiguous().view(ni2, nf, -1) + k = k.repeat(1, 1, scale**2) + k = k.contiguous().view([nf, ni, h, w]).transpose(0, 1) + x.data.copy_(k) + + +def conv1d(ni: int, + no: int, + ks: int = 1, + stride: int = 1, + padding: int = 0, + bias: bool = False): + 'Create and initialize a `nn.Conv1d` layer with spectral normalization.' + conv = nn.Conv1d(ni, no, ks, stride=stride, padding=padding, bias=bias) + nn.init.kaiming_normal_(conv.weight) + if bias: + conv.bias.data.zero_() + return nn.utils.spectral_norm(conv) + + +def custom_conv_layer( + ni: int, + nf: int, + ks: int = 3, + stride: int = 1, + padding: int = None, + bias: bool = None, + is_1d: bool = False, + norm_type=NormType.Batch, + use_activ: bool = True, + transpose: bool = False, + init=nn.init.kaiming_normal_, + self_attention: bool = False, + extra_bn: bool = False, +): + 'Create a sequence of convolutional (`ni` to `nf`), ReLU (if `use_activ`) and batchnorm (if `bn`) layers.' + if padding is None: + padding = (ks - 1) // 2 if not transpose else 0 + bn = norm_type in (NormType.Batch, NormType.BatchZero) or extra_bn + if bias is None: + bias = not bn + conv_func = nn.ConvTranspose2d if transpose else nn.Conv1d if is_1d else nn.Conv2d + conv = init_default( + conv_func( + ni, nf, kernel_size=ks, bias=bias, stride=stride, padding=padding), + init, + ) + + if norm_type == NormType.Weight: + conv = nn.utils.weight_norm(conv) + elif norm_type == NormType.Spectral: + conv = nn.utils.spectral_norm(conv) + layers = [conv] + if use_activ: + layers.append(nn.ReLU(True)) + if bn: + layers.append((nn.BatchNorm1d if is_1d else nn.BatchNorm2d)(nf)) + if self_attention: + layers.append(SelfAttention(nf)) + return nn.Sequential(*layers) + + +class CustomPixelShuffle_ICNR(nn.Module): + """ + Upsample by `scale` from `ni` filters to `nf` (default `ni`), + using `nn.PixelShuffle`, `icnr` init, and `weight_norm`. + """ + + def __init__(self, + ni: int, + nf: int = None, + scale: int = 2, + blur: bool = True, + norm_type=NormType.Spectral, + extra_bn=False): + super().__init__() + self.conv = custom_conv_layer( + ni, + nf * (scale**2), + ks=1, + use_activ=False, + norm_type=norm_type, + extra_bn=extra_bn) + icnr(self.conv[0].weight) + self.shuf = nn.PixelShuffle(scale) + self.do_blur = blur + # Blurring over (h*w) kernel + # "Super-Resolution using Convolutional Neural Networks without Any Checkerboard Artifacts" + # - https://arxiv.org/abs/1806.02658 + self.pad = nn.ReplicationPad2d((1, 0, 1, 0)) + self.blur = nn.AvgPool2d(2, stride=1) + self.relu = nn.ReLU(True) + + def forward(self, x): + x = self.shuf(self.relu(self.conv(x))) + return self.blur(self.pad(x)) if self.do_blur else x + + +class UnetBlockWide(nn.Module): + 'A quasi-UNet block, using `PixelShuffle_ICNR upsampling`.' + + def __init__(self, + up_in_c: int, + x_in_c: int, + n_out: int, + hook, + blur: bool = False, + self_attention: bool = False, + norm_type=NormType.Spectral): + super().__init__() + + self.hook = hook + up_out = n_out + self.shuf = CustomPixelShuffle_ICNR( + up_in_c, up_out, blur=blur, norm_type=norm_type, extra_bn=True) + self.bn = batchnorm_2d(x_in_c) + ni = up_out + x_in_c + self.conv = custom_conv_layer( + ni, + n_out, + norm_type=norm_type, + self_attention=self_attention, + extra_bn=True) + self.relu = nn.ReLU() + + def forward(self, up_in): + s = self.hook.feature + up_out = self.shuf(up_in) + cat_x = self.relu(torch.cat([up_out, self.bn(s)], dim=1)) + return self.conv(cat_x) diff --git a/modelscope/models/cv/image_colorization/unet/__init__.py b/modelscope/models/cv/image_colorization/unet/__init__.py new file mode 100644 index 00000000..66879422 --- /dev/null +++ b/modelscope/models/cv/image_colorization/unet/__init__.py @@ -0,0 +1,3 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +from .unet import DynamicUnetDeep, DynamicUnetWide +from .utils import NormType diff --git a/modelscope/models/cv/image_colorization/unet.py b/modelscope/models/cv/image_colorization/unet/unet.py similarity index 100% rename from modelscope/models/cv/image_colorization/unet.py rename to modelscope/models/cv/image_colorization/unet/unet.py diff --git a/modelscope/models/cv/image_colorization/utils.py b/modelscope/models/cv/image_colorization/unet/utils.py similarity index 100% rename from modelscope/models/cv/image_colorization/utils.py rename to modelscope/models/cv/image_colorization/unet/utils.py diff --git a/modelscope/pipelines/builder.py b/modelscope/pipelines/builder.py index fe2f1c1d..ae5b6603 100644 --- a/modelscope/pipelines/builder.py +++ b/modelscope/pipelines/builder.py @@ -163,8 +163,8 @@ DEFAULT_MODEL_FOR_PIPELINE = { 'damo/cv_csrnet_image-color-enhance-models'), Tasks.virtual_try_on: (Pipelines.virtual_try_on, 'damo/cv_daflow_virtual-try-on_base'), - Tasks.image_colorization: (Pipelines.image_colorization, - 'damo/cv_unet_image-colorization'), + Tasks.image_colorization: (Pipelines.ddcolor_image_colorization, + 'damo/cv_ddcolor_image-colorization'), Tasks.image_segmentation: (Pipelines.image_instance_segmentation, 'damo/cv_swin-b_image-instance-segmentation_coco'), diff --git a/modelscope/pipelines/cv/__init__.py b/modelscope/pipelines/cv/__init__.py index a2ccf462..d5e1dff2 100644 --- a/modelscope/pipelines/cv/__init__.py +++ b/modelscope/pipelines/cv/__init__.py @@ -74,6 +74,7 @@ if TYPE_CHECKING: from .pointcloud_sceneflow_estimation_pipeline import PointCloudSceneFlowEstimationPipeline from .maskdino_instance_segmentation_pipeline import MaskDINOInstanceSegmentationPipeline from .image_mvs_depth_estimation_pipeline import ImageMultiViewDepthEstimationPipeline + from .ddcolor_image_colorization_pipeline import DDColorImageColorizationPipeline else: _import_structure = { @@ -176,6 +177,9 @@ else: 'image_mvs_depth_estimation_pipeline': [ 'ImageMultiViewDepthEstimationPipeline' ], + 'ddcolor_image_colorization_pipeline': [ + 'DDColorImageColorizationPipeline' + ], } import sys diff --git a/modelscope/pipelines/cv/ddcolor_image_colorization_pipeline.py b/modelscope/pipelines/cv/ddcolor_image_colorization_pipeline.py new file mode 100644 index 00000000..43e505fd --- /dev/null +++ b/modelscope/pipelines/cv/ddcolor_image_colorization_pipeline.py @@ -0,0 +1,168 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +from typing import Any, Dict, Optional, Union + +import cv2 +import numpy as np +import torch +import torch.nn.functional as F +from torchvision import transforms + +from modelscope.metainfo import Pipelines +from modelscope.models.cv.image_colorization import DDColorForImageColorization +from modelscope.outputs import OutputKeys +from modelscope.pipelines.base import Input, Pipeline +from modelscope.pipelines.builder import PIPELINES +from modelscope.preprocessors import LoadImage +from modelscope.utils.constant import ModelFile, Tasks +from modelscope.utils.logger import get_logger + +logger = get_logger() + + +@PIPELINES.register_module( + Tasks.image_colorization, module_name=Pipelines.ddcolor_image_colorization) +class DDColorImageColorizationPipeline(Pipeline): + """ DDColor Image Colorization Pipeline. + + Example: + + ```python + >>> from modelscope.pipelines import pipeline + + >>> colorizer = pipeline('image-colorization', 'damo/cv_ddcolor_image-colorization') + >>> colorizer("data/test/images/audrey_hepburn.jpg") + {'output_img': array([[[198, 199, 193], + [198, 199, 193], + [197, 199, 195], + ..., + [197, 213, 206], + [197, 213, 206], + [197, 213, 207]], + + [[198, 199, 193], + [198, 199, 193], + [197, 199, 195], + ..., + [196, 212, 205], + [196, 212, 205], + [196, 212, 206]], + + [[198, 199, 193], + [198, 199, 193], + [197, 199, 195], + ..., + [193, 209, 202], + [193, 209, 202], + [193, 209, 203]], + + ..., + + [[ 56, 72, 103], + [ 56, 72, 103], + [ 56, 72, 102], + ..., + [233, 231, 232], + [233, 231, 232], + [233, 231, 232]], + + [[ 51, 62, 91], + [ 52, 63, 92], + [ 52, 64, 92], + ..., + [233, 232, 231], + [233, 232, 231], + [232, 232, 229]], + + [[ 60, 72, 101], + [ 59, 71, 100], + [ 57, 70, 99], + ..., + [233, 232, 231], + [233, 232, 231], + [232, 232, 229]]], dtype=uint8)} + >>> # + ``` + """ + + def __init__(self, model: Union[DDColorForImageColorization, str], + **kwargs): + """ + use `model` to create an image colorization pipeline for prediction + + Args: + model: model id on modelscope hub. + """ + super().__init__(model=model, **kwargs) + self.model.eval() + self.input_size = 512 + if torch.cuda.is_available(): + self._device = torch.device('cuda') + else: + self._device = torch.device('cpu') + + # self.model = DDColorForImageColorization( + # model_dir=model, + # encoder_name='convnext-l', + # input_size=[self.input_size, self.input_size], + # ).to(self.device) + + # model_path = f'{model}/{ModelFile.TORCH_MODEL_FILE}' + # logger.info(f'loading model from {model_path}') + # self.model.load_state_dict( + # torch.load(model_path, map_location=torch.device('cpu'))['params'], + # strict=True) + + logger.info('load model done') + + def preprocess(self, input: Input) -> Dict[str, Any]: + """preprocess the input image, extract L-channel and convert it back to RGB + + Args: + inputs: an input image from file or url + + Returns: + Dict[str, Any]: the pre-processed image + """ + img = LoadImage.convert_to_ndarray(input) + self.height, self.width = img.shape[:2] + + img = (img / 255.0).astype(np.float32) + self.orig_l = cv2.cvtColor(img, cv2.COLOR_BGR2Lab)[:, :, :1] + + img = cv2.resize(img, (self.input_size, self.input_size)) + img_l = cv2.cvtColor(img, cv2.COLOR_BGR2Lab)[:, :, :1] + img_gray_lab = np.concatenate( + (img_l, np.zeros_like(img_l), np.zeros_like(img_l)), axis=-1) + img_gray_rgb = cv2.cvtColor(img_gray_lab, cv2.COLOR_LAB2RGB) + tensor_gray_rgb = torch.from_numpy(img_gray_rgb.transpose( + (2, 0, 1))).float() + tensor_gray_rgb = tensor_gray_rgb.unsqueeze(0).to(self.device) + + result = {'img': tensor_gray_rgb} + return result + + @torch.no_grad() + def forward(self, input: Dict[str, Any]) -> Dict[str, Any]: + """call model to output the predictions and concatenate it with the original L-channel + + Args: + inputs: input image tensor + + Returns: + Dict[str, Any]: the result image + """ + + output_ab = self.model(input).cpu() + + output_ab_resize = F.interpolate( + output_ab, size=(self.height, self.width)) + output_ab_resize = output_ab_resize[0].float().numpy().transpose( + 1, 2, 0) + out_lab = np.concatenate((self.orig_l, output_ab_resize), axis=-1) + out_bgr = cv2.cvtColor(out_lab, cv2.COLOR_LAB2BGR) + output_img = (out_bgr * 255.0).round().astype(np.uint8) + + return {OutputKeys.OUTPUT_IMG: output_img} + + def postprocess(self, inputs: Dict[str, Any]) -> Dict[str, Any]: + return inputs diff --git a/tests/pipelines/test_ddcolor_image_colorization.py b/tests/pipelines/test_ddcolor_image_colorization.py new file mode 100644 index 00000000..e1876329 --- /dev/null +++ b/tests/pipelines/test_ddcolor_image_colorization.py @@ -0,0 +1,61 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import os.path as osp +import unittest + +import cv2 + +from modelscope.hub.snapshot_download import snapshot_download +from modelscope.models import Model +from modelscope.outputs import OutputKeys +from modelscope.pipelines import pipeline +from modelscope.pipelines.base import Pipeline +from modelscope.pipelines.cv import DDColorImageColorizationPipeline +from modelscope.utils.constant import Tasks +from modelscope.utils.demo_utils import DemoCompatibilityCheck +from modelscope.utils.test_utils import test_level + + +class DDColorImageColorizationTest(unittest.TestCase, DemoCompatibilityCheck): + + def setUp(self) -> None: + self.task = Tasks.image_colorization + self.model_id = 'damo/cv_ddcolor_image-colorization' + self.test_image = 'data/test/images/audrey_hepburn.jpg' + + def pipeline_inference(self, pipeline: Pipeline, test_image: str): + result = pipeline(test_image) + if result is not None: + cv2.imwrite('result.png', result[OutputKeys.OUTPUT_IMG]) + print(f'Output written to {osp.abspath("result.png")}') + + @unittest.skipUnless(test_level() >= 2, 'skip test in current test level') + def test_run_by_direct_model_download(self): + cache_path = snapshot_download(self.model_id) + image_colorization = DDColorImageColorizationPipeline(cache_path) + self.pipeline_inference(image_colorization, self.test_image) + + @unittest.skipUnless(test_level() >= 1, 'skip test in current test level') + def test_run_with_model_from_pretrained(self): + model = Model.from_pretrained(self.model_id) + image_colorization = pipeline( + task=Tasks.image_colorization, model=model) + self.pipeline_inference(image_colorization, self.test_image) + + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + def test_run_with_model_from_modelhub(self): + image_colorization = pipeline( + task=Tasks.image_colorization, model=self.model_id) + self.pipeline_inference(image_colorization, self.test_image) + + @unittest.skipUnless(test_level() >= 2, 'skip test in current test level') + def test_run_with_default_model(self): + image_colorization = pipeline(Tasks.image_colorization) + self.pipeline_inference(image_colorization, self.test_image) + + @unittest.skip('demo compatibility test is only enabled on a needed-basis') + def test_demo_compatibility(self): + self.compatibility_check() + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/run_config.yaml b/tests/run_config.yaml index 0b373fbc..8c84658b 100644 --- a/tests/run_config.yaml +++ b/tests/run_config.yaml @@ -45,6 +45,7 @@ isolated: # test cases that may require excessive anmount of GPU memory or run - test_video_super_resolution.py - test_kws_nearfield_trainer.py - test_gpt3_text_generation.py + - test_ddcolor_image_colorization.py envs: default: # default env, case not in other env will in default, pytorch.