import os import json import numpy as np import argparse import SimpleITK as sitk import cv2 from PIL import Image, ImageDraw, ImageFont import imageio class_tag_1 = {1: '磨玻璃结节', 2: '实性结节', 3: '混合结节', 4: '钙化结节'} class_tag_2 = {1: '良性', 2: '恶性'} def norm(image, hu_min=-1000.0, hu_max=600.0): image = (np.clip(image.astype(np.float32), hu_min, hu_max) - hu_min) / float(hu_max - hu_min) return image * 255. def vis(args): # 按病人id和时间整理匹配结果 followup_root = os.path.join(args.job_data_root, 'output/elastix') files = os.listdir(followup_root) nii_path = os.path.join(args.job_data_root, 'output/preprocess/preprocess_file_details/nii') nii_files = os.listdir(nii_path) working_dicts = {} for name in files: patientid = name.split('_')[0].split('-')[-1] if patientid not in working_dicts.keys(): working_dicts[patientid] = [] working_dicts[patientid].append(os.path.join(followup_root, name)) else: working_dicts[patientid].append(os.path.join(followup_root, name)) for key in list(working_dicts.keys()): working_dicts[key] = sorted(working_dicts[key]) # 计算每一个病人的信息，（size、类别） working_infos = {} for key in list(working_dicts.keys()): files = working_dicts[key] for file in files: working_infos[file] = [] temp= {'info': [], 'matched': []} match_results = json.load(open(os.path.join(file, 'match_results.json'), 'r')) temp['info'] = match_results['info'] if len(match_results['matched']) != 0: temp['matched'] = match_results['matched'] working_infos[file].append(temp) for file in list(working_infos.keys()): fix_uid = working_infos[file][0]['info'][0] fix_time = working_infos[file][0]['info'][1] matched_pairs = working_infos[file][0]['matched'] moving_uid = working_infos[file][0]['info'][2] moving_time = working_infos[file][0]['info'][3] fix_nii = [os.path.join(nii_path, file) for file in nii_files if fix_uid in file][0] moving_nii = [os.path.join(nii_path, file) for file in nii_files if moving_uid in file][0] fix_img = sitk.ReadImage(fix_nii) fix_array = sitk.GetArrayFromImage(fix_img) fix_origin = fix_img.GetOrigin() fix_spacing = fix_img.GetSpacing() moving_img = sitk.ReadImage(moving_nii) moving_array = sitk.GetArrayFromImage(moving_img) moving_origin = moving_img.GetOrigin() moving_spacing = moving_img.GetSpacing() for idx, match in enumerate(matched_pairs): p1 = match[0][:-2] p2 = match[1][:-2] p1[:3] = (np.array(p1[:3]) - np.array(list(fix_origin))) / np.array(list(fix_spacing)) p2[:3] = (np.array(p2[:3]) - np.array(list(moving_origin))) / np.array(list(moving_spacing)) fix_x1 = int(max(0, p1[0] - p1[3] / 2)) fix_y1 = int(max(0, p1[1] - p1[4] / 2)) fix_z1 = int(max(0, p1[2] - p1[5] / 2)) fix_x2 = int(min(fix_array.shape[2] - 1, p1[0] + p1[3] / 2)) fix_y2 = int(min(fix_array.shape[1] - 1, p1[1] + p1[4] / 2)) fix_z2 = int(min(fix_array.shape[0] - 1, p1[2] + p1[5] / 2)) fix_images = [] z_min = max(0, fix_z1-4) z_max = min(fix_array.shape[0]-1, fix_z2+4) for z in range(z_min, z_max, 1): img_z = fix_array.copy()[z, :, :] img_z = norm(img_z) img_z = cv2.cvtColor(img_z, cv2.COLOR_GRAY2BGR) # cv2.rectangle(img_z, (int(y1), int(x1)), (int(y2), int(x2)), (255,0,0), 2) if z>=fix_z1 and z<=fix_z2: cv2.rectangle(img_z, (int(fix_x1), int(fix_y1)), (int(fix_x2), int(fix_y2)), (255,0,0), 2) cut_x_min = max(0, fix_x1-128) cut_y_min = max(0, fix_y1-128) cut_x_max = min(fix_array.shape[2]-1, fix_x2+128) cut_y_max = min(fix_array.shape[1]-1, fix_y2+128) img_z = img_z[cut_y_min: cut_y_max, cut_x_min: cut_x_max, :] # img_z = img_z[cut_x_min: cut_x_max, cut_y_min: cut_y_max, :] img_pil = Image.fromarray(cv2.cvtColor(img_z, cv2.COLOR_BGR2RGB).astype(np.uint8)) draw = ImageDraw.Draw(img_pil) font = cv2.FONT_HERSHEY_SIMPLEX font_scale = 0.2 font_color = (255,0,0) font_thickness = 1 font_size = 10 text_1 = fix_uid (_, text_h_1), _ = cv2.getTextSize(text_1, font, font_scale, font_thickness) adjusted_position_1 = (0, 0+text_h_1) text_2 = fix_time (_, text_h_2), _ = cv2.getTextSize(text_2, font, font_scale, font_thickness) adjusted_position_2 = (0, 0+text_h_1+text_h_2+5) text_3 = str(0) (_, text_h_3), _ = cv2.getTextSize(text_3, font, font_scale, font_thickness) adjusted_position_3 = (0, 0+text_h_1+text_h_2+text_h_3+10) img_z = cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR) cv2.putText(img_z, text_1, adjusted_position_1, font, font_scale, font_color, font_thickness) cv2.putText(img_z, text_2, adjusted_position_2, font, font_scale, font_color, font_thickness) cv2.putText(img_z, text_3, adjusted_position_3, font, font_scale, font_color, font_thickness) fix_images.append(img_z) fps = 24.0 imageio.mimsave(os.path.join(file, fix_uid + ':' + fix_time + '-' + str(idx) + '.gif'), fix_images, 'GIF', duration=1/fps) moving_x1 = int(max(0, p2[0] - p2[3] / 2)) moving_y1 = int(max(0, p2[1] - p2[4] / 2)) moving_z1 = int(max(0, p2[2] - p2[5] / 2)) moving_x2 = int(min(moving_array.shape[2] - 1, p2[0] + p2[3] / 2)) moving_y2 = int(min(moving_array.shape[1] - 1, p2[1] + p2[4] / 2)) moving_z2 = int(min(moving_array.shape[0] - 1, p2[2] + p2[5] / 2)) moving_images = [] z_min = max(0, moving_z1-4) z_max = min(moving_array.shape[0]-1, moving_z2+4) for z in range(z_min, z_max, 1): img_z = moving_array.copy()[z, :, :] img_z = norm(img_z) img_z = cv2.cvtColor(img_z, cv2.COLOR_GRAY2BGR) # cv2.rectangle(img_z, (int(y1), int(x1)), (int(y2), int(x2)), (0,0,255), 2) if z>=moving_z1 and z<=moving_z2: cv2.rectangle(img_z, (int(moving_x1), int(moving_y1)), (int(moving_x2), int(moving_y2)), (255,0,0), 2) cut_x_min = max(0, moving_x1-128) cut_y_min = max(0, moving_y1-128) cut_x_max = min(moving_array.shape[2]-1, moving_x2+128) cut_y_max = min(moving_array.shape[1]-1, moving_y2+128) img_z = img_z[cut_y_min: cut_y_max, cut_x_min: cut_x_max, :] # img_z = img_z[cut_x_min: cut_x_max, cut_y_min: cut_y_max, :] img_pil = Image.fromarray(cv2.cvtColor(img_z, cv2.COLOR_BGR2RGB).astype(np.uint8)) draw = ImageDraw.Draw(img_pil) text_1 = moving_uid (_, text_h_1), _ = cv2.getTextSize(text_1, font, font_scale, font_thickness) adjusted_position_1 = (0, 0+text_h_1) text_2 = moving_time (_, text_h_2), _ = cv2.getTextSize(text_2, font, font_scale, font_thickness) adjusted_position_2 = (0, 0+text_h_1+text_h_2+5) text_3 = str(idx) (_, text_h_3), _ = cv2.getTextSize(text_3, font, font_scale, font_thickness) adjusted_position_3 = (0, 0+text_h_1+text_h_2+text_h_3+10) font = cv2.FONT_HERSHEY_SIMPLEX font_scale = 0.2 font_color = (255,0,0) font_thickness = 1 font_size = 10 img_z = cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR) cv2.putText(img_z, text_1, adjusted_position_1, font, font_scale, font_color, font_thickness) cv2.putText(img_z, text_2, adjusted_position_2, font, font_scale, font_color, font_thickness) cv2.putText(img_z, text_3, adjusted_position_3, font, font_scale, font_color, font_thickness) moving_images.append(img_z) fps = 24.0 imageio.mimsave(os.path.join(file, moving_uid + ':' + moving_time + '-' + str(idx) + '.gif'), moving_images, 'GIF', duration=1/fps) if __name__ == '__main__': parser = argparse.ArgumentParser(description='followup processing') parser.add_argument('--job_data_root', default='/data/job_715/job_data_preprocess', type=str, help='model') args = parser.parse_args() vis(args)