def main_FCNN(): # Step 01: Get Input Resources and Model Configuration parser = app_argparse() args = parser.parse_args() print(args) INPUT_IMAGE_PATH = args.input_RGB LABEL_IMAGE_PATH = args.input_GT WEIGHTS_FILE_PATH = args.output_model_path LOSS_PLOT_PATH = args.output_loss_plot use_gpu = args.use_gpu use_pretrain = args.use_pretrain epochs = args.epochs batch_size = args.batch_size tile_size = args.tile_size learning_rate = args.learning_rate weight_decay = args.weight_decay # Step 02: load the pretrained model device = utils.device(use_gpu=use_gpu) # init model structure model = FCNN() # model = utils.load_weights_from_disk(model) if use_pretrain: model = utils.load_entire_model(model, WEIGHTS_FILE_PATH, use_gpu) print("use pretrained model!") train_loader = dataset.training_loader( image_path=INPUT_IMAGE_PATH, label_path=LABEL_IMAGE_PATH, batch_size=batch_size, tile_size=tile_size, shuffle=True # use shuffle ) # turn the shuffle model, stats = train( model=model, train_loader=train_loader, device=device, epochs=epochs, batch_size=batch_size, tile_size=tile_size, learning_rate=learning_rate, weight_decay=weight_decay, ) # comment the following section to compare the results with 4 workers and pin_memory in dataloader. # Step 03: save the model # model_path = utils.save_weights_to_disk(model) model_path = utils.save_entire_model(model, WEIGHTS_FILE_PATH) # save the loss figure and data stats.save_loss_plot(LOSS_PLOT_PATH)
def main_UNet_II(): # TODO: Get through CLI arg # Step 01: Get Input Resources and Model Configuration parser = app_argparse() args = parser.parse_args() # print(args) use_gpu = args.use_gpu # tile_size = args.tile_size tile_size = (200, 200) INPUT_IMAGE_PATH = args.input_RGB LABEL_IMAGE_PATH = args.input_GT # WEIGHTS_FILE_PATH = args.output_model_path WEIGHTS_FILE_PATH = "weights/Adam.UNet.weights.II.pt" OUTPUT_IMAGE_PATH = args.output_images # Step 02: Get Input Resources and Model Configuration device = utils.device(use_gpu=use_gpu) model = UNet() # model = utils.load_weights_from_disk(model) model = utils.load_entire_model(model, WEIGHTS_FILE_PATH, use_gpu) print("use pretrained model!") # print(model) # summary(model, (3, tile_size[0], tile_size[1])) # this is issue !!! loader = dataset.full_image_loader( INPUT_IMAGE_PATH, LABEL_IMAGE_PATH, tile_size=tile_size) prediction = predict(model, loader, device=device, class_label=ClassLabel.house) # Step 03: save the output input_image = utils.input_image(INPUT_IMAGE_PATH) pred_image, mask_image = utils.overlay_class_prediction( input_image, prediction) pred_image_path = OUTPUT_IMAGE_PATH + "/prediction.png" pred_image.save(pred_image_path) pred_mask_path = OUTPUT_IMAGE_PATH + "/mask.png" mask_image.save(pred_mask_path) print("(i) Prediction and Mask image saved at {}".format(pred_image_path)) print("(ii) Mask image saved at {}".format(pred_mask_path)) # Step 04: Check the metrics img_gt = np.array(Image.open(LABEL_IMAGE_PATH), dtype=np.int32) img_mask = np.array(Image.open(pred_mask_path), dtype=np.int32) metricComputation(img_gt, img_mask)
def main_UNet_via_Folder(): # Step 01: Get Input Resources and Model Configuration parser = app_argparse() args = parser.parse_args() # INPUT_IMAGE_PATH = args.input_RGB # LABEL_IMAGE_PATH = args.input_GT INPUT_IMAGE_PATH = "data/Aerial/RGBRandom" LABEL_IMAGE_PATH = "data/Aerial/GTRandom" # WEIGHTS_FILE_PATH = args.output_model_path WEIGHTS_FILE_PATH = "weights/Adam.UNet.weights.II.pt" # LOSS_PLOT_PATH = args.output_loss_plot OUTPUT_IMAGE_PATH = args.output_images print(args) use_gpu = args.use_gpu use_pretrain = True # epochs = args.epochs epochs = 10 batch_size = args.batch_size tile_size = args.tile_size learning_rate = args.learning_rate weight_decay = args.weight_decay # Step 02: load the pretrained model device = utils.device(use_gpu=use_gpu) # init model structure model = UNet() # model = utils.load_weights_from_disk(model) if use_pretrain and Path(WEIGHTS_FILE_PATH).is_file(): model = utils.load_entire_model(model, WEIGHTS_FILE_PATH, use_gpu) print("use pretrained model!") else: print("build new model") train_loader = dataset.create_image_loader( image_path=INPUT_IMAGE_PATH, label_path=LABEL_IMAGE_PATH, batch_size=batch_size, shuffle=True # use shuffle ) # turn the shuffle model, stats = train( model=model, train_loader=train_loader, device=device, epochs=epochs, batch_size=batch_size, tile_size=tile_size, learning_rate=learning_rate, weight_decay=weight_decay, ) # comment the following section to compare the results with 4 workers and pin_memory in dataloader. # Step 03: save the model # model_path = utils.save_weights_to_disk(model) model_path = utils.save_entire_model(model, WEIGHTS_FILE_PATH) # save the loss figure and data stats.save_loss_plot(OUTPUT_IMAGE_PATH)
def matplotlib_imshow(img, one_channel=False): if one_channel: img = img.mean(dim=0) img = img / 2 + 0.5 # unnormalize npimg = img.numpy() if one_channel: plt.imshow(npimg, cmap="Greys") else: plt.imshow(np.transpose(npimg, (1, 2, 0))) if __name__ == "__main__": parser = app_argparse() args = parser.parse_args() print(args) use_gpu = args.use_gpu INPUT_IMAGE_PATH = args.input_RGB LABEL_IMAGE_PATH = args.input_GT WEIGHTS_FILE_PATH = args.output_model_path # Step 1: TensorBoard setup # default `log_dir` is "runs" - we'll be more specific here writer = SummaryWriter("runs/aerial_image_segmentation") # Step 2: Writing to TensorBoard # get some random training images tile_size = (250, 250) loader = dataset.full_image_loader(INPUT_IMAGE_PATH,