upload train code

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