def create_transform(self, args, is_train): """ Convert numpy array into Tensor if dataset is for validation. Apply data augmentation method to train dataset while cv or test if args.use_aug is 1. is_train: boolean flg that dataset is for validation in cv or test return: Compose of albumentations """ if is_train and args.use_aug == 1: transform = A.Compose([ trans.Resize(299, 299), trans.Normalize(mean=self.img_mean, std=self.img_std, max_pixel_value=1.0), ToTensorV2() ]) else: transform = A.Compose([ trans.Resize(299, 299), trans.Normalize(mean=self.img_mean, std=self.img_std, max_pixel_value=1.0), ToTensorV2() ]) return transform
def get_transformv2(opt): transform_list = [] # Transforms in opt.preprocess if 'fixsize' in opt.preprocess: transform_list.append(tr.Resize(286, 286, interpolation=2, p=1)) if 'resize' in opt.preprocess: transform_list.append( tr.Resize(opt.load_size, opt.load_size, interpolation=2, p=1)) if 'crop' in opt.preprocess: transform_list.append(tr.RandomCrop(opt.crop_size, opt.crop_size, p=1)) # Transforms in colorspace if 'color' in opt.preprocess: transform_list.extend([ tr.RandomContrast(limit=0.2, p=0.5), tr.RandomBrightness(limit=0.2, p=0.5), tr.HueSaturationValue(hue_shift_limit=20, sat_shift_limit=30, val_shift_limit=20, p=0.5), # tr.ISONoise() ]) # Necessary transforms transform_list.extend([ tr.HorizontalFlip(p=0.5), tr.VerticalFlip(p=0.5), tr.Normalize(p=1.0), ToTensorV2(p=1) ]) return Compose(transform_list, additional_targets={'imageB': 'image'})
def augmentation(mode, target_size, prob=0.5, aug_m=2): ''' description: augmentation mode: 'train' 'test' target_size: int or list, the shape of image , aug_m: Strength of transform ''' high_p = prob low_p = high_p / 2.0 M = aug_m first_size = [int(x / 0.7) for x in target_size] if mode == 'train': return composition.Compose([ transforms.Resize(first_size[0], first_size[1], interpolation=3), transforms.Flip(p=0.5), composition.OneOf([ RandomCenterCut(scale=0.1 * M), transforms.ShiftScaleRotate(shift_limit=0.05 * M, scale_limit=0.1 * M, rotate_limit=180, border_mode=cv2.BORDER_CONSTANT, value=0), albumentations.imgaug.transforms.IAAAffine( shear=(-10 * M, 10 * M), mode='constant') ], p=high_p), transforms.RandomBrightnessContrast( brightness_limit=0.1 * M, contrast_limit=0.03 * M, p=high_p), transforms.HueSaturationValue(hue_shift_limit=5 * M, sat_shift_limit=15 * M, val_shift_limit=10 * M, p=high_p), transforms.OpticalDistortion(distort_limit=0.03 * M, shift_limit=0, border_mode=cv2.BORDER_CONSTANT, value=0, p=low_p), composition.OneOf([ transforms.Blur(blur_limit=7), albumentations.imgaug.transforms.IAASharpen(), transforms.GaussNoise(var_limit=(2.0, 10.0), mean=0), transforms.ISONoise() ], p=low_p), transforms.Resize(target_size[0], target_size[1], interpolation=3), transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), max_pixel_value=255.0) ], p=1) else: return composition.Compose([ transforms.Resize(target_size[0], target_size[1], interpolation=3), transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), max_pixel_value=255.0) ], p=1)
def aug(self, image): ''' center crop, resize, affine, hue, saturation, ''' if self.phase == 'train': imgaug = Compose([ transforms.Resize(height=self.input_shape[0], width=self.input_shape[1]), RandomRotate90(), Flip(), Transpose(), OneOf([ IAAAdditiveGaussianNoise(), GaussNoise(), ], p=0.2), OneOf([ MotionBlur(p=.2), MedianBlur(blur_limit=3, p=.1), Blur(blur_limit=3, p=.1), ], p=0.2), ShiftScaleRotate(shift_limit=0.0625, scale_limit=0.2, rotate_limit=45, p=.2), ToTensorV2(), ], p=1) else: imgaug = Compose([transforms.Resize(height=self.input_shape[0], width=self.input_shape[1]), ToTensorV2()]) image = imgaug(image=image)['image'] return image
def transform(self, image, mask): H, W = image.shape[0], image.shape[1] new_h, new_w = self.input_shape[0], self.input_shape[1] if H > self.input_shape[0] and W > self.input_shape[1]: y_min, x_min = random.randint(0, H - self.input_shape[0]), random.randint( 0, W - self.input_shape[1]) compose = Compose([ transforms.Crop(x_min=x_min, y_min=y_min, x_max=x_min + self.input_shape[1], y_max=y_min + self.input_shape[0]), ToTensorV2() ]) new_image = compose(image=image)['image'] new_mask = compose(image=mask)['image'] else: resize_ratio = int( math.ceil(max([new_h / float(H), new_w / float(W)]))) h, w = H * resize_ratio, W * resize_ratio y_min, x_min = random.randint(0, h - self.input_shape[0]), random.randint( 0, w - self.input_shape[1]) compose = Compose([ transforms.Resize(h, w), transforms.Crop(x_min=x_min, y_min=y_min, x_max=x_min + self.input_shape[1], y_max=y_min + self.input_shape[0]), ToTensorV2() ]) new_image = compose(image=image)['image'] new_mask = compose(image=mask)['image'] assert new_image.shape[0] == 3 and new_mask.shape[0] == 2 return new_image, new_mask
def ben_valid_augmentation(): return ACompose([ atransforms.Resize(128, 128, interpolation=3), atransforms.Normalize(mean=BEN_BAND_STATS['mean'], std=BEN_BAND_STATS['std']), AToTensor(), ])
def __init__(self): self.transform = ACompose( [atransforms.Resize(128, 128, interpolation=3)]) self.post_transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean=BAND_STATS['mean'], std=BAND_STATS['std']) ])
def __init__(self, config): train_transform = Compose([ transforms.RandomRotate90(), transforms.Flip(), transforms.Resize(config['input_h'], config['input_w']), ]) val_transform = Compose([ transforms.Resize(config['input_h'], config['input_w']), ]) super().__init__() self.img_ext = config['img_ext'] self.mask_ext = config['mask_ext'] self.num_classes = config['num_classes'] self.num_workers = config['num_workers'] self.data_dir = os.path.join('data', config['dataset'], config['sub_dataset']) self.batch_size = config['batch_size'] self.pin_memory = config['pin_memory'] self.train_transform = train_transform self.val_transform = val_transform self.test_transform = None self.train_ids, self.val_ids, self.test_ids = None, None, None self.num_inputs = config['num_inputs']
def __call__(self, image: np.array): ''' :param image: H*W*3 :return: 299*299*3 ''' assert image.shape[2] == 3 H = image.shape[0] W = image.shape[1] scale = max([math.ceil(299 / H), math.ceil(299 / W)]) New_H, New_W = H * scale, W * scale transform = Compose([transforms.Resize(height=New_H, width=New_W), transforms.RandomCrop(height=299, width=299), ToTensorV2()], p=1.0) output = transform(image=image)['image'] return output
def get_transform(opt, params=None, grayscale=False, convert=True): transform_list = [] if grayscale: transform_list.append(transforms.ToGray()) if opt.preprocess == 'resize_and_crop': if params is None: transform_list.append(RandomResizedCrop(256, 256, scale=(0.5, 1.0))) elif opt.preprocess == 'resize': transform_list.append(transforms.Resize(256, 256)) if not opt.no_flip: if params is None: transform_list.append(HorizontalFlip()) return HCompose(transform_list)
def __init__(self, data_folder, cli_args): self.cli_args = cli_args self.root: str = data_folder self.image_names: List[str] = sorted(os.listdir(os.path.join(self.root, "test", "imgs"))) self.transform = Compose( [ # Normalize images to [0..1] tf.Normalize(mean=(0.0, 0.0, 0.0), std=(1.0, 1.0, 1.0), p=1), # Resize images to (image_size, image_size) tf.Resize(cli_args.image_size, cli_args.image_size), # Convert PIL images to torch.Tensor ToTensorV2(), ] )
def get_patched_input(img_path, config, gt_mask_flag): img_patch_set = [] p_size = config['patch_size'] img_size = config['input_w'] patch_overlap = config['patch_overlap'] if gt_mask_flag == True: label_path = img_path.replace('image', 'labels') img_input = cv2.imread(img_path) if gt_mask_flag == True: mask_input = cv2.imread(label_path) if gt_mask_flag == True: image_patch, mask_patch = patch_gen(img_input, mask_input, p_size, patch_overlap) else: image_patch, mask_patch = patch_gen(img_input, img_input, p_size, patch_overlap) val_transform = Compose([ transforms.Resize(config['input_h'], config['input_w']), transforms.Normalize(), ]) patch_len = len(image_patch) for idx in range(patch_len): img = image_patch[idx] img = cv2.resize(img, (img_size, img_size)) mask = img if val_transform is not None: augmented = val_transform(image=img, mask=mask) img = augmented['image'] img = img.astype('float32') / 255 img = img.transpose(2, 0, 1) img_patch_set.append(img) img_patch_set = np.array(img_patch_set) mask_patch_set = np.array(mask_patch) return img_input, img_patch_set, mask_patch_set
def main(): args = parse_args() # load model if args.arch == 'resnet18': model = resnet18(pretrained=True, progress=True) t_d = 448 d = 100 elif args.arch == 'wide_resnet50_2': model = wide_resnet50_2(pretrained=True, progress=True) t_d = 1792 d = 550 model.to(device) model.eval() random.seed(1024) torch.manual_seed(1024) """ os.removedirs(args.save_path+"/pictures_wide_resnet50_2") os.removedirs(args.save_path+"/pictures_temp_wide_resnet50_2") os.remove(args.save_path+"/pictures_temp_wide_resnet50_2/*")""" if use_cuda: torch.cuda.manual_seed_all(1024) idx = torch.tensor(sample(range(0, t_d), d)) # set model's intermediate outputs outputs = [] def hook(module, input, output): outputs.append(output) model.layer1[-1].register_forward_hook(hook) model.layer2[-1].register_forward_hook(hook) model.layer3[-1].register_forward_hook(hook) os.makedirs(os.path.join(args.save_path, 'temp_%s' % args.arch), exist_ok=True) fig, ax = plt.subplots(1, 2, figsize=(20, 10)) fig_img_rocauc = ax[0] fig_pixel_rocauc = ax[1] total_roc_auc = [] total_pixel_roc_auc = [] # load dataset train_transform = Compose([ transforms.Resize(224, 224), ]) val_transform = Compose([ transforms.Resize(224, 224), ]) file_path = 'assets_new_new/data/2021-03-05' train_file_name = 'train.json' val_file_name = "test.json" product_class = [args.product_class] for class_name in product_class: train_dataset = JsonDataset( json_file=file_path + "/" + class_name + "/" + train_file_name, image_data_root=default_config.Image_Data_Root, transform=train_transform, ) test_dataset = JsonDataset( json_file=file_path + "/" + class_name + "/" + val_file_name, image_data_root=default_config.Image_Data_Root, transform=val_transform, ) # use real data as validation if args.train_num_samples == 0: train_sampler = None else: train_sampler = RandomSampler(train_dataset, num_samples=args.train_num_samples, replacement=True) if args.test_num_samples == 0: test_sampler = None else: test_sampler = RandomSampler(test_dataset, num_samples=args.test_num_samples, replacement=True) train_dataloader = DataLoader( train_dataset, batch_size=default_config.Batch_Size, #shuffle=True, #pin_memory=True, sampler=train_sampler, drop_last=True) test_dataloader = DataLoader( test_dataset, batch_size=default_config.Batch_Size_Test, #shuffle=False, #pin_memory=True, sampler=test_sampler, collate_fn=collate_fn, drop_last=False) train_outputs = OrderedDict([('layer1', []), ('layer2', []), ('layer3', [])]) test_outputs = OrderedDict([('layer1', []), ('layer2', []), ('layer3', [])]) # extract train set features train_feature_filepath = os.path.join(args.save_path, 'temp_%s' % args.arch, 'train_%s.pkl' % class_name) if not os.path.exists(train_feature_filepath): for x in tqdm(train_dataloader, '| feature extraction | train | %s |' % class_name): # model prediction with torch.no_grad(): _ = model(x.to(device)) # get intermediate layer outputs for k, v in zip(train_outputs.keys(), outputs): train_outputs[k].append(v.cpu().detach()) # initialize hook outputs outputs = [] for k, v in train_outputs.items(): train_outputs[k] = torch.cat(v, 0) # Embedding concat embedding_vectors = train_outputs['layer1'] for layer_name in ['layer2', 'layer3']: embedding_vectors = embedding_concat(embedding_vectors, train_outputs[layer_name]) # randomly select d dimension embedding_vectors = torch.index_select(embedding_vectors, 1, idx) # calculate multivariate Gaussian distribution B, C, H, W = embedding_vectors.size() embedding_vectors = embedding_vectors.view(B, C, H * W) mean = torch.mean(embedding_vectors, dim=0).numpy() cov = torch.zeros(C, C, H * W).numpy() I = np.identity(C) for i in range(H * W): # cov[:, :, i] = LedoitWolf().fit(embedding_vectors[:, :, i].numpy()).covariance_ cov[:, :, i] = np.cov(embedding_vectors[:, :, i].numpy(), rowvar=False) + 0.01 * I # save learned distribution train_outputs = [mean, cov] with open(train_feature_filepath, 'wb') as f: pickle.dump(train_outputs, f) else: print('load train set feature from: %s' % train_feature_filepath) with open(train_feature_filepath, 'rb') as f: train_outputs = pickle.load(f) gt_list = [] gt_mask_list = [] test_img_list = [] template_img_list = [] anomaly_point_lists = [] anomaly_loc_and_label_list = [] image_name_list = [] original_img_shape_list = [] # extract test set features for kk, (x, template_img, anomaly_loc_and_label, image_name, original_img_shape) in tqdm( enumerate(test_dataloader), '| feature extraction | test | %s |' % class_name): test_img_list.extend(x) template_img_list.extend(template_img) anomaly_loc_and_label_list.extend(anomaly_loc_and_label) image_name_list.extend(image_name) original_img_shape_list.extend(original_img_shape) #gt_list.extend(y.cpu().detach().numpy()) #gt_mask_list.extend(mask.cpu().detach().numpy()) # model prediction with torch.no_grad(): _ = model(x.to(device)) """if kk==0: model=tensorrt_optimize_model(torch.stack(test_img_list),model)""" # get intermediate layer outputs # get intermediate layer outputs for k, v in zip(test_outputs.keys(), outputs): test_outputs[k].append(v.cpu().detach()) # initialize hook outputs outputs = [] for k, v in test_outputs.items(): test_outputs[k] = torch.cat(v, 0) # Embedding concat embedding_vectors = test_outputs['layer1'] for layer_name in ['layer2', 'layer3']: embedding_vectors = embedding_concat(embedding_vectors, test_outputs[layer_name]) # randomly select d dimension embedding_vectors = torch.index_select(embedding_vectors, 1, idx) # calculate distance matrix B, C, H, W = embedding_vectors.size() embedding_vectors = embedding_vectors.view(B, C, H * W).numpy() dist_list = [] for i in range(H * W): mean = train_outputs[0][:, i] conv_inv = np.linalg.inv(train_outputs[1][:, :, i]) dist = [ mahalanobis(sample[:, i], mean, conv_inv) for sample in embedding_vectors ] dist_list.append(dist) dist_list = np.array(dist_list).transpose(1, 0).reshape(B, H, W) # upsample dist_list = torch.tensor(dist_list) score_map = F.interpolate(dist_list.unsqueeze(1), size=x.size(2), mode='bilinear', align_corners=False).squeeze().numpy() # apply gaussian smoothing on the score map # for disturb anlysis, this can be deleted? for i in range(score_map.shape[0]): score_map[i] = gaussian_filter(score_map[i], sigma=1) # Normalization max_score = score_map.max() min_score = score_map.min() scores = (score_map - min_score) / (max_score - min_score) # calculate image-level ROC AUC score img_scores = scores.reshape(scores.shape[0], -1).max(axis=1) """gt_list = np.asarray(gt_list) fpr, tpr, _ = roc_curve(gt_list, img_scores) img_roc_auc = roc_auc_score(gt_list, img_scores) total_roc_auc.append(img_roc_auc) print('image ROCAUC: %.3f' % (img_roc_auc)) fig_img_rocauc.plot(fpr, tpr, label='%s img_ROCAUC: %.3f' % (class_name, img_roc_auc)) # get optimal threshold #gt_mask = np.asarray(gt_mask_list) precision, recall, thresholds = precision_recall_curve(gt_mask.flatten(), scores.flatten()) a = 2 * precision * recall b = precision + recall f1 = np.divide(a, b, out=np.zeros_like(a), where=b != 0) threshold = thresholds[np.argmax(f1)]""" """# calculate per-pixel level ROCAUC fpr, tpr, _ = roc_curve(gt_mask.flatten(), scores.flatten()) per_pixel_rocauc = roc_auc_score(gt_mask.flatten(), scores.flatten()) total_pixel_roc_auc.append(per_pixel_rocauc) print('pixel ROCAUC: %.3f' % (per_pixel_rocauc))""" #fig_pixel_rocauc.plot(fpr, tpr, label='%s ROCAUC: %.3f' % (class_name, per_pixel_rocauc)) save_picture_dir = args.save_path + '/' + f'pictures_{args.arch}' save_image_dir = args.save_path + '/' + f'segment_image_result_{args.arch}' os.makedirs(save_picture_dir, exist_ok=True) os.makedirs(save_image_dir, exist_ok=True) plot_fig(test_img_list, template_img_list, anomaly_loc_and_label_list, original_img_shape_list, scores, anomaly_point_lists, save_picture_dir, save_image_dir, class_name, args.threshold_coefficient, image_name_list) """print('Average ROCAUC: %.3f' % np.mean(total_roc_auc))
def main(): # config = vars(parse_args_func()) #config_file = "../configs/config_v1.json" args = vars(parse_args_func()) config_file = args['config'] config_dict = json.loads(open(config_file, 'rt').read()) # config_dict = json.loads(open(sys.argv[1], 'rt').read()) file_dict = config_dict['file_path'] config = config_dict['opt_config'] input_folder = file_dict['input_path'] # '../inputs' checkpoint_folder = file_dict['checkpoint_path'] # '../checkpoint' model_folder = file_dict['model_path'] # '../models' if 'False' in config['deep_supervision']: config['deep_supervision'] = False else: config['deep_supervision'] = True if 'False' in config['nesterov']: config['nesterov'] = False else: config['nesterov'] = True if 'None' in config['name']: config['name'] = None if config['name'] is None: config['name'] = '%s_%s_segmodel' % (config['dataset'], config['arch']) os.makedirs(os.path.join(model_folder, '%s' % config['name']), exist_ok=True) print('-' * 20) for key in config: print('%s: %s' % (key, config[key])) print('-' * 20) with open(os.path.join(model_folder, '%s/config.yml' % config['name']), 'w') as f: yaml.dump(config, f) # define loss function (criterion) if config['loss'] == 'BCEWithLogitsLoss': criterion = nn.BCEWithLogitsLoss().cuda() else: criterion = losses.__dict__[config['loss']]().cuda() cudnn.benchmark = True # create model if 'False' in config['resume']: config['resume'] = False else: config['resume'] = True # Data loading code img_ids = glob( os.path.join(input_folder, config['dataset'], 'images', 'training', '*' + config['img_ext'])) train_img_ids = [os.path.splitext(os.path.basename(p))[0] for p in img_ids] img_ids = glob( os.path.join(input_folder, config['val_dataset'], 'images', 'validation', '*' + config['img_ext'])) val_img_ids = [os.path.splitext(os.path.basename(p))[0] for p in img_ids] img_ids = glob( os.path.join(input_folder, config['val_dataset'], 'images', 'test', '*' + config['img_ext'])) test_img_ids = [os.path.splitext(os.path.basename(p))[0] for p in img_ids] mean = [0.485, 0.456, 0.406] std = [0.229, 0.224, 0.225] train_transform = Compose([ # transforms.RandomScale ([config['scale_min'], config['scale_max']]), # transforms.RandomRotate90(), transforms.Rotate([config['rotate_min'], config['rotate_max']], value=mean, mask_value=0), # transforms.GaussianBlur (), transforms.Flip(), # transforms.HorizontalFlip (), transforms.HueSaturationValue(hue_shift_limit=10, sat_shift_limit=10, val_shift_limit=10), transforms.RandomBrightnessContrast(brightness_limit=0.10, contrast_limit=0.10, brightness_by_max=True), transforms.Resize(config['input_h'], config['input_w']), transforms.Normalize(mean=mean, std=std), ]) val_transform = Compose([ transforms.Resize(config['input_h'], config['input_w']), transforms.Normalize(mean=mean, std=std), ]) train_dataset = Dataset(img_ids=train_img_ids, img_dir=os.path.join(input_folder, config['dataset'], 'images', 'training'), mask_dir=os.path.join(input_folder, config['dataset'], 'annotations', 'training'), img_ext=config['img_ext'], mask_ext=config['mask_ext'], num_classes=config['num_classes'], input_channels=config['input_channels'], transform=train_transform) val_dataset = Dataset(img_ids=val_img_ids, img_dir=os.path.join(input_folder, config['dataset'], 'images', 'validation'), mask_dir=os.path.join(input_folder, config['dataset'], 'annotations', 'validation'), img_ext=config['img_ext'], mask_ext=config['mask_ext'], num_classes=config['num_classes'], input_channels=config['input_channels'], transform=val_transform) test_dataset = Dataset(img_ids=val_img_ids, img_dir=os.path.join(input_folder, config['dataset'], 'images', 'test'), mask_dir=os.path.join(input_folder, config['dataset'], 'annotations', 'test'), img_ext=config['img_ext'], mask_ext=config['mask_ext'], num_classes=config['num_classes'], input_channels=config['input_channels'], transform=val_transform) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=config['batch_size'], shuffle=True, num_workers=config['num_workers'], drop_last=True) val_loader = torch.utils.data.DataLoader( val_dataset, batch_size=1, # config['batch_size'], shuffle=False, num_workers=config['num_workers'], drop_last=False) test_loader = torch.utils.data.DataLoader( test_dataset, batch_size=1, # config['batch_size'], shuffle=False, num_workers=config['num_workers'], drop_last=False) log = OrderedDict([ ('epoch', []), ('lr', []), ('loss', []), ('iou', []), ('dice', []), ('val_loss', []), ('val_iou', []), ('val_dice', []), ]) if not os.path.isdir(checkpoint_folder): os.mkdir(checkpoint_folder) # create generator model #val_config = config_dict['config'] generator_name = config['generator_name'] with open(os.path.join(model_folder, '%s/config.yml' % generator_name), 'r') as f: g_config = yaml.load(f, Loader=yaml.FullLoader) generator = Generator(g_config) generator.initialize_with_srresnet(model_folder, g_config) lr = config['gan_lr'] # Initialize generator's optimizer optimizer_g = torch.optim.Adam(params=filter(lambda p: p.requires_grad, generator.parameters()), lr=lr) #params = filter(lambda p: p.requires_grad, generator.parameters()) #optimizer_g, scheduler_g = optimizer_scheduler(params, config) # Discriminator # Discriminator parameters num_classes = config['num_classes'] kernel_size_d = 3 # kernel size in all convolutional blocks n_channels_d = 64 # number of output channels in the first convolutional block, after which it is doubled in every 2nd block thereafter n_blocks_d = 8 # number of convolutional blocks fc_size_d = 1024 # size of the first fully connected layer discriminator = Discriminator(num_classes, kernel_size=kernel_size_d, n_channels=n_channels_d, n_blocks=n_blocks_d, fc_size=fc_size_d) # Initialize discriminator's optimizer optimizer_d = torch.optim.Adam(params=filter(lambda p: p.requires_grad, discriminator.parameters()), lr=lr) #params = filter(lambda p: p.requires_grad, discriminator.parameters()) #optimizer_d, scheduler_d = optimizer_scheduler(params, config) adversarial_loss_criterion = nn.BCEWithLogitsLoss() content_loss_criterion = nn.MSELoss() generator = generator.cuda() discriminator = discriminator.cuda() #truncated_vgg19 = truncated_vgg19.cuda() content_loss_criterion = content_loss_criterion.cuda() adversarial_loss_criterion = adversarial_loss_criterion.cuda() generator = torch.nn.DataParallel(generator) discriminator = torch.nn.DataParallel(discriminator) if not os.path.isdir(checkpoint_folder): os.mkdir(checkpoint_folder) log_name = config['name'] log_dir = os.path.join(checkpoint_folder, log_name) writer = SummaryWriter(logdir=log_dir) best_iou = 0 trigger = 0 Best_dice = 0 iou_AtBestDice = 0 start_epoch = 0 for epoch in range(start_epoch, config['epochs']): print('Epoch [%d/%d]' % (epoch, config['epochs'])) # train for one epoch train_log = train(epoch, config, train_loader, generator, discriminator, criterion, adversarial_loss_criterion, content_loss_criterion, optimizer_g, optimizer_d) # evaluate on validation set val_log = validate(config, val_loader, generator, criterion) test_log = validate(config, test_loader, generator, criterion) if Best_dice < test_log['dice']: Best_dice = test_log['dice'] iou_AtBestDice = test_log['iou'] print( 'loss %.4f - iou %.4f - dice %.4f - val_loss %.4f - val_iou %.4f - val_dice %.4f - test_iou %.4f - test_dice %.4f - Best_dice %.4f - iou_AtBestDice %.4f' % (train_log['loss'], train_log['iou'], train_log['dice'], val_log['loss'], val_log['iou'], val_log['dice'], test_log['iou'], test_log['dice'], Best_dice, iou_AtBestDice)) save_tensorboard(writer, train_log, val_log, test_log, epoch) log['epoch'].append(epoch) log['lr'].append(config['lr']) log['loss'].append(train_log['loss']) log['iou'].append(train_log['iou']) log['dice'].append(train_log['dice']) log['val_loss'].append(val_log['loss']) log['val_iou'].append(val_log['iou']) log['val_dice'].append(val_log['dice']) pd.DataFrame(log).to_csv(os.path.join(model_folder, '%s/log.csv' % config['name']), index=False) trigger += 1 if test_log['iou'] > best_iou: torch.save( generator.state_dict(), os.path.join(model_folder, '%s/model.pth' % config['name'])) best_iou = test_log['iou'] print("=> saved best model") trigger = 0 # early stopping if config['early_stopping'] >= 0 and trigger >= config[ 'early_stopping']: print("=> early stopping") break torch.cuda.empty_cache()
def main(): args = parse_args() config_file = "../configs/config_SN7.json" config_dict = json.loads(open(config_file, 'rt').read()) #config_dict = json.loads(open(sys.argv[1], 'rt').read()) file_dict = config_dict['file_path'] val_config = config_dict['val_config'] name = val_config['name'] input_folder =file_dict['input_path'] # '../inputs' model_folder = file_dict['model_path'] # '../models' output_folder = file_dict['output_path'] # '../models' ss_unet_GAN = True # create model if ss_unet_GAN == False: path = os.path.join(model_folder, '%s/config.yml' % name) with open(os.path.join(model_folder, '%s/config.yml' % name), 'r') as f: config = yaml.load(f, Loader=yaml.FullLoader) config['name'] = name print('-' * 20) for key in config.keys(): print('%s: %s' % (key, str(config[key]))) print('-' * 20) cudnn.benchmark = True print("=> creating model %s" % config['arch']) model = archs.__dict__[config['arch']](config['num_classes'], config['input_channels'], config['deep_supervision']) model = model.cuda() #img_ids = glob(os.path.join(input_folder, config['dataset'], 'images', '*' + config['img_ext'])) #img_ids = [os.path.splitext(os.path.basename(p))[0] for p in img_ids] #_, val_img_ids = train_test_split(img_ids, test_size=0.2, random_state=41) model_dict = torch.load(os.path.join(model_folder,'%s/model.pth' %config['name'])) if "state_dict" in model_dict.keys(): model_dict = remove_prefix(model_dict['state_dict'], 'module.') else: model_dict = remove_prefix(model_dict, 'module.') model.load_state_dict(model_dict, strict=False) #model.load_state_dict(torch.load(os.path.join(model_folder,'%s/model.pth' %config['name']))) model.eval() else: val_config = config_dict['val_config'] generator_name = val_config['name'] with open(os.path.join(model_folder, '%s/config.yml' % generator_name), 'r') as f: config = yaml.load(f, Loader=yaml.FullLoader) generator = Generator(config) generator = generator.cuda() ''' with open(os.path.join(model_folder, '%s/config.yml' % name), 'r') as f: config = yaml.load(f, Loader=yaml.FullLoader) ''' config['name'] = name model_dict = torch.load(os.path.join(model_folder,'%s/model.pth' %config['name'])) if "state_dict" in model_dict.keys(): model_dict = remove_prefix(model_dict['state_dict'], 'module.') else: model_dict = remove_prefix(model_dict, 'module.') generator.load_state_dict(model_dict, strict=False) #model.load_state_dict(torch.load(os.path.join(model_folder,'%s/model.pth' %config['name']))) generator.eval() # Data loading code img_ids = glob(os.path.join(input_folder, config['val_dataset'], 'images','test', '*' + config['img_ext'])) val_img_ids = [os.path.splitext(os.path.basename(p))[0] for p in img_ids] mean = [0.485, 0.456, 0.406] std = [0.229, 0.224, 0.225] val_transform = Compose([ transforms.Resize(config['input_h'], config['input_w']), #transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5)), transforms.Normalize(mean=mean, std=std), ]) val_dataset = Dataset( img_ids=val_img_ids, img_dir=os.path.join(input_folder, config['val_dataset'], 'images','test'), mask_dir=os.path.join(input_folder, config['val_dataset'], 'annotations','test'), img_ext=config['img_ext'], mask_ext=config['mask_ext'], num_classes=config['num_classes'], input_channels=config['input_channels'], transform=val_transform) val_loader = torch.utils.data.DataLoader( val_dataset, batch_size=1, #config['batch_size'], shuffle=False, num_workers=config['num_workers'], drop_last=False) avg_meters = {'iou': AverageMeter(), 'dice' : AverageMeter()} num_classes = config['num_classes'] for c in range(config['num_classes']): os.makedirs(os.path.join( output_folder, config['name'], str(c)), exist_ok=True) csv_save_name = os.path.join(output_folder, config['name'] + '_result' + '.csv') result_submission = [] with torch.no_grad(): pbar = tqdm(total=len(val_loader)) for ori_img, input, target, targets, meta in val_loader: input = input.cuda() target = target.cuda() # compute output if ss_unet_GAN == True: if config['deep_supervision']: output = generator(input)[-1] else: output = generator(input) else: if config['deep_supervision']: output = model(input)[-1] else: output = model(input) out_m = output[:, 1:num_classes, :, :].clone() tar_m = target[:, 1:num_classes, :, :].clone() iou = iou_score(out_m, tar_m) dice = dice_coef(out_m, tar_m) result_submission.append([meta['img_id'][0], iou, dice]) avg_meters['iou'].update(iou, input.size(0)) avg_meters['dice'].update(dice, input.size(0)) output = torch.sigmoid(output).cpu().numpy() masks = target.cpu() for i in range(len(output)): for idx_c in range(num_classes): tmp_mask = np.array(masks[i][idx_c]) mask = np.array(255 * tmp_mask).astype('uint8') mask_out = np.array(255 * output[i][idx_c]).astype('uint8') mask_output = np.zeros((mask_out.shape[0], mask_out.shape[1])) mask_output = mask_output.astype('uint8') mask_ = mask_out > 127 mask_output[mask_] = 255 if idx_c >0: save_GT_RE_mask(output_folder, config, meta, idx_c, i, ori_img, mask, mask_output) postfix = OrderedDict([ ('iou', avg_meters['iou'].avg), ('dice', avg_meters['dice'].avg), ]) pbar.set_postfix(postfix) pbar.update(1) pbar.close() result_save_to_csv_filename(csv_save_name, result_submission) print('IoU: %.4f' % avg_meters['iou'].avg) print('dice: %.4f' % avg_meters['dice'].avg) torch.cuda.empty_cache()
def perform_validation(modelName, testNum, fileName): ''' params: modelName, fileName => modelname for loading models from model directory, and filename to store results, both generated as per patient indices in train, test and val set. (For identification later) testNum => patient indices in test set Trained model tested on test set and results stored in fileName. No objects returned. ''' fw = open('batch_results_val/' + fileName, 'w') with open('models/%s/config.yml' % modelName, 'r') as f: config = yaml.load(f, Loader=yaml.FullLoader) print('-'*20) fw.write('-'*20 + '\n') for key in config.keys(): print('%s: %s' % (key, str(config[key]))) fw.write('%s: %s' % (key, str(config[key])) + '\n') print('-'*20) fw.write('-'*20 + '\n') cudnn.benchmark = True # create model print("=> creating model %s" % config['arch']) fw.write("=> creating model %s" % config['arch'] + '\n') model = archs.__dict__[config['arch']](config['num_classes'], config['input_channels'], config['deep_supervision']) model = model.cuda() # Data loading code img_ids = glob(os.path.join('inputs', config['dataset'], 'images', '*' + config['img_ext'])) img_ids = [os.path.splitext(os.path.basename(p))[0] for p in img_ids] # 2 patients data used for validation. Filtering those images from the # entire dataset. val_idx = [testNum, testNum + 1] val_img_ids = [] for img in img_ids: im_begin = img.split('.')[0] if int(im_begin[-1]) in val_idx: val_img_ids.append(img) # Loading model and setting to evaluation model (since we only need forward pass) model.load_state_dict(torch.load('models/%s/model.pth' % config['name'])) model.eval() # Pytorch objects for transformation, dataset and dataloader val_transform = Compose([ transforms.Resize(config['input_h'], config['input_w']), transforms.Normalize(), ]) val_dataset = Dataset( img_ids=val_img_ids, img_dir=os.path.join('inputs', config['dataset'], 'images'), mask_dir=os.path.join('inputs', config['dataset'], 'masks'), img_ext=config['img_ext'], mask_ext=config['mask_ext'], num_classes=config['num_classes'], transform=val_transform) val_loader = torch.utils.data.DataLoader( val_dataset, batch_size=config['batch_size'], shuffle=False, num_workers=config['num_workers'], drop_last=False) avg_meter = AverageMeter() dice_avg_meter = AverageMeter() for c in range(config['num_classes']): os.makedirs(os.path.join('outputs', config['name'], str(c)), exist_ok=True) # Running forward pass and storing results including masks, IoU and dice score. with torch.no_grad(): for input, target, meta in tqdm(val_loader, total=len(val_loader)): input = input.cuda() target = target.cuda() # compute output if config['deep_supervision']: output = model(input)[-1] else: output = model(input) iou = iou_score(output, target) avg_meter.update(iou, input.size(0)) dice = dice_coef(output, target) dice_avg_meter.update(dice, input.size(0)) output = torch.sigmoid(output).cpu().numpy() for i in range(len(output)): for c in range(config['num_classes']): cv2.imwrite(os.path.join('outputs', config['name'], str(c), meta['img_id'][i] + '.jpg'), (output[i, c] * 255).astype('uint8')) print('IoU: %.4f' % avg_meter.avg) fw.write('IoU: %.4f' % avg_meter.avg) print('Dice: %.4f' % dice_avg_meter.avg) fw.write('Dice: %.4f' % dice_avg_meter.avg) torch.cuda.empty_cache()
def main(): config = vars(parse_args()) if config['name'] is None: if config['deep_supervision']: config['name'] = '%s_%s_wDS' % (config['dataset'], config['arch']) else: config['name'] = '%s_%s_woDS' % (config['dataset'], config['arch']) os.makedirs('models/%s' % config['name'], exist_ok=True) print('-' * 20) for key in config: print('%s: %s' % (key, config[key])) print('-' * 20) with open('models/%s/config.yml' % config['name'], 'w') as f: yaml.dump(config, f) # define loss function (criterion) if config['loss'] == 'BCEWithLogitsLoss': criterion = nn.BCEWithLogitsLoss().cuda() else: criterion = losses.__dict__[config['loss']]().cuda() cudnn.benchmark = True # create model print("=> creating model %s" % config['arch']) model = archs.__dict__[config['arch']](config['num_classes'], config['input_channels'], config['deep_supervision']) model = model.cuda() params = filter(lambda p: p.requires_grad, model.parameters()) if config['optimizer'] == 'Adam': optimizer = optim.Adam(params, lr=config['lr'], weight_decay=config['weight_decay']) elif config['optimizer'] == 'SGD': optimizer = optim.SGD(params, lr=config['lr'], momentum=config['momentum'], nesterov=config['nesterov'], weight_decay=config['weight_decay']) else: raise NotImplementedError if config['scheduler'] == 'CosineAnnealingLR': scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=config['epochs'], eta_min=config['min_lr']) elif config['scheduler'] == 'ReduceLROnPlateau': scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, factor=config['factor'], patience=config['patience'], verbose=1, min_lr=config['min_lr']) elif config['scheduler'] == 'MultiStepLR': scheduler = lr_scheduler.MultiStepLR( optimizer, milestones=[int(e) for e in config['milestones'].split(',')], gamma=config['gamma']) elif config['scheduler'] == 'ConstantLR': scheduler = None else: raise NotImplementedError # Data loading code img_ids = glob( os.path.join('inputs', config['dataset'], 'images', '*' + config['img_ext'])) img_ids = [os.path.splitext(os.path.basename(p))[0] for p in img_ids] train_img_ids, val_img_ids = train_test_split(img_ids, test_size=0.2, random_state=41) ####################### train_img_ids = img_ids ####################### train_transform = Compose([ transforms.RandomRotate90(), transforms.Flip(), OneOf([ transforms.HueSaturationValue(), transforms.RandomBrightness(), transforms.RandomContrast(), ], p=1), transforms.Resize(config['input_h'], config['input_w']), transforms.Normalize(), ]) val_transform = Compose([ transforms.Resize(config['input_h'], config['input_w']), transforms.Normalize(), ]) train_dataset = Dataset(img_ids=train_img_ids, img_dir=os.path.join('inputs', config['dataset'], 'images'), mask_dir=os.path.join('inputs', config['dataset'], 'masks'), img_ext=config['img_ext'], mask_ext=config['mask_ext'], num_classes=config['num_classes'], transform=train_transform) val_dataset = Dataset(img_ids=val_img_ids, img_dir=os.path.join('inputs', config['dataset'], 'images'), mask_dir=os.path.join('inputs', config['dataset'], 'masks'), img_ext=config['img_ext'], mask_ext=config['mask_ext'], num_classes=config['num_classes'], transform=val_transform) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=config['batch_size'], shuffle=True, num_workers=config['num_workers'], drop_last=True) val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=config['batch_size'], shuffle=False, num_workers=config['num_workers'], drop_last=False) log = OrderedDict([ ('epoch', []), ('lr', []), ('loss', []), ('iou', []), ('val_loss', []), ('val_iou', []), ]) best_iou = 0 trigger = 0 for epoch in range(config['epochs']): print('Epoch [%d/%d]' % (epoch, config['epochs'])) # train for one epoch train_log = train(config, train_loader, model, criterion, optimizer) # evaluate on validation set val_log = validate(config, val_loader, model, criterion) if config['scheduler'] == 'CosineAnnealingLR': scheduler.step() elif config['scheduler'] == 'ReduceLROnPlateau': scheduler.step(val_log['loss']) print('loss %.4f - iou %.4f - val_loss %.4f - val_iou %.4f' % (train_log['loss'], train_log['iou'], val_log['loss'], val_log['iou'])) log['epoch'].append(epoch) log['lr'].append(config['lr']) log['loss'].append(train_log['loss']) log['iou'].append(train_log['iou']) log['val_loss'].append(val_log['loss']) log['val_iou'].append(val_log['iou']) pd.DataFrame(log).to_csv('models/%s/log.csv' % config['name'], index=False) trigger += 1 if val_log['iou'] > best_iou: torch.save(model.state_dict(), 'models/%s/model.pth' % config['name']) best_iou = val_log['iou'] print("=> saved best model") trigger = 0 # early stopping if config['early_stopping'] >= 0 and trigger >= config[ 'early_stopping']: print("=> early stopping") break torch.cuda.empty_cache()
def resize_im(im, height, width): resize_fn = transforms.Resize(height=height, width=width) return resize_fn(image=im)['image']
def main(): config = vars(parse_args()) if config['name'] is None: config['name'] = '%s_%s' % (config['arch'], datetime.now().strftime('%m%d%H')) if not os.path.exists('models/pose/%s' % config['name']): os.makedirs('models/pose/%s' % config['name']) if config['resume']: with open('models/pose/%s/config.yml' % config['name'], 'r') as f: config = yaml.load(f, Loader=yaml.FullLoader) config['resume'] = True with open('models/pose/%s/config.yml' % config['name'], 'w') as f: yaml.dump(config, f) print('-'*20) for key in config.keys(): print('- %s: %s' % (key, str(config[key]))) print('-'*20) cudnn.benchmark = True df = pd.read_csv('inputs/train.csv') img_ids = df['ImageId'].values pose_df = pd.read_csv('processed/pose_train.csv') pose_df['img_path'] = 'processed/pose_images/train/' + pose_df['img_path'] if config['resume']: checkpoint = torch.load('models/pose/%s/checkpoint.pth.tar' % config['name']) if config['rot'] == 'eular': num_outputs = 3 elif config['rot'] == 'trig': num_outputs = 6 elif config['rot'] == 'quat': num_outputs = 4 else: raise NotImplementedError if config['loss'] == 'L1Loss': criterion = nn.L1Loss().cuda() elif config['loss'] == 'MSELoss': criterion = nn.MSELoss().cuda() else: raise NotImplementedError train_transform = Compose([ transforms.ShiftScaleRotate( shift_limit=config['shift_limit'], scale_limit=0, rotate_limit=0, border_mode=cv2.BORDER_CONSTANT, value=0, p=config['shift_p'] ) if config['shift'] else NoOp(), OneOf([ transforms.HueSaturationValue( hue_shift_limit=config['hue_limit'], sat_shift_limit=config['sat_limit'], val_shift_limit=config['val_limit'], p=config['hsv_p'] ) if config['hsv'] else NoOp(), transforms.RandomBrightness( limit=config['brightness_limit'], p=config['brightness_p'], ) if config['brightness'] else NoOp(), transforms.RandomContrast( limit=config['contrast_limit'], p=config['contrast_p'], ) if config['contrast'] else NoOp(), ], p=1), transforms.ISONoise( p=config['iso_noise_p'], ) if config['iso_noise'] else NoOp(), transforms.CLAHE( p=config['clahe_p'], ) if config['clahe'] else NoOp(), transforms.Resize(config['input_w'], config['input_h']), transforms.Normalize(), ToTensor(), ]) val_transform = Compose([ transforms.Resize(config['input_w'], config['input_h']), transforms.Normalize(), ToTensor(), ]) folds = [] best_losses = [] kf = KFold(n_splits=config['n_splits'], shuffle=True, random_state=41) for fold, (train_idx, val_idx) in enumerate(kf.split(img_ids)): print('Fold [%d/%d]' %(fold + 1, config['n_splits'])) if (config['resume'] and fold < checkpoint['fold'] - 1) or (not config['resume'] and os.path.exists('pose_models/%s/model_%d.pth' % (config['name'], fold+1))): log = pd.read_csv('models/pose/%s/log_%d.csv' %(config['name'], fold+1)) best_loss = log.loc[log['val_loss'].values.argmin(), 'val_loss'] # best_loss, best_score = log.loc[log['val_loss'].values.argmin(), ['val_loss', 'val_score']].values folds.append(str(fold + 1)) best_losses.append(best_loss) # best_scores.append(best_score) continue train_img_ids, val_img_ids = img_ids[train_idx], img_ids[val_idx] train_img_paths = [] train_labels = [] for img_id in train_img_ids: tmp = pose_df.loc[pose_df.ImageId == img_id] img_path = tmp['img_path'].values train_img_paths.append(img_path) yaw = tmp['yaw'].values pitch = tmp['pitch'].values roll = tmp['roll'].values roll = rotate(roll, np.pi) if config['rot'] == 'eular': label = np.array([ yaw, pitch, roll ]).T elif config['rot'] == 'trig': label = np.array([ np.cos(yaw), np.sin(yaw), np.cos(pitch), np.sin(pitch), np.cos(roll), np.sin(roll), ]).T elif config['rot'] == 'quat': raise NotImplementedError else: raise NotImplementedError train_labels.append(label) train_img_paths = np.hstack(train_img_paths) train_labels = np.vstack(train_labels) val_img_paths = [] val_labels = [] for img_id in val_img_ids: tmp = pose_df.loc[pose_df.ImageId == img_id] img_path = tmp['img_path'].values val_img_paths.append(img_path) yaw = tmp['yaw'].values pitch = tmp['pitch'].values roll = tmp['roll'].values roll = rotate(roll, np.pi) if config['rot'] == 'eular': label = np.array([ yaw, pitch, roll ]).T elif config['rot'] == 'trig': label = np.array([ np.cos(yaw), np.sin(yaw), np.cos(pitch), np.sin(pitch), np.cos(roll), np.sin(roll), ]).T elif config['rot'] == 'quat': raise NotImplementedError else: raise NotImplementedError val_labels.append(label) val_img_paths = np.hstack(val_img_paths) val_labels = np.vstack(val_labels) # train train_set = PoseDataset( train_img_paths, train_labels, transform=train_transform, ) train_loader = torch.utils.data.DataLoader( train_set, batch_size=config['batch_size'], shuffle=True, num_workers=config['num_workers'], # pin_memory=True, ) val_set = PoseDataset( val_img_paths, val_labels, transform=val_transform, ) val_loader = torch.utils.data.DataLoader( val_set, batch_size=config['batch_size'], shuffle=False, num_workers=config['num_workers'], # pin_memory=True, ) # create model model = get_pose_model(config['arch'], num_outputs=num_outputs, freeze_bn=config['freeze_bn']) model = model.cuda() params = filter(lambda p: p.requires_grad, model.parameters()) if config['optimizer'] == 'Adam': optimizer = optim.Adam(params, lr=config['lr'], weight_decay=config['weight_decay']) elif config['optimizer'] == 'AdamW': optimizer = optim.AdamW(params, lr=config['lr'], weight_decay=config['weight_decay']) elif config['optimizer'] == 'RAdam': optimizer = RAdam(params, lr=config['lr'], weight_decay=config['weight_decay']) elif config['optimizer'] == 'SGD': optimizer = optim.SGD(params, lr=config['lr'], momentum=config['momentum'], nesterov=config['nesterov'], weight_decay=config['weight_decay']) else: raise NotImplementedError if config['scheduler'] == 'CosineAnnealingLR': scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=config['epochs'], eta_min=config['min_lr']) elif config['scheduler'] == 'ReduceLROnPlateau': scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, factor=config['factor'], patience=config['patience'], verbose=1, min_lr=config['min_lr']) elif config['scheduler'] == 'MultiStepLR': scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=[int(e) for e in config['milestones'].split(',')], gamma=config['gamma']) else: raise NotImplementedError log = { 'epoch': [], 'loss': [], # 'score': [], 'val_loss': [], # 'val_score': [], } best_loss = float('inf') # best_score = float('inf') start_epoch = 0 if config['resume'] and fold == checkpoint['fold'] - 1: model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) scheduler.load_state_dict(checkpoint['scheduler']) start_epoch = checkpoint['epoch'] log = pd.read_csv('models/pose/%s/log_%d.csv' % (config['name'], fold+1)).to_dict(orient='list') best_loss = checkpoint['best_loss'] for epoch in range(start_epoch, config['epochs']): print('Epoch [%d/%d]' % (epoch + 1, config['epochs'])) # train for one epoch train_loss = train(config, train_loader, model, criterion, optimizer, epoch) # evaluate on validation set val_loss = validate(config, val_loader, model, criterion) if config['scheduler'] == 'CosineAnnealingLR': scheduler.step() elif config['scheduler'] == 'ReduceLROnPlateau': scheduler.step(val_loss) print('loss %.4f - val_loss %.4f' % (train_loss, val_loss)) # print('loss %.4f - score %.4f - val_loss %.4f - val_score %.4f' # % (train_loss, train_score, val_loss, val_score)) log['epoch'].append(epoch) log['loss'].append(train_loss) # log['score'].append(train_score) log['val_loss'].append(val_loss) # log['val_score'].append(val_score) pd.DataFrame(log).to_csv('models/pose/%s/log_%d.csv' % (config['name'], fold+1), index=False) if val_loss < best_loss: torch.save(model.state_dict(), 'models/pose/%s/model_%d.pth' % (config['name'], fold+1)) best_loss = val_loss # best_score = val_score print("=> saved best model") state = { 'fold': fold + 1, 'epoch': epoch + 1, 'state_dict': model.state_dict(), 'best_loss': best_loss, 'optimizer': optimizer.state_dict(), 'scheduler': scheduler.state_dict(), } torch.save(state, 'models/pose/%s/checkpoint.pth.tar' % config['name']) print('val_loss: %f' % best_loss) # print('val_score: %f' % best_score) folds.append(str(fold + 1)) best_losses.append(best_loss) # best_scores.append(best_score) results = pd.DataFrame({ 'fold': folds + ['mean'], 'best_loss': best_losses + [np.mean(best_losses)], # 'best_score': best_scores + [np.mean(best_scores)], }) print(results) results.to_csv('models/pose/%s/results.csv' % config['name'], index=False) del model torch.cuda.empty_cache() del train_set, train_loader del val_set, val_loader gc.collect() if not config['cv']: break
def main(): args = parse_args() with open('models/pose/%s/config.yml' % args.pose_name, 'r') as f: config = yaml.load(f, Loader=yaml.FullLoader) print('-' * 20) for key in config.keys(): print('%s: %s' % (key, str(config[key]))) print('-' * 20) cudnn.benchmark = True df = pd.read_csv('inputs/train.csv') img_ids = df['ImageId'].values img_paths = np.array('inputs/train_images/' + df['ImageId'].values + '.jpg') mask_paths = np.array('inputs/train_masks/' + df['ImageId'].values + '.jpg') labels = np.array( [convert_str_to_labels(s) for s in df['PredictionString']]) with open('outputs/decoded/val/%s.json' % args.det_name, 'r') as f: dets = json.load(f) if config['rot'] == 'eular': num_outputs = 3 elif config['rot'] == 'trig': num_outputs = 6 elif config['rot'] == 'quat': num_outputs = 4 else: raise NotImplementedError test_transform = Compose([ transforms.Resize(config['input_w'], config['input_h']), transforms.Normalize(), ToTensor(), ]) det_df = { 'ImageId': [], 'img_path': [], 'det': [], 'mask': [], } name = '%s_%.2f' % (args.det_name, args.score_th) if args.nms: name += '_nms%.2f' % args.nms_th output_dir = 'processed/pose_images/val/%s' % name os.makedirs(output_dir, exist_ok=True) df = [] kf = KFold(n_splits=config['n_splits'], shuffle=True, random_state=41) for fold, (train_idx, val_idx) in enumerate(kf.split(img_paths)): print('Fold [%d/%d]' % (fold + 1, config['n_splits'])) # create model model = get_pose_model(config['arch'], num_outputs=num_outputs, freeze_bn=config['freeze_bn']) model = model.cuda() model_path = 'models/pose/%s/model_%d.pth' % (config['name'], fold + 1) if not os.path.exists(model_path): print('%s is not exists.' % model_path) continue model.load_state_dict(torch.load(model_path)) model.eval() val_img_ids = img_ids[val_idx] val_img_paths = img_paths[val_idx] fold_det_df = { 'ImageId': [], 'img_path': [], 'det': [], 'mask': [], } for img_id, img_path in tqdm(zip(val_img_ids, val_img_paths), total=len(val_img_ids)): img = cv2.imread(img_path) height, width = img.shape[:2] det = np.array(dets[img_id]) det = det[det[:, 6] > args.score_th] if args.nms: det = nms(det, dist_th=args.nms_th) for k in range(len(det)): pitch, yaw, roll, x, y, z, score, w, h = det[k] fold_det_df['ImageId'].append(img_id) fold_det_df['det'].append(det[k]) output_path = '%s_%d.jpg' % (img_id, k) fold_det_df['img_path'].append(output_path) x, y = convert_3d_to_2d(x, y, z) w *= 1.1 h *= 1.1 xmin = int(round(x - w / 2)) xmax = int(round(x + w / 2)) ymin = int(round(y - h / 2)) ymax = int(round(y + h / 2)) cropped_img = img[ymin:ymax, xmin:xmax] if cropped_img.shape[0] > 0 and cropped_img.shape[1] > 0: cv2.imwrite(os.path.join(output_dir, output_path), cropped_img) fold_det_df['mask'].append(1) else: fold_det_df['mask'].append(0) fold_det_df = pd.DataFrame(fold_det_df) test_set = PoseDataset(output_dir + '/' + fold_det_df['img_path'].values, fold_det_df['det'].values, transform=test_transform, masks=fold_det_df['mask'].values) test_loader = torch.utils.data.DataLoader( test_set, batch_size=config['batch_size'], shuffle=False, num_workers=config['num_workers'], # pin_memory=True, ) fold_dets = [] with torch.no_grad(): for input, batch_det, mask in tqdm(test_loader, total=len(test_loader)): input = input.cuda() batch_det = batch_det.numpy() mask = mask.numpy() output = model(input) output = output.cpu() if config['rot'] == 'trig': yaw = torch.atan2(output[..., 1:2], output[..., 0:1]) pitch = torch.atan2(output[..., 3:4], output[..., 2:3]) roll = torch.atan2(output[..., 5:6], output[..., 4:5]) roll = rotate(roll, -np.pi) pitch = pitch.cpu().numpy()[:, 0] yaw = yaw.cpu().numpy()[:, 0] roll = roll.cpu().numpy()[:, 0] batch_det[mask, 0] = pitch[mask] batch_det[mask, 1] = yaw[mask] batch_det[mask, 2] = roll[mask] fold_dets.append(batch_det) fold_dets = np.vstack(fold_dets) fold_det_df['det'] = fold_dets.tolist() fold_det_df = fold_det_df.groupby('ImageId')['det'].apply(list) fold_det_df = pd.DataFrame({ 'ImageId': fold_det_df.index.values, 'PredictionString': fold_det_df.values, }) df.append(fold_det_df) break df = pd.concat(df).reset_index(drop=True) for i in tqdm(range(len(df))): img_id = df.loc[i, 'ImageId'] det = np.array(df.loc[i, 'PredictionString']) if args.show: img = cv2.imread('inputs/train_images/%s.jpg' % img_id) img_pred = visualize(img, det) plt.imshow(img_pred[..., ::-1]) plt.show() df.loc[i, 'PredictionString'] = convert_labels_to_str(det[:, :7]) name += '_%s' % args.pose_name df.to_csv('outputs/submissions/val/%s.csv' % name, index=False)
def main(): args = parse_args() with open('models/%s/config.yml' % args.name, 'r') as f: config = yaml.load(f, Loader=yaml.FullLoader) print('-' * 20) for key in config.keys(): print('%s: %s' % (key, str(config[key]))) print('-' * 20) cudnn.benchmark = True # create model print("=> creating model %s" % config['arch']) model = archs.__dict__[config['arch']](config['num_classes'], config['input_channels'], config['deep_supervision']) model = model.cuda() # Data loading code img_ids = glob( os.path.join('inputs', config['dataset'], 'test\\images', '*' + config['img_ext'])) ############ 바꿈 img_ids = [os.path.splitext(os.path.basename(p))[0] for p in img_ids] val_img_ids = img_ids model.load_state_dict(torch.load('models/%s/model.pth' % config['name'])) model.eval() val_transform = Compose([ transforms.Resize(config['input_h'], config['input_w']), transforms.Normalize(), ]) val_dataset = Dataset( img_ids=val_img_ids, img_dir=os.path.join('inputs', config['dataset'], 'test\\images'), ############ 바꿈 mask_dir=os.path.join('inputs', config['dataset'], 'test\\masks'), ############ 바꿈 img_ext=config['img_ext'], mask_ext=config['mask_ext'], num_classes=config['num_classes'], transform=val_transform) val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=config['batch_size'], shuffle=False, num_workers=config['num_workers'], drop_last=False) avg_meter = AverageMeter() for c in range(config['num_classes']): os.makedirs(os.path.join('outputs', config['name'], str(c)), exist_ok=True) with torch.no_grad(): for input, target, meta in tqdm(val_loader, total=len(val_loader)): input = input.cuda() target = target.cuda() # compute output if config['deep_supervision']: output = model(input)[-1] else: output = model(input) iou = iou_score(output, target) avg_meter.update(iou, input.size(0)) output = torch.sigmoid(output).cpu().numpy() for i in range(len(output)): for c in range(config['num_classes']): cv2.imwrite( os.path.join('outputs', config['name'], str(c), meta['img_id'][i] + '.jpg'), (output[i, c] * 255).astype('uint8')) print('IoU: %.4f' % avg_meter.avg) torch.cuda.empty_cache()
def main(): config = vars(parse_args()) now = datetime.datetime.now() if config["name"] is None: if config["deep_supervision"]: config["name"] = "%s_%s_wDS_%s" % ( config["dataset"], config["arch"], now.strftime("%Y%m%d_%H%M%S"), ) else: config["name"] = "%s_%s_woDS_%s" % ( config["dataset"], config["arch"], now.strftime("%Y%m%d_%H%M%S"), ) output_path = os.path.join(cfg.UNET_RESULTS_DIR, config["name"]) try: os.makedirs(output_path, exist_ok=True) except Exception as e: print(e) models_path = os.path.join(output_path, "models") os.mkdir(models_path) with open(os.path.join(models_path, "config.yml"), "w") as f: yaml.dump(config, f) print("-" * 20) for key in config: print("%s: %s" % (key, config[key])) print("-" * 20) # Tensorboad 用のログを記録するディレクトリパス log_dir = os.path.join(output_path, "log") os.mkdir(log_dir) writer = SummaryWriter(log_dir=log_dir) # define loss function(criterion) if config["loss"] == "BCEWithLogitsLoss": criterion = nn.BCEWithLogitsLoss().cuda() else: criterion = losses.__dict__[config["loss"]]().cuda() cudnn.benchmark = True # create model print("=> creating model %s" % config["arch"]) model = archs.__dict__[config["arch"]](config["num_classes"], config["input_channels"], config["deep_supervision"]) model = model.cuda() # モデルを TensorBorad で表示するため,ログに保存 # image = torch.randn(1, 3, 2224, 224) # writer.add_graph(model, image) params = filter(lambda p: p.requires_grad, model.parameters()) if config["optimizer"] == "Adam": optimizer = optim.Adam(params, lr=config["lr"], weight_decay=config["weight_decay"]) elif config["optimizer"] == "SGD": optimizer = optim.SGD( params, lr=config["lr"], momentum=config["momentum"], nesterov=config["nesterov"], weight_decay=config["weight_decay"], ) else: raise NotImplementedError # scheduler if config["scheduler"] == "CosineAnnealingLR": scheduler = lr_scheduler.CosineAnnealingLR(optimizer=optimizer, T_max=config["epochs"], eta_min=config["min_lr"]) elif config["scheduler"] == "ReduceLROnPlateau": scheduler = lr_scheduler( optimizer=optimizer, factor=config["factor"], patience=config["patience"], verbose=1, min_lr=config["min_lr"], ) elif config["scheduler"] == "MultiStepLR": scheduler = lr_scheduler.MultiStepLR( optimizer=optimizer, milestones=[int(e) for e in config["milestones"].split(",")], gamma=config["gamma"], ) elif config["scheduler"] == "ConstantLR": scheduler = None else: raise NotImplementedError # Data loading code if config["dataset"] == "dsb2018_96": input_dir = cfg.DSB2018_96_DIR img_ids = glob( os.path.join(input_dir, "images", "*" + config["img_ext"])) img_ids = [os.path.splitext(os.path.basename(p))[0] for p in img_ids] train_img_ids, val_img_ids = train_test_split(img_ids, test_size=0.2, random_state=41) train_transform = Compose([ transforms.RandomRotate90(), transforms.Flip(), OneOf( [ transforms.HueSaturationValue(), transforms.RandomBrightness(), transforms.RandomContrast(), ], p=1, ), transforms.Resize(config["input_h"], config["input_w"]), transforms.Normalize(), ]) val_transform = Compose([ transforms.Resize(config["input_h"], config["input_w"]), transforms.Normalize(), ]) train_dataset = Dataset( img_ids=train_img_ids, img_dir=os.path.join(input_dir, "images"), mask_dir=os.path.join(input_dir, "masks"), img_ext=config["img_ext"], mask_ext=config["mask_ext"], num_classes=config["num_classes"], transform=train_transform, ) val_dataset = Dataset( img_ids=val_img_ids, img_dir=os.path.join(input_dir, "images"), mask_dir=os.path.join(input_dir, "masks"), img_ext=config["img_ext"], mask_ext=config["mask_ext"], num_classes=config["num_classes"], transform=val_transform, ) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=config["batch_size"], shuffle=True, num_workers=config["num_workers"], drop_last=True, ) val_loader = torch.utils.data.DataLoader( val_dataset, batch_size=config["batch_size"], shuffle=False, num_workers=config["num_workers"], drop_last=False, ) log = OrderedDict([ ("epoch", []), ("lr", []), ("loss", []), ("iou", []), ("val_loss", []), ("val_iou", []), ]) best_iou = 0 trigger = 0 for epoch in range(config["epochs"] + 1): print("Epoch [%d/%d]" % (epoch, config["epochs"])) # train for one epoch train_log = train(config, train_loader, model, criterion, optimizer) # evaluate on validation set val_log = validate(config, val_loader, model, criterion) if config["scheduler"] == "CosineAnnealingLR": scheduler.step() elif config["scheduler"] == "ReduceLROnPlateau": scheduler.step(val_log["loss"]) print("loss %.4f - iou %.4f - val_loss %.4f - val_iou %.4f" % (train_log["loss"], train_log["iou"], val_log["loss"], val_log["iou"])) log["epoch"].append(epoch) log["lr"].append(config["lr"]) log["loss"].append(train_log["loss"]) log["iou"].append(train_log["iou"]) log["val_loss"].append(val_log["loss"]) log["val_iou"].append(val_log["iou"]) # Tensorboard用のデータ writer.add_scalar("training loss", train_log["loss"], epoch) writer.add_scalar("validation loss", val_log["loss"], epoch) pd.DataFrame(log).to_csv("%s/log.csv" % (log_dir), index=False) if epoch == 0: best_loss = val_log["loss"] trigger += 1 # Best Model Save # if val_log['iou'] > best_iou: if (val_log["iou"] > best_iou) and (val_log["loss"] <= best_loss): torch.save(model.state_dict(), "%s/model.pth" % (models_path)) best_iou = val_log["iou"] best_loss = val_log["loss"] print("=> saved best model") trigger = 0 # early stopping if (config["early_stopping"] >= 0 and trigger >= config["early_stopping"]) or val_log["loss"] < 1e-4: print("=> early stopping") break torch.cuda.empty_cache() # summary writer を必要としない場合,close()メソッドを呼び出す writer.close()
def main_func(train_idx, val_set, test_set, modelName, fileName): config = vars(parse_args()) config['name'] = modelName fw = open('batch_results_train/' + fileName, 'w') print('config of dataset is ' + str(config['dataset'])) fw.write('config of dataset is ' + str(config['dataset']) + '\n') if config['name'] is None: if config['deep_supervision']: config['name'] = '%s_%s_wDS' % (config['dataset'], config['arch']) else: config['name'] = '%s_%s_woDS' % (config['dataset'], config['arch']) os.makedirs('models/%s' % config['name'], exist_ok=True) print('-' * 20) fw.write('-' * 20 + '\n') for key in config: print('%s: %s' % (key, config[key])) fw.write('%s: %s' % (key, config[key]) + '\n') print('-' * 20) fw.write('-' * 20 + '\n') #TODO print parameters manually i think, all imports to function with open('models/%s/config.yml' % config['name'], 'w') as f: yaml.dump(config, f) # define loss function (criterion) if config['loss'] == 'BCEWithLogitsLoss': criterion = nn.BCEWithLogitsLoss().cuda() else: criterion = losses.__dict__[config['loss']]().cuda() cudnn.benchmark = True # create model print("=> creating model %s" % config['arch']) fw.write("=> creating model %s" % config['arch'] + '\n') model = archs.__dict__[config['arch']](config['num_classes'], config['input_channels'], config['deep_supervision']) model = model.cuda() params = filter(lambda p: p.requires_grad, model.parameters()) if config['optimizer'] == 'Adam': optimizer = optim.Adam(params, lr=config['lr'], weight_decay=config['weight_decay']) elif config['optimizer'] == 'SGD': optimizer = optim.SGD(params, lr=config['lr'], momentum=config['momentum'], nesterov=config['nesterov'], weight_decay=config['weight_decay']) else: raise NotImplementedError if config['scheduler'] == 'CosineAnnealingLR': scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=config['epochs'], eta_min=config['min_lr']) elif config['scheduler'] == 'ReduceLROnPlateau': scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, factor=config['factor'], patience=config['patience'], verbose=1, min_lr=config['min_lr']) elif config['scheduler'] == 'MultiStepLR': scheduler = lr_scheduler.MultiStepLR( optimizer, milestones=[int(e) for e in config['milestones'].split(',')], gamma=config['gamma']) elif config['scheduler'] == 'ConstantLR': scheduler = None else: raise NotImplementedError # Data loading code img_ids = glob( os.path.join('inputs', config['dataset'], 'images', '*' + config['img_ext'])) img_ids = [os.path.splitext(os.path.basename(p))[0] for p in img_ids] #train_img_ids, val_img_ids = train_test_split(img_ids, test_size=0.2, random_state=41 val_idx = [val_set, val_set + 1] #train_idx = [2, 3, 6, 7] val_img_ids = [] train_img_ids = [] for image in img_ids: im_begin = image.split('.')[0] if int(im_begin[-1]) in val_idx: val_img_ids.append(image) elif int(im_begin[-1]) in train_idx: train_img_ids.append(image) #print("train img ids size is " + str(len(train_img_ids))) '''train_transform = Compose([ transforms.RandomRotate90(), transforms.Flip(), OneOf([ transforms.HueSaturationValue(), transforms.RandomBrightness(), transforms.RandomContrast(), ], p=1), transforms.Resize(config['input_h'], config['input_w']), transforms.Normalize(), ]) val_transform = Compose([ transforms.Resize(config['input_h'], config['input_w']), transforms.Normalize(), ]) ''' train_transform = Compose([ #transforms.RandomRotate90(), #transforms.Flip(), #OneOf([ # transforms.HueSaturationValue(), # transforms.RandomBrightness(), # transforms.RandomContrast(), #], p=1), transforms.Resize(config['input_h'], config['input_w']), transforms.Normalize(), ]) train_transform2 = Compose([ transforms.Resize(config['input_h'], config['input_w']), transforms.Normalize(), transforms.ShiftScaleRotate( shift_limit=0.1, scale_limit=0, rotate_limit=0 ), # shift_limit_x = 0.1, shift_limit_y = 0.1, p = 1), ]) val_transform2 = Compose([ transforms.Resize(config['input_h'], config['input_w']), transforms.Normalize(), #transforms.RandomAffine(degrees = 0, translate = (10, 10)), transforms.ShiftScaleRotate( shift_limit=0.1, scale_limit=0, rotate_limit=0 ), # shift_limit_x = 0.1, shift_limit_y = 0.1, p = 1), ##TODO remove from validation ]) val_transform = Compose([ transforms.Resize(config['input_h'], config['input_w']), transforms.Normalize(), ]) train_dataset = Dataset(img_ids=train_img_ids, img_dir=os.path.join('inputs', config['dataset'], 'images'), mask_dir=os.path.join('inputs', config['dataset'], 'masks'), img_ext=config['img_ext'], mask_ext=config['mask_ext'], num_classes=config['num_classes'], transform=train_transform2) val_dataset = Dataset(img_ids=val_img_ids, img_dir=os.path.join('inputs', config['dataset'], 'images'), mask_dir=os.path.join('inputs', config['dataset'], 'masks'), img_ext=config['img_ext'], mask_ext=config['mask_ext'], num_classes=config['num_classes'], transform=val_transform2) #print("length of train dataset is " + str(len(train_dataset))) #print("length of val dataset is " + str(len(val_dataset))) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=config['batch_size'], shuffle=True, num_workers=config['num_workers'], drop_last=True) val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=config['batch_size'], shuffle=False, num_workers=config['num_workers'], drop_last=False) log = OrderedDict([ ('epoch', []), ('lr', []), ('loss', []), ('iou', []), ('val_loss', []), ('val_iou', []), ('dice', []), ]) best_iou = 0 trigger = 0 best_dice = 0 for epoch in range(config['epochs']): print('Epoch [%d/%d]' % (epoch, config['epochs'])) fw.write('Epoch [%d/%d]' % (epoch, config['epochs']) + '\n') # train for one epoch train_log = train(config, train_loader, model, criterion, optimizer) # evaluate on validation set val_log = validate(config, val_loader, model, criterion) if config['scheduler'] == 'CosineAnnealingLR': scheduler.step() elif config['scheduler'] == 'ReduceLROnPlateau': scheduler.step(val_log['loss']) print('loss %.4f - iou %.4f - val_loss %.4f - val_iou %.4f dice %.4f' % (train_log['loss'], train_log['iou'], val_log['loss'], val_log['iou'], val_log['dice'])) fw.write( 'loss %.4f - iou %.4f - val_loss %.4f - val_iou %.4f dice %.4f' % (train_log['loss'], train_log['iou'], val_log['loss'], val_log['iou'], val_log['dice']) + '\n') log['epoch'].append(epoch) log['lr'].append(config['lr']) log['loss'].append(train_log['loss']) log['iou'].append(train_log['iou']) log['val_loss'].append(val_log['loss']) log['val_iou'].append(val_log['iou']) log['dice'].append(val_log['dice']) pd.DataFrame(log).to_csv('models/%s/log.csv' % config['name'], index=False) trigger += 1 ''' if val_log['iou'] > best_iou: torch.save(model.state_dict(), 'models/%s/model.pth' % config['name']) best_iou = val_log['iou'] print("=> saved best model") trigger = 0 ''' if val_log['dice'] > best_dice: torch.save(model.state_dict(), 'models/%s/model.pth' % config['name']) best_dice = val_log['dice'] print("=> saved best model") fw.write("=> saved best model" + '\n') trigger = 0 # early stopping if config['early_stopping'] >= 0 and trigger >= config[ 'early_stopping']: print("=> early stopping") fw.write("=> early stopping" + '\n') break torch.cuda.empty_cache()
def main(): if 1: config_file = "../configs/config_v1.json" config_dict = json.loads(open(config_file, 'rt').read()) val_config = config_dict['val_config'] full_img_path = val_config['full_image_path'] gt_mask_flag = True if 'False' in val_config['gt_mask_flag']: gt_mask_flag = False else: gt_mask_flag = True image_folder = os.path.join(full_img_path, '*_image.*') output_folder = '../outputs' ## Test Image image_paths = glob(image_folder) model, config = load_segmentation_models(config_file) for idx in tqdm(range(len(image_paths))): img_path = image_paths[idx] save_image_name = os.path.basename(img_path) save_name, ext = os.path.splitext(save_image_name) ## Load Segmentation Model ## Get patch images img_input, img_patch_set, mask_patch_set = get_patched_input( img_path, config, gt_mask_flag) ## Operate the segmentation with patches re_class_mask, gt_class_mask = segmentation_inference( model, img_input, img_patch_set, mask_patch_set, config, gt_mask_flag) ## save the color mask of regions save_image_color_masking(output_folder, save_name, img_input, re_class_mask, gt_class_mask, config, gt_mask_flag) torch.cuda.empty_cache() if 0: args = parse_args() config_file = "../configs/config_v1.json" config_dict = json.loads(open(config_file, 'rt').read()) file_dict = config_dict['file_path'] val_config = config_dict['val_config'] name = val_config['name'] full_img_path = val_config['full_image_path'] p_overlap = val_config['patch_overlap'] input_folder = file_dict['input_path'] # '../inputs' model_folder = file_dict['model_path'] # '../models' output_folder = file_dict['output_path'] # '../models' name = val_config['name'] with open(os.path.join(model_folder, '%s/config.yml' % name), 'r') as f: config = yaml.load(f, Loader=yaml.FullLoader) print('-' * 20) for key in config.keys(): print('%s: %s' % (key, str(config[key]))) print('-' * 20) cudnn.benchmark = True # create model print("=> creating model %s" % config['arch']) model = archs.__dict__[config['arch']](config['num_classes'], config['input_channels'], config['deep_supervision']) model = model.cuda() # Data loading code ''' img_ids = glob(os.path.join(input_folder, config['dataset'], 'images', '*' + config['img_ext'])) img_ids = [os.path.splitext(os.path.basename(p))[0] for p in img_ids] _, val_img_ids = train_test_split(img_ids, test_size=0.2, random_state=41) ''' image_folder = os.path.join(full_img_path, '*_image.*') model.load_state_dict( torch.load( os.path.join(model_folder, '%s/model.pth' % config['name']))) model.eval() #seg_infer(model, config_file) val_transform = Compose([ transforms.Resize(config['input_h'], config['input_w']), transforms.Normalize(), ]) patch_size = 1024 infer_size = config['input_w'] val_dataset = DatasetPatch( img_ids=image_folder, #img_dir=os.path.join(input_folder, config['dataset'], 'images'), #mask_dir=os.path.join(input_folder, config['dataset'], 'masks'), img_ext=config['img_ext'], mask_ext=config['mask_ext'], num_classes=config['num_classes'], input_channels=config['input_channels'], image_w=infer_size, psize=patch_size, patch_overlap=p_overlap, transform=val_transform) val_loader = torch.utils.data.DataLoader( val_dataset, batch_size=1, #config['batch_size'], shuffle=False, num_workers=config['num_workers'], drop_last=False) with torch.no_grad(): for full_img, full_label, input, target, image_name in tqdm( val_loader, total=len(val_loader)): fname = os.path.splitext(image_name['img_name'][0]) input = np.squeeze(input, 0) input = input.cuda() # compute output full_output = [] for idx, data in enumerate(input): in_tmp = torch.unsqueeze(data, 0) output = model(in_tmp) output = torch.sigmoid(output).cpu().numpy() full_output.append(np.squeeze(output, 0)) full_img = np.squeeze(full_img, 0) gt_label = [] for idx, data in enumerate(target): mask_tmp = [] for c in range(config['num_classes']): data = torch.squeeze(data, 0) con_mask = mask_convert(data, c, infer_size) mask_tmp.append(con_mask) mask_patch = np.dstack(mask_tmp) mask_patch = mask_patch.transpose(2, 0, 1) gt_label.append(mask_patch / 255.0) all_class_mask = patch_merge(full_img, full_output, patch_size, config, p_overlap) gt_class_mask = patch_merge(full_img, gt_label, patch_size, config, p_overlap) for c in range(config['num_classes']): file_name = '{:s}_{:d}'.format(fname[0], c) full_img = np.array(full_img) if c > 0: save_name = os.path.join( output_folder, config['name'], file_name + '_GT_RE_masking_p075.jpg') save_name_GT = os.path.join( output_folder, config['name'], file_name + '_GT_masking_p075.jpg') save_name_RE = os.path.join( output_folder, config['name'], file_name + '_RE_masking_p075.jpg') mask_output = all_class_mask[c] mask_gt = gt_class_mask[c] #save_contour(np.array(full_img), mask_gt, mask_output, save_name) save_img = full_img.copy() save_masking_GT(np.array(save_img), mask_gt, mask_output, save_name_GT) save_img = full_img.copy() save_masking_RE(np.array(save_img), mask_gt, mask_output, save_name_RE) save_img = full_img.copy() save_masking(np.array(save_img), mask_gt, mask_output, save_name) torch.cuda.empty_cache()
def test_entry(config, ix_sum=10, ix=0): print('-' * 20) for key in config: print('%s: %s' % (key, config[key])) print('-' * 20) with open('models/%s/config.yml' % config['name'], 'w') as f: yaml.dump(config, f) cudnn.benchmark = True # create model arch = config['arch'].split('_')[0] print("=> creating model %s" % arch) model = archs.__dict__[arch](config['num_classes'], config['input_channels'], config['deep_supervision']) model = nn.DataParallel(model) # load trained model print("Reloading model 'models/%s/model.pth'..." % config['name']) model.load_state_dict(torch.load('models/%s/model.pth' % config['name'])) #model = model.cuda() # Data loading code data_path = os.path.join('data', config['dataset']) img_ids = glob( os.path.join(data_path, config['sub_dataset'] + '_testing', 'images', '*' + config['img_ext'])) img_ids = [os.path.splitext(os.path.basename(p))[0] for p in img_ids] split_filename = os.path.join(data_path, 'split.p') if config['load_split'] is True: split_file = open(split_filename, 'rb') _, val_img_ids = pickle.load(split_file) else: val_img_ids = img_ids # read statistics with open('data_processing/data_statistics.json') as json_file: data = json.load(json_file) entry = '%s_%s' % (config['dataset'], config['sub_dataset']) mean = data[entry]['mean'] std = data[entry]['std'] print("%s: mean=[%s] std=[%s]" % (config['dataset'], ', '.join( map(str, mean)), ', '.join(map(str, std)))) # data normalization normalize = transforms.Normalize(mean=mean, std=std) test_transform = Compose([ transforms.Resize(config['input_h'], config['input_w']), # normalize, ]) test_dataset = Dataset( img_ids=val_img_ids, img_dir=os.path.join('data', config['dataset'], config['sub_dataset'] + '_testing', 'images'), mask_dir=os.path.join('data', config['dataset'], config['sub_dataset'] + '_testing', 'masks'), img_ext=config['img_ext'], mask_ext=config['mask_ext'], num_classes=config['num_classes'], transform=test_transform) test_loader = torch.utils.data.DataLoader( test_dataset, batch_size=config['batch_size'], shuffle=False, num_workers=config['num_workers'], drop_last=False) if config['predict']: predict(config, test_loader, model) else: test_result = test(config, test_loader, model) df = pd.DataFrame(test_result.items(), columns=["metric", "result"]) print(df) test_result_per_class = test_per_class(config, test_loader, model) for class_id in range(config['num_classes']): print('[class %d]:' % class_id) df = pd.DataFrame(test_result_per_class[class_id].items(), columns=["metric", "result"]) print(df)
def main(): args = parse_args() if args.name is None: args.name = '%s_%s' % (args.arch, datetime.now().strftime('%m%d%H')) if not os.path.exists('models/%s' % args.name): os.makedirs('models/%s' % args.name) if args.resume: args = joblib.load('models/%s/args.pkl' % args.name) args.resume = True print('Config -----') for arg in vars(args): print('- %s: %s' % (arg, getattr(args, arg))) print('------------') with open('models/%s/args.txt' % args.name, 'w') as f: for arg in vars(args): print('- %s: %s' % (arg, getattr(args, arg)), file=f) joblib.dump(args, 'models/%s/args.pkl' % args.name) if args.seed is not None and not args.resume: print('set random seed') random.seed(args.seed) np.random.seed(args.seed) torch.manual_seed(args.seed) if args.loss == 'BCEWithLogitsLoss': criterion = BCEWithLogitsLoss().cuda() elif args.loss == 'WeightedBCEWithLogitsLoss': criterion = BCEWithLogitsLoss(weight=torch.Tensor([1., 1., 1., 1., 1., 2.]), smooth=args.label_smooth).cuda() elif args.loss == 'FocalLoss': criterion = FocalLoss().cuda() elif args.loss == 'WeightedFocalLoss': criterion = FocalLoss(weight=torch.Tensor([1., 1., 1., 1., 1., 2.])).cuda() else: raise NotImplementedError if args.pred_type == 'all': num_outputs = 6 elif args.pred_type == 'except_any': num_outputs = 5 else: raise NotImplementedError cudnn.benchmark = True # create model model = get_model(model_name=args.arch, num_outputs=num_outputs, freeze_bn=args.freeze_bn, dropout_p=args.dropout_p, pooling=args.pooling, lp_p=args.lp_p) model = model.cuda() train_transform = Compose([ transforms.Resize(args.img_size, args.img_size), transforms.HorizontalFlip() if args.hflip else NoOp(), transforms.VerticalFlip() if args.vflip else NoOp(), transforms.ShiftScaleRotate( shift_limit=args.shift_limit, scale_limit=args.scale_limit, rotate_limit=args.rotate_limit, border_mode=cv2.BORDER_CONSTANT, value=0, p=args.shift_scale_rotate_p ) if args.shift_scale_rotate else NoOp(), transforms.RandomContrast( limit=args.contrast_limit, p=args.contrast_p ) if args.contrast else NoOp(), RandomErase() if args.random_erase else NoOp(), transforms.CenterCrop(args.crop_size, args.crop_size) if args.center_crop else NoOp(), ForegroundCenterCrop(args.crop_size) if args.foreground_center_crop else NoOp(), transforms.RandomCrop(args.crop_size, args.crop_size) if args.random_crop else NoOp(), transforms.Normalize(mean=model.mean, std=model.std), ToTensor(), ]) if args.img_type: stage_1_train_dir = 'processed/stage_1_train_%s' %args.img_type else: stage_1_train_dir = 'processed/stage_1_train' df = pd.read_csv('inputs/stage_1_train.csv') img_paths = np.array([stage_1_train_dir + '/' + '_'.join(s.split('_')[:-1]) + '.png' for s in df['ID']][::6]) labels = np.array([df.loc[c::6, 'Label'].values for c in range(6)]).T.astype('float32') df = df[::6] df['img_path'] = img_paths for c in range(6): df['label_%d' %c] = labels[:, c] df['ID'] = df['ID'].apply(lambda s: '_'.join(s.split('_')[:-1])) meta_df = pd.read_csv('processed/stage_1_train_meta.csv') meta_df['ID'] = meta_df['SOPInstanceUID'] test_meta_df = pd.read_csv('processed/stage_1_test_meta.csv') df = pd.merge(df, meta_df, how='left') patient_ids = meta_df['PatientID'].unique() test_patient_ids = test_meta_df['PatientID'].unique() if args.remove_test_patient_ids: patient_ids = np.array([s for s in patient_ids if not s in test_patient_ids]) train_img_paths = np.hstack(df[['img_path', 'PatientID']].groupby(['PatientID'])['img_path'].apply(np.array).loc[patient_ids].to_list()).astype('str') train_labels = [] for c in range(6): train_labels.append(np.hstack(df[['label_%d' %c, 'PatientID']].groupby(['PatientID'])['label_%d' %c].apply(np.array).loc[patient_ids].to_list())) train_labels = np.array(train_labels).T if args.resume: checkpoint = torch.load('models/%s/checkpoint.pth.tar' % args.name) # train train_set = Dataset( train_img_paths, train_labels, transform=train_transform) train_loader = torch.utils.data.DataLoader( train_set, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, # pin_memory=True, ) if args.optimizer == 'Adam': optimizer = optim.Adam( filter(lambda p: p.requires_grad, model.parameters()), lr=args.lr, weight_decay=args.weight_decay) elif args.optimizer == 'AdamW': optimizer = optim.AdamW( filter(lambda p: p.requires_grad, model.parameters()), lr=args.lr, weight_decay=args.weight_decay) elif args.optimizer == 'RAdam': optimizer = RAdam( filter(lambda p: p.requires_grad, model.parameters()), lr=args.lr, weight_decay=args.weight_decay) elif args.optimizer == 'SGD': optimizer = optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay, nesterov=args.nesterov) else: raise NotImplementedError if args.apex: amp.initialize(model, optimizer, opt_level='O1') if args.scheduler == 'CosineAnnealingLR': scheduler = lr_scheduler.CosineAnnealingLR( optimizer, T_max=args.epochs, eta_min=args.min_lr) elif args.scheduler == 'MultiStepLR': scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=[int(e) for e in args.milestones.split(',')], gamma=args.gamma) else: raise NotImplementedError log = { 'epoch': [], 'loss': [], } start_epoch = 0 if args.resume: model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) scheduler.load_state_dict(checkpoint['scheduler']) start_epoch = checkpoint['epoch'] log = pd.read_csv('models/%s/log.csv' % args.name).to_dict(orient='list') for epoch in range(start_epoch, args.epochs): print('Epoch [%d/%d]' % (epoch + 1, args.epochs)) # train for one epoch train_loss = train(args, train_loader, model, criterion, optimizer, epoch) if args.scheduler == 'CosineAnnealingLR': scheduler.step() print('loss %.4f' % (train_loss)) log['epoch'].append(epoch) log['loss'].append(train_loss) pd.DataFrame(log).to_csv('models/%s/log.csv' % args.name, index=False) torch.save(model.state_dict(), 'models/%s/model.pth' % args.name) print("=> saved model") state = { 'epoch': epoch + 1, 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict(), 'scheduler': scheduler.state_dict(), } torch.save(state, 'models/%s/checkpoint.pth.tar' % args.name)
def train_entry(config, rank=-1, world_size=0): if (rank > 0): setup(rank, world_size) os.makedirs('models/%s' % config['name'], exist_ok=True) with open('models/%s/config.yml' % config['name'], 'w') as f: yaml.dump(config, f) # define loss function (criterion) if config['loss'] == 'BCEWithLogitsLoss': criterion = nn.BCEWithLogitsLoss().cuda() else: criterion = losses.__dict__[config['loss']]().cuda() cudnn.benchmark = True # create model arch = config['arch'].split('_')[0] print("=> creating model %s" % arch) model = getattr(archs, arch)(config['num_classes'], config['input_channels'], config['deep_supervision']) model = nn.DataParallel(model) print(f'rank={rank}, device={torch.cuda.current_device()}') if rank >= 0: model = DDP(model, device_ids=[rank]) else: model = model.cuda() # load trained model model_file = 'models/%s/model.pth' % config['name'] if os.path.isfile(model_file): print("Reloading model ...") model.load_state_dict(torch.load(model_file)) params = filter(lambda p: p.requires_grad, model.parameters()) if config['optimizer'] == 'Adam': optimizer = optim.Adam( params, lr=config['lr'], weight_decay=config['weight_decay']) elif config['optimizer'] == 'SGD': optimizer = optim.SGD(params, lr=config['lr'], momentum=config['momentum'], nesterov=config['nesterov'], weight_decay=config['weight_decay']) else: raise NotImplementedError if config['scheduler'] == 'CosineAnnealingLR': scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=config['epochs'], eta_min=config['min_lr']) elif config['scheduler'] == 'CosineAnnealingWarmRestarts': scheduler = lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=50) elif config['scheduler'] == 'ReduceLROnPlateau': scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', factor=config['factor'], patience=config['patience'], min_lr=config['min_lr']) elif config['scheduler'] == 'MultiStepLR': scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=[int(e) for e in config['milestones'].split(',')], gamma=config['gamma']) elif config['scheduler'] == 'ConstantLR': scheduler = None else: raise NotImplementedError # Data loading code data_path = os.path.join('data', config['dataset']) img_ids = glob(os.path.join(data_path, config['sub_dataset'], 'images', '*' + config['img_ext'])) # img_ids = [os.path.basename(p) for p in img_ids] img_ids = [os.path.splitext(os.path.basename(p))[0] for p in img_ids] if config['num_inputs'] > 0: img_ids = img_ids[:config['num_inputs']] train_img_ids, val_img_ids = train_test_split(img_ids, test_size=0.2, random_state=41) split_filename = os.path.join(data_path, 'split.p') if config['save_split'] is True: split_file = open(split_filename, 'wb') pickle.dump((train_img_ids, val_img_ids), split_file) if config['load_split'] is True: split_file = open(split_filename, 'rb') train_img_ids, val_img_ids = pickle.load(split_file) ''' Since it is gray image, some augmentation methods not used ''' # read statistics with open('data_processing/data_statistics.json') as json_file: data = json.load(json_file) entry = '%s_%s' % (config['dataset'], config['sub_dataset']) mean = data[entry]['mean'] std = data[entry]['std'] print("%s: mean=[%s] std=[%s]" % (config['dataset'], ', '.join(map(str, mean)), ', '.join(map(str, std)))) # data normalization normalize = transforms.Normalize(mean=mean, std=std) train_transform = Compose([ transforms.RandomRotate90(), transforms.Flip(), # transforms.HorizontalFlip(), # transforms.Rotate(limit=5), # OneOf([ # transforms.HueSaturationValue(), # transforms.RandomBrightnessContrast(), # # transforms.RandomContrast(), # ], p=1), transforms.Resize(config['input_h'], config['input_w']), # normalize, ]) val_transform = Compose([ transforms.Resize(config['input_h'], config['input_w']), # normalize, ]) train_dataset = Dataset( img_ids=train_img_ids, img_dir=os.path.join('data', config['dataset'], config['sub_dataset'], 'images'), mask_dir=os.path.join('data', config['dataset'], config['sub_dataset'], 'masks'), img_ext=config['img_ext'], mask_ext=config['mask_ext'], num_classes=config['num_classes'], transform=train_transform) val_dataset = Dataset( img_ids=val_img_ids, img_dir=os.path.join('data', config['dataset'], config['sub_dataset'], 'images'), mask_dir=os.path.join('data', config['dataset'], config['sub_dataset'], 'masks'), img_ext=config['img_ext'], mask_ext=config['mask_ext'], num_classes=config['num_classes'], transform=val_transform) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=config['batch_size'], shuffle=True, num_workers=config['num_workers'], drop_last=True) val_loader = torch.utils.data.DataLoader( val_dataset, batch_size=config['batch_size'], shuffle=False, num_workers=config['num_workers'], drop_last=False) log = OrderedDict([ ('epoch', []), ('lr', []), ('loss', []), ('iou', []), ('dice', []), ('val_loss', []), ('val_iou', []), ('val_dice', []) ]) log_per_class = [] for _ in range(config['num_classes']): log_per_class.append(OrderedDict([('val_iou', []), ('val_dice', [])])) best_iou = 0 best_iou_per_class = [0] * config['num_classes'] trigger = 0 for epoch in range(config['epochs']): print('Epoch [%d/%d]' % (epoch, config['epochs'])) # train for one epoch train_log = train(config, train_loader, model, criterion, optimizer) # evaluate on validation set val_log = validate(config, val_loader, model, criterion) # evaluate on validation set for each class val_per_class_log = test_per_class(config, val_loader, model) if config['scheduler'] == 'CosineAnnealingLR': scheduler.step() elif config['scheduler'] == 'ReduceLROnPlateau': scheduler.step(val_log['loss']) print('loss %.4f - iou %.4f - dice %.4f - val_loss %.4f - val_iou %.4f - val_dice %.4f' % (train_log['loss'], train_log['iou'], train_log['dice'], val_log['loss'], val_log['iou'], val_log['dice'])) log['epoch'].append(epoch) log['lr'].append(config['lr']) log['loss'].append(train_log['loss']) log['iou'].append(train_log['iou']) log['dice'].append(train_log['dice']) log['val_loss'].append(val_log['loss']) log['val_iou'].append(val_log['iou']) log['val_dice'].append(val_log['dice']) pd.DataFrame(log).to_csv('models/%s/log.csv' % config['name'], index=False) # write into log for each class for class_id in range(config['num_classes']): log_per_class[class_id]['val_iou'].append(val_per_class_log[class_id]['iou']) log_per_class[class_id]['val_dice'].append(val_per_class_log[class_id]['dice']) print('[class %d] val_iou %.4f - val_dice %.4f' % (class_id, val_per_class_log[class_id]['iou'], val_per_class_log[class_id]['dice'])) pd.DataFrame(log_per_class[class_id]).to_csv('models/%s/log_%d.csv' % (config['name'], class_id), index=False) if val_per_class_log[class_id]['iou'] > best_iou_per_class[class_id]: torch.save(model.state_dict(), 'models/%s/model_%d.pth' % (config['name'], class_id)) best_iou_per_class[class_id] = val_per_class_log[class_id]['iou'] print("===> saved best model for class %d" % class_id) trigger += 1 # torch.save(model.state_dict(), 'models/%s/model.pth' % config['name']) if val_log['iou'] > best_iou: torch.save(model.state_dict(), 'models/%s/model.pth' % config['name']) best_iou = val_log['iou'] print("=> saved best model") trigger = 0 # early stopping if config['early_stopping'] >= 0 and trigger >= config['early_stopping']: print("=> early stopping") break torch.cuda.empty_cache() print("\n[Best Results:]") print('- [Overall IoU] %.4f' % best_iou) for class_id in range(config['num_classes']): print('-- [class %d IoU] %.4f' % (class_id, best_iou_per_class[class_id])) if rank >= 0: cleanup() return best_iou
def main(): args = vars(parse_args_func()) #config_file = "../configs/config_SN7.json" config_file = args['config'] # "../configs/config_v1.json" config_dict = json.loads(open(config_file, 'rt').read()) #config_dict = json.loads(open(sys.argv[1], 'rt').read()) file_dict = config_dict['file_path'] config = config_dict['opt_config'] input_folder = file_dict['input_path'] # '../inputs' checkpoint_folder = file_dict['checkpoint_path'] # '../checkpoint' model_folder = file_dict['model_path'] # '../models' if 'False' in config['deep_supervision']: config['deep_supervision'] = False else: config['deep_supervision'] = True if 'False' in config['nesterov']: config['nesterov'] = False else: config['nesterov'] = True if 'None' in config['name']: config['name'] = None if config['name'] is None: config['name'] = '%s_%s_segmodel' % (config['dataset'], config['arch']) os.makedirs(os.path.join(model_folder, '%s' % config['name']), exist_ok=True) if not os.path.isdir(checkpoint_folder): os.mkdir(checkpoint_folder) log_name = config['name'] log_dir = os.path.join(checkpoint_folder, log_name) writer = SummaryWriter(logdir=log_dir) print('-' * 20) for key in config: print('%s: %s' % (key, config[key])) print('-' * 20) with open(os.path.join(model_folder, '%s/config.yml' % config['name']), 'w') as f: yaml.dump(config, f) # define loss function (criterion) if config['loss'] == 'BCEWithLogitsLoss': criterion = nn.BCEWithLogitsLoss().cuda() else: criterion = losses.__dict__[config['loss']]().cuda() cudnn.benchmark = True # create model print("=> creating model %s" % config['arch']) model = archs.__dict__[config['arch']](config['num_classes'], config['input_channels'], config['deep_supervision']) if 'False' in config['resume']: config['resume'] = False else: config['resume'] = True resume_flag = False if resume_flag == True: save_path = os.path.join(model_folder, config['name'], 'model.pth') weights = torch.load(save_path) model.load_state_dict(weights) name_yaml = config['name'] with open(os.path.join(model_folder, '%s/config.yml' % name_yaml), 'r') as f: config = yaml.load(f, Loader=yaml.FullLoader) #start_epoch = config['epochs'] start_epoch = 0 else: start_epoch = 0 model = model.cuda() if 'effnet' in config['arch']: eff_flag = True else: eff_flag = False if eff_flag == True: cnn_subs = list(model.encoder.eff_conv.children())[1:] #cnn_params = [list(sub_module.parameters()) for sub_module in cnn_subs] #cnn_params = [item for sublist in cnn_params for item in sublist] summary(model, (config['input_channels'], config['input_w'], config['input_h'])) params = filter(lambda p: p.requires_grad, model.parameters()) if eff_flag == True: params = list(params) + list(model.encoder.conv_a.parameters()) model = torch.nn.DataParallel(model) if config['optimizer'] == 'Adam': optimizer = optim.Adam(params, lr=config['lr'], weight_decay=config['weight_decay']) elif config['optimizer'] == 'SGD': optimizer = optim.SGD(params, lr=config['lr'], momentum=config['momentum'], nesterov=config['nesterov'], weight_decay=config['weight_decay']) else: raise NotImplementedError if eff_flag == True: cnn_params = [list(sub_module.parameters()) for sub_module in cnn_subs] cnn_params = [item for sublist in cnn_params for item in sublist] cnn_optimizer = torch.optim.Adam(cnn_params, lr=0.001, weight_decay=config['weight_decay']) #cnn_optimizer = None else: cnn_optimizer = None if config['optimizer'] == 'SGD': if config['scheduler'] == 'CosineAnnealingLR': scheduler = lr_scheduler.CosineAnnealingLR( optimizer, T_max=config['epochs'], eta_min=config['min_lr']) elif config['scheduler'] == 'ReduceLROnPlateau': scheduler = lr_scheduler.ReduceLROnPlateau( optimizer, factor=config['factor'], patience=config['patience'], verbose=1, min_lr=config['min_lr']) elif config['scheduler'] == 'MultiStepLR': scheduler = lr_scheduler.MultiStepLR( optimizer, milestones=[int(e) for e in config['milestones'].split(',')], gamma=config['gamma']) elif config['scheduler'] == 'ConstantLR': scheduler = None else: raise NotImplementedError else: scheduler = None # Data loading code img_ids = glob( os.path.join(input_folder, config['dataset'], 'images', 'training', '*' + config['img_ext'])) train_img_ids = [os.path.splitext(os.path.basename(p))[0] for p in img_ids] #img_dir = os.path.join(input_folder, config['dataset'], 'images', 'training') #mask_dir = os.path.join(input_folder, config['dataset'], 'annotations', 'training') #train_image_mask = image_to_afile(img_dir, mask_dir, None, train_img_ids, config) img_ids = glob( os.path.join(input_folder, config['val_dataset'], 'images', 'validation', '*' + config['img_ext'])) val_img_ids = [os.path.splitext(os.path.basename(p))[0] for p in img_ids] img_ids = glob( os.path.join(input_folder, config['val_dataset'], 'images', 'test', '*' + config['img_ext'])) test_img_ids = [os.path.splitext(os.path.basename(p))[0] for p in img_ids] mean = [0.485, 0.456, 0.406] std = [0.229, 0.224, 0.225] train_transform = Compose([ #transforms.RandomScale ([config['scale_min'], config['scale_max']]), #transforms.RandomRotate90(), transforms.Rotate([config['rotate_min'], config['rotate_max']], value=mean, mask_value=0), transforms.Flip(), #transforms.HorizontalFlip (), transforms.HueSaturationValue(hue_shift_limit=10, sat_shift_limit=10, val_shift_limit=10), transforms.RandomBrightnessContrast(brightness_limit=0.10, contrast_limit=0.10, brightness_by_max=True), transforms.Resize(config['input_h'], config['input_w']), transforms.Normalize(mean=mean, std=std), ]) val_transform = Compose([ transforms.Resize(config['input_h'], config['input_w']), transforms.Normalize(mean=mean, std=std), ]) train_dataset = Dataset(img_ids=train_img_ids, img_dir=os.path.join(input_folder, config['dataset'], 'images', 'training'), mask_dir=os.path.join(input_folder, config['dataset'], 'annotations', 'training'), img_ext=config['img_ext'], mask_ext=config['mask_ext'], num_classes=config['num_classes'], input_channels=config['input_channels'], transform=train_transform, from_file=None) val_dataset = Dataset(img_ids=val_img_ids, img_dir=os.path.join(input_folder, config['val_dataset'], 'images', 'validation'), mask_dir=os.path.join(input_folder, config['val_dataset'], 'annotations', 'validation'), img_ext=config['img_ext'], mask_ext=config['mask_ext'], num_classes=config['num_classes'], input_channels=config['input_channels'], transform=val_transform, from_file=None) test_dataset = Dataset(img_ids=test_img_ids, img_dir=os.path.join(input_folder, config['val_dataset'], 'images', 'test'), mask_dir=os.path.join(input_folder, config['val_dataset'], 'annotations', 'test'), img_ext=config['img_ext'], mask_ext=config['mask_ext'], num_classes=config['num_classes'], input_channels=config['input_channels'], transform=val_transform, from_file=None) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=config['batch_size'], shuffle=True, num_workers=config['num_workers'], drop_last=True) val_loader = torch.utils.data.DataLoader( val_dataset, batch_size=1, #config['batch_size'], shuffle=False, num_workers=config['num_workers'], drop_last=False) test_loader = torch.utils.data.DataLoader( test_dataset, batch_size=1, #config['batch_size'], shuffle=False, num_workers=config['num_workers'], drop_last=False) log = OrderedDict([ ('epoch', []), ('lr', []), ('loss', []), ('iou', []), ('dice', []), ('val_loss', []), ('val_iou', []), ('val_dice', []), ]) best_iou = 0 trigger = 0 Best_dice = 0 iou_AtBestDice = 0 for epoch in range(start_epoch, config['epochs']): print('{:s} Epoch [{:d}/{:d}]'.format(config['arch'], epoch, config['epochs'])) # train for one epoch train_log = train(epoch, config, train_loader, model, criterion, optimizer, cnn_optimizer) if config['optimizer'] == 'SGD': if config['scheduler'] == 'CosineAnnealingLR': scheduler.step() elif config['scheduler'] == 'ReduceLROnPlateau': scheduler.step(val_log['loss']) elif config['scheduler'] == 'MultiStepLR': scheduler.step() # evaluate on validation set val_log = validate(config, val_loader, model, criterion) test_log = validate(config, test_loader, model, criterion) if Best_dice < test_log['dice']: Best_dice = test_log['dice'] iou_AtBestDice = test_log['iou'] print( 'loss %.4f - iou %.4f - dice %.4f - val_loss %.4f - val_iou %.4f - val_dice %.4f - test_iou %.4f - test_dice %.4f - Best_dice %.4f - iou_AtBestDice %.4f' % (train_log['loss'], train_log['iou'], train_log['dice'], val_log['loss'], val_log['iou'], val_log['dice'], test_log['iou'], test_log['dice'], Best_dice, iou_AtBestDice)) save_tensorboard(writer, train_log, val_log, test_log, epoch) log['epoch'].append(epoch) log['lr'].append(config['lr']) log['loss'].append(train_log['loss']) log['iou'].append(train_log['iou']) log['dice'].append(train_log['dice']) log['val_loss'].append(val_log['loss']) log['val_iou'].append(val_log['iou']) log['val_dice'].append(val_log['dice']) pd.DataFrame(log).to_csv(os.path.join(model_folder, '%s/log.csv' % config['name']), index=False) trigger += 1 if val_log['iou'] > best_iou: torch.save( model.state_dict(), os.path.join(model_folder, '%s/model.pth' % config['name'])) best_iou = val_log['iou'] print("=> saved best model") trigger = 0 # early stopping if config['early_stopping'] >= 0 and trigger >= config[ 'early_stopping']: print("=> early stopping") break torch.cuda.empty_cache()
def Resize(height=cons.IMAGE_SIZE, width=cons.IMAGE_SIZE): return transforms.Resize(height=height, width=width)
def test(v, model, folds, visualize=False): gts = [] prs = [] for id in list(folds.keys()): test_fold = "fold" + str(id) _data_name = "Kvasir" data_path = "Kvasir_fold_new/" + "fold_" + str(id) # data_path = 'kvasir-seg/TestDataset/CVC-300' save_path = "results/PraNet_kfold/" model_path = ("./snapshots/_PraNetv" + str(v) + "_Res2Net_kfold/" + "PraNetDG-" + test_fold + "-" + str(folds[id]) + ".pth") print(model_path) try: model.load_state_dict(torch.load(model_path)["model_state_dict"]) except RuntimeError: model.load_state_dict(torch.load(model_path)) model.cuda() model.eval() os.makedirs(save_path + test_fold, exist_ok=True) X_test = glob("{}/images/*".format(data_path)) X_test.sort() y_test = glob("{}/masks/*".format(data_path)) y_test.sort() test_transform = Compose([ transforms.Resize(352, 352), # transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), ]) test_dataset = Dataset(X_test, y_test, aug=True, transform=test_transform) test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False, pin_memory=True, drop_last=False) print("TESTING " + test_fold) tp_all = 0 fp_all = 0 fn_all = 0 mean_iou = 0 mean_precision = 0 mean_recall = 0 mean_iou = 0 mean_dice = 0 for i, pack in tqdm.tqdm(enumerate(test_loader, start=1)): image, gt, filename, img = pack name = os.path.splitext(filename[0])[0] ext = os.path.splitext(filename[0])[1] # if(os.path.exists(os.path.join(save_path,test_fold,"v" + str(v),name+"_prv" + str(v) + ext))): # continue gt = gt[0][0] gt = np.asarray(gt, np.float32) res2 = 0 image = image.cuda() if v == 0 or v == 1 or v == 2 or v == 3 or v == 19 or v == 20: res5, res4, res3, res2 = model(image) elif (v == 4 or v == 5 or v == 13 or v == "GALD" or v == 14 or v == 15 or v == 16 or v == 17 or v == 18): res_head_out, res5, res4, res3, res2 = model(image) elif v == 6 or v == 7: res_gald_head_out, res_dual_head_out, res5, res4, res3, res2 = model( image) elif v == 8: res3_head_out, res2_head_out, res5, res4, res3, res2 = model( image) elif v == 12: res2 = model(image) else: print("Not have this version") break # res = res_head_out res = res2 res = F.upsample(res, size=gt.shape, mode="bilinear", align_corners=False) res = res.sigmoid().data.cpu().numpy().squeeze() res = (res - res.min()) / (res.max() - res.min() + 1e-8) if visualize: save_img( os.path.join(save_path, test_fold, "v" + str(v), name + "_prv" + str(v) + ext), res.round() * 255, "cv2", False, ) save_img( os.path.join( save_path, test_fold, "soft_v" + str(v), name + "_soft_prv" + str(v) + ext, ), res * 255, "cv2", False, ) # mask_img = np.asarray(img[0]) + cv2.cvtColor(res.round()*60, cv2.COLOR_GRAY2BGR) mask_img = (np.asarray(img[0]) + np.array(( np.zeros_like(res.round() * 60), res.round() * 60, np.zeros_like(res.round() * 60), )).transpose((1, 2, 0)) + np.array( (gt * 60, np.zeros_like(gt * 60), np.zeros_like( gt * 60))).transpose((1, 2, 0))) mask_img = mask_img[:, :, ::-1] save_img( os.path.join( save_path, test_fold, "mask_v" + str(v), name + "mask_prv" + str(v) + ext, ), mask_img, "cv2", False, ) # skimage.io.imsave("aaa.png", np.asarray(img[0]) + cv2.cvtColor(res.round()*60, cv2.COLOR_GRAY2BGR)) # import sys # sys.exit() # cv2.imwrite(save_path + test_fold + '/' + name + '_gt' + ext, gt*255) # print("cp {} {}".format('Kvasir_fold_new/' + 'fold_' + str(id)+'/images/' + name + ext, save_path + test_fold)) # os.system("cp {} {}".format('/content/Kvasir_fold_new/' + 'fold_' + str(id)+'/images/' + name + ext, save_path + test_fold )) pr = res.round() prs.append(pr) gts.append(gt) tp = np.sum(gt * pr) fp = np.sum(pr) - tp fn = np.sum(gt) - tp tp_all += tp fp_all += fp fn_all += fn mean_precision += precision_m(gt, pr) mean_recall += recall_m(gt, pr) mean_iou += jaccard_m(gt, pr) mean_dice += dice_m(gt, pr) if i % 100 == 0: print(i) precision_all = tp_all / (tp_all + fp_all + K.epsilon()) recall_all = tp_all / (tp_all + fn_all + K.epsilon()) dice_all = 2 * precision_all * recall_all / (precision_all + recall_all) iou_all = (recall_all * precision_all / (recall_all + precision_all - recall_all * precision_all)) # ious.append(iou_all) # precisions.append(precision_all) # recalls.append(recall_all) # dices.append(dice_all) print("scores ver2: {:.3f} {:.3f} {:.3f} {:.3f}".format( iou_all, precision_all, recall_all, dice_all)) mean_precision /= len(test_loader) mean_recall /= len(test_loader) mean_iou /= len(test_loader) mean_dice /= len(test_loader) print("scores ver1: {:.3f} {:.3f} {:.3f} {:.3f}".format( mean_iou, mean_precision, mean_recall, mean_dice)) # print(f"IoU: mean={round(np.mean(ious), 6)}, std={round(np.std(ious), 6)}, var={round(np.var(ious), 6)}") # print(f"dice: mean={round(np.mean(dices), 6)}, std={round(np.std(dices), 6)}, var={round(np.var(dices), 6)}") # print(f"precision: mean={round(np.mean(precisions), 6)}, std={round(np.std(precisions), 6)}, var={round(np.var(precisions), 6)}") # print(f"recall: mean={round(np.mean(recalls), 6)}, std={round(np.std(recalls), 6)}, var={round(np.var(recalls), 6)}") # print(f'iou = {iou_all}, precision = {precision_all}, recall = {recall_all}, dice = {dice_all}') return gts, prs