From d01af7bdd8ac283158ef78fc5d91bb5ff0bb96eb Mon Sep 17 00:00:00 2001 From: xming521 <1223398803@qq.com> Date: Mon, 17 Feb 2020 16:17:23 +0800 Subject: [PATCH] upload train code --- CTAI_model/cv/__init__.py | 0 CTAI_model/cv/get_ROI-all.py | 114 +++++++++++++++ CTAI_model/cv/get_ROI-one.py | 68 +++++++++ CTAI_model/cv/get_all_feature.py | 232 +++++++++++++++++++++++++++++++ CTAI_model/net/__init__.py | 0 CTAI_model/net/test.py | 53 +++++++ CTAI_model/net/train.py | 121 ++++++++++++++++ CTAI_model/net/unet.py | 68 +++++++++ CTAI_model/requirements.txt | 13 ++ CTAI_model/utils/dice_loss.py | 38 +++++ CTAI_model/utils/draw.py | 36 +++++ CTAI_model/utils/transform.py | 44 ++++++ CTAI_model/utils/校验文件.py | 44 ++++++ 13 files changed, 831 insertions(+) create mode 100644 CTAI_model/cv/__init__.py create mode 100644 CTAI_model/cv/get_ROI-all.py create mode 100644 CTAI_model/cv/get_ROI-one.py create mode 100644 CTAI_model/cv/get_all_feature.py create mode 100644 CTAI_model/net/__init__.py create mode 100644 CTAI_model/net/test.py create mode 100644 CTAI_model/net/train.py create mode 100644 CTAI_model/net/unet.py create mode 100644 CTAI_model/requirements.txt create mode 100644 CTAI_model/utils/dice_loss.py create mode 100644 CTAI_model/utils/draw.py create mode 100644 CTAI_model/utils/transform.py create mode 100644 CTAI_model/utils/校验文件.py diff --git a/CTAI_model/cv/__init__.py b/CTAI_model/cv/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/CTAI_model/cv/get_ROI-all.py b/CTAI_model/cv/get_ROI-all.py new file mode 100644 index 0000000..3490df8 --- /dev/null +++ b/CTAI_model/cv/get_ROI-all.py @@ -0,0 +1,114 @@ +import SimpleITK as sitk +import cv2 +import numpy as np + +from data_set.make import get_train_files + +# 跑train不加第二个train +train_data_path = '../data/train/train/' + + +# train_data_path = '../data/CT/' + + +def get_roi(path): + global w + file_name = path.split('/')[-3] + '-' + path.split('/')[-1].replace('.dcm', '') + image = sitk.ReadImage(path) + image = sitk.GetArrayFromImage(image)[0, :, :] + image[image < -300] = 0 + image[image > 300] = 0 + img_o = image.copy() + ROI = np.zeros(image.shape, np.uint8) + slices = [image] + img = slices[int(len(slices) / 2)].copy() + img = np.uint8(img) + # kernel = np.ones((3, 3), np.uint8) + # kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, 1)) + # img = cv2.morphologyEx(img, cv2.MORPH_OPEN, kernel) + # img = cv2.dilate(img, kernel, iterations=1) + + kernel = np.ones((4, 4), np.uint8) + img = cv2.dilate(img, kernel, iterations=1) + + # 对图像进行阈值分割 + ret, img = cv2.threshold(img, 150, 255, cv2.THRESH_BINARY_INV) + # 提取分割结果中的轮廓，并填充孔洞 + im2, contours, x = cv2.findContours(img.copy(), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE) + area = [] + for c in contours: + area.append(cv2.contourArea(c)) + cparea = area.copy() + + area.sort(reverse=True) + ROI_tmp = np.zeros(img.shape, np.uint8) + # for i in range(1, 10): + # # 选择最有可能的一个区域 + # ROI_tmp = np.zeros(img.shape, np.uint8) + # max_idx = cparea.index(area[i]) + # + # # 强度匹配 因为直接映射到dcm上不用转uint8 + # cv2.drawContours(ROI_tmp, contours, max_idx, (255, 255, 255), -1) + # + # index = np.nonzero(ROI_tmp) + # mean = np.mean(img_o[index]) + # std = np.std(img_o[index]) + # # if mean > 90 or mean < 10 or std > 70 or std < 15: + # # continue + # + # s = pd.Series(img_o[index]) + # piandu = s.skew() + # fengdu = s.kurt() + # # if piandu > 1 or fengdu < -5 or fengdu > 20: + # # continue + # + # # 生成矩 + # M = cv2.moments(contours[max_idx]) + # + # # 面积周长 + # perimeter = cv2.arcLength(contours[max_idx], True) + # # if area[i] > 2000 or area[i] < 500 or perimeter > 250 or perimeter < 80: + # # continue + # if area[i] > 2000 or perimeter > 400 : + # continue + # + # # if area[i] > 4000 or area[i] < 500 or perimeter > 400 : + # # continue + # # + # # 椭圆拟合 + # # try: + # # (x, y), (MA, ma), angle = cv2.fitEllipse(contours[max_idx]) + # # if ma - MA > 25: + # # continue + # # # ellipse.append(ma-MA) + # # except: + # # continue + # + # + # # img_o[] + # + # # 加矩形框 + # x, y, w, h = cv2.boundingRect(contours[max_idx]) + # ROI = cv2.rectangle(ROI, (x, y), (x + w + 10, y + h + 10), (255, 255, 255), -1) + # + # # 填充 + # # cv2.drawContours(ROI, contours, max_idx, (255, 255, 255), -1) + + ROI_tmp[270:430, 200:300] = image[270:430, 200:300] + + cv2.imshow("Image", image) + cv2.imshow("Image", ROI_tmp) + cv2.waitKey(0) + print(f"{train_data_path}ROI-{file_name}.png") + # cv2.imwrite(f"{train_data_path}ROI-{file_name}.png", ROI, [int(cv2.IMWRITE_PNG_COMPRESSION), 0]) + + +def main(): + global w + dcm_files, _ = get_train_files(train_data_path, file_type='dcm', all=False) + for i in dcm_files: + get_roi(i) + + +if __name__ == '__main__': + main() diff --git a/CTAI_model/cv/get_ROI-one.py b/CTAI_model/cv/get_ROI-one.py new file mode 100644 index 0000000..851a658 --- /dev/null +++ b/CTAI_model/cv/get_ROI-one.py @@ -0,0 +1,68 @@ +import SimpleITK as sitk +import cv2 +import numpy as np + +image = sitk.ReadImage('../data/train/train/10087.dcm') +image = sitk.GetArrayFromImage(image)[0, :, :] + +image[image < -300] = 0 +image[image > 300] = 0 + +ROI = np.zeros(image.shape, np.uint8) +# 获取图像中的像素数据 +slices = [image] + +# 复制Dicom图像中的像素数据 +img = slices[int(len(slices) / 2)].copy() +img = np.uint8(img) + +kernel = np.ones((4, 4), np.uint8) +img = cv2.dilate(img, kernel, iterations=1) + +# 对图像进行阈值分割 +ret, img = cv2.threshold(img, 150, 255, cv2.THRESH_BINARY_INV) + +xxx = img +# 提取分割结果中的轮廓，并填充孔洞 +im2, contours, x = cv2.findContours(img.copy(), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE) +# 需要反色处理一下 现在找的是白色的 应该是黑色的 + +# mask = np.zeros(img.shape, np.uint8) +# for contour in contours: +# cv2.fillPoly(mask, [contour], 255) +# img[(mask > 0)] = 255 + + +area = [] +for c in contours: + area.append(cv2.contourArea(c)) +cparea = area.copy() + +area.sort(reverse=True) + +for i in range(3, 8): + max_idx = cparea.index(area[i]) + perimeter = cv2.arcLength(contours[max_idx], True) + if area[i] > 5000 or perimeter > 500: + continue + print('周长', perimeter) + + cv2.drawContours(ROI, contours, max_idx, (220, 20, 60), -1) + +# max_idx = cparea.index(area[3]) +# cv2.drawContours(ROI, contours, max_idx, (220, 20, 60), -1) +# cv2.drawContours(ROI, contours, max_idx, (220, 20, 60), -1) + + +# 对分割结果进行形态学的开操作 +# kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2, 2)) +# img = cv2.morphologyEx(img, cv2.MORPH_OPEN, kernel) + +# plt.figure(figsize=(10, 7)) +# plt.imshow(img, 'gray') +# plt.title('Mask') +# plt.show() + + +cv2.imshow("Image", ROI) +cv2.waitKey(0) diff --git a/CTAI_model/cv/get_all_feature.py b/CTAI_model/cv/get_all_feature.py new file mode 100644 index 0000000..e966e95 --- /dev/null +++ b/CTAI_model/cv/get_all_feature.py @@ -0,0 +1,232 @@ +import sys + +sys.path.append("..") +from data_set.make import get_person_files +import cv2 +import numpy as np +import SimpleITK as sitk +import pandas as pd +from numba import jit +import inspect +import csv + +np.set_printoptions(suppress=True) # 输出时禁止科学表示法，直接输出小数值 + +mask_data_path = '../data/all/d2/' + +# ID为第一个字段名会会让xlsx报错 +# dict用表头 +column_all = ['ID', '性别', '年龄', '阴性/阳性', 'area', 'perimeter', 'focus_x', 'focus_y', 'ellipse', 'mean', 'std', 'piandu', + 'fengdu', + 'small_grads_dominance', + 'big_grads_dominance', 'gray_asymmetry', 'grads_asymmetry', 'energy', 'gray_mean', 'grads_mean', + 'gray_variance', 'grads_variance', 'corelation', 'gray_entropy', 'grads_entropy', 'entropy', 'inertia', + 'differ_moment'] +# 实际表头 +column_all_c = ['ID', '性别', '年龄', '阴性/阳性', '面积', '周长', '重心x', '重心y', '近圆度', '灰度均值', '灰度方差', '灰度偏度', + '灰度峰态', '小梯度优势', '大梯度优势', '灰度分布不均匀性', '梯度分布不均匀性', '能量', '灰度平均', '梯度平均', + '灰度均方差', '梯度均方差', '相关', '灰度熵', '梯度熵', '混合熵', '惯性', '逆差矩'] + +features_list = ['area', 'perimeter', 'focus_x', 'focus_y', 'ellipse', 'mean', 'std', 'piandu', 'fengdu', + 'small_grads_dominance', + 'big_grads_dominance', 'gray_asymmetry', 'grads_asymmetry', 'energy', 'gray_mean', 'grads_mean', + 'gray_variance', 'grads_variance', 'corelation', 'gray_entropy', 'grads_entropy', 'entropy', 'inertia', + 'differ_moment'] + + +# 最后俩偏度 峰度 + + +# 获取变量的名 +def get_variable_name(variable): + callers_local_vars = inspect.currentframe().f_back.f_locals.items() + return [var_name for var_name, var_val in callers_local_vars if var_val is variable] + + +def glcm(img_gray, ngrad=16, ngray=16): + '''Gray Level-Gradient Co-occurrence Matrix,取归一化后的灰度值、梯度值分别为16、16''' + # 利用sobel算子分别计算x-y方向上的梯度值 + gsx = cv2.Sobel(img_gray, cv2.CV_64F, 1, 0, ksize=3) + gsy = cv2.Sobel(img_gray, cv2.CV_64F, 0, 1, ksize=3) + height, width = img_gray.shape + grad = (gsx ** 2 + gsy ** 2) ** 0.5 # 计算梯度值 + grad = np.asarray(1.0 * grad * (ngrad - 1) / grad.max(), dtype=np.int16) + gray = np.asarray(1.0 * img_gray * (ngray - 1) / img_gray.max(), dtype=np.int16) # 0-255变换为0-15 + gray_grad = np.zeros([ngray, ngrad]) # 灰度梯度共生矩阵 + for i in range(height): + for j in range(width): + gray_value = gray[i][j] + grad_value = grad[i][j] + gray_grad[gray_value][grad_value] += 1 + gray_grad = 1.0 * gray_grad / (height * width) # 归一化灰度梯度矩阵，减少计算量 + get_glcm_features(gray_grad) + + +@jit +def get_gray_feature(): + # 灰度特征提取算法 + hist = cv2.calcHist([image_ROI_uint8[index]], [0], None, [256], [0, 256]) + + c_features['mean'].append(np.mean(image_ROI[index])) + c_features['std'].append(np.std(image_ROI[index])) + + s = pd.Series(image_ROI[index]) + c_features['piandu'].append(s.skew()) + c_features['fengdu'].append(s.kurt()) + + +def get_glcm_features(mat): + '''根据灰度梯度共生矩阵计算纹理特征量，包括小梯度优势，大梯度优势，灰度分布不均匀性，梯度分布不均匀性，能量，灰度平均，梯度平均， + 灰度方差，梯度方差，相关，灰度熵，梯度熵，混合熵，惯性，逆差矩''' + sum_mat = mat.sum() + small_grads_dominance = big_grads_dominance = gray_asymmetry = grads_asymmetry = energy = gray_mean = grads_mean = 0 + gray_variance = grads_variance = corelation = gray_entropy = grads_entropy = entropy = inertia = differ_moment = 0 + for i in range(mat.shape[0]): + gray_variance_temp = 0 + for j in range(mat.shape[1]): + small_grads_dominance += mat[i][j] / ((j + 1) ** 2) + big_grads_dominance += mat[i][j] * j ** 2 + energy += mat[i][j] ** 2 + if mat[i].sum() != 0: + gray_entropy -= mat[i][j] * np.log(mat[i].sum()) + if mat[:, j].sum() != 0: + grads_entropy -= mat[i][j] * np.log(mat[:, j].sum()) + if mat[i][j] != 0: + entropy -= mat[i][j] * np.log(mat[i][j]) + inertia += (i - j) ** 2 * np.log(mat[i][j]) + differ_moment += mat[i][j] / (1 + (i - j) ** 2) + gray_variance_temp += mat[i][j] ** 0.5 + + gray_asymmetry += mat[i].sum() ** 2 + gray_mean += i * mat[i].sum() ** 2 + gray_variance += (i - gray_mean) ** 2 * gray_variance_temp + for j in range(mat.shape[1]): + grads_variance_temp = 0 + for i in range(mat.shape[0]): + grads_variance_temp += mat[i][j] ** 0.5 + grads_asymmetry += mat[:, j].sum() ** 2 + grads_mean += j * mat[:, j].sum() ** 2 + grads_variance += (j - grads_mean) ** 2 * grads_variance_temp + small_grads_dominance /= sum_mat + big_grads_dominance /= sum_mat + gray_asymmetry /= sum_mat + grads_asymmetry /= sum_mat + gray_variance = gray_variance ** 0.5 + grads_variance = grads_variance ** 0.5 + for i in range(mat.shape[0]): + for j in range(mat.shape[1]): + corelation += (i - gray_mean) * (j - grads_mean) * mat[i][j] + glgcm_features = [small_grads_dominance, big_grads_dominance, gray_asymmetry, grads_asymmetry, energy, gray_mean, + grads_mean, + gray_variance, grads_variance, corelation, gray_entropy, grads_entropy, entropy, inertia, + differ_moment] + for i in range(len(glgcm_features)): + t = get_variable_name(glgcm_features[i])[0] + c_features[t].append(np.round(glgcm_features[i], 4)) + + +def get_geometry_feature(): + # 形态特征 分割mask获得一些特征 + im2, contours, x = cv2.findContours(mask_array.copy(), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE) + tarea = [] + tperimeter = [] + for c in contours: + # 生成矩 + try: + M = cv2.moments(c) + cx = int(M['m10'] / M['m00']) + cy = int(M['m01'] / M['m00']) + c_features['focus_x'].append(cx) + c_features['focus_y'].append(cy) + except: + print('error') + + # 椭圆拟合 + try: + (x, y), (MA, ma), angle = cv2.fitEllipse(c) + c_features['ellipse'].append((ma - MA)) + except: + continue + # 面积周长 + tarea.append(cv2.contourArea(c)) + tperimeter.append(cv2.arcLength(c, True)) + + # 将mask里的最大值追加 有黑洞 + try: + c_features['area'].append(max(tarea)) + c_features['perimeter'].append(round(max(tperimeter), 4)) + except: + print('area error') + + +# 提取肿瘤特征 +def get_feature(image, mask): + global w + global image_ROI_uint8, index, image_ROI_mini, image_ROI, mask_array + + mask_array = cv2.imread(mask, 0) + image = sitk.ReadImage(image) + image_arrary = sitk.GetArrayFromImage(image)[0, :, :] + # 映射到CT获得特征 + image_ROI = np.zeros(shape=image_arrary.shape) + index = np.nonzero(mask_array) + if not index[0].any(): + # c_features['no'] = True + return None + image_ROI[index] = image_arrary[index] + image_ROI_uint8 = np.uint8(image_ROI) + # 获得只有肿瘤的图片 + x, y, w, h = cv2.boundingRect(mask_array) + image_ROI_mini = np.uint8(image_arrary[y:y + h, x:x + w]) + w = image_ROI_mini + + get_geometry_feature() + + get_gray_feature() + + glcm(image_ROI_mini, 15, 15) + + return c_features + + +def main(): + global w + # 注意下面文件路径 格式 + csv_file = open('../data/all/res.csv', 'w', encoding='gbk', newline='') + writer1 = csv.writer(csv_file) + writer1.writerow(column_all_c) + writer2 = csv.DictWriter(csv_file, column_all) + + df = pd.read_csv('../data/all/临床数据.csv', encoding='gbk') + + image = get_person_files('../data/all/d2/') + mask = get_person_files('../data/out/') + # image1 mask2 + + for i in range(len(image)): + person_id = image[i][0] + global c_features + c_features = {} + for k in features_list: + c_features[k] = [] + if len(image[i][1]) != len(mask[i][2]): + print('文件有错') + for j in range(len(image[i][1])): + get_feature(image[i][1][j], mask[i][2][j]) + + for j in c_features: + if j == 'id': + continue + c_features[j] = np.round(np.mean(c_features[j]), 4) + + person_info = df[df['ID'].isin([person_id])].to_dict('index').values() + person_info = list(person_info)[0] + person_info.update(c_features) + writer2.writerow(person_info) + print(person_info) + + csv_file.close() + + +if __name__ == '__main__': + main() diff --git a/CTAI_model/net/__init__.py b/CTAI_model/net/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/CTAI_model/net/test.py b/CTAI_model/net/test.py new file mode 100644 index 0000000..4945610 --- /dev/null +++ b/CTAI_model/net/test.py @@ -0,0 +1,53 @@ +import os +import sys + +import cv2 + +sys.path.append("..") +import torch +from torch.utils.data import DataLoader +from data_set import make + +os.environ["CUDA_VISIBLE_DEVICES"] = "0" +torch.set_num_threads(4) +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +torch.cuda.empty_cache() +res = {'epoch': [], 'loss': [], 'dice': []} + +test_data_path = '../data/all/d2/' +rate = 0.5 + +test_dataset = make.get_d1_local(test_data_path) + +import os + + +def mkdir(path): + folder = os.path.exists(path) + if not folder: # 判断是否存在文件夹如果不存在则创建为文件夹 + os.makedirs(path) # makedirs 创建文件时如果路径不存在会创建这个路径 + + +def onlytest(): + unet = torch.load('../result/0.5unet.pkl').to(device) + global res, img_y, mask_arrary + epoch_dice = 0 + with torch.no_grad(): + dataloaders = DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=0) + for x in dataloaders: + id = x[1:] # ('1026',), ('10018',)]先病人号后片号 + print(id, 'id') + x = x[0].to(device) + y = unet(x) + img_y = torch.squeeze(y).cpu().numpy() + img_y[img_y >= rate] = 1 + img_y[img_y < rate] = 0 + img_y = img_y * 255 + mkdir(f'../data/out/{id[0][0]}/arterial phase/') + cv2.imwrite(f'../data/out/{id[0][0]}/arterial phase/{id[1][0]}_mask.png', img_y, + (cv2.IMWRITE_PNG_COMPRESSION, 0)) + + +if __name__ == '__main__': + # train() + onlytest() diff --git a/CTAI_model/net/train.py b/CTAI_model/net/train.py new file mode 100644 index 0000000..0bfa334 --- /dev/null +++ b/CTAI_model/net/train.py @@ -0,0 +1,121 @@ +import sys + +sys.path.append("..") +import torch +from torch.nn import init +from torch.utils.data import DataLoader +from data_set import make +from net import unet +from utils import dice_loss +import matplotlib.pyplot as plt +import numpy as np + +# os.environ["CUDA_VISIBLE_DEVICES"] = "1" +torch.set_num_threads(1) +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +torch.cuda.empty_cache() +res = {'epoch': [], 'loss': [], 'dice': []} + + +def weights_init(m): + classname = m.__class__.__name__ + # print(classname) + if classname.find('Conv3d') != -1: + init.xavier_normal(m.weight.data, 0.0) + init.constant_(m.bias.data, 0.0) + elif classname.find('Linear') != -1: + init.xavier_normal(m.weight.data, 0.0) + init.constant_(m.bias.data, 0.0) + + +# 参数 +rate = 0.50 +learn_rate = 0.001 +epochs = 1 +# train_dataset_path = '../data/all/d1/' +train_dataset_path = 'E:/projects/python projects/ct_data/' + +train_dataset, test_dataset = make.get_d1(train_dataset_path) +unet = unet.Unet(1, 1).to(device).apply(weights_init) +criterion = torch.nn.BCELoss().to(device) +optimizer = torch.optim.Adam(unet.parameters(), learn_rate) + + +def train(): + global res + dataloaders = DataLoader(train_dataset, batch_size=1, shuffle=True, num_workers=0) + for epoch in range(epochs): + dt_size = len(dataloaders.dataset) + epoch_loss, epoch_dice = 0, 0 + step = 0 + for x, y in dataloaders: + id = x[1:] + step += 1 + x = x[0].to(device) + y = y[1].to(device) + print(x.size()) + print(y.size()) + optimizer.zero_grad() + outputs = unet(x) + loss = criterion(outputs, y) + loss.backward() + optimizer.step() + # dice + # a = outputs.cpu().detach().squeeze(1).numpy() + # a[a >= rate] = 1 + # a[a < rate] = 0 + # b = y.cpu().detach().numpy() + # dice = dice_loss.dice(a, b) + # epoch_loss += float(loss.item()) + # epoch_dice += dice + + if step % 100 == 0: + res['epoch'].append((epoch + 1) * step) + res['loss'].append(loss.item()) + print("epoch%d step%d/%d train_loss:%0.3f" % ( + epoch, step, (dt_size - 1) // dataloaders.batch_size + 1, loss.item()), + end='') + test() + # print("epoch %d loss:%0.3f,dice %f" % (epoch, epoch_loss / step, epoch_dice / step)) + plt.plot(res['epoch'], np.squeeze(res['cost']), label='Train cost') + plt.ylabel('cost') + plt.xlabel('epochs') + plt.title("Model: train cost") + plt.legend() + + plt.plot(res['epoch'], np.squeeze(res), label='Validation cost', color='#FF9966') + plt.ylabel('loss') + plt.xlabel('epochs') + plt.title("Model:validation loss") + plt.legend() + + plt.savefig("examples.jpg") + + # torch.save(unet, 'unet.pkl') + # model = torch.load('unet.pkl') + test() + + +def test(): + global res, img_y, mask_arrary + epoch_dice = 0 + with torch.no_grad(): + dataloaders = DataLoader(test_dataset, batch_size=1, shuffle=True, num_workers=0) + for x, mask in dataloaders: + id = x[1:] # ('1026',), ('10018',)]先病人号后片号 + x = x[0].to(device) + y = unet(x) + mask_arrary = mask[1].cpu().squeeze(0).detach().numpy() + img_y = torch.squeeze(y).cpu().numpy() + img_y[img_y >= rate] = 1 + img_y[img_y < rate] = 0 + img_y = img_y * 255 + epoch_dice += dice_loss.dice(img_y, mask_arrary) + # cv.imwrite(f'data/out/{mask[0][0]}-result.png', img_y, (cv.IMWRITE_PNG_COMPRESSION, 0)) + print('test dice %f' % (epoch_dice / len(dataloaders))) + res['dice'].append(epoch_dice / len(dataloaders)) + + +if __name__ == '__main__': + train() + test() diff --git a/CTAI_model/net/unet.py b/CTAI_model/net/unet.py new file mode 100644 index 0000000..f67d754 --- /dev/null +++ b/CTAI_model/net/unet.py @@ -0,0 +1,68 @@ +import torch +import torch.nn as nn + + +class DoubleConv(nn.Module): + def __init__(self, in_ch, out_ch): + super(DoubleConv, self).__init__() + self.conv = nn.Sequential( + nn.Conv2d(in_ch, out_ch, 3, padding=1), + nn.BatchNorm2d(out_ch), # 归一 + nn.ReLU(inplace=True), + nn.Conv2d(out_ch, out_ch, 3, padding=1), + nn.BatchNorm2d(out_ch), + nn.ReLU(inplace=True) + ) + + def forward(self, input): + return self.conv(input) + + +class Unet(nn.Module): + def __init__(self, in_ch, out_ch): + super(Unet, self).__init__() + + self.conv1 = DoubleConv(in_ch, 64) + self.pool1 = nn.MaxPool2d(2) + self.conv2 = DoubleConv(64, 128) + self.pool2 = nn.MaxPool2d(2) + self.conv3 = DoubleConv(128, 256) + self.pool3 = nn.MaxPool2d(2) + self.conv4 = DoubleConv(256, 512) + self.pool4 = nn.MaxPool2d(2) + self.conv5 = DoubleConv(512, 1024) + self.up6 = nn.ConvTranspose2d(1024, 512, 2, stride=2) + self.conv6 = DoubleConv(1024, 512) + self.up7 = nn.ConvTranspose2d(512, 256, 2, stride=2) + self.conv7 = DoubleConv(512, 256) + self.up8 = nn.ConvTranspose2d(256, 128, 2, stride=2) + self.conv8 = DoubleConv(256, 128) + self.up9 = nn.ConvTranspose2d(128, 64, 2, stride=2) + self.conv9 = DoubleConv(128, 64) + self.conv10 = nn.Conv2d(64, out_ch, 1) + + def forward(self, x): + c1 = self.conv1(x) + p1 = self.pool1(c1) + c2 = self.conv2(p1) + p2 = self.pool2(c2) + c3 = self.conv3(p2) + p3 = self.pool3(c3) + c4 = self.conv4(p3) + p4 = self.pool4(c4) + c5 = self.conv5(p4) + up_6 = self.up6(c5) + merge6 = torch.cat([up_6, c4], dim=1) + c6 = self.conv6(merge6) + up_7 = self.up7(c6) + merge7 = torch.cat([up_7, c3], dim=1) + c7 = self.conv7(merge7) + up_8 = self.up8(c7) + merge8 = torch.cat([up_8, c2], dim=1) + c8 = self.conv8(merge8) + up_9 = self.up9(c8) + merge9 = torch.cat([up_9, c1], dim=1) + c9 = self.conv9(merge9) + c10 = self.conv10(c9) + out = nn.Sigmoid()(c10) + return out diff --git a/CTAI_model/requirements.txt b/CTAI_model/requirements.txt new file mode 100644 index 0000000..0ff1bed --- /dev/null +++ b/CTAI_model/requirements.txt @@ -0,0 +1,13 @@ +Flask==1.1.1 +Werkzeug==0.15.2 +pandas==0.23.4 +numpy==1.15.0 +Pillow==6.1.0 +SimpleITK==1.2.2 +matplotlib==3.1.1 +numba==0.45.1 +pydicom==1.3.0 +torch==1.2.0 +torchvision==0.4.0 +opencv-python==3.4.5.20 + diff --git a/CTAI_model/utils/dice_loss.py b/CTAI_model/utils/dice_loss.py new file mode 100644 index 0000000..81265b2 --- /dev/null +++ b/CTAI_model/utils/dice_loss.py @@ -0,0 +1,38 @@ +import numpy as np + + +def dice(im1, im2): + """ + Computes the Dice coefficient, a measure of set similarity. + Parameters + ---------- + im1 : array-like, bool + Any array of arbitrary size. If not boolean, will be converted. + im2 : array-like, bool + Any other array of identical size. If not boolean, will be converted. + Returns + ------- + dice : float + Dice coefficient as a float on range [0,1]. + Maximum similarity = 1 + No similarity = 0 + + Notes + ----- + The order of inputs for `dice` is irrelevant. The result will be + identical if `im1` and `im2` are switched. + """ + im1 = np.asarray(im1).astype(np.bool) + im2 = np.asarray(im2).astype(np.bool) + + if im1.shape != im2.shape: + raise ValueError("Shape mismatch: im1 and im2 must have the same shape.") + + # 俩都为全黑 + if not (im1.any() or im2.any()): + return 1.0 + + # Compute Dice coefficient + intersection = np.logical_and(im1, im2) + res = 2. * intersection.sum() / (im1.sum() + im2.sum()) + return np.round(res, 5) diff --git a/CTAI_model/utils/draw.py b/CTAI_model/utils/draw.py new file mode 100644 index 0000000..d947d75 --- /dev/null +++ b/CTAI_model/utils/draw.py @@ -0,0 +1,36 @@ +import matplotlib.pyplot as plt +import numpy as np + +data = [] +data_true = [] +with open('../result/0.50nohup50.txt', 'r') as f: + data = [i.replace('\n', '') for i in f.readlines()] + +for i in range(len(data)): + if i % 3 == 0: + x = data[i].split(' ') + data_true.append([x[2].replace('test', '').replace('train_loss:', ''), x[-1]]) + +print(data_true) + +ax = plt.gca() + +plt.rcParams['savefig.dpi'] = 300 # 图片像素 +plt.rcParams['figure.dpi'] = 200 # 分辨率 + +# plt.plot(range(1,51), np.squeeze([i[0] for i in data_true]), label='Train loss') +# plt.ylabel('loss') +# plt.xlabel('epochs') +# plt.title("Model: train loss") +# plt.legend() +# plt.show() +ax.invert_yaxis() + +plt.plot(range(1, 51), np.squeeze([i[0] for i in data_true]), label='Train loss') +plt.ylabel('loss') +plt.xlabel('epochs') +plt.title("Model: train loss") +plt.legend() +# plt.show() + +plt.savefig('plot123_2.png', dpi=200) # 指定分辨率保存 diff --git a/CTAI_model/utils/transform.py b/CTAI_model/utils/transform.py new file mode 100644 index 0000000..671a9a9 --- /dev/null +++ b/CTAI_model/utils/transform.py @@ -0,0 +1,44 @@ +import os + +import SimpleITK as sitk +import cv2 +import numpy as np + +from data_set import make + + +def mkdir(path): + folder = os.path.exists(path) + if not folder: # 判断是否存在文件夹如果不存在则创建为文件夹 + os.makedirs(path) # makedirs 创建文件时如果路径不存在会创建这个路径 + + +filename_list = make.get_person_files('../data/all/d2/') +for i in filename_list: + pid = i[0] + print(pid) + for j in i[1]: + image = sitk.ReadImage(j) + image_array = sitk.GetArrayFromImage(image).swapaxes(0, 2) + image_array = np.rot90(image_array, -1) + image_array = np.fliplr(image_array).squeeze() + + # ret, image_array = cv2.threshold(image_array, 150, 255, cv2.THRESH_BINARY) + mkdir(f'../data/png/{pid}/') + name = j.replace('.dcm', '').split('/')[-1] + # cv2.imwrite(f'../data/jpg/{pid}/{name}.jpg', image_array, [int(cv2.IMWRITE_JPEG_QUALITY), 100]) + cv2.imwrite(f'../data/png/{pid}/{name}.png', image_array, (cv2.IMWRITE_PNG_COMPRESSION, 0)) + +# print(filename_list) + +# for i in filename_list: +# if '.dcm' in i: +# image = sitk.ReadImage(data_path + '/' + i) +# image_array = sitk.GetArrayFromImage(image).swapaxes(0, 2) +# image_array = np.rot90(image_rray, -1) +# image_array = np.fliplr(image_array) +# name = i.replace('.dcm', '') +# cv2.imwrite(f'{data_path}/{name}_train.png', image_array, (cv2.IMWRITE_PNG_COMPRESSION, 0)) + +# t=cv2.imread('data/out/mask-tttt.png',cv2.IMREAD_GRAYSCALE) +# print(t) diff --git a/CTAI_model/utils/校验文件.py b/CTAI_model/utils/校验文件.py new file mode 100644 index 0000000..6048b91 --- /dev/null +++ b/CTAI_model/utils/校验文件.py @@ -0,0 +1,44 @@ +import os + + +def get_train_files(data_path, file_type='dcm', all=True): + file_type = '.' + file_type + image_list, mask_list, ROI_list = [], [], [] + dir_list = [data_path + i for i in os.listdir(data_path)] + filename_list = [] + for dir in dir_list: + # 所有数据跑 + if all: + temp = os.listdir(dir + '/arterial phase') + filename_list.extend([dir + '/arterial phase/' + name for name in temp]) + if not all: + filename_list.append(dir) + # temp = os.listdir(dir) + # filename_list.extend([dir + '/' + name for name in temp]) + + for i in filename_list: + if file_type in i: + image_list.append(i) + if '_mask' in i: + mask_list.append(i) + + # 校验文件正确 + + return image_list, mask_list + + +if __name__ == '__main__': + a, _ = get_train_files('../data/all/d2/') + _, b = get_train_files('../data/out/') + + for i in range(len(a)): + aa = a[i].split('/')[-1].replace('.dcm', '') + bb = b[i].split('/')[-1].replace('_mask.png', '') + if aa != bb: + print(aa, bb, b[i]) + print(a[i] + 'file list error!') + for i in range(len(b)): + aa = a[i].split('/')[-1].replace('.dcm', '') + bb = b[i].split('/')[-1].replace('_mask.png', '') + if aa != bb: + print(a[i] + 'file list error!')