Skip to content

Commit 8789c95

Browse files
procyontaoprocyontao
committed
slightly reduce the tile artifacts
1 parent 531ae7e commit 8789c95

File tree

3 files changed

+110
-52
lines changed

3 files changed

+110
-52
lines changed

‎models/unet/predict.py‎

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -83,10 +83,13 @@ def predict_one(args,one_tomo,output_file=None):
8383
with mrcfile.open(one_tomo,permissive=True) as mrcData:
8484
real_data = mrcData.data.astype(np.float32)*-1
8585
voxelsize = mrcData.voxel_size
86-
real_data = normalize(real_data,percentile=args.normalize_percentile)
87-
data=np.expand_dims(real_data,axis=-1)
88-
reform_ins = reform3D(data)
89-
data = reform_ins.pad_and_crop_new(args.cube_size,args.crop_size)
86+
data = normalize(real_data,percentile=args.normalize_percentile)
87+
#data=np.expand_dims(real_data,axis=-1)
88+
#reform_ins = reform3D(data)
89+
#data = reform_ins.pad_and_crop_new(args.cube_size,args.crop_size)
90+
#print(data.shape)
91+
reform_ins = reform3D(data,args.cube_size,args.crop_size,9)
92+
data = reform_ins.pad_and_crop()
9093
#to_predict_data_shape:(n,cropsize,cropsize,cropsize,1)
9194
#imposing wedge to every cubes
9295
#data=wedge_imposing(data)
@@ -105,11 +108,11 @@ def predict_one(args,one_tomo,output_file=None):
105108
in_data = data[i*N:(i+1)*N]
106109
# in_data_gen = get_gen_single(in_data,args.batch_size,shuffle=False)
107110
# in_data_tf = tf.data.Dataset.from_generator(in_data_gen,output_types=tf.float32)
108-
outData[i*N:(i+1)*N] = model.predict(in_data,verbose=0)
111+
outData[i*N:(i+1)*N] = model.predict(in_data,verbose=0).squeeze()
109112
outData = outData[0:num_patches]
110113

111-
outData=reform_ins.restore_from_cubes_new(outData.reshape(outData.shape[0:-1]), args.cube_size, args.crop_size)
112-
114+
#outData=reform_ins.restore_from_cubes_new(outData.reshape(outData.shape[0:-1]), args.cube_size, args.crop_size)
115+
outData=reform_ins.restore(outData)
113116
outData = normalize(outData,percentile=args.normalize_percentile)
114117
with mrcfile.new(output_file, overwrite=True) as output_mrc:
115118
output_mrc.set_data(-outData)

