--- a +++ b/vessel_segment.py @@ -0,0 +1,110 @@ +""" +This file is for plumonary vessels enhancement. +""" +import numpy as np +import SimpleITK as sitk +from scipy import ndimage +import cv2 + +from itertools import izip_longest + +class VesselSegment: + """ + This class is desigened for lung vasculature enhancement, including the methods of: + 1. downsampling + 2. thresholding + 3. 3D region growing + 4. filtering small struture + """ + + def __init__(self, original, closing): + """ + :param original: the orignal image in ITK formate + :param closing: the Closing result image + :param thval: threshold value (default: 130HU) + :param filter_vol: filter value for removing small struture after region growing (default: 2ml) + """ + self.original_img = original + self.closing_img = sitk.GetArrayFromImage(closing) + self.img = None + self.thval = None + self.filter_vol = None + self.temp_img = None + + def erosion(self, lunglabel=[201, 202]): + """ + Shrik the region of lung... + """ + temp = np.zeros_like(self.closing_img) + temp[self.closing_img == lunglabel[0]] = 1 + temp[self.closing_img == lunglabel[1]] = 1 + Erode_filter = sitk.BinaryErodeImageFilter() + Erode_filter.SetKernelRadius ( 2 ).SetForegroundValue ( 1 ) + self.closing_img = sitk.GetArrayFromImage( Erode_filter.Execute ( sitk.GetImageFromArray(temp) ) ) + + def generate_lung_mask(self, offset = 0): + """ + Generate lung mask + :return: None + """ + self.img = sitk.GetArrayFromImage(self.original_img).copy() - offset + self.img[self.closing_img == 0] = 0 + + def downsampling(self): + """ + Downsample the input image from [-1024, 0] to [0, 255] for reducing memory requirement. + / max(0, min(254, (Vorig+1024)/4), v belongs to Lung region + Vds = + \ 255, otherwise + :return: None + """ + temp = (self.img + 1024) / 4 + temp[temp > 254] = 254 + temp[temp < 0] = 0 + self.temp_img = temp + + def thresholding(self, thval=130): + """ + Threshold the image with given thresholding value. + :return: None + """ + self.thval = thval + self.temp_img[self.temp_img < thval] = thval + + def max_filter(self, filter_size=5): + """ + Implement maximum filter. + :param filter_size: filter size + :return: None + """ + temp = self.temp_img.copy() + temp[temp >= 254] = 0 + temp[temp <= self.thval] = 0 + self.temp_img = ndimage.filters.maximum_filter(temp, size=filter_size) + + def filtering(self, min_size=10, max_size=5000): + """ + Remove the small structure which volumn is less than given value (2ml) + :param min_size: + :return: None + """ + self.filter_vol = min_size + z, y, x = self.temp_img.shape + count_labels = [] + for i in range(z): + dist_transform = cv2.distanceTransform(np.uint8(self.temp_img[i, :, :]), cv2.DIST_L2, 5) + ret, sure_fg = cv2.threshold(dist_transform, 0.7 * dist_transform.max(), 255, 0) + sure_fg = np.uint8(sure_fg) + # Marker labelling + ret, markers = cv2.connectedComponents(sure_fg) + # Add one to all labels so that sure background is not 0, but 1 + markers += 1 + count_labels = np.asarray([x + y for x, y in izip_longest(count_labels, + np.bincount(markers.flatten()), + fillvalue=0)]) + labels = np.arange(0, len(count_labels)) + labels[count_labels < min_size] = 0 + labels[count_labels > max_size] = 0 + labels = np.asarray(list(set(labels))) + for label in labels: + self.temp_img[self.temp_img == label] = 0 \ No newline at end of file