整合去雾网页工具
This commit is contained in:
admin
49
GCANet/All_in_One.sh
Normal file
49
GCANet/All_in_One.sh
Normal file
@@ -0,0 +1,49 @@
|
||||
#!/bin/bash
|
||||
# 原图像位置
|
||||
Dir_src_pics="./imgs/src"
|
||||
Dir_result="./imgs/result"
|
||||
Dir_ori_src_pics="/root/Dehaze/SRC_files/src"
|
||||
|
||||
mkdir -p $Dir_src_pics $Dir_result
|
||||
|
||||
PS3='All in one choice : '
|
||||
applications=("Delete_src_pics" "Delete_generate_pics" "Copy_src_pics" "Run_program" "quit")
|
||||
select fav in "${applications[@]}"; do
|
||||
case $fav in
|
||||
# 删除原始文件选项
|
||||
"Delete_src_pics")
|
||||
# 删除src文件
|
||||
echo "Delete all src files in $Dir_src_pics"
|
||||
rm $Dir_src_pics/*
|
||||
;;
|
||||
|
||||
# 删除生成文件选项
|
||||
"Delete_generate_pics")
|
||||
# 删除result文件
|
||||
echo "Delete all src files in $Dir_result"
|
||||
rm $Dir_result/*
|
||||
;;
|
||||
|
||||
# 复制待处理文件选项
|
||||
"Copy_src_pics")
|
||||
# 删除src文件
|
||||
echo "Copy all src files in $Dir_ori_src_pics"
|
||||
ln -s $Dir_ori_src_pics/* $Dir_src_pics
|
||||
;;
|
||||
|
||||
# 运行程序
|
||||
"Run_program")
|
||||
source ~/miniconda/bin/activate Dehaze_GCANet
|
||||
python test.py --task dehaze --gpu_id 0 --indir ./imgs/src --outdir ./imgs/result
|
||||
;;
|
||||
|
||||
# 退出选项
|
||||
"quit")
|
||||
echo "User requested exit"
|
||||
exit
|
||||
;;
|
||||
|
||||
# 其他选项
|
||||
*) echo "invalid option $REPLY";;
|
||||
esac
|
||||
done
|
||||
102
GCANet/GCANet.py
Normal file
102
GCANet/GCANet.py
Normal file
@@ -0,0 +1,102 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class ShareSepConv(nn.Module):
|
||||
def __init__(self, kernel_size):
|
||||
super(ShareSepConv, self).__init__()
|
||||
assert kernel_size % 2 == 1, 'kernel size should be odd'
|
||||
self.padding = (kernel_size - 1)//2
|
||||
weight_tensor = torch.zeros(1, 1, kernel_size, kernel_size)
|
||||
weight_tensor[0, 0, (kernel_size-1)//2, (kernel_size-1)//2] = 1
|
||||
self.weight = nn.Parameter(weight_tensor)
|
||||
self.kernel_size = kernel_size
|
||||
|
||||
def forward(self, x):
|
||||
inc = x.size(1)
|
||||
expand_weight = self.weight.expand(inc, 1, self.kernel_size, self.kernel_size).contiguous()
|
||||
return F.conv2d(x, expand_weight,
|
||||
None, 1, self.padding, 1, inc)
|
||||
|
||||
|
||||
class SmoothDilatedResidualBlock(nn.Module):
|
||||
def __init__(self, channel_num, dilation=1, group=1):
|
||||
super(SmoothDilatedResidualBlock, self).__init__()
|
||||
self.pre_conv1 = ShareSepConv(dilation*2-1)
|
||||
self.conv1 = nn.Conv2d(channel_num, channel_num, 3, 1, padding=dilation, dilation=dilation, groups=group, bias=False)
|
||||
self.norm1 = nn.InstanceNorm2d(channel_num, affine=True)
|
||||
self.pre_conv2 = ShareSepConv(dilation*2-1)
|
||||
self.conv2 = nn.Conv2d(channel_num, channel_num, 3, 1, padding=dilation, dilation=dilation, groups=group, bias=False)
|
||||
self.norm2 = nn.InstanceNorm2d(channel_num, affine=True)
|
||||
|
||||
def forward(self, x):
|
||||
y = F.relu(self.norm1(self.conv1(self.pre_conv1(x))))
|
||||
y = self.norm2(self.conv2(self.pre_conv2(y)))
|
||||
return F.relu(x+y)
|
||||
|
||||
|
||||
class ResidualBlock(nn.Module):
|
||||
def __init__(self, channel_num, dilation=1, group=1):
|
||||
super(ResidualBlock, self).__init__()
|
||||
self.conv1 = nn.Conv2d(channel_num, channel_num, 3, 1, padding=dilation, dilation=dilation, groups=group, bias=False)
|
||||
self.norm1 = nn.InstanceNorm2d(channel_num, affine=True)
|
||||
self.conv2 = nn.Conv2d(channel_num, channel_num, 3, 1, padding=dilation, dilation=dilation, groups=group, bias=False)
|
||||
self.norm2 = nn.InstanceNorm2d(channel_num, affine=True)
|
||||
|
||||
def forward(self, x):
|
||||
y = F.relu(self.norm1(self.conv1(x)))
|
||||
y = self.norm2(self.conv2(y))
|
||||
return F.relu(x+y)
|
||||
|
||||
|
||||
class GCANet(nn.Module):
|
||||
def __init__(self, in_c=4, out_c=3, only_residual=True):
|
||||
super(GCANet, self).__init__()
|
||||
self.conv1 = nn.Conv2d(in_c, 64, 3, 1, 1, bias=False)
|
||||
self.norm1 = nn.InstanceNorm2d(64, affine=True)
|
||||
self.conv2 = nn.Conv2d(64, 64, 3, 1, 1, bias=False)
|
||||
self.norm2 = nn.InstanceNorm2d(64, affine=True)
|
||||
self.conv3 = nn.Conv2d(64, 64, 3, 2, 1, bias=False)
|
||||
self.norm3 = nn.InstanceNorm2d(64, affine=True)
|
||||
|
||||
self.res1 = SmoothDilatedResidualBlock(64, dilation=2)
|
||||
self.res2 = SmoothDilatedResidualBlock(64, dilation=2)
|
||||
self.res3 = SmoothDilatedResidualBlock(64, dilation=2)
|
||||
self.res4 = SmoothDilatedResidualBlock(64, dilation=4)
|
||||
self.res5 = SmoothDilatedResidualBlock(64, dilation=4)
|
||||
self.res6 = SmoothDilatedResidualBlock(64, dilation=4)
|
||||
self.res7 = ResidualBlock(64, dilation=1)
|
||||
|
||||
self.gate = nn.Conv2d(64 * 3, 3, 3, 1, 1, bias=True)
|
||||
|
||||
self.deconv3 = nn.ConvTranspose2d(64, 64, 4, 2, 1)
|
||||
self.norm4 = nn.InstanceNorm2d(64, affine=True)
|
||||
self.deconv2 = nn.Conv2d(64, 64, 3, 1, 1)
|
||||
self.norm5 = nn.InstanceNorm2d(64, affine=True)
|
||||
self.deconv1 = nn.Conv2d(64, out_c, 1)
|
||||
self.only_residual = only_residual
|
||||
|
||||
def forward(self, x):
|
||||
y = F.relu(self.norm1(self.conv1(x)))
|
||||
y = F.relu(self.norm2(self.conv2(y)))
|
||||
y1 = F.relu(self.norm3(self.conv3(y)))
|
||||
|
||||
y = self.res1(y1)
|
||||
y = self.res2(y)
|
||||
y = self.res3(y)
|
||||
y2 = self.res4(y)
|
||||
y = self.res5(y2)
|
||||
y = self.res6(y)
|
||||
y3 = self.res7(y)
|
||||
|
||||
gates = self.gate(torch.cat((y1, y2, y3), dim=1))
|
||||
gated_y = y1 * gates[:, [0], :, :] + y2 * gates[:, [1], :, :] + y3 * gates[:, [2], :, :]
|
||||
y = F.relu(self.norm4(self.deconv3(gated_y)))
|
||||
y = F.relu(self.norm5(self.deconv2(y)))
|
||||
if self.only_residual:
|
||||
y = self.deconv1(y)
|
||||
else:
|
||||
y = F.relu(self.deconv1(y))
|
||||
|
||||
return y
|
||||
102
GCANet/GCANet_train/GCANet.py
Normal file
102
GCANet/GCANet_train/GCANet.py
Normal file
@@ -0,0 +1,102 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class ShareSepConv(nn.Module):
|
||||
def __init__(self, kernel_size):
|
||||
super(ShareSepConv, self).__init__()
|
||||
assert kernel_size % 2 == 1, 'kernel size should be odd'
|
||||
self.padding = (kernel_size - 1)//2
|
||||
weight_tensor = torch.zeros(1, 1, kernel_size, kernel_size)
|
||||
weight_tensor[0, 0, (kernel_size-1)//2, (kernel_size-1)//2] = 1
|
||||
self.weight = nn.Parameter(weight_tensor)
|
||||
self.kernel_size = kernel_size
|
||||
|
||||
def forward(self, x):
|
||||
inc = x.size(1)
|
||||
expand_weight = self.weight.expand(inc, 1, self.kernel_size, self.kernel_size).contiguous()
|
||||
return F.conv2d(x, expand_weight,
|
||||
None, 1, self.padding, 1, inc)
|
||||
|
||||
|
||||
class SmoothDilatedResidualBlock(nn.Module):
|
||||
def __init__(self, channel_num, dilation=1, group=1):
|
||||
super(SmoothDilatedResidualBlock, self).__init__()
|
||||
self.pre_conv1 = ShareSepConv(dilation*2-1)
|
||||
self.conv1 = nn.Conv2d(channel_num, channel_num, 3, 1, padding=dilation, dilation=dilation, groups=group, bias=False)
|
||||
self.norm1 = nn.InstanceNorm2d(channel_num, affine=True)
|
||||
self.pre_conv2 = ShareSepConv(dilation*2-1)
|
||||
self.conv2 = nn.Conv2d(channel_num, channel_num, 3, 1, padding=dilation, dilation=dilation, groups=group, bias=False)
|
||||
self.norm2 = nn.InstanceNorm2d(channel_num, affine=True)
|
||||
|
||||
def forward(self, x):
|
||||
y = F.relu(self.norm1(self.conv1(self.pre_conv1(x))))
|
||||
y = self.norm2(self.conv2(self.pre_conv2(y)))
|
||||
return F.relu(x+y)
|
||||
|
||||
|
||||
class ResidualBlock(nn.Module):
|
||||
def __init__(self, channel_num, dilation=1, group=1):
|
||||
super(ResidualBlock, self).__init__()
|
||||
self.conv1 = nn.Conv2d(channel_num, channel_num, 3, 1, padding=dilation, dilation=dilation, groups=group, bias=False)
|
||||
self.norm1 = nn.InstanceNorm2d(channel_num, affine=True)
|
||||
self.conv2 = nn.Conv2d(channel_num, channel_num, 3, 1, padding=dilation, dilation=dilation, groups=group, bias=False)
|
||||
self.norm2 = nn.InstanceNorm2d(channel_num, affine=True)
|
||||
|
||||
def forward(self, x):
|
||||
y = F.relu(self.norm1(self.conv1(x)))
|
||||
y = self.norm2(self.conv2(y))
|
||||
return F.relu(x+y)
|
||||
|
||||
|
||||
class GCANet(nn.Module):
|
||||
def __init__(self, in_c=4, out_c=3, only_residual=True):
|
||||
super(GCANet, self).__init__()
|
||||
self.conv1 = nn.Conv2d(in_c, 64, 3, 1, 1, bias=False)
|
||||
self.norm1 = nn.InstanceNorm2d(64, affine=True)
|
||||
self.conv2 = nn.Conv2d(64, 64, 3, 1, 1, bias=False)
|
||||
self.norm2 = nn.InstanceNorm2d(64, affine=True)
|
||||
self.conv3 = nn.Conv2d(64, 64, 3, 2, 1, bias=False)
|
||||
self.norm3 = nn.InstanceNorm2d(64, affine=True)
|
||||
|
||||
self.res1 = SmoothDilatedResidualBlock(64, dilation=2)
|
||||
self.res2 = SmoothDilatedResidualBlock(64, dilation=2)
|
||||
self.res3 = SmoothDilatedResidualBlock(64, dilation=2)
|
||||
self.res4 = SmoothDilatedResidualBlock(64, dilation=4)
|
||||
self.res5 = SmoothDilatedResidualBlock(64, dilation=4)
|
||||
self.res6 = SmoothDilatedResidualBlock(64, dilation=4)
|
||||
self.res7 = ResidualBlock(64, dilation=1)
|
||||
|
||||
self.gate = nn.Conv2d(64 * 3, 3, 3, 1, 1, bias=True)
|
||||
|
||||
self.deconv3 = nn.ConvTranspose2d(64, 64, 4, 2, 1)
|
||||
self.norm4 = nn.InstanceNorm2d(64, affine=True)
|
||||
self.deconv2 = nn.Conv2d(64, 64, 3, 1, 1)
|
||||
self.norm5 = nn.InstanceNorm2d(64, affine=True)
|
||||
self.deconv1 = nn.Conv2d(64, out_c, 1)
|
||||
self.only_residual = only_residual
|
||||
|
||||
def forward(self, x):
|
||||
y = F.relu(self.norm1(self.conv1(x)))
|
||||
y = F.relu(self.norm2(self.conv2(y)))
|
||||
y1 = F.relu(self.norm3(self.conv3(y)))
|
||||
|
||||
y = self.res1(y1)
|
||||
y = self.res2(y)
|
||||
y = self.res3(y)
|
||||
y2 = self.res4(y)
|
||||
y = self.res5(y2)
|
||||
y = self.res6(y)
|
||||
y3 = self.res7(y)
|
||||
|
||||
gates = self.gate(torch.cat((y1, y2, y3), dim=1))
|
||||
gated_y = y1 * gates[:, [0], :, :] + y2 * gates[:, [1], :, :] + y3 * gates[:, [2], :, :]
|
||||
y = F.relu(self.norm4(self.deconv3(gated_y)))
|
||||
y = F.relu(self.norm5(self.deconv2(y)))
|
||||
if self.only_residual:
|
||||
y = self.deconv1(y)
|
||||
else:
|
||||
y = F.relu(self.deconv1(y))
|
||||
|
||||
return y
|
||||
104
GCANet/GCANet_train/ImagePairPrefixFolder.py
Normal file
104
GCANet/GCANet_train/ImagePairPrefixFolder.py
Normal file
@@ -0,0 +1,104 @@
|
||||
import os
|
||||
import bisect
|
||||
import threading
|
||||
import torch
|
||||
import numpy as np
|
||||
import numpy.random as random
|
||||
from PIL import Image
|
||||
from torch.utils.data import Dataset
|
||||
from folder_loader import FolderLoader
|
||||
import torchvision.transforms as transforms
|
||||
from utils import batch_edge_compute
|
||||
|
||||
def pil_loader(img_path):
|
||||
return Image.open(img_path).convert("RGB")
|
||||
|
||||
|
||||
class ImagePairPrefixFolder(Dataset):
|
||||
def __init__(self, input_folder, gt_folder, max_img_size=0, size_unit=1, force_rgb=False):
|
||||
super(ImagePairPrefixFolder, self).__init__()
|
||||
|
||||
self.gt_loader = FolderLoader(gt_folder)
|
||||
# build the map from image name to index
|
||||
self.gt_map = dict()
|
||||
for idx, img_name in enumerate(self.gt_loader.img_names):
|
||||
self.gt_map[os.path.splitext(img_name)[0].split('_')[0]] = idx
|
||||
|
||||
self.input_loader = FolderLoader(input_folder)
|
||||
assert all([os.path.splitext(x)[0].split('_')[0] in self.gt_map for x in self.input_loader.img_names]), \
|
||||
'cannot find corresponding gt names'
|
||||
|
||||
|
||||
self.input_folder = input_folder
|
||||
self.gt_folder = gt_folder
|
||||
self.max_img_size = max_img_size
|
||||
self.size_unit = size_unit
|
||||
self.force_rgb = force_rgb
|
||||
|
||||
def __getitem__(self, index):
|
||||
input_name, input_img = self.input_loader[index]
|
||||
input_basename = os.path.splitext(input_name)[0].split('_')[0]
|
||||
gt_idx = self.gt_map[input_basename]
|
||||
|
||||
gt_name, gt_img = self.gt_loader[gt_idx]
|
||||
if self.force_rgb:
|
||||
input_img = input_img.convert('RGB')
|
||||
gt_img = gt_img.convert('RGB')
|
||||
im_w, im_h = input_img.size
|
||||
gt_w, gt_h = gt_img.size
|
||||
assert im_w==gt_w and im_h==gt_h, 'input image and gt image size not match'
|
||||
|
||||
im_w, im_h = input_img.size
|
||||
if 0 < self.max_img_size < max(im_w, im_h):
|
||||
if im_w < im_h:
|
||||
out_h = int(self.max_img_size) // self.size_unit * self.size_unit
|
||||
out_w = int(im_w / im_h * out_h) // self.size_unit * self.size_unit
|
||||
else:
|
||||
out_w = int(self.max_img_size) // self.size_unit * self.size_unit
|
||||
out_h = int(im_h / im_w * out_w) // self.size_unit * self.size_unit
|
||||
else:
|
||||
out_w = im_w // self.size_unit * self.size_unit
|
||||
out_h = im_h // self.size_unit * self.size_unit
|
||||
|
||||
if im_w != out_w or im_h != out_h:
|
||||
input_img = input_img.resize((out_w, out_h), Image.BILINEAR)
|
||||
gt_img = gt_img.resize((out_w, out_h), Image.BILINEAR)
|
||||
|
||||
im_w, im_h = input_img.size
|
||||
|
||||
input_img = np.array(input_img).astype('float')
|
||||
gt_img = np.array(gt_img).astype('float')
|
||||
if len(input_img.shape) == 2:
|
||||
input_img = input_img[:, :, np.newaxis]
|
||||
if len(gt_img.shape) == 2:
|
||||
gt_img = gt_img[:, :, np.newaxis]
|
||||
return {'input_img': input_img, 'gt_img': gt_img, 'input_h': im_h, "input_w": im_w}
|
||||
|
||||
def get_input_info(self, index):
|
||||
image_name = os.path.splitext(self.input_loader.img_names[index])[0]
|
||||
return self.input_loader, image_name
|
||||
|
||||
def __len__(self):
|
||||
return len(self.input_loader)
|
||||
|
||||
|
||||
def var_custom_collate(batch):
|
||||
min_h, min_w = 10000, 10000
|
||||
for item in batch:
|
||||
min_h = min(min_h, item['input_h'])
|
||||
min_w = min(min_w, item['input_w'])
|
||||
inc = 1 if len(batch[0]['input_img'].shape)==2 else batch[0]['input_img'].shape[2]
|
||||
batch_input_images = torch.Tensor(len(batch), inc, min_h, min_w)
|
||||
batch_gt_images = torch.Tensor(len(batch), inc, min_h, min_w)
|
||||
|
||||
for idx, item in enumerate(batch):
|
||||
off_y = 0 if item['input_h']==min_h else random.randint(0, item['input_h'] - min_h)
|
||||
off_x = 0 if item['input_w']==min_w else random.randint(0, item['input_w'] - min_w)
|
||||
crop_input_img = item['input_img'][off_y:off_y + min_h, off_x:off_x + min_w, :]
|
||||
crop_gt_img = item['gt_img'][off_y:off_y + min_h, off_x:off_x + min_w, :]
|
||||
batch_input_images[idx] = torch.from_numpy(crop_input_img.transpose((2, 0, 1))) - 128
|
||||
batch_gt_images[idx] = torch.from_numpy(crop_gt_img.transpose((2, 0, 1)))
|
||||
|
||||
|
||||
batch_input_edges = batch_edge_compute(batch_input_images) - 128
|
||||
return batch_input_images, batch_input_edges, batch_gt_images
|
||||
19
GCANet/GCANet_train/folder_loader.py
Normal file
19
GCANet/GCANet_train/folder_loader.py
Normal file
@@ -0,0 +1,19 @@
|
||||
import io
|
||||
import os
|
||||
import utils
|
||||
import struct
|
||||
from PIL import Image
|
||||
|
||||
class FolderLoader(object):
|
||||
def __init__(self, fold_path):
|
||||
super(FolderLoader, self).__init__()
|
||||
self.fold_path = fold_path
|
||||
self.img_paths = utils.make_dataset(self.fold_path)
|
||||
self.img_names = [os.path.basename(x) for x in self.img_paths]
|
||||
|
||||
def __getitem__(self, index):
|
||||
img = Image.open(self.img_paths[index])#.convert('RGB')
|
||||
return self.img_names[index], img
|
||||
|
||||
def __len__(self):
|
||||
return len(self.img_names)
|
||||
66
GCANet/GCANet_train/test.py
Normal file
66
GCANet/GCANet_train/test.py
Normal file
@@ -0,0 +1,66 @@
|
||||
import os
|
||||
import argparse
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
import torch
|
||||
from torch.autograd import Variable
|
||||
|
||||
from utils import make_dataset, edge_compute
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--network', default='GCANet')
|
||||
parser.add_argument('--task', default='dehaze', help='dehaze | derain')
|
||||
parser.add_argument('--gpu_id', type=int, default=0)
|
||||
parser.add_argument('--indir', default='examples/')
|
||||
parser.add_argument('--outdir', default='output')
|
||||
opt = parser.parse_args()
|
||||
assert opt.task in ['dehaze', 'derain']
|
||||
## forget to regress the residue for deraining by mistake,
|
||||
## which should be able to produce better results
|
||||
opt.only_residual = opt.task == 'dehaze'
|
||||
opt.model = 'models/wacv_gcanet_%s.pth' % opt.task
|
||||
opt.use_cuda = opt.gpu_id >= 0
|
||||
if not os.path.exists(opt.outdir):
|
||||
os.makedirs(opt.outdir)
|
||||
test_img_paths = make_dataset(opt.indir)
|
||||
|
||||
if opt.network == 'GCANet':
|
||||
from GCANet import GCANet
|
||||
net = GCANet(in_c=4, out_c=3, only_residual=opt.only_residual)
|
||||
else:
|
||||
print('network structure %s not supported' % opt.network)
|
||||
raise ValueError
|
||||
|
||||
if opt.use_cuda:
|
||||
torch.cuda.set_device(opt.gpu_id)
|
||||
net.cuda()
|
||||
else:
|
||||
net.float()
|
||||
|
||||
net.load_state_dict(torch.load(opt.model, map_location='cpu'))
|
||||
net.eval()
|
||||
|
||||
for img_path in test_img_paths:
|
||||
img = Image.open(img_path).convert('RGB')
|
||||
im_w, im_h = img.size
|
||||
if im_w % 4 != 0 or im_h % 4 != 0:
|
||||
img = img.resize((int(im_w // 4 * 4), int(im_h // 4 * 4)))
|
||||
img = np.array(img).astype('float')
|
||||
img_data = torch.from_numpy(img.transpose((2, 0, 1))).float()
|
||||
edge_data = edge_compute(img_data)
|
||||
in_data = torch.cat((img_data, edge_data), dim=0).unsqueeze(0) - 128
|
||||
in_data = in_data.cuda() if opt.use_cuda else in_data.float()
|
||||
with torch.no_grad():
|
||||
pred = net(Variable(in_data))
|
||||
if opt.only_residual:
|
||||
out_img_data = (pred.data[0].cpu().float() + img_data).round().clamp(0, 255)
|
||||
else:
|
||||
out_img_data = pred.data[0].cpu().float().round().clamp(0, 255)
|
||||
out_img = Image.fromarray(out_img_data.numpy().astype(np.uint8).transpose(1, 2, 0))
|
||||
out_img.save(os.path.join(opt.outdir, os.path.splitext(os.path.basename(img_path))[0] + '_%s.png' % opt.task))
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
41
GCANet/GCANet_train/tf_visualizer.py
Normal file
41
GCANet/GCANet_train/tf_visualizer.py
Normal file
@@ -0,0 +1,41 @@
|
||||
import numpy as np
|
||||
import os
|
||||
import ntpath
|
||||
import time
|
||||
import utils
|
||||
from scipy.misc import imresize
|
||||
from tensorboardX import SummaryWriter
|
||||
|
||||
|
||||
class TFVisualizer():
|
||||
def __init__(self, opt):
|
||||
self.tf_visualizer = SummaryWriter(os.path.join(opt.logDir, opt.name))
|
||||
self.opt = opt
|
||||
self.saved = False
|
||||
self.ncols = 4
|
||||
|
||||
self.log_name = os.path.join(opt.checkpoints_dir, opt.name, 'loss_log.txt')
|
||||
with open(self.log_name, "a") as log_file:
|
||||
now = time.strftime("%c")
|
||||
log_file.write('================ Training Loss (%s) ================\n' % now)
|
||||
|
||||
def reset(self):
|
||||
self.saved = False
|
||||
|
||||
# |visuals|: dictionary of images to display or save
|
||||
def display_current_results(self, visuals, iter_mark, epoch, save_result):
|
||||
for label, image in visuals.items():
|
||||
img_gid = utils.tensor2imgrid(image)
|
||||
self.tf_visualizer.add_image(label, img_gid, iter_mark)
|
||||
|
||||
# losses: dictionary of error labels and values
|
||||
def plot_current_losses(self, iter_mark, losses):
|
||||
# for label, loss in losses.items():
|
||||
# self.tf_visualizer.add_scalar(label, loss, iter_mark)
|
||||
self.tf_visualizer.add_scalars('training loss', losses, iter_mark)
|
||||
|
||||
def print_logs(self, message):
|
||||
print(message)
|
||||
with open(self.log_name, "a") as log_file:
|
||||
log_file.write('%s\n' % message)
|
||||
|
||||
228
GCANet/GCANet_train/train.py
Normal file
228
GCANet/GCANet_train/train.py
Normal file
@@ -0,0 +1,228 @@
|
||||
import os
|
||||
import datetime
|
||||
import argparse
|
||||
import numpy as np
|
||||
|
||||
import torch
|
||||
import torch.optim as optim
|
||||
from torch.autograd import Variable
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
from ImagePairPrefixFolder import ImagePairPrefixFolder, var_custom_collate
|
||||
from utils import MovingAvg
|
||||
from tf_visualizer import TFVisualizer
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--network', default='GCANet')
|
||||
parser.add_argument('--name', default='default_exp')
|
||||
parser.add_argument('--gpu_ids', default='0')
|
||||
parser.add_argument('--epochs', type=int, default=100)
|
||||
parser.add_argument('--lr', type=float, default=0.001)
|
||||
parser.add_argument('--lr_step', type=int, default=40)
|
||||
parser.add_argument('--lr_gamma', type=float, default=0.1)
|
||||
parser.add_argument('--weight_decay', type=float, default=0.0005)
|
||||
parser.add_argument('--checkpoints_dir', default='checkpoint')
|
||||
parser.add_argument('--logDir', default='tblogdir')
|
||||
parser.add_argument('--resume_dir', default='')
|
||||
parser.add_argument('--resume_epoch', type=int, default=0)
|
||||
parser.add_argument('--save_epoch', type=int, default=5)
|
||||
parser.add_argument('--save_latest_freq', type=int, default=5000)
|
||||
parser.add_argument('--test_epoch', type=int, default=5)
|
||||
parser.add_argument('--test_max_size', type=int, default=1080)
|
||||
parser.add_argument('--size_unit', type=int, default=8)
|
||||
parser.add_argument('--print_iter', type=int, default=100)
|
||||
parser.add_argument('--input_folder', default='')
|
||||
parser.add_argument('--gt_folder', default='')
|
||||
parser.add_argument('--test_input_folder', default='')
|
||||
parser.add_argument('--test_gt_folder', default='')
|
||||
parser.add_argument('--num_workers', type=int, default=16)
|
||||
parser.add_argument('--batch_size', type=int, default=4)
|
||||
parser.add_argument('--only_residual', action='store_true', help='regress residual rather than image')
|
||||
parser.add_argument('--loss_func', default='l2', help='l2|l1')
|
||||
parser.add_argument('--inc', type=int, default=3)
|
||||
parser.add_argument('--outc', type=int, default=3)
|
||||
parser.add_argument('--force_rgb', action='store_true')
|
||||
parser.add_argument('--no_edge', action='store_true')
|
||||
|
||||
opt = parser.parse_args()
|
||||
|
||||
opt.input_folder = os.path.expanduser(opt.input_folder)
|
||||
opt.gt_folder = os.path.expanduser(opt.gt_folder)
|
||||
opt.test_input_folder = os.path.expanduser(opt.test_input_folder)
|
||||
opt.test_gt_folder = os.path.expanduser(opt.test_gt_folder)
|
||||
|
||||
|
||||
if not os.path.exists(os.path.join(opt.checkpoints_dir, opt.name)):
|
||||
os.makedirs(os.path.join(opt.checkpoints_dir, opt.name))
|
||||
opt.resume_dir = opt.resume_dir if opt.resume_dir != '' else os.path.join(opt.checkpoints_dir, opt.name)
|
||||
|
||||
visualizer = TFVisualizer(opt)
|
||||
### Log out
|
||||
with open(os.path.realpath(__file__), 'r') as fid:
|
||||
visualizer.print_logs(fid.read())
|
||||
|
||||
## print argument
|
||||
for key, val in vars(opt).items():
|
||||
visualizer.print_logs('%s: %s' % (key, val))
|
||||
|
||||
opt.gpu_ids = [int(x) for x in opt.gpu_ids.split(',')]
|
||||
assert all(0 <= x <= torch.cuda.device_count() for x in opt.gpu_ids), 'gpu id should ' \
|
||||
'be 0~{0}'.format(torch.cuda.device_count())
|
||||
torch.cuda.set_device(opt.gpu_ids[0])
|
||||
|
||||
|
||||
train_dataset = ImagePairPrefixFolder(opt.input_folder, opt.gt_folder, size_unit=opt.size_unit, force_rgb=opt.force_rgb)
|
||||
train_dataloader = DataLoader(train_dataset, batch_size=opt.batch_size, shuffle=True,
|
||||
collate_fn=var_custom_collate, pin_memory=True,
|
||||
num_workers=opt.num_workers)
|
||||
|
||||
opt.do_test = opt.test_gt_folder != ''
|
||||
if opt.do_test:
|
||||
test_dataset = ImagePairPrefixFolder(opt.test_input_folder, opt.test_gt_folder,
|
||||
max_img_size=opt.test_max_size, size_unit=opt.size_unit, force_rgb=opt.force_rgb)
|
||||
test_dataloader = DataLoader(test_dataset, batch_size=1, shuffle=False,
|
||||
collate_fn=var_custom_collate, pin_memory=True,
|
||||
num_workers=1)
|
||||
|
||||
total_inc = opt.inc if opt.no_edge else opt.inc + 1
|
||||
if opt.network == 'GCANet':
|
||||
from GCANet import GCANet
|
||||
net = GCANet(in_c=total_inc, out_c=3, only_residual=opt.only_residual)
|
||||
else:
|
||||
print('network structure %s not supported' % opt.network)
|
||||
raise ValueError
|
||||
|
||||
if opt.loss_func == 'l2':
|
||||
loss_crit = torch.nn.MSELoss()
|
||||
elif opt.loss_func == 'l1':
|
||||
loss_crit = torch.nn.SmoothL1Loss()
|
||||
else:
|
||||
print('loss_func %s not supported' % opt.loss_func)
|
||||
raise ValueError
|
||||
pnsr_crit = torch.nn.MSELoss()
|
||||
|
||||
if len(opt.gpu_ids) > 0:
|
||||
net.cuda()
|
||||
if len(opt.gpu_ids) > 1:
|
||||
net = torch.nn.DataParallel(net)
|
||||
loss_crit = loss_crit.cuda()
|
||||
pnsr_crit = pnsr_crit.cuda()
|
||||
|
||||
optimizer = optim.Adam(net.parameters(), lr=opt.lr)
|
||||
step_optim_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=opt.lr_step, gamma=opt.lr_gamma)
|
||||
loss_avg = MovingAvg(pool_size=50)
|
||||
|
||||
start_epoch = 0
|
||||
total_iter = 0
|
||||
|
||||
if os.path.exists(os.path.join(opt.checkpoints_dir, opt.name, 'latest.pth')):
|
||||
print('resuming from latest.pth')
|
||||
latest_info = torch.load(os.path.join(opt.checkpoints_dir, opt.name, 'latest.pth'))
|
||||
start_epoch = latest_info['epoch']
|
||||
total_iter = latest_info['total_iter']
|
||||
if isinstance(net, torch.nn.DataParallel):
|
||||
net.module.load_state_dict(latest_info['net_state'])
|
||||
else:
|
||||
net.load_state_dict(latest_info['net_state'])
|
||||
optimizer.load_state_dict(latest_info['optim_state'])
|
||||
|
||||
if opt.resume_epoch > 0:
|
||||
start_epoch = opt.resume_epoch
|
||||
total_iter = opt.resume_epoch * len(train_dataloader)
|
||||
resume_path = os.path.join(opt.resume_epoch, 'net_epoch_%d.pth') % opt.resume_epoch
|
||||
print('resume from : %s' % resume_path)
|
||||
assert os.path.exists(resume_path), 'cannot find the resume model: %s ' % resume_path
|
||||
if isinstance(net, torch.nn.DataParallel):
|
||||
net.module.load_state_dict(torch.load(resume_path))
|
||||
else:
|
||||
net.load_state_dict(torch.load(resume_path))
|
||||
|
||||
for epoch in range(start_epoch, opt.epochs):
|
||||
visualizer.print_logs("Start to train epoch %d" % epoch)
|
||||
net.train()
|
||||
for iter, data in enumerate(train_dataloader):
|
||||
total_iter += 1
|
||||
optimizer.zero_grad()
|
||||
step_optim_scheduler.step(epoch)
|
||||
|
||||
batch_input_img, batch_input_edge, batch_gt = data
|
||||
if len(opt.gpu_ids) > 0:
|
||||
batch_input_img, batch_input_edge, batch_gt = batch_input_img.cuda(), batch_input_edge.cuda(), batch_gt.cuda()
|
||||
|
||||
if opt.no_edge:
|
||||
batch_input = batch_input_img
|
||||
else:
|
||||
batch_input = torch.cat((batch_input_img, batch_input_edge), dim=1)
|
||||
batch_input_v = Variable(batch_input)
|
||||
if opt.only_residual:
|
||||
batch_gt_v = Variable(batch_gt - (batch_input_img+128))
|
||||
else:
|
||||
batch_gt_v = Variable(batch_gt)
|
||||
|
||||
pred = net(batch_input_v)
|
||||
|
||||
loss = loss_crit(pred, batch_gt_v)
|
||||
avg_loss = loss_avg.set_curr_val(loss.data)
|
||||
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
if iter % opt.print_iter == 0:
|
||||
visualizer.plot_current_losses(total_iter, { 'loss': loss})
|
||||
visualizer.print_logs('%s Step[%d/%d], lr: %f, mv_avg_loss: %f, loss: %f' %
|
||||
(str(datetime.datetime.now()).split(' ')[1], iter, len(train_dataloader),
|
||||
step_optim_scheduler.get_lr()[0], avg_loss, loss))
|
||||
|
||||
if total_iter % opt.save_latest_freq == 0:
|
||||
latest_info = {'total_iter': total_iter,
|
||||
'epoch': epoch,
|
||||
'optim_state': optimizer.state_dict()}
|
||||
if len(opt.gpu_ids) > 1:
|
||||
latest_info['net_state'] = net.module.state_dict()
|
||||
else:
|
||||
latest_info['net_state'] = net.state_dict()
|
||||
print('save lastest model.')
|
||||
torch.save(latest_info, os.path.join(opt.checkpoints_dir, opt.name, 'latest.pth'))
|
||||
|
||||
if (epoch+1) % opt.save_epoch == 0 :
|
||||
visualizer.print_logs('saving model for epoch %d' % epoch)
|
||||
if len(opt.gpu_ids) > 1:
|
||||
torch.save(net.module.state_dict(), os.path.join(opt.checkpoints_dir, opt.name, 'net_epoch_%d.pth' % (epoch+1)))
|
||||
else:
|
||||
torch.save(net.state_dict(), os.path.join(opt.checkpoints_dir, opt.name, 'net_epoch_%d.pth' % (epoch + 1)))
|
||||
|
||||
if opt.do_test:
|
||||
avg_psnr = 0
|
||||
task_cnt = 0
|
||||
net.eval()
|
||||
with torch.no_grad():
|
||||
for iter, data in enumerate(test_dataloader):
|
||||
batch_input_img, batch_input_edge, batch_gt = data
|
||||
if len(opt.gpu_ids) > 0:
|
||||
batch_input_img, batch_input_edge, batch_gt = batch_input_img.cuda(), batch_input_edge.cuda(), batch_gt.cuda()
|
||||
|
||||
if opt.no_edge:
|
||||
batch_input = batch_input_img
|
||||
else:
|
||||
batch_input = torch.cat((batch_input_img, batch_input_edge), dim=1)
|
||||
batch_input_v = Variable(batch_input)
|
||||
batch_gt_v = Variable(batch_gt)
|
||||
|
||||
|
||||
pred = net(batch_input_v)
|
||||
|
||||
if opt.only_residual:
|
||||
loss = pnsr_crit(pred+Variable(batch_input_img+128), batch_gt_v)
|
||||
else:
|
||||
loss = pnsr_crit(pred, batch_gt_v)
|
||||
avg_psnr += 10 * np.log10(255 * 255 / loss.item())
|
||||
task_cnt += 1
|
||||
|
||||
visualizer.print_logs('Testing for epoch: %d' % epoch)
|
||||
visualizer.print_logs('Average test PNSR is %f for %d images' % (avg_psnr/task_cnt, task_cnt))
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
198
GCANet/GCANet_train/utils.py
Normal file
198
GCANet/GCANet_train/utils.py
Normal file
@@ -0,0 +1,198 @@
|
||||
import os
|
||||
import torch
|
||||
import torch
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
import os
|
||||
|
||||
from scipy import signal
|
||||
from torchvision.utils import make_grid
|
||||
|
||||
IMG_EXTENSIONS = [
|
||||
'.jpg', '.JPG', '.jpeg', '.JPEG',
|
||||
'.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP',
|
||||
]
|
||||
|
||||
|
||||
def is_image_file(filename):
|
||||
return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
|
||||
|
||||
|
||||
def make_dataset(dir):
|
||||
images = []
|
||||
assert os.path.isdir(dir), '%s is not a valid directory' % dir
|
||||
|
||||
for root, _, fnames in sorted(os.walk(dir)):
|
||||
for fname in fnames:
|
||||
if is_image_file(fname):
|
||||
path = os.path.join(root, fname)
|
||||
images.append(path)
|
||||
|
||||
return images
|
||||
|
||||
|
||||
def edge_compute(x):
|
||||
x_diffx = torch.abs(x[:,:,1:] - x[:,:,:-1])
|
||||
x_diffy = torch.abs(x[:,1:,:] - x[:,:-1,:])
|
||||
|
||||
y = x.new(x.size())
|
||||
y.fill_(0)
|
||||
y[:,:,1:] += x_diffx
|
||||
y[:,:,:-1] += x_diffx
|
||||
y[:,1:,:] += x_diffy
|
||||
y[:,:-1,:] += x_diffy
|
||||
y = torch.sum(y,0,keepdim=True)/3
|
||||
y /= 4
|
||||
return y
|
||||
|
||||
def batch_edge_compute(x):
|
||||
x_diffx = torch.abs(x[:,:,:,1:] - x[:,:,:,:-1])
|
||||
x_diffy = torch.abs(x[:,:,1:,:] - x[:,:,:-1,:])
|
||||
|
||||
y = x.new(x.size())
|
||||
y.fill_(0)
|
||||
y[:,:,:,1:] += x_diffx
|
||||
y[:,:,:,:-1] += x_diffx
|
||||
y[:,:,1:,:] += x_diffy
|
||||
y[:,:,:-1,:] += x_diffy
|
||||
y = torch.sum(y,1,keepdim=True)/3
|
||||
y /= 4
|
||||
return y
|
||||
|
||||
# Converts a Tensor into an image array (numpy)
|
||||
# |imtype|: the desired type of the converted numpy array
|
||||
def tensor2im(input_image, imtype=np.uint8):
|
||||
if isinstance(input_image, torch.Tensor):
|
||||
image_tensor = input_image.data
|
||||
else:
|
||||
return input_image
|
||||
image_numpy = image_tensor[0].cpu().float().numpy()
|
||||
if image_numpy.shape[0] == 1:
|
||||
image_numpy = np.tile(image_numpy, (3, 1, 1))
|
||||
image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0
|
||||
image_numpy = image_numpy.clip(0, 255)
|
||||
return image_numpy.astype(imtype)
|
||||
|
||||
|
||||
def tensor2imgrid(input_image):
|
||||
im_grid = make_grid(input_image[:4, ...], nrow=2, normalize=True, range=(-128, 128))
|
||||
return im_grid
|
||||
# ndarr = im_grid.mul(255).clamp(0, 255).byte().permute(1, 2, 0).cpu().numpy()
|
||||
# im = Image.fromarray(ndarr)
|
||||
# return im
|
||||
|
||||
|
||||
def diagnose_network(net, name='network'):
|
||||
mean = 0.0
|
||||
count = 0
|
||||
for param in net.parameters():
|
||||
if param.grad is not None:
|
||||
mean += torch.mean(torch.abs(param.grad.data))
|
||||
count += 1
|
||||
if count > 0:
|
||||
mean = mean / count
|
||||
print(name)
|
||||
print(mean)
|
||||
|
||||
|
||||
def save_image(image_numpy, image_path):
|
||||
image_pil = Image.fromarray(image_numpy)
|
||||
image_pil.save(image_path)
|
||||
|
||||
|
||||
def print_numpy(x, val=True, shp=False):
|
||||
x = x.astype(np.float64)
|
||||
if shp:
|
||||
print('shape,', x.shape)
|
||||
if val:
|
||||
x = x.flatten()
|
||||
print('mean = %3.3f, min = %3.3f, max = %3.3f, median = %3.3f, std=%3.3f' % (
|
||||
np.mean(x), np.min(x), np.max(x), np.median(x), np.std(x)))
|
||||
|
||||
|
||||
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 fspecial_gauss(size, sigma):
|
||||
"""Function to mimic the 'fspecial' gaussian MATLAB function
|
||||
"""
|
||||
x, y = np.mgrid[-size // 2 + 1:size // 2 + 1, -size // 2 + 1:size // 2 + 1]
|
||||
g = np.exp(-((x ** 2 + y ** 2) / (2.0 * sigma ** 2)))
|
||||
return g / g.sum()
|
||||
|
||||
|
||||
def filter2(x, kernel, mode='same'):
|
||||
return signal.convolve2d(x, np.rot90(kernel, 2), mode=mode)
|
||||
|
||||
|
||||
def ssim(img1, img2, cs_map=False):
|
||||
"""Return the Structural Similarity Map corresponding to input images img1
|
||||
and img2 (images are assumed to be uint8)
|
||||
|
||||
This function attempts to mimic precisely the functionality of ssim.m a
|
||||
MATLAB provided by the author's of SSIM
|
||||
https://ece.uwaterloo.ca/~z70wang/research/ssim/ssim_index.m
|
||||
"""
|
||||
img1 = img1.astype(np.float64)
|
||||
img2 = img2.astype(np.float64)
|
||||
size = 11
|
||||
sigma = 1.5
|
||||
window = fspecial_gauss(size, sigma)
|
||||
K1 = 0.01
|
||||
K2 = 0.03
|
||||
L = 255 # bitdepth of image
|
||||
C1 = (K1 * L) ** 2
|
||||
C2 = (K2 * L) ** 2
|
||||
mu1 = filter2(img1, window, mode='valid')
|
||||
mu2 = filter2(img2, window, mode='valid')
|
||||
mu1_sq = mu1 * mu1
|
||||
mu2_sq = mu2 * mu2
|
||||
mu1_mu2 = mu1 * mu2
|
||||
sigma1_sq = filter2(img1 * img1, window, mode='valid') - mu1_sq
|
||||
sigma2_sq = filter2(img2 * img2, window, mode='valid') - mu2_sq
|
||||
sigma12 = filter2(img1 * img2, window, mode='valid') - mu1_mu2
|
||||
if cs_map:
|
||||
return np.mean(np.mean((((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) *
|
||||
(sigma1_sq + sigma2_sq + C2)),
|
||||
(2.0 * sigma12 + C2) / (sigma1_sq + sigma2_sq + C2))))
|
||||
else:
|
||||
return np.mean(np.mean(((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) *
|
||||
(sigma1_sq + sigma2_sq + C2))))
|
||||
|
||||
|
||||
|
||||
class MovingAvg(object):
|
||||
def __init__(self, pool_size=100):
|
||||
from queue import Queue
|
||||
self.pool = Queue(maxsize=pool_size)
|
||||
self.sum = 0
|
||||
self.curr_pool_size = 0
|
||||
self.pool_size = pool_size
|
||||
|
||||
def set_curr_val(self, val):
|
||||
if not self.pool.full():
|
||||
self.curr_pool_size += 1
|
||||
self.pool.put_nowait(val)
|
||||
else:
|
||||
last_first_val = self.pool.get_nowait()
|
||||
self.pool.put_nowait(val)
|
||||
self.sum -= last_first_val
|
||||
|
||||
self.sum += val
|
||||
return self.sum / self.curr_pool_size
|
||||
|
||||
def reset(self):
|
||||
from queue import Queue
|
||||
self.pool = Queue(maxsize=self.pool_size)
|
||||
self.sum = 0
|
||||
self.curr_pool_size = 0
|
||||
46
GCANet/README.md
Normal file
46
GCANet/README.md
Normal file
@@ -0,0 +1,46 @@
|
||||
Gated Context Aggregation Network for Image Dehazing and Deraining
|
||||
=======
|
||||
|
||||
|
||||
This is the implementation of our WACV 2019 paper *"Gated Context Aggregation Network for Image Dehazing and Deraining"* by [Dongdong Chen](<http://www.dongdongchen.bid/>), [Mingming He](<https://github.com/hmmlillian>), [Qingnan Fan](<https://fqnchina.github.io/>), *et al.*
|
||||
|
||||
In this paper, we propose a new end-to-end gated context aggregation network GCANet for image dehazing, in which the smoothed dilated convolution is used to avoid the gridding artifacts and a gated subnetwork is applied to fuse the features of different levels. Experiments show that GCANet can obtain much better performance than all the previous state-of-the-art image dehazing methods both qualitatively and quantitatively
|
||||
|
||||
|
||||
We further apply our proposed GCANet to the image deraining task, which also outperforms previous state-of-the-art image deraining methods and demonstrates its generality.
|
||||
|
||||
|
||||
|
||||
## Getting Started
|
||||
|
||||
This paper is implemented with Pytorch framework.
|
||||
|
||||
Demo
|
||||
----
|
||||
|
||||
Directly put all your test images under one directory. Then run:
|
||||
|
||||
```bash
|
||||
python test.py --task [dehaze | derain] --gpu_id [gpu_id] --indir [input directory] --outdir [output directory]
|
||||
```
|
||||
|
||||
For training, please download the training code from <https://drive.google.com/file/d/1T7X1HYztbz6S75vTRNtREgGEOI269KDk/view?usp=sharing>
|
||||
|
||||
Cite
|
||||
----
|
||||
|
||||
You can use our codes for research purpose only. And please cite our paper when you use our codes.
|
||||
```
|
||||
@article{chen2018gated,
|
||||
title={Gated Context Aggregation Network for Image Dehazing and Deraining},
|
||||
author={Chen, Dongdong and He, Mingming and Fan, Qingnan and Liao, Jing and Zhang, Liheng and Hou, Dongdong and Yuan, Lu and Hua, Gang},
|
||||
journal={WACV 2019},
|
||||
year={2018}
|
||||
}
|
||||
```
|
||||
Contact
|
||||
-------
|
||||
|
||||
If you find any bugs or have any ideas of optimizing these codes, please contact me via cddlyf [at] gmail [dot] com
|
||||
|
||||
|
||||
BIN
GCANet/models/wacv_gcanet_dehaze.pth
Normal file
BIN
GCANet/models/wacv_gcanet_dehaze.pth
Normal file
Binary file not shown.
BIN
GCANet/models/wacv_gcanet_derain.pth
Normal file
BIN
GCANet/models/wacv_gcanet_derain.pth
Normal file
Binary file not shown.
67
GCANet/test.py
Normal file
67
GCANet/test.py
Normal file
@@ -0,0 +1,67 @@
|
||||
import os
|
||||
import argparse
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
import torch
|
||||
from torch.autograd import Variable
|
||||
|
||||
from utils import make_dataset, edge_compute
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--network', default='GCANet')
|
||||
parser.add_argument('--task', default='dehaze', help='dehaze | derain')
|
||||
parser.add_argument('--gpu_id', type=int, default=0)
|
||||
parser.add_argument('--indir', default='examples/')
|
||||
parser.add_argument('--outdir', default='output')
|
||||
opt = parser.parse_args()
|
||||
assert opt.task in ['dehaze', 'derain']
|
||||
## forget to regress the residue for deraining by mistake,
|
||||
## which should be able to produce better results
|
||||
opt.only_residual = opt.task == 'dehaze'
|
||||
opt.model = 'models/wacv_gcanet_%s.pth' % opt.task
|
||||
opt.use_cuda = opt.gpu_id >= 0
|
||||
if not os.path.exists(opt.outdir):
|
||||
os.makedirs(opt.outdir)
|
||||
test_img_paths = make_dataset(opt.indir)
|
||||
|
||||
if opt.network == 'GCANet':
|
||||
from GCANet import GCANet
|
||||
net = GCANet(in_c=4, out_c=3, only_residual=opt.only_residual)
|
||||
else:
|
||||
print('network structure %s not supported' % opt.network)
|
||||
raise ValueError
|
||||
|
||||
if opt.use_cuda:
|
||||
torch.cuda.set_device(opt.gpu_id)
|
||||
net.cuda()
|
||||
else:
|
||||
net.float()
|
||||
|
||||
net.load_state_dict(torch.load(opt.model, map_location='cpu'))
|
||||
net.eval()
|
||||
|
||||
for img_path in test_img_paths:
|
||||
img = Image.open(img_path).convert('RGB')
|
||||
im_w, im_h = img.size
|
||||
if im_w % 4 != 0 or im_h % 4 != 0:
|
||||
img = img.resize((int(im_w // 4 * 4), int(im_h // 4 * 4)))
|
||||
img = np.array(img).astype('float')
|
||||
img_data = torch.from_numpy(img.transpose((2, 0, 1))).float()
|
||||
edge_data = edge_compute(img_data)
|
||||
in_data = torch.cat((img_data, edge_data), dim=0).unsqueeze(0) - 128
|
||||
in_data = in_data.cuda() if opt.use_cuda else in_data.float()
|
||||
with torch.no_grad():
|
||||
pred = net(Variable(in_data))
|
||||
if opt.only_residual:
|
||||
out_img_data = (pred.data[0].cpu().float() + img_data).round().clamp(0, 255)
|
||||
else:
|
||||
out_img_data = pred.data[0].cpu().float().round().clamp(0, 255)
|
||||
out_img = Image.fromarray(out_img_data.numpy().astype(np.uint8).transpose(1, 2, 0))
|
||||
print("-"*5,"图片存储在：",os.path.join(opt.outdir, os.path.splitext(os.path.basename(img_path))[0] + '_%s.png' % opt.task))
|
||||
out_img.save(os.path.join(opt.outdir, os.path.splitext(os.path.basename(img_path))[0] + '_%s.png' % opt.task))
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
39
GCANet/utils.py
Normal file
39
GCANet/utils.py
Normal file
@@ -0,0 +1,39 @@
|
||||
import os
|
||||
import torch
|
||||
|
||||
IMG_EXTENSIONS = [
|
||||
'.jpg', '.JPG', '.jpeg', '.JPEG',
|
||||
'.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP',
|
||||
]
|
||||
|
||||
|
||||
def is_image_file(filename):
|
||||
return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
|
||||
|
||||
|
||||
def make_dataset(dir):
|
||||
images = []
|
||||
assert os.path.isdir(dir), '%s is not a valid directory' % dir
|
||||
|
||||
for root, _, fnames in sorted(os.walk(dir)):
|
||||
for fname in fnames:
|
||||
if is_image_file(fname):
|
||||
path = os.path.join(root, fname)
|
||||
images.append(path)
|
||||
|
||||
return images
|
||||
|
||||
|
||||
def edge_compute(x):
|
||||
x_diffx = torch.abs(x[:,:,1:] - x[:,:,:-1])
|
||||
x_diffy = torch.abs(x[:,1:,:] - x[:,:-1,:])
|
||||
|
||||
y = x.new(x.size())
|
||||
y.fill_(0)
|
||||
y[:,:,1:] += x_diffx
|
||||
y[:,:,:-1] += x_diffx
|
||||
y[:,1:,:] += x_diffy
|
||||
y[:,:-1,:] += x_diffy
|
||||
y = torch.sum(y,0,keepdim=True)/3
|
||||
y /= 4
|
||||
return y
|
||||
Reference in New Issue
Block a user