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,