mmagic.models.editors.gca.gca_module
¶
Module Contents¶
Classes¶
Guided Contextual Attention Module. |
- class mmagic.models.editors.gca.gca_module.GCAModule(in_channels, out_channels, kernel_size=3, stride=1, rate=2, pad_args=dict(mode='reflect'), interpolation='nearest', penalty=- 10000.0, eps=0.0001)[source]¶
Bases:
torch.nn.Module
Guided Contextual Attention Module.
From https://arxiv.org/pdf/2001.04069.pdf. Based on https://github.com/nbei/Deep-Flow-Guided-Video-Inpainting. This module use image feature map to augment the alpha feature map with guided contextual attention score.
Image feature and alpha feature are unfolded to small patches and later used as conv kernel. Thus, we refer the unfolding size as kernel size. Image feature patches have a default kernel size 3 while the kernel size of alpha feature patches could be specified by rate (see rate below). The image feature patches are used to convolve with the image feature itself to calculate the contextual attention. Then the attention feature map is convolved by alpha feature patches to obtain the attention alpha feature. At last, the attention alpha feature is added to the input alpha feature.
- Parameters
in_channels (int) – Input channels of the guided contextual attention module.
out_channels (int) – Output channels of the guided contextual attention module.
kernel_size (int) – Kernel size of image feature patches. Default 3.
stride (int) – Stride when unfolding the image feature. Default 1.
rate (int) – The downsample rate of image feature map. The corresponding kernel size and stride of alpha feature patches will be rate x 2 and rate. It could be regarded as the granularity of the gca module. Default: 2.
pad_args (dict) – Parameters of padding when convolve image feature with image feature patches or alpha feature patches. Allowed keys are mode and value. See torch.nn.functional.pad() for more information. Default: dict(mode=’reflect’).
interpolation (str) – Interpolation method in upsampling and downsampling.
penalty (float) – Punishment hyperparameter to avoid a large correlation between each unknown patch and itself. Default: -1e4.
eps (float) – A small number to avoid dividing by 0 when calculating the normed image feature patch. Default: 1e-4.
- forward(img_feat, alpha_feat, unknown=None, softmax_scale=1.0)[source]¶
Forward function of GCAModule.
- Parameters
img_feat (Tensor) – Image feature map of shape (N, ori_c, ori_h, ori_w).
alpha_feat (Tensor) – Alpha feature map of shape (N, alpha_c, ori_h, ori_w).
unknown (Tensor, optional) – Unknown area map generated by trimap. If specified, this tensor should have shape (N, 1, ori_h, ori_w).
softmax_scale (float, optional) – The softmax scale of the attention if unknown area is not provided in forward. Default: 1.
- Returns
The augmented alpha feature.
- Return type
Tensor
- extract_feature_maps_patches(img_feat, alpha_feat, unknown)[source]¶
Extract image feature, alpha feature unknown patches.
- Parameters
img_feat (Tensor) – Image feature map of shape (N, img_c, img_h, img_w).
alpha_feat (Tensor) – Alpha feature map of shape (N, alpha_c, ori_h, ori_w).
unknown (Tensor, optional) – Unknown area map generated by trimap of shape (N, 1, img_h, img_w).
- Returns
3-tuple of
Tensor
: Image feature patches of shape (N, img_h*img_w, img_c, img_ks, img_ks).Tensor
: Guided contextual attention alpha feature map. (N, img_h*img_w, alpha_c, alpha_ks, alpha_ks).Tensor
: Unknown mask of shape (N, img_h*img_w, 1, 1).- Return type
tuple
- compute_similarity_map(img_feat, img_ps)[source]¶
Compute similarity between image feature patches.
- Parameters
img_feat (Tensor) – Image feature map of shape (1, img_c, img_h, img_w).
img_ps (Tensor) – Image feature patches tensor of shape (1, img_h*img_w, img_c, img_ks, img_ks).
- Returns
Similarity map between image feature patches with shape (1, img_h*img_w, img_h, img_w).
- Return type
Tensor
- compute_guided_attention_score(similarity_map, unknown_ps, scale, self_mask)[source]¶
Compute guided attention score.
- Parameters
similarity_map (Tensor) – Similarity map of image feature with shape (1, img_h*img_w, img_h, img_w).
unknown_ps (Tensor) – Unknown area patches tensor of shape (1, img_h*img_w, 1, 1).
scale (Tensor) – Softmax scale of known and unknown area: [unknown_scale, known_scale].
self_mask (Tensor) – Self correlation mask of shape (1, img_h*img_w, img_h, img_w). At (1, i*i, i, i) mask value equals -1e4 for i in [1, img_h*img_w] and other area is all zero.
- Returns
Similarity map between image feature patches with shape (1, img_h*img_w, img_h, img_w).
- Return type
Tensor
- propagate_alpha_feature(gca_score, alpha_ps)[source]¶
Propagate alpha feature based on guided attention score.
- Parameters
gca_score (Tensor) – Guided attention score map of shape (1, img_h*img_w, img_h, img_w).
alpha_ps (Tensor) – Alpha feature patches tensor of shape (1, img_h*img_w, alpha_c, alpha_ks, alpha_ks).
- Returns
Propagated alpha feature map of shape (1, alpha_c, alpha_h, alpha_w).
- Return type
Tensor
- process_unknown_mask(unknown, img_feat, softmax_scale)[source]¶
Process unknown mask.
- Parameters
unknown (Tensor, optional) – Unknown area map generated by trimap of shape (N, 1, ori_h, ori_w)
img_feat (Tensor) – The interpolated image feature map of shape (N, img_c, img_h, img_w).
softmax_scale (float, optional) – The softmax scale of the attention if unknown area is not provided in forward. Default: 1.
- Returns
2-tuple of
Tensor
: Interpolated unknown area map of shape (N, img_h*img_w, img_h, img_w).Tensor
: Softmax scale tensor of known and unknown area of shape (N, 2).- Return type
tuple
- extract_patches(x, kernel_size, stride)[source]¶
Extract feature patches.
The feature map will be padded automatically to make sure the number of patches is equal to (H / stride) * (W / stride).
- Parameters
x (Tensor) – Feature map of shape (N, C, H, W).
kernel_size (int) – Size of each patches.
stride (int) – Stride between patches.
- Returns
Extracted patches of shape (N, (H / stride) * (W / stride) , C, kernel_size, kernel_size).
- Return type
Tensor
- pad(x, kernel_size, stride)[source]¶
Pad input tensor.
- Parameters
x (Tensor) – Input tensor.
kernel_size (int) – Kernel size of conv layer.
stride (int) – Stride of conv layer.
- Returns
Padded tensor
- Return type
Tensor