-
Notifications
You must be signed in to change notification settings - Fork 0
/
reconstruct.py
72 lines (54 loc) · 2.53 KB
/
reconstruct.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
import matplotlib.pyplot as plt
import os
import numpy as np
import sys
import tensorflow as tf
from utils import utils
def reconstruct_img(model, iterator, dataio, is_plotting=False, img_folder="./", prefix_name=""):
steps_per_execution = dataio.data_dim[0] // dataio.batch_size
orig_imgs = np.array([]).reshape([0]+list(dataio.data_dim[1:]))
recon_imgs = np.array([]).reshape([0]+list(dataio.data_dim[1:]))
for i in range(steps_per_execution):
cur_data = next(iterator)
recon_img, _ = model.predict(cur_data)
orig_imgs = np.vstack([orig_imgs, cur_data])
recon_imgs = np.vstack([recon_imgs, recon_img])
orig_imgs = orig_imgs[:-dataio.num_added_data]
recon_imgs = recon_imgs[:-dataio.num_added_data]
normalizer = dataio.normalizer
padder = dataio.padder
padded_img_shape = list(padder.padded_img_shape) + [1]
num_tile_img = np.prod(padded_img_shape) // (np.square(padder.tile_size))
num_img = recon_imgs.shape[0] // num_tile_img
orig_imgs = np.array([padder.unsplit_image(orig_imgs[i*num_tile_img:(i+1)*num_tile_img], padded_img_shape)
for i in range(num_img)])
recon_imgs = np.array([padder.unsplit_image(recon_imgs[i*num_tile_img:(i+1)*num_tile_img], padded_img_shape)
for i in range(num_img)])
orig_imgs = padder.remove_pad_ds(orig_imgs)
recon_imgs = padder.remove_pad_ds(recon_imgs)
# Denormalizer
orig_imgs = normalizer.denormalize_log10(orig_imgs)
orig_imgs = normalizer.denormalize_minmax(orig_imgs)
recon_imgs = normalizer.denormalize_log10(recon_imgs)
recon_imgs = normalizer.denormalize_minmax(recon_imgs)
if is_plotting:
generate_plots(orig_imgs, recon_imgs, img_folder, num_img=3, prefix_name=prefix_name)
return (orig_imgs, recon_imgs)
def generate_plots(orig_imgs, recon_imgs, img_folder, num_img, prefix_name=""):
utils.mkdir_if_not_exist(img_folder)
fig1 = plt.figure(figsize=(15, 13))
cmap = 'viridis'
num_img = num_img
for i in range(num_img):
plt.subplot(num_img, 1, i + 1)
plt.imshow(orig_imgs[i], cmap=cmap)
plt.axis('off')
plt.savefig(os.path.join(img_folder, prefix_name+"original_images.png"))
fig2 = plt.figure(figsize=(15, 13))
for i in range(num_img):
plt.subplot(num_img, 1, i + 1)
plt.imshow(recon_imgs[i], cmap=cmap)
plt.axis('off')
# tight_layout minimizes the overlap between 2 sub-plots
plt.savefig(os.path.join(img_folder, prefix_name+"generate_image.png"))
plt.show()