Coverage for crip\shared.py: 75%

1''' 

2 Shared module of crip. 

3 

4 https://github.com/SEU-CT-Recon/crip 

5''' 

6 

7import numpy as np 

8from os import path 

9import cv2 

10import scipy.ndimage 

11 

12from .utils import * 

13from ._typing import * 

14from .io import imreadTiff 

15 

16 

17@ConvertListNDArray 

18def rotate(img: TwoOrThreeD, deg: int) -> TwoOrThreeD: 

19 ''' Rotate each image by deg [DEG] (multiple of 90) clockwise. 

20 ''' 

21 cripAssert(deg % 90 == 0, f'`deg` should be multiple of 90, but got {deg}.') 

22 

23 k = int(deg % 360) // 90 

24 axes = (1, 2) if is3D(img) else (0, 1) 

25 

26 return np.rot90(img, -k, axes) 

27 

28 

29@ConvertListNDArray 

30def verticalFlip(img: TwoOrThreeD, copy=False) -> TwoOrThreeD: 

31 ''' Vertical flip each image. 

32 ''' 

33 cripAssert(is2or3D(img), f'img should be 2D or 3D, but got {img.ndim}-D.') 

34 

35 if copy: 

36 return np.array(img[..., ::-1, :]) 

37 else: 

38 return img[..., ::-1, :] 

39 

40 

41@ConvertListNDArray 

42def horizontalFlip(img: TwoOrThreeD, copy=False) -> TwoOrThreeD: 

43 ''' Horizontal flip each image. 

44 ''' 

45 cripAssert(is2or3D(img), f'img should be 2D or 3D, but got {img.ndim}-D.') 

46 

47 if copy: 

48 return np.array(img[..., ::-1]) 

49 else: 

50 return img[..., ::-1] 

51 

52 

53@ConvertListNDArray 

54def stackFlip(img: ThreeD, copy=False) -> ThreeD: 

55 ''' Flip a stack w.r.t. z-axis, i.e., reverse slices. 

56 ''' 

57 cripAssert(is3D(img), 'img should be 3D.') 

58 

59 if copy: 

60 return np.array(np.flip(img, axis=0)) 

61 else: 

62 return np.flip(img, axis=0) 

63 

64 

65@ConvertListNDArray 

66def resizeTo(img: TwoOrThreeD, dsize: Tuple[int], interp='bicubic') -> TwoOrThreeD: 

67 ''' Resize each image to `dsize=(H, W)` using `interp` [bicubic, linear, nearest]. 

68 ''' 

69 cripAssert(interp in ['bicubic', 'linear', 'nearest'], f'Invalid interp method: {interp}.') 

70 cripAssert(len(dsize) == 2, '`dsize` should be 2D.') 

71 

72 dsize = dsize[::-1] # OpenCV dsize is in (W, H) form, so we reverse it. 

73 interp_ = {'bicubic': cv2.INTER_CUBIC, 'linear': cv2.INTER_LINEAR, 'nearest': cv2.INTER_NEAREST} 

74 

75 if is3D(img): 

76 res = [cv2.resize(img[i, ...], dsize, None, interpolation=interp_[interp]) for i in range(img.shape[0])] 

77 return np.array(res) 

78 else: 

79 return cv2.resize(img, dsize, None, interpolation=interp_[interp]) 

80 

81 

82@ConvertListNDArray 

83def resizeBy(img: TwoOrThreeD, factor: Or[Tuple[float], float], interp='bicubic') -> TwoOrThreeD: 

84 ''' Resize each slice in img by `factor=(fH, fW)` or `(f, f)` float using `interp` [bicubic, linear, nearest]. 

85 ''' 

86 cripAssert(interp in ['bicubic', 'linear', 'nearest'], f'Invalid interp method: {interp}.') 

87 if not isType(factor, Tuple): 

88 factor = (factor, factor) 

89 cripAssert(len(factor) == 2, '`factor` should be 2D.') 

90 

91 interp_ = {'bicubic': cv2.INTER_CUBIC, 'linear': cv2.INTER_LINEAR, 'nearest': cv2.INTER_NEAREST} 

92 fH, fW = factor 

93 

94 if is3D(img): 

95 res = [cv2.resize(img[i, ...], None, None, fW, fH, interpolation=interp_[interp]) for i in range(img.shape[0])] 

96 return np.array(res) 

97 else: 

98 print(img, fW, fH, interp_[interp]) 

99 return cv2.resize(img, None, None, fW, fH, interpolation=interp_[interp]) 

100 

101 

102@ConvertListNDArray 

103def resize3D(img: ThreeD, factor: Tuple[int], order=3) -> ThreeD: 

