|
a |
|
b/vessel_segment.py |
|
|
1 |
""" |
|
|
2 |
This file is for plumonary vessels enhancement. |
|
|
3 |
""" |
|
|
4 |
import numpy as np |
|
|
5 |
import SimpleITK as sitk |
|
|
6 |
from scipy import ndimage |
|
|
7 |
import cv2 |
|
|
8 |
|
|
|
9 |
from itertools import izip_longest |
|
|
10 |
|
|
|
11 |
class VesselSegment: |
|
|
12 |
""" |
|
|
13 |
This class is desigened for lung vasculature enhancement, including the methods of: |
|
|
14 |
1. downsampling |
|
|
15 |
2. thresholding |
|
|
16 |
3. 3D region growing |
|
|
17 |
4. filtering small struture |
|
|
18 |
""" |
|
|
19 |
|
|
|
20 |
def __init__(self, original, closing): |
|
|
21 |
""" |
|
|
22 |
:param original: the orignal image in ITK formate |
|
|
23 |
:param closing: the Closing result image |
|
|
24 |
:param thval: threshold value (default: 130HU) |
|
|
25 |
:param filter_vol: filter value for removing small struture after region growing (default: 2ml) |
|
|
26 |
""" |
|
|
27 |
self.original_img = original |
|
|
28 |
self.closing_img = sitk.GetArrayFromImage(closing) |
|
|
29 |
self.img = None |
|
|
30 |
self.thval = None |
|
|
31 |
self.filter_vol = None |
|
|
32 |
self.temp_img = None |
|
|
33 |
|
|
|
34 |
def erosion(self, lunglabel=[201, 202]): |
|
|
35 |
""" |
|
|
36 |
Shrik the region of lung... |
|
|
37 |
""" |
|
|
38 |
temp = np.zeros_like(self.closing_img) |
|
|
39 |
temp[self.closing_img == lunglabel[0]] = 1 |
|
|
40 |
temp[self.closing_img == lunglabel[1]] = 1 |
|
|
41 |
Erode_filter = sitk.BinaryErodeImageFilter() |
|
|
42 |
Erode_filter.SetKernelRadius ( 2 ).SetForegroundValue ( 1 ) |
|
|
43 |
self.closing_img = sitk.GetArrayFromImage( Erode_filter.Execute ( sitk.GetImageFromArray(temp) ) ) |
|
|
44 |
|
|
|
45 |
def generate_lung_mask(self, offset = 0): |
|
|
46 |
""" |
|
|
47 |
Generate lung mask |
|
|
48 |
:return: None |
|
|
49 |
""" |
|
|
50 |
self.img = sitk.GetArrayFromImage(self.original_img).copy() - offset |
|
|
51 |
self.img[self.closing_img == 0] = 0 |
|
|
52 |
|
|
|
53 |
def downsampling(self): |
|
|
54 |
""" |
|
|
55 |
Downsample the input image from [-1024, 0] to [0, 255] for reducing memory requirement. |
|
|
56 |
/ max(0, min(254, (Vorig+1024)/4), v belongs to Lung region |
|
|
57 |
Vds = |
|
|
58 |
\ 255, otherwise |
|
|
59 |
:return: None |
|
|
60 |
""" |
|
|
61 |
temp = (self.img + 1024) / 4 |
|
|
62 |
temp[temp > 254] = 254 |
|
|
63 |
temp[temp < 0] = 0 |
|
|
64 |
self.temp_img = temp |
|
|
65 |
|
|
|
66 |
def thresholding(self, thval=130): |
|
|
67 |
""" |
|
|
68 |
Threshold the image with given thresholding value. |
|
|
69 |
:return: None |
|
|
70 |
""" |
|
|
71 |
self.thval = thval |
|
|
72 |
self.temp_img[self.temp_img < thval] = thval |
|
|
73 |
|
|
|
74 |
def max_filter(self, filter_size=5): |
|
|
75 |
""" |
|
|
76 |
Implement maximum filter. |
|
|
77 |
:param filter_size: filter size |
|
|
78 |
:return: None |
|
|
79 |
""" |
|
|
80 |
temp = self.temp_img.copy() |
|
|
81 |
temp[temp >= 254] = 0 |
|
|
82 |
temp[temp <= self.thval] = 0 |
|
|
83 |
self.temp_img = ndimage.filters.maximum_filter(temp, size=filter_size) |
|
|
84 |
|
|
|
85 |
def filtering(self, min_size=10, max_size=5000): |
|
|
86 |
""" |
|
|
87 |
Remove the small structure which volumn is less than given value (2ml) |
|
|
88 |
:param min_size: |
|
|
89 |
:return: None |
|
|
90 |
""" |
|
|
91 |
self.filter_vol = min_size |
|
|
92 |
z, y, x = self.temp_img.shape |
|
|
93 |
count_labels = [] |
|
|
94 |
for i in range(z): |
|
|
95 |
dist_transform = cv2.distanceTransform(np.uint8(self.temp_img[i, :, :]), cv2.DIST_L2, 5) |
|
|
96 |
ret, sure_fg = cv2.threshold(dist_transform, 0.7 * dist_transform.max(), 255, 0) |
|
|
97 |
sure_fg = np.uint8(sure_fg) |
|
|
98 |
# Marker labelling |
|
|
99 |
ret, markers = cv2.connectedComponents(sure_fg) |
|
|
100 |
# Add one to all labels so that sure background is not 0, but 1 |
|
|
101 |
markers += 1 |
|
|
102 |
count_labels = np.asarray([x + y for x, y in izip_longest(count_labels, |
|
|
103 |
np.bincount(markers.flatten()), |
|
|
104 |
fillvalue=0)]) |
|
|
105 |
labels = np.arange(0, len(count_labels)) |
|
|
106 |
labels[count_labels < min_size] = 0 |
|
|
107 |
labels[count_labels > max_size] = 0 |
|
|
108 |
labels = np.asarray(list(set(labels))) |
|
|
109 |
for label in labels: |
|
|
110 |
self.temp_img[self.temp_img == label] = 0 |