Add files via upload

Author: Kent0n-Li

README.md

@@ -11,10 +11,6 @@
 The current code is designed for Windows, not for Linux. 
 
 
-If you only want to use nnSAM, please install [this](https://github.com/Kent0n-Li/nnSAM). <br>
-如果你只想运行nnSAM,请访问该代码仓库[this](https://github.com/Kent0n-Li/nnSAM)
-
-
 Install (安装步骤):
 
 ```bash
@@ -35,7 +31,7 @@ pip install timm
 pip install git+https://github.com/Kent0n-Li/nnSAM.git
 
 git clone https://github.com/Kent0n-Li/Medical-Image-Segmentation.git
-cd Medical-Image-Segmentation
+cd Medical-Image-Segmentation-Benchmark
 pip install -r requirements.txt
 ```
 
@@ -44,6 +40,8 @@ pip install -r requirements.txt
 python web.py
 ```
 
+If you only want to use nnSAM, please install [this](https://github.com/Kent0n-Li/nnSAM). <br>
+如果你只想运行nnSAM,请访问该代码仓库[this](https://github.com/Kent0n-Li/nnSAM)
 
 样例数据集：[Demo Dataset](https://github.com/Kent0n-Li/Medical-Image-Segmentation/tree/main/Demo_dataset)
 

config.py

@@ -79,3 +79,8 @@ def parse_args(type=None):
                             default=4, help='Batch Size')
 
     return parser.parse_args()
+
+
+def print_web(text):
+    with open(output_file, 'a') as f:
+        f.write(text + '\n')

data_process.py

@@ -0,0 +1,46 @@
+import os
+from config import *
+from web import *
+
+
+if __name__ == '__main__':
+    # Define target paths
+    target_paths = {
+        'training_image': os.path.join(os.environ['medseg_raw'], os.environ['current_dataset'], 'imagesTr'),
+        'training_label': os.path.join(os.environ['medseg_raw'], os.environ['current_dataset'], 'labelsTr'),
+        'testing_image': os.path.join(os.environ['medseg_raw'], os.environ['current_dataset'], 'imagesTs'),
+        'testing_label': os.path.join(os.environ['medseg_raw'], os.environ['current_dataset'], 'labelsTs'),
+        'validation_image': os.path.join(os.environ['medseg_raw'], os.environ['current_dataset'], 'imagesVal'),
+        'validation_label': os.path.join(os.environ['medseg_raw'], os.environ['current_dataset'], 'labelsVal'),
+    }
+    try:
+
+        #jpg to png
+        for path_each in target_paths.values():
+            convert_jpg_to_png_all_from_path(path_each)
+            
+
+        unique_values = find_unique_labels(target_paths['training_label'])
+
+        convert_label_by_searchsorted(target_paths['training_label'], unique_values)
+        convert_label_by_searchsorted(target_paths['validation_label'], unique_values)
+        convert_label_by_searchsorted(target_paths['testing_label'], unique_values)
+
+        # Generate dataset.json
+        print_web(f"Generating dataset.json for {os.environ['current_dataset']}")
+        npimg_path = os.path.join(target_paths['training_image'], os.listdir(target_paths['training_image'])[0])
+        npimg = cv2.imread(npimg_path, cv2.IMREAD_UNCHANGED)
+
+        img_channel = 3 if len(npimg.shape) == 3 else 1
+        label_class_num = len(unique_values)
+
+        dataset_id = len(os.listdir(os.environ['medseg_raw'])) + 1
+        dataset_id = "{:03}".format(dataset_id)
+
+        image_size = npimg.shape[0]
+
+        save_dataset_json(dataset_id, os.environ['current_dataset'], image_size, img_channel, label_class_num)
+
+    except Exception as e:
+        print_web(f"Error: {e}")
+        raise e

jpg_to_png.py

@@ -0,0 +1,36 @@
+import os
+from PIL import Image
+import numpy as np
+
+folder_path_list = ['E:\\Demo_dataset\\Kvasir-SEG\\train\\image','E:\\Demo_dataset\\Kvasir-SEG\\test\\image']
+for folder_path in folder_path_list:
+    for file in os.listdir(folder_path):
+        if file.endswith('.jpg'):
+            img = Image.open(os.path.join(folder_path, file))
+            folder_path_new = folder_path+'_png'
+            os.makedirs(folder_path_new, exist_ok=True)
+            img.save(os.path.join(folder_path_new, file.replace('.jpg', '.png')))
+            print(f'{file} converted to png')
+
+
+
+folder_path_list = ['E:\\Demo_dataset\\Kvasir-SEG\\train\\label','E:\\Demo_dataset\\Kvasir-SEG\\test\\label']
+
+for folder_path in folder_path_list:
+    for file in os.listdir(folder_path):
+        if file.endswith('.jpg'):
+            img = Image.open(os.path.join(folder_path, file))
+            # threshold the image
+            img = np.array(img)
+            img[img > 100] = 255
+            img[img <= 100] = 0
+            img = Image.fromarray(img)
+
+            img = img.convert('L')
+            folder_path_new = folder_path+'_png'
+            os.makedirs(folder_path_new, exist_ok=True)
+
+            img.save(os.path.join(folder_path_new, file.replace('.jpg', '.png')))
+            print(f'{file} converted to png')
+
+            

test.py

@@ -24,39 +24,18 @@
 import torch.optim as optim
 from networks.swin_config import get_swin_config
 import requests
-from  config import parse_args
 import gdown
 import matplotlib.pyplot as plt
-
-# parser = argparse.ArgumentParser()
-#
-# parser.add_argument('--max_iterations', type=int,
-#                     default=30000, help='maximum epoch number to train')
-# parser.add_argument('--max_epochs', type=int,
-#                     default=200, help='maximum epoch number to train')
-# parser.add_argument('--n_gpu', type=int, default=1, help='total gpu')
-# parser.add_argument('--deterministic', type=int, default=1,
-#                     help='whether use deterministic training')
-# parser.add_argument('--base_lr', type=float, default=0.01,
-#                     help='segmentation network learning rate')
-# parser.add_argument('--img_size', type=int,
-#                     default=224, help='input patch size of network input')
-# parser.add_argument('--seed', type=int,
-#                     default=1234, help='random seed')
-# parser.add_argument('--zip', action='store_true', help='use zipped dataset instead of folder dataset')
-# parser.add_argument('--cache-mode', type=str, default='part', choices=['no', 'full', 'part'],
-#                     help='no: no cache, '
-#                          'full: cache all data, '
-#                          'part: sharding the dataset into nonoverlapping pieces and only cache one piece')
-# parser.add_argument('--resume', help='resume from checkpoint')
-# parser.add_argument('--accumulation-steps', type=int, help="gradient accumulation steps")
-# parser.add_argument('--use-checkpoint', action='store_true',
-#                     help="whether to use gradient checkpointing to save memory")
-# parser.add_argument('--amp-opt-level', type=str, default='O1', choices=['O0', 'O1', 'O2'],
-#                     help='mixed precision opt level, if O0, no amp is used')
-# parser.add_argument('--tag', help='tag of experiment')
-# parser.add_argument('--eval', action='store_true', help='Perform evaluation only')
-# parser.add_argument('--throughput', action='store_true', help='Test throughput only')
+from logHelper import setup_logger
+from config import output_file, parse_args
+from networks.YourNet import Your_Net
+from networks.GT_UNet import GT_U_Net
+from networks.model.BiSeNet import BiSeNet
+from networks.model.DDRNet import DDRNet
+from networks.model.DeeplabV3Plus import Deeplabv3plus_res50
+from networks.model.FCN_ResNet import FCN_ResNet
+from networks.model.HRNet import HighResolutionNet
+from networks.SegNet import SegNet
 
 args = parse_args()
 
@@ -100,10 +79,7 @@ def download_model(url, destination):
 
 
 class DynamicDataset(data.Dataset):
-    def __init__(self, img_path, gt_path, data_end_json, size=None):
-
-        with open(data_end_json) as f:
-            self.file_end = json.load(f)['file_ending']
+    def __init__(self, img_path, gt_path, size=None):
 
         self.img_name = os.listdir(img_path)
         self.size = size
@@ -114,11 +90,13 @@ def __getitem__(self, item):
         imagename = self.img_name[item]
         img_path = os.path.join(self.img_path, imagename)
 
-        if self.file_end in ['.png', '.bmp', '.tif']:
+        file_end = imagename.split('.')[-1]
+
+        if file_end in ['png']:
             npimg = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)
             npimg = np.array(npimg)
 
-        elif self.file_end in ['.gz', '.nrrd', '.mha', '.nii.gz', '.nii']:
+        elif file_end in ['gz', 'nrrd', 'mha', 'nii.gz', 'nii']:
             npimg = sitk.ReadImage(img_path)
             npimg = sitk.GetArrayFromImage(npimg)
 
@@ -138,14 +116,14 @@ def __getitem__(self, item):
                                            antialias=None)
             npimg = adapt_size(npimg)
 
-        return npimg, imagename.replace('_0000', ''), ori_shape
+        return npimg, imagename, ori_shape
 
     def __len__(self):
         size = int(len(self.img_name))
         return size
 
 
-# if __name__ == "__main__":
+
 
 def test_model():
     cudnn.benchmark = False
@@ -155,37 +133,31 @@ def test_model():
     torch.manual_seed(args.seed)
     torch.cuda.manual_seed(args.seed)
 
+    data_json_file = os.path.join(os.environ['medseg_raw'], os.environ['current_dataset'], 'dataset.json')
+
+    with open(data_json_file) as f:
+        json_data = json.load(f)
+        num_classes = json_data['label_class_num']
+        in_channels = json_data['img_channel']
+        args.img_size = json_data['imgae_size']
+
+
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    if device.type == 'cuda':
-        total_memory = torch.cuda.get_device_properties(device).total_memory / (1024 ** 3)  # bytes to GB
-        args.batch_size = int(total_memory / 10) * 2
-    else:
-        args.batch_size = 2
 
-    fold = os.environ['current_fold']
 
-    data_json_file = os.path.join(os.environ['nnUNet_raw'], os.environ['current_dataset'], 'dataset.json')
-    # split_json_path = os.path.join(os.environ['nnUNet_preprocessed'], os.environ['current_dataset'], 'splits_final.json')
-    # base_json_path = os.path.join(os.environ['nnUNet_preprocessed'], os.environ['current_dataset'])
-    output_folder_test = os.path.join(os.environ['nnUNet_results'], os.environ['MODEL_NAME'],
-                                      os.environ['current_dataset'], 'nnUNetTrainer__nnUNetPlans__2d', 'test_pred')
-    output_folder_5fold = os.path.join(os.environ['nnUNet_results'], os.environ['MODEL_NAME'],
-                                       os.environ['current_dataset'], 'nnUNetTrainer__nnUNetPlans__2d', f'fold_{fold}')
+    args.batch_size = 1
+
+
+    output_folder_test = os.path.join(os.environ['medseg_results'], os.environ['current_dataset'], os.environ['MODEL_NAME'], 'test_pred')
 
     os.makedirs(output_folder_test, exist_ok=True)
-    os.makedirs(output_folder_5fold, exist_ok=True)
+    
 
-    imageTr_path = os.path.join(os.environ['nnUNet_raw'], os.environ['current_dataset'], 'imagesTr')
-    labelTr_path = os.path.join(os.environ['nnUNet_raw'], os.environ['current_dataset'], 'labelsTr')
-    imageTs_path = os.path.join(os.environ['nnUNet_raw'], os.environ['current_dataset'], 'imagesTs')
-    labelTs_path = os.path.join(os.environ['nnUNet_raw'], os.environ['current_dataset'], 'labelsTs')
 
-    with open(data_json_file) as f:
-        json_data = json.load(f)
-        num_classes = len(json_data['labels'])
-        in_channels = len(json_data['channel_names'])
+    imageTs_path = os.path.join(os.environ['medseg_raw'], os.environ['current_dataset'], 'imagesTs')
+    labelTs_path = os.path.join(os.environ['medseg_raw'], os.environ['current_dataset'], 'labelsTs')
+    weights_path = os.path.join(os.environ['medseg_results'], os.environ['current_dataset'], os.environ['MODEL_NAME'], 'checkpoint_final.pth') 
 
-    weights_path = os.path.join(output_folder_5fold, 'checkpoint_final.pth')
 
     model_name = os.environ['MODEL_NAME']
     if model_name == 'unet':
@@ -211,6 +183,7 @@ def test_model():
     elif model_name == 'swinunet':
         args.cfg = './networks/swin_tiny_patch4_window7_224_lite.yaml'
         args.opts = None
+        args.img_size = 224
         swin_config = get_swin_config(args)
         model = SwinUnet(swin_config, img_size=224, num_classes=num_classes).cuda()
         url = "https://drive.google.com/uc?id=1TyMf0_uvaxyacMmVzRfqvLLAWSOE2bJR"
@@ -226,37 +199,52 @@ def test_model():
     elif model_name == 'r2unet':
         model = R2U_Net(in_ch=in_channels, out_ch=num_classes).cuda()
 
+    elif model_name == 'gtunet':
+        model = GT_U_Net(in_ch=in_channels, out_ch=num_classes).to(device)
+        args.img_size = 256
+
+    elif model_name == 'bisenet':
+        model = BiSeNet(in_ch=in_channels, out_ch=num_classes).to(device)
+    
+    elif model_name == 'ddrnet':
+        model = DDRNet(in_ch=in_channels, out_ch=num_classes).to(device)
+    
+    elif model_name == 'deeplabv3plus':
+        model = Deeplabv3plus_res50(in_ch=in_channels, out_ch=num_classes).to(device)
+    
+    elif model_name == 'hrnet':
+        model = HighResolutionNet(in_ch=in_channels, out_ch=num_classes).to(device)
+
+    elif model_name == 'segnet':
+        model = SegNet(in_ch=in_channels, out_ch=num_classes).to(device)
+    
+    elif model_name == 'fcnresnet':
+        model = FCN_ResNet(in_ch=in_channels, out_ch=num_classes).to(device)
+    
+    elif model_name == 'yournet':
+        model = Your_Net(in_ch=in_channels, out_ch=num_classes).to(device)
+
+
     else:
         raise NotImplementedError(f"model_name {model_name} not supported")
 
     model.load_state_dict(torch.load(weights_path))
 
-    logging.basicConfig(filename="logging.txt", level=logging.INFO,
-                        format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S')
-    logging.getLogger().addHandler(logging.StreamHandler(sys.stdout))
-    logging.info(str(args))
+    logger = setup_logger("training_logger", output_file=output_file)
+    logger.info("Process started")
+    logger.info(str(args))
     base_lr = args.base_lr
-    batch_size = args.batch_size * args.n_gpu
-    if model_name == 'swinunet' or model_name == 'transunet':
-        db_test = DynamicDataset(img_path=imageTs_path, gt_path=labelTs_path, data_end_json=data_json_file,
-                                 size=args.img_size)
-    else:
-        db_test = DynamicDataset(img_path=imageTs_path, gt_path=labelTs_path, data_end_json=data_json_file)
+    
+    db_test = DynamicDataset(img_path=imageTs_path, gt_path=labelTs_path, size=args.img_size)
 
-    with open(data_json_file) as f:
-        file_end = json.load(f)['file_ending']
+        
 
-    testloader = DataLoader(db_test, batch_size=1, shuffle=True, num_workers=2, pin_memory=True)
+    testloader = DataLoader(db_test, batch_size=args.batch_size, shuffle=True, num_workers=2, pin_memory=True)
     if args.n_gpu > 1:
         model = nn.DataParallel(model)
-    model.train()
-    ce_loss = CrossEntropyLoss()
-    dice_loss = DiceLoss(num_classes)
-    optimizer = optim.SGD(model.parameters(), lr=base_lr, momentum=0.9, weight_decay=0.0001)
-
-    best_performance = 0.0
-    val_dice_scores = []
-    epoch_numbers = []
+    model.eval()
+
+
     for i_batch, (img, img_name, ori_shape) in enumerate(testloader):
         image_batch = img
         image_batch = image_batch.cuda()
@@ -268,11 +256,16 @@ def test_model():
                                                   align_corners=True)
         pred = outputs.data.max(1)[1].squeeze_(1).squeeze_(0).cpu().numpy()
         pred = pred.astype(np.uint8)
+        print(f"Processing {img_name[0]}")
 
-        if file_end in ['.png', '.bmp', '.tif']:
+        file_end = img_name[0].split('.')[-1]
+        if file_end in ['png', 'bmp', 'tif']:
             pred_img = Image.fromarray(pred)
             pred_img.save(os.path.join(output_folder_test, img_name[0]))
 
-        elif file_end in ['.gz', '.nrrd', '.mha', '.nii.gz', '.nii']:
+        elif file_end in ['gz', 'nrrd', 'mha', 'nii.gz', 'nii']:
             pred_img = sitk.GetImageFromArray(pred)
-            sitk.WriteImage(pred_img, os.path.join(output_folder_test, img_name[0]))
+            sitk.WriteImage(pred_img, os.path.join(output_folder_test, img_name[0]))
+
+if __name__ == "__main__":
+    test_model()