split_tibia_femur_spine_tiankuo.py

#! /home/spl/ml/sitk/bin/python

# -*- coding: utf-8 -*-

import os
import SimpleITK as sitk
from shubow_tools import *
import re
import time
import concurrent.futures
import logging
import cv2
import shutil
import glob

import matplotlib.pyplot as plt
import numpy as np


config = {
    "week" :" week 0" ,         # timepoint to scan
    "masterfolder" : r"D:\1b.  Aged Rad+Yoda1+WBV+ No- Tumor 062722\recon\week 0",
    "masterout" : r"D:\1b.  Aged Rad+Yoda1+WBV+ No- Tumor 062722\segmentation\week 0"
     
}


def knee_joint_z_index(limb, _min=1000, _max=2000, threshold=50):  #100
    '''
    Description:
        Given a 3D image of one limb, find the z_index of the knee where the largest contour 
        in this slice is the smallest among all 2D slices.
    Args:
        limb: ndarray, dimension = 3
    return: 
        z_index: integer
    '''

    min_contour_size = float('inf')
    smallest_largest_contour = None  # Initialize to None or a suitable default
    thresh_slice_image=None
    knee_slice_index = -1

    for i, slice in enumerate(limb):
        
        #apply Median Blur
        slice = cv2.medianBlur(slice, 5) 
        # Apply threshold
        ret, thresh_slice = cv2.threshold(slice, threshold, 255, cv2.THRESH_BINARY)
        # Define a kernel for morphological operations
        kernel = np.ones((15, 15), np.uint8)
        # Apply closing (dilation followed by erosion)
        thresh_slice = cv2.morphologyEx(thresh_slice, cv2.MORPH_CLOSE, kernel)
        # Find contoursc
        contours, _ = cv2.findContours(thresh_slice.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        if contours:           
            if _min<i<_max : #i == 1361: 
                # Find the contour with the largest perimeter
                largest_contour = max(contours, key=lambda contour: cv2.arcLength(contour, True))

                # Check if this is the smallest largest contour so far
                if cv2.arcLength(largest_contour, True) < min_contour_size:
                    # Update the smallest largest contour perimeter found
                    min_contour_size = cv2.arcLength(largest_contour, True)
                    thresh_slice_image=thresh_slice
                    smallest_largest_contour = largest_contour
                    knee_slice_index = i
                    
    # # Convert contour_img to a 3-channel RGB image if it's a grayscale image
    # if len(thresh_slice_image.shape) == 2:
    #     thresh_slice_image = cv2.cvtColor(thresh_slice_image, cv2.COLOR_GRAY2RGB)
    # # Draw the largest contour
    # cv2.drawContours(thresh_slice_image, [smallest_largest_contour], -1, (255, 0, 0), 3)  # Red contour
    # # Plotting
    # print(knee_slice_index,cv2.contourArea(smallest_largest_contour) )
    # plt.figure(figsize=(6, 6))
    # plt.imshow(thresh_slice_image, cmap='gray')
    # plt.title("smallest coutour among slices")
    # plt.show()
                  
    return knee_slice_index


def LR_mid_x(image):
    '''
    Description:
        Given a 3D reconstruction, find the y_index to split top and bottom,
        and then find the x_index to split left and right for the top section.
    Args:
        image: ndarray, dimension = 3 
    Returns: 
        y_index: integer, x_index: integer
    '''       

    # Project along x-axis to find y_index (top and bottom)
    x_project = image.mean(axis=2)          #image.size  (z,y) 
    yz_dex = np.vstack(np.nonzero(x_project))
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0)
    k = 2
    attempts = 5
    ret, label, center = cv2.kmeans(np.float32(yz_dex.transpose()), k, None, criteria, attempts, cv2.KMEANS_PP_CENTERS)
    y_index = int(center[:, 1].sum() / 2)
    # print("y_index",y_index)
    #y_index = int(center[:, 0].sum() / 2)
                                   
    # Isolate the top section based on y_index
    top_section = image[:, :y_index, :]

    # Project along z-axis to find x_index (left and right) for the top section
    z_project_top = top_section.mean(axis=0) 
    xy_dex_top = np.vstack(np.nonzero(z_project_top))
    ret, label, center_top = cv2.kmeans(np.float32(xy_dex_top.transpose()), k, None, criteria, attempts, cv2.KMEANS_PP_CENTERS)
    mid_x_index = int(center_top[:, 1].sum() / 2)

    return y_index, mid_x_index


