brain_tumor_segmentation_models.py 文件源码

def save_segmented_image(self, index, test_img, save=False):
        """
        Creates an image of original brain with segmentation overlay
        :param index: index of image to save
        :param test_img: filepath to test image for segmentation, including file extension
        :param save: If true, shows output image. (defaults to False)
        :return: if show is True, shows image of segmentation results
                 if show is false, returns segmented image.
        """

        segmentation = self.predict_image(test_img)

        img_mask = np.pad(segmentation, (16, 16), mode='edge')
        ones = np.argwhere(img_mask == 1)
        twos = np.argwhere(img_mask == 2)
        threes = np.argwhere(img_mask == 3)
        fours = np.argwhere(img_mask == 4)

        test_im = mpimg.imread(test_img).astype('float')
        test_back = rgb2gray(test_im).reshape(5, 216, 160)[-2]
        # overlay = mark_boundaries(test_back, img_mask)
        gray_img = img_as_float(test_back)

        # adjust gamma of image
        image = adjust_gamma(color.gray2rgb(gray_img), 0.65)
        sliced_image = image.copy()
        red_multiplier = [1, 0.2, 0.2]
        yellow_multiplier = [1, 1, 0.25]
        green_multiplier = [0.35, 0.75, 0.25]
        blue_multiplier = [0, 0.25, 0.9]

        # change colors of segmented classes
        for i in xrange(len(ones)):
            sliced_image[ones[i][0]][ones[i][1]] = red_multiplier
        for i in xrange(len(twos)):
            sliced_image[twos[i][0]][twos[i][1]] = green_multiplier
        for i in xrange(len(threes)):
            sliced_image[threes[i][0]][threes[i][1]] = blue_multiplier
        for i in xrange(len(fours)):
            sliced_image[fours[i][0]][fours[i][1]] = yellow_multiplier

        if save:

            try:
                mkdir_p('./results/')
                io.imsave('./results/result' + '_' + str(index) + '.png', sliced_image)
            except:
                io.imsave('./results/result' + '_' + str(index) + '.png', sliced_image)
        else:
            return sliced_image