Official Implementation of the CVPR 2021 Paper
[Project](https://junyonglee.me/projects/DMENet) | [Paper](https://openaccess.thecvf.com/content_CVPR_2019/papers/Lee_Deep_Defocus_Map_Estimation_Using_Domain_Adaptation_CVPR_2019_paper.pdf) | [Supp](https://www.dropbox.com/s/van0beau0npq3de/supp.zip?dl=1) | [Poster](https://www.dropbox.com/s/85umxea9uha3ptq/CVPR2019_poster.pdf?raw=1)
[-blue.svg?style=flat)](https://2JI532DIZN4TSYWF.anvil.app/BIEWGFSFTYML53VXPQZBRNTX) 2 | 3 | This repository contains the official matlab implementation of SYNDOF generation used in the following paper: 4 | 5 | > [**Deep Defocus Map Estimation using Domain Adaptation**](https://junyonglee.me/projects/DMENet)
6 | > [Junyong Lee](https://junyonglee.me)1, [Sungkil Lee](http://cg.skku.edu/slee/)2, [Sunghyun Cho](https://www.scho.pe.kr/)3, and [Seungyong Lee](http://cg.postech.ac.kr/leesy/)1
7 | > 1POSTECH, 2Sungkyunkwan University, 3DGIST
8 | > *IEEE Computer Vision and Pattern Recognition (**CVPR**) 2019*
9 | 10 | 11 |
12 |
13 | 
14 |
15 |
The provided checkpoint is the new checkpoint that shows the closest evaluation results as in the paper.* 96 | 97 | *Check out [updated performance](/evaluation) with the new checkpoint.* 98 | 99 | * Test the final model by: 100 | 101 | ```bash 102 | python main.py --mode DMENet_BDCS --test_set CUHK 103 | ``` 104 | > Testing results will be saved in `[LOG_ROOT]/DMENet_CVPR2019/[mode]/result/[test_set]/`: 105 | > 106 | > ``` 107 | > ... 108 | > [LOG_ROOT]/DMENet_CVPR2019/[mode]/result/ 109 | > └── [test_set] 110 | > ├── image # input defocused images 111 | > ├── defocus_map # defocus images (network's direct output in range [0, 1]) 112 | > ├── defocus_map_min_max_norm # min-max normalized defocus images in range [0, 1] for visualization 113 | > └── sigma_map_7_norm # sigma maps containing normalized standard deviations (in range [0, 1]) for a Gaussian kernel. For the actual standard deviation value, one should multiply 7 to this map. 114 | > ``` 115 | 116 | > Quantitative results are computed from matlab. (*e.g.*, [evaluation on the RTF dataset](https://github.com/codeslake/DMENet/tree/master/evaluation/RTF)). 117 | 118 | * Options 119 | * `--mode`: The name of a model to test. The logging folder named with the `[mode]` will be created as `[LOG_ROOT]/DMENet_CVPR2019/[mode]/`. Default: `DMENet_BDCS` 120 | * `--test_set`: The name of a dataset to evaluate. `CUHK` | `RTF0` | `RTF1` | `RTF1_6` | `random`. Default: `CUHK` 121 | * The folder structure can be modified in the function [`get_eval_path(..)`](https://github.com/codeslake/DMENet/blob/master/config.py#L85-L98) in `./config.py`. 122 | * `random` is for testing models with any images, which should be placed as `[DATASET_ROOT]/test/random/*.[jpg|png]`. 123 | 124 | * Check out [the evaluation code for the RTF dataset](https://github.com/codeslake/DMENet/tree/master/evaluation/RTF), and [the deconvolution code](https://github.com/codeslake/DMENet/tree/master/deconvolution). 125 | 126 | 127 | 128 | ## Training & testing the network 129 | 130 | * Train the network by: 131 | 132 | ```bash 133 | python main.py --is_train --mode [mode] 134 | ``` 135 | 136 | > **Note:** 137 | > 138 | > * If you train DMENet with newly generated SYNDOF dataset from [this repo](https://github.com/codeslake/SYNDOF), comment [this line](https://github.com/codeslake/DMENet/blob/master/utils.py#L43) and uncomment [this line](https://github.com/codeslake/DMENet/blob/master/utils.py#L49) before the training. 139 | 140 | * Test the network by: 141 | 142 | ```bash 143 | python main.py --mode [mode] --test_set [test_set] 144 | ``` 145 | 146 | * arguments 147 | * `--mode`: The name of a model to train. The logging folder named with the `[mode]` will be created as `[LOG_ROOT]/DMENet_CVPR2019/[mode]/`. Default: `DMENet_BDCS` 148 | * `--is_pretrain`: Pretrain the network with the MSE loss (`True` | `False`). Default: `False` 149 | * `--delete_log`: Deletes `[LOG_ROOT]/DMENet_CVPR2019/[mode]/*` before training begins (`True` | `False`). Default: `False` 150 | 151 | 152 | ## Contact 153 | Open an issue for any inquiries. 154 | You may also have contact with [junyonglee@postech.ac.kr](mailto:junyonglee@postech.ac.kr) 155 | 156 | ## Related Links 157 | * CVPR 2021: Iterative Filter Adaptive Network for Single Image Defocus Deblurring \[[paper](https://openaccess.thecvf.com/content/CVPR2021/papers/Lee_Iterative_Filter_Adaptive_Network_for_Single_Image_Defocus_Deblurring_CVPR_2021_paper.pdf)\]\[[code](https://github.com/codeslake/IFAN)\] 158 | * ICCV 2021: Single Image Defocus Deblurring Using Kernel-Sharing Parallel Atrous Convolutions \[[paper](https://arxiv.org/pdf/2108.09108.pdf)\]\[[code](https://github.com/HyeongseokSon1/KPAC)\] 159 | * SYNDOF dataset generation repo \[[link](https://github.com/codeslake/SYNDOF)\] 160 | 161 | ## License 162 | 
163 | This software is being made available under the terms in the [LICENSE](LICENSE) file. 164 | Any exemptions to these terms require a license from the Pohang University of Science and Technology. 165 | 166 | ## Citation 167 | If you find this code useful, please consider citing: 168 | 169 | ``` 170 | @InProceedings{Lee2019DMENet, 171 | author = {Junyong Lee and Sungkil Lee and Sunghyun Cho and Seungyong Lee}, 172 | title = {Deep Defocus Map Estimation Using Domain Adaptation}, 173 | booktitle = {Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR)}, 174 | year = {2019} 175 | } 176 | ``` 177 | 178 | -------------------------------------------------------------------------------- /assets/figure.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/codeslake/DMENet/d844e5a6ad3e7a1c9157d50935de1a6eb6bc4bf8/assets/figure.png -------------------------------------------------------------------------------- /assets/teaser_DMENet.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/codeslake/DMENet/d844e5a6ad3e7a1c9157d50935de1a6eb6bc4bf8/assets/teaser_DMENet.gif -------------------------------------------------------------------------------- /config.py: -------------------------------------------------------------------------------- 1 | from easydict import EasyDict as edict 2 | import json 3 | import os 4 | 5 | config = edict() 6 | config.TRAIN = edict() 7 | config.TEST = edict() 8 | 9 | ## Adam 10 | config.TRAIN.batch_size = 3 11 | config.TRAIN.batch_size_init = 8 12 | config.TRAIN.lr_init = 1e-4 13 | config.TRAIN.lr_init_init = 1e-4 14 | config.TRAIN.beta1 = 0.9 15 | 16 | # learning rate 17 | config.TRAIN.n_epoch = 10000 18 | config.TRAIN.n_epoch_init = 10 19 | config.TRAIN.lr_decay = 0.8 20 | config.TRAIN.decay_every = 20 21 | 22 | ## adversarial loss coefficient 23 | config.TRAIN.lambda_adv = 1e-3 24 | 25 | ## discriminator lr coefficient 26 | config.TRAIN.lambda_lr_d = 1 27 | 28 | ## binary loss coefficient 29 | config.TRAIN.lambda_binary = 2e-2 30 | 31 | ## perceptual loss coefficient 32 | config.TRAIN.lambda_perceptual = 1e-4 33 | 34 | ### TRAIN DATSET PATH 35 | data_offset = './datasets/' 36 | #data_offset = '/data1/junyonglee/defocus_map_estimation/DMENet/' 37 | config.TRAIN.synthetic_img_path = os.path.join(data_offset, 'train/SYNDOF/image/') 38 | config.TRAIN.defocus_map_path = os.path.join(data_offset, 'train/SYNDOF/blur_map/') 39 | config.TRAIN.defocus_map_norm_path = os.path.join(data_offset, 'train/SYNDOF/blur_map_norm/') 40 | config.TRAIN.synthetic_binary_map_path = os.path.join(data_offset, 'train/SYNDOF/blur_map_binary/') 41 | # Real 42 | config.TRAIN.real_img_path = os.path.join(data_offset, 'train/CUHK/image/') 43 | config.TRAIN.real_binary_map_path = os.path.join(data_offset, 'train/CUHK/gt/') 44 | config.TRAIN.real_img_no_label_path = os.path.join(data_offset, 'train/Flickr/') 45 | 46 | ### TEST DATSET PATH 47 | config.TEST.cuhk_img_path = os.path.join(data_offset, 'test/CUHK/image/') 48 | config.TEST.cuhk_binary_map_path = os.path.join(data_offset, 'test/CUHK/gt/') 49 | config.TEST.SYNDOF_img_path = os.path.join(data_offset, 'test/SYNDOF/image/') 50 | config.TEST.SYNDOF_gt_map_path = os.path.join(data_offset, 'test/SYNDOF/gt/') 51 | config.TEST.RTF0_img_path = os.path.join(data_offset, 'test/RTF/image/0/') 52 | config.TEST.RTF0_gt_map_path = os.path.join(data_offset, 'test/RTF/gt/0/') 53 | config.TEST.RTF1_img_path = os.path.join(data_offset, 'test/RTF/image/1/') 54 | config.TEST.RTF1_gt_map_path = os.path.join(data_offset, 'test/RTF/gt/1/') 55 | config.TEST.RTF1_6_img_path = os.path.join(data_offset, 'test/RTF/image/1.6/') 56 | config.TEST.RTF1_6_gt_map_path = os.path.join(data_offset, 'test/RTF/gt/1.6/') 57 | config.TEST.random_img_path = os.path.join(data_offset, 'test/random/') 58 | 59 | 60 | ## train image size 61 | config.TRAIN.height = 240 62 | config.TRAIN.width = 240 63 | 64 | ## log & checkpoint & samples 65 | # every global step 66 | config.TRAIN.write_log_every = 100 67 | config.TRAIN.write_ckpt_every = 1 68 | config.TRAIN.write_sample_every = 1000 69 | # every epoch 70 | config.TRAIN.refresh_image_log_every = 20 71 | 72 | # save dir 73 | config.root_offset = './logs/' 74 | #config.root_offset = '/Jarvis/logs/junyonglee' 75 | config.TRAIN.root_dir = os.path.join(config.root_offset, 'DMENet_CVPR2019/') 76 | 77 | config.TRAIN.max_coc = 15.; 78 | 79 | def log_config(path, cfg): 80 | with open(path + '/config.txt', 'w') as f: 81 | f.write('================================================\n') 82 | f.write(json.dumps(cfg, indent=4)) 83 | f.write('\n================================================\n') 84 | 85 | def get_eval_path(test_set, cfg): 86 | if test_set == 'CUHK': 87 | # return cfg.TEST.cuhk_img_path, cfg.TEST.cuhk_binary_map_path 88 | return cfg.TEST.cuhk_img_path, None 89 | elif test_set == 'SYNDOF': 90 | return cfg.TEST.SYNDOF_img_path, cfg.TEST.SYNDOF_gt_map_path 91 | elif test_set == 'RTF0': 92 | return cfg.TEST.RTF0_img_path, cfg.TEST.RTF0_gt_map_path 93 | elif test_set == 'RTF1': 94 | return cfg.TEST.RTF1_img_path, cfg.TEST.RTF1_gt_map_path 95 | elif test_set == 'RTF1_6': 96 | return cfg.TEST.RTF1_6_img_path, cfg.TEST.RTF1_6_gt_map_path 97 | elif test_set == 'random': 98 | return cfg.TEST.random_img_path, None 99 | -------------------------------------------------------------------------------- /deconvolution/DMENet_fast_deconv.m: -------------------------------------------------------------------------------- 1 | function [deconved, est_time] = fast_deconvolution(image, defocus_map, lambda, is_gpu) 2 | % Test the fast deconvolution method presented in the paper 3 | % D. Krishnan, R. Fergus: "Fast Image Deconvolution using Hyper-Laplacian 4 | % Priors", Proceedings of NIPS 2009. 5 | 6 | %% parameter values; other values such as the continuation regime of the parameter beta should be changed in fast_deconv.m 7 | alpha = 1.5; 8 | 9 | 10 | %% my params 11 | g_size = 101; 12 | 13 | %% deconv start 14 | 15 | unique_sigma = unique(defocus_map); 16 | 17 | g_center = (g_size + 1) / 2.; 18 | 19 | output = zeros(size(image)); 20 | if is_gpu 21 | output = gpuArray(output); 22 | end 23 | tic; 24 | parfor (c = 1:3) 25 | % for c = 1:3 26 | 27 | image_temp = image(:, :, c); 28 | output_temp = ones(size(image_temp)); 29 | if is_gpu 30 | output_temp = gpuArray(output_temp); 31 | end 32 | 33 | for j = 1:length(unique_sigma) 34 | % disp(sprintf('I[%02d/%02d], C[%d/%d], U[%03d/%03d]', i, length(image_file_paths), c, 3, j, length(unique_sigma))); 35 | s = unique_sigma(j); 36 | sigma = s; 37 | if s ~= 0 38 | G = fspecial('gaussian',[g_size, g_size], sigma); 39 | else 40 | G = zeros(g_size, g_size); 41 | G(g_center, g_center) = 1.0; 42 | end 43 | is_identity = G(g_center - 1:g_center+1, g_center - 1:g_center+1); 44 | is_identity(2, 2) = 0; 45 | is_identity = sum(is_identity(:)) == 0; 46 | 47 | if is_identity == false 48 | G_idx = find(G > 0); 49 | [y_G,x_G] = ind2sub(size(G),G_idx); 50 | kernel = G(min(y_G):max(y_G), min(x_G):max(x_G)); 51 | ks = floor((size(kernel, 1) - 1)/2); 52 | 53 | image_pad = padarray(image_temp, [1 1]*ks, 'replicate', 'both'); 54 | 55 | logical_G = logical(G > 0); 56 | if sum(logical_G(:)) > 1 57 | for a=1:4 58 | image_pad = edgetaper(image_pad, kernel); 59 | end 60 | end 61 | if is_gpu 62 | image_pad = gpuArray(image_pad); 63 | end 64 | 65 | % fast_deconv 66 | [x] = fast_deconv(image_pad, kernel, lambda, alpha); 67 | 68 | x = x(ks+1:end-ks, ks+1:end-ks); 69 | else 70 | x = image_temp; 71 | end 72 | 73 | output_temp(logical(defocus_map == s)) = x(logical(defocus_map == s)); 74 | end 75 | output(:, :, c) = output_temp; 76 | end 77 | est_time = toc(); 78 | 79 | output(logical(output > 1)) = 1; 80 | output(logical(output < 0)) = 0; 81 | if is_gpu 82 | output = gather(output); 83 | end 84 | deconved = output; 85 | 86 | -------------------------------------------------------------------------------- /deconvolution/dir2.m: -------------------------------------------------------------------------------- 1 | function full_path = dir2(varargin) 2 | if nargin == 0 3 | name = '.'; 4 | elseif nargin == 1 5 | name = varargin{1}; 6 | else 7 | error('Too many input arguments.') 8 | end 9 | 10 | listing = dir(name); 11 | 12 | inds = []; 13 | n = 0; 14 | k = 1; 15 | 16 | while k <= length(listing) 17 | if listing(k).isdir 18 | inds(end + 1) = k; 19 | end 20 | k = k + 1; 21 | end 22 | listing(inds) = []; 23 | 24 | full_path = []; 25 | for k = 1:length(listing) 26 | file_path = listing(k).folder; 27 | file_name = listing(k).name; 28 | full_path = [full_path, string(fullfile(file_path, file_name))]; 29 | end 30 | full_path = sort(full_path); 31 | 32 | end 33 | -------------------------------------------------------------------------------- /deconvolution/fast_deconv.m: -------------------------------------------------------------------------------- 1 | function [yout] = fast_deconv(yin, k, lambda, alpha, yout0) 2 | % 3 | % 4 | % fast_deconv solves the deconvolution problem in the paper (see Equation (1)) 5 | % D. Krishnan, R. Fergus: "Fast Image Deconvolution using Hyper-Laplacian 6 | % Priors", Proceedings of NIPS 2009. 7 | % 8 | % This paper and the code are related to the work and code of Wang 9 | % et. al.: 10 | % 11 | % Y. Wang, J. Yang, W. Yin and Y. Zhang, "A New Alternating Minimization 12 | % Algorithm for Total Variation Image Reconstruction", SIAM Journal on 13 | % Imaging Sciences, 1(3): 248:272, 2008. 14 | % and their FTVd code. 15 | 16 | % Input Parameters: 17 | % 18 | % yin: Observed blurry and noisy input grayscale image. 19 | % k: convolution kernel 20 | % lambda: parameter that balances likelihood and prior term weighting 21 | % alpha: parameter between 0 and 2 22 | % yout0: if this is passed in, it is used as an initialization for the 23 | % output deblurred image; if not passed in, then the input blurry image 24 | % is used as the initialization 25 | % 26 | % 27 | % Outputs: 28 | % yout: solution 29 | % 30 | % Note: for faster LUT interpolation, please download and install 31 | % matlabPyrTools of Eero Simoncelli from 32 | % www.cns.nyu.edu/~lcv/software.php. The specific MeX function required 33 | % is pointOp (used in solve_image.m). 34 | % 35 | % Copyright (C) 2009. Dilip Krishnan and Rob Fergus 36 | % Email: dilip,fergus@cs.nyu.edu 37 | 38 | % continuation parameters 39 | beta = 1; 40 | beta_rate = 2*sqrt(2); 41 | beta_max = 2^8; 42 | 43 | % number of inner iterations per outer iteration 44 | mit_inn = 1; 45 | 46 | [m n] = size(yin); 47 | % initialize with input or passed in initialization 48 | if (nargin == 5) 49 | yout = yout0; 50 | else 51 | yout = yin; 52 | end; 53 | 54 | % make sure k is a odd-sized 55 | if ((mod(size(k, 1), 2) ~= 1) | (mod(size(k, 2), 2) ~= 1)) 56 | fprintf('Error - blur kernel k must be odd-sized.\n'); 57 | return; 58 | end; 59 | ks = floor((size(k, 1)-1)/2); 60 | 61 | % compute constant quantities 62 | % see Eqn. (3) of paper 63 | [Nomin1, Denom1, Denom2] = computeDenominator(yin, k); 64 | 65 | % x and y gradients of yout (with circular boundary conditions) 66 | % other gradient filters may be used here and their transpose will then need to 67 | % be used within the inner loop (see comment below) and in the function 68 | % computeDenominator 69 | youtx = [diff(yout, 1, 2), yout(:,1) - yout(:,n)]; 70 | youty = [diff(yout, 1, 1); yout(1,:) - yout(m,:)]; 71 | 72 | % store some of the statistics 73 | costfun = []; 74 | Outiter = 0; 75 | 76 | %% Main loop 77 | while beta < beta_max 78 | Outiter = Outiter + 1; 79 | %fprintf('Outer iteration %d; beta %.3g\n',Outiter, beta); 80 | 81 | gamma = beta/lambda; 82 | Denom = Denom1 + gamma*Denom2; 83 | Inniter = 0; 84 | 85 | for Inniter = 1:mit_inn 86 | 87 | if (0) 88 | %%% Compute cost function - uncomment to see the original 89 | % minimization function costs at every iteration 90 | youtk = conv2(yout, k, 'same'); 91 | % likelihood term 92 | lh = sum(sum((youtk - yin).^2 )); 93 | 94 | if (alpha == 1) 95 | cost = (lambda/2)*lh + sum(abs(youtx(:))) + sum(abs(youty(:))); 96 | else 97 | cost = (lambda/2)*lh + sum(abs(youtx(:)).^alpha) + sum(abs(youty(:)).^alpha); 98 | end; 99 | %fprintf('Inniter iteration %d; cost %.3g\n', Inniter, cost); 100 | 101 | costfun = [costfun, cost]; 102 | end; 103 | % 104 | % w-subproblem: eqn (5) of paper 105 | % 106 | Wx = solve_image(youtx, beta, alpha); 107 | Wy = solve_image(youty, beta, alpha); 108 | 109 | % 110 | % x-subproblem: eqn (3) of paper 111 | % 112 | % The transpose of x and y gradients; if other gradient filters 113 | % (such as higher-order filters) are to be used, then add them 114 | % below the comment above as well 115 | 116 | Wxx = [Wx(:,n) - Wx(:, 1), -diff(Wx,1,2)]; 117 | Wxx = Wxx + [Wy(m,:) - Wy(1, :); -diff(Wy,1,1)]; 118 | 119 | Fyout = (Nomin1 + gamma*fft2(Wxx))./Denom; 120 | yout = real(ifft2(Fyout)); 121 | 122 | % update the gradient terms with new solution 123 | youtx = [diff(yout, 1, 2), yout(:,1) - yout(:,n)]; 124 | youty = [diff(yout, 1, 1); yout(1,:) - yout(m,:)]; 125 | 126 | end %inner 127 | beta = beta*beta_rate; 128 | end %Outer 129 | 130 | 131 | %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% 132 | function [Nomin1, Denom1, Denom2] = computeDenominator(y, k) 133 | % 134 | % computes denominator and part of the numerator for Equation (3) of the 135 | % paper 136 | % 137 | % Inputs: 138 | % y: blurry and noisy input 139 | % k: convolution kernel 140 | % 141 | % Outputs: 142 | % Nomin1 -- F(K)'*F(y) 143 | % Denom1 -- |F(K)|.^2 144 | % Denom2 -- |F(D^1)|.^2 + |F(D^2)|.^2 145 | % 146 | 147 | sizey = size(y); 148 | otfk = psf2otf(k, sizey); 149 | Nomin1 = conj(otfk).*fft2(y); 150 | Denom1 = abs(otfk).^2; 151 | % if higher-order filters are used, they must be added here too 152 | Denom2 = abs(psf2otf([1,-1],sizey)).^2 + abs(psf2otf([1;-1],sizey)).^2; 153 | -------------------------------------------------------------------------------- /deconvolution/readme.md: -------------------------------------------------------------------------------- 1 | # Deconvolution using defocus map estimated from DMENet 2 | * *This code is based on "[Fast Image Deconvolution using Hyper-Laplacian Priors](https://papers.nips.cc/paper/2009/file/3dd48ab31d016ffcbf3314df2b3cb9ce-Paper.pdf)", Krishnan *et al.*, In Proc. NIPS 2009.
Refer [here](https://dilipkay.wordpress.com/fast-deconvolution/) for the original code.* 3 | 4 | ## Getting Started 5 | 1. Place your dataset as: 6 | 7 | ``` 8 | ... 9 | ├── deconvolution 10 | │ ├── source 11 | │ │ ├── input 12 | │ │ │ ├── [DATASET] # the name of a dataset [`CUHK` | `DPDD` | `RealDOF`] 13 | │ │ │ │ ├── *.[jpg|png] # input images 14 | │ │ ├── defocus map 15 | │ │ │ ├── [DATASET] # the name of the dataset that DMENet ran on 16 | │ │ │ │ ├── *.[jpg|png] # defocus maps (results of DMENet in `[LOG_ROOT]/[mode]/result/[test_set]/defocus_map`) 17 | ... 18 | ``` 19 | 20 | > **Note:** 21 | > 22 | > * For the DPDD dataset, refer [here](https://www.eecs.yorku.ca/~abuolaim/eccv_2020_dp_defocus_deblurring/dataset.html). 23 | > * The RealDOF test set is the test set that we used for the defocus deblurring paper, which is provisionally accepted to CVPR2021. We will release the test set soon. 24 | 25 | 2. In the matlab console, type following for the evaluation. 26 | 27 | ``` 28 | >> run_DMENet_deconv([DATASET], is_gpu, gpu_num) 29 | 30 | # gpu example 31 | >> run_DMENet_deconv('CUHK', true, 1) 32 | 33 | # cpu example 34 | >> run_DMENet_deconv('CUHK') 35 | ``` 36 | 37 | * Parameters 38 | * [DATASET]: the name of a dataset. `CUHK` | `DPDD` | `RealDOF`. 39 | * is_gpu: whether to use gpu. `true` | `false`. Default: `false` 40 | * gpu_num: the device number of a gpu. 41 | 42 | * Results will be saved as: 43 | 44 | ``` 45 | ... 46 | ├── deconvolution 47 | ... 48 | │ ├── output 49 | │ │ ├── [DATASET] # the name of the dataset used for deconvolution 50 | │ │ │ ├── *.[jpg|png] # resulting deconvolution images 51 | ``` 52 | -------------------------------------------------------------------------------- /deconvolution/run_DMENet_deconv.m: -------------------------------------------------------------------------------- 1 | function run_DMENet_deconv(dataset, is_gpu, gpu_num) 2 | % dataset = CUHK | DPDD | RealDOF 3 | if nargin == 1 4 | is_gpu = false; 5 | gpu_num = 0; 6 | elseif nargin == 2 7 | gpu_num = 1; 8 | end 9 | 10 | g = 0; 11 | if is_gpu 12 | g = gpuDevice(gpu_num); 13 | end 14 | 15 | %% directory 16 | offset = './sources'; 17 | image_file_paths = dir2(fullfile(offset, 'input', dataset)); 18 | defocus_file_paths = dir2(fullfile(offset, 'defocus_map', dataset)); 19 | out_offset = fullfile('output', dataset); 20 | 21 | % remove output directory before begin 22 | if isdir(out_offset) 23 | rmdir(out_offset, 's'); 24 | end 25 | mkdir(out_offset) 26 | 27 | %% deconv parameter 28 | if contains(dataset, 'CUHK') 29 | % for the CUHK dataset 30 | lambda = 420; 31 | elseif contains(dataset, 'DPDD') 32 | % for the DPDD dataset 33 | lambda = 16.8; 34 | elseif contains(dataset, 'RealDOF') 35 | % for the RealDOF dataset 36 | lambda = 4.2; 37 | end 38 | 39 | % my parameter 40 | bin_num = 32; % bin_number 41 | 42 | %% deconv start 43 | est_time_mean = 0; 44 | for i = 1:length(image_file_paths) 45 | % read images 46 | input = read_img(image_file_paths(i)); 47 | 48 | %%%% Read defocus map and make it to sigma map 49 | %%% This is for defocus map result of DMENet. For the results of other 50 | %%% methods, modify them to have proper standard deviation value for 51 | %%% creating spatially varying Gaussian kernels. 52 | defocus_map = double(imread(char(defocus_file_paths(i))))./255.0; 53 | % DMENET 54 | defocus_map = (defocus_map * 15 - 1)/2; 55 | defocus_map(defocus_map < 0) = 0; 56 | % EBDB 57 | % defocus_map = defocus_map * 5; 58 | % JNB 59 | % defocus_map = defocus_map * 4; 60 | %%%% 61 | 62 | % make sure to have the same resolution 63 | [input, defocus_map] = refine_img(input, defocus_map); 64 | 65 | %%% quantize (for all compared methods, results are almost the same even without the qunatization when bin_num is more than 32) 66 | unique_sigma = unique(defocus_map); 67 | w_edges = quantile(unique_sigma, bin_num); % weighted edges 68 | w_edges = [min(unique_sigma), w_edges]; 69 | w_edges = [w_edges, max(unique_sigma)]; 70 | defocus_map_index = discretize(defocus_map, w_edges); 71 | defocus_map = w_edges(defocus_map_index); 72 | unique_sigma = unique(defocus_map); 73 | 74 | 75 | %%% deconvolution start 76 | [deconv_result, est_time] = DMENet_fast_deconv(input, defocus_map, lambda, is_gpu); 77 | if is_gpu 78 | reset(g); 79 | end 80 | %%% 81 | 82 | disp(sprintf('[%02d/%02d] (%.3f sec)', i, length(image_file_paths), est_time)); 83 | est_time_mean = est_time_mean + est_time; 84 | 85 | imwrite(uint8(deconv_result*255), fullfile(out_offset, sprintf('%02d.png', i))); 86 | end 87 | 88 | est_time_mean = est_time_mean / length(image_file_paths); 89 | disp(sprintf('Deconvolution done for %s dataset (%.3f sec)', dataset, est_time_mean)); 90 | end 91 | 92 | %% 93 | function image = read_img(path) 94 | image = imread(char(path)); 95 | image = im2double(image); 96 | image = double(uint8(image * 255)) / 255; 97 | end 98 | 99 | function [in1, in2] = refine_img(in1, in2) 100 | sz_in1 = size(in1); 101 | sz_in2 = size(in2); 102 | 103 | in1 = in1(1:min(sz_in1(1), sz_in2(1)), 1:min(sz_in1(2), sz_in2(2)), :); 104 | in2 = in2(1:min(sz_in1(1), sz_in2(1)), 1:min(sz_in1(2), sz_in2(2)), :); 105 | end 106 | -------------------------------------------------------------------------------- /deconvolution/solve_image.m: -------------------------------------------------------------------------------- 1 | function [w] = solve_image(v, beta, alpha) 2 | 3 | % 4 | % solve the following component-wise separable problem 5 | % min |w|^\alpha + \frac{\beta}{2} (w - v).^2 6 | % 7 | % A LUT is used to solve the problem; when the function is first called 8 | % for a new value of beta or alpha, a LUT is built for that beta/alpha 9 | % combination and for a range of values of v. The LUT stays persistent 10 | % between calls to solve_image. It will be recomputed the first time this 11 | % function is called. 12 | 13 | % range of input data and step size; increasing the range of decreasing 14 | % the step size will increase accuracy but also increase the size of the 15 | % LUT 16 | range = 10; 17 | step = 0.0001; 18 | 19 | persistent lookup_v known_beta xx known_alpha 20 | ind = find(known_beta==beta & known_alpha==alpha); 21 | if isempty(known_beta | known_alpha) 22 | xx = [-range:step:range]; 23 | end 24 | if any(ind) 25 | %fprintf('Reusing lookup table for beta %.3g and alpha %.3g\n', beta, alpha); 26 | %%% already computed 27 | if (exist('pointOp') == 3) 28 | % Use Eero Simoncelli's function to extrapolate 29 | w = pointOp(double(v),lookup_v(ind,:), -range, step, 0); 30 | else 31 | w = interp1(xx', lookup_v(ind,:)', v(:), 'linear', 'extrap'); 32 | w = reshape(w, size(v,1), size(v,2)); 33 | end; 34 | else 35 | %%% now go and recompute xx for new value of beta and alpha 36 | tmp = compute_w(xx, beta, alpha); 37 | lookup_v = [lookup_v; tmp(:)']; 38 | known_beta = [known_beta, beta]; 39 | known_alpha = [known_alpha, alpha]; 40 | 41 | %%% and lookup current v's in the new lookup table row. 42 | if (exist('pointOp') == 3) 43 | % Use Eero Simoncelli's function to extrapolate 44 | w = pointOp(double(v),lookup_v(end,:), -range, step, 0); 45 | else 46 | w = interp1(xx', lookup_v(end,:)', v(:), 'linear', 'extrap'); 47 | w = reshape(w, size(v,1), size(v,2)); 48 | end; 49 | 50 | %fprintf('Recomputing lookup table for new value of beta %.3g and alpha %.3g\n', beta, alpha); 51 | end 52 | 53 | %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% 54 | % 55 | % call different functions to solve the minimization problem 56 | % min |w|^\alpha + \frac{\beta}{2} (w - v).^2 for a fixed beta and alpha 57 | % 58 | function w = compute_w(v, beta, alpha) 59 | 60 | if (abs(alpha - 1) < 1e-9) 61 | % assume alpha = 1.0 62 | w = compute_w1(v, beta); 63 | return; 64 | end; 65 | 66 | if (abs(alpha - 2/3) < 1e-9) 67 | % assume alpha = 2/3 68 | w = compute_w23(v, beta); 69 | return; 70 | end; 71 | 72 | if (abs(alpha - 1/2) < 1e-9) 73 | % assume alpha = 1/2 74 | w = compute_w12(v, beta); 75 | return; 76 | end; 77 | 78 | % for any other value of alpha, plug in some other generic root-finder 79 | % here, we use Newton-Raphson 80 | w = newton_w(v, beta, alpha); 81 | 82 | %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% 83 | function w = compute_w23(v, beta) 84 | % solve a quartic equation 85 | % for alpha = 2/3 86 | 87 | epsilon = 1e-6; %% tolerance on imag part of real root 88 | 89 | k = 8/(27*beta^3); 90 | m = ones(size(v))*k; 91 | 92 | % Now use formula from 93 | % http://en.wikipedia.org/wiki/Quartic_equation (Ferrari's method) 94 | % running our coefficients through Mathmetica (quartic_solution.nb) 95 | % optimized to use as few operations as possible... 96 | 97 | %%% precompute certain terms 98 | v2 = v .* v; 99 | v3 = v2 .* v; 100 | v4 = v3 .* v; 101 | m2 = m .* m; 102 | m3 = m2 .* m; 103 | 104 | %% Compute alpha & beta 105 | alpha = -1.125*v2; 106 | beta2 = 0.25*v3; 107 | 108 | %%% Compute p,q,r and u directly. 109 | q = -0.125*(m.*v2); 110 | r1 = -q/2 + sqrt(-m3/27 + (m2.*v4)/256); 111 | 112 | u = exp(log(r1)/3); 113 | y = 2*(-5/18*alpha + u + (m./(3*u))); 114 | 115 | W = sqrt(alpha./3 + y); 116 | 117 | %%% now form all 4 roots 118 | root = zeros(size(v,1),size(v,2),4); 119 | root(:,:,1) = 0.75.*v + 0.5.*(W + sqrt(-(alpha + y + beta2./W ))); 120 | root(:,:,2) = 0.75.*v + 0.5.*(W - sqrt(-(alpha + y + beta2./W ))); 121 | root(:,:,3) = 0.75.*v + 0.5.*(-W + sqrt(-(alpha + y - beta2./W ))); 122 | root(:,:,4) = 0.75.*v + 0.5.*(-W - sqrt(-(alpha + y - beta2./W ))); 123 | 124 | 125 | %%%%%% Now pick the correct root, including zero option. 126 | 127 | %%% Clever fast approach that avoids lookups 128 | v2 = repmat(v,[1 1 4]); 129 | sv2 = sign(v2); 130 | rsv2 = real(root).*sv2; 131 | 132 | %%% condensed fast version 133 | %%% take out imaginary roots above v/2 but below v 134 | root_flag3 = sort(((abs(imag(root))
MSE / MAE | RTF
MSE / MAE | CUHK
acc / mAP | 22 | | :----: | :----: | :----: | :----: | 23 | | ... | ... | ... | ... | 24 | | DMENetBDCS | 0.013 / 0.084 | **0.009** / **0.077** | 0.878 / **0.987** | 25 | 26 | * Table 6 in the supplementary material (the last column) 27 | | Image # | ... | Ours | 28 | | :-------------: | :-------------: | :-------------: | 29 | | 01 | ... | **0.0643** | 30 | | 02 | ... | **0.0406** | 31 | | 03 | ... | **0.0863** | 32 | | 04 | ... | **0.0408** | 33 | | 05 | ... | **0.0335** | 34 | | 06 | ... | 0.0756 | 35 | | 07 | ... | 0.1129 | 36 | | 08 | ... | 0.1695 | 37 | | 09 | ... | **0.0427** | 38 | | 10 | ... | **0.0771** | 39 | | 11 | ... | **0.044** | 40 | | 12 | ... | 0.1347 | 41 | | 13 | ... | **0.0817** | 42 | | 14 | ... | **0.1123** | 43 | | 15 | ... | **0.0838** | 44 | | 16 | ... | **0.0881** | 45 | | 17 | ... | **0.0675** | 46 | | 18 | ... | **0.0786** | 47 | | 19 | ... | **0.0744** | 48 | | 20 | ... | **0.0841** | 49 | | 21 | ... | **0.0397** | 50 | | 22 | ... | **0.0535** | 51 | | Avg. MSE | ... | **0.0093** | 52 | | Avg. MAE | ... | **0.0767** | 53 | | Avg. MSE (s=1.0) | ... | **0.0207** | 54 | | Avg. MAE (s=1.0) | ... | **0.1006** | 55 | | Avg. MSE (s=1.6) | ... | 0.0491 | 56 | | Avg. MAE (s=1.6) | ... | 0.1579 | 57 | -------------------------------------------------------------------------------- /install_CUDA10.0.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | conda install -y cudatoolkit=10.0 3 | conda install -y cudnn=7.6 4 | pip install --no-cache tensorflow-gpu==1.15 5 | pip install --no-cache tensorlayer==1.11.1 6 | pip install --no-cache jupyterlab 7 | pip install --no-cache -r requirements.txt 8 | #apt-get install ffmpeg libsm6 libxext6 -y # cv2 error 9 | -------------------------------------------------------------------------------- /install_CUDA11.1.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | pip install --no-cache --user nvidia-pyindex 3 | conda install -y -c conda-forge openmpi 4 | export PATH="~/.local/bin:$PATH" 5 | export LD_LIBRARY_PATH="/anaconda3/envs/DMENet_CUDA11/lib/:$LD_LIBRARY_PATH" 6 | pip install --no-cache --user 'nvidia-tensorflow[horovod]' 7 | pip install --no-cache tensorlayer==1.11.1 8 | pip install --no-cache --upgrade numpy==1.16.0 9 | pip install --no-cache --upgrade warpt==1.11.1 10 | 11 | pip install --no-cache jupyterlab 12 | pip install --no-cache -r requirements.txt 13 | #apt-get install ffmpeg libsm6 libxext6 -y # cv2 error 14 | 15 | -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR) 3 | import tensorlayer as tl 4 | #tl.logging.set_verbosity(tl.logging.ERROR) 5 | 6 | import numpy as np 7 | from random import shuffle 8 | import matplotlib 9 | import datetime 10 | import time 11 | import shutil 12 | import os 13 | 14 | from config import config, log_config, get_eval_path 15 | from utils import * 16 | from model import * 17 | 18 | 19 | 20 | batch_size = config.TRAIN.batch_size 21 | batch_size_init = config.TRAIN.batch_size_init 22 | lr_init = config.TRAIN.lr_init 23 | lr_init_init = config.TRAIN.lr_init_init 24 | beta1 = config.TRAIN.beta1 25 | 26 | n_epoch = config.TRAIN.n_epoch 27 | n_epoch_init = config.TRAIN.n_epoch_init 28 | lr_decay = config.TRAIN.lr_decay 29 | decay_every = config.TRAIN.decay_every 30 | 31 | lambda_adv = config.TRAIN.lambda_adv 32 | lambda_lr_d = config.TRAIN.lambda_lr_d 33 | lambda_binary = config.TRAIN.lambda_binary 34 | lambda_perceptual = config.TRAIN.lambda_perceptual 35 | 36 | h = config.TRAIN.height 37 | w = config.TRAIN.width 38 | 39 | ni = int(np.ceil(np.sqrt(batch_size))) 40 | 41 | def train(): 42 | ## CREATE DIRECTORIES 43 | mode_dir = config.TRAIN.root_dir + '{}'.format(tl.global_flag['mode']) 44 | 45 | ckpt_dir = os.path.join(mode_dir, 'checkpoint') 46 | init_dir = os.path.join(mode_dir, 'init') 47 | log_dir_scalar_init = os.path.join(mode_dir, 'log/scalar_init') 48 | log_dir_image_init = os.path.join(mode_dir, 'log/image_init') 49 | log_dir_scalar = os.path.join(mode_dir, 'log/scalar') 50 | log_dir_image = os.path.join(mode_dir, 'log/image') 51 | sample_dir = os.path.join(mode_dir, 'samples/train') 52 | config_dir = os.path.join(mode_dir, 'config') 53 | 54 | if tl.global_flag['delete_log']: 55 | shutil.rmtree(ckpt_dir, ignore_errors = True) 56 | if tl.global_flag['is_pretrain']: 57 | shutil.rmtree(log_dir_scalar_init, ignore_errors = True) 58 | shutil.rmtree(log_dir_image_init, ignore_errors = True) 59 | shutil.rmtree(log_dir_scalar, ignore_errors = True) 60 | shutil.rmtree(log_dir_image, ignore_errors = True) 61 | shutil.rmtree(sample_dir, ignore_errors = True) 62 | shutil.rmtree(config_dir, ignore_errors = True) 63 | 64 | tl.files.exists_or_mkdir(ckpt_dir) 65 | tl.files.exists_or_mkdir(init_dir) 66 | tl.files.exists_or_mkdir(log_dir_scalar_init) 67 | tl.files.exists_or_mkdir(log_dir_image_init) 68 | tl.files.exists_or_mkdir(log_dir_scalar) 69 | tl.files.exists_or_mkdir(log_dir_image) 70 | tl.files.exists_or_mkdir(sample_dir) 71 | tl.files.exists_or_mkdir(config_dir) 72 | log_config(config_dir, config) 73 | 74 | ## DEFINE SESSION 75 | sess = tf.Session(config = tf.ConfigProto(allow_soft_placement = True, log_device_placement = False)) 76 | 77 | ## READ DATASET LIST 78 | # train 79 | train_real_img_list = np.array(sorted(tl.files.load_file_list(path = config.TRAIN.real_img_path, regx = '.*', printable = False))) 80 | train_real_binary_map_list = np.array(sorted(tl.files.load_file_list(path = config.TRAIN.real_binary_map_path, regx = '.*', printable = False))) 81 | train_real_img_no_label_list = np.array(sorted(tl.files.load_file_list(path = config.TRAIN.real_img_no_label_path, regx = '.*', printable = False))) 82 | 83 | # test 84 | test_blur_img_list = np.array(sorted(tl.files.load_file_list(path = config.TEST.cuhk_img_path, regx = '.*', printable = False))) 85 | test_gt_list = np.array(sorted(tl.files.load_file_list(path = config.TEST.cuhk_binary_map_path, regx = '.*', printable = False))) 86 | test_blur_imgs = read_all_imgs(test_blur_img_list, path = config.TEST.cuhk_img_path, mode = 'RGB') 87 | test_gt_imgs = read_all_imgs(test_gt_list, path = config.TEST.cuhk_binary_map_path, mode = 'GRAY') 88 | 89 | ## DEFINE MODEL 90 | # input 91 | with tf.variable_scope('input'): 92 | patches_synthetic = tf.placeholder('float32', [None, h, w, 3], name = 'input_synthetic') 93 | labels_synthetic_defocus = tf.placeholder('float32', [None, h, w, 1], name = 'labels_synthetic_defocus') 94 | 95 | patches_real = tf.placeholder('float32', [None, h, w, 3], name = 'input_real') 96 | labels_real_binary = tf.placeholder('float32', [None, h, w, 1], name = 'labels_real_binary') 97 | patches_real_no_label = tf.placeholder('float32', [None, h, w, 3], name = 'input_real_no_label') 98 | 99 | patches_real_test = tf.placeholder('float32', [None, h, w, 3], name = 'input_real') 100 | labels_real_binary_test = tf.placeholder('float32', [None, h, w, 1], name = 'labels_real_binary') 101 | 102 | # model 103 | with tf.variable_scope('main_net') as scope: 104 | with tf.variable_scope('defocus_net') as scope: 105 | with tf.variable_scope('encoder') as scope: 106 | net_vgg, feats_synthetic_down, _, _ = VGG19_down(patches_synthetic, reuse = False, scope = scope) 107 | _, feats_real_down, _, _ = VGG19_down(patches_real, reuse = True, scope = scope) 108 | _, feats_real_no_label_down, _, _ = VGG19_down(patches_real_no_label, reuse = True, scope = scope) 109 | with tf.variable_scope('decoder') as scope: 110 | output_synthetic_defocus, feats_synthetic_up_aux, feats_synthetic_da, _ = UNet_up(patches_synthetic, feats_synthetic_down, is_train = True, reuse = False, scope = scope) 111 | output_real_defocus, _, feats_real_da, _ = UNet_up(patches_real, feats_real_down, is_train = True, reuse = True, scope = scope) 112 | output_real_no_label_defocus, _, feats_real_no_label_da, _ = UNet_up(patches_real_no_label, feats_real_no_label_down, is_train = True, reuse = True, scope = scope) 113 | with tf.variable_scope('binary_net') as scope: 114 | output_real_binary_logits, output_real_binary = Binary_Net(output_real_defocus, is_train = True, reuse = False, scope = scope) 115 | 116 | with tf.variable_scope('discriminator') as scope: 117 | with tf.variable_scope('feature') as scope: 118 | d_feature_logits_synthetic, d_feature_synthetic = feature_discriminator(feats_synthetic_da, is_train = True, reuse = False, scope = scope) 119 | d_feature_logits_real, d_feature_real = feature_discriminator(feats_real_da, is_train = True, reuse = True, scope = scope) 120 | d_feature_logits_real_no_label, d_feature_real_no_label = feature_discriminator(feats_real_no_label_da, is_train = True, reuse = True, scope = scope) 121 | 122 | # fixed 123 | with tf.variable_scope('perceptual') as scope: 124 | output_synthetic_defocus_3c = tf.concat([output_synthetic_defocus, output_synthetic_defocus, output_synthetic_defocus], axis = 3) 125 | net_vgg_perceptual, _, perceptual_synthetic_out, logits_perceptual_out = VGG19_down(output_synthetic_defocus_3c, reuse = False, scope = scope) 126 | labels_synthetic_defocus_3c = tf.concat([labels_synthetic_defocus, labels_synthetic_defocus, labels_synthetic_defocus], axis = 3) 127 | _, _, perceptual_synthetic_label, logits_perceptual_label = VGG19_down(labels_synthetic_defocus_3c, reuse = True, scope = scope) 128 | 129 | # for test 130 | with tf.variable_scope('main_net') as scope: 131 | with tf.variable_scope('defocus_net') as scope: 132 | with tf.variable_scope('encoder') as scope: 133 | _, feats_real_down_test, _, _ = VGG19_down(patches_real_test, reuse = True, scope = scope) 134 | with tf.variable_scope('decoder') as scope: 135 | output_real_defocus_test, _, _, _ = UNet_up(patches_real, feats_real_down_test, is_train = True, reuse = True, scope = scope) 136 | 137 | ## DEFINE LOSS 138 | with tf.variable_scope('loss'): 139 | with tf.variable_scope('discriminator'): 140 | with tf.variable_scope('feature'): 141 | loss_synthetic_d_feature = tl.cost.sigmoid_cross_entropy(d_feature_logits_synthetic, tf.ones_like(d_feature_logits_synthetic), name = 'synthetic') 142 | loss_real_d_feature = tl.cost.sigmoid_cross_entropy(d_feature_logits_real, tf.zeros_like(d_feature_logits_real), name = 'real') 143 | loss_real_no_label_d_feature = tl.cost.sigmoid_cross_entropy(d_feature_logits_real_no_label, tf.zeros_like(d_feature_logits_real_no_label), name = 'real') 144 | loss_d_feature = tf.identity((2 * loss_synthetic_d_feature + loss_real_d_feature + loss_real_no_label_d_feature) / 2. * lambda_adv, name = 'total') 145 | 146 | loss_d = tf.identity(loss_d_feature, name = 'total') 147 | 148 | with tf.variable_scope('generator'): 149 | with tf.variable_scope('feature'): 150 | loss_real_g_feature = tl.cost.sigmoid_cross_entropy(d_feature_logits_real, tf.ones_like(d_feature_logits_real), name = 'real') 151 | loss_real_no_label_g_feature = tl.cost.sigmoid_cross_entropy(d_feature_logits_real_no_label, tf.ones_like(d_feature_logits_real_no_label), name = 'real') 152 | loss_g_feature = tf.identity((loss_real_g_feature + loss_real_no_label_g_feature) / 2., name = 'total') 153 | 154 | loss_g = tf.identity(loss_g_feature * lambda_adv, name = 'total') 155 | 156 | with tf.variable_scope('defocus'): 157 | loss_defocus = tl.cost.mean_squared_error(output_synthetic_defocus, labels_synthetic_defocus, is_mean = True, name = 'synthetic') 158 | 159 | with tf.variable_scope('auxilary'): 160 | labels_layer = InputLayer(labels_synthetic_defocus) 161 | loss_aux_1 = tl.cost.mean_squared_error(feats_synthetic_up_aux[0], 162 | DownSampling2dLayer(labels_layer, (int(h / 16), int(w / 16)), is_scale = False, method = 1, align_corners=True).outputs, is_mean = True, name = 'aux1') 163 | loss_aux_2 = tl.cost.mean_squared_error(feats_synthetic_up_aux[1], 164 | DownSampling2dLayer(labels_layer, (int(h / 8), int(w / 8)), is_scale = False, method = 1, align_corners=True).outputs, is_mean = True, name = 'aux2') 165 | loss_aux_3 = tl.cost.mean_squared_error(feats_synthetic_up_aux[2], 166 | DownSampling2dLayer(labels_layer, (int(h / 4), int(w / 4)), is_scale = False, method = 1, align_corners=True).outputs, is_mean = True, name = 'aux3') 167 | loss_aux_4 = tl.cost.mean_squared_error(feats_synthetic_up_aux[3], 168 | DownSampling2dLayer(labels_layer, (int(h / 2), int(w / 2)), is_scale = False, method = 1, align_corners=True).outputs, is_mean = True, name = 'aux4') 169 | loss_aux_5 = tl.cost.mean_squared_error(feats_synthetic_up_aux[4], labels_synthetic_defocus, is_mean = True, name = 'aux5') 170 | loss_aux = tf.identity(loss_aux_1 + loss_aux_2 + loss_aux_3 + loss_aux_4 + loss_aux_5, name = 'total') 171 | 172 | with tf.variable_scope('perceptual'): 173 | loss_synthetic_perceptual = tl.cost.mean_squared_error(perceptual_synthetic_out, perceptual_synthetic_label, is_mean = True, name = 'synthetic') 174 | loss_perceptual = tf.identity(loss_synthetic_perceptual * lambda_perceptual, name = 'total') 175 | 176 | with tf.variable_scope('binary'): 177 | loss_real_binary = tl.cost.sigmoid_cross_entropy(output_real_binary_logits, labels_real_binary, name = 'real') 178 | loss_binary = tf.identity(loss_real_binary * lambda_binary, name = 'total') 179 | 180 | loss_main = tf.identity(loss_defocus + loss_binary + loss_perceptual + loss_aux + loss_g, name = 'total') 181 | loss_init = tf.identity(loss_defocus, name = 'loss_init') 182 | 183 | ## DEFINE OPTIMIZER 184 | # variables to save / train 185 | d_vars = tl.layers.get_variables_with_name('discriminator', True, False) 186 | main_vars = tl.layers.get_variables_with_name('main_net', True, False) 187 | save_vars = tl.layers.get_variables_with_name('main_net', False, False) + tl.layers.get_variables_with_name('discriminator', False, False) 188 | 189 | init_vars = tl.layers.get_variables_with_name('defocus_net', True, False) 190 | save_init_vars = tl.layers.get_variables_with_name('defocus_net', False, False) 191 | 192 | # define optimizer 193 | with tf.variable_scope('Optimizer'): 194 | learning_rate = tf.Variable(lr_init, trainable = False) 195 | learning_rate_init = tf.Variable(lr_init_init, trainable = False) 196 | optim_d = tf.train.AdamOptimizer(learning_rate * lambda_lr_d, beta1 = beta1).minimize(loss_d, var_list = d_vars) 197 | optim_main = tf.train.AdamOptimizer(learning_rate, beta1 = beta1).minimize(loss_main, var_list = main_vars) 198 | optim_init = tf.train.AdamOptimizer(learning_rate_init, beta1 = beta1).minimize(loss_init, var_list = init_vars) 199 | 200 | ## DEFINE SUMMARY 201 | # writer 202 | writer_scalar = tf.summary.FileWriter(log_dir_scalar, sess.graph, flush_secs=30, filename_suffix = '.loss_log') 203 | writer_image = tf.summary.FileWriter(log_dir_image, sess.graph, flush_secs=30, filename_suffix = '.image_log') 204 | if tl.global_flag['is_pretrain']: 205 | writer_scalar_init = tf.summary.FileWriter(log_dir_scalar_init, sess.graph, flush_secs=30, filename_suffix = '.loss_log_init') 206 | writer_image_init = tf.summary.FileWriter(log_dir_image_init, sess.graph, flush_secs=30, filename_suffix = '.image_log_init') 207 | 208 | # for pretrain 209 | loss_sum_list_init = [] 210 | with tf.variable_scope('loss_init'): 211 | loss_sum_list_init.append(tf.summary.scalar('1_total_loss_init', loss_init)) 212 | loss_sum_list_init.append(tf.summary.scalar('2_defocus_loss_init', loss_defocus)) 213 | loss_sum_init = tf.summary.merge(loss_sum_list_init) 214 | 215 | image_sum_list_init = [] 216 | image_sum_list_init.append(tf.summary.image('1_synthetic_input_init', patches_synthetic)) 217 | image_sum_list_init.append(tf.summary.image('2_synthetic_defocus_out_init', fix_image_tf(output_synthetic_defocus, 1))) 218 | image_sum_list_init.append(tf.summary.image('3_synthetic_defocus_gt_init', fix_image_tf(labels_synthetic_defocus, 1))) 219 | image_sum_init = tf.summary.merge(image_sum_list_init) 220 | 221 | # for train 222 | loss_sum_g_list = [] 223 | with tf.variable_scope('loss_generator'): 224 | loss_sum_g_list.append(tf.summary.scalar('1_total', loss_main)) 225 | loss_sum_g_list.append(tf.summary.scalar('2_g', loss_g)) 226 | loss_sum_g_list.append(tf.summary.scalar('3_defocus', loss_defocus)) 227 | loss_sum_g_list.append(tf.summary.scalar('4_perceptual', loss_perceptual)) 228 | loss_sum_g_list.append(tf.summary.scalar('5_auxilary', loss_aux)) 229 | loss_sum_g_list.append(tf.summary.scalar('6_binary', loss_binary)) 230 | loss_sum_g = tf.summary.merge(loss_sum_g_list) 231 | 232 | loss_sum_d_list = [] 233 | with tf.variable_scope('loss_discriminator'): 234 | loss_sum_d_list.append(tf.summary.scalar('1_d', loss_d_feature)) 235 | loss_sum_d = tf.summary.merge(loss_sum_d_list) 236 | 237 | image_sum_list = [] 238 | image_sum_list.append(tf.summary.image('1_synthetic_input', patches_synthetic)) 239 | image_sum_list.append(tf.summary.image('2_synthetic_defocus_out', fix_image_tf(output_synthetic_defocus, 1))) 240 | image_sum_list.append(tf.summary.image('3_synthetic_defocus_gt', fix_image_tf(labels_synthetic_defocus, 1))) 241 | image_sum_list.append(tf.summary.image('4_real_input', patches_real)) 242 | image_sum_list.append(tf.summary.image('5_real_defocus_out', fix_image_tf(output_real_defocus, 1))) 243 | image_sum_list.append(tf.summary.image('6_real_binary_out', fix_image_tf(output_real_binary, 1))) 244 | image_sum_list.append(tf.summary.image('7_real_binary_gt', fix_image_tf(labels_real_binary, 1))) 245 | image_sum_list.append(tf.summary.image('8_real_binary_gt_no_label', patches_real_no_label)) 246 | image_sum_list.append(tf.summary.image('9_real_defocus_out_no_label', fix_image_tf(output_real_no_label_defocus, 1))) 247 | image_sum = tf.summary.merge(image_sum_list) 248 | 249 | ## INITIALIZE SESSION 250 | sess.run(tf.global_variables_initializer()) 251 | ## LOAD VGG 252 | vgg19_npy_path = "pretrained/vgg19.npy" 253 | if not os.path.isfile(vgg19_npy_path): 254 | print("Please download vgg19.npz from : https://github.com/machrisaa/tensorflow-vgg") 255 | exit() 256 | npz = np.load(vgg19_npy_path, encoding='latin1', allow_pickle=True).item() 257 | 258 | params = [] 259 | for val in sorted( npz.items() ): 260 | if val[0] == 'fc6': 261 | break; 262 | W = np.asarray(val[1][0]) 263 | b = np.asarray(val[1][1]) 264 | print(" Loading %s: %s, %s" % (val[0], W.shape, b.shape)) 265 | params.extend([W, b]) 266 | tl.files.assign_params(sess, params, net_vgg) 267 | tl.files.assign_params(sess, params, net_vgg_perceptual) 268 | 269 | init_name = init_dir + '/{}_init.npz'.format(tl.global_flag['mode']) 270 | if os.path.isfile(init_name): 271 | tl.files.load_and_assign_npz_dict(name = init_name, sess = sess) 272 | if tl.global_flag['is_pretrain']: 273 | print('*****************************************') 274 | print(' PRE-TRAINING START') 275 | print('*****************************************') 276 | global_step = 0 277 | for epoch in range(0, n_epoch_init): 278 | total_loss_init, n_iter = 0, 0 279 | # reload image list 280 | train_synthetic_img_list = np.array(sorted(tl.files.load_file_list(path = config.TRAIN.synthetic_img_path, regx = '.*', printable = False))) 281 | train_defocus_map_list = np.array(sorted(tl.files.load_file_list(path = config.TRAIN.defocus_map_path, regx = '.*', printable = False))) 282 | 283 | # shuffle datasets 284 | shuffle_index = np.arange(len(train_synthetic_img_list)) 285 | np.random.shuffle(shuffle_index) 286 | train_synthetic_img_list = train_synthetic_img_list[shuffle_index] 287 | train_defocus_map_list = train_defocus_map_list[shuffle_index] 288 | 289 | epoch_time = time.time() 290 | #for idx in range(0, len(train_synthetic_img_list), batch_size_init): 291 | for idx in range(0, 10): 292 | step_time = time.time() 293 | ## READ DATA 294 | # read synthetic data 295 | b_idx = (idx + np.arange(batch_size_init)) % len(train_synthetic_img_list) 296 | synthetic_images_blur = read_all_imgs(train_synthetic_img_list[b_idx], path = config.TRAIN.synthetic_img_path, mode = 'RGB') 297 | synthetic_defocus_maps = read_all_imgs(train_defocus_map_list[b_idx], path = config.TRAIN.defocus_map_path, mode = 'DEPTH') 298 | 299 | synthetic_images_blur, synthetic_defocus_maps = \ 300 | crop_pair_with_different_shape_images(synthetic_images_blur, synthetic_defocus_maps, [h, w], is_gaussian_noise = tl.global_flag['is_noise']) 301 | 302 | err_init, lr, _ = \ 303 | sess.run([loss_init, learning_rate_init, optim_init], {patches_synthetic: synthetic_images_blur, labels_synthetic_defocus: synthetic_defocus_maps}) 304 | 305 | print('[%s] Ep [%2d/%2d] %4d/%4d time: %4.2fs, err_init: %1.2e, lr: %1.2e' % \ 306 | (tl.global_flag['mode'], epoch, n_epoch_init, n_iter, len(train_synthetic_img_list)/batch_size_init, time.time() - step_time, err_init, lr)) 307 | 308 | if global_step % config.TRAIN.write_log_every == 0: 309 | summary_loss_init, summary_image_init = sess.run([loss_sum_init, image_sum_init], {patches_synthetic: synthetic_images_blur, labels_synthetic_defocus: synthetic_defocus_maps}) 310 | writer_scalar_init.add_summary(summary_loss_init, global_step) 311 | writer_image_init.add_summary(summary_image_init, global_step) 312 | 313 | total_loss_init += err_init 314 | n_iter += 1 315 | global_step += 1 316 | 317 | if epoch % config.TRAIN.refresh_image_log_every and epoch != n_epoch_init == 0: 318 | writer_image_init.close() 319 | remove_file_end_with(log_dir_image_init, '*.image_log') 320 | writer_image_init.reopen() 321 | 322 | if epoch % 2 or epoch == n_epoch_init - 1: 323 | tl.files.save_npz_dict(save_init_vars, name = init_dir + '/{}_init.npz'.format(tl.global_flag['mode']), sess = sess) 324 | 325 | writer_image_init.close() 326 | writer_scalar_init.close() 327 | 328 | ## START TRAINING 329 | print('*****************************************') 330 | print(' TRAINING START') 331 | print('*****************************************') 332 | global_step = 0 333 | for epoch in range(0, n_epoch + 1): 334 | total_loss, n_iter = 0, 0 335 | 336 | # reload synthetic datasets 337 | train_synthetic_img_list = np.array(sorted(tl.files.load_file_list(path = config.TRAIN.synthetic_img_path, regx = '.*', printable = False))) 338 | train_defocus_map_list = np.array(sorted(tl.files.load_file_list(path = config.TRAIN.defocus_map_path, regx = '.*', printable = False))) 339 | 340 | # shuffle datasets 341 | shuffle_index = np.arange(len(train_synthetic_img_list)) 342 | np.random.shuffle(shuffle_index) 343 | 344 | train_synthetic_img_list = train_synthetic_img_list[shuffle_index] 345 | train_defocus_map_list = train_defocus_map_list[shuffle_index] 346 | 347 | shuffle_index = np.arange(len(train_real_img_list)) 348 | np.random.shuffle(shuffle_index) 349 | train_real_img_list = train_real_img_list[shuffle_index] 350 | train_real_binary_map_list = train_real_binary_map_list[shuffle_index] 351 | 352 | shuffle_index = np.arange(len(train_real_img_no_label_list)) 353 | np.random.shuffle(shuffle_index) 354 | train_real_img_no_label_list = train_real_img_no_label_list[shuffle_index] 355 | 356 | # update learning rate 357 | if epoch != 0 and (epoch % decay_every == 0): 358 | new_lr_decay = lr_decay ** (epoch // decay_every) 359 | sess.run(tf.assign(learning_rate, lr_init * new_lr_decay)) 360 | elif epoch == 0: 361 | sess.run(tf.assign(learning_rate, lr_init)) 362 | 363 | epoch_time = time.time() 364 | for idx in range(0, len(train_synthetic_img_list), batch_size): 365 | step_time = time.time() 366 | 367 | ## READ DATA 368 | # read synthetic data 369 | b_idx = (idx + np.arange(batch_size)) % len(train_synthetic_img_list) 370 | synthetic_images_blur = read_all_imgs(train_synthetic_img_list[b_idx], path = config.TRAIN.synthetic_img_path, mode = 'RGB') 371 | synthetic_defocus_maps = read_all_imgs(train_defocus_map_list[b_idx], path = config.TRAIN.defocus_map_path, mode = 'DEPTH') 372 | 373 | synthetic_images_blur, synthetic_defocus_maps = crop_pair_with_different_shape_images(synthetic_images_blur, synthetic_defocus_maps, [h, w], is_gaussian_noise = tl.global_flag['is_noise']) 374 | 375 | # read real data # 376 | b_idx = (idx % len(train_real_img_list) + np.arange(batch_size)) % len(train_real_img_list) 377 | real_images_blur = read_all_imgs(train_real_img_list[b_idx], path = config.TRAIN.real_img_path, mode = 'RGB') 378 | real_binary_maps = read_all_imgs(train_real_binary_map_list[b_idx], path = config.TRAIN.real_binary_map_path, mode = 'GRAY') 379 | real_images_blur, real_binary_maps = crop_pair_with_different_shape_images(real_images_blur, real_binary_maps, [h, w]) 380 | real_images_no_label_blur = read_all_imgs(train_real_img_no_label_list[b_idx], path = config.TRAIN.real_img_no_label_path, mode = 'RGB') 381 | real_images_no_label_blur = random_crop(real_images_no_label_blur, [h, w]) 382 | 383 | ## RUN NETWORK 384 | #discriminator 385 | feed_dict = {patches_synthetic: synthetic_images_blur, patches_real: real_images_blur, labels_synthetic_defocus: synthetic_defocus_maps, patches_real_no_label: real_images_no_label_blur} 386 | _ = sess.run(optim_d, feed_dict) 387 | 388 | #generator 389 | feed_dict = {patches_synthetic: synthetic_images_blur, labels_synthetic_defocus: synthetic_defocus_maps, patches_real: real_images_blur, labels_real_binary: real_binary_maps, patches_real_no_label: real_images_no_label_blur} 390 | _ = sess.run(optim_main, feed_dict) 391 | 392 | #log 393 | err_main, err_g, err_d, d_synthetic, d_real, lr = \ 394 | sess.run([loss_main, loss_g, loss_d, d_feature_synthetic, d_feature_real, learning_rate], feed_dict) 395 | d_acc = get_disc_accuracy([d_synthetic, d_real], [0, 1]) 396 | g_acc = get_disc_accuracy([d_synthetic, d_real], [1, 0]) 397 | 398 | print('[%s] Ep [%2d/%2d] %4d/%4d time: %4.2fs, err[main: %1.2e, g(acc): %1.2e(%1.2f), d(acc): %1.2e(%1.2f)], lr: %1.2e' % \ 399 | (tl.global_flag['mode'], epoch, n_epoch, n_iter, len(train_synthetic_img_list)/batch_size, time.time() - step_time, err_main, err_g, g_acc, err_d, d_acc, lr)) 400 | 401 | ## SAVE LOGS 402 | # save loss & image log 403 | if global_step % config.TRAIN.write_log_every == 0: 404 | summary_loss_g, summary_loss_d, summary_image = sess.run([loss_sum_g, loss_sum_d, image_sum], {patches_synthetic: synthetic_images_blur, labels_synthetic_defocus: synthetic_defocus_maps, patches_real: real_images_blur, labels_real_binary: real_binary_maps, patches_real_no_label: real_images_no_label_blur}) 405 | writer_scalar.add_summary(summary_loss_d, global_step) 406 | writer_scalar.add_summary(summary_loss_g, global_step) 407 | writer_image.add_summary(summary_image, global_step) 408 | 409 | # save samples 410 | # if global_step != 0 and global_step % config.TRAIN.write_sample_every == 0: 411 | # synthetic_defocus_out, real_defocus_out, real_binary_out = sess.run([output_synthetic_defocus, output_real_defocus, output_real_binary], {patches_synthetic: synthetic_images_blur, patches_real: real_images_blur, labels_real_binary: real_binary_maps }) 412 | # save_images(synthetic_images_blur, [ni, ni], sample_dir + '/{}_{}_1_synthetic_input.png'.format(epoch, global_step)) 413 | # save_images(norm_image(synthetic_defocus_out), [ni, ni], sample_dir + '/{}_{}_2_synthetic_defocus_out_norm.png'.format(epoch, global_step)) 414 | # save_images(norm_image(synthetic_defocus_maps), [ni, ni], sample_dir + '/{}_{}_3_synthetic_defocus_gt.png'.format(epoch, global_step)) 415 | # save_images(real_images_blur, [ni, ni], sample_dir + '/{}_{}_4_real_input.png'.format(epoch, global_step)) 416 | # save_images(norm_image(real_defocus_out), [ni, ni], sample_dir + '/{}_{}_5_real_defocus_out_norm.png'.format(epoch, global_step)) 417 | # save_images(real_binary_out, [ni, ni], sample_dir + '/{}_{}_6_real_binary_out.png'.format(epoch, global_step)) 418 | # save_images(real_binary_maps, [ni, ni], sample_dir + '/{}_{}_7_real_binary_gt.png'.format(epoch, global_step)) 419 | 420 | total_loss += err_main 421 | n_iter += 1 422 | global_step += 1 423 | 424 | print('[TRAIN] Epoch: [%2d/%2d] time: %4.4fs, total_err: %1.2e' % (epoch, n_epoch, time.time() - epoch_time, total_loss/n_iter)) 425 | # reset image log 426 | if epoch % config.TRAIN.refresh_image_log_every == 0: 427 | writer_image.close() 428 | remove_file_end_with(log_dir_image, '*.image_log') 429 | writer_image.reopen() 430 | 431 | if epoch % config.TRAIN.write_ckpt_every == 0: 432 | #remove_file_end_with(ckpt_dir, '*.npz') 433 | tl.files.save_npz_dict(save_vars, name = ckpt_dir + '/{}_{}.npz'.format(tl.global_flag['mode'], epoch), sess = sess) 434 | 435 | 436 | def evaluate(): 437 | date = datetime.datetime.now().strftime('%Y_%m_%d_%H%M') 438 | # directories 439 | mode_dir = os.path.join(config.TRAIN.root_dir, tl.global_flag['mode']) 440 | ckpt_dir = os.path.join(mode_dir, 'checkpoint') 441 | sample_dir = os.path.join(mode_dir, 'results/{}/{}'.format(tl.global_flag['test_set'], date)) 442 | 443 | # input 444 | input_path, gt_path = get_eval_path(tl.global_flag['test_set'], config) 445 | test_blur_img_list = np.array(sorted(tl.files.load_file_list(path = input_path, regx = '.*', printable = False))) 446 | test_blur_imgs = read_all_imgs(test_blur_img_list, path = input_path, mode = 'RGB') 447 | # gt 448 | if gt_path is not None: 449 | test_gt_list = np.array(sorted(tl.files.load_file_list(path = gt_path, regx = '.*', printable = False))) 450 | mode = 'NPY' if 'RTF' in tl.global_flag['test_set'] else 'GRAY' 451 | test_gt_imgs = read_all_imgs(test_gt_list, path = gt_path, mode = mode) 452 | 453 | # define session 454 | sess = tf.Session(config = tf.ConfigProto(allow_soft_placement = True, log_device_placement = False)) 455 | # define model 456 | with tf.variable_scope('input'): 457 | patches_blurred = tf.placeholder('float32', [1, None, None, 3], name = 'input_patches') 458 | 459 | with tf.variable_scope('main_net') as scope: 460 | with tf.variable_scope('defocus_net') as scope: 461 | with tf.variable_scope('encoder') as scope: 462 | feats_down = VGG19_down(patches_blurred, reuse = False, scope = scope, is_test = True) 463 | with tf.variable_scope('decoder') as scope: 464 | output_defocus, feats_up, _, refine_lists = UNet_up(patches_blurred, feats_down, is_train = False, reuse = False, scope = scope) 465 | 466 | # init vars 467 | sess.run(tf.global_variables_initializer()) 468 | # load checkpoint 469 | ckpt_path = os.path.join(ckpt_dir, '{}.npz'.format(tl.global_flag['mode'])) 470 | if os.path.isfile(ckpt_path) is False: 471 | print('{} does not exit'.format(ckpt_path)) 472 | exit() 473 | tl.files.load_and_assign_npz_dict(name = ckpt_path, sess = sess) 474 | 475 | print('================') 476 | print('Evaluation Start') 477 | print('================') 478 | print('Results will be saved in: {}\n'.format(sample_dir)) 479 | avg_time = 0. 480 | MSE_total = 0 481 | MAE_total = 0 482 | for i in np.arange(len(test_blur_imgs)): 483 | test_blur_img = refine_image(test_blur_imgs[i]) 484 | 485 | if gt_path is not None: 486 | test_gt_img = refine_image(test_gt_imgs[i]) 487 | test_gt_img = np.squeeze(test_gt_img) 488 | 489 | # run network 490 | feed_dict = {patches_blurred: np.expand_dims(test_blur_img, axis = 0)} 491 | tic = time.time() 492 | defocus_map, feats_down_out, feats_up_out, refine_lists_out = sess.run([output_defocus, feats_down, feats_up, refine_lists], feed_dict) 493 | toc = time.time() 494 | 495 | defocus_map = np.squeeze(defocus_map) 496 | defocus_map_norm = defocus_map - defocus_map.min() 497 | defocus_map_norm = defocus_map_norm / defocus_map_norm.max() 498 | 499 | ################## 500 | sigma_map = ((defocus_map * 15) - 1) / 2 501 | sigma_map[np.where(sigma_map < 0)] = 0 502 | # when you read, multipy the image by 7 to get sigma 503 | sigma_map = sigma_map / 7. 504 | ################## 505 | 506 | # quantitative 507 | if gt_path is not None: 508 | if 'RTF' in tl.global_flag['test_set']: 509 | h, w = sigma_map.shape 510 | h_gt, w_gt = test_gt_img.shape 511 | in_temp = sigma_map[:min(h, h_gt), :min(w, w_gt)] * 7. 512 | in_temp[np.where(in_temp>3.275)] = 3.275 513 | in_temp = in_temp / 3.275 514 | gt_temp = test_gt_img[:min(h, h_gt), :min(w, w_gt)] 515 | elif 'SYNDOF' in tl.global_flag['test_set']: 516 | in_temp = defocus_map 517 | gt_temp = test_gt_img 518 | 519 | MSE_total = MSE_total + np.mean((in_temp - gt_temp)**2) 520 | MAE_total = MAE_total + np.mean(np.abs(in_temp - gt_temp)) 521 | 522 | # qualitative 523 | tl.files.exists_or_mkdir(sample_dir, verbose = False) 524 | tl.files.exists_or_mkdir(sample_dir + '/image') 525 | tl.files.exists_or_mkdir(sample_dir + '/defocus_map') 526 | tl.files.exists_or_mkdir(sample_dir + '/defocus_map_min_max_norm') 527 | tl.files.exists_or_mkdir(sample_dir + '/sigma_map_7_norm') 528 | scipy.misc.toimage(test_blur_img, cmin = 0., cmax = 1.).save(sample_dir + '/image/{0:04d}.png'.format(i)) 529 | scipy.misc.toimage(defocus_map, cmin = 0., cmax = 1.).save(sample_dir + '/defocus_map/{0:04d}.png'.format(i)) 530 | scipy.misc.toimage(defocus_map_norm, cmin = 0., cmax = 1.).save(sample_dir + '/defocus_map_min_max_norm/{0:04d}.png'.format(i)) 531 | scipy.misc.toimage(sigma_map, cmin = 0., cmax = 1.).save(sample_dir + '/sigma_map_7_norm/{0:04d}.png'.format(i)) 532 | 533 | if gt_path is not None: 534 | tl.files.exists_or_mkdir(sample_dir + '/gt') 535 | scipy.misc.toimage(np.squeeze(1 - refine_image(test_gt_imgs[i])), cmin = 0., cmax = 1.).save(sample_dir + '/gt/{0:04d}.png'.format(i)) 536 | 537 | avg_time = avg_time + (toc - tic) 538 | print('[{}/{}] {} [{:.3f}s]\n'.format(i+1, len(test_blur_imgs), test_blur_img_list[i], toc - tic)) 539 | 540 | avg_time = avg_time / len(test_blur_imgs) 541 | print('averge time: {:.3f}s'.format(avg_time)) 542 | if gt_path is not None: 543 | print('MSE: ', MSE_total / len(test_blur_imgs), ' MAE: ', MAE_total / len(test_blur_imgs)) 544 | 545 | if __name__ == '__main__': 546 | import argparse 547 | parser = argparse.ArgumentParser() 548 | 549 | parser.add_argument('--mode', type = str, default = 'DMENet', help = 'model name') 550 | parser.add_argument('--is_train', action = 'store_true', default = False, help = 'whether to train or not') 551 | parser.add_argument('--is_pretrain', action = 'store_true', default = False, help = 'whether to pretrain or not') 552 | parser.add_argument('--is_noise', action = 'store_true', default = False, help = 'whether to add noise to synthetic images') 553 | parser.add_argument('--delete_log', action = 'store_true', default = False, help = 'whether to delete log or not') 554 | parser.add_argument('--test_set', type = str , default = 'CUHK', help = 'test_set CUHK|SYNDOF|RTF0|RTF1|RTF1_6|random') 555 | 556 | args = parser.parse_args() 557 | 558 | tl.global_flag['mode'] = args.mode 559 | tl.global_flag['is_train'] = args.is_train 560 | tl.global_flag['is_pretrain'] = args.is_pretrain 561 | tl.global_flag['is_noise'] = args.is_noise 562 | tl.global_flag['delete_log'] = args.delete_log 563 | 564 | if tl.global_flag['is_train']: 565 | train() 566 | else: 567 | tl.global_flag['test_set'] = args.test_set 568 | evaluate() 569 | 570 | -------------------------------------------------------------------------------- /model.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import tensorlayer as tl 3 | import numpy as np 4 | from tensorlayer.layers import * 5 | 6 | def VGG19_down(rgb, reuse, scope, is_test = False): 7 | w_init_relu = tf.contrib.layers.variance_scaling_initializer() 8 | w_init_sigmoid = tf.contrib.layers.xavier_initializer() 9 | VGG_MEAN = [103.939, 116.779, 123.68] 10 | with tf.variable_scope(scope, reuse = reuse): 11 | rgb_scaled = rgb * 255.0 12 | if tf.__version__ <= '0.11': 13 | red, green, blue = tf.split(3, 3, rgb_scaled) 14 | else: 15 | red, green, blue = tf.split(rgb_scaled, 3, 3) 16 | if tf.__version__ <= '0.11': 17 | bgr = tf.concat(3, [ 18 | blue - VGG_MEAN[0], 19 | green - VGG_MEAN[1], 20 | red - VGG_MEAN[2], 21 | ]) 22 | else: 23 | bgr = tf.concat([ 24 | blue - VGG_MEAN[0], 25 | green - VGG_MEAN[1], 26 | red - VGG_MEAN[2], 27 | ], axis=3) 28 | 29 | """ input layer """ 30 | net_in = InputLayer(bgr, name='input') 31 | """ conv1 """ 32 | network = PadLayer(net_in, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='pad1_1') 33 | network = Conv2d(network, n_filter=64, filter_size=(3, 3), strides=(1, 1), act=tf.nn.relu,padding='VALID', name='conv1_1') 34 | network = PadLayer(network, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='pad1_2') 35 | network = Conv2d(network, n_filter=64, filter_size=(3, 3), strides=(1, 1), act=tf.nn.relu,padding='VALID', name='conv1_2') 36 | d0 = network 37 | network = MaxPool2d(network, filter_size=(2, 2), strides=(2, 2), padding='SAME', name='pool1') 38 | """ conv2 """ 39 | network = PadLayer(network, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='pad2_1') 40 | network = Conv2d(network, n_filter=128, filter_size=(3, 3), strides=(1, 1), act=tf.nn.relu,padding='VALID', name='conv2_1') 41 | network = PadLayer(network, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='pad2_2') 42 | network = Conv2d(network, n_filter=128, filter_size=(3, 3), strides=(1, 1), act=tf.nn.relu,padding='VALID', name='conv2_2') 43 | d1 = network 44 | network = MaxPool2d(network, filter_size=(2, 2), strides=(2, 2), padding='SAME', name='pool2') 45 | """ conv3 """ 46 | network = PadLayer(network, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='pad3_1') 47 | network = Conv2d(network, n_filter=256, filter_size=(3, 3), strides=(1, 1), act=tf.nn.relu,padding='VALID', name='conv3_1') 48 | network = PadLayer(network, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='pad3_2') 49 | network = Conv2d(network, n_filter=256, filter_size=(3, 3), strides=(1, 1), act=tf.nn.relu,padding='VALID', name='conv3_2') 50 | network = PadLayer(network, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='pad3_3') 51 | network = Conv2d(network, n_filter=256, filter_size=(3, 3), strides=(1, 1), act=tf.nn.relu,padding='VALID', name='conv3_3') 52 | network = PadLayer(network, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='pad3_4') 53 | network = Conv2d(network, n_filter=256, filter_size=(3, 3), strides=(1, 1), act=tf.nn.relu,padding='VALID', name='conv3_4') 54 | d2 = network 55 | network = MaxPool2d(network, filter_size=(2, 2), strides=(2, 2), padding='SAME', name='pool3') 56 | """ conv4 """ 57 | network = PadLayer(network, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='pad4_1') 58 | network = Conv2d(network, n_filter=512, filter_size=(3, 3), strides=(1, 1), act=tf.nn.relu,padding='VALID', name='conv4_1') 59 | network = PadLayer(network, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='pad4_2') 60 | network = Conv2d(network, n_filter=512, filter_size=(3, 3), strides=(1, 1), act=tf.nn.relu,padding='VALID', name='conv4_2') 61 | network = PadLayer(network, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='pad4_3') 62 | network = Conv2d(network, n_filter=512, filter_size=(3, 3), strides=(1, 1), act=tf.nn.relu,padding='VALID', name='conv4_3') 63 | network = PadLayer(network, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='pad4_4') 64 | network = Conv2d(network, n_filter=512, filter_size=(3, 3), strides=(1, 1), act=tf.nn.relu,padding='VALID', name='conv4_4') 65 | d3 = network 66 | network = MaxPool2d(network, filter_size=(2, 2), strides=(2, 2), padding='SAME', name='pool4') 67 | """ conv5 """ 68 | network = PadLayer(network, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='pad5_1') 69 | network = Conv2d(network, n_filter=512, filter_size=(3, 3), strides=(1, 1), act=tf.nn.relu,padding='VALID', name='conv5_1') 70 | network = PadLayer(network, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='pad5_2') 71 | network = Conv2d(network, n_filter=512, filter_size=(3, 3), strides=(1, 1), act=tf.nn.relu,padding='VALID', name='conv5_2') 72 | network = PadLayer(network, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='pad5_3') 73 | network = Conv2d(network, n_filter=512, filter_size=(3, 3), strides=(1, 1), act=tf.nn.relu,padding='VALID', name='conv5_3') 74 | network = PadLayer(network, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='pad5_4') 75 | network = Conv2d(network, n_filter=512, filter_size=(3, 3), strides=(1, 1), act=tf.nn.relu,padding='VALID', name='conv5_4') 76 | d4 = network 77 | 78 | if is_test == False: 79 | logits = PadLayer(network, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='pad6_1') 80 | logits = Conv2d(logits, n_filter=64, filter_size=(3, 3), strides=(1, 1), act=tf.nn.relu,padding='VALID', name='conv6_1') 81 | 82 | logits = logits.outputs 83 | size = logits.get_shape().as_list() 84 | logits = InputLayer(logits) 85 | logits = Conv2d(logits, n_filter=512, filter_size=(size[1], size[2]), strides=(1, 1), act=tf.nn.relu, padding='VALID', name='c_logits_1') 86 | logits = FlattenLayer(logits, name='flatten') 87 | logits = DenseLayer(logits, n_units=512, act=tf.nn.relu, W_init = w_init_relu, name='c_logits_1') 88 | logits = DenseLayer(logits, n_units=1, act=tf.identity, W_init = w_init_sigmoid, name='c_logits_2') 89 | 90 | return network, [d0.outputs, d1.outputs, d2.outputs, d3.outputs, d4.outputs], d3.outputs, logits.outputs 91 | else: 92 | return [d0.outputs, d1.outputs, d2.outputs, d3.outputs, d4.outputs] 93 | 94 | def UNet_up(images, feats, is_train=False, reuse=False, scope = 'unet_up'): 95 | w_init_relu = tf.contrib.layers.variance_scaling_initializer() 96 | w_init_sigmoid = tf.contrib.layers.xavier_initializer() 97 | g_init = None 98 | lrelu = lambda x: tf.nn.leaky_relu(x, 0.2) 99 | 100 | def UpSampling2dLayer_(input, scale, method, align_corners, name): 101 | input = input.outputs 102 | size = tf.shape(input) 103 | 104 | n = InputLayer(input, name = name + '_in') 105 | n = UpSampling2dLayer(n, size=[size[1] * scale[0], size[2] * scale[1]], is_scale = False, method = method, align_corners = align_corners, name = name) 106 | 107 | return n 108 | 109 | with tf.variable_scope(scope, reuse=reuse): 110 | d0 = InputLayer(feats[0], name='d0') 111 | d1 = InputLayer(feats[1], name='d1') 112 | d2 = InputLayer(feats[2], name='d2') 113 | d3 = InputLayer(feats[3], name='d3') 114 | d4 = InputLayer(feats[4], name='d4') 115 | 116 | u4 = d4 117 | u4 = PadLayer(u4, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='u4_aux/pad1') 118 | u4 = Conv2d(u4, 256, (3, 3), (1, 1), act=None, padding='VALID', W_init=w_init_relu, name='u4_aux/c1') 119 | u4 = BatchNormLayer(u4, act=lrelu, is_train = is_train, gamma_init = g_init, name='u4_aux/b1') 120 | u4 = PadLayer(u4, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='u4_aux/pad2') 121 | u4 = Conv2d(u4, 1, (3, 3), (1, 1), act=None, padding='VALID', W_init=w_init_sigmoid, name='u4_aux/c2') 122 | u4 = BatchNormLayer(u4, act=tf.nn.sigmoid, is_train = is_train, gamma_init = g_init, name='u4_aux/b2') 123 | u4 = u4.outputs 124 | 125 | n = UpSampling2dLayer_(d4, (2, 2), method = 1, align_corners=True, name='u3/u') 126 | n = ConcatLayer([n, d3], concat_dim = 3, name='u3/concat') 127 | n = PadLayer(n, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='u3/pad1') 128 | n = Conv2d(n, 256, (3, 3), (1, 1), act=None, padding='VALID', W_init=w_init_relu, name='u3/c1') 129 | n = BatchNormLayer(n, act=lrelu, is_train = is_train, gamma_init = g_init, name='u3/b1') 130 | n = PadLayer(n, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='u3/pad2') 131 | n = Conv2d(n, 256, (3, 3), (1, 1), act=None, padding='VALID', W_init=w_init_relu, name='u3/c2') 132 | n = BatchNormLayer(n, act=lrelu, is_train = is_train, gamma_init = g_init, name='u3/b2') 133 | n = PadLayer(n, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='u3/pad3') 134 | n = Conv2d(n, 256, (3, 3), (1, 1), act=None, padding='VALID', W_init=w_init_relu, name='u3/c3') 135 | n = BatchNormLayer(n, act=lrelu, is_train = is_train, gamma_init = g_init, name='u3/b3') 136 | 137 | u3 = PadLayer(n, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='u3_aux/pad1') 138 | u3 = Conv2d(u3, 128, (3, 3), (1, 1), act=None, padding='VALID', W_init=w_init_relu, name='u3_aux/c1') 139 | u3 = BatchNormLayer(u3, act=lrelu, is_train = is_train, gamma_init = g_init, name='u3_aux/b1') 140 | u3 = PadLayer(u3, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='u3_aux/pad2') 141 | u3 = Conv2d(u3, 1, (3, 3), (1, 1), act=None, padding='VALID', W_init=w_init_sigmoid, name='u3_aux/c2') 142 | u3 = BatchNormLayer(u3, act=tf.nn.sigmoid, is_train = is_train, gamma_init = g_init, name='u3_aux/b2') 143 | u3 = u3.outputs 144 | 145 | n = UpSampling2dLayer_(n, (2, 2), method = 1, align_corners=True, name='u2/u') 146 | n = ConcatLayer([n, d2], concat_dim = 3, name='u2/concat') 147 | n = PadLayer(n, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='u2/pad1') 148 | n = Conv2d(n, 128, (3, 3), (1, 1), act=None, padding='VALID', W_init=w_init_relu, name='u2/c1') 149 | n = BatchNormLayer(n, act=lrelu, is_train = is_train, gamma_init = g_init, name='u2/b1') 150 | n = PadLayer(n, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='u2/pad2') 151 | n = Conv2d(n, 128, (3, 3), (1, 1), act=None, padding='VALID', W_init=w_init_relu, name='u2/c2') 152 | n = BatchNormLayer(n, act=lrelu, is_train = is_train, gamma_init = g_init, name='u2/b2') 153 | n = PadLayer(n, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='u2/pad3') 154 | n = Conv2d(n, 128, (3, 3), (1, 1), act=None, padding='VALID', W_init=w_init_relu, name='u2/c3') 155 | n = BatchNormLayer(n, act=lrelu, is_train = is_train, gamma_init = g_init, name='u2/b3') 156 | 157 | u2 = PadLayer(n, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='u2_aux/pad1') 158 | u2 = Conv2d(u2, 64, (3, 3), (1, 1), act=None, padding='VALID', W_init=w_init_relu, name='u2_aux/c1') 159 | u2 = BatchNormLayer(u2, act=lrelu, is_train = is_train, gamma_init = g_init, name='u2_aux/b1') 160 | u2 = PadLayer(u2, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='u2_aux/pad2') 161 | u2 = Conv2d(u2, 1, (3, 3), (1, 1), act=None, padding='VALID', W_init=w_init_sigmoid, name='u2_aux/c2') 162 | u2 = BatchNormLayer(u2, act=tf.nn.sigmoid, is_train = is_train, gamma_init = g_init, name='u2_aux/b2') 163 | u2 = u2.outputs 164 | 165 | n = UpSampling2dLayer_(n, (2, 2), method = 1, align_corners=True, name='u1/u') 166 | n = ConcatLayer([n, d1], concat_dim = 3, name='u1/concat') 167 | n = PadLayer(n, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='u1/pad1') 168 | n = Conv2d(n, 64, (3, 3), (1, 1), act=None, padding='VALID', W_init=w_init_relu, name='u1/c1') 169 | n = BatchNormLayer(n, act=lrelu, is_train = is_train, gamma_init = g_init, name='u1/b1') 170 | n = PadLayer(n, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='u1/pad2') 171 | n = Conv2d(n, 64, (3, 3), (1, 1), act=None, padding='VALID', W_init=w_init_relu, name='u1/c2') 172 | n = BatchNormLayer(n, act=lrelu, is_train = is_train, gamma_init = g_init, name='u1/b2') 173 | n = PadLayer(n, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='u1/pad3') 174 | n = Conv2d(n, 64, (3, 3), (1, 1), act=None, padding='VALID', W_init=w_init_relu, name='u1/c3') 175 | n = BatchNormLayer(n, act=lrelu, is_train = is_train, gamma_init = g_init, name='u1/b3') 176 | 177 | u1 = PadLayer(n, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='u1_aux/pad1') 178 | u1 = Conv2d(u1, 32, (3, 3), (1, 1), act=None, padding='VALID', W_init=w_init_relu, name='u1_aux/c1') 179 | u1 = BatchNormLayer(u1, act=lrelu, is_train = is_train, gamma_init = g_init, name='u1_aux/b1') 180 | u1 = PadLayer(u1, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='u1_aux/pad2') 181 | u1 = Conv2d(u1, 1, (3, 3), (1, 1), act=None, padding='VALID', W_init=w_init_sigmoid, name='u1_aux/c2') 182 | u1 = BatchNormLayer(u1, act=tf.nn.sigmoid, is_train = is_train, gamma_init = g_init, name='u1_aux/b2') 183 | u1 = u1.outputs 184 | 185 | n = UpSampling2dLayer_(n, (2, 2), method = 1, align_corners=True, name='u0/u') 186 | n = ConcatLayer([n, d0], concat_dim = 3, name='u0/concat') 187 | n = PadLayer(n, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='u0/pad_init') 188 | n = Conv2d(n, 64, (3, 3), (1, 1), act=None, padding='VALID', W_init=w_init_relu, name='u0/c_init') 189 | n = BatchNormLayer(n, act=lrelu, is_train = is_train, gamma_init = g_init, name='u0/b_init') 190 | gan_feat = n.outputs 191 | 192 | u0 = PadLayer(n, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='u0_aux/pad1') 193 | u0 = Conv2d(u0, 32, (3, 3), (1, 1), act=None, padding='VALID', W_init=w_init_relu, name='u0_aux/c1') 194 | u0 = BatchNormLayer(u0, act=lrelu, is_train = is_train, gamma_init = g_init, name='u0_aux/b1') 195 | u0 = PadLayer(u0, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='u0_aux/pad2') 196 | u0 = Conv2d(u0, 1, (3, 3), (1, 1), act=None, padding='VALID', W_init=w_init_sigmoid, name='u0_aux/c2') 197 | u0 = BatchNormLayer(u0, act=tf.nn.sigmoid, is_train = is_train, gamma_init = g_init, name='u0_aux/b2') 198 | u0 = u0.outputs 199 | 200 | refine_lists = [] 201 | refine_lists.append(n.outputs) 202 | for i in np.arange(7): 203 | n_res = n 204 | n_res = Conv2d(n_res, 64, (1, 1), (1, 1), act=None, padding='VALID', W_init=w_init_relu, name='u0/c_res{}'.format(i))# 205 | n_res = BatchNormLayer(n_res, act=lrelu, is_train = is_train, gamma_init = g_init, name='u0/b_res{}'.format(i))# 206 | 207 | n = PadLayer(n, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='u0/pad{}_1'.format(i)) 208 | n = Conv2d(n, 64, (3, 3), (1, 1), act=None, padding='VALID', W_init=w_init_relu, name='u0/c{}_1'.format(i)) 209 | n = BatchNormLayer(n, act=lrelu, is_train = is_train, gamma_init = g_init, name='u0/b{}_1'.format(i)) 210 | n = PadLayer(n, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='u0/pad{}_2'.format(i)) 211 | n = Conv2d(n, 64, (3, 3), (1, 1), act=None, padding='VALID', W_init=w_init_relu, name='u0/c{}_2'.format(i)) 212 | n = BatchNormLayer(n, act=lrelu, is_train = is_train, gamma_init = g_init, name='u0/b{}_2'.format(i)) 213 | n = ElementwiseLayer([n, n_res], tf.add, name='u0/add{}'.format(i))# 214 | refine_lists.append(n.outputs) 215 | 216 | n = PadLayer(n, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='uf/pad1')# 217 | n = Conv2d(n, 64, (3, 3), (1, 1), act=None, padding='VALID', W_init=w_init_relu, name='uf/c1')# 218 | n = BatchNormLayer(n, act=lrelu, is_train = is_train, gamma_init = g_init, name='uf/b1')# 219 | n = PadLayer(n, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='uf/pad2')# 220 | n = Conv2d(n, 32, (3, 3), (1, 1), act=None, padding='VALID', W_init=w_init_relu, name='uf/c2')# 221 | n = BatchNormLayer(n, act=lrelu, is_train = is_train, gamma_init = g_init, name='uf/b2')# 222 | n = PadLayer(n, [[0, 0], [1, 1], [1, 1], [0, 0]], "Symmetric", name='uf/pad3')#pad1 223 | n = Conv2d(n, 1, (3, 3), (1, 1), act=None, padding='VALID', W_init=w_init_sigmoid, name='uf/c3')#c1 224 | 225 | return tf.nn.sigmoid(n.outputs), [u4, u3, u2, u1, u0], gan_feat, refine_lists 226 | 227 | def feature_discriminator(feats, is_train=True, reuse=False, scope = 'feature_discriminator'): 228 | w_init = tf.contrib.layers.variance_scaling_initializer() 229 | w_init_sigmoid = tf.contrib.layers.xavier_initializer() 230 | b_init = None 231 | g_init = None 232 | 233 | lrelu = lambda x: tf.nn.leaky_relu(x, 0.2) 234 | with tf.variable_scope(scope, reuse=reuse): 235 | n = InputLayer(feats, name='input_feature') 236 | 237 | n = Conv2d(n, 64, (4, 4), (2, 2), act=None, padding='SAME', W_init=w_init, b_init=b_init, name='h0/c1') 238 | n = BatchNormLayer(n, act=lrelu, is_train = is_train, gamma_init = g_init, name='h0/b1')# 239 | n = Conv2d(n, 128, (4, 4), (2, 2), act=None, padding='SAME', W_init=w_init, b_init=b_init, name='h1/c1') 240 | n = BatchNormLayer(n, act=lrelu, is_train = is_train, gamma_init = g_init, name='h1/b1')# 241 | n = Conv2d(n, 256, (4, 4), (2, 2), act=None, padding='SAME', W_init=w_init, b_init=b_init, name='h2/c1') 242 | n = BatchNormLayer(n, act=lrelu, is_train = is_train, gamma_init = g_init, name='h2/b1')# 243 | n = Conv2d(n, 512, (4, 4), (2, 2), act=None, padding='SAME', W_init=w_init, b_init=b_init, name='h3/c1') 244 | n = BatchNormLayer(n, act=lrelu, is_train = is_train, gamma_init = g_init, name='h3/b1')# 245 | n = Conv2d(n, 1, (4, 4), (2, 2), act=None, padding='SAME', W_init=w_init_sigmoid, b_init=b_init, name='h4/c1') 246 | 247 | logits = n.outputs 248 | 249 | return logits, tf.nn.sigmoid(logits) 250 | 251 | def Binary_Net(input_defocus, is_train=False, reuse=False, scope = 'Binary_Net'): 252 | w_init_relu = tf.contrib.layers.variance_scaling_initializer() 253 | w_init_sigmoid = tf.contrib.layers.xavier_initializer() 254 | b_init = None 255 | g_init = None 256 | lrelu = lambda x: tf.nn.leaky_relu(x, 0.2) 257 | with tf.variable_scope(scope, reuse=reuse): 258 | n = InputLayer(input_defocus, name='input_defocus') 259 | 260 | n = Conv2d(n, 64, (1, 1), (1, 1), act=None, padding='VALID', W_init=w_init_relu, name='l1/c1') 261 | n = BatchNormLayer(n, act=lrelu, is_train = is_train, gamma_init = g_init, name='l1/b1') 262 | n = Conv2d(n, 64, (1, 1), (1, 1), act=None, padding='VALID', W_init=w_init_relu, name='l1/c2') 263 | n = BatchNormLayer(n, act=lrelu, is_train = is_train, gamma_init = g_init, name='l1/b2') 264 | n = Conv2d(n, 64, (1, 1), (1, 1), act=None, padding='VALID', W_init=w_init_relu, name='l1/c3') 265 | n = BatchNormLayer(n, act=lrelu, is_train = is_train, gamma_init = g_init, name='l1/b3') 266 | 267 | n = Conv2d(n, 32, (1, 1), (1, 1), act=None, padding='VALID', W_init=w_init_relu, name='l2/c1') 268 | n = BatchNormLayer(n, act=lrelu, is_train = is_train, gamma_init = g_init, name='l2/b1') 269 | n = Conv2d(n, 32, (1, 1), (1, 1), act=None, padding='VALID', W_init=w_init_relu, name='l2/c2') 270 | n = BatchNormLayer(n, act=lrelu, is_train = is_train, gamma_init = g_init, name='l2/b2') 271 | n = Conv2d(n, 32, (1, 1), (1, 1), act=None, padding='VALID', W_init=w_init_relu, name='l2/c3') 272 | n = BatchNormLayer(n, act=lrelu, is_train = is_train, gamma_init = g_init, name='l2/b3') 273 | 274 | n = Conv2d(n, 16, (1, 1), (1, 1), act=None, padding='VALID', W_init=w_init_relu, name='l3/c1') 275 | n = BatchNormLayer(n, act=lrelu, is_train = is_train, gamma_init = g_init, name='l3/b1') 276 | n = Conv2d(n, 16, (1, 1), (1, 1), act=None, padding='VALID', W_init=w_init_relu, name='l3/c2') 277 | n = BatchNormLayer(n, act=lrelu, is_train = is_train, gamma_init = g_init, name='l3/b2') 278 | n = Conv2d(n, 16, (1, 1), (1, 1), act=None, padding='VALID', W_init=w_init_relu, name='l3/c3') 279 | n = BatchNormLayer(n, act=lrelu, is_train = is_train, gamma_init = g_init, name='l3/b3') 280 | 281 | n = Conv2d(n, 1, (1, 1), (1, 1), act=None, padding='VALID', W_init=w_init_sigmoid, name='l4/c1') 282 | logits = n.outputs 283 | 284 | return logits, tf.nn.sigmoid(n.outputs) 285 | 286 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | easydict==1.9 2 | opencv-python==4.5.1.48 3 | -------------------------------------------------------------------------------- /utils.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import tensorlayer as tl 3 | from tensorlayer.prepro import * 4 | from config import config, log_config 5 | from skimage import feature 6 | from skimage import color 7 | from scipy.ndimage.filters import gaussian_filter 8 | 9 | import scipy 10 | import numpy as np 11 | import cv2 12 | import math 13 | import random 14 | 15 | import os 16 | import fnmatch 17 | 18 | def read_all_imgs(file_name_list, path = '', mode = 'RGB'): 19 | imgs = [] 20 | for idx in range(0, len(file_name_list)): 21 | imgs.append(get_images(file_name_list[idx], path, mode)) 22 | 23 | return imgs 24 | 25 | def get_images(file_name, path, mode): 26 | """ Input an image path and name, return an image array """ 27 | # return scipy.misc.imread(path + file_name).astype(np.float) 28 | if mode is 'RGB': 29 | image = (scipy.misc.imread(path + file_name, mode='RGB')/255.).astype(np.float32) 30 | elif mode is 'GRAY': 31 | image = (scipy.misc.imread(path + file_name, mode='P')/255.).astype(np.float32) 32 | image = np.expand_dims(image, axis = 2) 33 | elif mode is 'NPY': 34 | image = np.load(path + file_name) 35 | image = image / 3.275 36 | image = np.expand_dims(image, axis = 2) 37 | elif mode is 'DEPTH': 38 | image = (np.float32(cv2.imread(path + file_name, cv2.IMREAD_UNCHANGED))/10.)[:, :, 1] 39 | ## If you train the network with the SYNDOF dataset (this is the original SYNDOF dataset) shared in this repository. 40 | ## The SYNDOF datasets's maximum COC value is 15 and we saved the defocus map with the COC value. 41 | ## (The paper say that maximum COC value is 28, becuase the blur kernel of orignal SYNDOF dataset visually had the maximaum coc value of 28 when it was generated with max_coc=15.) 42 | 43 | image = image / 15 44 | 45 | ## If you train the network with the new SYNDOF dataset generated with the codes in "https://github.com/codeslake/SYNDOF". 46 | ## We save the sigma value (max=7) in the code, where 47 | ## sigma = (max_coc-1)/4, when max_coc = 29, max_sigma = 7 48 | 49 | # image = image / 7 50 | 51 | image = np.expand_dims(image, axis = 2) 52 | 53 | return image 54 | 55 | def t_or_f(arg): 56 | ua = str(arg).upper() 57 | if 'TRUE'.startswith(ua): 58 | return True 59 | elif 'FALSE'.startswith(ua): 60 | return False 61 | else: 62 | pass 63 | 64 | def _tf_fspecial_gauss(size, sigma): 65 | """Function to mimic the 'fspecial' gaussian MATLAB function 66 | """ 67 | x_data, y_data = np.mgrid[-size//2 + 1:size//2 + 1, -size//2 + 1:size//2 + 1] 68 | 69 | x_data = np.expand_dims(x_data, axis=-1) 70 | x_data = np.expand_dims(x_data, axis=-1) 71 | 72 | y_data = np.expand_dims(y_data, axis=-1) 73 | y_data = np.expand_dims(y_data, axis=-1) 74 | 75 | x = tf.constant(x_data, dtype=tf.float32) 76 | y = tf.constant(y_data, dtype=tf.float32) 77 | 78 | g = tf.exp(-((x**2 + y**2)/(2.0*sigma**2))) 79 | return g / tf.reduce_sum(g) 80 | 81 | def refine_image(img): 82 | h, w = img.shape[:2] 83 | 84 | return img[0 : h - h % 16, 0 : w - w % 16] 85 | 86 | def random_crop(images, resize_shape, is_gaussian_noise = False): 87 | images_list = None 88 | h, w = resize_shape[:2] 89 | max_size_limit = 800 90 | 91 | for i in np.arange(len(images)): 92 | image = np.copy(images[i]) 93 | shape = np.array(image.shape[:2]) 94 | 95 | if shape.min() <= h: 96 | ratio = resize_shape[shape.argmin()]/float(shape.min()) 97 | resize_w = int(math.floor(shape[1] * ratio)) + 1 98 | resize_h = int(math.floor(shape[0] * ratio)) + 1 99 | image = cv2.resize(image, (resize_w, resize_h)) 100 | 101 | if shape.min() > max_size_limit: 102 | ratio = max_size_limit/float(shape.min()) 103 | resize_w = int(math.floor(shape[1] * ratio)) + 1 104 | resize_h = int(math.floor(shape[0] * ratio)) + 1 105 | image = cv2.resize(image, (resize_w, resize_h)) 106 | 107 | if is_gaussian_noise: 108 | image = add_gaussian_noise(image) 109 | 110 | cropped_image = tl.prepro.crop(image, wrg=w, hrg=h, is_random=True) 111 | augmented_image = _random_flip(cropped_image) 112 | angles = np.array([1, 2, 3, 4]) 113 | angle = np.random.choice(angles) 114 | augmented_image = _random_rotation(augmented_image, angle) 115 | image = np.expand_dims(augmented_image, axis=0) 116 | 117 | images_list = np.copy(image) if i == 0 else np.concatenate((images_list, image), axis = 0) 118 | 119 | return images_list 120 | 121 | def crop_pair_with_different_shape_images(images, labels, resize_shape, is_gaussian_noise = False): 122 | images_list = None 123 | labels_list = None 124 | h, w = resize_shape[:2] 125 | max_size_limit = 800 126 | 127 | for i in np.arange(len(images)): 128 | image = np.copy(images[i]) 129 | label = np.copy(labels[i]) 130 | shape = np.array(image.shape[:2]) 131 | 132 | if shape.min() <= h: 133 | ratio = resize_shape[shape.argmin()]/float(shape.min()) 134 | resize_w = int(math.floor(shape[1] * ratio)) + 1 135 | resize_h = int(math.floor(shape[0] * ratio)) + 1 136 | image = cv2.resize(image, (resize_w, resize_h)) 137 | label = np.expand_dims(cv2.resize(label[:, :, 0], (resize_w, resize_h)), axis = 2) 138 | 139 | if shape.min() > max_size_limit: 140 | ratio = max_size_limit/float(shape.min()) 141 | resize_w = int(math.floor(shape[1] * ratio)) + 1 142 | resize_h = int(math.floor(shape[0] * ratio)) + 1 143 | image = cv2.resize(image, (resize_w, resize_h)) 144 | label = np.expand_dims(cv2.resize(label[:, :, 0], (resize_w, resize_h)), axis = 2) 145 | 146 | if is_gaussian_noise: 147 | image = add_gaussian_noise(image) 148 | 149 | concatenated_images = np.concatenate((image, label), axis = 2) 150 | cropped_images = tl.prepro.crop(concatenated_images, wrg=w, hrg=h, is_random=True) 151 | augmented_images = _random_flip(cropped_images) 152 | angles = np.array([1, 2, 3, 4]) 153 | angle = np.random.choice(angles) 154 | augmented_images = _random_rotation(augmented_images, angle) 155 | 156 | image = np.expand_dims(augmented_images[:, :, 0:3], axis=0) 157 | label = np.expand_dims(np.expand_dims(augmented_images[:, :, 3], axis=3), axis=0) 158 | 159 | images_list = np.copy(image) if i == 0 else np.concatenate((images_list, image), axis = 0) 160 | labels_list = np.copy(label) if i == 0 else np.concatenate((labels_list, label), axis = 0) 161 | 162 | return images_list, labels_list 163 | 164 | def add_gaussian_noise(image): 165 | image = image.astype(np.float32) 166 | shape = image.shape[:2] 167 | 168 | mean = 0 169 | var = random.uniform(0,0.1) 170 | sigma = var ** 0.5 171 | gamma = 0.25 172 | alpha = 0.75 173 | beta = 1 - alpha 174 | 175 | gaussian = np.random.normal(loc=mean, scale = sigma, size = (shape[0], shape[1], 1)).astype(np.float32) 176 | gaussian = np.concatenate((gaussian, gaussian, gaussian), axis = 2) 177 | #gaussian_img = image * 0.75 + 0.25 * gaussian + 0.25 178 | gaussian_img = cv2.addWeighted(image, alpha, beta * gaussian, beta, gamma) 179 | 180 | return gaussian_img 181 | 182 | # noise_sigma = 0.01 183 | # h = image.shape[0] 184 | # w = image.shape[1] 185 | # noise = np.random.randn(h, w) * noise_sigma 186 | 187 | # noisy_image = np.zeros(image.shape, np.float64) 188 | # if len(image.shape) == 2: 189 | # noisy_image = image + noise 190 | # else: 191 | # noisy_image[:,:,0] = image[:,:,0] + noise 192 | # noisy_image[:,:,1] = image[:,:,1] + noise 193 | # noisy_image[:,:,2] = image[:,:,2] + noise 194 | 195 | # """ 196 | # print('min,max = ', np.min(noisy_image), np.max(noisy_image)) 197 | # print('type = ', type(noisy_image[0][0][0])) 198 | # """ 199 | 200 | # return noisy_image 201 | 202 | def _random_flip(images): 203 | flipped_images = tl.prepro.flip_axis(images, axis=0, is_random=True) 204 | 205 | return flipped_images 206 | 207 | def _random_rotation(images, angle): 208 | if angle != 4: 209 | rotated_images = np.rot90(images, angle) 210 | else: 211 | rotated_images = images 212 | 213 | return rotated_images 214 | 215 | def _get_file_path(path, regex): 216 | file_path = [] 217 | for root, dirnames, filenames in os.walk(path): 218 | for i in np.arange(len(regex)): 219 | for filename in fnmatch.filter(filenames, regex[i]): 220 | file_path.append(os.path.join(root, filename)) 221 | 222 | return file_path 223 | 224 | def remove_file_end_with(path, regex): 225 | file_paths = _get_file_path(path, [regex]) 226 | 227 | for i in np.arange(len(file_paths)): 228 | os.remove(file_paths[i]) 229 | 230 | def save_images(images, size, image_path='_temp.png'): 231 | if len(images.shape) == 3: # Greyscale [batch, h, w] --> [batch, h, w, 1] 232 | images = images[:, :, :, np.newaxis] 233 | 234 | def merge(images, size): 235 | h, w = images.shape[1], images.shape[2] 236 | img = np.zeros((h * size[0], w * size[1], 3)) 237 | for idx, image in enumerate(images): 238 | i = idx % size[1] 239 | j = idx // size[1] 240 | img[j * h:j * h + h, i * w:i * w + w, :] = image 241 | return img 242 | 243 | def imsave(images, size, path): 244 | return scipy.misc.toimage(merge(images, size), cmin = 0., cmax = 1.).save(path) 245 | 246 | assert len(images) <= size[0] * size[1], "number of images should be equal or less than size[0] * size[1] {}".format(len(images)) 247 | 248 | return imsave(images, size, image_path) 249 | 250 | def fix_image_tf(image, norm_value): 251 | return tf.cast(image / norm_value * 255., tf.uint8) 252 | 253 | def norm_image_tf(image): 254 | image = image - tf.reduce_min(image, axis = [1, 2, 3], keepdims=True) 255 | image = image / tf.reduce_max(image, axis = [1, 2, 3], keepdims=True) 256 | return tf.cast(image * 255., tf.uint8) 257 | 258 | def norm_image(image, axis = (1, 2, 3)): 259 | image = image - np.amin(image, axis = axis, keepdims=True) 260 | image = image / np.amax(image, axis = axis, keepdims=True) 261 | return image 262 | 263 | def get_disc_accuracy(logits, labels): 264 | acc = 0. 265 | for i in np.arange(len(logits)): 266 | tp = 0 267 | logits[i] = np.round(np.squeeze(logits[i])).astype(int) 268 | temp = logits[i] 269 | tp = tp + len(temp[np.where(temp == labels[i])]) 270 | acc = acc + (tp / float(len(logits[i]))) 271 | return acc / float(len(labels)) 272 | 273 | --------------------------------------------------------------------------------