-
Notifications
You must be signed in to change notification settings - Fork 0
/
main_step2c_generate_selected_plane_movies_low_res.py
93 lines (66 loc) · 3.92 KB
/
main_step2c_generate_selected_plane_movies_low_res.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
#!/usr/bin/env python
# this step is optional. It can give you better DL product.
# this script generates adequate plane movies using the model set selection you defined in main_step2b.
import function_list as ff
import os
import numpy as np
import nibabel as nib
import supplements
import pandas as pd
from PIL import Image
from Make_plane_movies import *
cg = supplements.Experiment()
WL = 500
WW = 800
######### Define patient_list
main_folder = os.path.join(cg.save_dir,'DL_prediction_low_res')
patient_list = ff.find_all_target_files(['Abnormal/CVC1803*'],main_folder)
######### Load the spreadsheet
sheet = pd.read_excel(os.path.join(cg.save_dir,'model_set_selection_example.xlsx'))
# main script (usually no need to change):
for patient in patient_list:
patient_id = os.path.basename(patient)
patient_class = os.path.basename(os.path.dirname(patient))
print(patient_class,patient_id)
# read model_set_selection spreadsheet info
info = sheet[sheet['Patient_ID'] == patient_id]
assert info.shape[0] == 1
model_set_pick = [int(info.iloc[0]['seg (default = 3)']), int(info.iloc[0]['2CH']),int(info.iloc[0]['3CH']),int(info.iloc[0]['4CH']),int(info.iloc[0]['SAX'])]
LAX_zoom = float(info.iloc[0]['Zoom_LAX (default = 1.2)'])
SAX_zoom = float(info.iloc[0]['Zoom_SAX (default = 1.4)'])
# define save folder
save_folder = os.path.join(patient,'planes_pred_low_res_Final')
ff.make_folder([save_folder, os.path.join(save_folder,'pngs')])
# check whether already done
if os.path.isfile(os.path.join(save_folder,patient_id+'_planes.mp4')) == 1:
print('already done for this patient')
continue
# load plane vectors
prepare = Prepare(main_folder,patient_class,patient_id, 0 ,cg.low_res_spacing,cg.low_res_dim)
image_center, vector_2C, vector_3C, vector_4C, vector_SA, normal_vector_SA = prepare.load_plane_vectors(batch_select = True, select_list = model_set_pick[1:])
# get affine matrices
volume_affine, A_2C, A_3C, A_4C = prepare.obtain_affine_matrix(vector_2C, vector_3C, vector_4C)
# get the range of SAX stack using the LV segmentation, should already obtain this info as txt file in main_step2
slice_num_info_file_name = os.path.join(patient,"slice_num_info_low_res_batch_" + str(model_set_pick[-1]) + ".txt")
if os.path.isfile(slice_num_info_file_name) == 1:
a, b = prepare.read_SAX_slice_num_file(slice_num_info_file_name)
else: # txt file has not been made
seg = nib.load(os.path.join(patient,'seg-pred/batch_'+str(model_set_pick[0]),'pred_s_0.nii.gz'))
seg_LV = seg.get_fdata()
a , b = prepare.define_SAX_range(vector_SA,image_center, seg_LV, model_set_pick[-1])
# get a center list of 9-plane SAX stack
_, center_list9, gap = prepare.define_SAX_planes_center_list(vector_SA, image_center, a, b, normal_vector_SA)
# get the image list
img_list = ff.sort_timeframe(ff.find_all_target_files(['img-nii-' + str(cg.low_res_spacing) + '/*.nii.gz'],os.path.join(cg.image_data_dir,patient_class,patient_id)),2)
# make the plane images for each time frame
for img in img_list:
volume_data = nib.load(img).get_fdata()
time = ff.find_timeframe(img,2)
save_path = os.path.join(save_folder,'pngs',str(time) +'.png')
make = Make_Planes(main_folder, patient_class, patient_id, volume_data,[WL,WW], image_center,[vector_2C, vector_3C, vector_4C, vector_SA], volume_affine, center_list9 ,cg.low_res_spacing, cg.low_res_dim, zoom_factor_LAX = LAX_zoom, zoom_factor_SAX = SAX_zoom)
make.plane_image(save_path, draw_lines = True) # draw_lines means use lines to represent SAX/LAX planes on LAX/SAX planes
print('finish time frame ',time)
# make cine movie
pngs = ff.sort_timeframe(ff.find_all_target_files(['*.png'],os.path.join(save_folder,'pngs')),1)
save_movie_path = os.path.join(save_folder,patient_id+'_planes.mp4')
ff.make_movies(save_movie_path,pngs)