Coverage for crip\mangoct.py: 27%

139 statements  

« prev     ^ index     » next       coverage.py v7.5.2, created at 2024-07-16 01:15 +0800

1''' 

2 MangoCT integration interface. 

3 

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

5''' 

6 

7import os 

8import io 

9import sys 

10import json 

11import tempfile 

12import subprocess 

13 

14from .utils import cripAssert, getAttrKeysOfObject, isType 

15from ._typing import * 

16 

17 

18class _MgCliConfig(object): 

19 ''' The base class for configuration of CLI version of mangoct. 

20 ''' 

21 

22 def __init__(self): 

23 pass 

24 

25 def dumpJSON(self): 

26 ''' Dump the configuration to JSON string. 

27 ''' 

28 dict_ = dict([(k, getattr(self, k)) for k in getAttrKeysOfObject(self)]) 

29 

30 return json.dumps(dict_, indent=2) 

31 

32 def dumpJSONFile(self, path: str): 

33 ''' Dump the configuration to JSON file. 

34 ''' 

35 with open(path, 'w') as fp: 

36 fp.write(self.dumpJSON()) 

37 

38 def fromJSON(self, json_: str): 

39 ''' Load the configuration from JSON string. 

40 ''' 

41 obj = json.loads(json_) 

42 for key in obj: 

43 self[key] = obj[key] 

44 

45 def fromJSONFile(self, path: str): 

46 ''' Load the configuration from JSON file. 

47 ''' 

48 with open(path, 'r') as fp: 

49 self.fromJSON(fp.read()) 

50 

51 

52class MgfbpCliConfig(_MgCliConfig): 

53 ''' Configuration class for CLI version `mgfbp`. 

54 ''' 

55 

56 def __init__(self): 

57 super().__init__() 

58 self.setIO(None, None, None, None, []) 

59 self.setGeometry(None, None, None, None, None, None, None, None) 

60 self.setSgmConeBeam(None, None, None, None, None, None, None) 

61 self.setRecConeBeam(None, None, None, None, 'HammingFilter', 1, None, None, None, None, None) 

62 

63 def setIO(self, 

64 InputDir: str, 

65 OutputDir: str, 

66 InputFiles: str, 

67 OutputFilePrefix: str = '', 

68 OutputFileReplace: List[str] = []): 

69 self.InputDir = InputDir 

70 self.OutputDir = OutputDir 

71 self.InputFiles = InputFiles 

72 self.OutputFilePrefix = OutputFilePrefix 

73 cripAssert(len(OutputFileReplace) % 2 == 0, '`OutputFileReplace` should be paired.') 

74 self.OutputFileReplace = OutputFileReplace 

75 

76 def setGeometry(self, 

77 SourceIsocenterDistance: Or[int, float], 

78 SourceDetectorDistance: Or[int, float], 

79 TotalScanAngle: Or[int, float], 

80 DetectorOffcenter: Or[int, float] = 0, 

81 PMatrixFile: Or[str, None] = None, 

82 SIDFile: Or[str, None] = None, 

83 SDDFile: Or[str, None] = None, 

84 ScanAngleFile: Or[str, None] = None, 

85 DetectorOffCenterFile: Or[str, None] = None): 

86 self.SourceIsocenterDistance = SourceIsocenterDistance 

87 self.SourceDetectorDistance = SourceDetectorDistance 

88 self.TotalScanAngle = TotalScanAngle 

89 self.DetectorOffcenter = DetectorOffcenter 

90 self.PMatrixFile = PMatrixFile 

91 self.SIDFile = SIDFile 

92 self.SDDFile = SDDFile 

93 self.ScanAngleFile = ScanAngleFile 

94 self.DetectorOffCenterFile = DetectorOffCenterFile 

95 

96 def setSgmFanBeam(self, 

97 SinogramWidth: int, 

98 SinogramHeight: int, 

99 Views: int, 

100 DetectorElementSize: Or[int, float], 

101 SliceCount: int = 1): 

102 self.ConeBeam = False 

103 self.SinogramWidth = SinogramWidth 

104 self.SinogramHeight = SinogramHeight 

105 self.Views = Views 

