-
Notifications
You must be signed in to change notification settings - Fork 1
/
prepare.py
100 lines (78 loc) · 2.87 KB
/
prepare.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
import numpy as np
from numpy import genfromtxt
from numpy import asarray
import math
import copy
import os
from PIL import Image
import cv2
patch_size = 64 #input = 64x64
label_size = 128 #output = 128x128
#get RGGB bayer image
def bayer_reverse(img):
height,width,c = img.shape;
tmp = np.zeros([height,width]);
for i in range( height ):
for j in range( width ):
if i % 2 == 0 :
if j % 2 == 0:
tmp[i][j] = img[i][j][0];#R
else:
tmp[i][j] = img[i][j][1];#G
else :
if j % 2 == 0:
tmp[i][j] = img[i][j][1];#G
else:
tmp[i][j] = img[i][j][2];#B
return tmp;
#split image to prepare the train set
def split(img,name):
height,width,c = img.shape;
print(img.shape)
count = 0;
for i in range(0,height,30):
for j in range(0,width,30):
if( i + label_size < height and j + label_size < width ):
tmp = np.zeros([label_size,label_size,3]);
tmp2 = np.zeros([label_size,label_size,3]);
tmp = img[ i : i + label_size, j : j + label_size,:];
#save splite label
path = 'label/'+name.split('.')[0] +'_'+str(count)+'.png';
#tmp = tmp/255
im = Image.fromarray(tmp)
#tmp2 = tmp
print(tmp2.shape)
tmp2[:,:,0] = tmp[:,:,2]
tmp2[:,:,1] = tmp[:,:,1]
tmp2[:,:,2] = tmp[:,:,0]
#im.save(path)
cv2.imwrite(path,tmp2)
zoom = im.resize((patch_size,patch_size))
zoom2 = np.zeros([patch_size,patch_size,3]);
gray = np.zeros([patch_size,patch_size]);
zoom = np.array(zoom)
zoom2[:,:,0] = zoom[:,:,2]
zoom2[:,:,1] = zoom[:,:,1]
zoom2[:,:,2] = zoom[:,:,0]
gray = bayer_reverse(zoom2)
#gray = gray/255
path = 'patch/'+name.split('.')[0] +'_'+str(count)+'.png';
im = Image.fromarray(gray)
#im = im.convert("L")
#im.save(path)
cv2.imwrite(path,gray)
count = count + 1
def main():
if not os.path.exists('patch'):
os.makedirs('patch')
if not os.path.exists('label'):
os.makedirs('label')
entries = os.listdir('origin/')
for entry in entries:
print(entry)
path = 'origin/'+entry
img = Image.open(path)
img = np.array(img)
split(img,entry)
if __name__ == '__main__':
main()