upload train code

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!')