-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathmake_grid.py
114 lines (81 loc) · 3.43 KB
/
make_grid.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
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
import argparse
import os
import sys
import cv2
import numpy as np
import torch
import torch.nn.functional as F
from torchvision import transforms
from torchvision.io import read_image, ImageReadMode
from torchvision.utils import make_grid
def edit(orig_img, recon_img):
recon_img[recon_img <= 0.05 * torch.max(orig_img)] = 0
orig_img[orig_img <= 0.05 * torch.max(orig_img)] = 0
orig_np = orig_img.squeeze().cpu().numpy()
recon_np = recon_img.squeeze().cpu().numpy()
mask = np.zeros(recon_np.shape, dtype=np.uint8)
cv2.circle(mask, (192, 192), 165, 255, -1)
recon_np = cv2.bitwise_and(recon_np, recon_np, mask=mask)
return orig_img, recon_img, orig_np, recon_np
def get_args():
parser = argparse.ArgumentParser()
parser.add_argument('--exp_dir', type=str, required=True, help='Path to experiment images')
args = parser.parse_args()
return args
def main():
args = get_args()
to_plot = []
for i in range(4):
origi = read_image(os.path.join(args.exp_dir, f'brain_T2_{i}_or.jpg'), mode=ImageReadMode.GRAY)
recon = read_image(os.path.join(args.exp_dir, f'brain_T2_{i}.jpg'))
_, _, orig_np, recon_np = edit(origi, recon)
pad = (1,1,1,1)
origi = F.pad(origi, pad, "constant", 255)
recon = F.pad(recon, pad, "constant", 255)
to_plot.append(origi)
to_plot.append(recon)
grid = make_grid(to_plot, nrow=2, ncol=4, padding=0)
grid = transforms.ToPILImage()(grid)
grid.save(os.path.join(args.exp_dir, f'brains_grid0.jpg'))
to_plot = []
for i in range(4,8):
origi = read_image(os.path.join(args.exp_dir, f'brain_T2_{i}_or.jpg'), mode=ImageReadMode.GRAY)
recon = read_image(os.path.join(args.exp_dir, f'brain_T2_{i}.jpg'))
_, _, orig_np, recon_np = edit(origi, recon)
pad = (1,1,1,1)
origi = F.pad(origi, pad, "constant", 255)
recon = F.pad(recon, pad, "constant", 255)
to_plot.append(origi)
to_plot.append(recon)
grid = make_grid(to_plot, nrow=2, ncol=4, padding=0)
grid = transforms.ToPILImage()(grid)
grid.save(os.path.join(args.exp_dir, f'brains_grid1.jpg'))
to_plot = []
for i in range(8, 12):
origi = read_image(os.path.join(args.exp_dir, f'brain_T2_{i}_or.jpg'), mode=ImageReadMode.GRAY)
recon = read_image(os.path.join(args.exp_dir, f'brain_T2_{i}.jpg'))
_, _, orig_np, recon_np = edit(origi, recon)
pad = (1,1,1,1)
origi = F.pad(origi, pad, "constant", 255)
recon = F.pad(recon, pad, "constant", 255)
to_plot.append(origi)
to_plot.append(recon)
grid = make_grid(to_plot, nrow=2, ncol=4, padding=0)
grid = transforms.ToPILImage()(grid)
grid.save(os.path.join(args.exp_dir, f'brains_grid2.jpg'))
to_plot = []
for i in range(12, 16):
origi = read_image(os.path.join(args.exp_dir, f'brain_T2_{i}_or.jpg'), mode=ImageReadMode.GRAY)
recon = read_image(os.path.join(args.exp_dir, f'brain_T2_{i}.jpg'))
_, _, orig_np, recon_np = edit(origi, recon)
pad = (1,1,1,1)
origi = F.pad(origi, pad, "constant", 255)
recon = F.pad(recon, pad, "constant", 255)
to_plot.append(origi)
to_plot.append(recon)
grid = make_grid(to_plot, nrow=2, ncol=4, padding=0)
grid = transforms.ToPILImage()(grid)
grid.save(os.path.join(args.exp_dir, f'brains_grid3.jpg'))
return 0
if __name__ == '__main__':
sys.exit(main())