main.py 5.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170
  1. import torch
  2. import torchvision
  3. # FOR DATA
  4. from utils.preprocess import prepare_datasets, prepare_predict
  5. from utils.show_image import show_image
  6. from utils.CNN import CNN_Net
  7. from torch.utils.data import DataLoader
  8. from torchvision import datasets
  9. from torch import nn
  10. from torchvision.transforms import ToTensor
  11. # import nonechucks as nc # Used to load data in pytorch even when images are corrupted / unavailable (skips them)
  12. # FOR IMAGE VISUALIZATION
  13. import nibabel as nib
  14. # GENERAL PURPOSE
  15. import os
  16. import pandas as pd
  17. import numpy as np
  18. import matplotlib.pyplot as plt
  19. import glob
  20. print("--- RUNNING ---")
  21. print("Pytorch Version: " + torch. __version__)
  22. # LOADING DATA
  23. # data & training properties:
  24. val_split = 0.2 # % of val and test, rest will be train
  25. seed = 12 # TODO Randomize seed
  26. # params = {
  27. # "target_rows": 91,
  28. # "target_cols": 109,
  29. # "depth": 91,
  30. # "axis": 1,
  31. # "num_clinical": 2,
  32. # "CNN_drop_rate": 0.3,
  33. # "RNN_drop_rate": 0.1,
  34. # # "CNN_w_regularizer": regularizers.l2(2e-2),
  35. # # "RNN_w_regularizer": regularizers.l2(1e-6),
  36. # "CNN_batch_size": 10,
  37. # "RNN_batch_size": 5,
  38. # "val_split": 0.2,
  39. # "final_layer_size": 5
  40. # }
  41. properties = {
  42. "batch_size":6,
  43. "padding":0,
  44. "dilation":1,
  45. "groups":1,
  46. "bias":True,
  47. "padding_mode":"zeros",
  48. "drop_rate":0
  49. }
  50. # Might have to replace datapaths or separate between training and testing
  51. model_filepath = '/data/data_wnx1/rschuurs/Pytorch_CNN-RNN'
  52. CNN_filepath = '/data/data_wnx1/rschuurs/Pytorch_CNN-RNN/cnn_net.pth' # cnn_net.pth
  53. # mri_datapath = '/data/data_wnx1/rschuurs/Pytorch_CNN-RNN/PET_volumes_customtemplate_float32/' # Small Test
  54. mri_datapath = '/data/data_wnx1/_Data/AlzheimersDL/CNN+RNN-2class-1cnn+data/PET_volumes_customtemplate_float32/' # Real data
  55. annotations_datapath = './data/data_wnx1/rschuurs/Pytorch_CNN-RNN/LP_ADNIMERGE.csv'
  56. # annotations_file = pd.read_csv(annotations_datapath) # DataFrame
  57. # show_image(17508)
  58. # TODO: Datasets include multiple labels, such as medical info
  59. training_data, val_data, test_data = prepare_datasets(mri_datapath, val_split, seed)
  60. # Create data loaders
  61. train_dataloader = DataLoader(training_data, batch_size=properties['batch_size'], shuffle=True, drop_last=True)
  62. test_dataloader = DataLoader(test_data, batch_size=properties['batch_size'], shuffle=True)
  63. val_dataloader = DataLoader(val_data, batch_size=properties['batch_size'], shuffle=True)
  64. # for X, y in train_dataloader:
  65. # print(f"Shape of X [Channels (colors), Y, X, Z]: {X.shape}") # X & Y are from TOP LOOKING DOWN
  66. # print(f"Shape of Y (Dataset?): {y.shape} {y.dtype}")
  67. # break
  68. # Display 4 images and labels.
  69. # x = 1
  70. # while x < 1:
  71. # train_features, train_labels = next(iter(train_dataloader))
  72. # print(f"Feature batch shape: {train_features.size()}")
  73. # img = train_features[0].squeeze()
  74. # print(f"Feature batch shape: {img.size()}")
  75. # image = img[:, :, 40]
  76. # print(f"Feature batch shape: {image.size()}")
  77. # label = train_labels[0]
  78. # print(f"Label: {label}")
  79. # plt.imshow(image, cmap="gray")
  80. # plt.show()
  81. # x = x+1
  82. train = False
  83. predict = False
  84. CNN = CNN_Net(train_dataloader, prps=properties, final_layer_size=2)
  85. CNN.cuda()
  86. # RUN CNN
  87. if(train):
  88. CNN.train_model(train_dataloader, test_dataloader, CNN_filepath, epochs=20)
  89. CNN.evaluate_model(val_dataloader, roc=True)
  90. else:
  91. CNN.load_state_dict(torch.load(CNN_filepath))
  92. CNN.evaluate_model(val_dataloader, roc=True)
  93. # PREDICT MODE TO TEST INDIVIDUAL IMAGES
  94. if(predict):
  95. on = True
  96. print("---- Predict mode ----")
  97. print("Integer for image")
  98. print("x or X for exit")
  99. while(on):
  100. inp = input("Next image: ")
  101. if(inp == None or inp.lower() == 'x' or not inp.isdigit()): on = False
  102. else:
  103. dataloader = DataLoader(prepare_predict(mri_datapath, [inp]), batch_size=properties['batch_size'], shuffle=True)
  104. prediction = CNN.predict(dataloader)
  105. features, labels = next(iter(dataloader), )
  106. img = features[0].squeeze()
  107. image = img[:, :, 40]
  108. print(f"Expected class: {labels}")
  109. print(f"Prediction: {prediction}")
  110. plt.imshow(image, cmap="gray")
  111. plt.show()
  112. print("--- END ---")
  113. # EXTRA
  114. # will I need these params?
  115. '''
  116. params_dict = { 'CNN_w_regularizer': CNN_w_regularizer, 'RNN_w_regularizer': RNN_w_regularizer,
  117. 'CNN_batch_size': CNN_batch_size, 'RNN_batch_size': RNN_batch_size,
  118. 'CNN_drop_rate': CNN_drop_rate, 'epochs': 30,
  119. 'gpu': "/gpu:0", 'model_filepath': model_filepath,
  120. 'image_shape': (target_rows, target_cols, depth, axis),
  121. 'num_clinical': num_clinical,
  122. 'final_layer_size': final_layer_size,
  123. 'optimizer': optimizer, 'RNN_drop_rate': RNN_drop_rate,}
  124. params = Parameters(params_dict)
  125. # WHAT WAS THIS AGAIN?
  126. seeds = [np.random.randint(1, 5000) for _ in range(1)]
  127. # READ THIS TO UNDERSTAND TRAIN VS VALIDATION DATA
  128. def evaluate_net (seed):
  129. n_classes = 2
  130. data_loader = DataLoader((target_rows, target_cols, depth, axis), seed = seed)
  131. train_data, val_data, test_data,rnn_HdataT1,rnn_HdataT2,rnn_HdataT3,rnn_AdataT1,rnn_AdataT2,rnn_AdataT3, test_mri_nonorm = data_loader.get_train_val_test(val_split, mri_datapath)
  132. print('Length Val Data[0]: ',len(val_data[0]))
  133. '''