split_tibia_femur_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

def knee_joint_z_index(limb, _min=1200, _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.GaussianBlur(slice, (5, 5), 0) 
        # 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 contours
        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):
    '''
    Desription:
        Given a 3D reconstruction, find the x_index to split left and right
    Args:
        limb: ndarray, dimension = 3 
    return: 
        center: ndarray, 
    '''       

    z_project = image.mean(axis = 0)
    xy_dex = np.vstack(np.nonzero(z_project))
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0) # kmeans parameters
    k=2
    attempts=15
    ret, label, center = cv2.kmeans(np.float32(xy_dex.transpose()),k,
                            None,criteria,attempts,cv2.KMEANS_PP_CENTERS)

    mid_x_index = int(center[:, 1].sum()/2)

    return mid_x_index


def splitLRTF(folder,imgtitle,outfd = None):
    logging.info("Reading {}".format(imgtitle))
    img = imreadseq_multithread(folder, sitkimg= False, rmbckgrd=60, thread=2)
    
    #logging.info("Processing...split LT & LF")
    width = img.shape[2]
    mid_idx = LR_mid_x(img)

    if not outfd is None:
        pass
    else:
        outfd = folder

    pathlist = []
    pathlist.append(os.path.join(outfd,imgtitle+' week 11 left tibia' ))
    pathlist.append(os.path.join(outfd,imgtitle+' week 11 left femur' ))
    pathlist.append(os.path.join(outfd,imgtitle+' week 11 right tibia'))
    pathlist.append(os.path.join(outfd,imgtitle+' week 11 right femur'))

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

    titlelist = [imgtitle+' week 11 left tibia', imgtitle+' week 11 left femur',
                imgtitle+' week 11 right tibia',imgtitle+' week 11 right femur']
    
    left = auto_crop(img[:,:,:mid_idx])
    right = auto_crop(img[:,:,-1:mid_idx:-1])
    del img

    ##### 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[:2],titlelist[:2])
    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[2:],titlelist[2:])

    del right_tibia,right_femur
    
if __name__ == "__main__":
    masterfolder = r'F:\PD mice\PD mice 120823\recon\week 11'
    masterout = r'F:\PD mice\PD mice 120823\seg\week 11 new'

    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] : #  #if folder[0:10] in ["555 week 2_Rec"]:      if folder[0:10]:      #####################edit
            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))