def set_model(self, params): 'Return a list with all classes' if params['model_type'].lower() == 'full': from Models.Full_Connected import Full_Connected self.__model = Full_Connected(params=params) elif params['model_type'].lower() == 'convolutional': from Models.Convolutional import Convolutional self.__model = Convolutional(params=params) elif params['model_type'].lower() == 'vgg16': from Models.VGG16 import VGG16 self.__model = VGG16(params=params) elif params['model_type'].lower() == 'mobile': from Models.Mobile_Net import Mobile_Net self.__model = Mobile_Net(params=params) elif params['model_type'].lower() == 'mobilev2': from Models.Mobile_NetV2 import Mobile_NetV2 self.__model = Mobile_NetV2(params=params) elif params['model_type'].lower() == 'resnet': from Models.ResNet import ResNet self.__model = ResNet(params=params) elif params['model_type'].lower() == 'unet': from Models.UNet import UNet self.__model = UNet(params=params) else: from Models.Custom_Network import Custom_Network self.__model = Custom_Network(params=params)
def evaluate_generator(model_path): ''' Evaluate performance of CT slices using generator :model_path: path of evaluated model ''' model = UNet(input_shape=(512, 512, 1)) model.compile(optimizer=Adam(lr=1e-4), loss=focal_tversky_loss, metrics=[dice_coef]) model.load_weights(model_path) image_path = 'generate_data/test.h5' # h5 file of evaluated patients TOTAL_IMAGE = 1814 # total images in h5 file BATCH_SIZE = 2 reader = HDF5DatasetGenerator(db_path=image_path, batch_size=BATCH_SIZE) generator = reader.generator() results = model.evaluate_generator(generator, steps=TOTAL_IMAGE // BATCH_SIZE, verbose=1) for i, metric in enumerate(model.metrics_names): print(metric, ':', results[i]) reader.close()
def train(): train_reader = HDF5DatasetGenerator(db_path=train_path, batch_size=BATCH_SIZE) train_iter = train_reader.generator() val_reader = HDF5DatasetGenerator(db_path=val_path, batch_size=BATCH_SIZE) val_iter = val_reader.generator() model = UNet(input_shape=(512, 512, 1)) model.compile(optimizer=Adam(lr=1e-4), loss=focal_tversky_loss, metrics=[dice_coef]) if not os.path.exists(save_path): os.mkdir(save_path) os.mkdir(save_path + '/model') os.mkdir(save_path + '/model/logs') model_checkpoint = ModelCheckpoint( save_path + '/model/weights.{epoch:02d}-{val_loss:.2f}.hdf5', monitor='val_loss', verbose=1, save_best_only=True) tensorboard = TensorBoard(log_dir=save_path + '/model/logs') reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=4, mode='auto') callbacks = [model_checkpoint, tensorboard, reduce_lr] model.fit_generator(train_iter, steps_per_epoch=TOTAL_TRAIN // BATCH_SIZE, epochs=20, validation_data=val_iter, validation_steps=TOTAL_VAL // BATCH_SIZE, verbose=1, callbacks=callbacks) train_reader.close() val_reader.close() model.save(save_path + '/model/model.h5') print('Finished training ......')
def predict(model_path, idx): ''' Predict segmentation mask of single CT slice :model_path: path of evaluated model :idx: index of evaluated patient ''' pred_dir = 'pred/' if not os.path.exists(pred_dir): os.mkdir(pred_dir) model = UNet((512, 512, 1)) model.compile(optimizer=Adam(lr=1e-4), loss=focal_tversky_loss, metrics=[dice_coef]) model.load_weights(model_path) # reading original images reader = sitk.ImageSeriesReader() dicom_names = reader.GetGDCMSeriesFileNames(data_path + 'volume/' + str(idx) + '/') reader.SetFileNames(dicom_names) images = reader.Execute() images = sitk.GetArrayFromImage(images) # images.shape: (slices, 512, 512) # HU values are restricted to lung window settings [-1100, 300] # namely: window width: 1400, window level: -400 images[images < -1100] = -1100 images[images > 300] = 300 # normalization max_value, min_value = images.max(), images.min() images = images / (max_value - min_value) pred_inputs = np.expand_dims( images, axis=-1) # pred_inputs.shape: (slices, 512, 512, 1) pred_masks = model.predict( pred_inputs, batch_size=1, verbose=1) # pred_masks.shape: (slices, 512, 512, 1) # values in pred_masks are probabilities in range [0, 1] # converting to binary mask using a threshold of 0.5 pred_masks[pred_masks >= 0.5] = 1 pred_masks[pred_masks < 0.5] = 0 pred_masks = pred_masks.astype(int) out = sitk.GetImageFromArray(pred_masks) sitk.WriteImage(out, pred_dir + 'pred_threshold_' + str(idx) + '.nii.gz')
def evaluate(model_path, idx): ''' Evaluate performance of single CT slice :model_path: path of evaluated model :idx: index of evaluated patient ''' model = UNet(input_shape=(512, 512, 1)) model.compile(optimizer=Adam(lr=1e-4), loss=focal_tversky_loss, metrics=[dice_coef]) model.load_weights(model_path) # reading original images reader = sitk.ImageSeriesReader() dicom_names = reader.GetGDCMSeriesFileNames(data_path + 'volume/' + str(idx) + '/') reader.SetFileNames(dicom_names) images = reader.Execute() images = sitk.GetArrayFromImage(images) # images.shape: (slices, 512, 512) # HU values are restricted to lung window settings [-1100, 300] # namely: window width: 1400, window level: -400 images[images < -1100] = -1100 images[images > 300] = 300 # normalization max_value, min_value = images.max(), images.min() images = images / (max_value - min_value) # reading ground truth labels labels = sitk.ReadImage(data_path + 'segmentation/' + str(idx) + '.nii') labels = sitk.GetArrayFromImage( labels) # eval_labels.shape: (slices, 512, 512) # (slices, 512, 512) => (slices, 512, 512, 1) images = np.expand_dims(images, axis=-1) labels = np.expand_dims(labels, axis=-1) results = model.evaluate(images, labels, batch_size=1) for i, metric in enumerate(model.metrics_names): print(metric, ':', results[i])
train_loader = DataLoader(train_loader, args.batch, True, num_workers=4) test_loader = DataLoader(test_loader, args.batch, False, num_workers=4) classes = getClasses(args.classes) weights = None loaded = False if os.path.isfile(args.weights): weights = np.load(args.weights) weights = weights / np.sum(weights) weights = 1 - weights print(weights) if os.path.isfile(args.model): model = torch.load(args.model).to(args.device) print("Model loaded!") loaded = True else: model = UNet(4, len(classes), args.depth, args.filters).to(args.device) model.initialize() print("Model initialized!") if weights is None: criterion = torch.nn.CrossEntropyLoss().to(args.device) else: weights = torch.FloatTensor(weights).to(args.device) criterion = torch.nn.CrossEntropyLoss(weight=weights).to(args.device) optimizer = torch.optim.Adam(model.parameters(), args.lr) train_matrix = np.zeros((len(classes), len(classes)), dtype=np.int32) def onTrainBatch(batch_id, features, labels, output, loss): global train_matrix output = torch.argmax(output, dim=1) mat = confusion_matrix(labels, output, len(classes))
def get_experiment_objs(exp_dir, create_model=True): settings = load_pickle_obj(exp_dir / 'settings.pkl') params = dict(num_categories=settings['num_categories'], pretrained_on=settings['pretrained_on'], num_epochs=settings['num_epochs'], device=settings['device']) model_name = settings['model_name'] model = None if create_model: if settings['continue_training']: assert settings['initialize_weights'] is False and settings[ 'continue_model_path'] is not None model = torch.load(settings['continue_model_path']) else: # base model: if model_name == 'unet': model = UNet(in_channels=settings['in_channels'], num_categories=settings['num_categories'], filter_sizes=settings['filter_sizes'], deep_supervision=settings['deep_supervision']) # attention models: elif model_name == 'attention_unet': model = AttentionUNet( in_channels=settings['in_channels'], num_categories=settings['num_categories'], filter_sizes=settings['filter_sizes'], deep_supervision=settings['deep_supervision']) elif model_name == 'cbam_unet': model = CBAM_UNet( in_channels=settings['in_channels'], num_categories=settings['num_categories'], filter_sizes=settings['filter_sizes'], deep_supervision=settings['deep_supervision']) elif model_name == 'dualattention_unet': model = DualAttentionUNet( in_channels=settings['in_channels'], num_categories=settings['num_categories'], filter_sizes=settings['filter_sizes'], deep_supervision=settings['deep_supervision']) elif model_name == 'residualattention_unet': model = ResidualAttentionUNet( in_channels=settings['in_channels'], num_categories=settings['num_categories'], filter_sizes=settings['filter_sizes'], deep_supervision=settings['deep_supervision']) elif model_name == 'scag_unet': model = scAG_UNet( in_channels=settings['in_channels'], num_categories=settings['num_categories'], filter_sizes=settings['filter_sizes'], deep_supervision=settings['deep_supervision']) # transfer learning models: elif model_name == 'densenet121_unet': model = DenseNet121_UNet( in_channels=settings['in_channels'], num_categories=settings['num_categories'], deep_supervision=settings['deep_supervision'], pretrained=True) elif model_name == 'mobilenetv2_unet': model = MobileNetV2_UNet( in_channels=settings['in_channels'], num_categories=settings['num_categories'], deep_supervision=settings['deep_supervision'], pretrained=True) elif model_name == 'resnet34_unet': model = ResNet34_UNet( in_channels=settings['in_channels'], num_categories=settings['num_categories'], deep_supervision=settings['deep_supervision'], pretrained=True) elif model_name == 'vgg11_unet': model = vgg11_UNet( in_channels=settings['in_channels'], num_categories=settings['num_categories'], deep_supervision=settings['deep_supervision'], pretrained=True) elif model_name == 'vgg16_unet': model = vgg16_UNet( in_channels=settings['in_channels'], num_categories=settings['num_categories'], deep_supervision=settings['deep_supervision'], pretrained=True) assert model is not None dataloader_path = settings['dataloader_path'] learning_rate_dict = dict( static_lr=settings['learning_rate_dict']['static_lr'], use_cyclic_learning_rate=settings['learning_rate_dict'] ['use_cyclic_learning_rate'], base_lr=settings['learning_rate_dict']['base_lr'], max_lr=settings['learning_rate_dict']['max_lr']) return params, model, dataloader_path, learning_rate_dict, settings
) if command == '--Help': print("Check the documentation.") elif command == "--Train" or command == "--Predict": try: architecture = sys.argv[2] except IndexError: print( "An Index error has occured! Check the documentation to make sure your passing all the required arguments." ) if command == "--Train": if architecture == '--UNet': UNet = UNet() UNet.train() elif architecture == '--UNet++': UNetPP = UNetPP() UNetPP.train() else: raise Exception( "You have passed an invalid argument.\nCheck the documentation for the allowed arguments." ) elif command == "--Predict": if architecture == '--UNet': UNet = UNet() UNet.predict() elif architecture == '--UNet++': UNetPP = UNetPP()
classes = getClasses(args.classes) weights = None loaded = False if os.path.isfile(args.weights): weights = np.load(args.weights) weights = weights / np.sum(weights) weights = 1 - weights print(weights) if os.path.isfile(args.model): model = torch.load(args.model).to(args.device) print("Model loaded!") loaded = True else: model = UNet(4, len(classes), args.depth, args.filters, convPerLayer=args.convperlayer, dropout=args.dropout).to(args.device) model.initialize() print("Model initialized!") if weights is None: criterion = torch.nn.CrossEntropyLoss().to(args.device) else: weights = torch.FloatTensor(weights).to(args.device) criterion = torch.nn.CrossEntropyLoss(weight=weights).to(args.device) optimizer = torch.optim.Adam(model.parameters(), args.lr) def print_matrix(matrix, classes): p = precision(matrix) r = recall(matrix) f = f1(p, r)
import matplotlib.pyplot as plt # Params: min_lr = 1e-9 max_lr = 1.0 window_size_smoothing = 100 num_epochs = 3 momentum = .95 params = dict( num_categories=2, pretrained_on='Imagenet', num_epochs=100, device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')) dataloader_path = Path('tmp/Dataloader_SN1_Buildings.pkl') model = UNet(in_channels=3, num_categories=params['num_categories'], filter_sizes=(32, 64, 128, 256, 512), deep_supervision=True) running_experiment = False initialize_weights = True freeze_encoder = False continue_training = False if len(sys.argv) > 1: exp_dir = Path(sys.argv[1]) running_experiment = True params, model, dataloader_path, learning_rate_dict, experiment_settings = get_experiment_objs( exp_dir) initialize_weights = experiment_settings['initialize_weights'] freeze_encoder = experiment_settings['freeze_encoder'] continue_training = experiment_settings['continue_training'] model.to(params['device']) if initialize_weights: