文章目录
1.关于depthwise-seperate conv 概念解释
主要包括depthwise conv 和 pointwise conv
优点是
- 参数少,避免过拟合
- 计算量小
以一个示例说明参数大小,输入输出
"""
Normal Convolution and depthwise-separable convolutions
should output a vector with the same dimensions
input shape = 3, 28, 38 (RGB image of 28x 28)
output shape = 10, 28, 28 (10 output channels with same width and height)
"""
import torch
import torch.nn as nn
input = torch.rand(3, 28, 28)
### Conv2d (normal)
conv_layer= nn.Conv2d(in_channels=3, out_channels=10, kernel_size=3, padding=1)
print([[p.shape, p.name, p.requires_grad, p.numel()] for p in conv_layer.parameters()])
print(conv_layer.weight.shape, conv_layer.bias.shape)
conv_layer_n_params = sum(p.numel() for p in conv_layer.parameters() if p.requires_grad)
conv_out = conv_layer(input)
print(f"Conv layer param numbers: {conv_layer_n_params}")
print(conv_out.shape)
print('********************************************************')
### Depthwise convolution
# by adding 'groups' param, you perform depthwise conv
depthwise_layer= nn.Conv2d(in_channels=3, out_channels=3, kernel_size=3, padding=1, groups=3)
depthwise_layer_n_params = sum(p.numel() for p in depthwise_layer.parameters() if p.requires_grad)
depthwise_out = depthwise_layer(input)
print([[p.shape, p.name, p.requires_grad, p.numel()] for p in depthwise_layer.parameters()])
print(f"DepthwiseConv layer param numbers: {depthwise_layer_n_params}")
print(depthwise_out.shape)
print('********************************************************')
### Pointwise Convolution (using depthwise output) <-- This is called DEPTHWISE-SEPARABLE CONVOLUTION
pointwise_layer = nn.Conv2d(in_channels=3, out_channels=10, kernel_size=1)
pointwise_layer_n_params = sum(p.numel() for p in pointwise_layer.parameters() if p.requires_grad)
pointwise_out = pointwise_layer(depthwise_out)
print([[p.shape, p.name, p.requires_grad, p.numel()] for p in pointwise_layer.parameters()])
print(f"PointwiseConv layer param numbers: {pointwise_layer_n_params}")
print(pointwise_out.shape)
print(f"\nConv params: {conv_layer_n_params} / Depthwise-separable
运行结果:
2.利用torch.nn.Conv2d来做图像滤波
类似于cv2.filter2D 和 ndimage.convolve
im_enhance5_torch 一次处理单通道。
im_enhance5_torch_depthwise 可以同时处理多通道。
2.1 pytorch 滤波函数,自定义滤波核
def im_enhance5_torch(img, amount = 1.6):
"""
img: h, w raw float32 torch.tensor
:param amount: 1-3
:return:
"""
#print(img.shape)
h, w = 400, 400
img = img.reshape(1, 1, h, w)
sharp_kernel = torch.FloatTensor([-0.0008, -0.0018, -0.0034, -0.0050, -0.0056, -0.0050, -0.0034, -0.0018, -0.0008
, -0.0018, -0.0044, -0.0082, -0.0119, -0.0135, -0.0119, -0.0082, -0.0044, -0.0018
, -0.0034, -0.0082, -0.0153, -0.0223, -0.0253, -0.0223, -0.0153, -0.0082, -0.0034
, -0.0050, -0.0119, -0.0223, -0.0325, -0.0368, -0.0325, -0.0223, -0.0119, -0.0050
, -0.0056, -0.0135, -0.0253, -0.0368, 0.9583, -0.0368, -0.0253, -0.0135, -0.0056
, -0.0050, -0.0119, -0.0223, -0.0325, -0.0368, -0.0325, -0.0223, -0.0119, -0.0050
, -0.0034, -0.0082, -0.0153, -0.0223, -0.0253, -0.0223, -0.0153, -0.0082, -0.0034
, -0.0018, -0.0044, -0.0082, -0.0119, -0.0135, -0.0119, -0.0082, -0.0044, -0.0018
, -0.0008, -0.0018, -0.0034, -0.0050, -0.0056, -0.0050, -0.0034, -0.0018,
-0.0008]).reshape(9, 9)
sharp_kernel = amount * sharp_kernel
sharp_kernel[4, 4] += 1
sharp_kernel = sharp_kernel.reshape(1, 1, 9, 9)
conv2d = torch.nn.Conv2d(1, 1, (9, 9), bias=False, padding=4, padding_mode='reflect') # 设置卷积网络
conv2d.weight.data = sharp_kernel # 初始化weight
conv2d.weight.requires_grad = False
im_dst = conv2d(img)
im_dst = im_dst.reshape(h, w)
# print(im_dst.shape)
# print(im_dst[90:100, 90:100])
return im_dst
def im_enhance5_torch_depthwise(img, amount = 1.6):
"""
# 输入[1,4, 400,400], kernel [4, 1, 9, 9] group : 4 , 输出 1, 4, 400, 400
img: h, w raw float32 torch.tensor
:param amount: 1-3
:return:
"""
#print(img.shape)
c, h, w = img.shape[0], img.shape[1], img.shape[2] # 4,400,400
img = img.reshape(1, c, h, w)
sharp_kernel = torch.FloatTensor([-0.0008, -0.0018, -0.0034, -0.0050, -0.0056, -0.0050, -0.0034, -0.0018, -0.0008
, -0.0018, -0.0044, -0.0082, -0.0119, -0.0135, -0.0119, -0.0082, -0.0044, -0.0018
, -0.0034, -0.0082, -0.0153, -0.0223, -0.0253, -0.0223, -0.0153, -0.0082, -0.0034
, -0.0050, -0.0119, -0.0223, -0.0325, -0.0368, -0.0325, -0.0223, -0.0119, -0.0050
, -0.0056, -0.0135, -0.0253, -0.0368, 0.9583, -0.0368, -0.0253, -0.0135, -0.0056
, -0.0050, -0.0119, -0.0223, -0.0325, -0.0368, -0.0325, -0.0223, -0.0119, -0.0050
, -0.0034, -0.0082, -0.0153, -0.0223, -0.0253, -0.0223, -0.0153, -0.0082, -0.0034
, -0.0018, -0.0044, -0.0082, -0.0119, -0.0135, -0.0119, -0.0082, -0.0044, -0.0018
, -0.0008, -0.0018, -0.0034, -0.0050, -0.0056, -0.0050, -0.0034, -0.0018,
-0.0008]).reshape(9, 9)
sharp_kernel = amount * sharp_kernel
sharp_kernel[4, 4] += 1
sharp_kernel = sharp_kernel.reshape(1, 1, 9, 9)
sharp_kernel = torch.tile(sharp_kernel, (4, 1, 1, 1))
print('sharp shape:', img.shape, sharp_kernel.shape) # 图像尺寸和kernel尺寸, 每个通道分别卷积, 和 cv2.filter2d, ndimage.convolve 类似
conv2d = torch.nn.Conv2d(4, 4, (9, 9), bias=False, padding=4, padding_mode='reflect', groups=4) # 设置卷积网络
conv2d.weight.data = sharp_kernel # 初始化weight
conv2d.weight.requires_grad = False
im_dst = conv2d(img)
im_dst = im_dst.reshape(c, h, w)
# print(im_dst.shape)
# print(im_dst[90:100, 90:100])
return im_dst
2.2 利用cupy做滤波
enhance_apply_YUV 是先将rgb转化为yuv,再对Y通道做滤波,再转换回rgb
enhance_apply_channel 是分别对每个通道同样的滤波处理。
import cv2
import numpy as np
import cupy as cp
from cupyx.scipy import ndimage
def get_thr_sharped_edge_mask_xy_cupy(img, index_shape=0.25):
h, w = img.shape
edge_mask = cp.ones((h, w), dtype=cp.float32)
dy, dx = cp.gradient(img)
ksize = 5
kernel = cp.ones((ksize, ksize)) / 25
dxm = ndimage.convolve(dx, kernel, mode='reflect')
dym = ndimage.convolve(dy, kernel, mode='reflect')
dxy = cp.abs(dxm + dym)
edge_mask[dxy < index_shape] = dxy[dxy < index_shape] / index_shape
return edge_mask
def im_enhance5_cupy(img, thr_sharp = 0, amount = 1.6):
"""
:param thr_sharp: 0-0.5
:param amount: 1-3
:return:
"""
img = cp.asarray(img)
sharp_kernel = cp.array([-0.0008, -0.0018, -0.0034, -0.0050, -0.0056, -0.0050, -0.0034, -0.0018, -0.0008
, -0.0018, -0.0044, -0.0082, -0.0119, -0.0135, -0.0119, -0.0082, -0.0044, -0.0018
, -0.0034, -0.0082, -0.0153, -0.0223, -0.0253, -0.0223, -0.0153, -0.0082, -0.0034
, -0.0050, -0.0119, -0.0223, -0.0325, -0.0368, -0.0325, -0.0223, -0.0119, -0.0050
, -0.0056, -0.0135, -0.0253, -0.0368, 0.9583, -0.0368, -0.0253, -0.0135, -0.0056
, -0.0050, -0.0119, -0.0223, -0.0325, -0.0368, -0.0325, -0.0223, -0.0119, -0.0050
, -0.0034, -0.0082, -0.0153, -0.0223, -0.0253, -0.0223, -0.0153, -0.0082, -0.0034
, -0.0018, -0.0044, -0.0082, -0.0119, -0.0135, -0.0119, -0.0082, -0.0044, -0.0018
, -0.0008, -0.0018, -0.0034, -0.0050, -0.0056, -0.0050, -0.0034, -0.0018,
-0.0008]).reshape(9, 9)
if thr_sharp > 0:
mask = get_thr_sharped_edge_mask_xy_cupy(img, index_shape=0.25)
# print(mask[90:100, 90:100])
im_dst = ndimage.convolve(img, sharp_kernel, mode='reflect')
# print(im_dst[90:100, 90:100])
im_dst = im_dst * mask # 高频,去除噪声(mask<1判为噪声)
im_dst = img + amount * im_dst
else:
sharp_kernel = amount * sharp_kernel
sharp_kernel[4, 4] += 1
im_dst = ndimage.convolve(img, sharp_kernel, mode='reflect')
# print(im_dst[90:100, 90:100])
return im_dst
def enhance_apply_YUV(img, amount=1.6):
# img = cv2.imread(r'D:\savedmodel\train_ret\02100_psnr_38.75_out_pair.png')
yuv = cv2.cvtColor(img, cv2.COLOR_BGR2YCrCb)
y2 = im_enhance5_cupy(cp.asarray(yuv[..., 0]), 0, amount)
yuv[..., 0] = cp.asnumpy(y2)
img2 = cv2.cvtColor(yuv, cv2.COLOR_YCrCb2BGR)
return img2
def enhance_apply_channel(img, amount):
h, w, c = img.shape
img2 = img.copy()
for iii in range(c):
img2[..., iii] = cp.asnumpy(im_enhance5_cupy(cp.asarray(img[..., iii]), 0, amount))
return img2
def usm_cupy(raw_4ch, detail_alpha=1):
"""
:param raw_4ch: h,w,4, uint16
:return:
"""
h, w, c = raw_4ch.shape
raw_4ch = raw_4ch.astype(cp.float32)
raw_4ch_blur = raw_4ch.copy()
siz = 3
hh = cp.ones([siz, siz]) / (siz * siz)
for i in range(c):
raw_4ch_blur[..., i] = ndimage.convolve(raw_4ch[..., i], hh, mode='nearest')
raw_4ch[..., i] = raw_4ch[..., i] + detail_alpha * (raw_4ch[..., i] - raw_4ch_blur[..., i])
siz = 11
hh = cp.ones([siz, siz]) / (siz * siz)
for i in range(c):
raw_4ch_blur[..., i] = ndimage.convolve(raw_4ch[..., i], hh, mode='nearest')
raw_4ch[..., i] = raw_4ch[..., i] + detail_alpha * (raw_4ch[..., i] - raw_4ch_blur[..., i])
return raw_4ch #np.clip(raw_4ch, 0, 1023).astype(np.uint16)
2.3比较效果
这里读取的是 raw图像,然后转化为4通道,
然后用pytorch的方法,cupy(ndimage.convolve方法)得到结果,结果是一致的
if __name__ == "__main__":
dir = r'D:\dataset\rawimage'
files = glob.glob(os.path.join(dir, '*.raw'))
print('len:', len(files))
files = sorted(files)
index = 6 * 202
file = files[index]
h, w = 800, 800
raw = np.fromfile(file, dtype=np.uint16).reshape(h, w).astype(np.float32) / 1023
# raw = np.load(file).astype(np.float32) / 1023
h, w = raw.shape
raw4 = raw.reshape(h // 2, 2, w // 2, 2).transpose(0, 2, 1, 3).reshape(h // 2, w // 2, 4)
img = torch.from_numpy(raw4.copy())
print(img.shape)
for i in range(4):
img[..., i] = im_enhance5_torch(img[..., i], 1.2)
img = img.detach().numpy()
img2 = torch.from_numpy(raw4.copy()).permute(2, 0, 1)
print(img2.shape)
img2 = im_enhance5_torch_depthwise(img2, 1.2)
img2 = img2.detach().numpy().transpose(1, 2, 0)
raw4_enhance = enhance_apply_channel(raw4, 1.2)
print(raw4.shape, img.shape, raw4_enhance.shape)
black_level = 16
white_level = 1023
wb = get_wb()
ccm = get_ccm()
srgb1 = apply_wb_ccm(img[..., [0, 1, 3]] * 1023, wb, ccm, 16, 1023)
srgb2 = apply_wb_ccm(raw4_enhance[..., [0, 1, 3]] * 1023, wb, ccm, 16, 1023)
srgb3 = apply_wb_ccm(img2[..., [0, 1, 3]] * 1023, wb, ccm, 16, 1023)
srgb1 = np.clip(srgb1 * 255 , 0 ,255).astype(np.uint8)
srgb2 = np.clip(srgb2 * 255, 0, 255).astype(np.uint8)
srgb3 = np.clip(srgb3 * 255, 0, 255).astype(np.uint8)
cv2.imwrite('im_enhance5_torch_1.png', srgb1[..., ::-1])
cv2.imwrite('im_enhance5_torch_2.png', srgb2[..., ::-1])
cv2.imwrite('im_enhance5_torch_3.png', srgb3[..., ::-1])
[1]https://faun.pub/depthwise-separable-convolutions-in-pytorch-fd41a97327d0