#!/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()
Datasets
Models