104 ''' Resize 3D `img` by `factor=(fSlice, fH, fW)` using `order`-spline interpolation. 

105 ''' 

106 cripAssert(len(factor) == 3, '`factor` should be 3D.') 

107 

108 return scipy.ndimage.zoom(img, factor, None, order, mode='mirror') 

109 

110 

111@ConvertListNDArray 

112def gaussianSmooth(img: TwoOrThreeD, 

113 sigma: Or[float, int, Tuple[Or[float, int]]], 

114 ksize: Or[Tuple[int], int, None] = None) -> TwoOrThreeD: 

115 ''' Perform Gaussian smooth with kernel size `ksize` and Gaussian sigma = `sigma` [int or tuple (row, col)]. 

116 Leave `ksize=None` to auto determine it to include the majority of kernel energy. 

117 ''' 

118 if ksize is not None: 

119 if not isType(ksize, Sequence): 

120 cripAssert(isInt(ksize), 'ksize should be int or sequence of int.') 

121 ksize = (ksize, ksize) 

122 else: 

123 cripAssert(len(ksize) == 2, 'ksize should have length 2.') 

124 cripAssert(isInt(ksize[0]) and isInt(ksize[1]), 'ksize should be int or sequence of int.') 

125 

126 if not isType(sigma, Sequence): 

127 sigma = (sigma, sigma) 

128 

129 if is3D(img): 

130 res = [cv2.GaussianBlur(img[i, ...], ksize, sigmaX=sigma[1], sigmaY=sigma[0]) for i in range(img.shape[0])] 

131 return np.array(res) 

132 else: 

133 return cv2.GaussianBlur(img, ksize, sigmaX=sigma[1], sigmaY=sigma[0]) 

134 

135 

136@ConvertListNDArray 

137def stackImages(imgList: ListNDArray, dtype=None) -> NDArray: 

138 ''' Stack seperate image into one numpy array. I.e., views * (h, w) -> (views, h, w). 

139 Convert dtype when `dtype` is not `None`. 

140 ''' 

141 if dtype is not None: 

142 imgList = imgList.astype(dtype) 

143 

144 return imgList 

145 

146 

147def splitImages(imgs: ThreeD, dtype=None) -> ListNDArray: 

148 ''' Split stacked image into seperate numpy arrays. I.e., (views, h, w) -> views * (h, w). 

149 Convert dtype when `dtype` is not `None`. 

150 ''' 

151 cripAssert(is3D(imgs), 'imgs should be 3D.') 

152 

153 if dtype is not None: 

154 imgs = imgs.astype(dtype) 

155 

156 return list(imgs) 

157 

158 

159@ConvertListNDArray 

160def binning(img: TwoOrThreeD, rates: Tuple[int]) -> TwoOrThreeD: 

161 ''' Perform binning with `rates = (c, h, w) / (h, w)`. 

162 ''' 

163 for rate in rates: 

164 cripAssert(isInt(rate) and rate > 0, 'rates should be positive int.') 

165 if is2D(img): 

166 cripAssert(len(rates) == 2, 'img is 2D, while rates is 3D.') 

167 if is3D(img): 

168 cripAssert(len(rates) == 3, 'img is 3D, while rates is 2D.') 

169 

170 if rates[-1] != 1: 

171 img = img[..., ::rates[-1]] 

172 if rates[-2] != 1: 

173 img = img[..., ::rates[-2], :] 

174 if len(rates) == 3 and rates[0] != 1: 

175 img = img[::rates[0], ...] 

176 

177 return img 

178 

179 

180@ConvertListNDArray 

181def transpose(vol: TwoOrThreeD, order: Tuple[int]) -> TwoOrThreeD: 

182 ''' Transpose vol with axes swapping `order`. 

183 ''' 

184 if is2D(vol): 

185 cripAssert(len(order) == 2, 'order should have length 2 for 2D input.') 

186 if is3D(vol): 

187 cripAssert(len(order) == 3, 'order should have length 3 for 3D input.') 

188 

189 return vol.transpose(order) 

190 

191 

192@ConvertListNDArray 

193def permute(vol: ThreeD, from_: str, to: str) -> ThreeD: 

194 ''' Permute axes (transpose) from `from_` to `to`, reslicing the reconstructed volume. 

195 

196 Valid directions are: 'sagittal', 'coronal', 'transverse'. 

197 ''' 

198 dirs = ['sagittal', 'coronal', 'transverse'] 

199 

200 cripAssert(from_ in dirs, f'Invalid direction string: {from_}') 

201 cripAssert(to in dirs, f'Invalid direction string: {to}') 