‎util/deconvolution.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ def wiener1d(angpix, voltage, cs, defocus, snrfalloff, deconvstrength, highpassn
4444
wiener = ctf/(ctf*ctf+1/snr)
4545
return ctf, wiener
4646

47-
def tom_deconv_tomo(vol_file, out_file,angpix, voltage, cs, defocus, snrfalloff, deconvstrength, highpassnyquist, phaseflipped, phaseshift, ncpu=8):
47+
def tom_deconv_tomo(vol_file, out_file, angpix, voltage, cs, defocus, snrfalloff, deconvstrength, highpassnyquist, phaseflipped, phaseshift, ncpu=8):
4848
with mrcfile.open(vol_file, permissive=True) as f:
4949
header_in = f.header
5050
vol = f.data

‎util/toTile.py‎

Lines changed: 99 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -1,75 +1,130 @@
11
import numpy as np
22

33
class reform3D:
4-
def __init__(self,data3D):
5-
self._sp = data3D.shape
4+
def __init__(self,data3D, cubesize, cropsize, edge_depth):
5+
self._sp = np.array(data3D.shape)
66
self._orig_data = data3D
7+
self.cubesize = cubesize
8+
self.cropsize = cropsize
9+
self.edge_depth = edge_depth
10+
self._sidelen = np.ceil((self._sp + edge_depth * 2)/self.cubesize).astype(int)
11+
#self._sidelen = np.ceil((1.*self._sp)/self.cubesize).astype(int)
712

8-
def pad_and_crop_new(self, cubesize=32, cropsize = 64):
9-
self._cropsize = cropsize
10-
sp = np.array(self._sp)
11-
self._sidelen = sp//cubesize+1
12-
padi = int((cropsize - cubesize)/2)
13-
padsize = (self._sidelen*cubesize + padi - sp).astype(int)
14-
data = np.pad(self._orig_data,((padi,padsize[0]),(padi,padsize[1]),(padi,padsize[2]),(0,0)),'symmetric')
13+
def pad_and_crop(self):
14+
15+
#----------------------------|---------------------------
16+
#| |
17+
#| ---------------|------|--------edge---------------
18+
#| | ------------|------|--------image_edge----------
19+
#| | | | |
20+
#| | | | |
21+
#| | | | |
22+
#| | | | |
23+
#| -----cube------- |
24+
#| | | |
25+
#| | | |
26+
#|----------crop--------------
27+
#| | |
28+
#| | |
29+
#| | |
30+
#| | |
31+
pad_left = int((self.cropsize - self.cubesize)/2 + self.edge_depth)
32+
33+
# pad_right + pad_left + shape = sidelen * cube_zize + (crop_size-cube_size)
34+
# pad_right + pad_left + shape >= (self._sp + edge_depth * 2) + (crop_size-cube_size)
35+
pad_right = (self._sidelen * self.cubesize + (self.cropsize-self.cubesize) - pad_left - self._sp).astype(int)
36+
37+
data = np.pad(self._orig_data,((pad_left,pad_right[0]),(pad_left,pad_right[1]),(pad_left,pad_right[2])),'symmetric')
1538
outdata=[]
1639

1740
for i in range(self._sidelen[0]):
1841
for j in range(self._sidelen[1]):
1942
for k in range(self._sidelen[2]):
20-
cube = data[i*cubesize:i*cubesize+cropsize,
21-
j*cubesize:j*cubesize+cropsize,
22-
k*cubesize:k*cubesize+cropsize]
43+
cube = data[i*self.cubesize:i*self.cubesize+self.cropsize,
44+
j*self.cubesize:j*self.cubesize+self.cropsize,
45+
k*self.cubesize:k*self.cubesize+self.cropsize]
2346
outdata.append(cube)
2447
outdata=np.array(outdata)
2548
return outdata
49+
50+
def mask(self, x_len, y_len, z_len):
51+
# need to consider should partisioned to len+1 so that left and right can add to one
52+
p = 2*self.edge_depth#(self.cropsize - self.cubesize)
53+
assert x_len > 2*p
54+
assert y_len > 2*p
55+
assert z_len > 2*p
2656

57+
array_x = np.minimum(np.arange(x_len+1), p) / p
58+
array_x = array_x * np.flip(array_x)
59+
array_x = array_x[np.newaxis,np.newaxis,:]
2760

61+
array_y = np.minimum(np.arange(y_len+1), p) / p
62+
array_y = array_y * np.flip(array_y)
63+
array_y = array_y[np.newaxis,:,np.newaxis]
2864

65+
array_z = np.minimum(np.arange(z_len+1), p) / p
66+
array_z = array_z * np.flip(array_z)
67+
array_z = array_z[:,np.newaxis,np.newaxis]
2968

30-
def pad_and_crop(self,cropsize=(64,64,64)):
31-
self._cropsize = cropsize
32-
sp = np.array(self._sp)
33-
padsize = (sp//64+1)*64-sp
34-
data = np.pad(self._orig_data,((0,padsize[0]),(0,padsize[1]),(0,padsize[2]),(0,0)),'edge')
35-
self._sidelen = (padsize+sp)//64
69+
out = array_x * array_y * array_z
70+
return out[:x_len,:y_len,:z_len]
3671

37-
outdata=[]
38-
for i in range(self._sidelen[0]):
39-
for j in range(self._sidelen[1]):
40-
for k in range(self._sidelen[2]):
41-
cube = data[i*cropsize[0]:(i+1)*cropsize[0],j*cropsize[0]:(j+1)*cropsize[0],k*cropsize[0]:(k+1)*cropsize[0]]
42-
outdata.append(cube)
43-
outdata=np.array(outdata)
44-
return outdata
4572

46-
def restore_from_cubes(self,cubes):
47-
if len(cubes.shape)==5 and cubes.shape[-1]==1:
48-
cubes = cubes.reshape(cubes.shape[0:-1])
49-
new = np.zeros((self._sidelen[0]*64,self._sidelen[1]*64,self._sidelen[2]*64))
73+
def restore(self,cubes):
74+
75+
start = (self.cropsize-self.cubesize)//2-self.edge_depth
76+
end = (self.cropsize-self.cubesize)//2+self.cubesize+self.edge_depth
77+
cubes = cubes[:,start:end,start:end,start:end]
78+
79+
restored = np.zeros((self._sidelen[0]*self.cubesize+self.edge_depth*2,
80+
self._sidelen[1]*self.cubesize+self.edge_depth*2,
81+
self._sidelen[2]*self.cubesize+self.edge_depth*2))
82+
print("size restored", restored.shape)
83+
mask_cube = self.mask(self.cubesize+self.edge_depth*2,self.cubesize+self.edge_depth*2,self.cubesize+self.edge_depth*2)
5084
for i in range(self._sidelen[0]):
5185
for j in range(self._sidelen[1]):
5286
for k in range(self._sidelen[2]):
53-
new[i*self._cropsize[0]:(i+1)*self._cropsize[0],j*self._cropsize[0]:(j+1)*self._cropsize[0],k*self._cropsize[0]:(k+1)*self._cropsize[0]] \
54-
= cubes[i*self._sidelen[1]*self._sidelen[2]+j*self._sidelen[1]+k]
55-
return new[0:self._sp[0],0:self._sp[1],0:self._sp[2]]
87+
restored[i*self.cubesize:(i+1)*self.cubesize+self.edge_depth*2,
88+
j*self.cubesize:(j+1)*self.cubesize+self.edge_depth*2,
89+
k*self.cubesize:(k+1)*self.cubesize+self.edge_depth*2] \
90+
+= cubes[i*self._sidelen[1]*self._sidelen[2]+j*self._sidelen[2]+k]\
91+
*mask_cube
92+
93+
94+
p =self.edge_depth*2 #int((self.cropsize-self.cubesize)/2+self.edge_depth)
95+
restored = restored[p:p+self._sp[0],p:p+self._sp[1],p:p+self._sp[2]]
96+
return restored
5697

57-
def restore_from_cubes_new(self,cubes, cubesize = 32, cropsize = 64):
58-
if len(cubes.shape)==5 and cubes.shape[-1]==1:
59-
cubes = cubes.reshape(cubes.shape[0:-1])
98+
def mask_old(self):
99+
from functools import reduce
100+
c = self.cropsize
101+
p = (self.cropsize - self.cubesize)
102+
mask = np.ones((c, c, c))
103+
f = lambda x: min(x, p)/p
104+
for i in range(c):
105+
for j in range(c):
106+
for k in range(c):
107+
d = [i, c-i, j, c-j, k, c-k]
108+
d = map(f,d)
109+
d = reduce(lambda a,b: a*b, d)
110+
mask[i,j,k] = d
111+
return mask
112+
def restore_from_cubes(self,cubes):
60113

61-
new = np.zeros((self._sidelen[0]*cubesize,self._sidelen[1]*cubesize,self._sidelen[2]*cubesize))
62-
start=int((cropsize-cubesize)/2)
63-
end=int((cropsize+cubesize)/2)
114+
new = np.zeros((self._sidelen[0]*self.cubesize,
115+
self._sidelen[1]*self.cubesize,
116+
self._sidelen[2]*self.cubesize))
117+
start=int((self.cropsize-self.cubesize)/2)
118+
end=int((self.cropsize+self.cubesize)/2)
64119

65120
for i in range(self._sidelen[0]):
66121
for j in range(self._sidelen[1]):
67122
for k in range(self._sidelen[2]):
68-
new[i*cubesize:(i+1)*cubesize,j*cubesize:(j+1)*cubesize,k*cubesize:(k+1)*cubesize] \
69-
= cubes[i*self._sidelen[1]*self._sidelen[2]+j*self._sidelen[2]+k][start:end,start:end,start:end]
123+
new[i*self.cubesize:(i+1)*self.cubesize,
124+
j*self.cubesize:(j+1)*self.cubesize,
125+
k*self.cubesize:(k+1)*self.cubesize] \
126+
= cubes[i*self._sidelen[1]*self._sidelen[2]+j*self._sidelen[2]+k][start:end,start:end,start:end]
70127
return new[0:self._sp[0],0:self._sp[1],0:self._sp[2]]
71-
72-
73128
def pad4times(self,time=4):
74129
sp = np.array(self._orig_data.shape)
75130
sp = np.expand_dims(sp,axis=0)
@@ -86,4 +141,4 @@ def cropback(self,padded):
86141
return padded[:orig_sp[0][0]][:orig_sp[0][1]][:orig_sp[0][2]]
87142

88143
if __name__ == '__main__':
89-
pass
144+
pass

0 commit comments

Comments
 (0)