/
astraOP.py
89 lines (81 loc) · 3.43 KB
/
astraOP.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
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
A class based on using ASTRA toolbox to perform projection/bakprojection and
also some reconstruction algorithms
-- currently 2D parallel beam
"""
import astra
class AstraTools:
"""A simple 2D parallel beam projection/backprojection class based on ASTRA toolbox"""
def __init__(self, DetectorsDim, AnglesVec, ObjSize, device):
self.DetectorsDim = DetectorsDim
self.AnglesVec = AnglesVec
self.ObjSize = ObjSize
self.proj_geom = astra.create_proj_geom('parallel', 1.0, DetectorsDim, AnglesVec)
self.vol_geom = astra.create_vol_geom(ObjSize, ObjSize)
if device == 'cpu':
self.proj_id = astra.create_projector('line', self.proj_geom, self.vol_geom) # for CPU
self.device = 1
elif device == 'gpu':
self.proj_id = astra.create_projector('cuda', self.proj_geom, self.vol_geom) # for GPU
self.device = 0
else:
print ("Select between 'cpu' or 'gpu' for device")
def forwproj(self, image):
"""Applying forward projection"""
sinogram_id, sinogram = astra.create_sino(image, self.proj_id)
astra.data2d.delete(sinogram_id)
astra.data2d.delete(self.proj_id)
return sinogram
def backproj(self, sinogram):
"""Applying backprojection"""
rec_id, image = astra.create_backprojection(sinogram, self.proj_id)
astra.data2d.delete(self.proj_id)
astra.data2d.delete(rec_id)
return image
def fbp2D(self, sinogram):
"""perform FBP reconstruction"""
rec_id = astra.data2d.create( '-vol', self.vol_geom)
# Create a data object to hold the sinogram data
sinogram_id = astra.data2d.create('-sino', self.proj_geom, sinogram)
if self.device == 1:
cfg = astra.astra_dict('FBP')
cfg['ProjectorId'] = self.proj_id
else:
cfg = astra.astra_dict('FBP_CUDA')
cfg['ReconstructionDataId'] = rec_id
cfg['ProjectionDataId'] = sinogram_id
cfg['FilterType'] = 'Ram-Lak'
# Create and run the algorithm object from the configuration structure
alg_id = astra.algorithm.create(cfg)
astra.algorithm.run(alg_id)
# Get the result
recFBP = astra.data2d.get(rec_id)
astra.algorithm.delete(alg_id)
astra.data2d.delete(rec_id)
astra.data2d.delete(sinogram_id)
astra.data2d.delete(self.proj_id)
return recFBP
def sirt2D(self, sinogram, iterations):
"""perform SIRT reconstruction"""
rec_id = astra.data2d.create( '-vol', self.vol_geom)
# Create a data object to hold the sinogram data
sinogram_id = astra.data2d.create('-sino', self.proj_geom, sinogram)
if self.device == 1:
cfg = astra.astra_dict('SIRT')
cfg['ProjectorId'] = self.proj_id
else:
cfg = astra.astra_dict('SIRT_CUDA')
cfg['ReconstructionDataId'] = rec_id
cfg['ProjectionDataId'] = sinogram_id
# Create and run the algorithm object from the configuration structure
alg_id = astra.algorithm.create(cfg)
astra.algorithm.run(alg_id, iterations)
# Get the result
recSIRT = astra.data2d.get(rec_id)
astra.algorithm.delete(alg_id)
astra.data2d.delete(rec_id)
astra.data2d.delete(sinogram_id)
astra.data2d.delete(self.proj_id)
return recSIRT