Dubhe/dubhe_data_process/program/exec/lung_segmentation/lung_segmentation.py

# !/usr/bin/env python
# -*- coding:utf-8 -*-

"""
Copyright 2020 Tianshu AI Platform. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

     http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
=============================================================
"""
import sched
import sys

import time
import json
import numpy as np
from abc import ABC
from program.abstract.algorithm import Algorithm
from skimage.morphology import disk, binary_erosion, binary_closing
from skimage.measure import label, regionprops, find_contours
from skimage.filters import roberts
from scipy import ndimage as ndi
from skimage.segmentation import clear_border
import pydicom as dicom
import os
import logging

schedule = sched.scheduler(time.time, time.sleep)

base_path = "/nfs/"
delayId = ""


class Lungsegmentation(Algorithm, ABC):

    def __init__(self):
        pass

    def execute(task):
        return Lungsegmentation.process(task)

    def process(task_dict):
        """Lung segmentation based on dcm task method.
            Args:
                task_dict: imagenet task details.
                key: imagenet task key.
        """
        global delayId
        base_path = task_dict["annotationPath"]
        if not os.path.exists(base_path):
            logging.info("make annotation path.")
            os.makedirs(base_path)

        for dcm in task_dict["dcms"]:
            image, image_path = Lungsegmentation.preprocesss_dcm_image(dcm)
            # segmentation and wirte coutours to result_path
            result_path = os.path.join(base_path, image_path)
            Lungsegmentation.contour(Lungsegmentation.segmentation(image), result_path)

        logging.info("all dcms in one task are processed.")
        finish_data = {"reTaskId": task_dict["reTaskId"]}
        return finish_data, True

    def preprocesss_dcm_image(path):
        """Load and preprocesss dcm image.
            Args:
                path: dcm file path.
        """
        # result_path = os.path.basename(path).split(".", 1)[0] + ".json"
        result_path = ".".join(os.path.basename(path).split(".")[0:-1]) + ".json"
        dcm = dicom.dcmread(path)
        image = dcm.pixel_array.astype(np.int16)

        # Set outside-of-scan pixels to 0.
        image[image == -2000] = 0

        # Convert to Hounsfield units (HU)
        intercept = dcm.RescaleIntercept
        slope = dcm.RescaleSlope

        if slope != 1:
            image = slope * image.astype(np.float64)
            image = image.astype(np.int16)

        image += np.int16(intercept)
        logging.info("preprocesss_dcm_image done.")
        return np.array(image, dtype=np.int16), result_path

    def segmentation(image):
        """Segments the lung from the given 2D slice.
            Args:
                image: single image in one dcm.
        """
        # Step 1: Convert into a binary image.
        binary = image < -350

        # Step 2: Remove the blobs connected to the border of the image.
        cleared = clear_border(binary)

        # Step 3: Label the image.
        label_image = label(cleared)

        # Step 4: Keep the labels with 2 largest areas.
        areas = [r.area for r in regionprops(label_image)]
        areas.sort()
        if len(areas) > 2:
            for region in regionprops(label_image):
                if region.area < areas[-2]:
                    for coordinates in region.coords:
                        label_image[coordinates[0], coordinates[1]] = 0
        binary = label_image > 0

        # Step 5: Erosion operation with a disk of radius 2. This operation is seperate the lung nodules attached to the blood vessels.
        selem = disk(1)
        binary = binary_erosion(binary, selem)

        # Step 6: Closure operation with a disk of radius 10. This operation is to keep nodules attached to the lung wall.
        selem = disk(16)
        binary = binary_closing(binary, selem)

        # Step 7: Fill in the small holes inside the binary mask of lungs.
        for _ in range(3):
            edges = roberts(binary)
            binary = ndi.binary_fill_holes(edges)
        logging.info("lung segmentation done.")
        return binary

    def contour(image, path):
        """Get contours of segmentation.
            Args:
                seg: segmentation of lung.
        """
        result = []
        contours = find_contours(image, 0.5)
        if len(contours) > 2:
            contours.sort(key=lambda x: int(x.shape[0]))
            contours = contours[-2:]

        for n, contour in enumerate(contours):
            # result.append({"type":n, "annotation":contour.tolist()})
            result.append({"type": n, "annotation": np.flip(contour, 1).tolist()})

        # write json
        with open(path, 'w') as f:
            json.dump(result, f)
        logging.info("write {} done.".format(path))