202 

203 if from_ == to: 

204 return vol 

205 

206 dirFrom = dirs.index(from_) 

207 dirTo = dirs.index(to) 

208 

209 # TODO check this matrix 

210 orders = [ 

211 # to sag cor tra # from 

212 [(0, 1, 2), (1, 2, 0), (2, 1, 0)], # sag 

213 [(1, 2, 0), (0, 1, 2), (1, 0, 2)], # cor 

214 [(2, 1, 0), (1, 0, 2), (0, 1, 2)], # tra 

215 ] 

216 order = orders[dirFrom][dirTo] 

217 

218 return transpose(vol, order) 

219 

220 

221def shepplogan(size: int = 512) -> TwoD: 

222 ''' Generate a `size x size` Shepp-Logan phantom whose pixel values are approximately 

223 LAC (mu) [mm-1]. 

224 ''' 

225 cripAssert(size in [256, 512, 1024], 'Shepp-Logan size should be in [256, 512, 1024]') 

226 

227 return imreadTiff(path.join(getAsset('shepplogan'), f'{size}.tif')) / 10 

228 

229 

230def fitPolyV2D2(x1: NDArray, x2: NDArray, y: NDArray, bias: bool = True) -> NDArray: 

231 ''' Fit a degree-2 polynomial using pseudo-inverse to a pair of variables `x1, x2`. 

232 Output 6 coefficients `c[0~5]`, minimizing the error of `y` and 

233 `c[0] * x1**2 + c[1] * x2**2 + c[2] * x1 * x2 + c[3] * x1 + c[4] * x2 + c[5]`. 

234 If `bias` is False, `c[5]` will be 0. 

235 ''' 

236 cripAssert(is1D(x1) and is1D(x2) and is1D(y), 'Inputs should be 1D sequence.') 

237 cripAssert(isOfSameShape(x1, x2) and isOfSameShape(x1, y), '`x1`, `x2` and `y` should have same shape.') 

238 

239 x1sq = x1.T**2 

240 x2sq = x2.T**2 

241 x1x2 = (x1 * x2).T 

242 x1 = x1.T 

243 x2 = x2.T 

244 const = (np.ones if bias else np.zeros)((x1.T.shape[0])) 

245 A = np.array([x1sq, x2sq, x1x2, x1, x2, const]).T 

246 

247 return np.linalg.pinv(A) @ y 

248 

249 

250def applyPolyV2D2(coeff: NDArray, x1: NDArray, x2: NDArray) -> NDArray: 

251 ''' Apply a degree-2 polynomial to a pair of variables `x1, x2`. 

252 `coeff` has 5 or 6 elements, expands to  

253 `c[0] * x1**2 + c[1] * x2**2 + c[2] * x1 * x2 + c[3] * x1 + c[4] * x2 + (c[5] or 0)`. 

254 ''' 

255 cripAssert(len(coeff) in [5, 6], '`coeff` should have length of 5 or 6.') 

256 

257 if len(coeff) == 5: 

258 bias = 0 

259 else: 

260 bias = coeff[5] 

261 

262 return coeff[0] * x1**2 + coeff[1] * x2**2 + coeff[2] * x1 * x2 + coeff[3] * x1 + coeff[4] * x2 + bias 

263 

264 

265def fitPolyV1D2(x1: NDArray, y: NDArray, bias: bool = True) -> NDArray: 

266 ''' Fit a degree-2 polynomial using pseudo-inverse to a variable `x1`. 

267 Output 3 coefficients `c[0~2]`, minimizing the error of `y` and 

268 `c[0] * x1**2 + c[1] * x1 + c[2]`. 

269 If `bias` is False, `c[2]` will be 0. 

270 ''' 

271 x1sq = x1.T**2 

272 x1 = x1.T 

273 const = (np.ones if bias else np.zeros)((x1.T.shape[0])) 

274 A = np.array([x1sq, x1, const]).T 

275 

276 return np.linalg.pinv(A) @ y 

277 

278 

279def applyPolyV1D2(coeff: NDArray, x1: NDArray) -> NDArray: 

280 ''' Apply a degree-2 polynomial to a variable `x1`. 

281 `coeff` has 2 or 3 elements, expands to  

282 `c[0] * x1**2 + c[1] * x1 + (c[2] or 0)`. 

283 ''' 

284 cripAssert(len(coeff) in [2, 3], 'coeff should have length 2 or 3.') 

285 

286 if len(coeff) == 2: 

287 bias = 0 

288 else: 

289 bias = coeff[2] 

290 

291 return coeff[0] * x1**2 + coeff[1] * x1 + bias