Showing preview only (1,634K chars total). The displayed content is truncated. Use the JSON API for full output.
Repository: MC-E/Deep-Generalized-Unfolding-Networks-for-Image-Restoration
Branch: main
Commit: bae845c2612d
Files: 197
Total size: 1.5 MB
Directory structure:
gitextract_g5u1c6i_/
├── Compressive-Sensing/
│ ├── DGUNet.py
│ ├── DGUNet_deblock.py
│ ├── DGUNet_plus.py
│ ├── DGUNet_plus_deblock.py
│ ├── Datasets/
│ │ └── README.md
│ ├── README.md
│ ├── config.py
│ ├── csdata_fast.py
│ ├── csdata_fast2.py
│ ├── losses.py
│ ├── sampling_matrix_new/
│ │ ├── README.md
│ │ ├── phi_sampling_10_32x32.npy
│ │ ├── phi_sampling_1_32x32.npy
│ │ ├── phi_sampling_25_32x32.npy
│ │ ├── phi_sampling_30_32x32.npy
│ │ ├── phi_sampling_40_32x32.npy
│ │ ├── phi_sampling_4_32x32.npy
│ │ └── phi_sampling_50_32x32.npy
│ ├── test.py
│ ├── test_plus.py
│ ├── train.py
│ ├── train_deblock.py
│ ├── train_plus.py
│ ├── train_plus_deblock.py
│ ├── training.yml
│ ├── training_cs.yml
│ └── utils/
│ ├── README.md
│ ├── __init__.py
│ ├── dataset_utils.py
│ ├── dir_utils.py
│ ├── image_utils.py
│ └── model_utils.py
├── Deblurring/
│ ├── .ipynb_checkpoints/
│ │ ├── DGUNet-checkpoint.py
│ │ ├── DGUNet_plus-checkpoint.py
│ │ ├── MPRNet_all_4gpu-checkpoint.out
│ │ ├── MPRNet_nochop_s5_v7_re-checkpoint.py
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_baseSAM-checkpoint.py
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft-checkpoint.py
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft_4gpu-checkpoint.out
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft_4gpu_2-checkpoint.out
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft_rp_4gpu_2-checkpoint.out
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft_v2-checkpoint.py
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft_v2_4gpu-checkpoint.out
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft_v2_all-checkpoint.out
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_rp_base-checkpoint.py
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_rp_base_sft_v2-checkpoint.py
│ │ ├── README-checkpoint.md
│ │ ├── config-checkpoint.py
│ │ ├── data_RGB-checkpoint.py
│ │ ├── dataset_RGB-checkpoint.py
│ │ ├── deblure_s7_v7_h3_u4_base-checkpoint.out
│ │ ├── deblure_s7_v7_h3_u4_rp_base-checkpoint.out
│ │ ├── deblure_s7_v7_h3_u4_rp_base_4gpu-checkpoint.out
│ │ ├── deblure_s7_v7_h3_u4_rp_base_4gpu_2e-4-checkpoint.out
│ │ ├── deblure_s7_v7_h3_u4_rp_base_4gpu_2e-4_b16-checkpoint.out
│ │ ├── losses-checkpoint.py
│ │ ├── test-checkpoint.py
│ │ ├── test_plus-checkpoint.py
│ │ ├── train-checkpoint.py
│ │ ├── train_plus-checkpoint.py
│ │ └── training-checkpoint.yml
│ ├── DGUNet.py
│ ├── DGUNet_plus.py
│ ├── Datasets/
│ │ └── README.md
│ ├── README.md
│ ├── config.py
│ ├── data_RGB.py
│ ├── dataset_RGB.py
│ ├── evaluate_GOPRO_HIDE.m
│ ├── evaluate_RealBlur.py
│ ├── losses.py
│ ├── test.py
│ ├── test_plus.py
│ ├── train.py
│ ├── train_plus.py
│ ├── training.yml
│ └── utils/
│ ├── __init__.py
│ ├── dataset_utils.py
│ ├── dir_utils.py
│ ├── image_utils.py
│ └── model_utils.py
├── Denoising/
│ ├── .ipynb_checkpoints/
│ │ ├── DGUNet_plus-checkpoint.py
│ │ ├── config-checkpoint.py
│ │ ├── test-checkpoint.py
│ │ ├── test_plus-checkpoint.py
│ │ ├── train-checkpoint.py
│ │ ├── train_plus-checkpoint.py
│ │ └── training-checkpoint.yml
│ ├── DGUNet.py
│ ├── DGUNet_plus.py
│ ├── Datasets/
│ │ └── README.md
│ ├── README.md
│ ├── config.py
│ ├── data_RGB.py
│ ├── dataset_RGB.py
│ ├── demo_DND.py
│ ├── generate_patches_SIDD.py
│ ├── losses.py
│ ├── test.py
│ ├── test_DND.py
│ ├── test_DND_plus.py
│ ├── test_plus.py
│ ├── train.py
│ ├── train_plus.py
│ ├── training.yml
│ └── utils/
│ ├── .ipynb_checkpoints/
│ │ ├── __init__-checkpoint.py
│ │ └── model_utils-checkpoint.py
│ ├── __init__.py
│ ├── dataset_utils.py
│ ├── dir_utils.py
│ ├── image_utils.py
│ └── model_utils.py
├── Deraining/
│ ├── .ipynb_checkpoints/
│ │ ├── DGUNet-checkpoint.py
│ │ ├── DGUNet_plus-checkpoint.py
│ │ ├── MPRNet_nochop_s3-checkpoint.py
│ │ ├── MPRNet_nochop_s3_v7_hi3_u4-checkpoint.py
│ │ ├── MPRNet_nochop_s5-checkpoint.py
│ │ ├── MPRNet_nochop_s5_dense-checkpoint.py
│ │ ├── MPRNet_nochop_s5_dense_v3-checkpoint.py
│ │ ├── MPRNet_nochop_s5_rb-checkpoint.py
│ │ ├── MPRNet_nochop_s5_sk-checkpoint.py
│ │ ├── MPRNet_nochop_s5_sk_dense-checkpoint.py
│ │ ├── MPRNet_nochop_s5_v3-checkpoint.py
│ │ ├── MPRNet_nochop_s5_v7-checkpoint.py
│ │ ├── MPRNet_nochop_s5_v7_hi3_u4-checkpoint.py
│ │ ├── MPRNet_nochop_s5_v7_re-checkpoint.py
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4-checkpoint.py
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_baseSAM-checkpoint.py
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft-checkpoint.py
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft_v2-checkpoint.py
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft_v2_sf-checkpoint.py
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_nointer-checkpoint.py
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_nophi-checkpoint.py
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_rp_base-checkpoint.py
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_rp_base_sft_v2-checkpoint.py
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_rp_base_v2-checkpoint.py
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_rp_base_v3-checkpoint.py
│ │ ├── MPRNet_nochop_s9_v7_hi3_u4_baseSAM_sft_v2-checkpoint.py
│ │ ├── MPRNet_v2-checkpoint.py
│ │ ├── MPRNet_v3-checkpoint.py
│ │ ├── MPRNet_v4-checkpoint.py
│ │ ├── MPRNet_v5-checkpoint.py
│ │ ├── MPRNet_v6-checkpoint.py
│ │ ├── data_RGB-checkpoint.py
│ │ ├── dataset_RGB-checkpoint.py
│ │ ├── re_ceckpoint-checkpoint.py
│ │ ├── test-checkpoint.py
│ │ ├── test_plus-checkpoint.py
│ │ ├── train-checkpoint.py
│ │ ├── train2-checkpoint.py
│ │ ├── train_MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft-checkpoint.out
│ │ ├── train_MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft_v2-checkpoint.out
│ │ ├── train_MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft_v2_scale1-checkpoint.out
│ │ ├── train_MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft_v2_scale2-checkpoint.out
│ │ ├── train_MPRNet_nochop_s7_v7_hi3_u4_rp_baseSAM_sft_v2-checkpoint.out
│ │ ├── train_MPRNet_nochop_s9_v7_hi3_u4_baseSAM_sft_v2-checkpoint.out
│ │ ├── train_plus-checkpoint.py
│ │ └── training-checkpoint.yml
│ ├── DGUNet.py
│ ├── DGUNet_plus.py
│ ├── Datasets/
│ │ ├── .ipynb_checkpoints/
│ │ │ └── README-checkpoint.md
│ │ └── README.md
│ ├── README.md
│ ├── SKmodel.py
│ ├── config.py
│ ├── data_RGB.py
│ ├── dataset_RGB.py
│ ├── evaluate_PSNR_SSIM.m
│ ├── losses.py
│ ├── test.py
│ ├── test_plus.py
│ ├── train.py
│ ├── train_plus.py
│ ├── training.yml
│ ├── utils/
│ │ ├── __init__.py
│ │ ├── dataset_utils.py
│ │ ├── dir_utils.py
│ │ ├── image_utils.py
│ │ └── model_utils.py
│ └── warmup_scheduler.egg-info/
│ ├── PKG-INFO
│ ├── SOURCES.txt
│ ├── dependency_links.txt
│ └── top_level.txt
├── README.md
├── figs/
│ └── README.md
└── pytorch-gradual-warmup-lr/
├── build/
│ └── lib/
│ └── warmup_scheduler/
│ ├── __init__.py
│ ├── run.py
│ └── scheduler.py
├── dist/
│ └── warmup_scheduler-0.3-py3.6.egg
├── setup.py
├── warmup_scheduler/
│ ├── __init__.py
│ ├── run.py
│ └── scheduler.py
└── warmup_scheduler.egg-info/
├── PKG-INFO
├── SOURCES.txt
├── dependency_links.txt
└── top_level.txt
================================================
FILE CONTENTS
================================================
================================================
FILE: Compressive-Sensing/DGUNet.py
================================================
import torch
import torch.nn as nn
import torch.nn.functional as F
from pdb import set_trace as stx
# from pdb import set_trace as stx
from torch.nn import init
import numpy as np
##########################################################################
def conv(in_channels, out_channels, kernel_size, bias=False, stride=1):
return nn.Conv2d(
in_channels, out_channels, kernel_size,
padding=(kernel_size // 2), bias=bias, stride=stride)
def conv_down(in_chn, out_chn, bias=False):
layer = nn.Conv2d(in_chn, out_chn, kernel_size=4, stride=2, padding=1, bias=bias)
return layer
##########################################################################
## Channel Attention Layer
class CALayer(nn.Module):
def __init__(self, channel, reduction=16, bias=False):
super(CALayer, self).__init__()
# global average pooling: feature --> point
self.avg_pool = nn.AdaptiveAvgPool2d(1)
# feature channel downscale and upscale --> channel weight
self.conv_du = nn.Sequential(
nn.Conv2d(channel, channel // reduction, 1, padding=0, bias=bias),
nn.ReLU(inplace=True),
nn.Conv2d(channel // reduction, channel, 1, padding=0, bias=bias),
nn.Sigmoid()
)
def forward(self, x):
y = self.avg_pool(x) # 计算注意力权重
y = self.conv_du(y)
return x * y # 加权
##########################################################################
## Channel Attention Block (CAB)
class CAB(nn.Module):
def __init__(self, n_feat, kernel_size, reduction, bias, act):
super(CAB, self).__init__()
modules_body = []
modules_body.append(conv(n_feat, n_feat, kernel_size, bias=bias))
modules_body.append(act)
modules_body.append(conv(n_feat, n_feat, kernel_size, bias=bias))
self.CA = CALayer(n_feat, reduction, bias=bias)
self.body = nn.Sequential(*modules_body)
def forward(self, x):
res = self.body(x)
res = self.CA(res)
res += x
return res
##########################################################################
## Supervised Attention Module
class SAM(nn.Module):
def __init__(self, n_feat, kernel_size, bias,in_c):
super(SAM, self).__init__()
self.conv1 = conv(n_feat, n_feat, kernel_size, bias=bias)
self.conv2 = conv(n_feat, in_c, kernel_size, bias=bias)
self.conv3 = conv(in_c, n_feat, kernel_size, bias=bias)
def forward(self, x, x_img):
x1 = self.conv1(x)
img = self.conv2(x) + x_img
x2 = torch.sigmoid(self.conv3(img))
x1 = x1 * x2
x1 = x1 + x
return x1, img
class UNetConvBlock(nn.Module):
def __init__(self, in_size, out_size, downsample, relu_slope, use_csff=False, use_HIN=False):
super(UNetConvBlock, self).__init__()
self.downsample = downsample
self.identity = nn.Conv2d(in_size, out_size, 1, 1, 0)
self.use_csff = use_csff
self.conv_1 = nn.Conv2d(in_size, out_size, kernel_size=3, padding=1, bias=True)
self.relu_1 = nn.LeakyReLU(relu_slope, inplace=False)
self.conv_2 = nn.Conv2d(out_size, out_size, kernel_size=3, padding=1, bias=True)
self.relu_2 = nn.LeakyReLU(relu_slope, inplace=False)
if downsample and use_csff:
self.csff_enc = nn.Conv2d(out_size, out_size, 3, 1, 1)
self.csff_dec = nn.Conv2d(in_size, out_size, 3, 1, 1)
if use_HIN:
self.norm = nn.InstanceNorm2d(out_size//2, affine=True)
self.use_HIN = use_HIN
if downsample:
self.downsample = conv_down(out_size, out_size, bias=False)
def forward(self, x, enc=None, dec=None):
out = self.conv_1(x)
if self.use_HIN:
out_1, out_2 = torch.chunk(out, 2, dim=1)
out = torch.cat([self.norm(out_1), out_2], dim=1)
out = self.relu_1(out)
out = self.relu_2(self.conv_2(out))
out += self.identity(x)
if enc is not None and dec is not None:
assert self.use_csff
# print('enc+dec:',out.shape,enc.shape,dec.shape)
out = out + self.csff_enc(enc) + self.csff_dec(dec)
if self.downsample:
out_down = self.downsample(out)
return out_down, out
else:
return out
class UNetUpBlock(nn.Module):
def __init__(self, in_size, out_size, relu_slope):
super(UNetUpBlock, self).__init__()
self.up = nn.ConvTranspose2d(in_size, out_size, kernel_size=2, stride=2, bias=True)
self.conv_block = UNetConvBlock(out_size*2, out_size, False, relu_slope)
def forward(self, x, bridge):
up = self.up(x)
out = torch.cat([up, bridge], 1)
out = self.conv_block(out)
return out
##########################################################################
## U-Net
class Encoder(nn.Module):
def __init__(self, n_feat, kernel_size, reduction, act, bias, scale_unetfeats, csff,depth=4):
super(Encoder, self).__init__()
self.body=nn.ModuleList()#[]
self.depth=depth
for i in range(depth-1):
self.body.append(UNetConvBlock(in_size=n_feat+scale_unetfeats*i, out_size=n_feat+scale_unetfeats*(i+1), downsample=True, relu_slope=0.2, use_csff=True, use_HIN=True))
self.body.append(UNetConvBlock(in_size=n_feat+scale_unetfeats*(depth-1), out_size=n_feat+scale_unetfeats*(depth-1), downsample=False, relu_slope=0.2, use_csff=True, use_HIN=True))
def forward(self, x, encoder_outs=None, decoder_outs=None):
res=[]
if encoder_outs is not None and decoder_outs is not None:
for i,down in enumerate(self.body):
if (i+1) < self.depth:
x, x_up = down(x,encoder_outs[i],decoder_outs[-i-1])
res.append(x_up)
else:
x = down(x)
else:
for i,down in enumerate(self.body):
if (i+1) < self.depth:
x, x_up = down(x)
res.append(x_up)
else:
x = down(x)
return res,x
class Decoder(nn.Module):
def __init__(self, n_feat, kernel_size, reduction, act, bias, scale_unetfeats,depth=4):
super(Decoder, self).__init__()
self.body=nn.ModuleList()
self.skip_conv=nn.ModuleList()#[]
for i in range(depth-1):
self.body.append(UNetUpBlock(in_size=n_feat+scale_unetfeats*(depth-i-1), out_size=n_feat+scale_unetfeats*(depth-i-2), relu_slope=0.2))
self.skip_conv.append(nn.Conv2d(n_feat+scale_unetfeats*(depth-i-1), n_feat+scale_unetfeats*(depth-i-2), 3, 1, 1))
def forward(self, x, bridges):
res=[]
for i,up in enumerate(self.body):
x=up(x,self.skip_conv[i](bridges[-i-1]))
res.append(x)
return res
##########################################################################
##---------- Resizing Modules ----------
class DownSample(nn.Module):
def __init__(self, in_channels, s_factor):
super(DownSample, self).__init__()
self.down = nn.Sequential(nn.Upsample(scale_factor=0.5, mode='bilinear', align_corners=False),
nn.Conv2d(in_channels, in_channels + s_factor, 1, stride=1, padding=0, bias=False))
def forward(self, x):
x = self.down(x)
return x
class UpSample(nn.Module):
def __init__(self, in_channels, out_channels):
super(UpSample, self).__init__()
self.up = nn.Sequential(nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
nn.Conv2d(in_channels, out_channels, 1, stride=1, padding=0, bias=False))
def forward(self, x):
x = self.up(x)
return x
class SkipUpSample(nn.Module):
def __init__(self, in_channels, s_factor):
super(SkipUpSample, self).__init__()
self.up = nn.Sequential(nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
nn.Conv2d(in_channels + s_factor, in_channels, 1, stride=1, padding=0, bias=False))
def forward(self, x, y):
x = self.up(x)
x = x + y
return x
##########################################################################
## Original Resolution Block (ORB)
class ORB(nn.Module):
def __init__(self, n_feat, kernel_size, reduction, act, bias, num_cab):
super(ORB, self).__init__()
modules_body = []
modules_body = [CAB(n_feat, kernel_size, reduction, bias=bias, act=act) for _ in range(num_cab)]
modules_body.append(conv(n_feat, n_feat, kernel_size))
self.body = nn.Sequential(*modules_body)
def forward(self, x):
res = self.body(x)
res += x
return res
##########################################################################
class ORSNet(nn.Module):
def __init__(self, n_feat, scale_orsnetfeats, kernel_size, reduction, act, bias, scale_unetfeats, num_cab):
super(ORSNet, self).__init__()
self.orb1 = ORB(n_feat + scale_orsnetfeats, kernel_size, reduction, act, bias, num_cab)
self.orb2 = ORB(n_feat + scale_orsnetfeats, kernel_size, reduction, act, bias, num_cab)
self.orb3 = ORB(n_feat + scale_orsnetfeats, kernel_size, reduction, act, bias, num_cab)
self.up_enc1 = UpSample(n_feat+2*scale_unetfeats, n_feat+scale_unetfeats)
self.up_dec1 = UpSample(n_feat+scale_unetfeats, n_feat)
self.up_enc2 = nn.Sequential(UpSample(n_feat + 3*scale_unetfeats, n_feat+2*scale_unetfeats),
UpSample(n_feat+2*scale_unetfeats, n_feat+scale_unetfeats))
self.up_dec2 = nn.Sequential(UpSample(n_feat + 2*scale_unetfeats, n_feat+scale_unetfeats),
UpSample(n_feat+scale_unetfeats, n_feat))
self.conv_enc1 = nn.Conv2d(n_feat+scale_unetfeats, n_feat + scale_orsnetfeats, kernel_size=1, bias=bias)
self.conv_enc2 = nn.Conv2d(n_feat+scale_unetfeats, n_feat + scale_orsnetfeats, kernel_size=1, bias=bias)
self.conv_enc3 = nn.Conv2d(n_feat+scale_unetfeats, n_feat + scale_orsnetfeats, kernel_size=1, bias=bias)
self.conv_dec1 = nn.Conv2d(n_feat, n_feat + scale_orsnetfeats, kernel_size=1, bias=bias)
self.conv_dec2 = nn.Conv2d(n_feat, n_feat + scale_orsnetfeats, kernel_size=1, bias=bias)
self.conv_dec3 = nn.Conv2d(n_feat, n_feat + scale_orsnetfeats, kernel_size=1, bias=bias)
def forward(self, x, encoder_outs, decoder_outs):
x = self.orb1(x)
x = x + self.conv_enc1(encoder_outs[0]) + self.conv_dec1(decoder_outs[-1])
x = self.orb2(x)
x = x + self.conv_enc2(self.up_enc1(encoder_outs[1])) + self.conv_dec2(self.up_dec1(decoder_outs[-2]))
x = self.orb3(x)
x = x + self.conv_enc3(self.up_enc2(encoder_outs[2])) + self.conv_dec3(self.up_dec2(decoder_outs[-3]))
return x
def default_conv(in_channels, out_channels, kernel_size,stride=1, bias=True):
return nn.Conv2d(
in_channels, out_channels, kernel_size,
padding=(kernel_size//2),stride=stride, bias=bias)
class ResBlock(nn.Module):
def __init__(
self, conv, n_feats, kernel_size,
bias=True, bn=False, act=nn.PReLU(), res_scale=1):
super(ResBlock, self).__init__()
m = []
for i in range(2):
if i == 0:
m.append(conv(n_feats, 64, kernel_size, bias=bias))
else:
m.append(conv(64, n_feats, kernel_size, bias=bias))
if bn:
m.append(nn.BatchNorm2d(n_feats))
if i == 0:
m.append(act)
self.body = nn.Sequential(*m)
self.res_scale = res_scale
def forward(self, x):
res = self.body(x).mul(self.res_scale)
res += x
return res
##########################################################################
class MPRBlock(nn.Module):
def __init__(self, in_c=3, out_c=3, n_feat=80, scale_unetfeats=48, kernel_size=3,
reduction=4, bias=False):
super(MPRBlock, self).__init__()
act = nn.PReLU()
self.shallow_feat = nn.Sequential(conv(in_c, n_feat, kernel_size, bias=bias),
CAB(n_feat, kernel_size, reduction, bias=bias, act=act))
# Cross Stage Feature Fusion (CSFF)
self.stage_encoder = Encoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats, csff=True)
self.stage_decoder = Decoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats)
self.sam = SAM(n_feat, kernel_size=1, bias=bias,in_c=in_c)
self.phi = ResBlock(default_conv, 3, 3)
self.phit = ResBlock(default_conv, 3, 3)
self.r = nn.Parameter(torch.Tensor([0.5]))
self.concat = conv(n_feat * 2, n_feat, kernel_size, bias=bias)
def forward(self, stage_img, img,x_samfeats,f_encoder,f_decoder,PhiTPhi,PhiTb):
b,c,w,h = stage_img.shape
x_k_1=stage_img.view(b,-1)
# compute r_k
x = x_k_1 - self.r * torch.mm(x_k_1, PhiTPhi)
r_k = x + self.r * PhiTb
r_k = r_k.view(b,c,w,h)
# compute x_k
x = self.shallow_feat(r_k)
## Concatenate SAM features of Stage 1 with shallow features of Stage 2
x_cat = self.concat(torch.cat([x, x_samfeats], 1))
## Process features of both patches with Encoder of Stage 2
feat1, f_encoder = self.stage_encoder(x_cat, f_encoder, f_decoder)
## Pass features through Decoder of Stage 2
f_decoder = self.stage_decoder(f_encoder, feat1)
## Apply SAM
x_samfeats, stage_img = self.sam(f_decoder[-1], r_k)
return stage_img, x_samfeats,feat1 ,f_decoder
##########################################################################
class DGUNet(nn.Module):
def __init__(self, in_c=3, out_c=3, n_feat=32, scale_unetfeats=16, scale_orsnetfeats=16, num_cab=8, kernel_size=3,
reduction=4, bias=False, nums_stages=5, cs_ratio=25):
super(DGUNet, self).__init__()
print('CS Ratio: ',cs_ratio)
self.Phi = nn.Parameter(init.xavier_normal_(torch.Tensor(np.ceil(cs_ratio*0.01*1024.).astype(np.int), 1024)))
act = nn.PReLU()
self.shallow_feat1 = nn.Sequential(conv(in_c, n_feat, kernel_size, bias=bias),
CAB(n_feat, kernel_size, reduction, bias=bias, act=act))
self.body = nn.ModuleList()
self.nums_stages = nums_stages
self.stage_model=MPRBlock(
in_c=in_c, out_c=in_c, n_feat=n_feat, scale_unetfeats=scale_unetfeats, kernel_size=kernel_size,
reduction=reduction, bias=bias
)
self.shallow_feat_final = nn.Sequential(conv(in_c, n_feat, kernel_size, bias=bias),
CAB(n_feat, kernel_size, reduction, bias=bias, act=act))
# Cross Stage Feature Fusion (CSFF)
self.stage1_encoder = Encoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats, csff=False)
self.stage1_decoder = Decoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats)
self.sam12 = SAM(n_feat, kernel_size=1, bias=bias,in_c=in_c)
self.r0 = nn.Parameter(torch.Tensor([0.5]))
self.r_final = nn.Parameter(torch.Tensor([0.5]))
self.concat_final = conv(n_feat * 2, n_feat + scale_orsnetfeats, kernel_size, bias=bias)
self.tail = conv(n_feat + scale_orsnetfeats, out_c, kernel_size, bias=bias)
def forward(self, img):
output_=[]
##-------------------------------------------
##-------------- Stage 1---------------------
##-------------------------------------------
b, c, w, h = img.shape
PhiTPhi = torch.mm(torch.transpose(self.Phi, 0, 1), self.Phi) # torch.mm(Phix, Phi)
Phix = torch.mm(img.view(b,-1), torch.transpose(self.Phi, 0, 1)) # compression result
PhiTb = torch.mm(Phix,self.Phi)
# compute r_0
x_0=PhiTb.view(b,-1)
x = x_0 - self.r0 * torch.mm(x_0, PhiTPhi)
r_0 = x + self.r0 * PhiTb
r_0=r_0.view(b,c,w,h)
# compute x_k
x = self.shallow_feat1(r_0)
## Process features of all 4 patches with Encoder of Stage 1
feat1, f_encoder = self.stage1_encoder(x)
## Pass features through Decoder of Stage 1
f_decoder = self.stage1_decoder(f_encoder,feat1)
## Apply Supervised Attention Module (SAM)
x_samfeats, stage_img = self.sam12(f_decoder[-1], r_0)
output_.append(stage_img)
##-------------------------------------------
##-------------- Stage 2_k-1---------------------
##-------------------------------------------
for _ in range(self.nums_stages):
stage_img, x_samfeats, feat1, f_decoder = self.stage_model(stage_img, img,x_samfeats,feat1,f_decoder,PhiTPhi,PhiTb)
output_.append(stage_img)
##-------------------------------------------
##-------------- Stage k---------------------
##-------------------------------------------
# compute r_k
x_k_1 = stage_img.view(b,-1)
x = x_k_1 - self.r_final * torch.mm(x_k_1, PhiTPhi)
r_k = x + self.r_final * PhiTb
r_k = r_k.view(b,c,w,h)
# compute x_k
x = self.shallow_feat_final(r_k)
## Concatenate SAM features
x_cat = self.concat_final(torch.cat([x, x_samfeats], 1))
stage_img = self.tail(x_cat)+r_k
output_.append(stage_img)
return output_[::-1]
================================================
FILE: Compressive-Sensing/DGUNet_deblock.py
================================================
import torch
import torch.nn as nn
import torch.nn.functional as F
from pdb import set_trace as stx
from torch.nn import init
import numpy as np
##########################################################################
def conv(in_channels, out_channels, kernel_size, bias=False, stride=1):
return nn.Conv2d(
in_channels, out_channels, kernel_size,
padding=(kernel_size // 2), bias=bias, stride=stride)
def conv_down(in_chn, out_chn, bias=False):
layer = nn.Conv2d(in_chn, out_chn, kernel_size=4, stride=2, padding=1, bias=bias)
return layer
##########################################################################
## Channel Attention Layer
class CALayer(nn.Module):
def __init__(self, channel, reduction=16, bias=False):
super(CALayer, self).__init__()
# global average pooling: feature --> point
self.avg_pool = nn.AdaptiveAvgPool2d(1)
# feature channel downscale and upscale --> channel weight
self.conv_du = nn.Sequential(
nn.Conv2d(channel, channel // reduction, 1, padding=0, bias=bias),
nn.ReLU(inplace=True),
nn.Conv2d(channel // reduction, channel, 1, padding=0, bias=bias),
nn.Sigmoid()
)
def forward(self, x):
y = self.avg_pool(x) # 计算注意力权重
y = self.conv_du(y)
return x * y # 加权
##########################################################################
## Channel Attention Block (CAB)
class CAB(nn.Module):
def __init__(self, n_feat, kernel_size, reduction, bias, act):
super(CAB, self).__init__()
modules_body = []
modules_body.append(conv(n_feat, n_feat, kernel_size, bias=bias))
modules_body.append(act)
modules_body.append(conv(n_feat, n_feat, kernel_size, bias=bias))
self.CA = CALayer(n_feat, reduction, bias=bias)
self.body = nn.Sequential(*modules_body)
def forward(self, x):
res = self.body(x)
res = self.CA(res)
res += x
return res
##########################################################################
## Supervised Attention Module
class SAM(nn.Module):
def __init__(self, n_feat, kernel_size, bias,in_c):
super(SAM, self).__init__()
self.conv1 = conv(n_feat, n_feat, kernel_size, bias=bias)
self.conv2 = conv(n_feat, in_c, kernel_size, bias=bias)
self.conv3 = conv(in_c, n_feat, kernel_size, bias=bias)
def forward(self, x, x_img):
x1 = self.conv1(x)
img = self.conv2(x) + x_img
x2 = torch.sigmoid(self.conv3(img))
x1 = x1 * x2
x1 = x1 + x
return x1, img
class UNetConvBlock(nn.Module):
def __init__(self, in_size, out_size, downsample, relu_slope, use_csff=False, use_HIN=False):
super(UNetConvBlock, self).__init__()
self.downsample = downsample
self.identity = nn.Conv2d(in_size, out_size, 1, 1, 0)
self.use_csff = use_csff
self.conv_1 = nn.Conv2d(in_size, out_size, kernel_size=3, padding=1, bias=True)
self.relu_1 = nn.LeakyReLU(relu_slope, inplace=False)
self.conv_2 = nn.Conv2d(out_size, out_size, kernel_size=3, padding=1, bias=True)
self.relu_2 = nn.LeakyReLU(relu_slope, inplace=False)
if downsample and use_csff:
self.csff_enc = nn.Conv2d(out_size, out_size, 3, 1, 1)
self.csff_dec = nn.Conv2d(in_size, out_size, 3, 1, 1)
if use_HIN:
self.norm = nn.InstanceNorm2d(out_size//2, affine=True)
self.use_HIN = use_HIN
if downsample:
self.downsample = conv_down(out_size, out_size, bias=False)
def forward(self, x, enc=None, dec=None):
out = self.conv_1(x)
if self.use_HIN:
out_1, out_2 = torch.chunk(out, 2, dim=1)
out = torch.cat([self.norm(out_1), out_2], dim=1)
out = self.relu_1(out)
out = self.relu_2(self.conv_2(out))
out += self.identity(x)
if enc is not None and dec is not None:
assert self.use_csff
# print('enc+dec:',out.shape,enc.shape,dec.shape)
out = out + self.csff_enc(enc) + self.csff_dec(dec)
if self.downsample:
out_down = self.downsample(out)
return out_down, out
else:
return out
class UNetUpBlock(nn.Module):
def __init__(self, in_size, out_size, relu_slope):
super(UNetUpBlock, self).__init__()
self.up = nn.ConvTranspose2d(in_size, out_size, kernel_size=2, stride=2, bias=True)
self.conv_block = UNetConvBlock(out_size*2, out_size, False, relu_slope)
def forward(self, x, bridge):
# print('merge',x.shape,bridge.shape)
up = self.up(x)
out = torch.cat([up, bridge], 1)
out = self.conv_block(out)
return out
##########################################################################
## U-Net
class Encoder(nn.Module):
def __init__(self, n_feat, kernel_size, reduction, act, bias, scale_unetfeats, csff,depth=4):
super(Encoder, self).__init__()
self.body=nn.ModuleList()#[]
self.depth=depth
for i in range(depth-1):
# downsample = True if (i+1) < depth else False
self.body.append(UNetConvBlock(in_size=n_feat+scale_unetfeats*i, out_size=n_feat+scale_unetfeats*(i+1), downsample=True, relu_slope=0.2, use_csff=True, use_HIN=True))
self.body.append(UNetConvBlock(in_size=n_feat+scale_unetfeats*(depth-1), out_size=n_feat+scale_unetfeats*(depth-1), downsample=False, relu_slope=0.2, use_csff=True, use_HIN=True))
def forward(self, x, encoder_outs=None, decoder_outs=None):
res=[]
if encoder_outs is not None and decoder_outs is not None:
for i,down in enumerate(self.body):
if (i+1) < self.depth:
x, x_up = down(x,encoder_outs[i],decoder_outs[-i-1])
res.append(x_up)
else:
x = down(x)
else:
for i,down in enumerate(self.body):
if (i+1) < self.depth:
x, x_up = down(x)
res.append(x_up)
else:
x = down(x)
return res,x
class Decoder(nn.Module):
def __init__(self, n_feat, kernel_size, reduction, act, bias, scale_unetfeats,depth=4):
super(Decoder, self).__init__()
self.body=nn.ModuleList()
self.skip_conv=nn.ModuleList()#[]
for i in range(depth-1):
self.body.append(UNetUpBlock(in_size=n_feat+scale_unetfeats*(depth-i-1), out_size=n_feat+scale_unetfeats*(depth-i-2), relu_slope=0.2))
self.skip_conv.append(nn.Conv2d(n_feat+scale_unetfeats*(depth-i-1), n_feat+scale_unetfeats*(depth-i-2), 3, 1, 1))
def forward(self, x, bridges):
res=[]
for i,up in enumerate(self.body):
x=up(x,self.skip_conv[i](bridges[-i-1]))
res.append(x)
return res
##########################################################################
##---------- Resizing Modules ----------
class DownSample(nn.Module):
def __init__(self, in_channels, s_factor):
super(DownSample, self).__init__()
self.down = nn.Sequential(nn.Upsample(scale_factor=0.5, mode='bilinear', align_corners=False),
nn.Conv2d(in_channels, in_channels + s_factor, 1, stride=1, padding=0, bias=False))
def forward(self, x):
x = self.down(x)
return x
class UpSample(nn.Module):
def __init__(self, in_channels, out_channels):
super(UpSample, self).__init__()
self.up = nn.Sequential(nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
nn.Conv2d(in_channels, out_channels, 1, stride=1, padding=0, bias=False))
def forward(self, x):
x = self.up(x)
return x
class SkipUpSample(nn.Module):
def __init__(self, in_channels, s_factor):
super(SkipUpSample, self).__init__()
self.up = nn.Sequential(nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
nn.Conv2d(in_channels + s_factor, in_channels, 1, stride=1, padding=0, bias=False))
def forward(self, x, y):
x = self.up(x)
x = x + y
return x
##########################################################################
## Original Resolution Block (ORB)
class ORB(nn.Module):
def __init__(self, n_feat, kernel_size, reduction, act, bias, num_cab):
super(ORB, self).__init__()
modules_body = []
modules_body = [CAB(n_feat, kernel_size, reduction, bias=bias, act=act) for _ in range(num_cab)]
modules_body.append(conv(n_feat, n_feat, kernel_size))
self.body = nn.Sequential(*modules_body)
def forward(self, x):
res = self.body(x)
res += x
return res
##########################################################################
class ORSNet(nn.Module):
def __init__(self, n_feat, scale_orsnetfeats, kernel_size, reduction, act, bias, scale_unetfeats, num_cab):
super(ORSNet, self).__init__()
self.orb1 = ORB(n_feat + scale_orsnetfeats, kernel_size, reduction, act, bias, num_cab)
self.orb2 = ORB(n_feat + scale_orsnetfeats, kernel_size, reduction, act, bias, num_cab)
self.orb3 = ORB(n_feat + scale_orsnetfeats, kernel_size, reduction, act, bias, num_cab)
self.up_enc1 = UpSample(n_feat+2*scale_unetfeats, n_feat+scale_unetfeats)
self.up_dec1 = UpSample(n_feat+scale_unetfeats, n_feat)
self.up_enc2 = nn.Sequential(UpSample(n_feat + 3*scale_unetfeats, n_feat+2*scale_unetfeats),
UpSample(n_feat+2*scale_unetfeats, n_feat+scale_unetfeats))
self.up_dec2 = nn.Sequential(UpSample(n_feat + 2*scale_unetfeats, n_feat+scale_unetfeats),
UpSample(n_feat+scale_unetfeats, n_feat))
self.conv_enc1 = nn.Conv2d(n_feat+scale_unetfeats, n_feat + scale_orsnetfeats, kernel_size=1, bias=bias)
self.conv_enc2 = nn.Conv2d(n_feat+scale_unetfeats, n_feat + scale_orsnetfeats, kernel_size=1, bias=bias)
self.conv_enc3 = nn.Conv2d(n_feat+scale_unetfeats, n_feat + scale_orsnetfeats, kernel_size=1, bias=bias)
self.conv_dec1 = nn.Conv2d(n_feat, n_feat + scale_orsnetfeats, kernel_size=1, bias=bias)
self.conv_dec2 = nn.Conv2d(n_feat, n_feat + scale_orsnetfeats, kernel_size=1, bias=bias)
self.conv_dec3 = nn.Conv2d(n_feat, n_feat + scale_orsnetfeats, kernel_size=1, bias=bias)
def forward(self, x, encoder_outs, decoder_outs):
x = self.orb1(x)
x = x + self.conv_enc1(encoder_outs[0]) + self.conv_dec1(decoder_outs[-1])
x = self.orb2(x)
x = x + self.conv_enc2(self.up_enc1(encoder_outs[1])) + self.conv_dec2(self.up_dec1(decoder_outs[-2]))
x = self.orb3(x)
x = x + self.conv_enc3(self.up_enc2(encoder_outs[2])) + self.conv_dec3(self.up_dec2(decoder_outs[-3]))
return x
def default_conv(in_channels, out_channels, kernel_size,stride=1, bias=True):
return nn.Conv2d(
in_channels, out_channels, kernel_size,
padding=(kernel_size//2),stride=stride, bias=bias)
class ResBlock(nn.Module):
def __init__(
self, conv, n_feats, kernel_size,
bias=True, bn=False, act=nn.PReLU(), res_scale=1):
super(ResBlock, self).__init__()
m = []
for i in range(2):
if i == 0:
m.append(conv(n_feats, 64, kernel_size, bias=bias))
else:
m.append(conv(64, n_feats, kernel_size, bias=bias))
if bn:
m.append(nn.BatchNorm2d(n_feats))
if i == 0:
m.append(act)
self.body = nn.Sequential(*m)
self.res_scale = res_scale
def forward(self, x):
res = self.body(x).mul(self.res_scale)
res += x
return res
##########################################################################
class MPRBlock(nn.Module):
def __init__(self, in_c=3, out_c=3, n_feat=80, scale_unetfeats=48, kernel_size=3,
reduction=4, bias=False,block_size=32):
super(MPRBlock, self).__init__()
act = nn.PReLU()
self.block_size=block_size
self.shallow_feat = nn.Sequential(conv(in_c, n_feat, kernel_size, bias=bias),
CAB(n_feat, kernel_size, reduction, bias=bias, act=act))
# Cross Stage Feature Fusion (CSFF)
self.stage_encoder = Encoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats, csff=True)
self.stage_decoder = Decoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats)
self.sam = SAM(n_feat, kernel_size=1, bias=bias,in_c=in_c)
self.phi = ResBlock(default_conv, 3, 3)
self.phit = ResBlock(default_conv, 3, 3)
self.r = nn.Parameter(torch.Tensor([0.5]))
self.concat = conv(n_feat * 2, n_feat, kernel_size, bias=bias)
def forward(self, stage_img, img,x_samfeats,f_encoder,f_decoder,PhiTPhi,PhiTb):
b,c,w,h = stage_img.shape
x_k_1 = F.unfold(stage_img,kernel_size=self.block_size, stride=self.block_size).permute(0, 2, 1).contiguous() # (b,l,cwh)
l=x_k_1.shape[1]
x_k_1=x_k_1.view(b*l,-1)
x = x_k_1 - self.r * torch.mm(x_k_1, PhiTPhi)
r_k = x + self.r * PhiTb
r_k=r_k.view(b,l,-1).permute(0, 2, 1).contiguous()#view(b,l,w,h)
r_k = F.fold(r_k, output_size=(w, h), kernel_size=self.block_size, stride=self.block_size).contiguous()
# compute x_k
x = self.shallow_feat(r_k)
## Concatenate SAM features of Stage 1 with shallow features of Stage 2
x_cat = self.concat(torch.cat([x, x_samfeats], 1))
## Process features of both patches with Encoder of Stage 2
feat1, f_encoder = self.stage_encoder(x_cat, f_encoder, f_decoder)
## Pass features through Decoder of Stage 2
f_decoder = self.stage_decoder(f_encoder, feat1)
## Apply SAM
x_samfeats, stage_img = self.sam(f_decoder[-1], r_k)
return stage_img, x_samfeats,feat1 ,f_decoder
##########################################################################
class DGUNet(nn.Module):
def __init__(self, in_c=3, out_c=3, n_feat=32, scale_unetfeats=16, scale_orsnetfeats=16, num_cab=8, kernel_size=3,
reduction=4, bias=False, nums_stages=5, cs_ratio=25,block_size=32):
super(DGUNet, self).__init__()
print('CS Ratio: ',cs_ratio)
self.block_size=block_size
self.Phi = nn.Parameter(init.xavier_normal_(torch.Tensor(np.ceil(cs_ratio*0.01*1024.).astype(np.int), 1024)))
act = nn.PReLU()
self.shallow_feat1 = nn.Sequential(conv(in_c, n_feat, kernel_size, bias=bias),
CAB(n_feat, kernel_size, reduction, bias=bias, act=act))
self.body = nn.ModuleList()
self.nums_stages = nums_stages
self.stage_model=MPRBlock(
in_c=in_c, out_c=in_c, n_feat=n_feat, scale_unetfeats=scale_unetfeats, kernel_size=kernel_size,
reduction=reduction, bias=bias
)
self.shallow_feat_final = nn.Sequential(conv(in_c, n_feat, kernel_size, bias=bias),
CAB(n_feat, kernel_size, reduction, bias=bias, act=act))
# Cross Stage Feature Fusion (CSFF)
self.stage1_encoder = Encoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats, csff=False)
self.stage1_decoder = Decoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats)
self.sam12 = SAM(n_feat, kernel_size=1, bias=bias,in_c=in_c)
self.r0 = nn.Parameter(torch.Tensor([0.5]))
self.r_final = nn.Parameter(torch.Tensor([0.5]))
self.concat_final = conv(n_feat * 2, n_feat + scale_orsnetfeats, kernel_size, bias=bias)
self.tail = conv(n_feat + scale_orsnetfeats, out_c, kernel_size, bias=bias)
def forward(self, img):
output_=[]
##-------------------------------------------
##-------------- Stage 1---------------------
##-------------------------------------------
b, c, w, h = img.shape
blocks = F.unfold(img,kernel_size=self.block_size, stride=self.block_size).permute(0, 2, 1).contiguous() # (b,l,cwh)
l = blocks.shape[1]
blocks=blocks.view(b*l,-1)
PhiTPhi = torch.mm(torch.transpose(self.Phi, 0, 1), self.Phi) # torch.mm(Phix, Phi)
Phix = torch.mm(blocks, torch.transpose(self.Phi, 0, 1)) # compression result
PhiTb = torch.mm(Phix,self.Phi)#.view(b,l,-1).permute(0, 2, 1).contiguous()
# compute r_0
x_0=PhiTb#.view(b,-1)
x = x_0 - self.r0 * torch.mm(x_0, PhiTPhi)
r_0 = x + self.r0 * PhiTb
r_0=r_0.view(b,l,-1).permute(0, 2, 1).contiguous()#view(b,l,w,h)
r_0 = F.fold(r_0, output_size=(w, h), kernel_size=self.block_size, stride=self.block_size).contiguous()
# compute x_k
x = self.shallow_feat1(r_0)
## Process features of all 4 patches with Encoder of Stage 1
feat1, f_encoder = self.stage1_encoder(x)
## Pass features through Decoder of Stage 1
f_decoder = self.stage1_decoder(f_encoder,feat1)
## Apply Supervised Attention Module (SAM)
x_samfeats, stage_img = self.sam12(f_decoder[-1], r_0)
output_.append(stage_img)
##-------------------------------------------
##-------------- Stage 2_k-1---------------------
##-------------------------------------------
for _ in range(self.nums_stages):
stage_img, x_samfeats, feat1, f_decoder = self.stage_model(stage_img, img,x_samfeats,feat1,f_decoder,PhiTPhi,PhiTb)
output_.append(stage_img)
##-------------------------------------------
##-------------- Stage k---------------------
##-------------------------------------------
# compute r_k
x_k_1 = F.unfold(stage_img,kernel_size=self.block_size, stride=self.block_size).permute(0, 2, 1).contiguous() # (b,l,cwh)
x_k_1=x_k_1.view(b*l,-1)
x = x_k_1 - self.r_final * torch.mm(x_k_1, PhiTPhi)
r_k = x + self.r_final * PhiTb
r_k=r_k.view(b,l,-1).permute(0, 2, 1).contiguous()#view(b,l,w,h)
r_k = F.fold(r_k, output_size=(w, h), kernel_size=self.block_size, stride=self.block_size).contiguous()
# compute x_k
x = self.shallow_feat_final(r_k)
## Concatenate SAM features
x_cat = self.concat_final(torch.cat([x, x_samfeats], 1))
stage_img = self.tail(x_cat)+r_k
output_.append(stage_img)
return output_[::-1]
================================================
FILE: Compressive-Sensing/DGUNet_plus.py
================================================
import torch
import torch.nn as nn
import torch.nn.functional as F
from pdb import set_trace as stx
# from pdb import set_trace as stx
from torch.nn import init
import numpy as np
##########################################################################
def conv(in_channels, out_channels, kernel_size, bias=False, stride=1):
return nn.Conv2d(
in_channels, out_channels, kernel_size,
padding=(kernel_size // 2), bias=bias, stride=stride)
def conv_down(in_chn, out_chn, bias=False):
layer = nn.Conv2d(in_chn, out_chn, kernel_size=4, stride=2, padding=1, bias=bias)
return layer
##########################################################################
## Channel Attention Layer
class CALayer(nn.Module):
def __init__(self, channel, reduction=16, bias=False):
super(CALayer, self).__init__()
# global average pooling: feature --> point
self.avg_pool = nn.AdaptiveAvgPool2d(1)
# feature channel downscale and upscale --> channel weight
self.conv_du = nn.Sequential(
nn.Conv2d(channel, channel // reduction, 1, padding=0, bias=bias),
nn.ReLU(inplace=True),
nn.Conv2d(channel // reduction, channel, 1, padding=0, bias=bias),
nn.Sigmoid()
)
def forward(self, x):
y = self.avg_pool(x) # 计算注意力权重
y = self.conv_du(y)
return x * y # 加权
##########################################################################
## Channel Attention Block (CAB)
class CAB(nn.Module):
def __init__(self, n_feat, kernel_size, reduction, bias, act):
super(CAB, self).__init__()
modules_body = []
modules_body.append(conv(n_feat, n_feat, kernel_size, bias=bias))
modules_body.append(act)
modules_body.append(conv(n_feat, n_feat, kernel_size, bias=bias))
self.CA = CALayer(n_feat, reduction, bias=bias)
self.body = nn.Sequential(*modules_body)
def forward(self, x):
res = self.body(x)
res = self.CA(res)
res += x
return res
##########################################################################
## Supervised Attention Module
class SAM(nn.Module):
def __init__(self, n_feat, kernel_size, bias,in_c):
super(SAM, self).__init__()
self.conv1 = conv(n_feat, n_feat, kernel_size, bias=bias)
self.conv2 = conv(n_feat, in_c, kernel_size, bias=bias)
self.conv3 = conv(in_c, n_feat, kernel_size, bias=bias)
def forward(self, x, x_img):
x1 = self.conv1(x)
img = self.conv2(x) + x_img
x2 = torch.sigmoid(self.conv3(img))
x1 = x1 * x2
x1 = x1 + x
return x1, img
class UNetConvBlock(nn.Module):
def __init__(self, in_size, out_size, downsample, relu_slope, use_csff=False, use_HIN=False):
super(UNetConvBlock, self).__init__()
self.downsample = downsample
self.identity = nn.Conv2d(in_size, out_size, 1, 1, 0)
self.use_csff = use_csff
self.conv_1 = nn.Conv2d(in_size, out_size, kernel_size=3, padding=1, bias=True)
self.relu_1 = nn.LeakyReLU(relu_slope, inplace=False)
self.conv_2 = nn.Conv2d(out_size, out_size, kernel_size=3, padding=1, bias=True)
self.relu_2 = nn.LeakyReLU(relu_slope, inplace=False)
if downsample and use_csff:
self.csff_enc = nn.Conv2d(out_size, out_size, 3, 1, 1)
self.csff_dec = nn.Conv2d(in_size, out_size, 3, 1, 1)
if use_HIN:
self.norm = nn.InstanceNorm2d(out_size//2, affine=True)
self.use_HIN = use_HIN
if downsample:
self.downsample = conv_down(out_size, out_size, bias=False)
def forward(self, x, enc=None, dec=None):
out = self.conv_1(x)
if self.use_HIN:
out_1, out_2 = torch.chunk(out, 2, dim=1)
out = torch.cat([self.norm(out_1), out_2], dim=1)
out = self.relu_1(out)
out = self.relu_2(self.conv_2(out))
out += self.identity(x)
if enc is not None and dec is not None:
assert self.use_csff
# print('enc+dec:',out.shape,enc.shape,dec.shape)
out = out + self.csff_enc(enc) + self.csff_dec(dec)
if self.downsample:
out_down = self.downsample(out)
return out_down, out
else:
return out
class UNetUpBlock(nn.Module):
def __init__(self, in_size, out_size, relu_slope):
super(UNetUpBlock, self).__init__()
self.up = nn.ConvTranspose2d(in_size, out_size, kernel_size=2, stride=2, bias=True)
self.conv_block = UNetConvBlock(out_size*2, out_size, False, relu_slope)
def forward(self, x, bridge):
up = self.up(x)
out = torch.cat([up, bridge], 1)
out = self.conv_block(out)
return out
##########################################################################
## U-Net
class Encoder(nn.Module):
def __init__(self, n_feat, kernel_size, reduction, act, bias, scale_unetfeats, csff,depth=4):
super(Encoder, self).__init__()
self.body=nn.ModuleList()#[]
self.depth=depth
for i in range(depth-1):
self.body.append(UNetConvBlock(in_size=n_feat+scale_unetfeats*i, out_size=n_feat+scale_unetfeats*(i+1), downsample=True, relu_slope=0.2, use_csff=True, use_HIN=True))
self.body.append(UNetConvBlock(in_size=n_feat+scale_unetfeats*(depth-1), out_size=n_feat+scale_unetfeats*(depth-1), downsample=False, relu_slope=0.2, use_csff=True, use_HIN=True))
def forward(self, x, encoder_outs=None, decoder_outs=None):
res=[]
if encoder_outs is not None and decoder_outs is not None:
for i,down in enumerate(self.body):
if (i+1) < self.depth:
x, x_up = down(x,encoder_outs[i],decoder_outs[-i-1])
res.append(x_up)
else:
x = down(x)
else:
for i,down in enumerate(self.body):
if (i+1) < self.depth:
x, x_up = down(x)
res.append(x_up)
else:
x = down(x)
return res,x
class Decoder(nn.Module):
def __init__(self, n_feat, kernel_size, reduction, act, bias, scale_unetfeats,depth=4):
super(Decoder, self).__init__()
self.body=nn.ModuleList()
self.skip_conv=nn.ModuleList()#[]
for i in range(depth-1):
self.body.append(UNetUpBlock(in_size=n_feat+scale_unetfeats*(depth-i-1), out_size=n_feat+scale_unetfeats*(depth-i-2), relu_slope=0.2))
self.skip_conv.append(nn.Conv2d(n_feat+scale_unetfeats*(depth-i-1), n_feat+scale_unetfeats*(depth-i-2), 3, 1, 1))
def forward(self, x, bridges):
res=[]
for i,up in enumerate(self.body):
x=up(x,self.skip_conv[i](bridges[-i-1]))
res.append(x)
return res
##########################################################################
##---------- Resizing Modules ----------
class DownSample(nn.Module):
def __init__(self, in_channels, s_factor):
super(DownSample, self).__init__()
self.down = nn.Sequential(nn.Upsample(scale_factor=0.5, mode='bilinear', align_corners=False),
nn.Conv2d(in_channels, in_channels + s_factor, 1, stride=1, padding=0, bias=False))
def forward(self, x):
x = self.down(x)
return x
class UpSample(nn.Module):
def __init__(self, in_channels, out_channels):
super(UpSample, self).__init__()
self.up = nn.Sequential(nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
nn.Conv2d(in_channels, out_channels, 1, stride=1, padding=0, bias=False))
def forward(self, x):
x = self.up(x)
return x
class SkipUpSample(nn.Module):
def __init__(self, in_channels, s_factor):
super(SkipUpSample, self).__init__()
self.up = nn.Sequential(nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
nn.Conv2d(in_channels + s_factor, in_channels, 1, stride=1, padding=0, bias=False))
def forward(self, x, y):
x = self.up(x)
x = x + y
return x
##########################################################################
## Original Resolution Block (ORB)
class ORB(nn.Module):
def __init__(self, n_feat, kernel_size, reduction, act, bias, num_cab):
super(ORB, self).__init__()
modules_body = []
modules_body = [CAB(n_feat, kernel_size, reduction, bias=bias, act=act) for _ in range(num_cab)]
modules_body.append(conv(n_feat, n_feat, kernel_size))
self.body = nn.Sequential(*modules_body)
def forward(self, x):
res = self.body(x)
res += x
return res
##########################################################################
class ORSNet(nn.Module):
def __init__(self, n_feat, scale_orsnetfeats, kernel_size, reduction, act, bias, scale_unetfeats, num_cab):
super(ORSNet, self).__init__()
self.orb1 = ORB(n_feat + scale_orsnetfeats, kernel_size, reduction, act, bias, num_cab)
self.orb2 = ORB(n_feat + scale_orsnetfeats, kernel_size, reduction, act, bias, num_cab)
self.orb3 = ORB(n_feat + scale_orsnetfeats, kernel_size, reduction, act, bias, num_cab)
self.up_enc1 = UpSample(n_feat+2*scale_unetfeats, n_feat+scale_unetfeats)
self.up_dec1 = UpSample(n_feat+scale_unetfeats, n_feat)
self.up_enc2 = nn.Sequential(UpSample(n_feat + 3*scale_unetfeats, n_feat+2*scale_unetfeats),
UpSample(n_feat+2*scale_unetfeats, n_feat+scale_unetfeats))
self.up_dec2 = nn.Sequential(UpSample(n_feat + 2*scale_unetfeats, n_feat+scale_unetfeats),
UpSample(n_feat+scale_unetfeats, n_feat))
self.conv_enc1 = nn.Conv2d(n_feat+scale_unetfeats, n_feat + scale_orsnetfeats, kernel_size=1, bias=bias)
self.conv_enc2 = nn.Conv2d(n_feat+scale_unetfeats, n_feat + scale_orsnetfeats, kernel_size=1, bias=bias)
self.conv_enc3 = nn.Conv2d(n_feat+scale_unetfeats, n_feat + scale_orsnetfeats, kernel_size=1, bias=bias)
self.conv_dec1 = nn.Conv2d(n_feat, n_feat + scale_orsnetfeats, kernel_size=1, bias=bias)
self.conv_dec2 = nn.Conv2d(n_feat, n_feat + scale_orsnetfeats, kernel_size=1, bias=bias)
self.conv_dec3 = nn.Conv2d(n_feat, n_feat + scale_orsnetfeats, kernel_size=1, bias=bias)
def forward(self, x, encoder_outs, decoder_outs):
x = self.orb1(x)
x = x + self.conv_enc1(encoder_outs[0]) + self.conv_dec1(decoder_outs[-1])
x = self.orb2(x)
x = x + self.conv_enc2(self.up_enc1(encoder_outs[1])) + self.conv_dec2(self.up_dec1(decoder_outs[-2]))
x = self.orb3(x)
x = x + self.conv_enc3(self.up_enc2(encoder_outs[2])) + self.conv_dec3(self.up_dec2(decoder_outs[-3]))
return x
def default_conv(in_channels, out_channels, kernel_size,stride=1, bias=True):
return nn.Conv2d(
in_channels, out_channels, kernel_size,
padding=(kernel_size//2),stride=stride, bias=bias)
class ResBlock(nn.Module):
def __init__(
self, conv, n_feats, kernel_size,
bias=True, bn=False, act=nn.PReLU(), res_scale=1):
super(ResBlock, self).__init__()
m = []
for i in range(2):
if i == 0:
m.append(conv(n_feats, 64, kernel_size, bias=bias))
else:
m.append(conv(64, n_feats, kernel_size, bias=bias))
if bn:
m.append(nn.BatchNorm2d(n_feats))
if i == 0:
m.append(act)
self.body = nn.Sequential(*m)
self.res_scale = res_scale
def forward(self, x):
res = self.body(x).mul(self.res_scale)
res += x
return res
##########################################################################
class MPRBlock(nn.Module):
def __init__(self, in_c=3, out_c=3, n_feat=80, scale_unetfeats=48, kernel_size=3,
reduction=4, bias=False):
super(MPRBlock, self).__init__()
act = nn.PReLU()
self.shallow_feat = nn.Sequential(conv(in_c, n_feat, kernel_size, bias=bias),
CAB(n_feat, kernel_size, reduction, bias=bias, act=act))
# Cross Stage Feature Fusion (CSFF)
self.stage_encoder = Encoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats, csff=True)
self.stage_decoder = Decoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats)
self.sam = SAM(n_feat, kernel_size=1, bias=bias,in_c=in_c)
self.phi = ResBlock(default_conv, 3, 3)
self.phit = ResBlock(default_conv, 3, 3)
self.r = nn.Parameter(torch.Tensor([0.5]))
self.concat = conv(n_feat * 2, n_feat, kernel_size, bias=bias)
def forward(self, stage_img, img,x_samfeats,f_encoder,f_decoder,PhiTPhi,PhiTb):
b,c,w,h = stage_img.shape
x_k_1=stage_img.view(b,-1)
# compute r_k
x = x_k_1 - self.r * torch.mm(x_k_1, PhiTPhi)
r_k = x + self.r * PhiTb
r_k = r_k.view(b,c,w,h)
# compute x_k
x = self.shallow_feat(r_k)
## Concatenate SAM features of Stage 1 with shallow features of Stage 2
x_cat = self.concat(torch.cat([x, x_samfeats], 1))
## Process features of both patches with Encoder of Stage 2
feat1, f_encoder = self.stage_encoder(x_cat, f_encoder, f_decoder)
## Pass features through Decoder of Stage 2
f_decoder = self.stage_decoder(f_encoder, feat1)
## Apply SAM
x_samfeats, stage_img = self.sam(f_decoder[-1], r_k)
return stage_img, x_samfeats,feat1 ,f_decoder
##########################################################################
class DGUNet(nn.Module):
def __init__(self, in_c=3, out_c=3, n_feat=32, scale_unetfeats=16, scale_orsnetfeats=16, num_cab=8, kernel_size=3,
reduction=4, bias=False, nums_stages=5, cs_ratio=25):
super(DGUNet, self).__init__()
print('CS Ratio: ',cs_ratio)
self.Phi = nn.Parameter(init.xavier_normal_(torch.Tensor(np.ceil(cs_ratio*0.01*1024.).astype(np.int), 1024)))
act = nn.PReLU()
self.shallow_feat1 = nn.Sequential(conv(in_c, n_feat, kernel_size, bias=bias),
CAB(n_feat, kernel_size, reduction, bias=bias, act=act))
self.body = nn.ModuleList()
for _ in range(nums_stages):
self.body.append(MPRBlock(
in_c=in_c, out_c=in_c, n_feat=n_feat, scale_unetfeats=scale_unetfeats, kernel_size=kernel_size,
reduction=reduction, bias=bias
))
self.shallow_feat_final = nn.Sequential(conv(in_c, n_feat, kernel_size, bias=bias),
CAB(n_feat, kernel_size, reduction, bias=bias, act=act))
# Cross Stage Feature Fusion (CSFF)
self.stage1_encoder = Encoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats, csff=False)
self.stage1_decoder = Decoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats)
self.sam12 = SAM(n_feat, kernel_size=1, bias=bias,in_c=in_c)
self.r0 = nn.Parameter(torch.Tensor([0.5]))
self.r_final = nn.Parameter(torch.Tensor([0.5]))
self.concat_final = conv(n_feat * 2, n_feat + scale_orsnetfeats, kernel_size, bias=bias)
self.tail = conv(n_feat + scale_orsnetfeats, out_c, kernel_size, bias=bias)
def forward(self, img):
output_=[]
##-------------------------------------------
##-------------- Stage 1---------------------
##-------------------------------------------
b, c, w, h = img.shape
PhiTPhi = torch.mm(torch.transpose(self.Phi, 0, 1), self.Phi) # torch.mm(Phix, Phi)
Phix = torch.mm(img.view(b,-1), torch.transpose(self.Phi, 0, 1)) # compression result
PhiTb = torch.mm(Phix,self.Phi)
# compute r_0
x_0=PhiTb.view(b,-1)
x = x_0 - self.r0 * torch.mm(x_0, PhiTPhi)
r_0 = x + self.r0 * PhiTb
r_0=r_0.view(b,c,w,h)
# compute x_k
x = self.shallow_feat1(r_0)
## Process features of all 4 patches with Encoder of Stage 1
feat1, f_encoder = self.stage1_encoder(x)
## Pass features through Decoder of Stage 1
f_decoder = self.stage1_decoder(f_encoder,feat1)
## Apply Supervised Attention Module (SAM)
x_samfeats, stage_img = self.sam12(f_decoder[-1], r_0)
output_.append(stage_img)
##-------------------------------------------
##-------------- Stage 2_k-1---------------------
##-------------------------------------------
for stage_model in self.body:
stage_img, x_samfeats, feat1, f_decoder = stage_model(stage_img, img,x_samfeats,feat1,f_decoder,PhiTPhi,PhiTb)
output_.append(stage_img)
##-------------------------------------------
##-------------- Stage k---------------------
##-------------------------------------------
# compute r_k
x_k_1 = stage_img.view(b,-1)
x = x_k_1 - self.r_final * torch.mm(x_k_1, PhiTPhi)
r_k = x + self.r_final * PhiTb
r_k = r_k.view(b,c,w,h)
# compute x_k
x = self.shallow_feat_final(r_k)
## Concatenate SAM features
x_cat = self.concat_final(torch.cat([x, x_samfeats], 1))
stage_img = self.tail(x_cat)+r_k
output_.append(stage_img)
return output_[::-1]
================================================
FILE: Compressive-Sensing/DGUNet_plus_deblock.py
================================================
import torch
import torch.nn as nn
import torch.nn.functional as F
from pdb import set_trace as stx
# from pdb import set_trace as stx
from torch.nn import init
import numpy as np
##########################################################################
def conv(in_channels, out_channels, kernel_size, bias=False, stride=1):
return nn.Conv2d(
in_channels, out_channels, kernel_size,
padding=(kernel_size // 2), bias=bias, stride=stride)
def conv_down(in_chn, out_chn, bias=False):
layer = nn.Conv2d(in_chn, out_chn, kernel_size=4, stride=2, padding=1, bias=bias)
return layer
##########################################################################
## Channel Attention Layer
class CALayer(nn.Module):
def __init__(self, channel, reduction=16, bias=False):
super(CALayer, self).__init__()
# global average pooling: feature --> point
self.avg_pool = nn.AdaptiveAvgPool2d(1)
# feature channel downscale and upscale --> channel weight
self.conv_du = nn.Sequential(
nn.Conv2d(channel, channel // reduction, 1, padding=0, bias=bias),
nn.ReLU(inplace=True),
nn.Conv2d(channel // reduction, channel, 1, padding=0, bias=bias),
nn.Sigmoid()
)
def forward(self, x):
y = self.avg_pool(x) # 计算注意力权重
y = self.conv_du(y)
return x * y # 加权
##########################################################################
## Channel Attention Block (CAB)
class CAB(nn.Module):
def __init__(self, n_feat, kernel_size, reduction, bias, act):
super(CAB, self).__init__()
modules_body = []
modules_body.append(conv(n_feat, n_feat, kernel_size, bias=bias))
modules_body.append(act)
modules_body.append(conv(n_feat, n_feat, kernel_size, bias=bias))
self.CA = CALayer(n_feat, reduction, bias=bias)
self.body = nn.Sequential(*modules_body)
def forward(self, x):
res = self.body(x)
res = self.CA(res)
res += x
return res
##########################################################################
## Supervised Attention Module
class SAM(nn.Module):
def __init__(self, n_feat, kernel_size, bias,in_c):
super(SAM, self).__init__()
self.conv1 = conv(n_feat, n_feat, kernel_size, bias=bias)
self.conv2 = conv(n_feat, in_c, kernel_size, bias=bias)
self.conv3 = conv(in_c, n_feat, kernel_size, bias=bias)
def forward(self, x, x_img):
x1 = self.conv1(x)
img = self.conv2(x) + x_img
x2 = torch.sigmoid(self.conv3(img))
x1 = x1 * x2
x1 = x1 + x
return x1, img
class UNetConvBlock(nn.Module):
def __init__(self, in_size, out_size, downsample, relu_slope, use_csff=False, use_HIN=False):
super(UNetConvBlock, self).__init__()
self.downsample = downsample
self.identity = nn.Conv2d(in_size, out_size, 1, 1, 0)
self.use_csff = use_csff
self.conv_1 = nn.Conv2d(in_size, out_size, kernel_size=3, padding=1, bias=True)
self.relu_1 = nn.LeakyReLU(relu_slope, inplace=False)
self.conv_2 = nn.Conv2d(out_size, out_size, kernel_size=3, padding=1, bias=True)
self.relu_2 = nn.LeakyReLU(relu_slope, inplace=False)
if downsample and use_csff:
self.csff_enc = nn.Conv2d(out_size, out_size, 3, 1, 1)
self.csff_dec = nn.Conv2d(in_size, out_size, 3, 1, 1)
if use_HIN:
self.norm = nn.InstanceNorm2d(out_size//2, affine=True)
self.use_HIN = use_HIN
if downsample:
self.downsample = conv_down(out_size, out_size, bias=False)
def forward(self, x, enc=None, dec=None):
out = self.conv_1(x)
if self.use_HIN:
out_1, out_2 = torch.chunk(out, 2, dim=1)
out = torch.cat([self.norm(out_1), out_2], dim=1)
out = self.relu_1(out)
out = self.relu_2(self.conv_2(out))
out += self.identity(x)
if enc is not None and dec is not None:
assert self.use_csff
# print('enc+dec:',out.shape,enc.shape,dec.shape)
out = out + self.csff_enc(enc) + self.csff_dec(dec)
if self.downsample:
out_down = self.downsample(out)
return out_down, out
else:
return out
class UNetUpBlock(nn.Module):
def __init__(self, in_size, out_size, relu_slope):
super(UNetUpBlock, self).__init__()
self.up = nn.ConvTranspose2d(in_size, out_size, kernel_size=2, stride=2, bias=True)
self.conv_block = UNetConvBlock(out_size*2, out_size, False, relu_slope)
def forward(self, x, bridge):
# print('merge',x.shape,bridge.shape)
up = self.up(x)
out = torch.cat([up, bridge], 1)
out = self.conv_block(out)
return out
##########################################################################
## U-Net
class Encoder(nn.Module):
def __init__(self, n_feat, kernel_size, reduction, act, bias, scale_unetfeats, csff,depth=4):
super(Encoder, self).__init__()
self.body=nn.ModuleList()#[]
self.depth=depth
for i in range(depth-1):
# downsample = True if (i+1) < depth else False
self.body.append(UNetConvBlock(in_size=n_feat+scale_unetfeats*i, out_size=n_feat+scale_unetfeats*(i+1), downsample=True, relu_slope=0.2, use_csff=True, use_HIN=True))
self.body.append(UNetConvBlock(in_size=n_feat+scale_unetfeats*(depth-1), out_size=n_feat+scale_unetfeats*(depth-1), downsample=False, relu_slope=0.2, use_csff=True, use_HIN=True))
def forward(self, x, encoder_outs=None, decoder_outs=None):
res=[]
if encoder_outs is not None and decoder_outs is not None:
for i,down in enumerate(self.body):
# print(encoder_outs[i].shape,decoder_outs[-i-1].shape)
# exit(0)
if (i+1) < self.depth:
x, x_up = down(x,encoder_outs[i],decoder_outs[-i-1])
res.append(x_up)
else:
# print(i,len(encoder_outs),len(decoder_outs))
x = down(x)#,encoder_outs[i],decoder_outs[-i-1])
else:
for i,down in enumerate(self.body):
if (i+1) < self.depth:
x, x_up = down(x)
res.append(x_up)
else:
x = down(x)
return res,x
class Decoder(nn.Module):
def __init__(self, n_feat, kernel_size, reduction, act, bias, scale_unetfeats,depth=4):
super(Decoder, self).__init__()
self.body=nn.ModuleList()
self.skip_conv=nn.ModuleList()#[]
for i in range(depth-1):
self.body.append(UNetUpBlock(in_size=n_feat+scale_unetfeats*(depth-i-1), out_size=n_feat+scale_unetfeats*(depth-i-2), relu_slope=0.2))
self.skip_conv.append(nn.Conv2d(n_feat+scale_unetfeats*(depth-i-1), n_feat+scale_unetfeats*(depth-i-2), 3, 1, 1))
def forward(self, x, bridges):
# for b in bridges:
# print(b.shape)
res=[]
for i,up in enumerate(self.body):
x=up(x,self.skip_conv[i](bridges[-i-1]))
# print(x.shape)
res.append(x)
return res
##########################################################################
##---------- Resizing Modules ----------
class DownSample(nn.Module):
def __init__(self, in_channels, s_factor):
super(DownSample, self).__init__()
self.down = nn.Sequential(nn.Upsample(scale_factor=0.5, mode='bilinear', align_corners=False),
nn.Conv2d(in_channels, in_channels + s_factor, 1, stride=1, padding=0, bias=False))
def forward(self, x):
x = self.down(x)
return x
class UpSample(nn.Module):
def __init__(self, in_channels, out_channels):
super(UpSample, self).__init__()
self.up = nn.Sequential(nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
nn.Conv2d(in_channels, out_channels, 1, stride=1, padding=0, bias=False))
def forward(self, x):
x = self.up(x)
return x
class SkipUpSample(nn.Module):
def __init__(self, in_channels, s_factor):
super(SkipUpSample, self).__init__()
self.up = nn.Sequential(nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
nn.Conv2d(in_channels + s_factor, in_channels, 1, stride=1, padding=0, bias=False))
def forward(self, x, y):
x = self.up(x)
x = x + y
return x
##########################################################################
## Original Resolution Block (ORB)
class ORB(nn.Module):
def __init__(self, n_feat, kernel_size, reduction, act, bias, num_cab):
super(ORB, self).__init__()
modules_body = []
modules_body = [CAB(n_feat, kernel_size, reduction, bias=bias, act=act) for _ in range(num_cab)]
modules_body.append(conv(n_feat, n_feat, kernel_size))
self.body = nn.Sequential(*modules_body)
def forward(self, x):
res = self.body(x)
res += x
return res
##########################################################################
class ORSNet(nn.Module):
def __init__(self, n_feat, scale_orsnetfeats, kernel_size, reduction, act, bias, scale_unetfeats, num_cab):
super(ORSNet, self).__init__()
self.orb1 = ORB(n_feat + scale_orsnetfeats, kernel_size, reduction, act, bias, num_cab)
self.orb2 = ORB(n_feat + scale_orsnetfeats, kernel_size, reduction, act, bias, num_cab)
self.orb3 = ORB(n_feat + scale_orsnetfeats, kernel_size, reduction, act, bias, num_cab)
self.up_enc1 = UpSample(n_feat+2*scale_unetfeats, n_feat+scale_unetfeats)
self.up_dec1 = UpSample(n_feat+scale_unetfeats, n_feat)
self.up_enc2 = nn.Sequential(UpSample(n_feat + 3*scale_unetfeats, n_feat+2*scale_unetfeats),
UpSample(n_feat+2*scale_unetfeats, n_feat+scale_unetfeats))
self.up_dec2 = nn.Sequential(UpSample(n_feat + 2*scale_unetfeats, n_feat+scale_unetfeats),
UpSample(n_feat+scale_unetfeats, n_feat))
self.conv_enc1 = nn.Conv2d(n_feat+scale_unetfeats, n_feat + scale_orsnetfeats, kernel_size=1, bias=bias)
self.conv_enc2 = nn.Conv2d(n_feat+scale_unetfeats, n_feat + scale_orsnetfeats, kernel_size=1, bias=bias)
self.conv_enc3 = nn.Conv2d(n_feat+scale_unetfeats, n_feat + scale_orsnetfeats, kernel_size=1, bias=bias)
self.conv_dec1 = nn.Conv2d(n_feat, n_feat + scale_orsnetfeats, kernel_size=1, bias=bias)
self.conv_dec2 = nn.Conv2d(n_feat, n_feat + scale_orsnetfeats, kernel_size=1, bias=bias)
self.conv_dec3 = nn.Conv2d(n_feat, n_feat + scale_orsnetfeats, kernel_size=1, bias=bias)
def forward(self, x, encoder_outs, decoder_outs):
x = self.orb1(x)
# print(encoder_outs[0].shape,decoder_outs[-1].shape)
x = x + self.conv_enc1(encoder_outs[0]) + self.conv_dec1(decoder_outs[-1])
x = self.orb2(x)
# print(encoder_outs[1].shape,decoder_outs[-2].shape)
x = x + self.conv_enc2(self.up_enc1(encoder_outs[1])) + self.conv_dec2(self.up_dec1(decoder_outs[-2]))
x = self.orb3(x)
# print(encoder_outs[2].shape,decoder_outs[-3].shape)
x = x + self.conv_enc3(self.up_enc2(encoder_outs[2])) + self.conv_dec3(self.up_dec2(decoder_outs[-3]))
return x
def default_conv(in_channels, out_channels, kernel_size,stride=1, bias=True):
return nn.Conv2d(
in_channels, out_channels, kernel_size,
padding=(kernel_size//2),stride=stride, bias=bias)
class ResBlock(nn.Module):
def __init__(
self, conv, n_feats, kernel_size,
bias=True, bn=False, act=nn.PReLU(), res_scale=1):
super(ResBlock, self).__init__()
m = []
for i in range(2):
if i == 0:
m.append(conv(n_feats, 64, kernel_size, bias=bias))
else:
m.append(conv(64, n_feats, kernel_size, bias=bias))
if bn:
m.append(nn.BatchNorm2d(n_feats))
if i == 0:
m.append(act)
self.body = nn.Sequential(*m)
self.res_scale = res_scale
def forward(self, x):
res = self.body(x).mul(self.res_scale)
res += x
return res
##########################################################################
class MPRBlock(nn.Module):
def __init__(self, in_c=3, out_c=3, n_feat=80, scale_unetfeats=48, kernel_size=3,
reduction=4, bias=False,block_size=32):
super(MPRBlock, self).__init__()
act = nn.PReLU()
self.block_size=block_size
self.shallow_feat = nn.Sequential(conv(in_c, n_feat, kernel_size, bias=bias),
CAB(n_feat, kernel_size, reduction, bias=bias, act=act))
# Cross Stage Feature Fusion (CSFF)
self.stage_encoder = Encoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats, csff=True)
self.stage_decoder = Decoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats)
self.sam = SAM(n_feat, kernel_size=1, bias=bias,in_c=in_c)
self.phi = ResBlock(default_conv, 3, 3)
self.phit = ResBlock(default_conv, 3, 3)
self.r = nn.Parameter(torch.Tensor([0.5]))
self.concat = conv(n_feat * 2, n_feat, kernel_size, bias=bias)
def forward(self, stage_img, img,x_samfeats,f_encoder,f_decoder,PhiTPhi,PhiTb):
b,c,w,h = stage_img.shape
x_k_1 = F.unfold(stage_img,kernel_size=self.block_size, stride=self.block_size).permute(0, 2, 1).contiguous() # (b,l,cwh)
l=x_k_1.shape[1]
x_k_1=x_k_1.view(b*l,-1)
x = x_k_1 - self.r * torch.mm(x_k_1, PhiTPhi)
r_k = x + self.r * PhiTb
r_k=r_k.view(b,l,-1).permute(0, 2, 1).contiguous()#view(b,l,w,h)
r_k = F.fold(r_k, output_size=(w, h), kernel_size=self.block_size, stride=self.block_size).contiguous()
# compute x_k
x = self.shallow_feat(r_k)
## Concatenate SAM features of Stage 1 with shallow features of Stage 2
x_cat = self.concat(torch.cat([x, x_samfeats], 1))
## Process features of both patches with Encoder of Stage 2
feat1, f_encoder = self.stage_encoder(x_cat, f_encoder, f_decoder)
## Pass features through Decoder of Stage 2
f_decoder = self.stage_decoder(f_encoder, feat1)
## Apply SAM
x_samfeats, stage_img = self.sam(f_decoder[-1], r_k)
return stage_img, x_samfeats,feat1 ,f_decoder
##########################################################################
class DGUNet(nn.Module):
def __init__(self, in_c=3, out_c=3, n_feat=32, scale_unetfeats=16, scale_orsnetfeats=16, num_cab=8, kernel_size=3,
reduction=4, bias=False, nums_stages=5, cs_ratio=25,block_size=32):
super(DGUNet, self).__init__()
print('CS Ratio: ',cs_ratio)
self.block_size=block_size
self.Phi = nn.Parameter(init.xavier_normal_(torch.Tensor(np.ceil(cs_ratio*0.01*1024.).astype(np.int), 1024)))
act = nn.PReLU()
self.shallow_feat1 = nn.Sequential(conv(in_c, n_feat, kernel_size, bias=bias),
CAB(n_feat, kernel_size, reduction, bias=bias, act=act))
self.body = nn.ModuleList()
for _ in range(nums_stages):
self.body.append(MPRBlock(
in_c=in_c, out_c=in_c, n_feat=n_feat, scale_unetfeats=scale_unetfeats, kernel_size=kernel_size,
reduction=reduction, bias=bias
))
self.shallow_feat_final = nn.Sequential(conv(in_c, n_feat, kernel_size, bias=bias),
CAB(n_feat, kernel_size, reduction, bias=bias, act=act))
# Cross Stage Feature Fusion (CSFF)
self.stage1_encoder = Encoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats, csff=False)
self.stage1_decoder = Decoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats)
# self.stage_orsnet = ORSNet(n_feat, scale_orsnetfeats, kernel_size, reduction, act, bias, scale_unetfeats,
# num_cab)
self.sam12 = SAM(n_feat, kernel_size=1, bias=bias,in_c=in_c)
self.r0 = nn.Parameter(torch.Tensor([0.5]))
self.r_final = nn.Parameter(torch.Tensor([0.5]))
self.concat_final = conv(n_feat * 2, n_feat + scale_orsnetfeats, kernel_size, bias=bias)
self.tail = conv(n_feat + scale_orsnetfeats, out_c, kernel_size, bias=bias)
def forward(self, img):
output_=[]
##-------------------------------------------
##-------------- Stage 1---------------------
##-------------------------------------------
b, c, w, h = img.shape
blocks = F.unfold(img,kernel_size=self.block_size, stride=self.block_size).permute(0, 2, 1).contiguous() # (b,l,cwh)
l = blocks.shape[1]
blocks=blocks.view(b*l,-1)
PhiTPhi = torch.mm(torch.transpose(self.Phi, 0, 1), self.Phi) # torch.mm(Phix, Phi)
Phix = torch.mm(blocks, torch.transpose(self.Phi, 0, 1)) # compression result
PhiTb = torch.mm(Phix,self.Phi)#.view(b,l,-1).permute(0, 2, 1).contiguous()
# compute r_0
x_0=PhiTb#.view(b,-1)
# print(x_0.shape, PhiTPhi.shape)
# exit(0)
x = x_0 - self.r0 * torch.mm(x_0, PhiTPhi)
r_0 = x + self.r0 * PhiTb
r_0=r_0.view(b,l,-1).permute(0, 2, 1).contiguous()#view(b,l,w,h)
r_0 = F.fold(r_0, output_size=(w, h), kernel_size=self.block_size, stride=self.block_size).contiguous()
# compute x_k
x = self.shallow_feat1(r_0)
## Process features of all 4 patches with Encoder of Stage 1
feat1, f_encoder = self.stage1_encoder(x)
## Pass features through Decoder of Stage 1
f_decoder = self.stage1_decoder(f_encoder,feat1)
## Apply Supervised Attention Module (SAM)
x_samfeats, stage_img = self.sam12(f_decoder[-1], r_0)
output_.append(stage_img)
##-------------------------------------------
##-------------- Stage 2_k-1---------------------
##-------------------------------------------
for stage_model in self.body:
stage_img, x_samfeats, feat1, f_decoder = stage_model(stage_img, img,x_samfeats,feat1,f_decoder,PhiTPhi,PhiTb)
output_.append(stage_img)
##-------------------------------------------
##-------------- Stage k---------------------
##-------------------------------------------
# compute r_k
x_k_1 = F.unfold(stage_img,kernel_size=self.block_size, stride=self.block_size).permute(0, 2, 1).contiguous() # (b,l,cwh)
x_k_1=x_k_1.view(b*l,-1)
x = x_k_1 - self.r_final * torch.mm(x_k_1, PhiTPhi)
r_k = x + self.r_final * PhiTb
r_k=r_k.view(b,l,-1).permute(0, 2, 1).contiguous()#view(b,l,w,h)
r_k = F.fold(r_k, output_size=(w, h), kernel_size=self.block_size, stride=self.block_size).contiguous()
# compute x_k
x = self.shallow_feat_final(r_k)
## Concatenate SAM features
x_cat = self.concat_final(torch.cat([x, x_samfeats], 1))
# x_cat = self.stage_orsnet(x_cat, f_encoder, f_decoder)
stage_img = self.tail(x_cat)+r_k
output_.append(stage_img)
return output_[::-1]#[stage5_img + x3_img, stage4_img, stage3_img, stage2_img, stage1_img]
================================================
FILE: Compressive-Sensing/Datasets/README.md
================================================
Download datasets from the provided links and place them in this directory. Your directory structure should look something like this
`├──`[DIV2K_BSD400](https://drive.google.com/file/d/1eW4eC3eazqRYnjXoIOXNX-xOYSwoIoHf/view?usp=sharing) <br/>
`├──`[Set11](https://drive.google.com/file/d/11WXTqqD4cZtymJJ_WObYb2s1vv1Dlr-2/view?usp=sharing) <br/>
`└──`[BSD68](https://drive.google.com/file/d/11WXTqqD4cZtymJJ_WObYb2s1vv1Dlr-2/view?usp=sharing) <br/>
================================================
FILE: Compressive-Sensing/README.md
================================================
## Training
- Download the [Datasets](Datasets/README.md)
- Train the model with default arguments by running
### For DGUNet
```
Step 1: python train.py
Step 2: python train_deblock.py
```
### For DGUNet_plus
```
Step 1: python train_plus.py
Step 2: python train_plus_deblock.py
```
## Evaluation
### Download the [model](https://drive.google.com/file/d/1euV8SnXHuYswwbm9FaV1y5RQRLd58ou_/view?usp=sharing) and place it in ./model/
- Download [images](https://drive.google.com/drive/folders/1a2qKfXWpNuTGOm2-Jex8kfNSzYJLbqkf?usp=sharing) of Set11&BSD68 and place them in `./Datasets/`
- Run
```
python test_deblock.py
python test_plus_deblock.py
```
================================================
FILE: Compressive-Sensing/config.py
================================================
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Jul 23 14:35:48 2019
@author: aditya
"""
r"""This module provides package-wide configuration management."""
from typing import Any, List
from yacs.config import CfgNode as CN
class Config(object):
r"""
A collection of all the required configuration parameters. This class is a nested dict-like
structure, with nested keys accessible as attributes. It contains sensible default values for
all the parameters, which may be overriden by (first) through a YAML file and (second) through
a list of attributes and values.
Extended Summary
----------------
This class definition contains default values corresponding to ``joint_training`` phase, as it
is the final training phase and uses almost all the configuration parameters. Modification of
any parameter after instantiating this class is not possible, so you must override required
parameter values in either through ``config_yaml`` file or ``config_override`` list.
Parameters
----------
config_yaml: str
Path to a YAML file containing configuration parameters to override.
config_override: List[Any], optional (default= [])
A list of sequential attributes and values of parameters to override. This happens after
overriding from YAML file.
Examples
--------
Let a YAML file named "config.yaml" specify these parameters to override::
ALPHA: 1000.0
BETA: 0.5
>>> _C = Config("config.yaml", ["OPTIM.BATCH_SIZE", 2048, "BETA", 0.7])
>>> _C.ALPHA # default: 100.0
1000.0
>>> _C.BATCH_SIZE # default: 256
2048
>>> _C.BETA # default: 0.1
0.7
Attributes
----------
"""
def __init__(self, config_yaml: str, config_override: List[Any] = []):
self._C = CN()
self._C.GPU = [0]
self._C.VERBOSE = False
self._C.MODEL = CN()
self._C.MODEL.MODE = 'global'
self._C.MODEL.SESSION = 'ps128_bs1'
self._C.OPTIM = CN()
self._C.OPTIM.BATCH_SIZE = 1
self._C.OPTIM.NUM_EPOCHS = 100
self._C.OPTIM.NEPOCH_DECAY = [100]
self._C.OPTIM.LR_INITIAL = 0.0002
self._C.OPTIM.LR_MIN = 0.0002
self._C.OPTIM.BETA1 = 0.5
self._C.TRAINING = CN()
self._C.TRAINING.VAL_AFTER_EVERY = 3
self._C.TRAINING.RESUME = False
self._C.TRAINING.SAVE_IMAGES = False
self._C.TRAINING.TRAIN_DIR = 'images_dir/train'
self._C.TRAINING.VAL_DIR = 'images_dir/val'
self._C.TRAINING.SAVE_DIR = 'checkpoints'
self._C.TRAINING.TRAIN_PS = 64
self._C.TRAINING.VAL_PS = 64
# Override parameter values from YAML file first, then from override list.
self._C.merge_from_file(config_yaml)
self._C.merge_from_list(config_override)
# Make an instantiated object of this class immutable.
self._C.freeze()
def dump(self, file_path: str):
r"""Save config at the specified file path.
Parameters
----------
file_path: str
(YAML) path to save config at.
"""
self._C.dump(stream=open(file_path, "w"))
def __getattr__(self, attr: str):
return self._C.__getattr__(attr)
def __repr__(self):
return self._C.__repr__()
================================================
FILE: Compressive-Sensing/csdata_fast.py
================================================
from torch.utils.data import Dataset
import imageio
import os
import torch
import glob
import random
import numpy as np
from functools import lru_cache
ycbcr_from_rgb = torch.tensor([[ 65.481, 128.553, 24.966],
[ -37.797, -74.203, 112.0],
[ 112.0, -93.786, -18.214]])
def rgb2ycbcr(rgb):
arr = rgb.float() / 255.0 @ ycbcr_from_rgb.transpose(1,0)
arr[..., 0] += 16
arr[..., 1] += 128
arr[..., 2] += 128
return arr
class SlowDataset(Dataset):
def __init__(self, args, train=True): # __init__是初始化该类的一些基础参数
super(SlowDataset, self).__init__()
self.args = args
self.train = train
self.image_folder = os.path.join('.', args.data_dir)
self.bin_image_folder = os.path.join('.', args.data_dir+'bin')
if not os.path.exists(self.bin_image_folder): os.makedirs(self.bin_image_folder, exist_ok=True)
self.ext = '/*%s' % args.ext
self.file_names = glob.glob(self.image_folder + self.ext)
self.bin_file_names = list()
self.prepare_cache()
def prepare_cache(self):
for fname in self.file_names:
bin_fname = fname.replace(self.image_folder, self.bin_image_folder).replace(self.args.ext, '.npy')
self.bin_file_names.append(bin_fname)
if not os.path.exists(bin_fname):
img = imageio.imread(fname)
np.save(bin_fname, img)
print(f'{bin_fname} prepared!')
def __len__(self): # 返回整个数据集的大小
return len(self.file_names) * 200
@lru_cache(maxsize=400)
def get_ndarray(self, fname):
return np.load(fname)
def __getitem__(self, index):
rgb_range = self.args.rgb_range
n_channels = self.args.n_channels
# img = torch.Tensor(np.load(self.file_names[index]))
img = torch.Tensor(self.get_ndarray(self.bin_file_names[index % len(self.file_names)]))
if img.numpy().ndim == 2:
img = img.unsqueeze(2)
c = img.shape[2]
# input rgb image output y chanel
if n_channels == 1 and c == 3:
img = rgb2ycbcr(img)[:, :, 0].unsqueeze(2)
elif n_channels == 3 and c == 1:
img = img.repeat(1, 1, 3)
w, h, _ = img.shape
th = tw = self.args.patch_size
i = random.randint(0, w - tw)
j = random.randint(0, h - th)
img = img[i:i + tw, j:j + th, :]
img_tensor = img.permute(2, 0, 1)
img_tensor = img_tensor * rgb_range / 255.0
img_tensor = self.augment(img_tensor)
return img_tensor
def augment(self, img, hflip=True, rot=True):
hflip = hflip and random.random() < 0.5
vflip = rot and random.random() < 0.5
rot90 = rot and random.random() < 0.5
if hflip:
img = img.flip([1])
if vflip:
img = img.flip([0])
if rot90:
img = img.permute(0, 2, 1)
return img
class SlowDataset2(Dataset):
def __init__(self, args, train=True): # __init__是初始化该类的一些基础参数
super(SlowDataset2, self).__init__()
self.args = args
self.train = train
self.image_folder = os.path.join('.', args.data_dir)
self.bin_image_folder = os.path.join('.', args.data_dir+'bin')
if not os.path.exists(self.bin_image_folder): os.makedirs(self.bin_image_folder, exist_ok=True)
self.ext = '/*%s' % args.ext
self.file_names = glob.glob(self.image_folder + self.ext)
self.bin_file_names = list()
self.prepare_cache()
def prepare_cache(self):
for fname in self.file_names:
bin_fname = fname.replace(self.image_folder, self.bin_image_folder).replace(self.args.ext, '.npy')
self.bin_file_names.append(bin_fname)
if not os.path.exists(bin_fname):
img = imageio.imread(fname)
np.save(bin_fname, img)
print(f'{bin_fname} prepared!')
def __len__(self): # 返回整个数据集的大小
return len(self.file_names) * 200
@lru_cache(maxsize=400)
def get_ndarray(self, fname):
return np.load(fname)
def __getitem__(self, index):
rgb_range = self.args.rgb_range
n_channels = self.args.n_channels
# img = torch.Tensor(np.load(self.file_names[index]))
img = torch.Tensor(self.get_ndarray(self.bin_file_names[index % len(self.file_names)]))
if img.numpy().ndim == 2:
img = img.unsqueeze(2)
c = img.shape[2]
# input rgb image output y chanel
if n_channels == 1 and c == 3:
img = rgb2ycbcr(img)[:, :, 0].unsqueeze(2)
elif n_channels == 3 and c == 1:
img = img.repeat(1, 1, 3)
w, h, _ = img.shape
th = tw = self.args.patch_size
i = random.randint(0, w - tw)
j = random.randint(0, h - th)
img = img[i:i + tw, j:j + th, :]
img_tensor = img.permute(2, 0, 1)
img_tensor = img_tensor * rgb_range / 255.0
img_tensor = self.augment(img_tensor)
return img_tensor
def augment(self, img, hflip=True, rot=True):
hflip = hflip and random.random() < 0.5
vflip = rot and random.random() < 0.5
rot90 = rot and random.random() < 0.5
if hflip:
img = img.flip([1])
if vflip:
img = img.flip([0])
if rot90:
img = img.permute(0, 2, 1)
return img
================================================
FILE: Compressive-Sensing/csdata_fast2.py
================================================
from torch.utils.data import Dataset
import imageio
import os
import torch
import glob
import random
import numpy as np
from functools import lru_cache
ycbcr_from_rgb = torch.tensor([[ 65.481, 128.553, 24.966],
[ -37.797, -74.203, 112.0],
[ 112.0, -93.786, -18.214]])
def rgb2ycbcr(rgb):
arr = rgb.float() / 255.0 @ ycbcr_from_rgb.transpose(1,0)
arr[..., 0] += 16
arr[..., 1] += 128
arr[..., 2] += 128
return arr
class SlowDataset(Dataset):
def __init__(self, args, train=True): # __init__是初始化该类的一些基础参数
super(SlowDataset, self).__init__()
self.args = args
self.train = train
self.image_folder = os.path.join('.', args.data_dir)
self.bin_image_folder = os.path.join('.', args.data_dir+'bin')
if not os.path.exists(self.bin_image_folder): os.makedirs(self.bin_image_folder, exist_ok=True)
self.ext = '/*%s' % args.ext
self.file_names = glob.glob(self.image_folder + self.ext)
self.bin_file_names = list()
self.prepare_cache()
def prepare_cache(self):
for fname in self.file_names:
bin_fname = fname.replace(self.image_folder, self.bin_image_folder).replace(self.args.ext, '.npy')
self.bin_file_names.append(bin_fname)
if not os.path.exists(bin_fname):
img = imageio.imread(fname)
np.save(bin_fname, img)
print(f'{bin_fname} prepared!')
def __len__(self): # 返回整个数据集的大小
return len(self.file_names) * 200
@lru_cache(maxsize=400)
def get_ndarray(self, fname):
return np.load(fname)
def __getitem__(self, index):
rgb_range = self.args.rgb_range
n_channels = self.args.n_channels
# img = torch.Tensor(np.load(self.file_names[index]))
img = torch.Tensor(self.get_ndarray(self.bin_file_names[index % len(self.file_names)]))
if img.numpy().ndim == 2:
img = img.unsqueeze(2)
c = img.shape[2]
# input rgb image output y chanel
if n_channels == 1 and c == 3:
img = rgb2ycbcr(img)[:, :, 0].unsqueeze(2)
elif n_channels == 3 and c == 1:
img = img.repeat(1, 1, 3)
w, h, _ = img.shape
th = tw = self.args.im_size
i = random.randint(0, w - tw)
j = random.randint(0, h - th)
img = img[i:i + tw, j:j + th, :]
img_tensor = img.permute(2, 0, 1)
img_tensor = img_tensor * rgb_range / 255.0
img_tensor = self.augment(img_tensor)
return img_tensor
def augment(self, img, hflip=True, rot=True):
hflip = hflip and random.random() < 0.5
vflip = rot and random.random() < 0.5
rot90 = rot and random.random() < 0.5
if hflip:
img = img.flip([1])
if vflip:
img = img.flip([0])
if rot90:
img = img.permute(0, 2, 1)
return img
================================================
FILE: Compressive-Sensing/losses.py
================================================
import torch
import torch.nn as nn
import torch.nn.functional as F
class CharbonnierLoss(nn.Module):
"""Charbonnier Loss (L1)"""
def __init__(self, eps=1e-3):
super(CharbonnierLoss, self).__init__()
self.eps = eps
def forward(self, x, y):
diff = x - y
# loss = torch.sum(torch.sqrt(diff * diff + self.eps))
loss = torch.mean(torch.sqrt((diff * diff) + (self.eps*self.eps)))
return loss
class EdgeLoss(nn.Module):
def __init__(self):
super(EdgeLoss, self).__init__()
k = torch.Tensor([[.05, .25, .4, .25, .05]])
self.kernel = torch.matmul(k.t(),k).unsqueeze(0).repeat(3,1,1,1)
if torch.cuda.is_available():
self.kernel = self.kernel.cuda()
self.loss = CharbonnierLoss()
def conv_gauss(self, img):
n_channels, _, kw, kh = self.kernel.shape
img = F.pad(img, (kw//2, kh//2, kw//2, kh//2), mode='replicate')
return F.conv2d(img, self.kernel, groups=n_channels)
def laplacian_kernel(self, current):
filtered = self.conv_gauss(current) # filter
down = filtered[:,:,::2,::2] # downsample
new_filter = torch.zeros_like(filtered)
new_filter[:,:,::2,::2] = down*4 # upsample
filtered = self.conv_gauss(new_filter) # filter
diff = current - filtered
return diff
def forward(self, x, y):
loss = self.loss(self.laplacian_kernel(x), self.laplacian_kernel(y))
return loss
================================================
FILE: Compressive-Sensing/sampling_matrix_new/README.md
================================================
================================================
FILE: Compressive-Sensing/test.py
================================================
import torch
import torch.nn as nn
from torch.nn import init
import torch.nn.functional as F
import scipy.io as sio
import numpy as np
import os
from torch.utils.data import Dataset, DataLoader
import platform
from argparse import ArgumentParser
import random
import csdata_fast2 as csdata_fast
import cv2
import glob
import math
from skimage.measure import compare_ssim as ssim
# from skimage import measure
from DGUNet_deblock import DGUNet
import utils
parser = ArgumentParser(description='ISTA-Net-plus')
parser.add_argument('--start_epoch', type=int, default=0, help='epoch number of start training')
parser.add_argument('--end_epoch', type=int, default=400, help='epoch number of end training')
parser.add_argument('--layer_num_ICNN', type=int, default=15, help='phase number of ISTA-Net-plus')
parser.add_argument('--layer_num_IFC', type=int, default=5, help='phase number of ISTA-Net-plus')
parser.add_argument('--learning_rate', type=float, default=1e-4, help='learning rate')
parser.add_argument('--group_num', type=int, default=1, help='group number for training')
parser.add_argument('--batch_size', type=int, default=32, help='from {3-10}')
parser.add_argument('--cs_ratio', type=int, default=50, help='from {10, 25, 30, 40, 50}')
parser.add_argument('--gpu_list', type=str, default='0', help='gpu index')
parser.add_argument('--patch_size', type=int, default=32, help='from {1, 4, 10, 25, 40, 50}')
parser.add_argument('--im_size', type=int, default=128, help='from {1, 4, 10, 25, 40, 50}') # 32 or 128
parser.add_argument('--rgb_range', type=int, default=1, help='value range 1 or 255')
parser.add_argument('--n_channels', type=int, default=1, help='1 for gray, 3 for color')
parser.add_argument('--rb_type', type=int, default=1, help='from {1, 2}')
parser.add_argument('--rb_num', type=int, default=2, help='from {3-10}')
parser.add_argument('--matrix_dir', type=str, default='sampling_matrix_new', help='sampling matrix directory')
parser.add_argument('--model_dir', type=str, default='model', help='trained or pre-trained model directory')
parser.add_argument('--data_dir', type=str, default='Datasets/train/BSD400', help='training data directory')
parser.add_argument('--ext', type=str, default='.png', help='training data directory')
parser.add_argument('--log_dir', type=str, default='log', help='log directory')
parser.add_argument('--algo_name', type=str, default='DGUNet_P32_s7_v7_hi3_u4_rp_lm2_learn_BSD_DIV2K_fine', help='log directory')
parser.add_argument('--test_name', type=str, default='BSD68', help='name of test set')
parser.add_argument('--num_gpu',type=int,default=1,help='Number of GPU')
parser.add_argument('--loss_mod',type=int,default=1,help='loss mode')
args = parser.parse_args()
start_epoch = args.start_epoch
end_epoch = args.end_epoch
learning_rate = args.learning_rate
layer_num_ICNN = args.layer_num_ICNN
layer_num_IFC = args.layer_num_IFC
group_num = args.group_num
cs_ratio = args.cs_ratio
gpu_list = args.gpu_list
rb_type = args.rb_type
rb_num = args.rb_num
batch_size = args.batch_size
# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
# os.environ["CUDA_VISIBLE_DEVICES"] = gpu_list
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
ratio_dict = {10: 0, 25: 1, 30: 2, 40: 3, 50: 4}
# n_input = ratio_dict[cs_ratio]
n_output = 1089
nrtrain = 88912 # number of training blocks
psnr_best = 0
best_epoch = 0
# test data
test_name = args.test_name
test_dir = os.path.join('Datasets', test_name)
# filepaths = glob.glob(test_dir + '/*.tif')
filepaths = glob.glob(test_dir + '/*.png')
ImgNum = len(filepaths)
# Load CS Sampling Matrix: phi
Phi_data_Name = os.path.join(args.matrix_dir,
'phi_sampling_%d_%dx%d.npy' % (args.cs_ratio, args.patch_size, args.patch_size))
Phi_input = np.load(Phi_data_Name)
Phi = torch.from_numpy(Phi_input).type(torch.FloatTensor)
Phi = Phi.to(device) # .unsqueeze(0).unsqueeze(0)
length, in_channels = Phi.shape
print('Length:', length, ' In_channels:', in_channels)
model = DGUNet(in_c=1, out_c=1,cs_ratio=args.cs_ratio)#ISTANetplus(layer_num_ICNN, layer_num_IFC, length=length, in_channels=in_channels)
print('GPU: ',list(range(args.num_gpu)))
model = nn.DataParallel(model,device_ids=list(range(args.num_gpu)))
# model = nn.DataParallel(model,device_ids=[0,1])
model = model.to(device)
model_dir = "./%s/CS_%s_layerICNN_%d_layerIFC_%s_group_%d_ratio_%d" % (
args.model_dir, args.algo_name, layer_num_ICNN, layer_num_IFC, group_num, cs_ratio)
result_dir="./results/%s_CS_%s_layerICNN_%d_layerIFC_%s_group_%d_ratio_%d" % (args.test_name,args.algo_name, layer_num_ICNN, layer_num_IFC, group_num, cs_ratio)
if not os.path.exists(result_dir):
os.makedirs(result_dir)
# Load pre-trained model with epoch number
model.load_state_dict(torch.load('%s/net_best.pkl' % (model_dir)))
def col2im_CS_py(X_col, row, col, row_new, col_new):
block_size = args.im_size
X0_rec = np.zeros([row_new, col_new])
count = 0
for x in range(0, row_new - block_size + 1, block_size):
for y in range(0, col_new - block_size + 1, block_size):
X0_rec[x:x + block_size, y:y + block_size] = X_col[:, count].reshape([block_size, block_size])
count = count + 1
X_rec = X0_rec[:row, :col]
return X_rec
def psnr(img1, img2):
img1.astype(np.float32)
img2.astype(np.float32)
mse = np.mean((img1 - img2) ** 2)
if mse == 0:
return 100
PIXEL_MAX = 255.0
return 20 * math.log10(PIXEL_MAX / math.sqrt(mse))
def imread_CS_py(Iorg):
block_size = args.im_size
[row, col] = Iorg.shape
row_pad = block_size - np.mod(row, block_size)
col_pad = block_size - np.mod(col, block_size)
Ipad = np.concatenate((Iorg, np.zeros([row, col_pad])), axis=1)
Ipad = np.concatenate((Ipad, np.zeros([row_pad, col + col_pad])), axis=0)
[row_new, col_new] = Ipad.shape
return [Iorg, row, col, Ipad, row_new, col_new]
def img2col_py(Ipad, block_size):
[row, col] = Ipad.shape
row_block = row / block_size
col_block = col / block_size
block_num = int(row_block * col_block)
img_col = np.zeros([block_size ** 2, block_num])
count = 0
for x in range(0, row - block_size + 1, block_size):
for y in range(0, col - block_size + 1, block_size):
img_col[:, count] = Ipad[x:x + block_size, y:y + block_size].reshape([-1])
count = count + 1
return img_col
print_flag = 1 # print parameter number
if print_flag:
num_count = 0
num_params = 0
for para in model.parameters():
num_count += 1
num_params += para.numel()
# print('Layer %d' % num_count)
# print(para.size())
print("total para num: %d" % num_params)
class RandomDataset(Dataset):
def __init__(self, data, length):
self.data = data
self.len = length
def __getitem__(self, index):
return torch.Tensor(self.data[index, :]).float()
def __len__(self):
return self.len
def augment_img_tensor(img, mode=0):
img_size = img.size()
img_np = img.data.cpu().numpy()
if len(img_size) == 3:
img_np = np.transpose(img_np, (1, 2, 0))
elif len(img_size) == 4:
img_np = np.transpose(img_np, (2, 3, 1, 0))
img_np = augment_img(img_np, mode=mode)
img_tensor = torch.from_numpy(np.ascontiguousarray(img_np))
if len(img_size) == 3:
img_tensor = img_tensor.permute(2, 0, 1)
elif len(img_size) == 4:
img_tensor = img_tensor.permute(3, 2, 0, 1)
return img_tensor.type_as(img)
def augment_img(img, mode=0):
if mode == 0:
return img
elif mode == 1:
return np.flipud(np.rot90(img))
elif mode == 2:
return np.flipud(img)
elif mode == 3:
return np.rot90(img, k=3)
elif mode == 4:
return np.flipud(np.rot90(img, k=2))
elif mode == 5:
return np.rot90(img)
elif mode == 6:
return np.rot90(img, k=2)
elif mode == 7:
return np.flipud(np.rot90(img, k=3))
def test_x8(model, L):
E_list = [model(augment_img_tensor(L, mode=i))[0] for i in range(8)]
for i in range(len(E_list)):
if i == 3 or i == 5:
E_list[i] = augment_img_tensor(E_list[i], mode=8 - i)
else:
E_list[i] = augment_img_tensor(E_list[i], mode=i)
output_cat = torch.stack(E_list, dim=0)
E = output_cat.mean(dim=0, keepdim=False)
return E
# Testing
model.eval()
with torch.no_grad():
psnr_ave = 0
ssim_ave = 0
for img_no in range(ImgNum):
imgName = filepaths[img_no]
Img = cv2.imread(imgName, 1)
Img_yuv = cv2.cvtColor(Img, cv2.COLOR_BGR2YCrCb)
Iorg_y = Img_yuv[:, :, 0]
[Iorg, row, col, Ipad, row_new, col_new] = imread_CS_py(Iorg_y)
Icol = img2col_py(Ipad, args.im_size).transpose() / 255.0
Img_output = Icol
batch_x = torch.from_numpy(Img_output)
batch_x = batch_x.type(torch.FloatTensor)
batch_x = batch_x.to(device)
batch_x = batch_x.view(batch_x.shape[0], 1, args.im_size, args.im_size)
x_output = model(batch_x)#test_x8(model,batch_x)#model(batch_x)#[0] # torch.mm(batch_x,
x_output=x_output[0]
x_output = x_output.view(x_output.shape[0], -1)
Prediction_value = x_output.cpu().data.numpy()
X_rec = np.clip(col2im_CS_py(Prediction_value.transpose(), row, col, row_new, col_new), 0, 1)
utils.save_img((os.path.join(result_dir, imgName.split('/')[-1])), (X_rec*255.).astype(np.uint8))
rec_PSNR = psnr(X_rec * 255, Iorg.astype(np.float64))
rec_ssim = ssim(X_rec * 255, Iorg.astype(np.float64), data_range=255)
print(imgName.split('/')[-1],rec_PSNR)
psnr_ave += rec_PSNR
ssim_ave += rec_ssim
del x_output
psnr_ave /= ImgNum
ssim_ave /= ImgNum
print('best ssim is %.4f best psnr is %.4f in epoch %d psnr_rec: %.4f' % (ssim_ave,psnr_best, best_epoch, psnr_ave))
================================================
FILE: Compressive-Sensing/test_plus.py
================================================
import torch
import torch.nn as nn
from torch.nn import init
import torch.nn.functional as F
import scipy.io as sio
import numpy as np
import os
from torch.utils.data import Dataset, DataLoader
import platform
from argparse import ArgumentParser
import random
import csdata_fast2 as csdata_fast
import cv2
import glob
import math
from skimage.measure import compare_ssim as ssim
from DGUNet_plus_deblock import DGUNet
import utils
parser = ArgumentParser(description='ISTA-Net-plus')
parser.add_argument('--start_epoch', type=int, default=0, help='epoch number of start training')
parser.add_argument('--end_epoch', type=int, default=400, help='epoch number of end training')
parser.add_argument('--layer_num_ICNN', type=int, default=15, help='phase number of ISTA-Net-plus')
parser.add_argument('--layer_num_IFC', type=int, default=5, help='phase number of ISTA-Net-plus')
parser.add_argument('--learning_rate', type=float, default=1e-4, help='learning rate')
parser.add_argument('--group_num', type=int, default=1, help='group number for training')
parser.add_argument('--batch_size', type=int, default=32, help='from {3-10}')
parser.add_argument('--cs_ratio', type=int, default=50, help='from {10, 25, 30, 40, 50}')
parser.add_argument('--gpu_list', type=str, default='0', help='gpu index')
parser.add_argument('--patch_size', type=int, default=32, help='from {1, 4, 10, 25, 40, 50}')
parser.add_argument('--im_size', type=int, default=128, help='from {1, 4, 10, 25, 40, 50}') # 32 or 128
parser.add_argument('--rgb_range', type=int, default=1, help='value range 1 or 255')
parser.add_argument('--n_channels', type=int, default=1, help='1 for gray, 3 for color')
parser.add_argument('--rb_type', type=int, default=1, help='from {1, 2}')
parser.add_argument('--rb_num', type=int, default=2, help='from {3-10}')
parser.add_argument('--matrix_dir', type=str, default='sampling_matrix_new', help='sampling matrix directory')
parser.add_argument('--model_dir', type=str, default='model', help='trained or pre-trained model directory')
parser.add_argument('--data_dir', type=str, default='Datasets/train/BSD400', help='training data directory')
parser.add_argument('--ext', type=str, default='.png', help='training data directory')
parser.add_argument('--log_dir', type=str, default='log', help='log directory')
parser.add_argument('--algo_name', type=str, default='DGUNet_P32_s7_v7_hi3_u4_lm2_learn_BSD_DIV2K_fine', help='log directory')
parser.add_argument('--test_name', type=str, default='BSD68', help='name of test set')
# parser.add_argument('--test_name', type=str, default='Set11', help='name of test set')
parser.add_argument('--num_gpu',type=int,default=1,help='Number of GPU')
parser.add_argument('--loss_mod',type=int,default=1,help='loss mode')
args = parser.parse_args()
start_epoch = args.start_epoch
end_epoch = args.end_epoch
learning_rate = args.learning_rate
layer_num_ICNN = args.layer_num_ICNN
layer_num_IFC = args.layer_num_IFC
group_num = args.group_num
cs_ratio = args.cs_ratio
gpu_list = args.gpu_list
rb_type = args.rb_type
rb_num = args.rb_num
batch_size = args.batch_size
# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
# os.environ["CUDA_VISIBLE_DEVICES"] = gpu_list
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
ratio_dict = {10: 0, 25: 1, 30: 2, 40: 3, 50: 4}
# n_input = ratio_dict[cs_ratio]
n_output = 1089
nrtrain = 88912 # number of training blocks
psnr_best = 0
best_epoch = 0
# test data
test_name = args.test_name
test_dir = os.path.join('Datasets', test_name)
# filepaths = glob.glob(test_dir + '/*.tif')
filepaths = glob.glob(test_dir + '/*.png')
ImgNum = len(filepaths)
# Load CS Sampling Matrix: phi
Phi_data_Name = os.path.join(args.matrix_dir,
'phi_sampling_%d_%dx%d.npy' % (args.cs_ratio, args.patch_size, args.patch_size))
Phi_input = np.load(Phi_data_Name)
Phi = torch.from_numpy(Phi_input).type(torch.FloatTensor)
Phi = Phi.to(device) # .unsqueeze(0).unsqueeze(0)
length, in_channels = Phi.shape
print('Length:', length, ' In_channels:', in_channels)
model = DGUNet(in_c=1, out_c=1,cs_ratio=args.cs_ratio)
print('GPU: ',list(range(args.num_gpu)))
model = nn.DataParallel(model,device_ids=list(range(args.num_gpu)))
# model = nn.DataParallel(model,device_ids=[0,1])
model = model.to(device)
model_dir = "./%s/CS_%s_layerICNN_%d_layerIFC_%s_group_%d_ratio_%d" % (
args.model_dir, args.algo_name, layer_num_ICNN, layer_num_IFC, group_num, cs_ratio)
result_dir="./results/%s_CS_%s_layerICNN_%d_layerIFC_%s_group_%d_ratio_%d" % (args.test_name,args.algo_name, layer_num_ICNN, layer_num_IFC, group_num, cs_ratio)
if not os.path.exists(result_dir):
os.makedirs(result_dir)
# Load pre-trained model with epoch number
model.load_state_dict(torch.load('%s/net_best.pkl' % (model_dir)))
def col2im_CS_py(X_col, row, col, row_new, col_new):
block_size = args.im_size
X0_rec = np.zeros([row_new, col_new])
count = 0
for x in range(0, row_new - block_size + 1, block_size):
for y in range(0, col_new - block_size + 1, block_size):
X0_rec[x:x + block_size, y:y + block_size] = X_col[:, count].reshape([block_size, block_size])
count = count + 1
X_rec = X0_rec[:row, :col]
return X_rec
def psnr(img1, img2):
img1.astype(np.float32)
img2.astype(np.float32)
mse = np.mean((img1 - img2) ** 2)
if mse == 0:
return 100
PIXEL_MAX = 255.0
return 20 * math.log10(PIXEL_MAX / math.sqrt(mse))
def imread_CS_py(Iorg):
block_size = args.im_size
[row, col] = Iorg.shape
row_pad = block_size - np.mod(row, block_size)
col_pad = block_size - np.mod(col, block_size)
Ipad = np.concatenate((Iorg, np.zeros([row, col_pad])), axis=1)
Ipad = np.concatenate((Ipad, np.zeros([row_pad, col + col_pad])), axis=0)
[row_new, col_new] = Ipad.shape
return [Iorg, row, col, Ipad, row_new, col_new]
def img2col_py(Ipad, block_size):
[row, col] = Ipad.shape
row_block = row / block_size
col_block = col / block_size
block_num = int(row_block * col_block)
img_col = np.zeros([block_size ** 2, block_num])
count = 0
for x in range(0, row - block_size + 1, block_size):
for y in range(0, col - block_size + 1, block_size):
img_col[:, count] = Ipad[x:x + block_size, y:y + block_size].reshape([-1])
count = count + 1
return img_col
print_flag = 1 # print parameter number
if print_flag:
num_count = 0
num_params = 0
for para in model.parameters():
num_count += 1
num_params += para.numel()
# print('Layer %d' % num_count)
# print(para.size())
print("total para num: %d" % num_params)
class RandomDataset(Dataset):
def __init__(self, data, length):
self.data = data
self.len = length
def __getitem__(self, index):
return torch.Tensor(self.data[index, :]).float()
def __len__(self):
return self.len
def augment_img_tensor(img, mode=0):
img_size = img.size()
img_np = img.data.cpu().numpy()
if len(img_size) == 3:
img_np = np.transpose(img_np, (1, 2, 0))
elif len(img_size) == 4:
img_np = np.transpose(img_np, (2, 3, 1, 0))
img_np = augment_img(img_np, mode=mode)
img_tensor = torch.from_numpy(np.ascontiguousarray(img_np))
if len(img_size) == 3:
img_tensor = img_tensor.permute(2, 0, 1)
elif len(img_size) == 4:
img_tensor = img_tensor.permute(3, 2, 0, 1)
return img_tensor.type_as(img)
def augment_img(img, mode=0):
if mode == 0:
return img
elif mode == 1:
return np.flipud(np.rot90(img))
elif mode == 2:
return np.flipud(img)
elif mode == 3:
return np.rot90(img, k=3)
elif mode == 4:
return np.flipud(np.rot90(img, k=2))
elif mode == 5:
return np.rot90(img)
elif mode == 6:
return np.rot90(img, k=2)
elif mode == 7:
return np.flipud(np.rot90(img, k=3))
def test_x8(model, L):
E_list = [model(augment_img_tensor(L, mode=i))[0] for i in range(8)]
for i in range(len(E_list)):
if i == 3 or i == 5:
E_list[i] = augment_img_tensor(E_list[i], mode=8 - i)
else:
E_list[i] = augment_img_tensor(E_list[i], mode=i)
output_cat = torch.stack(E_list, dim=0)
E = output_cat.mean(dim=0, keepdim=False)
return E
# Testing
model.eval()
with torch.no_grad():
psnr_ave = 0
ssim_ave = 0
for img_no in range(ImgNum):
imgName = filepaths[img_no]
Img = cv2.imread(imgName, 1)
Img_yuv = cv2.cvtColor(Img, cv2.COLOR_BGR2YCrCb)
Iorg_y = Img_yuv[:, :, 0]
[Iorg, row, col, Ipad, row_new, col_new] = imread_CS_py(Iorg_y)
Icol = img2col_py(Ipad, args.im_size).transpose() / 255.0
Img_output = Icol
batch_x = torch.from_numpy(Img_output)
batch_x = batch_x.type(torch.FloatTensor)
batch_x = batch_x.to(device)
batch_x = batch_x.view(batch_x.shape[0], 1, args.im_size, args.im_size)
x_output = model(batch_x)#test_x8(model,batch_x)#model(batch_x)#[0] # torch.mm(batch_x,
x_output=x_output[0]
x_output = x_output.view(x_output.shape[0], -1)
Prediction_value = x_output.cpu().data.numpy()
X_rec = np.clip(col2im_CS_py(Prediction_value.transpose(), row, col, row_new, col_new), 0, 1)
utils.save_img((os.path.join(result_dir, imgName.split('/')[-1])), (X_rec*255.).astype(np.uint8))
rec_PSNR = psnr(X_rec * 255, Iorg.astype(np.float64))
rec_ssim = ssim(X_rec * 255, Iorg.astype(np.float64), data_range=255)
print(imgName.split('/')[-1],rec_PSNR)
psnr_ave += rec_PSNR
ssim_ave += rec_ssim
del x_output
psnr_ave /= ImgNum
ssim_ave /= ImgNum
print('best ssim is %.4f best psnr is %.4f in epoch %d psnr_rec: %.4f' % (ssim_ave,psnr_best, best_epoch, psnr_ave))
================================================
FILE: Compressive-Sensing/train.py
================================================
import torch
import torch.nn as nn
from torch.nn import init
import torch.nn.functional as F
import scipy.io as sio
import numpy as np
import os
from torch.utils.data import Dataset, DataLoader
import platform
from argparse import ArgumentParser
import random
import csdata_fast
import cv2
import glob
import math
from DGUNet import DGUNet
parser = ArgumentParser(description='ISTA-Net-plus')
parser.add_argument('--start_epoch', type=int, default=0, help='epoch number of start training')
parser.add_argument('--end_epoch', type=int, default=400, help='epoch number of end training')
parser.add_argument('--layer_num_ICNN', type=int, default=15, help='phase number of ISTA-Net-plus')
parser.add_argument('--layer_num_IFC', type=int, default=5, help='phase number of ISTA-Net-plus')
parser.add_argument('--learning_rate', type=float, default=1e-4, help='learning rate')
parser.add_argument('--group_num', type=int, default=1, help='group number for training')
parser.add_argument('--batch_size', type=int, default=32, help='from {3-10}')
parser.add_argument('--cs_ratio', type=int, default=25, help='from {10, 25, 30, 40, 50}')
parser.add_argument('--gpu_list', type=str, default='0', help='gpu index')
parser.add_argument('--patch_size', type=int, default=32, help='from {1, 4, 10, 25, 40, 50}')
parser.add_argument('--rgb_range', type=int, default=1, help='value range 1 or 255')
parser.add_argument('--n_channels', type=int, default=1, help='1 for gray, 3 for color')
parser.add_argument('--rb_type', type=int, default=1, help='from {1, 2}')
parser.add_argument('--rb_num', type=int, default=2, help='from {3-10}')
parser.add_argument('--matrix_dir', type=str, default='sampling_matrix_new', help='sampling matrix directory')
parser.add_argument('--model_dir', type=str, default='model', help='trained or pre-trained model directory')
parser.add_argument('--data_dir', type=str, default='Datasets/train/DIV2K_BSD400', help='training data directory')
parser.add_argument('--ext', type=str, default='.png', help='training data directory')
parser.add_argument('--log_dir', type=str, default='log', help='log directory')
parser.add_argument('--algo_name', type=str, default='DGUNet_P32_s7_v7_hi3_u4_rp_lm2_learn_BSD_DIV2K', help='log directory')
parser.add_argument('--test_name', type=str, default='Set11', help='name of test set')
parser.add_argument('--num_gpu',type=int,default=1,help='Number of GPU')
parser.add_argument('--loss_mod',type=int,default=1,help='loss mode')
args = parser.parse_args()
start_epoch = args.start_epoch
end_epoch = args.end_epoch
learning_rate = args.learning_rate
layer_num_ICNN = args.layer_num_ICNN
layer_num_IFC = args.layer_num_IFC
group_num = args.group_num
cs_ratio = args.cs_ratio
gpu_list = args.gpu_list
rb_type = args.rb_type
rb_num = args.rb_num
batch_size = args.batch_size
# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
# os.environ["CUDA_VISIBLE_DEVICES"] = gpu_list
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
ratio_dict = {10: 0, 25: 1, 30: 2, 40: 3, 50: 4}
# n_input = ratio_dict[cs_ratio]
n_output = 1089
nrtrain = 88912 # number of training blocks
psnr_best = 0
best_epoch = 0
# test data
test_name = args.test_name
test_dir = os.path.join('Datasets', test_name)
filepaths = glob.glob(test_dir + '/*.tif')
ImgNum = len(filepaths)
# Load CS Sampling Matrix: phi
Phi_data_Name = os.path.join(args.matrix_dir,
'phi_sampling_%d_%dx%d.npy' % (args.cs_ratio, args.patch_size, args.patch_size))
Phi_input = np.load(Phi_data_Name)
Phi = torch.from_numpy(Phi_input).type(torch.FloatTensor)
Phi = Phi.to(device) # .unsqueeze(0).unsqueeze(0)
length, in_channels = Phi.shape
print('Length:', length, ' In_channels:', in_channels)
model = DGUNet(in_c=1, out_c=1,cs_ratio=args.cs_ratio)
print('GPU: ',list(range(args.num_gpu)))
model = nn.DataParallel(model,device_ids=list(range(args.num_gpu)))
# model = nn.DataParallel(model,device_ids=[0,1])
model = model.to(device)
def col2im_CS_py(X_col, row, col, row_new, col_new):
block_size = args.patch_size
X0_rec = np.zeros([row_new, col_new])
count = 0
for x in range(0, row_new - block_size + 1, block_size):
for y in range(0, col_new - block_size + 1, block_size):
X0_rec[x:x + block_size, y:y + block_size] = X_col[:, count].reshape([block_size, block_size])
count = count + 1
X_rec = X0_rec[:row, :col]
return X_rec
def psnr(img1, img2):
img1.astype(np.float32)
img2.astype(np.float32)
mse = np.mean((img1 - img2) ** 2)
if mse == 0:
return 100
PIXEL_MAX = 255.0
return 20 * math.log10(PIXEL_MAX / math.sqrt(mse))
def imread_CS_py(Iorg):
block_size = args.patch_size
[row, col] = Iorg.shape
row_pad = block_size - np.mod(row, block_size)
col_pad = block_size - np.mod(col, block_size)
Ipad = np.concatenate((Iorg, np.zeros([row, col_pad])), axis=1)
Ipad = np.concatenate((Ipad, np.zeros([row_pad, col + col_pad])), axis=0)
[row_new, col_new] = Ipad.shape
return [Iorg, row, col, Ipad, row_new, col_new]
def img2col_py(Ipad, block_size):
[row, col] = Ipad.shape
row_block = row / block_size
col_block = col / block_size
block_num = int(row_block * col_block)
img_col = np.zeros([block_size ** 2, block_num])
count = 0
for x in range(0, row - block_size + 1, block_size):
for y in range(0, col - block_size + 1, block_size):
img_col[:, count] = Ipad[x:x + block_size, y:y + block_size].reshape([-1])
count = count + 1
return img_col
print_flag = 1 # print parameter number
if print_flag:
num_count = 0
num_params = 0
for para in model.parameters():
num_count += 1
num_params += para.numel()
print('Layer %d' % num_count)
print(para.size())
print("total para num: %d" % num_params)
class RandomDataset(Dataset):
def __init__(self, data, length):
self.data = data
self.len = length
def __getitem__(self, index):
return torch.Tensor(self.data[index, :]).float()
def __len__(self):
return self.len
training_data = csdata_fast.SlowDataset(args)
if (platform.system() == "Windows"):
rand_loader = DataLoader(dataset=training_data, batch_size=batch_size, num_workers=0,
shuffle=True)
else:
print('linux')
rand_loader = DataLoader(dataset=training_data, batch_size=batch_size, num_workers=24,
shuffle=True)
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
model_dir = "./%s/CS_%s_layerICNN_%d_layerIFC_%s_group_%d_ratio_%d" % (
args.model_dir, args.algo_name, layer_num_ICNN, layer_num_IFC, group_num, cs_ratio)
log_file_name = "./%s/Log_CS_%s_layerICNN_%d_layerIFC_%d_group_%d_ratio_%d.txt" % (
args.log_dir, args.algo_name, layer_num_ICNN, layer_num_IFC, group_num, cs_ratio)
if not os.path.exists(model_dir):
os.makedirs(model_dir)
if start_epoch > 0:
pre_model_dir = model_dir
model.load_state_dict(torch.load('./%s/net_params_%d.pkl' % (pre_model_dir, start_epoch)))
step_all = len(rand_loader)
# Training loop
for epoch_i in range(start_epoch + 1, end_epoch + 1):
model.train()
for step, data in enumerate(rand_loader):
batch_x = data
batch_x = batch_x.to(device)
batch_x = batch_x.view(batch_x.shape[0], args.patch_size * args.patch_size)
batch_x = batch_x.view(batch_x.shape[0], 1, args.patch_size, args.patch_size)
x_output_f = model(batch_x)
if args.loss_mod==1:
loss_all = np.sum([torch.mean(torch.pow(torch.clamp(x_output_f[j], 0, 1) - batch_x,2)) for j in range(len(x_output_f))])
else:
loss_all = torch.mean(torch.pow(torch.clamp(x_output_f[0], 0, 1) - batch_x,2))
optimizer.zero_grad()
loss_all.backward()
optimizer.step()
if step % 40 == 0:
output_data = "[Epoch: %02d/%02d Step: %d/%d] Total Loss: %.4f" % (
epoch_i, end_epoch, step, step_all, loss_all.item())
print(output_data)
output_file = open(log_file_name, 'a')
output_file.write(output_data)
output_file.close()
model.eval()
with torch.no_grad():
psnr_ave = 0
for img_no in range(ImgNum):
imgName = filepaths[img_no]
Img = cv2.imread(imgName, 1)
Img_yuv = cv2.cvtColor(Img, cv2.COLOR_BGR2YCrCb)
Iorg_y = Img_yuv[:, :, 0]
[Iorg, row, col, Ipad, row_new, col_new] = imread_CS_py(Iorg_y)
Icol = img2col_py(Ipad, args.patch_size).transpose() / 255.0
Img_output = Icol
batch_x = torch.from_numpy(Img_output)
batch_x = batch_x.type(torch.FloatTensor)
batch_x = batch_x.to(device)
# Phix = torch.mm(batch_x, torch.transpose(Phi, 0, 1)) # compression result
# PhixPhiT = torch.mm(Phix, Phi)
batch_x = batch_x.view(batch_x.shape[0], 1, args.patch_size, args.patch_size)
x_output = model(batch_x)[0] # torch.mm(batch_x,
x_output = x_output.view(x_output.shape[0], -1)
Prediction_value = x_output.cpu().data.numpy()
X_rec = np.clip(col2im_CS_py(Prediction_value.transpose(), row, col, row_new, col_new), 0, 1)
rec_PSNR = psnr(X_rec * 255, Iorg.astype(np.float64))
psnr_ave += rec_PSNR
del x_output
psnr_ave /= ImgNum
if psnr_ave > psnr_best:
best_epoch = epoch_i
psnr_best = psnr_ave
torch.save(model.state_dict(), "./%s/net_best.pkl" % (model_dir)) # save only the parameters
torch.save(model.state_dict(), "./%s/net_last.pkl" % (model_dir)) # save only the parameters
print('best psnr is %.4f in epoch %d psnr_rec: %.4f' % (psnr_best, best_epoch, psnr_ave))
================================================
FILE: Compressive-Sensing/train_deblock.py
================================================
import torch
import torch.nn as nn
from torch.nn import init
import torch.nn.functional as F
import scipy.io as sio
import numpy as np
import os
from torch.utils.data import Dataset, DataLoader
import platform
from argparse import ArgumentParser
import random
import csdata_fast2 as csdata_fast
import cv2
import glob
import math
from DGUNet_deblock import DGUNet
parser = ArgumentParser(description='ISTA-Net-plus')
parser.add_argument('--start_epoch', type=int, default=0, help='epoch number of start training')
parser.add_argument('--end_epoch', type=int, default=100, help='epoch number of end training')
parser.add_argument('--layer_num_ICNN', type=int, default=15, help='phase number of ISTA-Net-plus')
parser.add_argument('--layer_num_IFC', type=int, default=5, help='phase number of ISTA-Net-plus')
parser.add_argument('--learning_rate', type=float, default=1e-4, help='learning rate')
parser.add_argument('--group_num', type=int, default=1, help='group number for training')
parser.add_argument('--batch_size', type=int, default=32, help='from {3-10}')
parser.add_argument('--cs_ratio', type=int, default=25, help='from {10, 25, 30, 40, 50}')
parser.add_argument('--gpu_list', type=str, default='0', help='gpu index')
parser.add_argument('--patch_size', type=int, default=64, help='from {1, 4, 10, 25, 40, 50}')
parser.add_argument('--im_size', type=int, default=32, help='from {1, 4, 10, 25, 40, 50}')
parser.add_argument('--rgb_range', type=int, default=1, help='value range 1 or 255')
parser.add_argument('--n_channels', type=int, default=1, help='1 for gray, 3 for color')
parser.add_argument('--rb_type', type=int, default=1, help='from {1, 2}')
parser.add_argument('--rb_num', type=int, default=2, help='from {3-10}')
parser.add_argument('--matrix_dir', type=str, default='sampling_matrix_new', help='sampling matrix directory')
parser.add_argument('--model_dir', type=str, default='model', help='trained or pre-trained model directory')
parser.add_argument('--ckp_dir', type=str, default='DGUNet_P32_s7_v7_hi3_u4_rp_lm2_learn_BSD_DIV2K', help='trained or pre-trained model directory')
parser.add_argument('--data_dir', type=str, default='Datasets/train/DIV2K_BSD400', help='training data directory')
parser.add_argument('--ext', type=str, default='.png', help='training data directory')
parser.add_argument('--log_dir', type=str, default='log', help='log directory')
parser.add_argument('--algo_name', type=str, default='DGUNet_P32_s7_v7_hi3_u4_rp_lm2_learn_BSD_DIV2K_fine', help='log directory')
parser.add_argument('--test_name', type=str, default='Set11', help='name of test set')
parser.add_argument('--num_gpu',type=int,default=1,help='Number of GPU')
parser.add_argument('--loss_mod',type=int,default=1,help='loss mode')
args = parser.parse_args()
start_epoch = args.start_epoch
end_epoch = args.end_epoch
learning_rate = args.learning_rate
layer_num_ICNN = args.layer_num_ICNN
layer_num_IFC = args.layer_num_IFC
group_num = args.group_num
cs_ratio = args.cs_ratio
gpu_list = args.gpu_list
rb_type = args.rb_type
rb_num = args.rb_num
batch_size = args.batch_size
# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
# os.environ["CUDA_VISIBLE_DEVICES"] = gpu_list
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
ratio_dict = {10: 0, 25: 1, 30: 2, 40: 3, 50: 4}
# n_input = ratio_dict[cs_ratio]
n_output = 1089
nrtrain = 88912 # number of training blocks
psnr_best = 0
best_epoch = 0
# test data
test_name = args.test_name
test_dir = os.path.join('Datasets', test_name)
filepaths = glob.glob(test_dir + '/*.tif')
ImgNum = len(filepaths)
# Load CS Sampling Matrix: phi
Phi_data_Name = os.path.join(args.matrix_dir,
'phi_sampling_%d_%dx%d.npy' % (args.cs_ratio, args.patch_size, args.patch_size))
Phi_input = np.load(Phi_data_Name)
Phi = torch.from_numpy(Phi_input).type(torch.FloatTensor)
Phi = Phi.to(device) # .unsqueeze(0).unsqueeze(0)
length, in_channels = Phi.shape
print('Length:', length, ' In_channels:', in_channels)
model = DGUNet(in_c=1, out_c=1,cs_ratio=args.cs_ratio)
print('GPU: ',list(range(args.num_gpu)))
model = nn.DataParallel(model,device_ids=list(range(args.num_gpu)))
# model = nn.DataParallel(model,device_ids=[0,1])
model = model.to(device)
ckp_path="./%s/CS_%s_layerICNN_%d_layerIFC_%s_group_%d_ratio_%d" % (
args.model_dir, args.ckp_dir, layer_num_ICNN, layer_num_IFC, group_num, cs_ratio)
model.load_state_dict(torch.load('./%s/net_best.pkl' % (ckp_path)))
def col2im_CS_py(X_col, row, col, row_new, col_new):
block_size = args.im_size
X0_rec = np.zeros([row_new, col_new])
count = 0
for x in range(0, row_new - block_size + 1, block_size):
for y in range(0, col_new - block_size + 1, block_size):
X0_rec[x:x + block_size, y:y + block_size] = X_col[:, count].reshape([block_size, block_size])
count = count + 1
X_rec = X0_rec[:row, :col]
return X_rec
def psnr(img1, img2):
img1.astype(np.float32)
img2.astype(np.float32)
mse = np.mean((img1 - img2) ** 2)
if mse == 0:
return 100
PIXEL_MAX = 255.0
return 20 * math.log10(PIXEL_MAX / math.sqrt(mse))
def imread_CS_py(Iorg):
block_size = args.im_size
[row, col] = Iorg.shape
row_pad = block_size - np.mod(row, block_size)
col_pad = block_size - np.mod(col, block_size)
Ipad = np.concatenate((Iorg, np.zeros([row, col_pad])), axis=1)
Ipad = np.concatenate((Ipad, np.zeros([row_pad, col + col_pad])), axis=0)
[row_new, col_new] = Ipad.shape
return [Iorg, row, col, Ipad, row_new, col_new]
def img2col_py(Ipad, block_size):
[row, col] = Ipad.shape
row_block = row / block_size
col_block = col / block_size
block_num = int(row_block * col_block)
img_col = np.zeros([block_size ** 2, block_num])
count = 0
for x in range(0, row - block_size + 1, block_size):
for y in range(0, col - block_size + 1, block_size):
img_col[:, count] = Ipad[x:x + block_size, y:y + block_size].reshape([-1])
count = count + 1
return img_col
print_flag = 1 # print parameter number
if print_flag:
num_count = 0
num_params = 0
for para in model.parameters():
num_count += 1
num_params += para.numel()
print('Layer %d' % num_count)
print(para.size())
print("total para num: %d" % num_params)
class RandomDataset(Dataset):
def __init__(self, data, length):
self.data = data
self.len = length
def __getitem__(self, index):
return torch.Tensor(self.data[index, :]).float()
def __len__(self):
return self.len
training_data = csdata_fast.SlowDataset(args)
if (platform.system() == "Windows"):
rand_loader = DataLoader(dataset=training_data, batch_size=batch_size, num_workers=0,
shuffle=True)
else:
rand_loader = DataLoader(dataset=training_data, batch_size=batch_size, num_workers=8,
shuffle=True)
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[60, 120, 180, 240], gamma=0.5)
model_dir = "./%s/CS_%s_layerICNN_%d_layerIFC_%s_group_%d_ratio_%d" % (
args.model_dir, args.algo_name, layer_num_ICNN, layer_num_IFC, group_num, cs_ratio)
log_file_name = "./%s/Log_CS_%s_layerICNN_%d_layerIFC_%d_group_%d_ratio_%d.txt" % (
args.log_dir, args.algo_name, layer_num_ICNN, layer_num_IFC, group_num, cs_ratio)
if not os.path.exists(model_dir):
os.makedirs(model_dir)
if start_epoch > 0:
pre_model_dir = model_dir
model.load_state_dict(torch.load('./%s/net_params_%d.pkl' % (pre_model_dir, start_epoch)))
step_all = len(rand_loader)
# Training loop
for epoch_i in range(start_epoch + 1, end_epoch + 1):
model.train()
print('Learning rate: ',scheduler.get_lr()[0])
for step, data in enumerate(rand_loader):
batch_x = data
batch_x = batch_x.to(device)
batch_x = batch_x.view(batch_x.shape[0], 1, args.im_size, args.im_size)
x_output_f = model(batch_x)
if args.loss_mod==1:
loss_all = np.sum([torch.mean(torch.pow(torch.clamp(x_output_f[j], 0, 1) - batch_x,2)) for j in range(len(x_output_f))])
else:
loss_all = torch.mean(torch.pow(torch.clamp(x_output_f[0], 0, 1) - batch_x,2))
optimizer.zero_grad()
loss_all.backward()
optimizer.step()
if step % 40 == 0:
output_data = "[Epoch: %02d/%02d Step: %d/%d] Total Loss: %.4f" % (
epoch_i, end_epoch, step, step_all, loss_all.item())
print(output_data)
scheduler.step()
output_file = open(log_file_name, 'a')
output_file.write(output_data)
output_file.close()
model.eval()
with torch.no_grad():
psnr_ave = 0
for img_no in range(ImgNum):
imgName = filepaths[img_no]
Img = cv2.imread(imgName, 1)
Img_yuv = cv2.cvtColor(Img, cv2.COLOR_BGR2YCrCb)
Iorg_y = Img_yuv[:, :, 0]
[Iorg, row, col, Ipad, row_new, col_new] = imread_CS_py(Iorg_y)
Icol = img2col_py(Ipad, args.im_size).transpose() / 255.0
Img_output = Icol
batch_x = torch.from_numpy(Img_output)
batch_x = batch_x.type(torch.FloatTensor)
batch_x = batch_x.to(device)
# Phix = torch.mm(batch_x, torch.transpose(Phi, 0, 1)) # compression result
# PhixPhiT = torch.mm(Phix, Phi)
batch_x = batch_x.view(batch_x.shape[0], 1, args.im_size, args.im_size)
x_output = model(batch_x)[0] # torch.mm(batch_x,
x_output = x_output.view(x_output.shape[0], -1)
Prediction_value = x_output.cpu().data.numpy()
X_rec = np.clip(col2im_CS_py(Prediction_value.transpose(), row, col, row_new, col_new), 0, 1)
rec_PSNR = psnr(X_rec * 255, Iorg.astype(np.float64))
psnr_ave += rec_PSNR
del x_output
psnr_ave /= ImgNum
if psnr_ave > psnr_best:
best_epoch = epoch_i
psnr_best = psnr_ave
torch.save(model.state_dict(), "./%s/net_best.pkl" % (model_dir)) # save only the parameters
torch.save(model.state_dict(), "./%s/net_last.pkl" % (model_dir)) # save only the parameters
print('best psnr is %.4f in epoch %d psnr_rec: %.4f' % (psnr_best, best_epoch, psnr_ave))
================================================
FILE: Compressive-Sensing/train_plus.py
================================================
import torch
import torch.nn as nn
from torch.nn import init
import torch.nn.functional as F
import scipy.io as sio
import numpy as np
import os
from torch.utils.data import Dataset, DataLoader
import platform
from argparse import ArgumentParser
import random
import csdata_fast
import cv2
import glob
import math
from DGUNet_plus import DGUNet
parser = ArgumentParser(description='ISTA-Net-plus')
parser.add_argument('--start_epoch', type=int, default=0, help='epoch number of start training')
parser.add_argument('--end_epoch', type=int, default=400, help='epoch number of end training')
parser.add_argument('--layer_num_ICNN', type=int, default=15, help='phase number of ISTA-Net-plus')
parser.add_argument('--layer_num_IFC', type=int, default=5, help='phase number of ISTA-Net-plus')
parser.add_argument('--learning_rate', type=float, default=1e-4, help='learning rate')
parser.add_argument('--group_num', type=int, default=1, help='group number for training')
parser.add_argument('--batch_size', type=int, default=32, help='from {3-10}')
parser.add_argument('--cs_ratio', type=int, default=25, help='from {10, 25, 30, 40, 50}')
parser.add_argument('--gpu_list', type=str, default='0', help='gpu index')
parser.add_argument('--patch_size', type=int, default=32, help='from {1, 4, 10, 25, 40, 50}')
parser.add_argument('--rgb_range', type=int, default=1, help='value range 1 or 255')
parser.add_argument('--n_channels', type=int, default=1, help='1 for gray, 3 for color')
parser.add_argument('--rb_type', type=int, default=1, help='from {1, 2}')
parser.add_argument('--rb_num', type=int, default=2, help='from {3-10}')
parser.add_argument('--matrix_dir', type=str, default='sampling_matrix_new', help='sampling matrix directory')
parser.add_argument('--model_dir', type=str, default='model', help='trained or pre-trained model directory')
parser.add_argument('--data_dir', type=str, default='Datasets/train/DIV2K_BSD400', help='training data directory')
parser.add_argument('--ext', type=str, default='.png', help='training data directory')
parser.add_argument('--log_dir', type=str, default='log', help='log directory')
parser.add_argument('--algo_name', type=str, default='DGUNet_P32_s7_v7_hi3_u4_lm2_learn_BSD_DIV2K', help='log directory')
parser.add_argument('--test_name', type=str, default='Set11', help='name of test set')
parser.add_argument('--num_gpu',type=int,default=1,help='Number of GPU')
parser.add_argument('--loss_mod',type=int,default=1,help='loss mode')
args = parser.parse_args()
start_epoch = args.start_epoch
end_epoch = args.end_epoch
learning_rate = args.learning_rate
layer_num_ICNN = args.layer_num_ICNN
layer_num_IFC = args.layer_num_IFC
group_num = args.group_num
cs_ratio = args.cs_ratio
gpu_list = args.gpu_list
rb_type = args.rb_type
rb_num = args.rb_num
batch_size = args.batch_size
# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
# os.environ["CUDA_VISIBLE_DEVICES"] = gpu_list
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
ratio_dict = {10: 0, 25: 1, 30: 2, 40: 3, 50: 4}
# n_input = ratio_dict[cs_ratio]
n_output = 1089
nrtrain = 88912 # number of training blocks
psnr_best = 0
best_epoch = 0
# test data
test_name = args.test_name
test_dir = os.path.join('Datasets', test_name)
filepaths = glob.glob(test_dir + '/*.tif')
ImgNum = len(filepaths)
# Load CS Sampling Matrix: phi
Phi_data_Name = os.path.join(args.matrix_dir,
'phi_sampling_%d_%dx%d.npy' % (args.cs_ratio, args.patch_size, args.patch_size))
Phi_input = np.load(Phi_data_Name)
Phi = torch.from_numpy(Phi_input).type(torch.FloatTensor)
Phi = Phi.to(device) # .unsqueeze(0).unsqueeze(0)
length, in_channels = Phi.shape
print('Length:', length, ' In_channels:', in_channels)
model = DGUNet(in_c=1, out_c=1,cs_ratio=args.cs_ratio)
print('GPU: ',list(range(args.num_gpu)))
model = nn.DataParallel(model,device_ids=list(range(args.num_gpu)))
# model = nn.DataParallel(model,device_ids=[0,1])
model = model.to(device)
def col2im_CS_py(X_col, row, col, row_new, col_new):
block_size = args.patch_size
X0_rec = np.zeros([row_new, col_new])
count = 0
for x in range(0, row_new - block_size + 1, block_size):
for y in range(0, col_new - block_size + 1, block_size):
X0_rec[x:x + block_size, y:y + block_size] = X_col[:, count].reshape([block_size, block_size])
count = count + 1
X_rec = X0_rec[:row, :col]
return X_rec
def psnr(img1, img2):
img1.astype(np.float32)
img2.astype(np.float32)
mse = np.mean((img1 - img2) ** 2)
if mse == 0:
return 100
PIXEL_MAX = 255.0
return 20 * math.log10(PIXEL_MAX / math.sqrt(mse))
def imread_CS_py(Iorg):
block_size = args.patch_size
[row, col] = Iorg.shape
row_pad = block_size - np.mod(row, block_size)
col_pad = block_size - np.mod(col, block_size)
Ipad = np.concatenate((Iorg, np.zeros([row, col_pad])), axis=1)
Ipad = np.concatenate((Ipad, np.zeros([row_pad, col + col_pad])), axis=0)
[row_new, col_new] = Ipad.shape
return [Iorg, row, col, Ipad, row_new, col_new]
def img2col_py(Ipad, block_size):
[row, col] = Ipad.shape
row_block = row / block_size
col_block = col / block_size
block_num = int(row_block * col_block)
img_col = np.zeros([block_size ** 2, block_num])
count = 0
for x in range(0, row - block_size + 1, block_size):
for y in range(0, col - block_size + 1, block_size):
img_col[:, count] = Ipad[x:x + block_size, y:y + block_size].reshape([-1])
count = count + 1
return img_col
print_flag = 1 # print parameter number
if print_flag:
num_count = 0
num_params = 0
for para in model.parameters():
num_count += 1
num_params += para.numel()
print('Layer %d' % num_count)
print(para.size())
print("total para num: %d" % num_params)
class RandomDataset(Dataset):
def __init__(self, data, length):
self.data = data
self.len = length
def __getitem__(self, index):
return torch.Tensor(self.data[index, :]).float()
def __len__(self):
return self.len
training_data = csdata_fast.SlowDataset(args)
if (platform.system() == "Windows"):
rand_loader = DataLoader(dataset=training_data, batch_size=batch_size, num_workers=0,
shuffle=True)
else:
print('linux')
rand_loader = DataLoader(dataset=training_data, batch_size=batch_size, num_workers=24,
shuffle=True)
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
model_dir = "./%s/CS_%s_layerICNN_%d_layerIFC_%s_group_%d_ratio_%d" % (
args.model_dir, args.algo_name, layer_num_ICNN, layer_num_IFC, group_num, cs_ratio)
log_file_name = "./%s/Log_CS_%s_layerICNN_%d_layerIFC_%d_group_%d_ratio_%d.txt" % (
args.log_dir, args.algo_name, layer_num_ICNN, layer_num_IFC, group_num, cs_ratio)
if not os.path.exists(model_dir):
os.makedirs(model_dir)
if start_epoch > 0:
pre_model_dir = model_dir
model.load_state_dict(torch.load('./%s/net_params_%d.pkl' % (pre_model_dir, start_epoch)))
step_all = len(rand_loader)
# Training loop
for epoch_i in range(start_epoch + 1, end_epoch + 1):
model.train()
for step, data in enumerate(rand_loader):
batch_x = data
batch_x = batch_x.to(device)
batch_x = batch_x.view(batch_x.shape[0], args.patch_size * args.patch_size)
batch_x = batch_x.view(batch_x.shape[0], 1, args.patch_size, args.patch_size)
x_output_f = model(batch_x)
if args.loss_mod==1:
loss_all = np.sum([torch.mean(torch.pow(torch.clamp(x_output_f[j], 0, 1) - batch_x,2)) for j in range(len(x_output_f))])
else:
loss_all = torch.mean(torch.pow(torch.clamp(x_output_f[0], 0, 1) - batch_x,2))
optimizer.zero_grad()
loss_all.backward()
optimizer.step()
if step % 40 == 0:
output_data = "[Epoch: %02d/%02d Step: %d/%d] Total Loss: %.4f" % (
epoch_i, end_epoch, step, step_all, loss_all.item())
print(output_data)
output_file = open(log_file_name, 'a')
output_file.write(output_data)
output_file.close()
model.eval()
with torch.no_grad():
psnr_ave = 0
for img_no in range(ImgNum):
imgName = filepaths[img_no]
Img = cv2.imread(imgName, 1)
Img_yuv = cv2.cvtColor(Img, cv2.COLOR_BGR2YCrCb)
Iorg_y = Img_yuv[:, :, 0]
[Iorg, row, col, Ipad, row_new, col_new] = imread_CS_py(Iorg_y)
Icol = img2col_py(Ipad, args.patch_size).transpose() / 255.0
Img_output = Icol
batch_x = torch.from_numpy(Img_output)
batch_x = batch_x.type(torch.FloatTensor)
batch_x = batch_x.to(device)
# Phix = torch.mm(batch_x, torch.transpose(Phi, 0, 1)) # compression result
# PhixPhiT = torch.mm(Phix, Phi)
batch_x = batch_x.view(batch_x.shape[0], 1, args.patch_size, args.patch_size)
x_output = model(batch_x)[0] # torch.mm(batch_x,
x_output = x_output.view(x_output.shape[0], -1)
Prediction_value = x_output.cpu().data.numpy()
X_rec = np.clip(col2im_CS_py(Prediction_value.transpose(), row, col, row_new, col_new), 0, 1)
rec_PSNR = psnr(X_rec * 255, Iorg.astype(np.float64))
psnr_ave += rec_PSNR
del x_output
psnr_ave /= ImgNum
if psnr_ave > psnr_best:
best_epoch = epoch_i
psnr_best = psnr_ave
torch.save(model.state_dict(), "./%s/net_best.pkl" % (model_dir)) # save only the parameters
torch.save(model.state_dict(), "./%s/net_last.pkl" % (model_dir)) # save only the parameters
print('best psnr is %.4f in epoch %d psnr_rec: %.4f' % (psnr_best, best_epoch, psnr_ave))
================================================
FILE: Compressive-Sensing/train_plus_deblock.py
================================================
import torch
import torch.nn as nn
from torch.nn import init
import torch.nn.functional as F
import scipy.io as sio
import numpy as np
import os
from torch.utils.data import Dataset, DataLoader
import platform
from argparse import ArgumentParser
import random
import csdata_fast2 as csdata_fast
import cv2
import glob
import math
from DGUNet_plus_deblock import DGUNet
parser = ArgumentParser(description='ISTA-Net-plus')
parser.add_argument('--start_epoch', type=int, default=0, help='epoch number of start training')
parser.add_argument('--end_epoch', type=int, default=100, help='epoch number of end training')
parser.add_argument('--layer_num_ICNN', type=int, default=15, help='phase number of ISTA-Net-plus')
parser.add_argument('--layer_num_IFC', type=int, default=5, help='phase number of ISTA-Net-plus')
parser.add_argument('--learning_rate', type=float, default=1e-4, help='learning rate')
parser.add_argument('--group_num', type=int, default=1, help='group number for training')
parser.add_argument('--batch_size', type=int, default=32, help='from {3-10}')
parser.add_argument('--cs_ratio', type=int, default=25, help='from {10, 25, 30, 40, 50}')
parser.add_argument('--gpu_list', type=str, default='0', help='gpu index')
parser.add_argument('--patch_size', type=int, default=64, help='from {1, 4, 10, 25, 40, 50}')
parser.add_argument('--im_size', type=int, default=32, help='from {1, 4, 10, 25, 40, 50}')
parser.add_argument('--rgb_range', type=int, default=1, help='value range 1 or 255')
parser.add_argument('--n_channels', type=int, default=1, help='1 for gray, 3 for color')
parser.add_argument('--rb_type', type=int, default=1, help='from {1, 2}')
parser.add_argument('--rb_num', type=int, default=2, help='from {3-10}')
parser.add_argument('--matrix_dir', type=str, default='sampling_matrix_new', help='sampling matrix directory')
parser.add_argument('--model_dir', type=str, default='model', help='trained or pre-trained model directory')
parser.add_argument('--ckp_dir', type=str, default='DGUNet_P32_s7_v7_hi3_u4_lm2_learn_BSD_DIV2K', help='trained or pre-trained model directory')
parser.add_argument('--data_dir', type=str, default='Datasets/train/DIV2K_BSD400', help='training data directory')
parser.add_argument('--ext', type=str, default='.png', help='training data directory')
parser.add_argument('--log_dir', type=str, default='log', help='log directory')
parser.add_argument('--algo_name', type=str, default='DGUNet_P32_s7_v7_hi3_u4_lm2_learn_BSD_DIV2K_fine', help='log directory')
parser.add_argument('--test_name', type=str, default='Set11', help='name of test set')
parser.add_argument('--num_gpu',type=int,default=1,help='Number of GPU')
parser.add_argument('--loss_mod',type=int,default=1,help='loss mode')
args = parser.parse_args()
start_epoch = args.start_epoch
end_epoch = args.end_epoch
learning_rate = args.learning_rate
layer_num_ICNN = args.layer_num_ICNN
layer_num_IFC = args.layer_num_IFC
group_num = args.group_num
cs_ratio = args.cs_ratio
gpu_list = args.gpu_list
rb_type = args.rb_type
rb_num = args.rb_num
batch_size = args.batch_size
# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
# os.environ["CUDA_VISIBLE_DEVICES"] = gpu_list
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
ratio_dict = {10: 0, 25: 1, 30: 2, 40: 3, 50: 4}
# n_input = ratio_dict[cs_ratio]
n_output = 1089
nrtrain = 88912 # number of training blocks
psnr_best = 0
best_epoch = 0
# test data
test_name = args.test_name
test_dir = os.path.join('Datasets', test_name)
filepaths = glob.glob(test_dir + '/*.tif')
ImgNum = len(filepaths)
# Load CS Sampling Matrix: phi
Phi_data_Name = os.path.join(args.matrix_dir,
'phi_sampling_%d_%dx%d.npy' % (args.cs_ratio, args.patch_size, args.patch_size))
Phi_input = np.load(Phi_data_Name)
Phi = torch.from_numpy(Phi_input).type(torch.FloatTensor)
Phi = Phi.to(device) # .unsqueeze(0).unsqueeze(0)
length, in_channels = Phi.shape
print('Length:', length, ' In_channels:', in_channels)
model = DGUNet(in_c=1, out_c=1,cs_ratio=args.cs_ratio)
print('GPU: ',list(range(args.num_gpu)))
model = nn.DataParallel(model,device_ids=list(range(args.num_gpu)))
# model = nn.DataParallel(model,device_ids=[0,1])
model = model.to(device)
ckp_path="./%s/CS_%s_layerICNN_%d_layerIFC_%s_group_%d_ratio_%d" % (
args.model_dir, args.ckp_dir, layer_num_ICNN, layer_num_IFC, group_num, cs_ratio)
model.load_state_dict(torch.load('./%s/net_best.pkl' % (ckp_path)))
def col2im_CS_py(X_col, row, col, row_new, col_new):
block_size = args.im_size
X0_rec = np.zeros([row_new, col_new])
count = 0
for x in range(0, row_new - block_size + 1, block_size):
for y in range(0, col_new - block_size + 1, block_size):
X0_rec[x:x + block_size, y:y + block_size] = X_col[:, count].reshape([block_size, block_size])
count = count + 1
X_rec = X0_rec[:row, :col]
return X_rec
def psnr(img1, img2):
img1.astype(np.float32)
img2.astype(np.float32)
mse = np.mean((img1 - img2) ** 2)
if mse == 0:
return 100
PIXEL_MAX = 255.0
return 20 * math.log10(PIXEL_MAX / math.sqrt(mse))
def imread_CS_py(Iorg):
block_size = args.im_size
[row, col] = Iorg.shape
row_pad = block_size - np.mod(row, block_size)
col_pad = block_size - np.mod(col, block_size)
Ipad = np.concatenate((Iorg, np.zeros([row, col_pad])), axis=1)
Ipad = np.concatenate((Ipad, np.zeros([row_pad, col + col_pad])), axis=0)
[row_new, col_new] = Ipad.shape
return [Iorg, row, col, Ipad, row_new, col_new]
def img2col_py(Ipad, block_size):
[row, col] = Ipad.shape
row_block = row / block_size
col_block = col / block_size
block_num = int(row_block * col_block)
img_col = np.zeros([block_size ** 2, block_num])
count = 0
for x in range(0, row - block_size + 1, block_size):
for y in range(0, col - block_size + 1, block_size):
img_col[:, count] = Ipad[x:x + block_size, y:y + block_size].reshape([-1])
count = count + 1
return img_col
print_flag = 1 # print parameter number
if print_flag:
num_count = 0
num_params = 0
for para in model.parameters():
num_count += 1
num_params += para.numel()
print('Layer %d' % num_count)
print(para.size())
print("total para num: %d" % num_params)
class RandomDataset(Dataset):
def __init__(self, data, length):
self.data = data
self.len = length
def __getitem__(self, index):
return torch.Tensor(self.data[index, :]).float()
def __len__(self):
return self.len
training_data = csdata_fast.SlowDataset(args)
if (platform.system() == "Windows"):
rand_loader = DataLoader(dataset=training_data, batch_size=batch_size, num_workers=0,
shuffle=True)
else:
rand_loader = DataLoader(dataset=training_data, batch_size=batch_size, num_workers=8,
shuffle=True)
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[60, 120, 180, 240], gamma=0.5)
model_dir = "./%s/CS_%s_layerICNN_%d_layerIFC_%s_group_%d_ratio_%d" % (
args.model_dir, args.algo_name, layer_num_ICNN, layer_num_IFC, group_num, cs_ratio)
log_file_name = "./%s/Log_CS_%s_layerICNN_%d_layerIFC_%d_group_%d_ratio_%d.txt" % (
args.log_dir, args.algo_name, layer_num_ICNN, layer_num_IFC, group_num, cs_ratio)
if not os.path.exists(model_dir):
os.makedirs(model_dir)
if start_epoch > 0:
pre_model_dir = model_dir
model.load_state_dict(torch.load('./%s/net_params_%d.pkl' % (pre_model_dir, start_epoch)))
step_all = len(rand_loader)
# Training loop
for epoch_i in range(start_epoch + 1, end_epoch + 1):
model.train()
print('Learning rate: ',scheduler.get_lr()[0])
for step, data in enumerate(rand_loader):
batch_x = data
batch_x = batch_x.to(device)
batch_x = batch_x.view(batch_x.shape[0], 1, args.im_size, args.im_size)
x_output_f = model(batch_x)
if args.loss_mod==1:
loss_all = np.sum([torch.mean(torch.pow(torch.clamp(x_output_f[j], 0, 1) - batch_x,2)) for j in range(len(x_output_f))])
else:
loss_all = torch.mean(torch.pow(torch.clamp(x_output_f[0], 0, 1) - batch_x,2))
optimizer.zero_grad()
loss_all.backward()
optimizer.step()
if step % 40 == 0:
output_data = "[Epoch: %02d/%02d Step: %d/%d] Total Loss: %.4f" % (
epoch_i, end_epoch, step, step_all, loss_all.item())
print(output_data)
scheduler.step()
output_file = open(log_file_name, 'a')
output_file.write(output_data)
output_file.close()
model.eval()
with torch.no_grad():
psnr_ave = 0
for img_no in range(ImgNum):
imgName = filepaths[img_no]
Img = cv2.imread(imgName, 1)
Img_yuv = cv2.cvtColor(Img, cv2.COLOR_BGR2YCrCb)
Iorg_y = Img_yuv[:, :, 0]
[Iorg, row, col, Ipad, row_new, col_new] = imread_CS_py(Iorg_y)
Icol = img2col_py(Ipad, args.im_size).transpose() / 255.0
Img_output = Icol
batch_x = torch.from_numpy(Img_output)
batch_x = batch_x.type(torch.FloatTensor)
batch_x = batch_x.to(device)
# Phix = torch.mm(batch_x, torch.transpose(Phi, 0, 1)) # compression result
# PhixPhiT = torch.mm(Phix, Phi)
batch_x = batch_x.view(batch_x.shape[0], 1, args.im_size, args.im_size)
x_output = model(batch_x)[0] # torch.mm(batch_x,
x_output = x_output.view(x_output.shape[0], -1)
Prediction_value = x_output.cpu().data.numpy()
X_rec = np.clip(col2im_CS_py(Prediction_value.transpose(), row, col, row_new, col_new), 0, 1)
rec_PSNR = psnr(X_rec * 255, Iorg.astype(np.float64))
psnr_ave += rec_PSNR
del x_output
psnr_ave /= ImgNum
if psnr_ave > psnr_best:
best_epoch = epoch_i
psnr_best = psnr_ave
torch.save(model.state_dict(), "./%s/net_best.pkl" % (model_dir)) # save only the parameters
torch.save(model.state_dict(), "./%s/net_last.pkl" % (model_dir)) # save only the parameters
print('best psnr is %.4f in epoch %d psnr_rec: %.4f' % (psnr_best, best_epoch, psnr_ave))
================================================
FILE: Compressive-Sensing/training.yml
================================================
# ###############
# ##
# ####
# GPU: [0,1,2,3]
# VERBOSE: True
# MODEL:
# MODE: 'Denoising'
# SESSION: 'MPRNet'
# # Optimization arguments.
# OPTIM:
# BATCH_SIZE: 16
# NUM_EPOCHS: 80
# # NEPOCH_DECAY: [10]
# LR_INITIAL: 2e-4
# LR_MIN: 1e-6
# # BETA1: 0.9
# TRAINING:
# VAL_AFTER_EVERY: 1
# RESUME: False
# TRAIN_PS: 128
# VAL_PS: 196
# TRAIN_DIR: './Datasets/train' # path to training data
# VAL_DIR: './Datasets/val' # path to validation data
# SAVE_DIR: './checkpoints' # path to save models and images
# # SAVE_IMAGES: False
###############
##
####
GPU: [0,1,2,3]
VERBOSE: True
MODEL:
MODE: 'Inpaint'
SESSION: 'MPRNet'
# Optimization arguments.
OPTIM:
BATCH_SIZE: 16
NUM_EPOCHS: 250
# NEPOCH_DECAY: [10]
LR_INITIAL: 2e-4
LR_MIN: 1e-6
# BETA1: 0.9
TRAINING:
VAL_AFTER_EVERY: 5
RESUME: False
TRAIN_PS: 128
VAL_PS: 128
TRAIN_DIR: './Datasets/train' # path to training data
VAL_DIR: './Datasets/val' # path to validation data
SAVE_DIR: './checkpoints' # path to save models and images
# SAVE_IMAGES: False
================================================
FILE: Compressive-Sensing/training_cs.yml
================================================
# ###############
# ##
# ####
# GPU: [0,1,2,3]
# VERBOSE: True
# MODEL:
# MODE: 'Denoising'
# SESSION: 'MPRNet'
# # Optimization arguments.
# OPTIM:
# BATCH_SIZE: 16
# NUM_EPOCHS: 80
# # NEPOCH_DECAY: [10]
# LR_INITIAL: 2e-4
# LR_MIN: 1e-6
# # BETA1: 0.9
# TRAINING:
# VAL_AFTER_EVERY: 1
# RESUME: False
# TRAIN_PS: 128
# VAL_PS: 196
# TRAIN_DIR: './Datasets/train' # path to training data
# VAL_DIR: './Datasets/val' # path to validation data
# SAVE_DIR: './checkpoints' # path to save models and images
# # SAVE_IMAGES: False
###############
##
####
GPU: [0,1,2,3]
VERBOSE: True
MODEL:
MODE: 'CS'
SESSION: 'MPRNet'
# Optimization arguments.
OPTIM:
BATCH_SIZE: 16
NUM_EPOCHS: 250
# NEPOCH_DECAY: [10]
LR_INITIAL: 2e-4
LR_MIN: 1e-6
# BETA1: 0.9
TRAINING:
VAL_AFTER_EVERY: 5
RESUME: False
TRAIN_PS: 128
VAL_PS: 128
TRAIN_DIR: './Datasets/train' # path to training data
VAL_DIR: './Datasets/val' # path to validation data
SAVE_DIR: './checkpoints' # path to save models and images
# SAVE_IMAGES: False
================================================
FILE: Compressive-Sensing/utils/README.md
================================================
================================================
FILE: Compressive-Sensing/utils/__init__.py
================================================
from .dir_utils import *
from .image_utils import *
from .model_utils import *
from .dataset_utils import *
================================================
FILE: Compressive-Sensing/utils/dataset_utils.py
================================================
import torch
class MixUp_AUG:
def __init__(self):
self.dist = torch.distributions.beta.Beta(torch.tensor([0.6]), torch.tensor([0.6]))
def aug(self, rgb_gt, rgb_noisy):
bs = rgb_gt.size(0)
indices = torch.randperm(bs)
rgb_gt2 = rgb_gt[indices]
rgb_noisy2 = rgb_noisy[indices]
lam = self.dist.rsample((bs,1)).view(-1,1,1,1).cuda()
rgb_gt = lam * rgb_gt + (1-lam) * rgb_gt2
rgb_noisy = lam * rgb_noisy + (1-lam) * rgb_noisy2
return rgb_gt, rgb_noisy
================================================
FILE: Compressive-Sensing/utils/dir_utils.py
================================================
import os
from natsort import natsorted
from glob import glob
def mkdirs(paths):
if isinstance(paths, list) and not isinstance(paths, str):
for path in paths:
mkdir(path)
else:
mkdir(paths)
def mkdir(path):
if not os.path.exists(path):
os.makedirs(path)
def get_last_path(path, session):
x = natsorted(glob(os.path.join(path,'*%s'%session)))[-1]
return x
================================================
FILE: Compressive-Sensing/utils/image_utils.py
================================================
import torch
import numpy as np
import cv2
def torchPSNR(tar_img, prd_img):
imdff = torch.clamp(prd_img,0,1) - torch.clamp(tar_img,0,1)
rmse = (imdff**2).mean().sqrt()
ps = 20*torch.log10(1/rmse)
return ps
def save_img(filepath, img):
cv2.imwrite(filepath,cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
def numpyPSNR(tar_img, prd_img):
imdff = np.float32(prd_img) - np.float32(tar_img)
rmse = np.sqrt(np.mean(imdff**2))
ps = 20*np.log10(255/rmse)
return ps
================================================
FILE: Compressive-Sensing/utils/model_utils.py
================================================
import torch
import os
from collections import OrderedDict
def freeze(model):
for p in model.parameters():
p.requires_grad=False
def unfreeze(model):
for p in model.parameters():
p.requires_grad=True
def is_frozen(model):
x = [p.requires_grad for p in model.parameters()]
return not all(x)
def save_checkpoint(model_dir, state, session):
epoch = state['epoch']
model_out_path = os.path.join(model_dir,"model_epoch_{}_{}.pth".format(epoch,session))
torch.save(state, model_out_path)
def load_checkpoint(model, weights):
checkpoint = torch.load(weights)
try:
model.load_state_dict(checkpoint["state_dict"])
except:
state_dict = checkpoint["state_dict"]
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k[7:] # remove `module.`
new_state_dict[name] = v
model.load_state_dict(new_state_dict)
def load_checkpoint_multigpu(model, weights):
checkpoint = torch.load(weights)
state_dict = checkpoint["state_dict"]
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k[7:] # remove `module.`
new_state_dict[name] = v
model.load_state_dict(new_state_dict)
def load_start_epoch(weights):
checkpoint = torch.load(weights)
epoch = checkpoint["epoch"]
return epoch
def load_optim(optimizer, weights):
checkpoint = torch.load(weights)
optimizer.load_state_dict(checkpoint['optimizer'])
# for p in optimizer.param_groups: lr = p['lr']
# return lr
================================================
FILE: Deblurring/.ipynb_checkpoints/DGUNet-checkpoint.py
================================================
import torch
import torch.nn as nn
import torch.nn.functional as F
from pdb import set_trace as stx
##########################################################################
# Basic modules
def conv(in_channels, out_channels, kernel_size, bias=False, stride=1):
return nn.Conv2d(
in_channels, out_channels, kernel_size,
padding=(kernel_size // 2), bias=bias, stride=stride)
def conv_down(in_chn, out_chn, bias=False):
layer = nn.Conv2d(in_chn, out_chn, kernel_size=4, stride=2, padding=1, bias=bias)
return layer
def default_conv(in_channels, out_channels, kernel_size,stride=1, bias=True):
return nn.Conv2d(
in_channels, out_channels, kernel_size,
padding=(kernel_size//2),stride=stride, bias=bias)
class ResBlock(nn.Module):
def __init__(
self, conv, n_feats, kernel_size,
bias=True, bn=False, act=nn.PReLU(), res_scale=1):
super(ResBlock, self).__init__()
m = []
for i in range(2):
if i == 0:
m.append(conv(n_feats, 64, kernel_size, bias=bias))
else:
m.append(conv(64, n_feats, kernel_size, bias=bias))
if bn:
m.append(nn.BatchNorm2d(n_feats))
if i == 0:
m.append(act)
self.body = nn.Sequential(*m)
self.res_scale = res_scale
def forward(self, x):
res = self.body(x).mul(self.res_scale)
res += x
return res
##########################################################################
## Channel Attention Block (CAB)
class CAB(nn.Module):
def __init__(self, n_feat, kernel_size, reduction, bias, act):
super(CAB, self).__init__()
modules_body = []
modules_body.append(conv(n_feat, n_feat, kernel_size, bias=bias))
modules_body.append(act)
modules_body.append(conv(n_feat, n_feat, kernel_size, bias=bias))
self.CA = CALayer(n_feat, reduction, bias=bias)
self.body = nn.Sequential(*modules_body)
def forward(self, x):
res = self.body(x)
res = self.CA(res)
res += x
return res
class CALayer(nn.Module):
def __init__(self, channel, reduction=16, bias=False):
super(CALayer, self).__init__()
# global average pooling: feature --> point
self.avg_pool = nn.AdaptiveAvgPool2d(1)
# feature channel downscale and upscale --> channel weight
self.conv_du = nn.Sequential(
nn.Conv2d(channel, channel // reduction, 1, padding=0, bias=bias),
nn.ReLU(inplace=True),
nn.Conv2d(channel // reduction, channel, 1, padding=0, bias=bias),
nn.Sigmoid()
)
def forward(self, x):
y = self.avg_pool(x)
y = self.conv_du(y)
return x * y
##########################################################################
## Compute inter-stage features
class SAM(nn.Module):
def __init__(self, n_feat, kernel_size, bias):
super(SAM, self).__init__()
self.conv1 = conv(n_feat, n_feat, kernel_size, bias=bias)
self.conv2 = conv(n_feat, 3, kernel_size, bias=bias)
def forward(self, x, x_img):
x1 = self.conv1(x)
img = self.conv2(x) + x_img
x1 = x1 + x
return x1, img
class mergeblock(nn.Module):
def __init__(self, n_feat, kernel_size, bias, subspace_dim=16):
super(mergeblock, self).__init__()
self.conv_block = conv(n_feat * 2, n_feat, kernel_size, bias=bias)
self.num_subspace = subspace_dim
self.subnet = conv(n_feat * 2, self.num_subspace, kernel_size, bias=bias)
def forward(self, x, bridge):
out = torch.cat([x, bridge], 1)
b_, c_, h_, w_ = bridge.shape
sub = self.subnet(out)
V_t = sub.view(b_, self.num_subspace, h_*w_)
V_t = V_t / (1e-6 + torch.abs(V_t).sum(axis=2, keepdims=True))
V = V_t.permute(0, 2, 1)
mat = torch.matmul(V_t, V)
mat_inv = torch.inverse(mat)
project_mat = torch.matmul(mat_inv, V_t)
bridge_ = bridge.view(b_, c_, h_*w_)
project_feature = torch.matmul(project_mat, bridge_.permute(0, 2, 1))
bridge = torch.matmul(V, project_feature).permute(0, 2, 1).view(b_, c_, h_, w_)
out = torch.cat([x, bridge], 1)
out = self.conv_block(out)
return out+x
##########################################################################
## U-Net
class Encoder(nn.Module):
def __init__(self, n_feat, kernel_size, reduction, act, bias, scale_unetfeats, csff,depth=5):
super(Encoder, self).__init__()
self.body=nn.ModuleList()#[]
self.depth=depth
for i in range(depth-1):
self.body.append(UNetConvBlock(in_size=n_feat+scale_unetfeats*i, out_size=n_feat+scale_unetfeats*(i+1), downsample=True, relu_slope=0.2, use_csff=csff, use_HIN=True))
self.body.append(UNetConvBlock(in_size=n_feat+scale_unetfeats*(depth-1), out_size=n_feat+scale_unetfeats*(depth-1), downsample=False, relu_slope=0.2, use_csff=csff, use_HIN=True))
def forward(self, x, encoder_outs=None, decoder_outs=None):
res=[]
if encoder_outs is not None and decoder_outs is not None:
for i,down in enumerate(self.body):
if (i+1) < self.depth:
x, x_up = down(x,encoder_outs[i],decoder_outs[-i-1])
res.append(x_up)
else:
x = down(x)
else:
for i,down in enumerate(self.body):
if (i+1) < self.depth:
x, x_up = down(x)
res.append(x_up)
else:
x = down(x)
return res,x
class UNetConvBlock(nn.Module):
def __init__(self, in_size, out_size, downsample, relu_slope, use_csff=False, use_HIN=False):
super(UNetConvBlock, self).__init__()
self.downsample = downsample
self.identity = nn.Conv2d(in_size, out_size, 1, 1, 0)
self.use_csff = use_csff
self.conv_1 = nn.Conv2d(in_size, out_size, kernel_size=3, padding=1, bias=True)
self.relu_1 = nn.LeakyReLU(relu_slope, inplace=False)
self.conv_2 = nn.Conv2d(out_size, out_size, kernel_size=3, padding=1, bias=True)
self.relu_2 = nn.LeakyReLU(relu_slope, inplace=False)
if downsample and use_csff:
self.csff_enc = nn.Conv2d(out_size, out_size, 3, 1, 1)
self.csff_dec = nn.Conv2d(in_size, out_size, 3, 1, 1)
self.phi = nn.Conv2d(out_size, out_size, 3, 1, 1)
self.gamma = nn.Conv2d(out_size, out_size, 3, 1, 1)
if use_HIN:
self.norm = nn.InstanceNorm2d(out_size//2, affine=True)
self.use_HIN = use_HIN
if downsample:
self.downsample = conv_down(out_size, out_size, bias=False)
def forward(self, x, enc=None, dec=None):
out = self.conv_1(x)
if self.use_HIN:
out_1, out_2 = torch.chunk(out, 2, dim=1)
out = torch.cat([self.norm(out_1), out_2], dim=1)
out = self.relu_1(out)
out = self.relu_2(self.conv_2(out))
out += self.identity(x)
if enc is not None and dec is not None:
assert self.use_csff
skip_ = F.leaky_relu(self.csff_enc(enc) + self.csff_dec(dec), 0.1, inplace=True)
out = out*F.sigmoid(self.phi(skip_)) + self.gamma(skip_) + out
if self.downsample:
out_down = self.downsample(out)
return out_down, out
else:
return out
class UNetUpBlock(nn.Module):
def __init__(self, in_size, out_size, relu_slope):
super(UNetUpBlock, self).__init__()
self.up = nn.ConvTranspose2d(in_size, out_size, kernel_size=2, stride=2, bias=True)
self.conv_block = UNetConvBlock(out_size*2, out_size, False, relu_slope)
def forward(self, x, bridge):
up = self.up(x)
out = torch.cat([up, bridge], 1)
out = self.conv_block(out)
return out
class Decoder(nn.Module):
def __init__(self, n_feat, kernel_size, reduction, act, bias, scale_unetfeats,depth=5):
super(Decoder, self).__init__()
self.body=nn.ModuleList()
self.skip_conv=nn.ModuleList()#[]
for i in range(depth-1):
self.body.append(UNetUpBlock(in_size=n_feat+scale_unetfeats*(depth-i-1), out_size=n_feat+scale_unetfeats*(depth-i-2), relu_slope=0.2))
self.skip_conv.append(nn.Conv2d(n_feat+scale_unetfeats*(depth-i-1), n_feat+scale_unetfeats*(depth-i-2), 3, 1, 1))
def forward(self, x, bridges):
res=[]
for i,up in enumerate(self.body):
x=up(x,self.skip_conv[i](bridges[-i-1]))
res.append(x)
return res
##########################################################################
##---------- Resizing Modules ----------
class DownSample(nn.Module):
def __init__(self, in_channels, s_factor):
super(DownSample, self).__init__()
self.down = nn.Sequential(nn.Upsample(scale_factor=0.5, mode='bilinear', align_corners=False),
nn.Conv2d(in_channels, in_channels + s_factor, 1, stride=1, padding=0, bias=False))
def forward(self, x):
x = self.down(x)
return x
class UpSample(nn.Module):
def __init__(self, in_channels, out_channels):
super(UpSample, self).__init__()
self.up = nn.Sequential(nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
nn.Conv2d(in_channels, out_channels, 1, stride=1, padding=0, bias=False))
def forward(self, x):
x = self.up(x)
return x
##########################################################################
## One Stage
class Basic_block(nn.Module):
def __init__(self, in_c=3, out_c=3, n_feat=80, scale_unetfeats=48, scale_orsnetfeats=32, num_cab=8, kernel_size=3, reduction=4, bias=False):
super(Basic_block, self).__init__()
act = nn.PReLU()
self.phi_1 = ResBlock(default_conv,3,3)
self.phit_1 = ResBlock(default_conv,3,3)
self.shallow_feat2 = nn.Sequential(conv(3, n_feat, kernel_size, bias=bias),
CAB(n_feat, kernel_size, reduction, bias=bias, act=act))
self.stage2_encoder = Encoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats,depth=4, csff=True)
self.stage2_decoder = Decoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats,depth=4)
self.sam23 = SAM(n_feat, kernel_size=1, bias=bias)
self.r1 = nn.Parameter(torch.Tensor([0.5]))
self.concat12 = conv(n_feat * 2, n_feat, kernel_size, bias=bias)
self.merge12=mergeblock(n_feat,3,True)
def forward(self, img,stage1_img,feat1,res1,x2_samfeats):
## GDM
phixsy_2 = self.phi_1(stage1_img) - img
x2_img = stage1_img - self.r1*self.phit_1(phixsy_2)
## PMM
x2 = self.shallow_feat2(x2_img)
x2_cat = self.merge12(x2, x2_samfeats)
feat2,feat_fin2 = self.stage2_encoder(x2_cat, feat1, res1)
res2 = self.stage2_decoder(feat_fin2,feat2)
x3_samfeats, stage2_img = self.sam23(res2[-1], x2_img)
return x3_samfeats, stage2_img, feat2, res2
##########################################################################
## DGUNet_plus
class DGUNet(nn.Module):
def __init__(self, in_c=3, out_c=3, n_feat=80, scale_unetfeats=48, scale_orsnetfeats=32, num_cab=8, kernel_size=3,
reduction=4, bias=False, depth=5):
super(DGUNet, self).__init__()
act = nn.PReLU()
self.depth=depth
self.basic=Basic_block(in_c, out_c, n_feat, scale_unetfeats, scale_orsnetfeats, num_cab, kernel_size, reduction, bias)
self.shallow_feat1 = nn.Sequential(conv(3, n_feat, kernel_size, bias=bias),
CAB(n_feat, kernel_size, reduction, bias=bias, act=act))
self.shallow_feat7 = nn.Sequential(conv(3, n_feat, kernel_size, bias=bias),
CAB(n_feat, kernel_size, reduction, bias=bias, act=act))
# Cross Stage Feature Fusion (CSFF)
self.stage1_encoder = Encoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats,depth=4, csff=False)
self.stage1_decoder = Decoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats,depth=4)
self.sam12 = SAM(n_feat, kernel_size=1, bias=bias)
self.phi_0 = ResBlock(default_conv,3,3)
self.phit_0 = ResBlock(default_conv,3,3)
self.phi_6 = ResBlock(default_conv,3,3)
self.phit_6 = ResBlock(default_conv,3,3)
self.r0 = nn.Parameter(torch.Tensor([0.5]))
self.r6 = nn.Parameter(torch.Tensor([0.5]))
self.concat67 = conv(n_feat * 2, n_feat + scale_orsnetfeats, kernel_size, bias=bias)
self.tail = conv(n_feat + scale_orsnetfeats, 3, kernel_size, bias=bias)
def forward(self, img):
res=[]
##-------------------------------------------
##-------------- Stage 1---------------------
##-------------------------------------------
## GDM
phixsy_1 = self.phi_0(img) - img
x1_img = img - self.r0*self.phit_0(phixsy_1)
## PMM
x1 = self.shallow_feat1(x1_img)
feat1,feat_fin1 = self.stage1_encoder(x1)
res1 = self.stage1_decoder(feat_fin1,feat1)
x2_samfeats, stage1_img = self.sam12(res1[-1], x1_img)
res.append(stage1_img)
##-------------------------------------------
##-------------- Stage 2-6 ---------------------
##-------------------------------------------
## Compute Shallow Features
for _ in range(self.depth):
x2_samfeats, stage1_img, feat1, res1 = self.basic(img,stage1_img,feat1,res1,x2_samfeats)
res.append(stage1_img)
# ##-------------------------------------------
# ##-------------- Stage 7---------------------
# ##-------------------------------------------
# ## Compute Shallow Features
phixsy_7 = self.phi_6(stage1_img) - img
x7_img = stage1_img - self.r6*self.phit_6(phixsy_7)
x7 = self.shallow_feat7(x7_img)
# ## Concatenate SAM features of Stage 2 with shallow features of Stage 3
x7_cat = self.concat67(torch.cat([x7, x2_samfeats], 1))
stage7_img = self.tail(x7_cat)+ img
res.append(stage7_img)
return res[::-1]
================================================
FILE: Deblurring/.ipynb_checkpoints/DGUNet_plus-checkpoint.py
================================================
import torch
import torch.nn as nn
import torch.nn.functional as F
from pdb import set_trace as stx
import numpy as np
import cv2
##########################################################################
# Basic modules
def conv(in_channels, out_channels, kernel_size, bias=False, stride=1):
return nn.Conv2d(
in_channels, out_channels, kernel_size,
padding=(kernel_size // 2), bias=bias, stride=stride)
def conv_down(in_chn, out_chn, bias=False):
layer = nn.Conv2d(in_chn, out_chn, kernel_size=4, stride=2, padding=1, bias=bias)
return layer
def default_conv(in_channels, out_channels, kernel_size,stride=1, bias=True):
return nn.Conv2d(
in_channels, out_channels, kernel_size,
padding=(kernel_size//2),stride=stride, bias=bias)
class ResBlock(nn.Module):
def __init__(
self, conv, n_feats, kernel_size,
bias=True, bn=False, act=nn.PReLU(), res_scale=1):
super(ResBlock, self).__init__()
m = []
for i in range(2):
if i == 0:
m.append(conv(n_feats, 64, kernel_size, bias=bias))
else:
m.append(conv(64, n_feats, kernel_size, bias=bias))
if bn:
m.append(nn.BatchNorm2d(n_feats))
if i == 0:
m.append(act)
self.body = nn.Sequential(*m)
self.res_scale = res_scale
def forward(self, x):
res = self.body(x).mul(self.res_scale)
res += x
return res
##########################################################################
## Channel Attention Block (CAB)
class CAB(nn.Module):
def __init__(self, n_feat, kernel_size, reduction, bias, act):
super(CAB, self).__init__()
modules_body = []
modules_body.append(conv(n_feat, n_feat, kernel_size, bias=bias))
modules_body.append(act)
modules_body.append(conv(n_feat, n_feat, kernel_size, bias=bias))
self.CA = CALayer(n_feat, reduction, bias=bias)
self.body = nn.Sequential(*modules_body)
def forward(self, x):
res = self.body(x)
res = self.CA(res)
res += x
return res
class CALayer(nn.Module):
def __init__(self, channel, reduction=16, bias=False):
super(CALayer, self).__init__()
# global average pooling: feature --> point
self.avg_pool = nn.AdaptiveAvgPool2d(1)
# feature channel downscale and upscale --> channel weight
self.conv_du = nn.Sequential(
nn.Conv2d(channel, channel // reduction, 1, padding=0, bias=bias),
nn.ReLU(inplace=True),
nn.Conv2d(channel // reduction, channel, 1, padding=0, bias=bias),
nn.Sigmoid()
)
def forward(self, x):
y = self.avg_pool(x)
y = self.conv_du(y)
return x * y
##########################################################################
## Compute inter-stage features
class SAM(nn.Module):
def __init__(self, n_feat, kernel_size, bias):
super(SAM, self).__init__()
self.conv1 = conv(n_feat, n_feat, kernel_size, bias=bias)
self.conv2 = conv(n_feat, 3, kernel_size, bias=bias)
def forward(self, x, x_img):
x1 = self.conv1(x)
img = self.conv2(x) + x_img
x1 = x1 + x
return x1, img
class mergeblock(nn.Module):
def __init__(self, n_feat, kernel_size, bias, subspace_dim=16):
super(mergeblock, self).__init__()
self.conv_block = conv(n_feat * 2, n_feat, kernel_size, bias=bias)
self.num_subspace = subspace_dim
self.subnet = conv(n_feat * 2, self.num_subspace, kernel_size, bias=bias)
def forward(self, x, bridge):
out = torch.cat([x, bridge], 1)
b_, c_, h_, w_ = bridge.shape
sub = self.subnet(out)
V_t = sub.view(b_, self.num_subspace, h_*w_)
V_t = V_t / (1e-6 + torch.abs(V_t).sum(axis=2, keepdims=True))
V = V_t.permute(0, 2, 1)
mat = torch.matmul(V_t, V)
mat_inv = torch.inverse(mat)
project_mat = torch.matmul(mat_inv, V_t)
bridge_ = bridge.view(b_, c_, h_*w_)
project_feature = torch.matmul(project_mat, bridge_.permute(0, 2, 1))
bridge = torch.matmul(V, project_feature).permute(0, 2, 1).view(b_, c_, h_, w_)
out = torch.cat([x, bridge], 1)
out = self.conv_block(out)
return out+x
##########################################################################
## U-Net
class Encoder(nn.Module):
def __init__(self, n_feat, kernel_size, reduction, act, bias, scale_unetfeats, csff,depth=5):
super(Encoder, self).__init__()
self.body=nn.ModuleList()#[]
self.depth=depth
for i in range(depth-1):
self.body.append(UNetConvBlock(in_size=n_feat+scale_unetfeats*i, out_size=n_feat+scale_unetfeats*(i+1), downsample=True, relu_slope=0.2, use_csff=csff, use_HIN=True))
self.body.append(UNetConvBlock(in_size=n_feat+scale_unetfeats*(depth-1), out_size=n_feat+scale_unetfeats*(depth-1), downsample=False, relu_slope=0.2, use_csff=csff, use_HIN=True))
def forward(self, x, encoder_outs=None, decoder_outs=None):
res=[]
if encoder_outs is not None and decoder_outs is not None:
for i,down in enumerate(self.body):
if (i+1) < self.depth:
x, x_up = down(x,encoder_outs[i],decoder_outs[-i-1])
res.append(x_up)
else:
x = down(x)
else:
for i,down in enumerate(self.body):
if (i+1) < self.depth:
x, x_up = down(x)
res.append(x_up)
else:
x = down(x)
return res,x
class UNetConvBlock(nn.Module):
def __init__(self, in_size, out_size, downsample, relu_slope, use_csff=False, use_HIN=False):
super(UNetConvBlock, self).__init__()
self.downsample = downsample
self.identity = nn.Conv2d(in_size, out_size, 1, 1, 0)
self.use_csff = use_csff
self.conv_1 = nn.Conv2d(in_size, out_size, kernel_size=3, padding=1, bias=True)
self.relu_1 = nn.LeakyReLU(relu_slope, inplace=False)
self.conv_2 = nn.Conv2d(out_size, out_size, kernel_size=3, padding=1, bias=True)
self.relu_2 = nn.LeakyReLU(relu_slope, inplace=False)
if downsample and use_csff:
self.csff_enc = nn.Conv2d(out_size, out_size, 3, 1, 1)
self.csff_dec = nn.Conv2d(in_size, out_size, 3, 1, 1)
self.phi = nn.Conv2d(out_size, out_size, 3, 1, 1)
self.gamma = nn.Conv2d(out_size, out_size, 3, 1, 1)
if use_HIN:
self.norm = nn.InstanceNorm2d(out_size//2, affine=True)
self.use_HIN = use_HIN
if downsample:
self.downsample = conv_down(out_size, out_size, bias=False)
def forward(self, x, enc=None, dec=None):
out = self.conv_1(x)
if self.use_HIN:
out_1, out_2 = torch.chunk(out, 2, dim=1)
out = torch.cat([self.norm(out_1), out_2], dim=1)
out = self.relu_1(out)
out = self.relu_2(self.conv_2(out))
out += self.identity(x)
if enc is not None and dec is not None:
assert self.use_csff
skip_ = F.leaky_relu(self.csff_enc(enc) + self.csff_dec(dec), 0.1, inplace=True)
out = out*F.sigmoid(self.phi(skip_)) + self.gamma(skip_) + out
if self.downsample:
out_down = self.downsample(out)
return out_down, out
else:
return out
class UNetUpBlock(nn.Module):
def __init__(self, in_size, out_size, relu_slope):
super(UNetUpBlock, self).__init__()
self.up = nn.ConvTranspose2d(in_size, out_size, kernel_size=2, stride=2, bias=True)
self.conv_block = UNetConvBlock(out_size*2, out_size, False, relu_slope)
def forward(self, x, bridge):
up = self.up(x)
out = torch.cat([up, bridge], 1)
out = self.conv_block(out)
return out
class Decoder(nn.Module):
def __init__(self, n_feat, kernel_size, reduction, act, bias, scale_unetfeats,depth=5):
super(Decoder, self).__init__()
self.body=nn.ModuleList()
self.skip_conv=nn.ModuleList()#[]
for i in range(depth-1):
self.body.append(UNetUpBlock(in_size=n_feat+scale_unetfeats*(depth-i-1), out_size=n_feat+scale_unetfeats*(depth-i-2), relu_slope=0.2))
self.skip_conv.append(nn.Conv2d(n_feat+scale_unetfeats*(depth-i-1), n_feat+scale_unetfeats*(depth-i-2), 3, 1, 1))
def forward(self, x, bridges):
res=[]
for i,up in enumerate(self.body):
x=up(x,self.skip_conv[i](bridges[-i-1]))
res.append(x)
return res
##########################################################################
##---------- Resizing Modules ----------
class DownSample(nn.Module):
def __init__(self, in_channels, s_factor):
super(DownSample, self).__init__()
self.down = nn.Sequential(nn.Upsample(scale_factor=0.5, mode='bilinear', align_corners=False),
nn.Conv2d(in_channels, in_channels + s_factor, 1, stride=1, padding=0, bias=False))
def forward(self, x):
x = self.down(x)
return x
class UpSample(nn.Module):
def __init__(self, in_channels, out_channels):
super(UpSample, self).__init__()
self.up = nn.Sequential(nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
nn.Conv2d(in_channels, out_channels, 1, stride=1, padding=0, bias=False))
def forward(self, x):
x = self.up(x)
return x
##########################################################################
## DGUNet
class DGUNet(nn.Module):
def __init__(self, in_c=3, out_c=3, n_feat=96, scale_unetfeats=48, scale_orsnetfeats=32, num_cab=8, kernel_size=3, reduction=4, bias=False):
super(DGUNet, self).__init__()
# Extract Shallow Features
act = nn.PReLU()
self.shallow_feat1 = nn.Sequential(conv(3, n_feat, kernel_size, bias=bias),
CAB(n_feat, kernel_size, reduction, bias=bias, act=act))
self.shallow_feat2 = nn.Sequential(conv(3, n_feat, kernel_size, bias=bias),
CAB(n_feat, kernel_size, reduction, bias=bias, act=act))
self.shallow_feat3 = nn.Sequential(conv(3, n_feat, kernel_size, bias=bias),
CAB(n_feat, kernel_size, reduction, bias=bias, act=act))
self.shallow_feat4 = nn.Sequential(conv(3, n_feat, kernel_size, bias=bias),
CAB(n_feat, kernel_size, reduction, bias=bias, act=act))
self.shallow_feat5 = nn.Sequential(conv(3, n_feat, kernel_size, bias=bias),
CAB(n_feat, kernel_size, reduction, bias=bias, act=act))
self.shallow_feat6 = nn.Sequential(conv(3, n_feat, kernel_size, bias=bias),
CAB(n_feat, kernel_size, reduction, bias=bias, act=act))
self.shallow_feat7 = nn.Sequential(conv(3, n_feat, kernel_size, bias=bias),
CAB(n_feat, kernel_size, reduction, bias=bias, act=act))
# Gradient Descent Module (GDM)
self.phi_0 = ResBlock(default_conv,3,3)
self.phit_0 = ResBlock(default_conv,3,3)
self.phi_1 = ResBlock(default_conv,3,3)
self.phit_1 = ResBlock(default_conv,3,3)
self.phi_2 = ResBlock(default_conv,3,3)
self.phit_2 = ResBlock(default_conv,3,3)
self.phi_3 = ResBlock(default_conv,3,3)
self.phit_3 = ResBlock(default_conv,3,3)
self.phi_4 = ResBlock(default_conv,3,3)
self.phit_4 = ResBlock(default_conv,3,3)
self.phi_5 = ResBlock(default_conv,3,3)
self.phit_5 = ResBlock(default_conv,3,3)
self.phi_6 = ResBlock(default_conv,3,3)
self.phit_6 = ResBlock(default_conv,3,3)
self.r0 = nn.Parameter(torch.Tensor([0.5]))
self.r1 = nn.Parameter(torch.Tensor([0.5]))
self.r2 = nn.Parameter(torch.Tensor([0.5]))
self.r3 = nn.Parameter(torch.Tensor([0.5]))
self.r4 = nn.Parameter(torch.Tensor([0.5]))
self.r5 = nn.Parameter(torch.Tensor([0.5]))
self.r6 = nn.Parameter(torch.Tensor([0.5]))
# Informative Proximal Mapping Module (IPMM)
self.stage1_encoder = Encoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats,depth=4, csff=False)
self.stage1_decoder = Decoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats,depth=4)
self.stage2_encoder = Encoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats,depth=4, csff=True)
self.stage2_decoder = Decoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats,depth=4)
self.stage3_encoder = Encoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats,depth=4, csff=True)
self.stage3_decoder = Decoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats,depth=4)
self.stage4_encoder = Encoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats,depth=4, csff=True)
self.stage4_decoder = Decoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats,depth=4)
self.stage5_encoder = Encoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats,depth=4, csff=True)
self.stage5_decoder = Decoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats,depth=4)
self.stage6_encoder = Encoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats,depth=4, csff=True)
self.stage6_decoder = Decoder(n_feat, kernel_size, reduction, act, bias, scale_unetfeats,depth=4)
self.sam12 = SAM(n_feat, kernel_size=1, bias=bias)
self.merge12=mergeblock(n_feat,3,True)
self.sam23 = SAM(n_feat, kernel_size=1, bias=bias)
self.merge23=mergeblock(n_feat,3,True)
self.sam34 = SAM(n_feat, kernel_size=1, bias=bias)
self.merge34=mergeblock(n_feat,3,True)
self.sam45 = SAM(n_feat, kernel_size=1, bias=bias)
self.merge45=mergeblock(n_feat,3,True)
self.sam56 = SAM(n_feat, kernel_size=1, bias=bias)
self.merge56=mergeblock(n_feat,3,True)
self.sam67 = SAM(n_feat, kernel_size=1, bias=bias)
self.merge67=mergeblock(n_feat,3,True)
self.tail = conv(n_feat, 3, kernel_size, bias=bias)
def forward(self, img):
##-------------------------------------------
##-------------- Stage 1---------------------
##-------------------------------------------
## GDM
phixsy_1 = self.phi_0(img) - img
x1_img = img - self.r0*self.phit_0(phixsy_1)
# PMM
x1 = self.shallow_feat1(x1_img)
feat1,feat_fin1 = self.stage1_encoder(x1)
res1 = self.stage1_decoder(feat_fin1,feat1)
x2_samfeats, stage1_img = self.sam12(res1[-1], x1_img)
##-------------------------------------------
##-------------- Stage 2---------------------
##-------------------------------------------
## GDM
phixsy_2 = self.phi_1(stage1_img) - img
x2_img = stage1_img - self.r1*self.phit_1(phixsy_2)
x2 = self.shallow_feat2(x2_img)
## PMM
x2_cat = self.merge12(x2, x2_samfeats)
feat2,feat_fin2 = self.stage2_encoder(x2_cat, feat1, res1)
res2 = self.stage2_decoder(feat_fin2,feat2)
x3_samfeats, stage2_img = self.sam23(res2[-1], x2_img)
##-------------------------------------------
##-------------- Stage 3---------------------
##-------------------------------------------
## GDM
phixsy_3 = self.phi_2(stage2_img) - img
x3_img = stage2_img - self.r2*self.phit_2(phixsy_3)
## PMM
x3 = self.shallow_feat3(x3_img)
x3_cat = self.merge23(x3, x3_samfeats)
feat3,feat_fin3 = self.stage3_encoder(x3_cat, feat2, res2)
res3 = self.stage3_decoder(feat_fin3,feat3)
x4_samfeats, stage3_img = self.sam34(res3[-1], x3_img)
##-------------------------------------------
##-------------- Stage 4---------------------
##-------------------------------------------
## GDM
phixsy_4 = self.phi_3(stage3_img) - img
x4_img = stage3_img - self.r3*self.phit_3(phixsy_4)
## PMM
x4 = self.shallow_feat4(x4_img)
x4_cat = self.merge34(x4, x4_samfeats)
feat4,feat_fin4 = self.stage4_encoder(x4_cat, feat3, res3)
res4 = self.stage4_decoder(feat_fin4,feat4)
x5_samfeats, stage4_img = self.sam45(res4[-1], x4_img)
##-------------------------------------------
##-------------- Stage 5---------------------
##-------------------------------------------
## GDM
phixsy_5 = self.phi_4(stage4_img) - img
x5_img = stage4_img - self.r4*self.phit_4(phixsy_5)
## PMM
x5 = self.shallow_feat5(x5_img)
x5_cat = self.merge45(x5, x5_samfeats)
feat5,feat_fin5 = self.stage5_encoder(x5_cat, feat4, res4)
res5 = self.stage5_decoder(feat_fin5,feat5)
x6_samfeats, stage5_img = self.sam56(res5[-1], x5_img)
##-------------------------------------------
##-------------- Stage 6---------------------
##-------------------------------------------
## GDM
phixsy_6 = self.phi_5(stage5_img) - img
x6_img = stage5_img - self.r5*self.phit_5(phixsy_6)
## PMM
x6 = self.shallow_feat6(x6_img)
x6_cat = self.merge56(x6, x6_samfeats)
feat6,feat_fin6 = self.stage6_encoder(x6_cat, feat5, res5)
res6 = self.stage6_decoder(feat_fin6,feat6)
x7_samfeats, stage6_img = self.sam67(res6[-1], x6_img)
##-------------------------------------------
##-------------- Stage 7---------------------
##-------------------------------------------
## GDM
phixsy_7 = self.phi_6(stage6_img) - img
x7_img = stage6_img - self.r6*self.phit_6(phixsy_7)
## PMM
x7 = self.shallow_feat7(x7_img)
x7_cat = self.merge67(x7, x7_samfeats)
stage7_img = self.tail(x7_cat)+ img
return [stage7_img,stage6_img,stage5_img,stage4_img, stage3_img, stage2_img, stage1_img]
================================================
FILE: Deblurring/.ipynb_checkpoints/MPRNet_all_4gpu-checkpoint.out
================================================
Let's use 2 GPUs!
===> Start Epoch 1 End Epoch 3001
===> Loading datasets
------------------------------------------------------------------
Epoch: 1 Time: 445.3842 Loss: 42.0905 LearningRate 0.000033
------------------------------------------------------------------
------------------------------------------------------------------
Epoch: 2 Time: 427.1102 Loss: 35.7552 LearningRate 0.000033
------------------------------------------------------------------
------------------------------------------------------------------
Epoch: 3 Time: 421.8640 Loss: 33.9910 LearningRate 0.000033
------------------------------------------------------------------
------------------------------------------------------------------
Epoch: 4 Time: 424.1221 Loss: 32.7013 LearningRate 0.000033
------------------------------------------------------------------
------------------------------------------------------------------
Epoch: 5 T
gitextract_g5u1c6i_/
├── Compressive-Sensing/
│ ├── DGUNet.py
│ ├── DGUNet_deblock.py
│ ├── DGUNet_plus.py
│ ├── DGUNet_plus_deblock.py
│ ├── Datasets/
│ │ └── README.md
│ ├── README.md
│ ├── config.py
│ ├── csdata_fast.py
│ ├── csdata_fast2.py
│ ├── losses.py
│ ├── sampling_matrix_new/
│ │ ├── README.md
│ │ ├── phi_sampling_10_32x32.npy
│ │ ├── phi_sampling_1_32x32.npy
│ │ ├── phi_sampling_25_32x32.npy
│ │ ├── phi_sampling_30_32x32.npy
│ │ ├── phi_sampling_40_32x32.npy
│ │ ├── phi_sampling_4_32x32.npy
│ │ └── phi_sampling_50_32x32.npy
│ ├── test.py
│ ├── test_plus.py
│ ├── train.py
│ ├── train_deblock.py
│ ├── train_plus.py
│ ├── train_plus_deblock.py
│ ├── training.yml
│ ├── training_cs.yml
│ └── utils/
│ ├── README.md
│ ├── __init__.py
│ ├── dataset_utils.py
│ ├── dir_utils.py
│ ├── image_utils.py
│ └── model_utils.py
├── Deblurring/
│ ├── .ipynb_checkpoints/
│ │ ├── DGUNet-checkpoint.py
│ │ ├── DGUNet_plus-checkpoint.py
│ │ ├── MPRNet_all_4gpu-checkpoint.out
│ │ ├── MPRNet_nochop_s5_v7_re-checkpoint.py
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_baseSAM-checkpoint.py
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft-checkpoint.py
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft_4gpu-checkpoint.out
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft_4gpu_2-checkpoint.out
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft_rp_4gpu_2-checkpoint.out
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft_v2-checkpoint.py
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft_v2_4gpu-checkpoint.out
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft_v2_all-checkpoint.out
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_rp_base-checkpoint.py
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_rp_base_sft_v2-checkpoint.py
│ │ ├── README-checkpoint.md
│ │ ├── config-checkpoint.py
│ │ ├── data_RGB-checkpoint.py
│ │ ├── dataset_RGB-checkpoint.py
│ │ ├── deblure_s7_v7_h3_u4_base-checkpoint.out
│ │ ├── deblure_s7_v7_h3_u4_rp_base-checkpoint.out
│ │ ├── deblure_s7_v7_h3_u4_rp_base_4gpu-checkpoint.out
│ │ ├── deblure_s7_v7_h3_u4_rp_base_4gpu_2e-4-checkpoint.out
│ │ ├── deblure_s7_v7_h3_u4_rp_base_4gpu_2e-4_b16-checkpoint.out
│ │ ├── losses-checkpoint.py
│ │ ├── test-checkpoint.py
│ │ ├── test_plus-checkpoint.py
│ │ ├── train-checkpoint.py
│ │ ├── train_plus-checkpoint.py
│ │ └── training-checkpoint.yml
│ ├── DGUNet.py
│ ├── DGUNet_plus.py
│ ├── Datasets/
│ │ └── README.md
│ ├── README.md
│ ├── config.py
│ ├── data_RGB.py
│ ├── dataset_RGB.py
│ ├── evaluate_GOPRO_HIDE.m
│ ├── evaluate_RealBlur.py
│ ├── losses.py
│ ├── test.py
│ ├── test_plus.py
│ ├── train.py
│ ├── train_plus.py
│ ├── training.yml
│ └── utils/
│ ├── __init__.py
│ ├── dataset_utils.py
│ ├── dir_utils.py
│ ├── image_utils.py
│ └── model_utils.py
├── Denoising/
│ ├── .ipynb_checkpoints/
│ │ ├── DGUNet_plus-checkpoint.py
│ │ ├── config-checkpoint.py
│ │ ├── test-checkpoint.py
│ │ ├── test_plus-checkpoint.py
│ │ ├── train-checkpoint.py
│ │ ├── train_plus-checkpoint.py
│ │ └── training-checkpoint.yml
│ ├── DGUNet.py
│ ├── DGUNet_plus.py
│ ├── Datasets/
│ │ └── README.md
│ ├── README.md
│ ├── config.py
│ ├── data_RGB.py
│ ├── dataset_RGB.py
│ ├── demo_DND.py
│ ├── generate_patches_SIDD.py
│ ├── losses.py
│ ├── test.py
│ ├── test_DND.py
│ ├── test_DND_plus.py
│ ├── test_plus.py
│ ├── train.py
│ ├── train_plus.py
│ ├── training.yml
│ └── utils/
│ ├── .ipynb_checkpoints/
│ │ ├── __init__-checkpoint.py
│ │ └── model_utils-checkpoint.py
│ ├── __init__.py
│ ├── dataset_utils.py
│ ├── dir_utils.py
│ ├── image_utils.py
│ └── model_utils.py
├── Deraining/
│ ├── .ipynb_checkpoints/
│ │ ├── DGUNet-checkpoint.py
│ │ ├── DGUNet_plus-checkpoint.py
│ │ ├── MPRNet_nochop_s3-checkpoint.py
│ │ ├── MPRNet_nochop_s3_v7_hi3_u4-checkpoint.py
│ │ ├── MPRNet_nochop_s5-checkpoint.py
│ │ ├── MPRNet_nochop_s5_dense-checkpoint.py
│ │ ├── MPRNet_nochop_s5_dense_v3-checkpoint.py
│ │ ├── MPRNet_nochop_s5_rb-checkpoint.py
│ │ ├── MPRNet_nochop_s5_sk-checkpoint.py
│ │ ├── MPRNet_nochop_s5_sk_dense-checkpoint.py
│ │ ├── MPRNet_nochop_s5_v3-checkpoint.py
│ │ ├── MPRNet_nochop_s5_v7-checkpoint.py
│ │ ├── MPRNet_nochop_s5_v7_hi3_u4-checkpoint.py
│ │ ├── MPRNet_nochop_s5_v7_re-checkpoint.py
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4-checkpoint.py
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_baseSAM-checkpoint.py
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft-checkpoint.py
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft_v2-checkpoint.py
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft_v2_sf-checkpoint.py
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_nointer-checkpoint.py
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_nophi-checkpoint.py
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_rp_base-checkpoint.py
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_rp_base_sft_v2-checkpoint.py
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_rp_base_v2-checkpoint.py
│ │ ├── MPRNet_nochop_s7_v7_hi3_u4_rp_base_v3-checkpoint.py
│ │ ├── MPRNet_nochop_s9_v7_hi3_u4_baseSAM_sft_v2-checkpoint.py
│ │ ├── MPRNet_v2-checkpoint.py
│ │ ├── MPRNet_v3-checkpoint.py
│ │ ├── MPRNet_v4-checkpoint.py
│ │ ├── MPRNet_v5-checkpoint.py
│ │ ├── MPRNet_v6-checkpoint.py
│ │ ├── data_RGB-checkpoint.py
│ │ ├── dataset_RGB-checkpoint.py
│ │ ├── re_ceckpoint-checkpoint.py
│ │ ├── test-checkpoint.py
│ │ ├── test_plus-checkpoint.py
│ │ ├── train-checkpoint.py
│ │ ├── train2-checkpoint.py
│ │ ├── train_MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft-checkpoint.out
│ │ ├── train_MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft_v2-checkpoint.out
│ │ ├── train_MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft_v2_scale1-checkpoint.out
│ │ ├── train_MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft_v2_scale2-checkpoint.out
│ │ ├── train_MPRNet_nochop_s7_v7_hi3_u4_rp_baseSAM_sft_v2-checkpoint.out
│ │ ├── train_MPRNet_nochop_s9_v7_hi3_u4_baseSAM_sft_v2-checkpoint.out
│ │ ├── train_plus-checkpoint.py
│ │ └── training-checkpoint.yml
│ ├── DGUNet.py
│ ├── DGUNet_plus.py
│ ├── Datasets/
│ │ ├── .ipynb_checkpoints/
│ │ │ └── README-checkpoint.md
│ │ └── README.md
│ ├── README.md
│ ├── SKmodel.py
│ ├── config.py
│ ├── data_RGB.py
│ ├── dataset_RGB.py
│ ├── evaluate_PSNR_SSIM.m
│ ├── losses.py
│ ├── test.py
│ ├── test_plus.py
│ ├── train.py
│ ├── train_plus.py
│ ├── training.yml
│ ├── utils/
│ │ ├── __init__.py
│ │ ├── dataset_utils.py
│ │ ├── dir_utils.py
│ │ ├── image_utils.py
│ │ └── model_utils.py
│ └── warmup_scheduler.egg-info/
│ ├── PKG-INFO
│ ├── SOURCES.txt
│ ├── dependency_links.txt
│ └── top_level.txt
├── README.md
├── figs/
│ └── README.md
└── pytorch-gradual-warmup-lr/
├── build/
│ └── lib/
│ └── warmup_scheduler/
│ ├── __init__.py
│ ├── run.py
│ └── scheduler.py
├── dist/
│ └── warmup_scheduler-0.3-py3.6.egg
├── setup.py
├── warmup_scheduler/
│ ├── __init__.py
│ ├── run.py
│ └── scheduler.py
└── warmup_scheduler.egg-info/
├── PKG-INFO
├── SOURCES.txt
├── dependency_links.txt
└── top_level.txt
Condensed preview — 188 files, each showing path, character count, and a content snippet. Download the .json file or copy for the full structured content (1,654K chars).
[
{
"path": "Compressive-Sensing/DGUNet.py",
"chars": 17727,
"preview": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom pdb import set_trace as stx\n# from pdb import se..."
},
{
"path": "Compressive-Sensing/DGUNet_deblock.py",
"chars": 18849,
"preview": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom pdb import set_trace as stx\nfrom torch.nn import..."
},
{
"path": "Compressive-Sensing/DGUNet_plus.py",
"chars": 17712,
"preview": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom pdb import set_trace as stx\n# from pdb import se..."
},
{
"path": "Compressive-Sensing/DGUNet_plus_deblock.py",
"chars": 19712,
"preview": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom pdb import set_trace as stx\n# from pdb import se..."
},
{
"path": "Compressive-Sensing/Datasets/README.md",
"chars": 450,
"preview": "Download datasets from the provided links and place them in this directory. Your directory structure should look somethi..."
},
{
"path": "Compressive-Sensing/README.md",
"chars": 655,
"preview": "## Training\n- Download the [Datasets](Datasets/README.md)\n\n- Train the model with default arguments by running\n\n### For..."
},
{
"path": "Compressive-Sensing/config.py",
"chars": 3336,
"preview": "#!/usr/bin/env python3\n# -*- coding: utf-8 -*-\n\"\"\"\nCreated on Tue Jul 23 14:35:48 2019\n\n@author: aditya\n\"\"\"\n\nr\"\"\"This mo..."
},
{
"path": "Compressive-Sensing/csdata_fast.py",
"chars": 5504,
"preview": "from torch.utils.data import Dataset\nimport imageio\nimport os\nimport torch\nimport glob\nimport random\nimport numpy as np..."
},
{
"path": "Compressive-Sensing/csdata_fast2.py",
"chars": 3009,
"preview": "from torch.utils.data import Dataset\nimport imageio\nimport os\nimport torch\nimport glob\nimport random\nimport numpy as np..."
},
{
"path": "Compressive-Sensing/losses.py",
"chars": 1535,
"preview": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nclass CharbonnierLoss(nn.Module):\n \"\"\"Charbonnier..."
},
{
"path": "Compressive-Sensing/sampling_matrix_new/README.md",
"chars": 1,
"preview": "\n"
},
{
"path": "Compressive-Sensing/test.py",
"chars": 9931,
"preview": "import torch\nimport torch.nn as nn\nfrom torch.nn import init\nimport torch.nn.functional as F\nimport scipy.io as sio\nimpo..."
},
{
"path": "Compressive-Sensing/test_plus.py",
"chars": 9909,
"preview": "import torch\nimport torch.nn as nn\nfrom torch.nn import init\nimport torch.nn.functional as F\nimport scipy.io as sio\nimpo..."
},
{
"path": "Compressive-Sensing/train.py",
"chars": 9870,
"preview": "import torch\nimport torch.nn as nn\nfrom torch.nn import init\nimport torch.nn.functional as F\nimport scipy.io as sio\nimpo..."
},
{
"path": "Compressive-Sensing/train_deblock.py",
"chars": 10397,
"preview": "import torch\nimport torch.nn as nn\nfrom torch.nn import init\nimport torch.nn.functional as F\nimport scipy.io as sio\nimpo..."
},
{
"path": "Compressive-Sensing/train_plus.py",
"chars": 9872,
"preview": "import torch\nimport torch.nn as nn\nfrom torch.nn import init\nimport torch.nn.functional as F\nimport scipy.io as sio\nimpo..."
},
{
"path": "Compressive-Sensing/train_plus_deblock.py",
"chars": 10396,
"preview": "import torch\nimport torch.nn as nn\nfrom torch.nn import init\nimport torch.nn.functional as F\nimport scipy.io as sio\nimpo..."
},
{
"path": "Compressive-Sensing/training.yml",
"chars": 1108,
"preview": "# ###############\n# ## \n# ####\n\n# GPU: [0,1,2,3]\n\n# VERBOSE: True\n\n# MODEL:\n# MODE: 'Denoising'\n# SESSION: 'MPRNet'..."
},
{
"path": "Compressive-Sensing/training_cs.yml",
"chars": 1103,
"preview": "# ###############\n# ## \n# ####\n\n# GPU: [0,1,2,3]\n\n# VERBOSE: True\n\n# MODEL:\n# MODE: 'Denoising'\n# SESSION: 'MPRNet'..."
},
{
"path": "Compressive-Sensing/utils/README.md",
"chars": 1,
"preview": "\n"
},
{
"path": "Compressive-Sensing/utils/__init__.py",
"chars": 108,
"preview": "from .dir_utils import *\nfrom .image_utils import *\nfrom .model_utils import *\nfrom .dataset_utils import *\n"
},
{
"path": "Compressive-Sensing/utils/dataset_utils.py",
"chars": 534,
"preview": "import torch\n\nclass MixUp_AUG:\n def __init__(self):\n self.dist = torch.distributions.beta.Beta(torch.tensor([0..."
},
{
"path": "Compressive-Sensing/utils/dir_utils.py",
"chars": 407,
"preview": "import os\nfrom natsort import natsorted\nfrom glob import glob\n\ndef mkdirs(paths):\n if isinstance(paths, list) and not..."
},
{
"path": "Compressive-Sensing/utils/image_utils.py",
"chars": 487,
"preview": "import torch\nimport numpy as np\nimport cv2\n\ndef torchPSNR(tar_img, prd_img):\n imdff = torch.clamp(prd_img,0,1) - torc..."
},
{
"path": "Compressive-Sensing/utils/model_utils.py",
"chars": 1562,
"preview": "import torch\nimport os\nfrom collections import OrderedDict\n\ndef freeze(model):\n for p in model.parameters():..."
},
{
"path": "Deblurring/.ipynb_checkpoints/DGUNet-checkpoint.py",
"chars": 14542,
"preview": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom pdb import set_trace as stx\n\n\n##################..."
},
{
"path": "Deblurring/.ipynb_checkpoints/DGUNet_plus-checkpoint.py",
"chars": 18517,
"preview": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom pdb import set_trace as stx\nimport numpy as np\ni..."
},
{
"path": "Deblurring/.ipynb_checkpoints/MPRNet_all_4gpu-checkpoint.out",
"chars": 17085,
"preview": "\n\nLet's use 2 GPUs!\n\n\n===> Start Epoch 1 End Epoch 3001\n===> Loading datasets\n------------------------------------------..."
},
{
"path": "Deblurring/.ipynb_checkpoints/MPRNet_nochop_s5_v7_re-checkpoint.py",
"chars": 19517,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deblurring/.ipynb_checkpoints/MPRNet_nochop_s7_v7_hi3_u4_baseSAM-checkpoint.py",
"chars": 25126,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deblurring/.ipynb_checkpoints/MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft-checkpoint.py",
"chars": 25359,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deblurring/.ipynb_checkpoints/MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft_4gpu-checkpoint.out",
"chars": 721,
"preview": "73,473,587 total parameters.\n73,473,587 training parameters.\n\n\nLet's use 4 GPUs!\n\n\n===> Start Epoch 1 End Epoch 3001\n===..."
},
{
"path": "Deblurring/.ipynb_checkpoints/MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft_4gpu_2-checkpoint.out",
"chars": 5712,
"preview": "73,473,587 total parameters.\n73,473,587 training parameters.\n\n\nLet's use 4 GPUs!\n\n\n===> Start Epoch 1 End Epoch 3001\n===..."
},
{
"path": "Deblurring/.ipynb_checkpoints/MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft_rp_4gpu_2-checkpoint.out",
"chars": 139,
"preview": "20,616,824 total parameters.\n20,616,824 training parameters.\n\n\nLet's use 4 GPUs!\n\n\n===> Start Epoch 1 End Epoch 3001\n===..."
},
{
"path": "Deblurring/.ipynb_checkpoints/MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft_v2-checkpoint.py",
"chars": 25356,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deblurring/.ipynb_checkpoints/MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft_v2_4gpu-checkpoint.out",
"chars": 30554,
"preview": "73,473,587 total parameters.\n73,473,587 training parameters.\n\n\nLet's use 4 GPUs!\n\n\n===> Start Epoch 1 End Epoch 3001\n===..."
},
{
"path": "Deblurring/.ipynb_checkpoints/MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft_v2_all-checkpoint.out",
"chars": 76747,
"preview": "\n\nLet's use 2 GPUs!\n\n\n===> Start Epoch 1 End Epoch 3001\n===> Loading datasets\n------------------------------------------..."
},
{
"path": "Deblurring/.ipynb_checkpoints/MPRNet_nochop_s7_v7_hi3_u4_rp_base-checkpoint.py",
"chars": 18647,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deblurring/.ipynb_checkpoints/MPRNet_nochop_s7_v7_hi3_u4_rp_base_sft_v2-checkpoint.py",
"chars": 18947,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deblurring/.ipynb_checkpoints/README-checkpoint.md",
"chars": 1507,
"preview": "## Training\n- Download the [Datasets](Datasets/README.md)\n\n- Train the model with default arguments by running\n\n```\npyth..."
},
{
"path": "Deblurring/.ipynb_checkpoints/config-checkpoint.py",
"chars": 3951,
"preview": "#!/usr/bin/env python3\n# -*- coding: utf-8 -*-\n\"\"\"\nCreated on Tue Jul 23 14:35:48 2019\n\n@author: aditya\n\"\"\"\n\nr\"\"\"This mo..."
},
{
"path": "Deblurring/.ipynb_checkpoints/data_RGB-checkpoint.py",
"chars": 649,
"preview": "import os\nfrom dataset_RGB import DataLoaderTrain, DataLoaderVal, DataLoaderTest, DataLoaderTrain_all\n\ndef get_training_..."
},
{
"path": "Deblurring/.ipynb_checkpoints/dataset_RGB-checkpoint.py",
"chars": 11652,
"preview": "import os\nimport numpy as np\nfrom torch.utils.data import Dataset\nimport torch\nfrom PIL import Image\nimport torchvision...."
},
{
"path": "Deblurring/.ipynb_checkpoints/deblure_s7_v7_h3_u4_rp_base-checkpoint.out",
"chars": 117,
"preview": "20,657,800 total parameters.\n20,657,800 training parameters.\n===> Start Epoch 1 End Epoch 3001\n===> Loading datasets\n"
},
{
"path": "Deblurring/.ipynb_checkpoints/deblure_s7_v7_h3_u4_rp_base_4gpu-checkpoint.out",
"chars": 334,
"preview": "20,657,800 total parameters.\n20,657,800 training parameters.\n\n\nLet's use 4 GPUs!\n\n\n===> Start Epoch 1 End Epoch 3001\n===..."
},
{
"path": "Deblurring/.ipynb_checkpoints/deblure_s7_v7_h3_u4_rp_base_4gpu_2e-4-checkpoint.out",
"chars": 3116,
"preview": "20,657,800 total parameters.\n20,657,800 training parameters.\n\n\nLet's use 4 GPUs!\n\n\n===> Start Epoch 1 End Epoch 3001\n===..."
},
{
"path": "Deblurring/.ipynb_checkpoints/deblure_s7_v7_h3_u4_rp_base_4gpu_2e-4_b16-checkpoint.out",
"chars": 78109,
"preview": "20,657,800 total parameters.\n20,657,800 training parameters.\n\n\nLet's use 4 GPUs!\n\n\n===> Start Epoch 1 End Epoch 3001\n===..."
},
{
"path": "Deblurring/.ipynb_checkpoints/losses-checkpoint.py",
"chars": 1535,
"preview": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nclass CharbonnierLoss(nn.Module):\n \"\"\"Charbonnier..."
},
{
"path": "Deblurring/.ipynb_checkpoints/test-checkpoint.py",
"chars": 2818,
"preview": "import numpy as np\nimport os\nimport argparse\nfrom tqdm import tqdm\n\nimport torch.nn as nn\nimport torch\nfrom torch.utils...."
},
{
"path": "Deblurring/.ipynb_checkpoints/test_plus-checkpoint.py",
"chars": 2881,
"preview": "import numpy as np\nimport os\nimport argparse\nfrom tqdm import tqdm\n\nimport torch.nn as nn\nimport torch\nfrom torch.utils...."
},
{
"path": "Deblurring/.ipynb_checkpoints/train-checkpoint.py",
"chars": 5907,
"preview": "import os\nfrom config import Config \nopt = Config('training.yml')\n\ngpus = ','.join([str(i) for i in opt.GPU])\nos.environ..."
},
{
"path": "Deblurring/.ipynb_checkpoints/train_plus-checkpoint.py",
"chars": 5912,
"preview": "import os\nfrom config import Config \nopt = Config('training.yml')\n\ngpus = ','.join([str(i) for i in opt.GPU])\nos.environ..."
},
{
"path": "Deblurring/.ipynb_checkpoints/training-checkpoint.yml",
"chars": 709,
"preview": "###############\n## \n####\n\nGPU: [0,1,2,3]\n\nVERBOSE: True\n\nMODEL:\n MODE: 'Deblurring'\n SESSION: 'DGUNet'\n\n# Optimization..."
},
{
"path": "Deblurring/DGUNet.py",
"chars": 14366,
"preview": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom pdb import set_trace as stx\n\n\n##################..."
},
{
"path": "Deblurring/DGUNet_plus.py",
"chars": 18522,
"preview": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom pdb import set_trace as stx\nimport numpy as np\ni..."
},
{
"path": "Deblurring/Datasets/README.md",
"chars": 460,
"preview": "Download datasets from the google drive links and place them in this directory. Your directory tree should look like thi..."
},
{
"path": "Deblurring/README.md",
"chars": 985,
"preview": "## Training\n- Download the [Datasets](Datasets/README.md)\n\n- Train the model with default arguments by running\n\n```\npyth..."
},
{
"path": "Deblurring/config.py",
"chars": 3951,
"preview": "#!/usr/bin/env python3\n# -*- coding: utf-8 -*-\n\"\"\"\nCreated on Tue Jul 23 14:35:48 2019\n\n@author: aditya\n\"\"\"\n\nr\"\"\"This mo..."
},
{
"path": "Deblurring/data_RGB.py",
"chars": 649,
"preview": "import os\nfrom dataset_RGB import DataLoaderTrain, DataLoaderVal, DataLoaderTest, DataLoaderTrain_all\n\ndef get_training_..."
},
{
"path": "Deblurring/dataset_RGB.py",
"chars": 11652,
"preview": "import os\nimport numpy as np\nfrom torch.utils.data import Dataset\nimport torch\nfrom PIL import Image\nimport torchvision...."
},
{
"path": "Deblurring/evaluate_GOPRO_HIDE.m",
"chars": 1287,
"preview": "%% Multi-Stage Progressive Image Restoration\n%% Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Shahba..."
},
{
"path": "Deblurring/evaluate_RealBlur.py",
"chars": 3882,
"preview": "## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Shahba..."
},
{
"path": "Deblurring/losses.py",
"chars": 1535,
"preview": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nclass CharbonnierLoss(nn.Module):\n \"\"\"Charbonnier..."
},
{
"path": "Deblurring/test.py",
"chars": 2825,
"preview": "import numpy as np\nimport os\nimport argparse\nfrom tqdm import tqdm\n\nimport torch.nn as nn\nimport torch\nfrom torch.utils...."
},
{
"path": "Deblurring/test_plus.py",
"chars": 2888,
"preview": "import numpy as np\nimport os\nimport argparse\nfrom tqdm import tqdm\n\nimport torch.nn as nn\nimport torch\nfrom torch.utils...."
},
{
"path": "Deblurring/train.py",
"chars": 5907,
"preview": "import os\nfrom config import Config \nopt = Config('training.yml')\n\ngpus = ','.join([str(i) for i in opt.GPU])\nos.environ..."
},
{
"path": "Deblurring/train_plus.py",
"chars": 5912,
"preview": "import os\nfrom config import Config \nopt = Config('training.yml')\n\ngpus = ','.join([str(i) for i in opt.GPU])\nos.environ..."
},
{
"path": "Deblurring/training.yml",
"chars": 709,
"preview": "###############\n## \n####\n\nGPU: [0,1,2,3]\n\nVERBOSE: True\n\nMODEL:\n MODE: 'Deblurring'\n SESSION: 'DGUNet'\n\n# Optimization..."
},
{
"path": "Deblurring/utils/__init__.py",
"chars": 108,
"preview": "from .dir_utils import *\nfrom .image_utils import *\nfrom .model_utils import *\nfrom .dataset_utils import *\n"
},
{
"path": "Deblurring/utils/dataset_utils.py",
"chars": 534,
"preview": "import torch\n\nclass MixUp_AUG:\n def __init__(self):\n self.dist = torch.distributions.beta.Beta(torch.tensor([0..."
},
{
"path": "Deblurring/utils/dir_utils.py",
"chars": 407,
"preview": "import os\nfrom natsort import natsorted\nfrom glob import glob\n\ndef mkdirs(paths):\n if isinstance(paths, list) and not..."
},
{
"path": "Deblurring/utils/image_utils.py",
"chars": 487,
"preview": "import torch\nimport numpy as np\nimport cv2\n\ndef torchPSNR(tar_img, prd_img):\n imdff = torch.clamp(prd_img,0,1) - torc..."
},
{
"path": "Deblurring/utils/model_utils.py",
"chars": 1562,
"preview": "import torch\nimport os\nfrom collections import OrderedDict\n\ndef freeze(model):\n for p in model.parameters():..."
},
{
"path": "Denoising/.ipynb_checkpoints/DGUNet_plus-checkpoint.py",
"chars": 18539,
"preview": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom pdb import set_trace as stx\nimport numpy as np\ni..."
},
{
"path": "Denoising/.ipynb_checkpoints/config-checkpoint.py",
"chars": 3328,
"preview": "#!/usr/bin/env python3\n# -*- coding: utf-8 -*-\n\"\"\"\nCreated on Tue Jul 23 14:35:48 2019\n\n@author: aditya\n\"\"\"\n\nr\"\"\"This mo..."
},
{
"path": "Denoising/.ipynb_checkpoints/test-checkpoint.py",
"chars": 4344,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Denoising/.ipynb_checkpoints/test_plus-checkpoint.py",
"chars": 4354,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Denoising/.ipynb_checkpoints/train-checkpoint.py",
"chars": 6521,
"preview": "import os\nfrom config import Config \nopt = Config('training.yml')\n\ngpus = ','.join([str(i) for i in opt.GPU])\nos.environ..."
},
{
"path": "Denoising/.ipynb_checkpoints/train_plus-checkpoint.py",
"chars": 6526,
"preview": "import os\nfrom config import Config \nopt = Config('training.yml')\n\ngpus = ','.join([str(i) for i in opt.GPU])\nos.environ..."
},
{
"path": "Denoising/.ipynb_checkpoints/training-checkpoint.yml",
"chars": 574,
"preview": "###############\n## \n####\n\nGPU: [0,1,2,3]\n\nVERBOSE: True\n\nMODEL:\n MODE: 'Denoising'\n SESSION: 'DGUNet'\n\n# Optimization..."
},
{
"path": "Denoising/DGUNet.py",
"chars": 14366,
"preview": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom pdb import set_trace as stx\n\n\n##################..."
},
{
"path": "Denoising/DGUNet_plus.py",
"chars": 18675,
"preview": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom pdb import set_trace as stx\nimport numpy as np\ni..."
},
{
"path": "Denoising/Datasets/README.md",
"chars": 741,
"preview": "Download datasets from the provided links and place them in this directory. Your directory structure should look somethi..."
},
{
"path": "Denoising/README.md",
"chars": 1124,
"preview": "## Training\n- Download the [Datasets](Datasets/README.md)\n\n- Generate image patches from full-resolution training images..."
},
{
"path": "Denoising/config.py",
"chars": 3328,
"preview": "#!/usr/bin/env python3\n# -*- coding: utf-8 -*-\n\"\"\"\nCreated on Tue Jul 23 14:35:48 2019\n\n@author: aditya\n\"\"\"\n\nr\"\"\"This mo..."
},
{
"path": "Denoising/data_RGB.py",
"chars": 466,
"preview": "import os\nfrom dataset_RGB import DataLoaderTrain, DataLoaderVal, DataLoaderTest\n\ndef get_training_data(rgb_dir, img_opt..."
},
{
"path": "Denoising/dataset_RGB.py",
"chars": 5079,
"preview": "import os\nimport numpy as np\nfrom torch.utils.data import Dataset\nimport torch\nfrom PIL import Image\nimport torchvision...."
},
{
"path": "Denoising/demo_DND.py",
"chars": 3620,
"preview": "import numpy as np\nimport os\nimport argparse\nfrom tqdm import tqdm\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.f..."
},
{
"path": "Denoising/generate_patches_SIDD.py",
"chars": 2156,
"preview": "from glob import glob\nfrom tqdm import tqdm\nimport numpy as np\nimport os\nfrom natsort import natsorted\nimport cv2\nfrom j..."
},
{
"path": "Denoising/losses.py",
"chars": 1535,
"preview": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nclass CharbonnierLoss(nn.Module):\n \"\"\"Charbonnier..."
},
{
"path": "Denoising/test.py",
"chars": 4351,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Denoising/test_DND.py",
"chars": 4415,
"preview": "import numpy as np\nimport os\nimport argparse\nfrom tqdm import tqdm\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.f..."
},
{
"path": "Denoising/test_DND_plus.py",
"chars": 4430,
"preview": "import numpy as np\nimport os\nimport argparse\nfrom tqdm import tqdm\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.f..."
},
{
"path": "Denoising/test_plus.py",
"chars": 4361,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Denoising/train.py",
"chars": 6521,
"preview": "import os\nfrom config import Config \nopt = Config('training.yml')\n\ngpus = ','.join([str(i) for i in opt.GPU])\nos.environ..."
},
{
"path": "Denoising/train_plus.py",
"chars": 6526,
"preview": "import os\nfrom config import Config \nopt = Config('training.yml')\n\ngpus = ','.join([str(i) for i in opt.GPU])\nos.environ..."
},
{
"path": "Denoising/training.yml",
"chars": 575,
"preview": "###############\n## \n####\n\nGPU: [0,1,2,3]\n\nVERBOSE: True\n\nMODEL:\n MODE: 'Denoising'\n SESSION: 'DGUNet'\n\n# Optimization..."
},
{
"path": "Denoising/utils/.ipynb_checkpoints/__init__-checkpoint.py",
"chars": 108,
"preview": "from .dir_utils import *\nfrom .image_utils import *\nfrom .model_utils import *\nfrom .dataset_utils import *\n"
},
{
"path": "Denoising/utils/.ipynb_checkpoints/model_utils-checkpoint.py",
"chars": 1544,
"preview": "import torch\nimport os\nfrom collections import OrderedDict\n\ndef freeze(model):\n for p in model.parameters():..."
},
{
"path": "Denoising/utils/__init__.py",
"chars": 108,
"preview": "from .dir_utils import *\nfrom .image_utils import *\nfrom .model_utils import *\nfrom .dataset_utils import *\n"
},
{
"path": "Denoising/utils/dataset_utils.py",
"chars": 534,
"preview": "import torch\n\nclass MixUp_AUG:\n def __init__(self):\n self.dist = torch.distributions.beta.Beta(torch.tensor([0..."
},
{
"path": "Denoising/utils/dir_utils.py",
"chars": 407,
"preview": "import os\nfrom natsort import natsorted\nfrom glob import glob\n\ndef mkdirs(paths):\n if isinstance(paths, list) and not..."
},
{
"path": "Denoising/utils/image_utils.py",
"chars": 487,
"preview": "import torch\nimport numpy as np\nimport cv2\n\ndef torchPSNR(tar_img, prd_img):\n imdff = torch.clamp(prd_img,0,1) - torc..."
},
{
"path": "Denoising/utils/model_utils.py",
"chars": 1544,
"preview": "import torch\nimport os\nfrom collections import OrderedDict\n\ndef freeze(model):\n for p in model.parameters():..."
},
{
"path": "Deraining/.ipynb_checkpoints/DGUNet-checkpoint.py",
"chars": 14349,
"preview": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom pdb import set_trace as stx\n\n\n##################..."
},
{
"path": "Deraining/.ipynb_checkpoints/DGUNet_plus-checkpoint.py",
"chars": 18531,
"preview": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom pdb import set_trace as stx\nimport numpy as np\ni..."
},
{
"path": "Deraining/.ipynb_checkpoints/MPRNet_nochop_s3-checkpoint.py",
"chars": 14400,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deraining/.ipynb_checkpoints/MPRNet_nochop_s3_v7_hi3_u4-checkpoint.py",
"chars": 17005,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deraining/.ipynb_checkpoints/MPRNet_nochop_s5-checkpoint.py",
"chars": 17203,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deraining/.ipynb_checkpoints/MPRNet_nochop_s5_dense-checkpoint.py",
"chars": 17547,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deraining/.ipynb_checkpoints/MPRNet_nochop_s5_dense_v3-checkpoint.py",
"chars": 18006,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deraining/.ipynb_checkpoints/MPRNet_nochop_s5_rb-checkpoint.py",
"chars": 17669,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deraining/.ipynb_checkpoints/MPRNet_nochop_s5_sk-checkpoint.py",
"chars": 17359,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deraining/.ipynb_checkpoints/MPRNet_nochop_s5_sk_dense-checkpoint.py",
"chars": 17704,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deraining/.ipynb_checkpoints/MPRNet_nochop_s5_v3-checkpoint.py",
"chars": 17611,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deraining/.ipynb_checkpoints/MPRNet_nochop_s5_v7-checkpoint.py",
"chars": 19441,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deraining/.ipynb_checkpoints/MPRNet_nochop_s5_v7_hi3_u4-checkpoint.py",
"chars": 20100,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deraining/.ipynb_checkpoints/MPRNet_nochop_s5_v7_re-checkpoint.py",
"chars": 19686,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deraining/.ipynb_checkpoints/MPRNet_nochop_s7_v7_hi3_u4-checkpoint.py",
"chars": 23160,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deraining/.ipynb_checkpoints/MPRNet_nochop_s7_v7_hi3_u4_baseSAM-checkpoint.py",
"chars": 25126,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deraining/.ipynb_checkpoints/MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft-checkpoint.py",
"chars": 25350,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deraining/.ipynb_checkpoints/MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft_v2-checkpoint.py",
"chars": 25356,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deraining/.ipynb_checkpoints/MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft_v2_sf-checkpoint.py",
"chars": 25505,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deraining/.ipynb_checkpoints/MPRNet_nochop_s7_v7_hi3_u4_nointer-checkpoint.py",
"chars": 23358,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deraining/.ipynb_checkpoints/MPRNet_nochop_s7_v7_hi3_u4_nophi-checkpoint.py",
"chars": 23268,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deraining/.ipynb_checkpoints/MPRNet_nochop_s7_v7_hi3_u4_rp_base-checkpoint.py",
"chars": 18647,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deraining/.ipynb_checkpoints/MPRNet_nochop_s7_v7_hi3_u4_rp_base_sft_v2-checkpoint.py",
"chars": 18930,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deraining/.ipynb_checkpoints/MPRNet_nochop_s7_v7_hi3_u4_rp_base_v2-checkpoint.py",
"chars": 18647,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deraining/.ipynb_checkpoints/MPRNet_nochop_s7_v7_hi3_u4_rp_base_v3-checkpoint.py",
"chars": 18647,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deraining/.ipynb_checkpoints/MPRNet_nochop_s9_v7_hi3_u4_baseSAM_sft_v2-checkpoint.py",
"chars": 27955,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deraining/.ipynb_checkpoints/MPRNet_v2-checkpoint.py",
"chars": 15815,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deraining/.ipynb_checkpoints/MPRNet_v3-checkpoint.py",
"chars": 15882,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deraining/.ipynb_checkpoints/MPRNet_v4-checkpoint.py",
"chars": 15886,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deraining/.ipynb_checkpoints/MPRNet_v5-checkpoint.py",
"chars": 16374,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deraining/.ipynb_checkpoints/MPRNet_v6-checkpoint.py",
"chars": 17073,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deraining/.ipynb_checkpoints/data_RGB-checkpoint.py",
"chars": 621,
"preview": "import os\nfrom dataset_RGB import DataLoaderTrain, DataLoaderVal, DataLoaderTest, DataLoaderTest_fine\n\ndef get_training_..."
},
{
"path": "Deraining/.ipynb_checkpoints/dataset_RGB-checkpoint.py",
"chars": 6068,
"preview": "import os\nfrom torch.utils.data import Dataset\nimport torch\nfrom PIL import Image\nimport torchvision.transforms.function..."
},
{
"path": "Deraining/.ipynb_checkpoints/re_ceckpoint-checkpoint.py",
"chars": 3558,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deraining/.ipynb_checkpoints/test-checkpoint.py",
"chars": 4040,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deraining/.ipynb_checkpoints/test_plus-checkpoint.py",
"chars": 4055,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deraining/.ipynb_checkpoints/train-checkpoint.py",
"chars": 6149,
"preview": "import os\nfrom config import Config \nopt = Config('training.yml')\n\ngpus = ','.join([str(i) for i in opt.GPU])\nos.environ..."
},
{
"path": "Deraining/.ipynb_checkpoints/train2-checkpoint.py",
"chars": 6237,
"preview": "import os\nfrom config import Config \nopt = Config('training.yml')\n\ngpus = ','.join([str(i) for i in opt.GPU])\nos.environ..."
},
{
"path": "Deraining/.ipynb_checkpoints/train_MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft-checkpoint.out",
"chars": 21649,
"preview": "12,853,867 total parameters.\n12,853,867 training parameters.\n\n\nLet's use 2 GPUs!\n\n\n===> Start Epoch 1 End Epoch 251\n===>..."
},
{
"path": "Deraining/.ipynb_checkpoints/train_MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft_v2-checkpoint.out",
"chars": 29291,
"preview": "12,853,867 total parameters.\n12,853,867 training parameters.\n\n\nLet's use 2 GPUs!\n\n\n===> Start Epoch 1 End Epoch 251\n===>..."
},
{
"path": "Deraining/.ipynb_checkpoints/train_MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft_v2_scale1-checkpoint.out",
"chars": 4948,
"preview": "12,853,867 total parameters.\n12,853,867 training parameters.\n\n\nLet's use 2 GPUs!\n\n\n===> Start Epoch 1 End Epoch 251\n===>..."
},
{
"path": "Deraining/.ipynb_checkpoints/train_MPRNet_nochop_s7_v7_hi3_u4_baseSAM_sft_v2_scale2-checkpoint.out",
"chars": 21249,
"preview": "12,853,867 total parameters.\n12,853,867 training parameters.\n\n\nLet's use 2 GPUs!\n\n\n===> Start Epoch 1 End Epoch 251\n===>..."
},
{
"path": "Deraining/.ipynb_checkpoints/train_MPRNet_nochop_s7_v7_hi3_u4_rp_baseSAM_sft_v2-checkpoint.out",
"chars": 114,
"preview": "4,339,672 total parameters.\n4,339,672 training parameters.\n===> Start Epoch 1 End Epoch 251\n===> Loading datasets\n"
},
{
"path": "Deraining/.ipynb_checkpoints/train_MPRNet_nochop_s9_v7_hi3_u4_baseSAM_sft_v2-checkpoint.out",
"chars": 4948,
"preview": "17,125,553 total parameters.\n17,125,553 training parameters.\n\n\nLet's use 4 GPUs!\n\n\n===> Start Epoch 1 End Epoch 251\n===>..."
},
{
"path": "Deraining/.ipynb_checkpoints/train_plus-checkpoint.py",
"chars": 6154,
"preview": "import os\nfrom config import Config \nopt = Config('training.yml')\n\ngpus = ','.join([str(i) for i in opt.GPU])\nos.environ..."
},
{
"path": "Deraining/.ipynb_checkpoints/training-checkpoint.yml",
"chars": 542,
"preview": "###############\n## \n####\n\nGPU: [0,1,2,3]\n\nVERBOSE: True\n\nMODEL:\n MODE: 'Deraining'\n SESSION: 'DGUNet'\n# Optimization a..."
},
{
"path": "Deraining/DGUNet.py",
"chars": 14349,
"preview": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom pdb import set_trace as stx\n\n\n##################..."
},
{
"path": "Deraining/DGUNet_plus.py",
"chars": 18531,
"preview": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom pdb import set_trace as stx\nimport numpy as np\ni..."
},
{
"path": "Deraining/Datasets/.ipynb_checkpoints/README-checkpoint.md",
"chars": 517,
"preview": "Download datasets from the google drive links and place them in this directory. Your directory structure should look som..."
},
{
"path": "Deraining/Datasets/README.md",
"chars": 517,
"preview": "Download datasets from the google drive links and place them in this directory. Your directory structure should look som..."
},
{
"path": "Deraining/README.md",
"chars": 693,
"preview": "\n## Training\n- Download the [Datasets](Datasets/README.md)\n\n- Train the model with default arguments by running\n\n```\npyt..."
},
{
"path": "Deraining/SKmodel.py",
"chars": 5028,
"preview": "import torch\nfrom torch import nn\n\n\nclass SKConv(nn.Module):\n def __init__(self, features, M, r, L=32):\n \"\"\" C..."
},
{
"path": "Deraining/config.py",
"chars": 3336,
"preview": "#!/usr/bin/env python3\n# -*- coding: utf-8 -*-\n\"\"\"\nCreated on Tue Jul 23 14:35:48 2019\n\n@author: aditya\n\"\"\"\n\nr\"\"\"This mo..."
},
{
"path": "Deraining/data_RGB.py",
"chars": 621,
"preview": "import os\nfrom dataset_RGB import DataLoaderTrain, DataLoaderVal, DataLoaderTest, DataLoaderTest_fine\n\ndef get_training_..."
},
{
"path": "Deraining/dataset_RGB.py",
"chars": 6068,
"preview": "import os\nfrom torch.utils.data import Dataset\nimport torch\nfrom PIL import Image\nimport torchvision.transforms.function..."
},
{
"path": "Deraining/evaluate_PSNR_SSIM.m",
"chars": 7800,
"preview": "clc;close all;clear all;addpath(genpath('./'));\n\n% datasets = {'Rain100L'};\ndatasets = {'Rain100L', 'Rain100H', 'Test100..."
},
{
"path": "Deraining/losses.py",
"chars": 1535,
"preview": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nclass CharbonnierLoss(nn.Module):\n \"\"\"Charbonnier..."
},
{
"path": "Deraining/test.py",
"chars": 4047,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deraining/test_plus.py",
"chars": 4062,
"preview": "\"\"\"\n## Multi-Stage Progressive Image Restoration\n## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Sh..."
},
{
"path": "Deraining/train.py",
"chars": 6149,
"preview": "import os\nfrom config import Config \nopt = Config('training.yml')\n\ngpus = ','.join([str(i) for i in opt.GPU])\nos.environ..."
},
{
"path": "Deraining/train_plus.py",
"chars": 6154,
"preview": "import os\nfrom config import Config \nopt = Config('training.yml')\n\ngpus = ','.join([str(i) for i in opt.GPU])\nos.environ..."
},
{
"path": "Deraining/training.yml",
"chars": 542,
"preview": "###############\n## \n####\n\nGPU: [0,1,2,3]\n\nVERBOSE: True\n\nMODEL:\n MODE: 'Deraining'\n SESSION: 'DGUNet'\n# Optimization a..."
},
{
"path": "Deraining/utils/__init__.py",
"chars": 108,
"preview": "from .dir_utils import *\nfrom .image_utils import *\nfrom .model_utils import *\nfrom .dataset_utils import *\n"
},
{
"path": "Deraining/utils/dataset_utils.py",
"chars": 534,
"preview": "import torch\n\nclass MixUp_AUG:\n def __init__(self):\n self.dist = torch.distributions.beta.Beta(torch.tensor([0..."
},
{
"path": "Deraining/utils/dir_utils.py",
"chars": 407,
"preview": "import os\nfrom natsort import natsorted\nfrom glob import glob\n\ndef mkdirs(paths):\n if isinstance(paths, list) and not..."
},
{
"path": "Deraining/utils/image_utils.py",
"chars": 487,
"preview": "import torch\nimport numpy as np\nimport cv2\n\ndef torchPSNR(tar_img, prd_img):\n imdff = torch.clamp(prd_img,0,1) - torc..."
},
{
"path": "Deraining/utils/model_utils.py",
"chars": 1562,
"preview": "import torch\nimport os\nfrom collections import OrderedDict\n\ndef freeze(model):\n for p in model.parameters():..."
},
{
"path": "Deraining/warmup_scheduler.egg-info/PKG-INFO",
"chars": 273,
"preview": "Metadata-Version: 1.0\nName: warmup-scheduler\nVersion: 0.3\nSummary: Gradually Warm-up LR Scheduler for Pytorch\nHome-page:..."
},
{
"path": "Deraining/warmup_scheduler.egg-info/SOURCES.txt",
"chars": 271,
"preview": "README.md\nutils/__init__.py\nutils/dataset_utils.py\nutils/dir_utils.py\nutils/image_utils.py\nutils/model_utils.py\nwarmup_s..."
},
{
"path": "Deraining/warmup_scheduler.egg-info/dependency_links.txt",
"chars": 1,
"preview": "\n"
},
{
"path": "Deraining/warmup_scheduler.egg-info/top_level.txt",
"chars": 6,
"preview": "utils\n"
},
{
"path": "README.md",
"chars": 3561,
"preview": "# Deep Generalized Unfolding Networks for Image Restoration (CVPR 2022)\n[Chong Mou](https://scholar.google.com.h..."
},
{
"path": "figs/README.md",
"chars": 1,
"preview": "\n"
},
{
"path": "pytorch-gradual-warmup-lr/build/lib/warmup_scheduler/__init__.py",
"chars": 63,
"preview": "\nfrom warmup_scheduler.scheduler import GradualWarmupScheduler\n"
},
{
"path": "pytorch-gradual-warmup-lr/build/lib/warmup_scheduler/run.py",
"chars": 817,
"preview": "import torch\nfrom torch.optim.lr_scheduler import StepLR, ExponentialLR\nfrom torch.optim.sgd import SGD\n\nfrom warmup_sch..."
},
{
"path": "pytorch-gradual-warmup-lr/build/lib/warmup_scheduler/scheduler.py",
"chars": 3069,
"preview": "from torch.optim.lr_scheduler import _LRScheduler\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau\n\n\nclass Gradual..."
},
{
"path": "pytorch-gradual-warmup-lr/setup.py",
"chars": 578,
"preview": "from __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport set..."
},
{
"path": "pytorch-gradual-warmup-lr/warmup_scheduler/__init__.py",
"chars": 63,
"preview": "\nfrom warmup_scheduler.scheduler import GradualWarmupScheduler\n"
},
{
"path": "pytorch-gradual-warmup-lr/warmup_scheduler/run.py",
"chars": 817,
"preview": "import torch\nfrom torch.optim.lr_scheduler import StepLR, ExponentialLR\nfrom torch.optim.sgd import SGD\n\nfrom warmup_sch..."
},
{
"path": "pytorch-gradual-warmup-lr/warmup_scheduler/scheduler.py",
"chars": 3069,
"preview": "from torch.optim.lr_scheduler import _LRScheduler\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau\n\n\nclass Gradual..."
},
{
"path": "pytorch-gradual-warmup-lr/warmup_scheduler.egg-info/PKG-INFO",
"chars": 273,
"preview": "Metadata-Version: 1.0\nName: warmup-scheduler\nVersion: 0.3\nSummary: Gradually Warm-up LR Scheduler for Pytorch\nHome-page:..."
},
{
"path": "pytorch-gradual-warmup-lr/warmup_scheduler.egg-info/SOURCES.txt",
"chars": 251,
"preview": "setup.py\nwarmup_scheduler/__init__.py\nwarmup_scheduler/run.py\nwarmup_scheduler/scheduler.py\nwarmup_scheduler.egg-info/PK..."
},
{
"path": "pytorch-gradual-warmup-lr/warmup_scheduler.egg-info/dependency_links.txt",
"chars": 1,
"preview": "\n"
},
{
"path": "pytorch-gradual-warmup-lr/warmup_scheduler.egg-info/top_level.txt",
"chars": 17,
"preview": "warmup_scheduler\n"
}
]
About this extraction
This page contains the full source code of the MC-E/Deep-Generalized-Unfolding-Networks-for-Image-Restoration GitHub repository, extracted and formatted as plain text for AI agents and large language models (LLMs). The extraction includes 197 files (1.5 MB), approximately 406.9k tokens. Use this with OpenClaw, Claude, ChatGPT, Cursor, Windsurf, or any other AI tool that accepts text input. You can copy the full output to your clipboard or download it as a .txt file.
Extracted by GitExtract — free GitHub repo to text converter for AI. Built by Nikandr Surkov.