106 self.DetectorElementSize = DetectorElementSize 

107 self.SliceCount = SliceCount 

108 

109 def setSgmConeBeam(self, 

110 SinogramWidth: int, 

111 SinogramHeight: int, 

112 Views: int, 

113 DetectorElementSize: Or[int, float], 

114 SliceCount: int, 

115 SliceThickness: Or[int, float], 

116 SliceOffCenter: Or[int, float] = 0): 

117 self.setSgmFanBeam(SinogramWidth, SinogramHeight, Views, DetectorElementSize, SliceCount) 

118 self.ConeBeam = True 

119 self.SliceThickness = SliceThickness 

120 self.SliceOffCenter = SliceOffCenter 

121 

122 def setRecFanBeam(self, 

123 ImageDimension: int, 

124 PixelSize: Or[int, float], 

125 _Filter: str, 

126 _FilterParam: Or[float, List[float]], 

127 ImageRotation: Or[int, float] = 0, 

128 ImageCenter: List[float] = [0, 0], 

129 WaterMu: Or[float, None] = None, 

130 SaveFilteredSinogram: bool = False): 

131 self.ImageDimension = ImageDimension 

132 self.PixelSize = PixelSize 

133 cripAssert(_Filter in ['HammingFilter', 'QuadraticFilter', 'Polynomial', 'GaussianApodizedRamp'], 

134 f'Invalid _Filter: {_Filter}') 

135 if _Filter != 'HammingFilter': 

136 del self.HammingFilter 

137 exec(f'self.{_Filter} = _FilterParam') 

138 self.ImageRotation = ImageRotation 

139 self.ImageCenter = ImageCenter 

140 self.WaterMu = WaterMu 

141 self.SaveFilteredSinogram = SaveFilteredSinogram 

142 

143 def setRecConeBeam(self, 

144 ImageDimension: int, 

145 PixelSize: Or[int, float], 

146 ImageSliceCount: int, 

147 ImageSliceThickness: Or[int, float], 

148 _Filter: str, 

149 _FilterParam: Or[float, List[float]], 

150 ImageRotation: Or[int, float] = 0, 

151 ImageCenter: List[float] = [0, 0], 

152 ImageCenterZ: Or[int, float] = 0, 

153 WaterMu: Or[float, None] = None, 

154 SaveFilteredSinogram: bool = False): 

155 self.setRecFanBeam(ImageDimension, PixelSize, _Filter, _FilterParam, ImageRotation, ImageCenter, WaterMu, 

156 SaveFilteredSinogram) 

157 self.ImageSliceCount = ImageSliceCount 

158 self.ImageSliceThickness = ImageSliceThickness 

159 self.ImageCenterZ = ImageCenterZ 

160 

161 

162class MgfpjCliConfig(_MgCliConfig): 

163 ''' Configuration class for CLI version `mgfpj`. 

164 ''' 

165 

166 def __init__(self): 

167 super().__init__() 

168 self.setIO(None, None, None, None, []) 

169 self.setGeometry(None, None, None, None) 

170 self.setRecConeBeam(None, None, None, None) 

171 self.setSgmConeBeam(None, None, None, None, None, None, None, None) 

172 

173 def setIO(self, 

174 InputDir: str, 

175 OutputDir: str, 

176 InputFiles: str, 

177 OutputFilePrefix: str = '', 

178 OutputFileReplace: List[str] = []): 

179 self.InputDir = InputDir 

180 self.OutputDir = OutputDir 

181 self.InputFiles = InputFiles 

182 self.OutputFilePrefix = OutputFilePrefix 

183 cripAssert(len(OutputFileReplace) % 2 == 0, '`OutputFileReplace` should be paired.') 

184 self.OutputFileReplace = OutputFileReplace 

185 

186 def setGeometry(self, SourceIsocenterDistance: Or[int, float], SourceDetectorDistance: Or[int, float], 

187 StartAngle: Or[int, float], TotalScanAngle: Or[int, float]): 

188 self.SourceIsocenterDistance = SourceIsocenterDistance 

189 self.SourceDetectorDistance = SourceDetectorDistance 

190 self.StartAngle = StartAngle 

191 self.TotalScanAngle = TotalScanAngle 

