def load_model(data, model_path, cuda=True): if cuda and not torch.cuda.is_available(): raise Exception("No GPU found, please run without --cuda") unet = UNet() if cuda: unet = unet.cuda() if not cuda: unet.load_state_dict( torch.load(model_path, map_location=lambda storage, loc: storage)) else: unet.load_state_dict(torch.load(model_path)) if cuda: data = Variable(data.cuda()) else: data = Variable(data) data = torch.unsqueeze(data, 0) output = unet(data) if cuda: output = output.cuda() return output
def see_results(n_channels, n_classes, load_weights, dir_img, dir_cmp, savedir, title): # Use GPU or not use_cuda = torch.cuda.is_available() device = torch.device("cuda" if use_cuda else "cpu") # Create the model net = UNet(n_channels, n_classes).to(device) net = torch.nn.DataParallel( net, device_ids=list(range(torch.cuda.device_count()))).to(device) # Load old weights checkpoint = torch.load(load_weights, map_location='cpu') net.load_state_dict(checkpoint['state_dict']) # Load the dataset loader = get_dataloader_show(dir_img, dir_cmp) # If savedir does not exists make folder if not os.path.exists(savedir): os.makedirs(savedir) net.eval() with torch.no_grad(): for (data, gt) in loader: # Use GPU or not data, gt = data.to(device), gt.to(device) # Forward predictions = net(data) save_image(predictions, savedir + title + "_pred.png") save_image(gt, savedir + title + "_gt.png")
def main(): # width_in = 284 # height_in = 284 # width_out = 196 # height_out = 196 # PATH = './unet.pt' # x_train, y_train, x_val, y_val = get_dataset(width_in, height_in, width_out, height_out) # print(x_train.shape, y_train.shape, x_val.shape, y_val.shape) batch_size = 3 epochs = 1 epoch_lapse = 50 threshold = 0.5 learning_rate = 0.01 unet = UNet(in_channel=1, out_channel=2) if use_gpu: unet = unet.cuda() criterion = torch.nn.CrossEntropyLoss() optimizer = torch.optim.SGD(unet.parameters(), lr=0.01, momentum=0.99) if sys.argv[1] == 'train': train(unet, batch_size, epochs, epoch_lapse, threshold, learning_rate, criterion, optimizer, x_train, y_train, x_val, y_val, width_out, height_out) pass else: if use_gpu: unet.load_state_dict(torch.load(PATH)) else: unet.load_state_dict(torch.load(PATH, map_location='cpu')) print(unet.eval())
def lr_find(model: UNet, data_loader, optimizer: Optimizer, criterion, use_gpu, min_lr=0.0001, max_lr=0.1): # Save model and optimizer states to revert model_state = model.state_dict() optimizer_state = optimizer.state_dict() losses = [] lrs = [] scheduler = CyclicExpLR(optimizer, min_lr, max_lr, step_size_up=100, mode='triangular', cycle_momentum=True) model.train() for i, (data, target, class_ids) in enumerate(data_loader): data, target = data, target if use_gpu: data = data.cuda() target = target.cuda() optimizer.zero_grad() output_raw = model(data) # This step is specific for this project output = torch.zeros(output_raw.shape[0], 1, output_raw.shape[2], output_raw.shape[3]) if use_gpu: output = output.cuda() # This step is specific for this project for idx, (raw_o, class_id) in enumerate(zip(output_raw, class_ids)): output[idx] = raw_o[class_id - 1] loss = criterion(output, target) loss.backward() current_lr = optimizer.param_groups[0]['lr'] # Stop if lr stopped increasing if len(lrs) > 0 and current_lr < lrs[-1]: break lrs.append(current_lr) losses.append(loss.item()) optimizer.step() scheduler.step() # Plot in log scale plt.plot(lrs, losses) plt.xscale('log') plt.show() model.load_state_dict(model_state) optimizer.load_state_dict(optimizer_state)
def load_model(key: str): if key == BRAIN_DENOISING_MODEL: checkpoint = torch.load(BRAIN_DENOISING_MODEL) unet = UNet(in_channels=1, n_classes=1, depth=3, wf=8, padding=True, batch_norm=False, up_mode='upconv', grid=True, bias=True) unet.load_state_dict(checkpoint["model_state_dict"]) return unet if key == BRAIN_SEG_MODEL: checkpoint = torch.load(BRAIN_SEG_MODEL) unet = UNet(in_channels=1, n_classes=1, depth=3, wf=8, padding=True, batch_norm=False, up_mode='upconv', grid=False, bias=True) unet.load_state_dict(checkpoint["model_state_dict"]) if key == ABDOM_DENOISING_MODEL: checkpoint = torch.load(ABDOM_DENOISING_MODEL) unet = UNet(in_channels=1, n_classes=1, depth=3, wf=8, padding=True, batch_norm=False, up_mode='upconv', grid=True, bias=True) unet.load_state_dict(checkpoint["model_state_dict"]) if key == ABDOM_SEG_MODEL: checkpoint = torch.load(ABDOM_SEG_MODEL) unet = UNet(in_channels=1, n_classes=1, depth=3, wf=8, padding=True, batch_norm=False, up_mode='upconv', grid=False, bias=True) unet.load_state_dict(checkpoint["model_state_dict"]) return unet
def test(path): """Count the input image""" device = torch.device( 'cuda') if torch.cuda.is_available() else torch.device('cpu') # Image image = np.array(Image.open(path), dtype=np.float32) / 255 image = torch.Tensor(np.transpose(image, (2, 0, 1))).unsqueeze(0) # Ground Truth header = ".".join(path.split('/')[-1].split('.')[:2]) label_path = opt.label_path + header + '.label.png' label = np.array(Image.open(label_path)) if opt.color == 'red': labels = 100.0 * (label[:, :, 0] > 0) else: labels = 100.0 * (label[:, :, 1] > 0) labels = ndimage.gaussian_filter(labels, sigma=(1, 1), order=0) labels = torch.Tensor(labels).unsqueeze(0) if opt.model.find("UNet") != -1: model = UNet(input_filters=3, filters=opt.unet_filters, N=opt.conv).to(device) else: model = FCRN_A(input_filters=3, filters=opt.unet_filters, N=opt.conv).to(device) model = torch.nn.DataParallel(model) if os.path.exists('{}.pth'.format(opt.model)): model.load_state_dict(torch.load('{}.pth'.format(opt.model))) model.eval() image = image.to(device) labels = labels.to(device) out = model(image) predicted_counts = torch.sum(out).item() / 100 real_counts = torch.sum(labels).item() / 100 print(predicted_counts, real_counts) label = np.zeros((image.shape[2], image.shape[2], 3)) if opt.color == 'red': label[:, :, 0] = out[0][0].cpu().detach().numpy() else: label[:, :, 1] = out[0][0].cpu().detach().numpy() imageio.imwrite('example/test_results/density_map_{}.png'.format(header), label) return header, predicted_counts, real_counts
def load_model(): checkpoint = get_model() model = UNet( backbone="mobilenetv2", num_classes=2, pretrained_backbone=None ) trained_dict = torch.load(checkpoint, map_location="cpu")['state_dict'] model.load_state_dict(trained_dict, strict=False) model.eval() print("model is loaded") return model
def main(): opt = parser.parse_args() print(torch.__version__) print(opt) enc_layers = [1, 2, 2, 4] dec_layers = [1, 1, 1, 1] number_of_channels = [ int(8 * 2**i) for i in range(1, 1 + len(enc_layers)) ] #[16,32,64,128] model = UNet(depth=len(enc_layers), encoder_layers=enc_layers, decoder_layers=dec_layers, number_of_channels=number_of_channels, number_of_outputs=3) s = torch.load(os.path.join(opt.models_path, opt.name, opt.name + 'best_model.pth'), map_location='cpu') new_state_dict = OrderedDict() for k, v in s['model'].state_dict().items(): name = k[7:] # remove 'module' word in the beginning of keys. new_state_dict[name] = v model.load_state_dict(new_state_dict) model.eval() x = torch.randn(1, 4, opt.input_size[0], opt.input_size[1], opt.input_size[2], requires_grad=True) register_op("group_norm", group_norm_symbolic, "", 10) torch_out = torch.onnx.export( model, # model being run [ x, ], # model input (or a tuple for multiple inputs) os.path.join( opt.models_path, opt.name, opt.name + ".onnx" ), # where to save the model (can be a file or file-like object) export_params=True, verbose= True, # store the trained parameter weights inside the model file opset_version=10)
def test(args): """ Test some data from trained UNet """ image = load_test_image(args.test_image) # 1 c w h net = UNet(in_channels=3, out_channels=5) if args.cuda: net = net.cuda() image = image.cuda() print('Loading model param from {}'.format(args.model_state_dict)) net.load_state_dict(torch.load(args.model_state_dict)) net.eval() print('Predicting for {}...'.format(args.test_image)) ys_pred = net(image) # 1 ch w h colors = [] with open(args.mask_json_path, 'r', encoding='utf-8') as mask: print('Reading mask colors list from {}'.format(args.mask_json_path)) colors = json.loads(mask.read()) colors = [tuple(c) for c in colors] print('Mask colors: {}'.format(colors)) ys_pred = ys_pred.cpu().detach().numpy()[0] ys_pred[ys_pred < 0.5] = 0 ys_pred[ys_pred >= 0.5] = 1 ys_pred = ys_pred.astype(np.int) image_w = ys_pred.shape[1] image_h = ys_pred.shape[2] out_image = np.zeros((image_w, image_h, 3)) for w in range(image_w): for h in range(image_h): for ch in range(ys_pred.shape[0]): if ys_pred[ch][w][h] == 1: out_image[w][h][0] = colors[ch][0] out_image[w][h][1] = colors[ch][1] out_image[w][h][2] = colors[ch][2] out_image = out_image.astype(np.uint8) # w h c out_image = out_image.transpose((1, 0, 2)) # h w c out_image = Image.fromarray(out_image) out_image.save(args.test_save_path) print('Segmentation result has been saved to {}'.format( args.test_save_path))
class modelLoader(): def __init__(self, model_folder, in_chnl, out_chnl, checkpoint=None): self.checkpoint = checkpoint self.model_folder = model_folder self.in_chnl = in_chnl self.out_chnl = out_chnl self.model = None def set_model(self): self.model = UNet(self.in_chnl, self.out_chnl).to(device) PATH = "./{}/checkpoint_{}.pth".format(self.model_folder, self.checkpoint) self.model.load_state_dict(torch.load(PATH)) self.model.eval() def model_infos(self): print( "Checkpoint: {}, Model folder: ./{}, Input channel: {}, Output channel: {}" .format(self.checkpoint, self.model_folder, self.in_chnl, self.out_chnl))
def test(args): model = UNet(n_channels=125, n_classes=10).to(device) model.load_state_dict( torch.load(args.ckpt % args.num_epochs, map_location='cpu')) liver_dataset = LiverDataset('data', transform=x_transform, target_transform=y_transform, train=False) dataloaders = DataLoader(liver_dataset, batch_size=1) model.eval() with torch.no_grad(): for x, y in dataloaders: x = x.permute(0, 2, 1, 3) print(x.shape) x = x.float() label = y.float() x = x.to(device) label = label.to(device) outputs = model(x) # print(outputs.shape) label = label.squeeze(1)
def test(n_class=1, in_channel=1, load=False, img_size=None, directory='../Data/train/'): global original, dataset, model original = Segmentation(directory, 'training.json', ToTensor()) if img_size is None: dataset = Segmentation(directory, 'test.json', ToTensor()) else: dataset = Segmentation(directory, 'test.json', T.Compose([\ RandomCrop((img_size, img_size)), ToTensor()])) model = UNet(n_class=n_class, in_channel=in_channel) if load: filename = "unet.pth" map_location = 'cuda:0' if torch.cuda.is_available() else 'cpu' try: checkpoint = torch.load(filename, map_location=map_location) model.load_state_dict(checkpoint['state_dict']) print("Loaded saved model") except: print("Unable to load saved model")
def run_inference(args): model = UNet(input_channels=3, num_classes=3) model.load_state_dict(torch.load(args.model_path, map_location='cpu'), strict=False) print('Log: Loaded pretrained {}'.format(args.model_path)) model.eval() # annus/Desktop/palsar/ test_image_path = '/home/annus/Desktop/palsar/palsar_dataset_full/palsar_dataset/palsar_{}_region_{}.tif'.format( args.year, args.region) test_label_path = '/home/annus/Desktop/palsar/palsar_dataset_full/palsar_dataset/fnf_{}_region_{}.tif'.format( args.year, args.region) inference_loader = get_inference_loader(image_path=test_image_path, label_path=test_label_path, model_input_size=128, num_classes=4, one_hot=True, batch_size=args.bs, num_workers=4) # we need to fill our new generated test image generated_map = np.empty(shape=inference_loader.dataset.get_image_size()) weights = torch.Tensor([1, 1, 2]) focal_criterion = FocalLoss2d(weight=weights) un_confusion_meter = tnt.meter.ConfusionMeter(2, normalized=False) confusion_meter = tnt.meter.ConfusionMeter(2, normalized=True) total_correct, total_examples = 0, 0 net_loss = [] for idx, data in enumerate(inference_loader): coordinates, test_x, label = data['coordinates'].tolist( ), data['input'], data['label'] out_x, softmaxed = model.forward(test_x) pred = torch.argmax(softmaxed, dim=1) not_one_hot_target = torch.argmax(label, dim=1) # convert to binary classes # 0-> noise, 1-> forest, 2-> non-forest, 3-> water pred[pred == 0] = 2 pred[pred == 3] = 2 not_one_hot_target[not_one_hot_target == 0] = 2 not_one_hot_target[not_one_hot_target == 3] = 2 # now convert 1, 2 to 0, 1 pred -= 1 not_one_hot_target -= 1 pred_numpy = pred.numpy().transpose(1, 2, 0) for k in range(test_x.shape[0]): x, x_, y, y_ = coordinates[k] generated_map[x:x_, y:y_] = pred_numpy[:, :, k] loss = focal_criterion( softmaxed, not_one_hot_target) # dice_criterion(softmaxed, label) # accurate = (pred == not_one_hot_target).sum().item() numerator = float(accurate) denominator = float( pred.view(-1).size(0)) # test_x.size(0) * dimension ** 2) total_correct += numerator total_examples += denominator net_loss.append(loss.item()) un_confusion_meter.add(predicted=pred.view(-1), target=not_one_hot_target.view(-1)) confusion_meter.add(predicted=pred.view(-1), target=not_one_hot_target.view(-1)) # if idx % 5 == 0: accuracy = float(numerator) * 100 / denominator print( '{}, {} -> ({}/{}) output size = {}, loss = {}, accuracy = {}/{} = {:.2f}%' .format(args.year, args.region, idx, len(inference_loader), out_x.size(), loss.item(), numerator, denominator, accuracy)) ################################# mean_accuracy = total_correct * 100 / total_examples mean_loss = np.asarray(net_loss).mean() print('$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$') print('log: test:: total loss = {:.5f}, total accuracy = {:.5f}%'.format( mean_loss, mean_accuracy)) print('$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$') print('---> Confusion Matrix:') print(confusion_meter.value()) # class_names = ['background/clutter', 'buildings', 'trees', 'cars', # 'low_vegetation', 'impervious_surfaces', 'noise'] with open('normalized.pkl', 'wb') as this: pkl.dump(confusion_meter.value(), this, protocol=pkl.HIGHEST_PROTOCOL) with open('un_normalized.pkl', 'wb') as this: pkl.dump(un_confusion_meter.value(), this, protocol=pkl.HIGHEST_PROTOCOL) # save_path = 'generated_maps/generated_{}_{}.npy'.format(args.year, args.region) save_path = '/home/annus/Desktop/palsar/generated_maps/using_separate_models/generated_{}_{}.npy'.format( args.year, args.region) np.save(save_path, generated_map) #########################################################################################3 inference_loader.dataset.clear_mem() pass
if not os.path.exists(save_dir): os.mkdir(save_dir) if __name__ == '__main__': # model selection print('===> Building model') model = UNet(input_channels=2, image_channels=1) pre_model = DnCNN(image_channels=1) pre_model = torch.load(os.path.join(args.load_model_dir, 'model.pth')) initial_epoch = findLastCheckpoint(save_dir=save_dir)# load the last model in matconvnet style # initial_epoch = 150 if initial_epoch > 0: print('resuming by loading epoch %03d' % initial_epoch) model.load_state_dict(torch.load(os.path.join(save_dir, 'model_%03d.pth' % initial_epoch))) # model = torch.load(os.path.join(save_dir, 'model_%03d.pth' % initial_epoch)) model.train() pre_model.eval() # criterion = nn.MSELoss(reduction = 'sum') # PyTorch 0.4.1 # criterion = sum_squared_error() # criterion = nn.MSELoss() criterion = Intensity_loss() cri_chk = nn.MSELoss() if cuda: model = model.cuda() pre_model = pre_model.cuda() # device_ids = [0]
import numpy as np import io import cv2 import os from werkzeug.utils import secure_filename import base64 from model import UNet import json from PIL import Image import re MODEL_PATH = "model/0.pth" UPLOAD_PATH = "../backend/uploads/" model = UNet() model.load_state_dict(torch.load(MODEL_PATH, map_location="cpu")) app = flask.Flask(__name__) app.secret_key = "super secret key" cors = CORS(app) app.config["UPLOAD_FOLDER"] = UPLOAD_PATH app.config['CORS_HEADERS'] = 'Content-Type' ctr = 0 @app.route("/", methods=["GET", "POST"]) @cross_origin() def upload_file(): global ctr if request.method == "POST":
def run_inference(args): model = UNet(topology=args.model_topology, input_channels=len(args.bands), num_classes=len(args.classes)) model.load_state_dict(torch.load(args.model_path, map_location='cpu'), strict=False) print('Log: Loaded pretrained {}'.format(args.model_path)) model.eval() if args.cuda: print('log: Using GPU') model.cuda(device=args.device) # all_districts = ["abbottabad", "battagram", "buner", "chitral", "hangu", "haripur", "karak", "kohat", "kohistan", "lower_dir", "malakand", "mansehra", # "nowshehra", "shangla", "swat", "tor_ghar", "upper_dir"] all_districts = ["abbottabad"] # years = [2014, 2016, 2017, 2018, 2019, 2020] years = [2016] # change this to do this for all the images in that directory for district in all_districts: for year in years: print("(LOG): On District: {} @ Year: {}".format(district, year)) # test_image_path = os.path.join(args.data_path, 'landsat8_4326_30_{}_region_{}.tif'.format(year, district)) test_image_path = os.path.join(args.data_path, 'landsat8_{}_region_{}.tif'.format( year, district)) #added(nauman) inference_loader, adjustment_mask = get_inference_loader( rasterized_shapefiles_path=args.rasterized_shapefiles_path, district=district, image_path=test_image_path, model_input_size=128, bands=args.bands, num_classes=len(args.classes), batch_size=args.bs, num_workers=4) # inference_loader = get_inference_loader(rasterized_shapefiles_path=args.rasterized_shapefiles_path, district=district, # image_path=test_image_path, model_input_size=128, bands=args.bands, # num_classes=len(args.classes), batch_size=args.bs, num_workers=4) # we need to fill our new generated test image generated_map = np.empty( shape=inference_loader.dataset.get_image_size()) for idx, data in enumerate(inference_loader): coordinates, test_x = data['coordinates'].tolist( ), data['input'] test_x = test_x.cuda( device=args.device) if args.cuda else test_x out_x, softmaxed = model.forward(test_x) pred = torch.argmax(softmaxed, dim=1) pred_numpy = pred.cpu().numpy().transpose(1, 2, 0) if idx % 5 == 0: print('LOG: on {} of {}'.format(idx, len(inference_loader))) for k in range(test_x.shape[0]): x, x_, y, y_ = coordinates[k] generated_map[x:x_, y:y_] = pred_numpy[:, :, k] # adjust the inferred map generated_map += 1 # to make forest pixels: 2, non-forest pixels: 1, null pixels: 0 generated_map = np.multiply(generated_map, adjustment_mask) # save generated map as png image, not numpy array forest_map_rband = np.zeros_like(generated_map) forest_map_gband = np.zeros_like(generated_map) forest_map_bband = np.zeros_like(generated_map) forest_map_gband[generated_map == FOREST_LABEL] = 255 forest_map_rband[generated_map == NON_FOREST_LABEL] = 255 forest_map_for_visualization = np.dstack( [forest_map_rband, forest_map_gband, forest_map_bband]).astype(np.uint8) save_this_map_path = os.path.join( args.dest, '{}_{}_inferred_map.png'.format(district, year)) matimg.imsave(save_this_map_path, forest_map_for_visualization) print('Saved: {} @ {}'.format(save_this_map_path, forest_map_for_visualization.shape))
model = UNet(in_channels=n_channels, n_classes=n_classes) elif args.model == 'SegNet': model = SegNet(in_channels=n_channels, n_classes=n_classes) elif args.model == 'DenseNet': model = DenseNet(in_channels=n_channels, n_classes=n_classes) else: print("wrong model : must be UNet, SegNet, or DenseNet") raise SystemExit summary(model, input_size=(n_channels, args.height, args.width), device='cpu') model.load_state_dict(torch.load(args.model_path)) adversarial_examples = DAG_Attack(model, test_dataset, args) if args.attack_path is None: adversarial_path = 'data/' + args.model + '_' + args.attacks + '.pickle' else: adversarial_path = args.attack_path # save adversarial examples([adversarial examples, labels]) with open(adversarial_path, 'wb') as fp: pickle.dump(adversarial_examples, fp)
optimizer = optim.Adam(model.parameters(), lr=opt.learning_rate, weight_decay=opt.weight_decay) criterion = nn.BCELoss() vis = Visualizer(env=opt.env) if opt.is_cuda: model.cuda() criterion.cuda() if opt.n_gpu > 1: model = nn.DataParallel(model) run(model, train_loader, val_loader, criterion, vis) else: if opt.is_cuda: model.cuda() if opt.n_gpu > 1: model = nn.DataParallel(model) test_loader = get_test_loader(batch_size=20, shuffle=True, num_workers=opt.num_workers, pin_memory=opt.pin_memory) # load the model and run test model.load_state_dict( torch.load( os.path.join(opt.checkpoint_dir, 'RSNA_UNet_0.895_09210122'))) img_ids, images, pred_masks = run_test(model, test_loader) save_pred_result(img_ids, images, pred_masks)
# %% n_sticks = 40 # %% # Model # model = UNet().cuda() model = UNet() model.eval() model_dir = Path('model') model_path = model_dir.joinpath('model.pth') param = torch.load(model_path) model.load_state_dict(param) # %% # Dataset for inference # test_dataset = SSSDataset(train=False, n_sticks=n_sticks, data_size=16) test_dataset = DataLoaderInstanceSegmentation(train = False) # test_dataloader = DataLoader(test_dataset, batch_size=16, # shuffle=False, num_workers=0, # pin_memory=True) test_dataloader = DataLoader(test_dataset, batch_size=20, shuffle=False, num_workers=0, pin_memory=True) # %%
def main(args): writer = SummaryWriter(os.path.join('./logs')) # torch.backends.cudnn.benchmark = True if not os.path.isdir(args.checkpoint_dir): os.mkdir(args.checkpoint_dir) device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") print('[MODEL] CUDA DEVICE : {}'.format(device)) # TODO DEFINE TRAIN AND TEST TRANSFORMS train_tf = None test_tf = None # Channel wise mean calculated on adobe240-fps training dataset mean = [0.429, 0.431, 0.397] std = [1, 1, 1] normalize = transforms.Normalize(mean=mean, std=std) transform = transforms.Compose([transforms.ToTensor(), normalize]) test_valid = 'validation' if args.valid else 'test' train_data = BlurDataset(os.path.join(args.dataset_root, 'train'), seq_len=args.sequence_length, tau=args.num_frame_blur, delta=5, transform=train_tf) test_data = BlurDataset(os.path.join(args.dataset_root, test_valid), seq_len=args.sequence_length, tau=args.num_frame_blur, delta=5, transform=train_tf) train_loader = DataLoader(train_data, batch_size=args.train_batch_size, shuffle=True) test_loader = DataLoader(test_data, batch_size=args.test_batch_size, shuffle=False) # TODO IMPORT YOUR CUSTOM MODEL model = UNet(3, 3, device, decode_mode=args.decode_mode) if args.checkpoint: store_dict = torch.load(args.checkpoint) try: model.load_state_dict(store_dict['state_dict']) except KeyError: model.load_state_dict(store_dict) if args.train_continue: store_dict = torch.load(args.checkpoint) model.load_state_dict(store_dict['state_dict']) else: store_dict = {'loss': [], 'valLoss': [], 'valPSNR': [], 'epoch': -1} model.to(device) model.train(True) # model = nn.DataParallel(model) # TODO DEFINE MORE CRITERIA # input(True if device == torch.device('cuda:0') else False) criterion = { 'MSE': nn.MSELoss(), 'L1': nn.L1Loss(), # 'Perceptual': PerceptualLoss(model='net-lin', net='vgg', dataparallel=True, # use_gpu=True if device == torch.device('cuda:0') else False) } criterion_w = {'MSE': 1.0, 'L1': 10.0, 'Perceptual': 10.0} # Define optimizers # optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9,weight_decay=5e-4) optimizer = optim.Adam(model.parameters(), lr=args.init_learning_rate) # Define lr scheduler scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=args.milestones, gamma=0.1) # best_acc = 0.0 # start = time.time() cLoss = store_dict['loss'] valLoss = store_dict['valLoss'] valPSNR = store_dict['valPSNR'] checkpoint_counter = 0 loss_tracker = {} loss_tracker_test = {} psnr_old = 0.0 dssim_old = 0.0 for epoch in range(1, 10 * args.epochs): # loop over the dataset multiple times # Append and reset cLoss.append([]) valLoss.append([]) valPSNR.append([]) running_loss = 0 # Increment scheduler count scheduler.step() tqdm_loader = tqdm(range(len(train_loader)), ncols=150) loss = 0.0 psnr_ = 0.0 dssim_ = 0.0 loss_tracker = {} for loss_fn in criterion.keys(): loss_tracker[loss_fn] = 0.0 # Train model.train(True) total_steps = 0.01 total_steps_test = 0.01 '''for train_idx, data in enumerate(train_loader, 1): loss = 0.0 blur_data, sharpe_data = data #import pdb; pdb.set_trace() # input(sharpe_data.shape) #import pdb; pdb.set_trace() interp_idx = int(math.ceil((args.num_frame_blur/2) - 0.49)) #input(interp_idx) if args.decode_mode == 'interp': sharpe_data = sharpe_data[:, :, 1::2, :, :] elif args.decode_mode == 'deblur': sharpe_data = sharpe_data[:, :, 0::2, :, :] else: #print('\nBoth\n') sharpe_data = sharpe_data #print(sharpe_data.shape) #input(blur_data.shape) blur_data = blur_data.to(device)[:, :, :, :352, :].permute(0, 1, 2, 4, 3) try: sharpe_data = sharpe_data.squeeze().to(device)[:, :, :, :352, :].permute(0, 1, 2, 4, 3) except: sharpe_data = sharpe_data.squeeze(3).to(device)[:, :, :, :352, :].permute(0, 1, 2, 4, 3) # clear gradient optimizer.zero_grad() # forward pass sharpe_out = model(blur_data) # import pdb; pdb.set_trace() # input(sharpe_out.shape) # compute losses # import pdb; # pdb.set_trace() sharpe_out = sharpe_out.permute(0, 2, 1, 3, 4) B, C, S, Fx, Fy = sharpe_out.shape for loss_fn in criterion.keys(): loss_tmp = 0.0 if loss_fn == 'Perceptual': for bidx in range(B): loss_tmp += criterion_w[loss_fn] * \ criterion[loss_fn](sharpe_out[bidx].permute(1, 0, 2, 3), sharpe_data[bidx].permute(1, 0, 2, 3)).sum() # loss_tmp /= B else: loss_tmp = criterion_w[loss_fn] * \ criterion[loss_fn](sharpe_out, sharpe_data) # try: # import pdb; pdb.set_trace() loss += loss_tmp # if # except : try: loss_tracker[loss_fn] += loss_tmp.item() except KeyError: loss_tracker[loss_fn] = loss_tmp.item() # Backpropagate loss.backward() optimizer.step() # statistics # import pdb; pdb.set_trace() sharpe_out = sharpe_out.detach().cpu().numpy() sharpe_data = sharpe_data.cpu().numpy() for sidx in range(S): for bidx in range(B): psnr_ += psnr(sharpe_out[bidx, :, sidx, :, :], sharpe_data[bidx, :, sidx, :, :]) #, peak=1.0) """dssim_ += dssim(np.moveaxis(sharpe_out[bidx, :, sidx, :, :], 0, 2), np.moveaxis(sharpe_data[bidx, :, sidx, :, :], 0, 2) )""" """sharpe_out = sharpe_out.reshape(-1,3, sx, sy).detach().cpu().numpy() sharpe_data = sharpe_data.reshape(-1, 3, sx, sy).cpu().numpy() for idx in range(sharpe_out.shape[0]): # import pdb; pdb.set_trace() psnr_ += psnr(sharpe_data[idx], sharpe_out[idx]) dssim_ += dssim(np.swapaxes(sharpe_data[idx], 2, 0), np.swapaxes(sharpe_out[idx], 2, 0))""" # psnr_ /= sharpe_out.shape[0] # dssim_ /= sharpe_out.shape[0] running_loss += loss.item() loss_str = '' total_steps += B*S for key in loss_tracker.keys(): loss_str += ' {0} : {1:6.4f} '.format(key, 1.0*loss_tracker[key] / total_steps) # set display info if train_idx % 5 == 0: tqdm_loader.set_description(('\r[Training] [Ep {0:6d}] loss: {1:6.4f} PSNR: {2:6.4f} SSIM: {3:6.4f} '.format (epoch, running_loss / total_steps, psnr_ / total_steps, dssim_ / total_steps) + loss_str )) tqdm_loader.update(5) tqdm_loader.close()''' # Validation running_loss_test = 0.0 psnr_test = 0.0 dssim_test = 0.0 # print('len', len(test_loader)) tqdm_loader_test = tqdm(range(len(test_loader)), ncols=150) # import pdb; pdb.set_trace() loss_tracker_test = {} for loss_fn in criterion.keys(): loss_tracker_test[loss_fn] = 0.0 with torch.no_grad(): model.eval() total_steps_test = 0.0 for test_idx, data in enumerate(test_loader, 1): loss = 0.0 blur_data, sharpe_data = data interp_idx = int(math.ceil((args.num_frame_blur / 2) - 0.49)) # input(interp_idx) if args.decode_mode == 'interp': sharpe_data = sharpe_data[:, :, 1::2, :, :] elif args.decode_mode == 'deblur': sharpe_data = sharpe_data[:, :, 0::2, :, :] else: # print('\nBoth\n') sharpe_data = sharpe_data # print(sharpe_data.shape) # input(blur_data.shape) blur_data = blur_data.to(device)[:, :, :, :352, :].permute( 0, 1, 2, 4, 3) try: sharpe_data = sharpe_data.squeeze().to( device)[:, :, :, :352, :].permute(0, 1, 2, 4, 3) except: sharpe_data = sharpe_data.squeeze(3).to( device)[:, :, :, :352, :].permute(0, 1, 2, 4, 3) # clear gradient optimizer.zero_grad() # forward pass sharpe_out = model(blur_data) # import pdb; pdb.set_trace() # input(sharpe_out.shape) # compute losses sharpe_out = sharpe_out.permute(0, 2, 1, 3, 4) B, C, S, Fx, Fy = sharpe_out.shape for loss_fn in criterion.keys(): loss_tmp = 0.0 if loss_fn == 'Perceptual': for bidx in range(B): loss_tmp += criterion_w[loss_fn] * \ criterion[loss_fn](sharpe_out[bidx].permute(1, 0, 2, 3), sharpe_data[bidx].permute(1, 0, 2, 3)).sum() # loss_tmp /= B else: loss_tmp = criterion_w[loss_fn] * \ criterion[loss_fn](sharpe_out, sharpe_data) loss += loss_tmp try: loss_tracker_test[loss_fn] += loss_tmp.item() except KeyError: loss_tracker_test[loss_fn] = loss_tmp.item() if ((test_idx % args.progress_iter) == args.progress_iter - 1): itr = test_idx + epoch * len(test_loader) # itr_train writer.add_scalars( 'Loss', { 'trainLoss': running_loss / total_steps, 'validationLoss': running_loss_test / total_steps_test }, itr) writer.add_scalar('Train PSNR', psnr_ / total_steps, itr) writer.add_scalar('Test PSNR', psnr_test / total_steps_test, itr) # import pdb; pdb.set_trace() # writer.add_image('Validation', sharpe_out.permute(0, 2, 3, 1), itr) # statistics sharpe_out = sharpe_out.detach().cpu().numpy() sharpe_data = sharpe_data.cpu().numpy() for sidx in range(S): for bidx in range(B): psnr_test += psnr( sharpe_out[bidx, :, sidx, :, :], sharpe_data[bidx, :, sidx, :, :]) #, peak=1.0) dssim_test += dssim( np.moveaxis(sharpe_out[bidx, :, sidx, :, :], 0, 2), np.moveaxis(sharpe_data[bidx, :, sidx, :, :], 0, 2)) #,range=1.0 ) running_loss_test += loss.item() total_steps_test += B * S loss_str = '' for key in loss_tracker.keys(): loss_str += ' {0} : {1:6.4f} '.format( key, 1.0 * loss_tracker_test[key] / total_steps_test) # set display info tqdm_loader_test.set_description(( '\r[Test ] [Ep {0:6d}] loss: {1:6.4f} PSNR: {2:6.4f} SSIM: {3:6.4f} ' .format(epoch, running_loss_test / total_steps_test, psnr_test / total_steps_test, dssim_test / total_steps_test) + loss_str)) tqdm_loader_test.update(1) tqdm_loader_test.close() # save model if psnr_old < (psnr_test / total_steps_test): if epoch != 1: os.remove( os.path.join( args.checkpoint_dir, 'epoch-{}-test-psnr-{}-ssim-{}.ckpt'.format( epoch_old, str(round(psnr_old, 4)).replace('.', 'pt'), str(round(dssim_old, 4)).replace('.', 'pt')))) epoch_old = epoch psnr_old = psnr_test / total_steps_test dssim_old = dssim_test / total_steps_test checkpoint_dict = { 'epoch': epoch_old, 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict(), 'train_psnr': psnr_ / total_steps, 'train_dssim': dssim_ / total_steps, 'train_mse': loss_tracker['MSE'] / total_steps, 'train_l1': loss_tracker['L1'] / total_steps, # 'train_percp': loss_tracker['Perceptual'] / total_steps, 'test_psnr': psnr_old, 'test_dssim': dssim_old, 'test_mse': loss_tracker_test['MSE'] / total_steps_test, 'test_l1': loss_tracker_test['L1'] / total_steps_test, # 'test_percp': loss_tracker_test['Perceptual'] / total_steps_test, } torch.save( checkpoint_dict, os.path.join( args.checkpoint_dir, 'epoch-{}-test-psnr-{}-ssim-{}.ckpt'.format( epoch_old, str(round(psnr_old, 4)).replace('.', 'pt'), str(round(dssim_old, 4)).replace('.', 'pt')))) # if epoch % args.checkpoint_epoch == 0: # torch.save(model.state_dict(),args.checkpoint_dir + str(int(epoch/100))+".ckpt") return None
class Solver: def __init__(self, config, train_loader=None, val_loader=None, test_loader=None): self.cfg = config self.device = torch.device( "cuda" if torch.cuda.is_available() else "cpu") self.n_gpus = self.cfg.n_gpus if self.cfg.mode in ['train', 'test']: self.train_loader = train_loader self.val_loader = val_loader else: self.test_loader = test_loader # Build model self.build_model() if self.cfg.resume: self.load_pre_model() else: self.start_epoch = 0 # Trigger Tensorboard Logger if self.cfg.use_tensorboard: try: from tensorboardX import SummaryWriter self.writer = SummaryWriter() except ImportError: print( '=> There is no module named tensorboardX, tensorboard disabled' ) self.cfg.use_tensorboard = False def train_val(self): # Build record objs self.build_recorder() iter_per_epoch = len( self.train_loader.dataset) // self.cfg.train_batch_size if len(self.train_loader.dataset) % self.cfg.train_batch_size != 0: iter_per_epoch += 1 for epoch in range(self.start_epoch, self.start_epoch + self.cfg.n_epochs): self.model.train() self.train_time.reset() self.train_loss.reset() self.train_cls_acc.reset() self.train_pix_acc.reset() self.train_mIoU.reset() for i, (image, label) in enumerate(self.train_loader): start_time = time.time() image_var = image.to(self.device) label_var = label.to(self.device) output = self.model(image_var) loss = self.criterion(output, label_var) self.optim.zero_grad() loss.backward() self.optim.step() end_time = time.time() self.train_time.update(end_time - start_time) self.train_loss.update(loss.item()) if self.cfg.task == 'cls': # Record classification accuracy cls_acc = cal_acc(output, label_var) # Update recorder self.train_cls_acc.update(cls_acc.item()) if (i + 1) % self.cfg.log_step == 0: print( 'Epoch[{0}][{1}/{2}]\t' 'Time {train_time.val:.3f} ({train_time.avg:.3f})\t' 'Loss {train_loss.val:.4f} ({train_loss.avg:.4f})\t' 'Accuracy {train_cls_acc.val:.4f} ({train_cls_acc.avg:.4f})' .format(epoch + 1, i + 1, iter_per_epoch, train_time=self.train_time, train_loss=self.train_loss, train_cls_acc=self.train_cls_acc)) if self.cfg.use_tensorboard: self.writer.add_scalar('train/loss', loss.item(), epoch * iter_per_epoch + i) self.writer.add_scalar('train/accuracy', cls_acc.item(), epoch * iter_per_epoch + i) elif self.cfg.task == 'seg': # Record mIoU and pixel-wise accuracy pix_acc = cal_pixel_acc(output, label_var) mIoU = cal_mIoU(output, label_var)[-1] mIoU = torch.mean(mIoU) # Update recorders self.train_pix_acc.update(pix_acc.item()) self.train_mIoU.update(mIoU.item()) if (i + 1) % self.cfg.log_step == 0: print( 'Epoch[{0}][{1}/{2}]\t' 'Time {train_time.val:.3f} ({train_time.avg:.3f})\t' 'Loss {train_loss.val:.4f} ({train_loss.avg:.4f})\t' 'Pixel-Acc {train_pix_acc.val:.4f} ({train_pix_acc.avg:.4f})\t' 'mIoU {train_mIoU.val:.4f} ({train_mIoU.avg:.4f})'. format(epoch + 1, i + 1, iter_per_epoch, train_time=self.train_time, train_loss=self.train_loss, train_pix_acc=self.train_pix_acc, train_mIoU=self.train_mIoU)) if self.cfg.use_tensorboard: self.writer.add_scalar('train/loss', loss.item(), epoch * iter_per_epoch + i) self.writer.add_scalar('train/pix_acc', pix_acc.item(), epoch * iter_per_epoch + i) self.writer.add_scalar('train/mIoU', mIoU.item(), epoch * iter_per_epoch + i) #FIXME currently test validation code #if (i + 1) % 100 == 0: if (epoch + 1) % self.cfg.val_step == 0: self.validate(epoch) # Close logging self.writer.close() def validate(self, epoch): """ Validate with validation dataset """ self.model.eval() self.val_time.reset() self.val_loss.reset() self.val_cls_acc.reset() self.val_mIoU.reset() self.val_pix_acc.reset() iter_per_epoch = len( self.val_loader.dataset) // self.cfg.val_batch_size if len(self.val_loader.dataset) % self.cfg.val_batch_size != 0: iter_per_epoch += 1 for i, (image, label) in enumerate(self.val_loader): start_time = time.time() image_var = image.to(self.device) label_var = label.to(self.device) output = self.model(image_var) loss = self.criterion(output, label_var) end_time = time.time() self.val_time.update(end_time - start_time) self.val_loss.update(loss.item()) if self.cfg.task == 'cls': # Record classification accuracy cls_acc = cal_acc(output, label_var) # Update recorder self.val_cls_acc.update(cls_acc.item()) if (i + 1) % self.cfg.log_step == 0: print( 'Epoch[{0}][{1}/{2}]\t' 'Time {val_time.val:.3f} ({val_time.avg:.3f})\t' 'Loss {val_loss.val:.4f} ({val_loss.avg:.4f})\t' 'Accuracy {val_cls_acc.val:.4f} ({val_cls_acc.avg:.4f})' .format(epoch + 1, i + 1, iter_per_epoch, val_time=self.val_time, val_loss=self.val_loss, val_cls_acc=self.val_cls_acc)) if self.cfg.use_tensorboard: self.writer.add_scalar('val/loss', loss.item(), epoch * iter_per_epoch + i) self.writer.add_scalar('val/accuracy', cls_acc.item(), epoch * iter_per_epoch + i) elif self.cfg.task == 'seg': # Record mIoU and pixel-wise accuracy pix_acc = cal_pixel_acc(output, label_var) mIoU = cal_mIoU(output, label_var)[-1] mIoU = torch.mean(mIoU) # Update recorders self.val_pix_acc.update(pix_acc.item()) self.val_mIoU.update(mIoU.item()) if (i + 1) % self.cfg.log_step == 0: print( ' ##### Validation\t' 'Epoch[{0}][{1}/{2}]\t' 'Time {val_time.val:.3f} ({val_time.avg:.3f})\t' 'Loss {val_loss.val:.4f} ({val_loss.avg:.4f})\t' 'Pixel-Acc {val_pix_acc.val:.4f} ({val_pix_acc.avg:.4f})\t' 'mIoU {val_mIoU.val:.4f} ({val_mIoU.avg:.4f})'.format( epoch + 1, i + 1, iter_per_epoch, val_time=self.val_time, val_loss=self.val_loss, val_pix_acc=self.val_pix_acc, val_mIoU=self.val_mIoU)) if self.cfg.use_tensorboard: self.writer.add_scalar('val/loss', loss.item(), epoch * iter_per_epoch + i) self.writer.add_scalar('val/pix_acc', pix_acc.item(), epoch * iter_per_epoch + i) self.writer.add_scalar('val/mIoU', mIoU.item(), epoch * iter_per_epoch + i) if self.cfg.task == 'cls': if (epoch + 1) % self.cfg.model_save_epoch == 0: state = { 'epoch': epoch + 1, 'state_dict': self.model.state_dict(), 'optim': self.optim.state_dict() } if self.best_cls < self.val_cls_acc.avg: self.best_cls = self.val_cls_acc.avg torch.save( state, './model/cls_model_' + str(epoch + 1) + '_' + str(self.val_cls_acc.avg)[0:5] + '.pth') elif self.cfg.task == 'seg': # Save segmentation samples and model if (epoch + 1) % self.cfg.sample_save_epoch == 0: pred = torch.argmax(output, dim=1) save_image(image, './sample/ori_' + str(epoch + 1) + '.png') save_image(label.unsqueeze(1), './sample/true_' + str(epoch + 1) + '.png') save_image(pred.cpu().unsqueeze(1), './sample/pred_' + str(epoch + 1) + '.png') if (epoch + 1) % self.cfg.model_save_epoch == 0: state = { 'epoch': epoch + 1, 'state_dict': self.model.state_dict(), 'optim': self.optim.state_dict() } if self.best_seg < self.val_pix_acc.avg: self.best_seg = self.val_pix_acc.avg torch.save( state, './model/seg_model_' + str(epoch + 1) + '_' + str(self.val_pix_acc.avg)[0:5] + '.pth') if self.cfg.use_tensorboard: image = make_grid(image) label = make_grid(label.unsqueeze(1)) pred = make_grid(pred.cpu().unqueeze(1)) self.writer.add_image('Origianl', image, epoch + 1) self.writer.add_image('Labels', label, epoch + 1) self.writer.add_image('Predictions', pred, epoch + 1) def build_model(self): """ Rough """ if self.cfg.task == 'cls': self.model = BinaryClassifier(num_classes=2) elif self.cfg.task == 'seg': self.model = UNet(num_classes=2) self.criterion = nn.CrossEntropyLoss() self.optim = torch.optim.Adam(self.model.parameters(), lr=self.cfg.lr, betas=(self.cfg.beta0, self.cfg.beta1)) if self.n_gpus > 1: print('### {} of gpus are used!!!'.format(self.n_gpus)) self.model = nn.DataParallel(self.model) self.model = self.model.to(self.device) def build_recorder(self): # Train recorder self.train_time = AverageMeter() self.train_loss = AverageMeter() # For classification self.train_cls_acc = AverageMeter() # For segmentation self.train_mIoU = AverageMeter() self.train_pix_acc = AverageMeter() # Validation recorder self.val_time = AverageMeter() self.val_loss = AverageMeter() # For classification self.val_cls_acc = AverageMeter() # For segmentation self.val_mIoU = AverageMeter() self.val_pix_acc = AverageMeter() # self.logger = Logger('./logs') self.best_cls = 0 self.best_seg = 0 def load_pre_model(self): """ Load pretrained model """ print('=> loading checkpoint {}'.format(self.cfg.pre_model)) checkpoint = torch.load(self.cfg.pre_model) self.start_epoch = checkpoint['epoch'] self.model.load_state_dict(checkpoint['state_dict']) self.optim.load_state_dict(checkpoint['optim']) print('=> loaded checkpoint {}(epoch {})'.format( self.cfg.pre_model, self.start_epoch)) #TODO:Inference part: def infer(self, data): """ input @data: iterable 256 x 256 patches output @output : segmentation results from each patch i) If classifier's result is that there is a tissue inside of patch, outcome is a masked result. ii) Otherwise, output is segmentated mask which all of pixels are background """ # Data Loading # Load models of classification and segmetation and freeze them self.freeze() # Forward images to Classification model / Select targeted images # Forward images to Segmentation model # Record Loss / Accuracy / Pixel-Accuracy # Print samples out.. def freeze(self): pass print('{}, {} have frozen!!!'.format('model_name_1', 'model_name_2'))
base_transform = torchvision.transforms.Compose([ torchvision.transforms.ToTensor(), ]) test_dataset = KneeDataset('./data/test_knee/', transform=base_transform) test_data_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size_test, shuffle=False) # Load model net = UNet().to(device) optimizer = torch.optim.SGD(net.parameters(), lr=lr, momentum=momentum) ckpt = torch.load('best_unet.pt') net.load_state_dict(ckpt['model_state_dict']) optimizer.load_state_dict(ckpt['optimizer_state_dict']) with torch.no_grad(): net.eval() i = 0 for img, mask in test_data_loader: inputs = img.to(device) labels = mask.to(device) outputs = net(inputs) inputs = inputs.squeeze() labels = labels.squeeze() outputs = outputs.squeeze() pred = np.logical_not(outputs.cpu() < 0.5)
from model import UNet path = "./checkpoints/checkpoint_SAR.pt" model_test = UNet() model_test.load_state_dict(torch.load(path,map_location=torch.device('cpu'))) for filename in os.listdir('./testimages'): if not filename.startswith('.'): image = Image.open('./testimages/' + filename).convert("LA") mean, std = np.mean(image), np.std(image) x=9 data_transforms = transforms.Compose([ transforms.RandomResizedCrop(2**x), transforms.ToTensor(), transforms.Normalize(0.5, 0.5) ]) img = data_transforms(image)[0] img = torch.reshape(img, (1, 1, 2**x, 2**x)) image_log = lintolog(img) img_out = logtolin(model_test(image_log)) plt.imsave('./output/'+filename, img_out[0][0], cmap='gray') fig2 = plt.figure(figsize = (10,10)) # create a 5 x 5 figure ax2 = fig2.add_subplot(121) ax2.imshow(img_out[0][0], interpolation='none', cmap='gray') ax2 = fig2.add_subplot(122) ax2.imshow(img[0][0], interpolation='none', cmap='gray') ax2.set_title(filename)
class Instructor: ''' Model training and evaluation ''' def __init__(self, opt): self.opt = opt if opt.inference: self.testset = TestImageDataset(fdir=opt.impaths['test'], imsize=opt.imsize) else: self.trainset = ImageDataset(fdir=opt.impaths['train'], bdir=opt.impaths['btrain'], imsize=opt.imsize, mode='train', aug_prob=opt.aug_prob, prefetch=opt.prefetch) self.valset = ImageDataset(fdir=opt.impaths['val'], bdir=opt.impaths['bval'], imsize=opt.imsize, mode='val', aug_prob=opt.aug_prob, prefetch=opt.prefetch) self.model = UNet(n_channels=3, n_classes=1, bilinear=self.opt.use_bilinear) if opt.checkpoint: self.model.load_state_dict( torch.load('./state_dict/{:s}'.format(opt.checkpoint), map_location=self.opt.device)) print('checkpoint {:s} has been loaded'.format(opt.checkpoint)) if opt.multi_gpu == 'on': self.model = torch.nn.DataParallel(self.model) self.model = self.model.to(opt.device) self._print_args() def _print_args(self): n_trainable_params, n_nontrainable_params = 0, 0 for p in self.model.parameters(): n_params = torch.prod(torch.tensor(p.shape)) if p.requires_grad: n_trainable_params += n_params else: n_nontrainable_params += n_params self.info = 'n_trainable_params: {0}, n_nontrainable_params: {1}\n'.format( n_trainable_params, n_nontrainable_params) self.info += 'training arguments:\n' + '\n'.join([ '>>> {0}: {1}'.format(arg, getattr(self.opt, arg)) for arg in vars(self.opt) ]) if self.opt.device.type == 'cuda': print('cuda memory allocated:', torch.cuda.memory_allocated(opt.device.index)) print(self.info) def _reset_records(self): self.records = { 'best_epoch': 0, 'best_dice': 0, 'train_loss': list(), 'val_loss': list(), 'val_dice': list(), 'checkpoints': list() } def _update_records(self, epoch, train_loss, val_loss, val_dice): if val_dice > self.records['best_dice']: path = './state_dict/{:s}_dice{:.4f}_temp{:s}.pt'.format( self.opt.model_name, val_dice, str(time.time())[-6:]) if self.opt.multi_gpu == 'on': torch.save(self.model.module.state_dict(), path) else: torch.save(self.model.state_dict(), path) self.records['best_epoch'] = epoch self.records['best_dice'] = val_dice self.records['checkpoints'].append(path) self.records['train_loss'].append(train_loss) self.records['val_loss'].append(val_loss) self.records['val_dice'].append(val_dice) def _draw_records(self): timestamp = str(int(time.time())) print('best epoch: {:d}'.format(self.records['best_epoch'])) print('best train loss: {:.4f}, best val loss: {:.4f}'.format( min(self.records['train_loss']), min(self.records['val_loss']))) print('best val dice {:.4f}'.format(self.records['best_dice'])) os.rename( self.records['checkpoints'][-1], './state_dict/{:s}_dice{:.4f}_save{:s}.pt'.format( self.opt.model_name, self.records['best_dice'], timestamp)) for path in self.records['checkpoints'][0:-1]: os.remove(path) # Draw figures plt.figure() trainloss, = plt.plot(self.records['train_loss']) valloss, = plt.plot(self.records['val_loss']) plt.legend([trainloss, valloss], ['train', 'val'], loc='upper right') plt.title('{:s} loss curve'.format(timestamp)) plt.savefig('./figs/{:s}_loss.png'.format(timestamp), format='png', transparent=True, dpi=300) plt.figure() valdice, = plt.plot(self.records['val_dice']) plt.title('{:s} dice curve'.format(timestamp)) plt.savefig('./figs/{:s}_dice.png'.format(timestamp), format='png', transparent=True, dpi=300) # Save report report = '\t'.join( ['val_dice', 'train_loss', 'val_loss', 'best_epoch', 'timestamp']) report += "\n{:.4f}\t{:.4f}\t{:.4f}\t{:d}\t{:s}\n{:s}".format( self.records['best_dice'], min(self.records['train_loss']), min(self.records['val_loss']), self.records['best_epoch'], timestamp, self.info) with open('./logs/{:s}_log.txt'.format(timestamp), 'w') as f: f.write(report) print('report saved:', './logs/{:s}_log.txt'.format(timestamp)) def _train(self, train_dataloader, criterion, optimizer): self.model.train() train_loss, n_total, n_batch = 0, 0, len(train_dataloader) for i_batch, sample_batched in enumerate(train_dataloader): inputs, target = sample_batched[0].to( self.opt.device), sample_batched[1].to(self.opt.device) predict = self.model(inputs) optimizer.zero_grad() loss = criterion(predict, target) loss.backward() optimizer.step() train_loss += loss.item() * len(sample_batched) n_total += len(sample_batched) ratio = int((i_batch + 1) * 50 / n_batch) sys.stdout.write("\r[" + ">" * ratio + " " * (50 - ratio) + "] {}/{} {:.2f}%".format(i_batch + 1, n_batch, (i_batch + 1) * 100 / n_batch)) sys.stdout.flush() print() return train_loss / n_total def _evaluation(self, val_dataloader, criterion): self.model.eval() val_loss, val_dice, n_total = 0, 0, 0 with torch.no_grad(): for sample_batched in val_dataloader: inputs, target = sample_batched[0].to( self.opt.device), sample_batched[1].to(self.opt.device) predict = self.model(inputs) loss = criterion(predict, target) dice = dice_coeff(predict, target) val_loss += loss.item() * len(sample_batched) val_dice += dice.item() * len(sample_batched) n_total += len(sample_batched) return val_loss / n_total, val_dice / n_total def run(self): _params = filter(lambda p: p.requires_grad, self.model.parameters()) optimizer = torch.optim.Adam(_params, lr=self.opt.lr, weight_decay=self.opt.l2reg) criterion = BCELoss2d() train_dataloader = DataLoader(dataset=self.trainset, batch_size=self.opt.batch_size, shuffle=True) val_dataloader = DataLoader(dataset=self.valset, batch_size=self.opt.batch_size, shuffle=False) self._reset_records() for epoch in range(self.opt.num_epoch): train_loss = self._train(train_dataloader, criterion, optimizer) val_loss, val_dice = self._evaluation(val_dataloader, criterion) self._update_records(epoch, train_loss, val_loss, val_dice) print( '{:d}/{:d} > train loss: {:.4f}, val loss: {:.4f}, val dice: {:.4f}' .format(epoch + 1, self.opt.num_epoch, train_loss, val_loss, val_dice)) self._draw_records() def inference(self): test_dataloader = DataLoader(dataset=self.testset, batch_size=1, shuffle=False) n_batch = len(test_dataloader) with torch.no_grad(): for i_batch, sample_batched in enumerate(test_dataloader): index, inputs = sample_batched[0], sample_batched[1].to( self.opt.device) predict = self.model(inputs) self.testset.save_img(index.item(), predict, self.opt.use_crf) ratio = int((i_batch + 1) * 50 / n_batch) sys.stdout.write( "\r[" + ">" * ratio + " " * (50 - ratio) + "] {}/{} {:.2f}%".format(i_batch + 1, n_batch, (i_batch + 1) * 100 / n_batch)) sys.stdout.flush() print()
default=False, help='use cuda') parser.add_option('-l', '--load', dest='load', default=False, help='load file model') (options, args) = parser.parse_args() return options if __name__ == '__main__': args = get_args() net = UNet(n_classes=args.n_classes).cuda() if args.load: net.load_state_dict(torch.load(args.load)) print('Model loaded from %s' % (args.load)) if args.gpu: net.cuda() cudnn.benchmark = True train_net(net=net, epochs=args.epochs, n_classes=args.n_classes, gpu=args.gpu, data_dir=args.data_dir)
'--load', dest='load', default=False, help='load file model') (options, args) = parser.parse_args() return options if __name__ == '__main__': args = get_args() net = UNet(n_classes=args.n_classes) if args.load: if args.gpu: net.load_state_dict(torch.load(args.load)) else: net.load_state_dict(torch.load(args.load, map_location='cpu')) print('Model loaded from %s' % (args.load)) if args.gpu: net.cuda() cudnn.benchmark = True train_net(net=net, epochs=args.epochs, n_classes=args.n_classes, gpu=args.gpu, data_dir=args.data_dir)
dct['down4.mpconv.1.conv.4.bias'].data.copy_(dctvgg['features.38.bias']) dct['down4.mpconv.1.conv.4.running_mean'].data.copy_( dctvgg['features.38.running_mean']) # dct['down4.mpconv.1.conv.4.running_var'].data.copy_( dctvgg['features.38.running_var']) dct['down4.mpconv.1.conv.6.weight'].data.copy_(dctvgg['features.40.weight']) dct['down4.mpconv.1.conv.6.bias'].data.copy_(dctvgg['features.40.bias']) dct['down4.mpconv.1.conv.7.weight'].data.copy_(dctvgg['features.41.weight']) # dct['down4.mpconv.1.conv.7.bias'].data.copy_(dctvgg['features.41.bias']) dct['down4.mpconv.1.conv.7.running_mean'].data.copy_( dctvgg['features.41.running_mean']) # dct['down4.mpconv.1.conv.7.running_var'].data.copy_( dctvgg['features.41.running_var']) model.load_state_dict(dct) writer = SummaryWriter() def compute_overlaps_masks(masks1, masks2): """Computes IoU overlaps between two sets of masks. masks1, masks2: [Height, Width, instances] """ # If either set of masks is empty return empty result if masks1.shape[-1] == 0 or masks2.shape[-1] == 0: return np.zeros((masks1.shape[-1], masks2.shape[-1])) # flatten masks and compute their areas masks1 = np.reshape(masks1 > .5, (-1, masks1.shape[-1])).astype(np.float32) masks2 = np.reshape(masks2 > .5, (-1, masks2.shape[-1])).astype(np.float32)
if __name__ == "__main__": args = get_args() in_files = args.input out_files = get_output_filenames(args) net = UNet(3, 1) net = torch.nn.DataParallel(net) print("Loading model {}".format(args.model)) if not args.cpu: print("Using CUDA version of the net, prepare your GPU !") net.cuda() net.load_state_dict(torch.load(args.model)) else: net.cpu() net.load_state_dict(torch.load(args.model, map_location='cpu')) print("Using CPU version of the net, this may be very slow") print("Model loaded !") for i, fn in enumerate(in_files): print("\nPredicting image {} ...".format(fn)) img = Image.open(fn) if img.size[0] < img.size[1]: print("Error: image height larger than the width") mask = predict_img(net=net,
# model = NestedUNet(n_channels=1, n_classes=20) # model = SCSE_UNet(n_channels=1, n_classes=20) # model = SCSENestedUNet(n_channels=1, n_classes=18) # model = DilatedUNet(in_channels=1, classes=20) # model = ResUNet(in_channel=1, n_classes=5) # stage1_model_whole = ResUNet(in_channel=1, n_classes=2) # # root_dir = "E:\CSI2019\AASCE_stage4_Process/" # stage1_model_whole_load_dir = "saved_model/stage1_resunet_whole_line/CP200.pth" # stage1_model_whole.load_state_dict(torch.load(os.path.join(root_dir, stage1_model_whole_load_dir))) # logger.info('Stage1_Model loaded from {}'.format(stage1_model_whole_load_dir)) stage1_model_whole = UNet(n_channels=1, n_classes=2) # root_dir = "./" stage1_model_whole_load_dir = "saved_model/stage1_unet_whole_line/CP324.pth" stage1_model_whole.load_state_dict( torch.load(os.path.join(root_dir, stage1_model_whole_load_dir))) logger.info( 'Stage1_Model loaded from {}'.format(stage1_model_whole_load_dir)) stage1_model_segm = ResUNet(in_channel=1, n_classes=2) # stage1_model_segm_load_dir = "saved_model/stage1_resunet_segm_line/CP1127.pth" stage1_model_segm.load_state_dict( torch.load(os.path.join(root_dir, stage1_model_segm_load_dir))) logger.info( 'Stage1_Model loaded from {}'.format(stage1_model_segm_load_dir)) stage2_model_box = UNet(n_channels=1, n_classes=18) # stage2_model_box_load_dir = "saved_model/stage2__unet_box/Bestmodel_568.pth" stage2_model_box.load_state_dict( torch.load(os.path.join(root_dir, stage2_model_box_load_dir))) logger.info(
args = parser.parse_args() criterionMSE = nn.MSELoss() #.to(device) transform = transforms.Compose(transform_list) img_dir = open('/n/holyscratch01/wadduwage_lab/uom_bme/dataset_static_2020/20200105_synthBeads_1/test.txt','r') avg_mse = 0 avg_psnr = 0 for epochs in range(55,56): my_model = 'ckpt/train_deep_tfm_loss_mse/fcn_deep_' + str(epochs) + '.pth' netG = UNet(n_classes=args.output_nc) netG.load_state_dict(torch.load(my_model)) netG.eval() p = 0 f_path = '/n/holyscratch01/wadduwage_lab/uom_bme/dataset_static_2020/20200105_synthBeads_1/tr_data_1sls/' for line in img_dir: print(line) GT_ = io.imread(f_path + str(line[0:-1]) + '_gt.png') modalities = np.zeros((32,128,128)) for i in range(0,32): modalities[i,:,:] = io.imread(f_path + str(line[0:-1]) +'_'+str(i+1) +'.png') depth = modalities.shape[2] predicted_im = np.zeros((128,128,1)) if np.min(np.array(GT_))==np.max(np.array(GT_)): print('Yes') GT = torch.from_numpy(np.divide(GT_,max_gt)) img = torch.from_numpy(np.divide(modalities,max_im)[None, :, :]).float()