Diff of /seg_crop_pad_scale.py [000000] .. [3b7fea]

Switch to side-by-side view

--- a
+++ b/seg_crop_pad_scale.py
@@ -0,0 +1,445 @@
+#!/usr/bin/env python3
+"""
+Author : briancottle <briancottle@localhost>
+Date   : 2023-01-02
+Purpose: Cropping, padding, and scaling segmentations in preparation for 3D modeling
+"""
+
+import argparse
+from typing import NamedTuple
+import numpy as np
+import cv2 as cv
+import os
+import matplotlib.pyplot as plt
+import tqdm
+from natsort import natsorted
+from glob import glob
+import magic
+import re
+from PIL import Image
+
+class Args(NamedTuple):
+    """ Command-line arguments """
+    tissue_chainID: str
+    cropping: bool
+    old_files: bool
+    
+    fid_seg: bool
+    node_seg: bool
+    unet_seg: bool
+    high_res_seg: bool
+
+    node_white: bool
+    fid_white: bool
+
+    reduction_size: int
+
+
+# --------------------------------------------------
+def get_args() -> Args:
+    """ Get command-line arguments """
+
+    parser = argparse.ArgumentParser(
+        description='Cropping, padding, and scaling segmentations in preparation for 3D modeling',
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+
+    parser.add_argument('-t_id',
+                        '--tissue_chainID',
+                        type=str,
+                        metavar='chainID',
+                        help='ID used to identify where the tissue files are stored')
+
+    parser.add_argument('-c',
+                        '--cropping',
+                        help='Boolean: Do you need to crop the uNet dataset first?',
+                        action='store_true')
+
+    parser.add_argument('-o',
+                        '--old_files',
+                        help='Is the dataset of interest from the old set?',
+                        action='store_true')
+
+    parser.add_argument('-f',
+                        '--fid_seg',
+                        help='Do you want to process the fiduciary segmentations?',
+                        action='store_true')
+
+    parser.add_argument('-n',
+                        '--node_seg',
+                        help='Do you want to process the nodal segmentations?',
+                        action='store_true')
+
+    parser.add_argument('-u',
+                        '--unet_seg',
+                        help='Do you want to process the uNet segmentations?',
+                        action='store_true')
+
+    parser.add_argument('-d',
+                        '--high_res_seg',
+                        help='Do you want to process the high_res segmentations?',
+                        action='store_true')
+
+    parser.add_argument('-i',
+                        '--node_white',
+                        help='Is the segmentation for the node inverted?',
+                        action='store_true')
+
+    parser.add_argument('-j',
+                        '--fid_white',
+                        help='Is the segmentation for the fiduciary inverted?',
+                        action='store_true')
+
+    parser.add_argument('-r',
+                        '--reduction_size',
+                        help='what scalar to use for reducing the size(4 or 8), default is 4',
+                        metavar='int',
+                        type=int,
+                        default=4)
+
+
+
+    args = parser.parse_args()
+
+    return Args(args.tissue_chainID,
+                args.cropping,
+                args.old_files,
+                args.fid_seg,
+                args.node_seg,
+                args.unet_seg,
+                args.high_res_seg,
+                args.node_white,
+                args.fid_white,
+                args.reduction_size)
+
+
+# --------------------------------------------------
+def main() -> None:
+    """ Make a jazz noise here """
+
+    args = get_args()
+
+    TissueChainID = args.tissue_chainID
+    cropping = args.cropping
+    old_files = args.old_files
+    fiduciary_files = args.fid_seg
+    nodal_files = args.node_seg
+    seg_files = args.unet_seg
+    high_res_files = args.high_res_seg
+    nodal_white = args.node_white
+    fiduciary_white = args.fid_white
+    reduction_size = args.reduction_size
+
+    if reduction_size == 4:
+        reduction_name = 'QuarterScale'
+
+    if reduction_size == 8:
+        reduction_name = 'EighthScale'
+
+
+    base_directory = '/var/confocaldata/HumanNodal/HeartData/'+ TissueChainID
+    JPG_directory = '/var/confocaldata/HumanNodal/HeartData/'+ TissueChainID +'JPG/'
+    jpg_file_names = glob(JPG_directory + '*.jpg')
+
+
+    if old_files:
+        ML_directory = '/var/confocaldata/HumanNodal/HeartData/'+ TissueChainID +'uNet_Segmentations/'
+        Nodal_Seg_directory = '/var/confocaldata/HumanNodal/HeartData/'+ TissueChainID +'Segmentations/Nodal Segmentation FullScale_NoPad/'
+        if fiduciary_files:
+            Fiduciary_Seg_directory = '/var/confocaldata/HumanNodal/HeartData/'+ TissueChainID +'Segmentations/Fiduciary Segmentation FullScale_NoPad/'
+    else:
+        ML_directory = '/var/confocaldata/HumanNodal/HeartData/'+ TissueChainID +'uNet_Segmentations/'
+        Nodal_Seg_directory = '/var/confocaldata/HumanNodal/HeartData/'+ TissueChainID +'Segmentations/Nodal Segmentation/'
+        if fiduciary_files:
+            Fiduciary_Seg_directory = '/var/confocaldata/HumanNodal/HeartData/'+ TissueChainID +'Segmentations/Fiduciary Segmentation/'
+
+    high_res_directory = '/var/confocaldata/HumanNodal/HeartData/'+ TissueChainID +'HighResSeg/'
+
+    os.chdir(ML_directory)
+    jpg_file_names = natsorted(jpg_file_names)
+    padding_size = 4000 # + 1536
+
+    # %% USE THIS SECTION FOR CROPPING THE SEGMENTATIONS AFTER THE UNET HAS AT IT
+    if cropping:
+        for idx in tqdm.tqdm(range(len(jpg_file_names))):
+
+            out_directory = './../Cropped_uNet_Segmentations/'
+            
+            # create the directory for saving if it doesn't already exist
+            if not os.path.isdir(out_directory):
+                os.mkdir(out_directory)
+
+            os.chdir(out_directory)
+
+            jpg_file = jpg_file_names[idx]
+            id = jpg_file.split('/')[-1].split('.')[0]
+            id = id.split('_')[0] + '_' + id.split('_')[1] + '_' + id.split('_')[2]
+            ml_file = glob(ML_directory + f'{id}_*.png')[0]
+
+            jpg_image1 = cv.imread(jpg_file)
+            ml_image1 = cv.imread(ml_file)[:,:,0]
+            [x,y,z] = jpg_image1.shape
+
+            cropped_ml1 = ml_image1[padding_size:padding_size+x,
+                                    padding_size:padding_size+y]
+
+            cv.imwrite(
+                id + 
+                f'_CroppedSeg.png',
+                cropped_ml1
+                )
+
+    # %%
+    ML_directory = '/var/confocaldata/HumanNodal/HeartData/'+ TissueChainID +'Cropped_uNet_Segmentations/'
+    ml_file_names = glob(ML_directory + '*.png')
+    all_image_sizes = []
+    for file_name in ml_file_names:
+        header = magic.from_file(file_name)
+        size = re.search('(\d+) x (\d+)',header).groups()
+        sizes = [int(a) for a in size]
+        all_image_sizes.append(sizes)
+
+    max_width = np.max(np.asarray(all_image_sizes)[:,0])
+    max_height = np.max(np.asarray(all_image_sizes)[:,1])
+
+    idx = 200
+    additional_padding = 4000
+    os.chdir(JPG_directory)
+    out_big_directory = base_directory + 'Padded_Images/'
+    out_small_directory = base_directory + 'Padded_Images_' + reduction_name + '/'
+    out_parent_list = [out_big_directory,out_small_directory]
+    out_list = []
+    if not os.path.isdir(out_big_directory):
+        os.mkdir(out_big_directory)
+
+    if not os.path.isdir(out_small_directory):
+        os.mkdir(out_small_directory)
+
+    for idx, out_directory in enumerate(out_parent_list):
+        jpg_out = out_directory + 'JPG'
+        seg_out = out_directory + 'Seg'
+        nodal_out = out_directory + 'Nodal'
+        fiduciary_out = out_directory + 'Fiduciary'
+        high_res_out = out_directory + 'HighRes'
+        out_list.append([jpg_out,seg_out,nodal_out,high_res_out,fiduciary_out])
+
+        for directory in out_list[idx]:
+            if not os.path.isdir(directory):
+                os.mkdir(directory)
+
+
+
+    for idx in tqdm.tqdm(range(len(jpg_file_names))):
+
+        jpg_file = jpg_file_names[idx]
+        id = jpg_file.split('/')[-1].split('.')[0]
+        id = id.split('_')[0] + '_' + id.split('_')[1] + '_' + id.split('_')[2]
+        
+        
+        
+        # Create a separate section for the nodal tissue stuff, as it looks like 
+        # nodal segmentation will actually happen after the registration using the
+        # segmentations
+        # change this to _*.png if you are not using the FullScale_NoPad segmentations
+        # will need to scale the segmentations for newer segmentations that haven't been 
+        # performed using previously padded images. This section is below, starting with
+
+
+        jpg_image = cv.imread(jpg_file)
+        [height,width,z] = jpg_image.shape
+
+        height_diff = max_height - height
+        width_diff = max_width - width
+
+        if height_diff%2 == 1:
+            pad_top = np.floor(height_diff/2) + additional_padding
+            pad_bottom = np.floor(height_diff/2) + additional_padding
+            pad_bottom += 1
+        else:
+            pad_top = np.floor(height_diff/2) + additional_padding
+            pad_bottom = np.floor(height_diff/2) + additional_padding
+
+        if width_diff%2 == 1:
+            pad_left = np.floor(width_diff/2) + additional_padding
+            pad_right = np.floor(width_diff/2) + additional_padding
+            pad_right += 1
+        else:
+            pad_left = np.floor(width_diff/2) + additional_padding
+            pad_right = np.floor(width_diff/2) + additional_padding
+
+        padded_jpg = cv.copyMakeBorder(jpg_image,
+                                    int(pad_top),
+                                    int(pad_bottom),
+                                    int(pad_left),
+                                    int(pad_right),
+                                    borderType=cv.cv2.BORDER_CONSTANT,
+                                    value=[255,255,255])
+
+        os.chdir(out_list[0][0])
+        cv.imwrite(
+            id + 
+            f'_Padded.png',
+            padded_jpg
+            )
+
+        [pad_height,pad_width,z] = padded_jpg.shape
+        
+
+        width_small = int(pad_width/reduction_size)
+        height_small = int(pad_height/reduction_size)
+        jpg_small = cv.resize(padded_jpg,[width_small,height_small],cv.INTER_AREA)
+        
+        os.chdir(out_list[1][0])
+        cv.imwrite(
+            id + 
+            f'_Padded_' + reduction_name + '.png',
+            jpg_small
+            )
+
+
+        if seg_files:
+            ml_file = glob(ML_directory + f'{id}_*.png')[0]
+            ml_image = cv.imread(ml_file)[:,:,0]
+            padded_seg = cv.copyMakeBorder(ml_image,
+                                        int(pad_top),
+                                        int(pad_bottom),
+                                        int(pad_left),
+                                        int(pad_right),
+                                        borderType=cv.cv2.BORDER_CONSTANT,
+                                        value=[0,0,0])
+            os.chdir(out_list[0][1])
+            cv.imwrite(
+                id + 
+                f'_Padded_Seg.png',
+                padded_seg
+                )
+
+            seg_small = np.array(Image.fromarray(padded_seg).resize((width_small,height_small), Image.NEAREST))
+            os.chdir(out_list[1][1])
+            cv.imwrite(
+                id + 
+                f'_Padded_Seg_' + reduction_name + '.png',
+                seg_small
+                )
+
+
+        if nodal_files:
+
+            if old_files:
+                nodal_file = glob(Nodal_Seg_directory + f'{id}-*.png')[0]
+            else:
+                nodal_file = glob(Nodal_Seg_directory + f'{id}_*.png')[0]
+
+            nodal_image = cv.imread(nodal_file)[:,:,0]
+            # be warry of this, you may need to use this later, though I'm not sure what
+            # it was originally used for.
+            if nodal_white:
+                if sum(sum(nodal_image)) > 0:
+                    nodal_image = ~nodal_image
+            # This accounts for the nodal segmentation images being a quarter the
+            # original size, but you should make sure that you haven't already done the 
+            # fullscale noPad stuff yet
+            if ~old_files:
+                nodal_image = np.array(Image.fromarray(nodal_image).resize((width,height), Image.NEAREST))
+
+            padded_nodal = cv.copyMakeBorder(nodal_image,
+                                        int(pad_top),
+                                        int(pad_bottom),
+                                        int(pad_left),
+                                        int(pad_right),
+                                        borderType=cv.cv2.BORDER_CONSTANT,
+                                        value=[0,0,0])
+
+            os.chdir(out_list[0][2])
+            cv.imwrite(
+                id + 
+                f'_Padded_Nodal.png',
+                padded_nodal
+                )
+
+            nodal_small = np.array(Image.fromarray(padded_nodal).resize((width_small,height_small), Image.NEAREST))
+            os.chdir(out_list[1][2])
+            cv.imwrite(
+                id + 
+                f'_Padded_Nodal_' + reduction_name + '.png',
+                nodal_small
+                )
+
+
+
+
+
+        if high_res_files:
+            high_res_file = glob(high_res_directory + f'{id}_*.png')[0]
+            high_res_image = cv.imread(high_res_file)[:,:,0]
+
+            padded_high_res = cv.copyMakeBorder(high_res_image,
+                                        int(pad_top),
+                                        int(pad_bottom),
+                                        int(pad_left),
+                                        int(pad_right),
+                                        borderType=cv.cv2.BORDER_CONSTANT,
+                                        value=[0,0,0])
+
+
+            os.chdir(out_list[0][3])
+            cv.imwrite(
+                id + 
+                f'_Padded_HighRes.png',
+                padded_high_res
+                )
+
+            high_res_small = np.array(Image.fromarray(padded_high_res).resize((width_small,height_small), Image.NEAREST))
+            os.chdir(out_list[1][3])
+            cv.imwrite(
+                id + 
+                f'_Padded_HighRes_' + reduction_name + '.png',
+                high_res_small
+                )
+
+
+
+
+
+        if fiduciary_files:
+            if old_files:
+                fiduciary_file = glob(Fiduciary_Seg_directory + f'{id}-*.png')[0]
+            else:
+                fiduciary_file = glob(Fiduciary_Seg_directory + f'{id}_*.png')[0]
+
+            fiduciary_image = cv.imread(fiduciary_file)[:,:,0]
+
+            if fiduciary_white:
+                if sum(sum(fiduciary_image)) > 0:
+                    fiduciary_image = ~fiduciary_image
+
+            if ~old_files:
+                fiduciary_image = np.array(Image.fromarray(fiduciary_image).resize((width,height), Image.NEAREST))
+
+            padded_fiduciary = cv.copyMakeBorder(fiduciary_image,
+                                        int(pad_top),
+                                        int(pad_bottom),
+                                        int(pad_left),
+                                        int(pad_right),
+                                        borderType=cv.cv2.BORDER_CONSTANT,
+                                        value=[0,0,0])
+            os.chdir(out_list[0][4])
+            cv.imwrite(
+                id + 
+                f'_Padded_Fiduciary.png',
+                padded_fiduciary
+                )
+
+            fiduciary_small = np.array(Image.fromarray(padded_fiduciary).resize((width_small,height_small), Image.NEAREST))
+            os.chdir(out_list[1][4])
+            cv.imwrite(
+                id + 
+                f'_Padded_Fiduciary_' + reduction_name + '.png',
+                fiduciary_small
+                )
+
+    # %%
+
+
+# --------------------------------------------------
+if __name__ == '__main__':
+    main()