def splitLRTF(folder,imgtitle,outfd = None):
    logging.info("Reading {}".format(imgtitle))
    img = imreadseq_multithread(folder, sitkimg= False, rmbckgrd=60, thread=2)
    
    depth = img.shape[0]
    height = img.shape[1]   #(z,y,x)
    width = img.shape[2]
    
    y_index, mid_x_index = LR_mid_x(img)  #y_index, mid_x_index     mid_idx

    if not outfd is None:
        pass
    else:
        outfd = folder
    
    
    week = config["week"]
    
    
    pathlist = []
    pathlist.append(os.path.join(outfd,imgtitle + week + ' spine'))
    pathlist.append(os.path.join(outfd,imgtitle + week +' left tibia' ))
    pathlist.append(os.path.join(outfd,imgtitle + week +' left femur' ))
    pathlist.append(os.path.join(outfd,imgtitle + week +' right tibia'))
    pathlist.append(os.path.join(outfd,imgtitle + week +' right femur'))

    for fd in pathlist:
        if os.path.exists(fd):
            shutil.rmtree(fd)
            os.mkdir(fd)
        else:
            os.mkdir(fd)

    titlelist = [imgtitle + week +' spine',
                imgtitle + week +' left tibia', 
                imgtitle + week +' left femur',
                imgtitle + week +' right tibia',
                imgtitle + week +' right femur']
    
    # print(y_index)
    # print("total index", len(range(y_index, img.shape[1])))
    
    
    ####################################
    y_index=y_index #+170 #move downward to leave space for femoral head

    bottom = auto_crop(img[:,y_index:,:])
    top = auto_crop(img[:,:y_index,:]) #split the top into left and right 
    left = auto_crop(img[:,:y_index,:mid_x_index])
    right = auto_crop(img[:,:y_index,-1:mid_x_index:-1])
    del img
    
    ##### Save spine #####   
    logging.info("Processing...split Spine & limb")
    spine = sitk.GetImageFromArray(auto_crop(bottom[:]))
    del bottom
    # use multiple threads to accelerate writng images to disk.
    # create an iterable to be passed to imreadseq(img,fd,title)
    imagelist = [spine]
    logging.info("Writing...")
    with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
        executor.map(imsaveseq,imagelist,pathlist[:1],titlelist[:1])
    del spine
    
    ##### Save left tibia and femur #####
    logging.info("Processing...split LT & LF")
    z_index_splt_left=knee_joint_z_index(left)
    left_tibia = sitk.GetImageFromArray(auto_crop(left[:z_index_splt_left]))
    left_femur = sitk.GetImageFromArray(auto_crop(left[z_index_splt_left:]))
    del left
    # use multiple threads to accelerate writng images to disk.
    # create an iterable to be passed to imreadseq(img,fd,title)
    imagelist = [left_tibia,left_femur]
    logging.info("Writing...")
    with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
        executor.map(imsaveseq,imagelist,pathlist[1:3],titlelist[1:3])
    del left_tibia,left_femur
    
    ##### Save right tibia and femur #####
    logging.info("Processing...split RT & RF")
    z_index_splt_right=knee_joint_z_index(right)
    right_tibia = sitk.GetImageFromArray(auto_crop(right[:z_index_splt_right]))
    right_femur = sitk.GetImageFromArray(auto_crop(right[z_index_splt_right:]))
    del right
    imagelist = [right_tibia,right_femur]
    logging.info("Writing...")
    with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
        executor.map(imsaveseq,imagelist,pathlist[3:],titlelist[3:])

    del right_tibia,right_femur
    
if __name__ == "__main__":

    masterfolder = config["masterfolder"]
    masterout = config["masterout"]

    if not os.path.exists(masterout):
        os.mkdir(masterout)

    time1 = time.time()   
    count = 0
    '''
    with open(r"/media/spl/D/MicroCT_data/Shubo/Reconstruction image/retry.txt", "r") as file:
        retry = file.readlines()
    retry = [i[:-1] for i in retry]
    '''
    format = "%(asctime)s: %(message)s"
    logging.basicConfig(format=format, level=logging.INFO,
                        datefmt="%H:%M:%S")
    failed = []

    #for inputfd in glob.glob(os.path.join(masterfolder,"*reconstruction")):
    for folder in sorted(os.listdir(masterfolder))[:]:
        if folder [0:10]:
          # wk1 ["865","875"]:
          # wk2 ["851","862","875"]:    853
          # wk3 ["853","860","862","864"]:  856R,859R     866R,869R,875L     "856","859","866","869","875"
          # wk4 ["853","864","866","868","871","873"]:  
            count += 1
            ID = os.path.basename(folder)[0:12]
            logging.info('Cropping for {} started.'.format(ID)) 
            try:
                splitLRTF(os.path.join(masterfolder,folder),ID,masterout)
                logging.info('Cropping for {} is completed.'.format(ID))
            except Exception as err:
                print(err)
                failed.append(folder)
                logging.exception('Cropping for {} failed.'.format(ID))
                pass

    print(failed)
    time2 = time.time()
    logging.info("Total time used: {: >8.1f} seconds".format(time2-time1))
    logging.info("Average time used: {: >8.1f} seconds".format((time2-time1)/count))