2021-1-capstone-design2 / 2017103084

Go to a project
Toggle navigation Toggle navigation pinning
서민정
db2fc7ef 1 parent 99c6f4d8

docs&code: 중간보고서 수정 및 code 업로드

Expand all
Inline Side-by-side
Showing 18 changed files with 376 additions and 383 deletions
The MIT License (MIT)
Copyright (c) 2017- Jiu XU
Copyright (c) 2017- Rakuten, Inc
Copyright (c) 2017- Rakuten Institute of Technology
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
\ No newline at end of file
# PyTorch VDSR
Implementation of CVPR2016 Paper: "Accurate Image Super-Resolution Using
Very Deep Convolutional Networks"(http://cv.snu.ac.kr/research/VDSR/) in PyTorch
## Usage
### Training
```
usage: main_vdsr.py [-h] [--batchSize BATCHSIZE] [--nEpochs NEPOCHS] [--lr LR]
[--step STEP] [--cuda] [--resume RESUME]
[--start-epoch START_EPOCH] [--clip CLIP] [--threads THREADS]
[--momentum MOMENTUM] [--weight-decay WEIGHT_DECAY]
[--pretrained PRETRAINED] [--gpus GPUS]
optional arguments:
-h, --help Show this help message and exit
--batchSize Training batch size
--nEpochs Number of epochs to train for
--lr Learning rate. Default=0.01
--step Learning rate decay, Default: n=10 epochs
--cuda Use cuda
--resume Path to checkpoint
--clip Clipping Gradients. Default=0.4
--threads Number of threads for data loader to use Default=1
--momentum Momentum, Default: 0.9
--weight-decay Weight decay, Default: 1e-4
--pretrained PRETRAINED
path to pretrained model (default: none)
--gpus GPUS gpu ids (default: 0)
```
An example of training usage is shown as follows:
```
python main_vdsr.py --cuda --gpus 0
```
### Evaluation
```
usage: eval.py [-h] [--cuda] [--model MODEL] [--dataset DATASET]
[--scale SCALE] [--gpus GPUS]
PyTorch VDSR Eval
optional arguments:
-h, --help show this help message and exit
--cuda use cuda?
--model MODEL model path
--dataset DATASET dataset name, Default: Set5
--gpus GPUS gpu ids (default: 0)
```
An example of training usage is shown as follows:
```
python eval.py --cuda --dataset Set5
```
### Demo
```
usage: demo.py [-h] [--cuda] [--model MODEL] [--image IMAGE] [--scale SCALE] [--gpus GPUS]
optional arguments:
-h, --help Show this help message and exit
--cuda Use cuda
--model Model path. Default=model/model_epoch_50.pth
--image Image name. Default=butterfly_GT
--scale Scale factor, Default: 4
--gpus GPUS gpu ids (default: 0)
```
An example of usage is shown as follows:
```
python eval.py --model model/model_epoch_50.pth --dataset Set5 --cuda
```
### Prepare Training dataset
- We provide a simple hdf5 format training sample in data folder with 'data' and 'label' keys, the training data is generated with Matlab Bicubic Interplotation, please refer [Code for Data Generation](https://github.com/twtygqyy/pytorch-vdsr/tree/master/data) for creating training files.
### Performance
- We provide a pretrained VDSR model trained on [291](https://drive.google.com/open?id=1Rt3asDLuMgLuJvPA1YrhyjWhb97Ly742) images with data augmentation
- No bias is used in this implementation, and the gradient clipping's implementation is different from paper
- Performance in PSNR on Set5
| Scale | VDSR Paper | VDSR PyTorch|
| ------------- |:-------------:| -----:|
| 2x | 37.53 | 37.65 |
| 3x | 33.66 | 33.77|
| 4x | 31.35 | 31.45 |
### Result
From left to right are ground truth, bicubic and vdsr
<p>
<img src='Set5/butterfly_GT.bmp' height='200' width='200'/>
<img src='result/input.bmp' height='200' width='200'/>
<img src='result/output.bmp' height='200' width='200'/>
</p>
import torch.utils.data as data
import torch
import h5py
class DatasetFromHdf5(data.Dataset):
def __init__(self, file_path):
super(DatasetFromHdf5, self).__init__()
hf = h5py.File(file_path)
self.data = hf.get('data')
self.target = hf.get('label')
def __getitem__(self, index):
return torch.from_numpy(self.data[index,:,:,:]).float(), torch.from_numpy(self.target[index,:,:,:]).float()
def __len__(self):
return self.data.shape[0]
\ No newline at end of file
import argparse, os
import torch
from torch.autograd import Variable
from scipy.ndimage import imread
from PIL import Image
import numpy as np
import time, math
import matplotlib.pyplot as plt
parser = argparse.ArgumentParser(description="PyTorch VDSR Demo")
parser.add_argument("--cuda", action="store_true", help="use cuda?")
parser.add_argument("--model", default="model/model_epoch_50.pth", type=str, help="model path")
parser.add_argument("--image", default="butterfly_GT", type=str, help="image name")
parser.add_argument("--scale", default=4, type=int, help="scale factor, Default: 4")
parser.add_argument("--gpus", default="0", type=str, help="gpu ids (default: 0)")
def PSNR(pred, gt, shave_border=0):
height, width = pred.shape[:2]
pred = pred[shave_border:height - shave_border, shave_border:width - shave_border]
gt = gt[shave_border:height - shave_border, shave_border:width - shave_border]
imdff = pred - gt
rmse = math.sqrt(np.mean(imdff ** 2))
if rmse == 0:
return 100
return 20 * math.log10(255.0 / rmse)
def colorize(y, ycbcr):
img = np.zeros((y.shape[0], y.shape[1], 3), np.uint8)
img[:,:,0] = y
img[:,:,1] = ycbcr[:,:,1]
img[:,:,2] = ycbcr[:,:,2]
img = Image.fromarray(img, "YCbCr").convert("RGB")
return img
opt = parser.parse_args()
cuda = opt.cuda
if cuda:
print("=> use gpu id: '{}'".format(opt.gpus))
os.environ["CUDA_VISIBLE_DEVICES"] = opt.gpus
if not torch.cuda.is_available():
raise Exception("No GPU found or Wrong gpu id, please run without --cuda")
model = torch.load(opt.model, map_location=lambda storage, loc: storage)["model"]
im_gt_ycbcr = imread("Set5/" + opt.image + ".bmp", mode="YCbCr")
im_b_ycbcr = imread("Set5/"+ opt.image + "_scale_"+ str(opt.scale) + ".bmp", mode="YCbCr")
im_gt_y = im_gt_ycbcr[:,:,0].astype(float)
im_b_y = im_b_ycbcr[:,:,0].astype(float)
psnr_bicubic = PSNR(im_gt_y, im_b_y,shave_border=opt.scale)
im_input = im_b_y/255.
im_input = Variable(torch.from_numpy(im_input).float()).view(1, -1, im_input.shape[0], im_input.shape[1])
if cuda:
model = model.cuda()
im_input = im_input.cuda()
else:
model = model.cpu()
start_time = time.time()
out = model(im_input)
elapsed_time = time.time() - start_time
out = out.cpu()
im_h_y = out.data[0].numpy().astype(np.float32)
im_h_y = im_h_y * 255.
im_h_y[im_h_y < 0] = 0
im_h_y[im_h_y > 255.] = 255.
psnr_predicted = PSNR(im_gt_y, im_h_y[0,:,:], shave_border=opt.scale)
im_h = colorize(im_h_y[0,:,:], im_b_ycbcr)
im_gt = Image.fromarray(im_gt_ycbcr, "YCbCr").convert("RGB")
im_b = Image.fromarray(im_b_ycbcr, "YCbCr").convert("RGB")
print("Scale=",opt.scale)
print("PSNR_predicted=", psnr_predicted)
print("PSNR_bicubic=", psnr_bicubic)
print("It takes {}s for processing".format(elapsed_time))
fig = plt.figure()
ax = plt.subplot("131")
ax.imshow(im_gt)
ax.set_title("GT")
ax = plt.subplot("132")
ax.imshow(im_b)
ax.set_title("Input(bicubic)")
ax = plt.subplot("133")
ax.imshow(im_h)
ax.set_title("Output(vdsr)")
plt.show()
import argparse, os
import torch
from torch.autograd import Variable
import numpy as np
import time, math, glob
import scipy.io as sio
parser = argparse.ArgumentParser(description="PyTorch VDSR Eval")
parser.add_argument("--cuda", action="store_true", help="use cuda?")
parser.add_argument("--model", default="model/model_epoch_50.pth", type=str, help="model path")
parser.add_argument("--dataset", default="Set5", type=str, help="dataset name, Default: Set5")
parser.add_argument("--gpus", default="0", type=str, help="gpu ids (default: 0)")
def PSNR(pred, gt, shave_border=0):
height, width = pred.shape[:2]
pred = pred[shave_border:height - shave_border, shave_border:width - shave_border]
gt = gt[shave_border:height - shave_border, shave_border:width - shave_border]
imdff = pred - gt
rmse = math.sqrt(np.mean(imdff ** 2))
if rmse == 0:
return 100
return 20 * math.log10(255.0 / rmse)
opt = parser.parse_args()
cuda = opt.cuda
if cuda:
print("=> use gpu id: '{}'".format(opt.gpus))
os.environ["CUDA_VISIBLE_DEVICES"] = opt.gpus
if not torch.cuda.is_available():
raise Exception("No GPU found or Wrong gpu id, please run without --cuda")
model = torch.load(opt.model, map_location=lambda storage, loc: storage)["model"]
scales = [2,3,4]
image_list = glob.glob(opt.dataset+"_mat/*.*")
for scale in scales:
avg_psnr_predicted = 0.0
avg_psnr_bicubic = 0.0
avg_elapsed_time = 0.0
count = 0.0
for image_name in image_list:
if str(scale) in image_name:
count += 1
print("Processing ", image_name)
im_gt_y = sio.loadmat(image_name)['im_gt_y']
im_b_y = sio.loadmat(image_name)['im_b_y']
im_gt_y = im_gt_y.astype(float)
im_b_y = im_b_y.astype(float)
psnr_bicubic = PSNR(im_gt_y, im_b_y,shave_border=scale)
avg_psnr_bicubic += psnr_bicubic
im_input = im_b_y/255.
im_input = Variable(torch.from_numpy(im_input).float()).view(1, -1, im_input.shape[0], im_input.shape[1])
if cuda:
model = model.cuda()
im_input = im_input.cuda()
else:
model = model.cpu()
start_time = time.time()
HR = model(im_input)
elapsed_time = time.time() - start_time
avg_elapsed_time += elapsed_time
HR = HR.cpu()
im_h_y = HR.data[0].numpy().astype(np.float32)
im_h_y = im_h_y * 255.
im_h_y[im_h_y < 0] = 0
im_h_y[im_h_y > 255.] = 255.
im_h_y = im_h_y[0,:,:]
psnr_predicted = PSNR(im_gt_y, im_h_y,shave_border=scale)
avg_psnr_predicted += psnr_predicted
print("Scale=", scale)
print("Dataset=", opt.dataset)
print("PSNR_predicted=", avg_psnr_predicted/count)
print("PSNR_bicubic=", avg_psnr_bicubic/count)
print("It takes average {}s for processing".format(avg_elapsed_time/count))
# Feature SR
1. train
`!python main.py --dataRoot /content/drive/MyDrive/feature/HR_trainset/features --scaleFactor 4 --featureType p6 --batchSize 16 --cuda --nEpochs 20`
2. inference
`!python inference.py --cuda --model "model.pth" --dataset "/content/drive/MyDrive/feature/features/LR_2" --featureType "p3" --scaleFactor 4`
3. calculate mAP
```
# [1]
# install dependencies:
!pip install pyyaml==5.1
import torch, torchvision
print(torch.__version__, torch.cuda.is_available())
!gcc --version
# opencv is pre-installed on colab
# [2]
# install detectron2: (Colab has CUDA 10.1 + torch 1.8)
# See https://detectron2.readthedocs.io/tutorials/install.html for instructions
import torch
assert torch.__version__.startswith("1.8") # need to manually install torch 1.8 if Colab changes its default version
!pip install detectron2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu101/torch1.8/index.html
# exit(0) # After installation, you need to "restart runtime" in Colab. This line can also restart runtime
# [3]
# Some basic setup:
# Setup detectron2 logger
import detectron2
from detectron2.utils.logger import setup_logger
setup_logger()
!python calculate_mAP.py --valid_data_path /content/drive/MyDrive/dataset/validset_100/ --model_name VDSR --loss_type MSE --batch_size 16
```
\ No newline at end of file
This diff is collapsed. Click to expand it.
import os
from PIL import Image
def crop_feature(datapath, feature_type, scale_factor, print_message=False):
data_path = datapath
datatype = feature_type
rescale_factor = scale_factor
if not os.path.exists(data_path):
raise Exception(f"[!] {data_path} not existed")
hr_imgs = []
w, h = Image.open(datapath).size
width = int(w / 16)
height = int(h / 16)
lwidth = int(width / rescale_factor)
lheight = int(height / rescale_factor)
if print_message:
print("lr: ({} {}), hr: ({} {})".format(lwidth, lheight, width, height))
hr_image = Image.open(datapath) # .convert('RGB')\
for i in range(16):
for j in range(16):
(left, upper, right, lower) = (
i * width, j * height, (i + 1) * width, (j + 1) * height)
crop = hr_image.crop((left, upper, right, lower))
crop = crop.resize((lwidth,lheight), Image.BICUBIC)
crop = crop.resize((width, height), Image.BICUBIC)
hr_imgs.append(crop)
return hr_imgs
\ No newline at end of file
......@@ -4,64 +4,87 @@ import os
from glob import glob
from torchvision import transforms
from torch.utils.data.dataset import Dataset
from torchvision import transforms
import torch
import pdb
import math
import numpy as np
class FeatureDataset(Dataset):
def __init__(self, data_path, datatype, rescale_factor,valid):
def __init__(self, data_path, datatype, rescale_factor, valid):
self.data_path = data_path
self.datatype = datatype
self.rescale_factor = rescale_factor
if not os.path.exists(data_path):
raise Exception(f"[!] {self.data_path} not existed")
if(valid):
self.hr_path = os.path.join(self.data_path,'valid')
self.hr_path = os.path.join(self.hr_path,self.datatype)
if (valid):
self.hr_path = os.path.join(self.data_path, 'valid')
self.hr_path = os.path.join(self.hr_path, self.datatype)
else:
self.hr_path = os.path.join(self.data_path,'LR_2')
self.hr_path = os.path.join(self.hr_path,self.datatype)
self.hr_path = os.path.join(self.data_path, 'LR_2')
self.hr_path = os.path.join(self.hr_path, self.datatype)
print(self.hr_path)
self.hr_path = sorted(glob(os.path.join(self.hr_path, "*.*")))
self.hr_imgs = []
w,h = Image.open(self.hr_path[0]).size
self.width = int(w/16)
self.height = int(h/16)
self.lwidth = int(self.width/self.rescale_factor)
self.lheight = int(self.height/self.rescale_factor)
print("lr: ({} {}), hr: ({} {})".format(self.lwidth,self.lheight,self.width,self.height))
for hr in self.hr_path:
hr_image = Image.open(hr)#.convert('RGB')\
w, h = Image.open(self.hr_path[0]).size
self.width = int(w / 16)
self.height = int(h / 16)
self.lwidth = int(self.width / self.rescale_factor) # rescale_factor만큼 크기를 줄인다.
self.lheight = int(self.height / self.rescale_factor)
print("lr: ({} {}), hr: ({} {})".format(self.lwidth, self.lheight, self.width, self.height))
for hr in self.hr_path: # 256개의 피쳐로 나눈다.
hr_image = Image.open(hr) # .convert('RGB')\
for i in range(16):
for j in range(16):
(left,upper,right,lower) = (i*self.width,j*self.height,(i+1)*self.width,(j+1)*self.height)
crop = hr_image.crop((left,upper,right,lower))
(left, upper, right, lower) = (
i * self.width, j * self.height, (i + 1) * self.width, (j + 1) * self.height)
crop = hr_image.crop((left, upper, right, lower))
self.hr_imgs.append(crop)
def __getitem__(self, idx):
hr_image = self.hr_imgs[idx]
transform = transforms.Compose([
transforms.Resize((self.lheight,self.lwidth),3),
transforms.Resize((self.lheight, self.lwidth), Image.BICUBIC),
transforms.Resize((self.height, self.width), Image.BICUBIC),
transforms.ToTensor()
])
return transform(hr_image), transforms.ToTensor()(hr_image)
return transform(hr_image), transforms.ToTensor()(hr_image) # hr_image를 변환한 것과, 변환하지 않은 것을 Tensor로 각각 반환
def __len__(self):
return len(self.hr_path*16*16)
return len(self.hr_path * 16 * 16)
def get_data_loader_test_version(data_path, feature_type, rescale_factor, batch_size, num_workers):
full_dataset = FeatureDataset(data_path, feature_type, rescale_factor, False)
print("dataset의 사이즈는 {}".format(len(full_dataset)))
for f in full_dataset:
print(type(f))
def get_data_loader(data_path, feature_type, rescale_factor, batch_size, num_workers):
full_dataset = FeatureDataset(data_path,feature_type,rescale_factor,False)
full_dataset = FeatureDataset(data_path, feature_type, rescale_factor, False)
train_size = int(0.9 * len(full_dataset))
test_size = len(full_dataset) - train_size
train_dataset, test_dataset = torch.utils.data.random_split(full_dataset, [train_size, test_size])
torch.manual_seed(3334)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,batch_size=batch_size,shuffle=True,num_workers=num_workers, pin_memory=False)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset,batch_size=batch_size,shuffle=False,num_workers=num_workers, pin_memory=True)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True,
num_workers=num_workers, pin_memory=False)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False,
num_workers=num_workers, pin_memory=True)
return train_loader, test_loader
def get_training_data_loader(data_path, feature_type, rescale_factor, batch_size, num_workers):
full_dataset = FeatureDataset(data_path, feature_type, rescale_factor, False)
torch.manual_seed(3334)
train_loader = torch.utils.data.DataLoader(dataset=full_dataset, batch_size=batch_size, shuffle=True,
num_workers=num_workers, pin_memory=False)
return train_loader
def get_infer_dataloader(data_path, feature_type, rescale_factor, batch_size, num_workers):
dataset = FeatureDataset(data_path,feature_type,rescale_factor,True)
data_loader = torch.utils.data.DataLoader(dataset=dataset,batch_size=batch_size,shuffle=False,num_workers=num_workers,pin_memory=False)
dataset = FeatureDataset(data_path, feature_type, rescale_factor, True)
data_loader = torch.utils.data.DataLoader(dataset=dataset, batch_size=batch_size, shuffle=False,
num_workers=num_workers, pin_memory=False)
return data_loader
\ No newline at end of file
......
import argparse, os
import torch
from torch.autograd import Variable
import numpy as np
import time, math, glob
import scipy.io as sio
from crop_feature import crop_feature
from PIL import Image
import cv2
from matplotlib import pyplot as plt
from math import log10, sqrt
parser = argparse.ArgumentParser(description="PyTorch VDSR Eval")
parser.add_argument("--cuda", action="store_true", help="use cuda?")
parser.add_argument("--model", default="model/model_epoch_50.pth", type=str, help="model path")
parser.add_argument("--dataset", default="Set5", type=str, help="dataset name, Default: Set5")
parser.add_argument("--gpus", default="0", type=str, help="gpu ids (default: 0)")
parser.add_argument("--featureType", default="p3", type=str)
parser.add_argument("--scaleFactor", default=4, type=int, help="scale factor")
parser.add_argument("--singleImage", type=str, default="N", help="if it is a single image, enter \"y\"")
def PSNR(pred, gt, shave_border=0):
height, width = pred.shape[:2]
pred = pred[shave_border:height - shave_border, shave_border:width - shave_border]
gt = gt[shave_border:height - shave_border, shave_border:width - shave_border]
imdff = pred - gt
rmse = math.sqrt(np.mean(imdff ** 2))
if rmse == 0:
return 100
return 20 * math.log10(255.0 / rmse)
def concatFeatures(features, image_name, bicubic=False):
features_0 = features[:16]
features_1 = features[16:32]
features_2 = features[32:48]
features_3 = features[48:64]
features_4 = features[64:80]
features_5 = features[80:96]
features_6 = features[96:112]
features_7 = features[112:128]
features_8 = features[128:144]
features_9 = features[144:160]
features_10 = features[160:176]
features_11 = features[176:192]
features_12 = features[192:208]
features_13 = features[208:224]
features_14 = features[224:240]
features_15 = features[240:256]
features_new = list()
features_new.extend([
concat_vertical(features_0),
concat_vertical(features_1),
concat_vertical(features_2),
concat_vertical(features_3),
concat_vertical(features_4),
concat_vertical(features_5),
concat_vertical(features_6),
concat_vertical(features_7),
concat_vertical(features_8),
concat_vertical(features_9),
concat_vertical(features_10),
concat_vertical(features_11),
concat_vertical(features_12),
concat_vertical(features_13),
concat_vertical(features_14),
concat_vertical(features_15)
])
final_concat_feature = concat_horizontal(features_new)
if bicubic:
save_path = "features/LR_2/LR/" + opt.featureType + "/" + image_name
if not os.path.exists("features/"):
os.makedirs("features/")
if not os.path.exists("features/LR_2/"):
os.makedirs("features/LR_2/")
if not os.path.exists("features/LR_2/LR/"):
os.makedirs("features/LR_2/LR/")
if not os.path.exists("features/LR_2/LR/" + opt.featureType):
os.makedirs("features/LR_2/LR/" + opt.featureType)
cv2.imwrite(save_path, final_concat_feature)
else:
save_path = "features/LR_2/" + opt.featureType + "/" + image_name
if not os.path.exists("features/"):
os.makedirs("features/")
if not os.path.exists("features/LR_2/"):
os.makedirs("features/LR_2/")
if not os.path.exists("features/LR_2/" + opt.featureType):
os.makedirs("features/LR_2/" + opt.featureType)
cv2.imwrite(save_path, final_concat_feature)
def concat_horizontal(feature):
result = cv2.hconcat([feature[0], feature[1]])
for i in range(2, len(feature)):
result = cv2.hconcat([result, feature[i]])
return result
def concat_vertical(feature):
result = cv2.vconcat([feature[0], feature[1]])
for i in range(2, len(feature)):
result = cv2.vconcat([result, feature[i]])
return result
opt = parser.parse_args()
cuda = opt.cuda
if cuda:
print("=> use gpu id: '{}'".format(opt.gpus))
os.environ["CUDA_VISIBLE_DEVICES"] = opt.gpus
if not torch.cuda.is_available():
raise Exception("No GPU found or Wrong gpu id, please run without --cuda")
model = torch.load(opt.model, map_location=lambda storage, loc: storage)["model"]
scales = [opt.scaleFactor]
# image_list = glob.glob(opt.dataset+"/*.*")
if opt.singleImage == "Y" :
# image_list = crop_feature(opt.dataset, opt.featureType, opt.scaleFactor)
image_list = opt.dataset
else:
image_path = os.path.join(opt.dataset, opt.featureType)
image_list = os.listdir(image_path)
print(image_path)
print(image_list)
for scale in scales:
for image in image_list:
avg_psnr_predicted = 0.0
avg_psnr_bicubic = 0.0
avg_elapsed_time = 0.0
count = 0.0
image_name_cropped = crop_feature(os.path.join(image_path, image), opt.featureType, opt.scaleFactor)
features = []
features_bicubic = []
for image_name in image_name_cropped:
count += 1
f_gt = image_name
w, h = image_name.size
f_bi = image_name.resize((w//scale,h//scale), Image.BICUBIC)
f_bi = f_bi.resize((w,h), Image.BICUBIC)
f_gt = np.array(f_gt)
f_bi = np.array(f_bi)
f_gt = f_gt.astype(float)
f_bi = f_bi.astype(float)
features_bicubic.append(f_bi)
psnr_bicubic = PSNR(f_bi, f_gt, shave_border=scale)
# psnr_bicubic = PSNR_ver2(cv2.imread(f_gt), cv2.imread(f_bi))
avg_psnr_bicubic += psnr_bicubic
f_input = f_bi/255.
f_input = Variable(torch.from_numpy(f_input).float()).view(1, -1, f_input.shape[0], f_input.shape[1])
if cuda:
model = model.cuda()
f_input = f_input.cuda()
else:
model = model.cpu()
start_time = time.time()
SR = model(f_input)
elapsed_time = time.time() - start_time
avg_elapsed_time += elapsed_time
SR = SR.cpu()
f_sr = SR.data[0].numpy().astype(np.float32)
f_sr = f_sr * 255
f_sr[f_sr<0] = 0
f_sr[f_sr>255.] = 255.
f_sr = f_sr[0,:,:]
psnr_predicted = PSNR(f_sr, f_gt, shave_border=scale)
# psnr_predicted = PSNR_ver2(cv2.imread(f_gt), cv2.imread(f_sr))
avg_psnr_predicted += psnr_predicted
features.append(f_sr)
concatFeatures(features, image)
concatFeatures(features_bicubic, image, True)
print("Scale=", scale)
print("Dataset=", opt.dataset)
print("Average PSNR_predicted=", avg_psnr_predicted/count)
print("Average PSNR_bicubic=", avg_psnr_bicubic/count)
# Show graph
# f_gt = Image.fromarray(f_gt)
# f_b = Image.fromarray(f_bi)
# f_sr = Image.fromarray(f_sr)
# fig = plt.figure(figsize=(18, 16), dpi= 80)
# ax = plt.subplot("131")
# ax.imshow(f_gt)
# ax.set_title("GT")
# ax = plt.subplot("132")
# ax.imshow(f_bi)
# ax.set_title("Input(bicubic)")
# ax = plt.subplot("133")
# ax.imshow(f_sr)
# ax.set_title("Output(vdsr)")
# plt.show()
This diff could not be displayed because it is too large.
import argparse, os
from datasets import get_data_loader
import torch
import random
import torch.backends.cudnn as cudnn
......@@ -7,33 +8,37 @@ import torch.optim as optim
from torch.autograd import Variable
from torch.utils.data import DataLoader
from vdsr import Net
from dataset import DatasetFromHdf5
## Custom
from data import FeatureDataset
from datasets import get_training_data_loader
# from datasets import get_data_loader_test_version
# from feature_dataset import get_training_data_loader
# from make_dataset import make_dataset
import numpy as np
from dataFromH5 import Read_dataset_h5
import matplotlib.pyplot as plt
import math
# Training settings
parser = argparse.ArgumentParser(description="PyTorch VDSR")
parser.add_argument("--batchSize", type=int, default=128, help="Training batch size")
parser.add_argument("--nEpochs", type=int, default=50, help="Number of epochs to train for")
parser.add_argument("--lr", type=float, default=0.1, help="Learning Rate. Default=0.1")
parser.add_argument("--dataRoot", type=str)
parser.add_argument("--featureType", type=str)
parser.add_argument("--scaleFactor", type=int, default=4)
parser.add_argument("--batchSize", type=int, default=64, help="Training batch size")
parser.add_argument("--nEpochs", type=int, default=20, help="Number of epochs to train for")
parser.add_argument("--lr", type=float, default=0.001, help="Learning Rate. Default=0.1")
parser.add_argument("--step", type=int, default=10, help="Sets the learning rate to the initial LR decayed by momentum every n epochs, Default: n=10")
parser.add_argument("--cuda", action="store_true", help="Use cuda?")
parser.add_argument("--resume", default="", type=str, help="Path to checkpoint (default: none)")
parser.add_argument("--start-epoch", default=1, type=int, help="Manual epoch number (useful on restarts)")
parser.add_argument("--clip", type=float, default=0.4, help="Clipping Gradients. Default=0.4")
# 1->3 custom
parser.add_argument("--threads", type=int, default=3, help="Number of threads for data loader to use, Default: 1")
parser.add_argument("--threads", type=int, default=1, help="Number of threads for data loader to use, Default: 1")
parser.add_argument("--momentum", default=0.9, type=float, help="Momentum, Default: 0.9")
parser.add_argument("--weight-decay", "--wd", default=1e-4, type=float, help="Weight decay, Default: 1e-4")
parser.add_argument('--pretrained', default='', type=str, help='path to pretrained model (default: none)')
parser.add_argument("--gpus", default="0", type=str, help="gpu ids (default: 0)")
## custom
parser.add_argument("--dataPath", type=str)
parser.add_argument("--featureType", type=str, default="p2")
parser.add_argument("--scaleFactor",type=int, default=4)
# parser.add_argument("--trainingData", type=DataLoader)
total_loss_for_plot = list()
total_pnsr = list()
def main():
global opt, model
......@@ -55,21 +60,17 @@ def main():
cudnn.benchmark = True
print("===> Loading datasets")
################## Loading Datasets ##########################
if os.path.isfile('dataloader/training_data_loader.pth'):
training_data_loader = torch.load('dataloader/training_data_loader.pth')
else:
train_set = FeatureDataset(opt.dataPath,opt.featureType,opt.scaleFactor,False)
train_size = 100 #우선은 100개만
test_size = len(train_set) - train_size
train_dataset, test_dataset = torch.utils.data.random_split(train_set, [train_size, test_size])
training_data_loader = DataLoader(dataset=train_dataset, num_workers=3, batch_size=8, shuffle=True, pin_memory=False)
torch.save(training_data_loader, 'dataloader/training_data_loader.pth'.format(DataLoader))
print("===> Loading datasets")
# train_set = DatasetFromHdf5("data/train.h5")
# training_data_loader = DataLoader(dataset=train_set, num_workers=opt.threads, batch_size=opt.batchSize, shuffle=True)
training_data_loader = get_training_data_loader(opt.dataRoot, opt.featureType, opt.scaleFactor, opt.batchSize, opt.threads)
# training_data_loader = make_dataset(opt.dataRoot, opt.featureType, opt.scaleFactor, opt.batchSize, opt.threads)
print("===> Building model")
model = Net(opt.scaleFactor)
criterion = nn.MSELoss(size_average=False)
model = Net()
criterion = nn.MSELoss(reduction='sum')
print("===> Setting GPU")
if cuda:
......@@ -91,6 +92,7 @@ def main():
if os.path.isfile(opt.pretrained):
print("=> loading model '{}'".format(opt.pretrained))
weights = torch.load(opt.pretrained)
opt.start_epoch = weights["epoch"] + 1
model.load_state_dict(weights['model'].state_dict())
else:
print("=> no model found at '{}'".format(opt.pretrained))
......@@ -99,15 +101,21 @@ def main():
optimizer = optim.SGD(model.parameters(), lr=opt.lr, momentum=opt.momentum, weight_decay=opt.weight_decay)
print("===> Training")
for epoch in range(opt.start_epoch, opt.nEpochs + 1):
train(training_data_loader, optimizer, model, criterion, epoch)
save_checkpoint(model, epoch)
save_checkpoint(model, epoch, optimizer)
def adjust_learning_rate(optimizer, epoch):
"""Sets the learning rate to the initial LR decayed by 10 every 10 epochs"""
lr = opt.lr * (0.1 ** (epoch // opt.step))
return lr
def PSNR(loss):
psnr = 10 * np.log10(1 / (loss + 1e-10))
# psnr = 20 * math.log10(255.0 / (math.sqrt(loss)))
return psnr
def train(training_data_loader, optimizer, model, criterion, epoch):
lr = adjust_learning_rate(optimizer, epoch-1)
......@@ -119,25 +127,31 @@ def train(training_data_loader, optimizer, model, criterion, epoch):
model.train()
for iteration, batch in enumerate(training_data_loader, 1):
input, target = Variable(batch[0]), Variable(batch[1], requires_grad=False)
optimizer.zero_grad()
input, target = Variable(batch[0], requires_grad=False), Variable(batch[1], requires_grad=False)
total_loss = 0
if opt.cuda:
input = input.cuda()
target = target.cuda()
loss = criterion(model(input), target)
optimizer.zero_grad()
total_loss += loss.item()
loss.backward()
nn.utils.clip_grad_norm(model.parameters(),opt.clip)
nn.utils.clip_grad_norm_(model.parameters(), opt.clip)
optimizer.step()
if iteration%10 == 0:
# loss.data[0] --> loss.data
print("===> Epoch[{}]({}/{}): Loss: {:.10f}".format(epoch, iteration, len(training_data_loader), loss.data))
def save_checkpoint(model, epoch):
model_out_path = "checkpoint/" + "model_epoch_{}.pth".format(epoch)
state = {"epoch": epoch ,"model": model}
epoch_loss = total_loss / len(training_data_loader)
total_loss_for_plot.append(epoch_loss)
psnr = PSNR(epoch_loss)
total_pnsr.append(psnr)
print("===> Epoch[{}]: loss : {:.10f} ,PSNR : {:.10f}".format(epoch, epoch_loss, psnr))
# if iteration%100 == 0:
# print("===> Epoch[{}]({}/{}): Loss: {:.10f}".format(epoch, iteration, len(training_data_loader), loss.item()))
def save_checkpoint(model, epoch, optimizer):
model_out_path = "checkpoint/" + "model_epoch_{}_{}.pth".format(epoch, opt.featureType)
state = {"epoch": epoch ,"model": model, "model_state_dict":model.state_dict(), "optimizer_state_dict":optimizer.state_dict(),
"loss": total_loss_for_plot, "psnr":total_pnsr}
if not os.path.exists("checkpoint/"):
os.makedirs("checkpoint/")
......
......@@ -12,12 +12,11 @@ class Conv_ReLU_Block(nn.Module):
return self.relu(self.conv(x))
class Net(nn.Module):
def __init__(self,upscale_factor):
def __init__(self):
super(Net, self).__init__()
self.residual_layer = self.make_layer(Conv_ReLU_Block, 18)
self.input = nn.Conv2d(in_channels=1, out_channels=64, kernel_size=3, stride=1, padding=1, bias=False)
self.output = nn.Conv2d(in_channels=64, out_channels=1, kernel_size=3, stride=1, padding=1, bias=False)
self.upsample = nn.Upsample(scale_factor=upscale_factor, mode='bicubic')
self.relu = nn.ReLU(inplace=True)
for m in self.modules():
......@@ -32,7 +31,6 @@ class Net(nn.Module):
return nn.Sequential(*layers)
def forward(self, x):
x = self.upsample(x)
residual = x
out = self.relu(self.input(x))
out = self.residual_layer(out)
......
No preview for this file type
No preview for this file type