--- 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()