192 

193 def setRecFanBeam(self, ImageDimension: int, PixelSize: Or[int, float], SliceCount: int = 1): 

194 self.ConeBeam = False 

195 self.ImageDimension = ImageDimension 

196 self.PixelSize = PixelSize 

197 self.SliceCount = SliceCount 

198 

199 def setRecConeBeam(self, ImageDimension: int, PixelSize: Or[int, float], SliceCount: int, 

200 ImageSliceThickness: Or[int, float]): 

201 self.setRecFanBeam(ImageDimension, PixelSize, SliceCount) 

202 self.ConeBeam = True 

203 self.ImageSliceThickness = ImageSliceThickness 

204 

205 def setSgmFanBeam(self, 

206 Views: int, 

207 DetectorElementCount: int, 

208 DetectorElementSize: Or[int, float], 

209 DetectorOffcenter: Or[int, float] = 0, 

210 OversampleSize: int = 2): 

211 self.Views = Views 

212 self.DetectorElementCount = DetectorElementCount 

213 self.DetectorElementSize = DetectorElementSize 

214 self.DetectorOffcenter = DetectorOffcenter 

215 self.OversampleSize = OversampleSize 

216 

217 def setSgmConeBeam(self, 

218 Views: int, 

219 DetectorElementCount: int, 

220 DetectorElementSize: Or[int, float], 

221 DetectorZElementCount: int, 

222 DetectorElementHeight: Or[int, float], 

223 DetectorOffcenter: Or[int, float] = 0, 

224 DetectorZOffcenter: Or[int, float] = 0, 

225 OversampleSize: int = 2): 

226 self.setSgmFanBeam(Views, DetectorElementCount, DetectorElementSize, DetectorOffcenter, OversampleSize) 

227 self.DetectorZElementCount = DetectorZElementCount 

228 self.DetectorElementHeight = DetectorElementHeight 

229 self.DetectorZOffcenter = DetectorZOffcenter 

230 

231 

232class _MgCliBin(object): 

233 ''' The base class for execution of CLI version of mangoct. 

234 ''' 

235 

236 def __init__(self, exe: str, name: str, cudaDevice: int = 0, tempDir: str = None): 

237 self.exe = exe 

238 self.name = name 

239 self.cudaDevice = cudaDevice 

240 self.tempDir = tempDir 

241 self.cmd = [f'{self.exe}', '<1>'] 

242 

243 def exec(self, conf: Or[str, _MgCliConfig], verbose=True): 

244 if isType(conf, _MgCliConfig): 

245 tmp = tempfile.NamedTemporaryFile('w', 

246 prefix='crip_mangoct_', 

247 suffix=f'.{self.name}.jsonc', 

248 dir=self.tempDir, 

249 delete=False) 

250 tmp.write(conf.dumpJSON()) 

251 conf = tmp.name 

252 tmp.close() 

253 

254 os.environ['CUDA_VISIBLE_DEVICES'] = str(self.cudaDevice) 

255 self.cmd[1] = self.cmd[1].replace('<1>', conf) 

256 proc = subprocess.Popen(self.cmd, stdout=subprocess.PIPE) 

257 for line in io.TextIOWrapper(proc.stdout, encoding="utf-8"): 

258 print(line, flush=True, file=sys.stdout) 

259 

260 

261class MgCliFbp(_MgCliBin): 

262 

263 def __init__(self, exe='mgfbp', cudaDevice=0, tempDir: Or[str, None] = None): 

264 ''' Initialize a handler object to use the CLI version FBP tool in mangoct. 

265 `exe` is the path to the executable. 

266 ''' 

267 super().__init__(exe, 'mgfbp', cudaDevice, tempDir) 

268 

269 

270class MgCliFpj(_MgCliBin): 

271 

272 def __init__(self, exe='mgfpj', cudaDevice=0, tempDir: Or[str, None] = None): 

273 ''' Initialize a handler object to use the CLI version FPJ tool in mangoct. 

274 `exe` is the path to the executable. 

275 ''' 

276 super().__init__(exe, 'mgfpj', cudaDevice, tempDir) 

277 

278 

279# TODO Add supports for Taichi version mangoct.