mindspore-ci-bot
5 years ago
committed by
Gitee
commit
fe4cde6b1e
@ -0,0 +1,197 @@
|
|||||||
|
# Copyright 2020 Huawei Technologies Co., Ltd
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
# ============================================================================
|
||||||
|
"""image"""
|
||||||
|
import numpy as np
|
||||||
|
import mindspore.common.dtype as mstype
|
||||||
|
from mindspore.common.tensor import Tensor
|
||||||
|
from mindspore.ops import operations as P
|
||||||
|
from mindspore.ops import functional as F
|
||||||
|
from mindspore.ops.primitive import constexpr
|
||||||
|
from mindspore._checkparam import ParamValidator as validator
|
||||||
|
from mindspore._checkparam import Rel
|
||||||
|
from ..cell import Cell
|
||||||
|
|
||||||
|
|
||||||
|
class ImageGradients(Cell):
|
||||||
|
r"""
|
||||||
|
Returns two tensors, the first is along the height dimension and the second is along the width dimension.
|
||||||
|
|
||||||
|
Assume an image shape is :math:`h*w`. The gradients along the height and the width are :math:`dy` and :math:`dx`,
|
||||||
|
respectively.
|
||||||
|
|
||||||
|
.. math::
|
||||||
|
dy[i] = \begin{cases} image[i+1, :]-image[i, :], &if\ 0<=i<h-1 \cr
|
||||||
|
0, &if\ i==h-1\end{cases}
|
||||||
|
|
||||||
|
dx[i] = \begin{cases} image[:, i+1]-image[:, i], &if\ 0<=i<w-1 \cr
|
||||||
|
0, &if\ i==w-1\end{cases}
|
||||||
|
|
||||||
|
Inputs:
|
||||||
|
- **images** (Tensor) - The input image data, with format 'NCHW'.
|
||||||
|
|
||||||
|
Outputs:
|
||||||
|
- **dy** (Tensor) - vertical image gradients, the same type and shape as input.
|
||||||
|
- **dx** (Tensor) - horizontal image gradients, the same type and shape as input.
|
||||||
|
|
||||||
|
Examples:
|
||||||
|
>>> net = nn.ImageGradients()
|
||||||
|
>>> image = Tensor(np.array([[[[1,2],[3,4]]]]), dtype=mstype.int32)
|
||||||
|
>>> net(image)
|
||||||
|
[[[[2,2]
|
||||||
|
[0,0]]]]
|
||||||
|
[[[[1,0]
|
||||||
|
[1,0]]]]
|
||||||
|
"""
|
||||||
|
def __init__(self):
|
||||||
|
super(ImageGradients, self).__init__()
|
||||||
|
|
||||||
|
def construct(self, images):
|
||||||
|
batch_size, depth, height, width = P.Shape()(images)
|
||||||
|
dy = images[:, :, 1:, :] - images[:, :, :height - 1, :]
|
||||||
|
dy_last = P.Fill()(P.DType()(images), (batch_size, depth, 1, width), 0)
|
||||||
|
dy = P.Concat(2)((dy, dy_last))
|
||||||
|
|
||||||
|
dx = images[:, :, :, 1:] - images[:, :, :, :width - 1]
|
||||||
|
dx_last = P.Fill()(P.DType()(images), (batch_size, depth, height, 1), 0)
|
||||||
|
dx = P.Concat(3)((dx, dx_last))
|
||||||
|
return dy, dx
|
||||||
|
|
||||||
|
|
||||||
|
@constexpr
|
||||||
|
def _gauss_kernel_helper(filter_size):
|
||||||
|
"""gauss kernel helper"""
|
||||||
|
filter_size = F.scalar_cast(filter_size, mstype.int32)
|
||||||
|
coords = ()
|
||||||
|
for i in range(filter_size):
|
||||||
|
i_cast = F.scalar_cast(i, mstype.float32)
|
||||||
|
offset = F.scalar_cast(filter_size-1, mstype.float32)/2.0
|
||||||
|
element = i_cast-offset
|
||||||
|
coords = coords+(element,)
|
||||||
|
g = np.square(coords).astype(np.float32)
|
||||||
|
g = Tensor(g)
|
||||||
|
return filter_size, g
|
||||||
|
|
||||||
|
|
||||||
|
class SSIM(Cell):
|
||||||
|
r"""
|
||||||
|
Returns SSIM index between img1 and img2.
|
||||||
|
|
||||||
|
Its implementation is based on Wang, Z., Bovik, A. C., Sheikh, H. R., & Simoncelli, E. P. (2004). `Image quality
|
||||||
|
assessment: from error visibility to structural similarity <https://ieeexplore.ieee.org/document/1284395>`_.
|
||||||
|
IEEE transactions on image processing.
|
||||||
|
|
||||||
|
.. math::
|
||||||
|
|
||||||
|
l(x,y)&=\frac{2\mu_x\mu_y+C_1}{\mu_x^2+\mu_y^2+C_1}, C_1=(K_1L)^2.\\
|
||||||
|
c(x,y)&=\frac{2\sigma_x\sigma_y+C_2}{\sigma_x^2+\sigma_y^2+C_2}, C_2=(K_2L)^2.\\
|
||||||
|
s(x,y)&=\frac{\sigma_{xy}+C_3}{\sigma_x\sigma_y+C_3}, C_3=C_2/2.\\
|
||||||
|
SSIM(x,y)&=l*c*s\\&=\frac{(2\mu_x\mu_y+C_1)(2\sigma_{xy}+C_2}{(\mu_x^2+\mu_y^2+C_1)(\sigma_x^2+\sigma_y^2+C_2)}.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
max_val (Union[int, float]): The dynamic range of the pixel values (255 for 8-bit grayscale images).
|
||||||
|
Default: 1.0.
|
||||||
|
filter_size (int): The size of the Gaussian filter. Default: 11.
|
||||||
|
filter_sigma (float): The standard deviation of Gaussian kernel. Default: 1.5.
|
||||||
|
k1 (float): The constant used to generate c1 in the luminance comparison function. Default: 0.01.
|
||||||
|
k2 (float): The constant used to generate c2 in the contrast comparison function. Default: 0.03.
|
||||||
|
|
||||||
|
Inputs:
|
||||||
|
- **img1** (Tensor) - The first image batch with format 'NCHW'. It should be the same shape and dtype as img2.
|
||||||
|
- **img2** (Tensor) - The second image batch with format 'NCHW'. It should be the same shape and dtype as img1.
|
||||||
|
|
||||||
|
Outputs:
|
||||||
|
Tensor, has the same dtype as img1. It is a 1-D tensor with shape N, where N is the batch num of img1.
|
||||||
|
|
||||||
|
Examples:
|
||||||
|
>>> net = nn.SSIM()
|
||||||
|
>>> img1 = Tensor(np.random.random((1,3,16,16)))
|
||||||
|
>>> img2 = Tensor(np.random.random((1,3,16,16)))
|
||||||
|
>>> ssim = net(img1, img2)
|
||||||
|
"""
|
||||||
|
def __init__(self, max_val=1.0, filter_size=11, filter_sigma=1.5, k1=0.01, k2=0.03):
|
||||||
|
super(SSIM, self).__init__()
|
||||||
|
validator.check_type('max_val', max_val, [int, float])
|
||||||
|
validator.check('max_val', max_val, '', 0.0, Rel.GT)
|
||||||
|
self.max_val = max_val
|
||||||
|
self.filter_size = validator.check_integer('filter_size', filter_size, 1, Rel.GE)
|
||||||
|
self.filter_sigma = validator.check_float_positive('filter_sigma', filter_sigma)
|
||||||
|
validator.check_type('k1', k1, [float])
|
||||||
|
self.k1 = validator.check_number_range('k1', k1, 0.0, 1.0, Rel.INC_NEITHER)
|
||||||
|
validator.check_type('k2', k2, [float])
|
||||||
|
self.k2 = validator.check_number_range('k2', k2, 0.0, 1.0, Rel.INC_NEITHER)
|
||||||
|
self.mean = P.DepthwiseConv2dNative(channel_multiplier=1, kernel_size=filter_size)
|
||||||
|
|
||||||
|
def construct(self, img1, img2):
|
||||||
|
max_val = self._convert_img_dtype_to_float32(self.max_val, self.max_val)
|
||||||
|
img1 = self._convert_img_dtype_to_float32(img1, self.max_val)
|
||||||
|
img2 = self._convert_img_dtype_to_float32(img2, self.max_val)
|
||||||
|
|
||||||
|
kernel = self._fspecial_gauss(self.filter_size, self.filter_sigma)
|
||||||
|
kernel = P.Tile()(kernel, (1, P.Shape()(img1)[1], 1, 1))
|
||||||
|
|
||||||
|
mean_ssim = self._calculate_mean_ssim(img1, img2, kernel, max_val, self.k1, self.k2)
|
||||||
|
|
||||||
|
return mean_ssim
|
||||||
|
|
||||||
|
def _convert_img_dtype_to_float32(self, img, max_val):
|
||||||
|
"""convert img dtype to float32"""
|
||||||
|
# Ususally max_val is 1.0 or 255, we will do the scaling if max_val > 1.
|
||||||
|
# We will scale img pixel value if max_val > 1. and just cast otherwise.
|
||||||
|
ret = P.Cast()(img, mstype.float32)
|
||||||
|
max_val = F.scalar_cast(max_val, mstype.float32)
|
||||||
|
if max_val > 1.:
|
||||||
|
scale = 1./max_val
|
||||||
|
ret = ret*scale
|
||||||
|
return ret
|
||||||
|
|
||||||
|
def _calculate_mean_ssim(self, x, y, kernel, max_val, k1, k2):
|
||||||
|
"""calculate mean ssim"""
|
||||||
|
c1 = (k1*max_val)*(k1*max_val)
|
||||||
|
c2 = (k2*max_val)*(k2*max_val)
|
||||||
|
|
||||||
|
# SSIM luminance formula
|
||||||
|
# (2 * mean_{x} * mean_{y} + c1) / (mean_{x}**2 + mean_{y}**2 + c1)
|
||||||
|
mean_x = self.mean(x, kernel)
|
||||||
|
mean_y = self.mean(y, kernel)
|
||||||
|
square_sum = F.square(mean_x)+F.square(mean_y)
|
||||||
|
luminance = (2*mean_x*mean_y+c1)/(square_sum+c1)
|
||||||
|
|
||||||
|
# SSIM contrast*structure formula (when c3 = c2/2)
|
||||||
|
# (2 * conv_{xy} + c2) / (conv_{xx} + conv_{yy} + c2), equals to
|
||||||
|
# (2 * (mean_{xy} - mean_{x}*mean_{y}) + c2) / (mean_{xx}-mean_{x}**2 + mean_{yy}-mean_{y}**2 + c2)
|
||||||
|
mean_xy = self.mean(x*y, kernel)
|
||||||
|
mean_square_add = self.mean(F.square(x)+F.square(y), kernel)
|
||||||
|
|
||||||
|
cs = (2*(mean_xy-mean_x*mean_y)+c2)/(mean_square_add-square_sum+c2)
|
||||||
|
|
||||||
|
# SSIM formula
|
||||||
|
# luminance * cs
|
||||||
|
ssim = luminance*cs
|
||||||
|
|
||||||
|
mean_ssim = P.ReduceMean()(ssim, (-3, -2, -1))
|
||||||
|
|
||||||
|
return mean_ssim
|
||||||
|
|
||||||
|
def _fspecial_gauss(self, filter_size, filter_sigma):
|
||||||
|
"""get gauss kernel"""
|
||||||
|
filter_size, g = _gauss_kernel_helper(filter_size)
|
||||||
|
|
||||||
|
square_sigma_scale = -0.5/(filter_sigma * filter_sigma)
|
||||||
|
g = g*square_sigma_scale
|
||||||
|
g = F.reshape(g, (1, -1))+F.reshape(g, (-1, 1))
|
||||||
|
g = F.reshape(g, (1, -1))
|
||||||
|
g = P.Softmax()(g)
|
||||||
|
ret = F.reshape(g, (1, 1, filter_size, filter_size))
|
||||||
|
return ret
|
||||||
Loading…
Reference in new issue