def test_train_functions(train_dataset, val_dataset, test_dataset):
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=64,
shuffle=True,
num_workers=0,
)
val_loader = torch.utils.data.DataLoader(
val_dataset,
batch_size=64,
shuffle=False,
num_workers=0,
)
test_loader = torch.utils.data.DataLoader(
test_dataset,
batch_size=64,
shuffle=False,
num_workers=0,
)
device = 'cuda:0'
model = FCNN().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
loss_fn = torch.nn.BCELoss()
for i in range(5):
print('\n' + '=' * 30)
print(f'Training epoch {i + 1}')
train(
model=model,
loader=train_loader,
optimizer=optimizer,
loss_fn=loss_fn,
device=device,
)
test(
model=model,
loader=val_loader,
device=device,
)
print('\n\n\n' + '+' * 30)
print('Test set results')
test(
model=model,
loader=test_loader,
device=device,
)
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_FCNN():
# 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
INPUT_IMAGE_PATH = args.input_RGB
LABEL_IMAGE_PATH = args.input_GT
WEIGHTS_FILE_PATH = args.output_model_path
OUTPUT_IMAGE_PATH = args.output_images
# Step 02: Get Input Resources and Model Configuration
device = utils.device(use_gpu=use_gpu)
model = FCNN()
# 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 dev_model(args): # modified from __main__ in train.py
# Get the arguments from GUI
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_height, args.tile_size_width)
learning_rate = args.learning_rate
weight_decay = args.weight_decay
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,
)
# 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)
print("[>>>] Passed!")
optimizer.step()
return model, training_stats
if __name__ == '__main__':
#TODO: Get through CLI args
epochs = 10
batch_size = 1
use_gpu = False
tile_size = (256, 256)
device = utils.device(use_gpu=use_gpu)
model = FCNN()
train_loader = dataset.training_loader(batch_size=batch_size,
tile_size=tile_size)
model, stats = train(model=model,
train_loader=train_loader,
device=device,
epochs=epochs,
batch_size=batch_size,
tile_size=tile_size)
model_path = utils.save_weights_to_disk(model)
print('(i) Model saved at {}'.format(model_path))
loss_plot_path = './images/output/loss_plot.png'
import torch import torchvision import utils import dataset from model import FCNN from utils import ClassLabel from torchsummary import summary dummy_input = torch.randn(10, 3, 250, 250) model = FCNN() # model = utils.load_weights_from_disk(model) model = utils.load_entire_model(model) # call model.eval() to set dropout and batch normalization layers to evaluation mode before running inference. model.eval() print(summary(model, (3, 250, 250))) # Providing input and output names sets the display names for values # within the model's graph. Setting these does not change the semantics # of the graph; it is only for readability. # # The inputs to the network consist of the flat list of inputs (i.e. # the values you would pass to the forward() method) followed by the # flat list of parameters. You can partially specify names, i.e. provide # a list here shorter than the number of inputs to the model, and we will # only set that subset of names, starting from the beginning. # input_names = ["actual_input_1"] # output_names = ["output1"] # torch.onnx.export(model, dummy_input, "model.onnx", verbose=True,
def dev_predit(args):
use_gpu = args.use_gpu
tile_size = tile_size = (args.tile_size_height, args.tile_size_width)
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
OUTPUT_IMAGE_PATH = args.output_images
# Step 02: Get Input Resources and Model Configuration
device = utils.device(use_gpu=use_gpu)
model = FCNN()
# model = utils.load_weights_from_disk(model)
model = utils.load_entire_model(model, WEIGHTS_FILE_PATH, use_gpu)
# 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=utils.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) Prediction and Mask image saved at {}".format(pred_mask_path))
# Show Metrics Computation
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)
# show images
img_rgb = cv.imread(INPUT_IMAGE_PATH)
img_gt = cv.imread(LABEL_IMAGE_PATH)
img_pred = cv.imread(pred_mask_path) # pred_image_path
img_lost = cv.imread(LOSS_PLOT_PATH)
images = [img_rgb, img_gt, img_pred, img_lost]
titles = ["RGB", "GT", "Prediction", "Training Loss"]
plt.figure(num=None, figsize=(7, 7), dpi=80, facecolor="w", edgecolor="k")
for i in range(4):
plt.subplot(2, 2, i + 1), plt.imshow(images[i],
"gray",
vmin=0,
vmax=255)
plt.title(titles[i])
plt.xticks([]), plt.yticks([])
plt.show()
return pred_image_path, pred_mask_path