import tensorflow as tf
import numpy as np
from cityscapes import CityscapesDatset
from model import CGNet
from pipeline import batch_generator
from tqdm import tqdm
from pipeline import load_image_test
from numpy import matlib
import cv2
model = CGNet(classes=19)
## TODO we need to make argument input from command line
DATA_DIR = '/home/soojin/UOS-SSaS Dropbox/05. Data/00. Benchmarks/01. cityscapes'
tf.executing_eagerly()
# choose 'val' for validation or 'test' for test
cityscapes_dataset = CityscapesDatset(DATA_DIR, data_type = 'train')
TEST_LENGTH = len(cityscapes_dataset)
print("Length of the dataset : {}".format(TEST_LENGTH))
model_weight_path = '/home/soojin/UOS-SSaS Dropbox/05. Data/03. Checkpoints/#cgnet/2021.07.28 single_train/epoch_240.h5'
model.build((1, 680, 680, 3))
model.load_weights(model_weight_path)
########################################################################### # Created by: Tianyi Wu # Email: [email protected] # Copyright (c) 2018 ########################################################################### from utils.summary import summary import model.CGNet as net model = net.Context_Guided_Network(19, M=3, N=21) model.cuda() summary(model, (3, 640, 640))
def test_func(args):
"""
main function for testing
param args: global arguments
return: None
"""
print(args)
global network_type
if args.cuda:
print("=====> use gpu id: '{}'".format(args.gpus))
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
if not torch.cuda.is_available():
raise Exception(
"no GPU found or wrong gpu id, please run without --cuda")
args.seed = random.randint(1, 10000)
print("Random Seed: ", args.seed)
torch.manual_seed(args.seed)
if args.cuda:
torch.cuda.manual_seed(args.seed)
print('=====> checking if processed cached_data_file exists')
if not os.path.isfile(args.inform_data_file):
dataCollect = CityscapesTrainInform(
args.data_dir,
args.classes,
train_set_file=args.dataset_list,
inform_data_file=args.inform_data_file
) #collect mean std, weigth_class information
data = dataCollect.collectDataAndSave()
if data is None:
print("error while pickling data, please check")
exit(-1)
else:
data = pickle.load(open(args.inform_data_file, "rb"))
M = args.M
N = args.N
model = CGNet.Context_Guided_Network(classes=args.classes, M=M, N=N)
network_type = "CGNet"
print("Arch: CGNet")
# define optimization criteria
weight = torch.from_numpy(
data['classWeights']) # convert the numpy array to torch
if args.cuda:
weight = weight.cuda()
criteria = CrossEntropyLoss2d(weight) #weight
if args.cuda:
model = model.cuda() # using GPU for inference
criteria = criteria.cuda()
cudnn.benchmark = True
print('Dataset statistics')
print('mean and std: ', data['mean'], data['std'])
print('classWeights: ', data['classWeights'])
if args.save_seg_dir:
if not os.path.exists(args.save_seg_dir):
os.makedirs(args.save_seg_dir)
# validation set
testLoader = torch.utils.data.DataLoader(CityscapesTestDataSet(
args.data_dir, args.test_data_list, mean=data['mean']),
batch_size=1,
shuffle=False,
num_workers=args.num_workers,
pin_memory=True)
if args.resume:
if os.path.isfile(args.resume):
print("=====> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
#model.load_state_dict(checkpoint['model'])
model.load_state_dict(convert_state_dict(checkpoint['model']))
else:
print("=====> no checkpoint found at '{}'".format(args.resume))
print("=====> beginning testing")
print("test set length: ", len(testLoader))
test(args, testLoader, model)
def train_model(args):
"""
Main function for training
Args:
args: global arguments
"""
h, w = map(int, args.input_size.split(','))
input_size = (h, w)
print("=====> Check if processed data file exists or not")
if not os.path.isfile(args.inform_data_file):
print("%s is not found" % (args.inform_data_file))
dataCollect = CityscapesTrainInform(
args.data_dir,
args.classes,
train_set_file=args.dataset_list,
inform_data_file=args.inform_data_file
) #collect mean std, weigth_class information
datas = dataCollect.collectDataAndSave()
if datas is None:
print('Error while pickling data. Please check.')
exit(-1)
else:
datas = pickle.load(open(args.inform_data_file, "rb"))
print(args)
global network_type
if args.cuda:
print("=====> Use gpu id: '{}'".format(args.gpus))
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
if not torch.cuda.is_available():
raise Exception(
"No GPU found or Wrong gpu id, please run without --cuda")
args.seed = random.randint(1, 10000)
print("Random Seed: ", args.seed)
torch.manual_seed(args.seed)
if args.cuda:
torch.cuda.manual_seed(args.seed)
cudnn.enabled = True
M = args.M
N = args.N
# load the model
print('=====> Building network')
model = CGNet.Context_Guided_Network(classes=args.classes, M=M, N=N)
network_type = "CGNet"
print("current architeture: CGNet")
args.savedir = args.savedir + network_type + "_M" + str(M) + 'N' + str(
N) + '/'
# create the directory of checkpoint if not exist
if not os.path.exists(args.savedir):
os.makedirs(args.savedir)
print('=====> Computing network parameters')
total_paramters = netParams(model)
print('Total network parameters: ' + str(total_paramters))
print("data['classWeights']: ", datas['classWeights'])
weight = torch.from_numpy(datas['classWeights'])
criteria = CrossEntropyLoss2d(weight)
criteria = criteria.cuda()
print('=====> Dataset statistics')
print('mean and std: ', datas['mean'], datas['std'])
if args.cuda:
if torch.cuda.device_count() > 1:
print("torch.cuda.device_count()=", torch.cuda.device_count())
model = torch.nn.DataParallel(
model).cuda() #multi-card data parallel
else:
print("single GPU for training")
model = model.cuda() #single card
start_epoch = 0
#DataLoader
trainLoader = data.DataLoader(CityscapesDataSet(args.data_dir,
args.train_data_list,
crop_size=input_size,
scale=args.random_scale,
mirror=args.random_mirror,
mean=datas['mean']),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True,
drop_last=True)
valLoader = data.DataLoader(CityscapesValDataSet(args.data_dir,
args.val_data_list,
f_scale=1,
mean=datas['mean']),
batch_size=1,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True,
drop_last=True)
if args.resume:
if os.path.isfile(args.resume):
print("=====> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
start_epoch = checkpoint['epoch']
model.load_state_dict(checkpoint['model'])
print("=====> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("=====> no checkpoint found at '{}'".format(args.resume))
model.train()
cudnn.benchmark = True
logFileLoc = args.savedir + args.logFile
if os.path.isfile(logFileLoc):
logger = open(logFileLoc, 'a')
else:
logger = open(logFileLoc, 'w')
logger.write("Parameters: %s" % (str(total_paramters)))
logger.write(
"\n%s\t\t%s\t\t%s\t\t%s\t\t%s\t\t" %
('Epoch', 'Loss(Tr)', 'Loss(val)', 'mIOU (tr)', 'mIOU (val)'))
logger.flush()
optimizer = torch.optim.Adam(model.parameters(),
args.lr, (0.9, 0.999),
eps=1e-08,
weight_decay=5e-4)
print('=====> beginning training')
for epoch in range(start_epoch, args.max_epochs):
lossTr, per_class_iu_tr, mIOU_tr, lr = train(args, trainLoader, model,
criteria, optimizer,
epoch)
# evaluate on validation set
if epoch % 50 == 0:
mIOU_val, per_class_iu = val(args, valLoader, model, criteria)
logger.write("\n%d\t\t%.4f\t\t%.4f\t\t%.4f\t\t%.7f" %
(epoch, lossTr, mIOU_tr, mIOU_val, lr))
logger.flush()
print("Epoch : " + str(epoch) + ' Details')
print(
"\nEpoch No.: %d\tTrain Loss = %.4f\t mIOU(tr) = %.4f\t mIOU(val) = %.4f\t lr= %.6f"
% (epoch, lossTr, mIOU_tr, mIOU_val, lr))
else:
logger.write("\n%d\t\t%.4f\t\t%.4f\t\t%.7f" %
(epoch, lossTr, mIOU_tr, lr))
logger.flush()
print("Epoch : " + str(epoch) + ' Details')
print(
"\nEpoch No.: %d\tTrain Loss = %.4f\t mIOU(tr) = %.4f\t lr= %.6f"
% (epoch, lossTr, mIOU_tr, lr))
#save the model
model_file_name = args.savedir + '/model_' + str(epoch + 1) + '.pth'
state = {"epoch": epoch + 1, "model": model.state_dict()}
torch.save(state, model_file_name)
logger.close()
def test_model(args):
"""
main function for testing
args:
args: global arguments
"""
print("=====> Check if the cached file exists ")
if not os.path.isfile(args.inform_data_file):
print("%s is not found" % (args.inform_data_file))
dataCollect = CamVidTrainInform(
args.data_dir,
args.classes,
train_set_file=args.dataset_list,
inform_data_file=args.inform_data_file
) #collect mean std, weigth_class information
datas = dataCollect.collectDataAndSave()
if datas is None:
print('Error while pickling data. Please check.')
exit(-1)
else:
print("%s exists" % (args.inform_data_file))
datas = pickle.load(open(args.inform_data_file, "rb"))
print(args)
global network_type
if args.cuda:
print("=====> Use gpu id: '{}'".format(args.gpus))
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
if not torch.cuda.is_available():
raise Exception(
"No GPU found or Wrong gpu id, please run without --cuda")
args.seed = random.randint(1, 10000)
print("Random Seed: ", args.seed)
torch.manual_seed(args.seed)
if args.cuda:
torch.cuda.manual_seed(args.seed)
cudnn.enabled = True
M = args.M
N = args.N
model = CGNet.Context_Guided_Network(classes=args.classes, M=M, N=N)
network_type = "CGNet"
print("=====> current architeture: CGNet_M%sN%s" % (M, N))
total_paramters = netParams(model)
print("the number of parameters: " + str(total_paramters))
print("data['classWeights']: ", datas['classWeights'])
weight = torch.from_numpy(datas['classWeights'])
print("=====> Dataset statistics")
print("mean and std: ", datas['mean'], datas['std'])
# define optimization criteria
criteria = CrossEntropyLoss2d(weight, args.ignore_label)
if args.cuda:
model = model.cuda()
criteria = criteria.cuda()
#load test set
train_transform = transforms.Compose([transforms.ToTensor()])
testLoader = data.DataLoader(CamVidValDataSet(args.data_dir,
args.test_data_list,
f_scale=1,
mean=datas['mean']),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True,
drop_last=True)
if args.resume:
if os.path.isfile(args.resume):
print("=====> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
#model.load_state_dict(convert_state_dict(checkpoint['model']))
model.load_state_dict(checkpoint['model'])
else:
print("=====> no checkpoint found at '{}'".format(args.resume))
cudnn.benchmark = True
print("=====> beginning test")
print("length of test set:", len(testLoader))
mIOU_val, per_class_iu = test(args, testLoader, model, criteria)
print(mIOU_val)
print(per_class_iu)
epochs = config.get('EPOCHS', None)
num_gpu = config.get('NUM_GPU', None)
num_m_blocks = config.get('NUM_M_BLOCKS', None)
num_n_blocks = config.get('NUM_N_BLOCKS', None)
init_learn_rate = config.get('INIT_LEARN_RATE', None)
end_learn_rate = config.get('END_LEARN_RATE', None)
power = config.get('POWER', None)
class_weight = config.get('CLASS_WEIGHT', None)
log_template = 'Epoch: {}/{}, steps:{}/{}, Loss: {:2f}, Accuracy: {:2f}, MeanIoU: {:2f}'
model = CGNet(num_classes = num_classes, M= num_m_blocks, N=num_n_blocks)
train_loss = tf.keras.metrics.Mean(name='train_loss')
train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy')
train_iou = tf.keras.metrics.MeanIoU(num_classes= num_classes, name='train_miou')
if dataset == 'Cityscapes':
train_dataset = CityscapesDatset(data_dir)
elif dataset == 'Concrete_Damage_Cityscapes':
train_dataset = Concrete_Damage_Dataset_as_Cityscapes(data_dir)
num_steps = len(train_dataset)//batch_size
optimizer = load_optimizer(init_learn_rate, end_learn_rate, power)
def test_func(args):
'''
Main function for testing
param args: global arguments
return: None
'''
print(args)
global network_type
if args.cuda:
print("=====> use gpu id: '{}'".format(args.gpus))
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
if not torch.cuda.is_available():
raise Exception(
"No GPU found or Wrong gpu id, please run without --cuda")
args.seed = random.randint(1, 10000)
print("Random Seed: ", args.seed)
torch.manual_seed(args.seed)
if args.cuda:
torch.cuda.manual_seed(args.seed)
print('checking if processed cached_data_file exists or not')
if not os.path.isfile(args.cached_data_file):
dataCollect = StatisticalInformDataset(args.data_dir, args.classes,
args.cached_data_file)
data = dataCollect.collectDataAndSave()
if data is None:
print("error while pickling data, please check")
exit(-1)
else:
data = pickle.load(open(args.cached_data_file, "rb"))
M = args.M
N = args.N
# load the model
print('====> Building network')
model = net.Context_Guided_Network(classes=args.classes, M=M, N=N)
network_type = "CGNet"
print("Arch: CGNet")
if args.cuda:
#model = torch.nn.DataParallel(model).cuda() # multi-card testing
model = model.cuda() # single-card testing
print('Dataset statistics')
print('mean and std: ', data['mean'], data['std'])
print('classWeights: ', data['classWeights'])
if args.save_seg_dir:
if not os.path.exists(args.save_seg_dir):
os.makedirs(args.save_seg_dir)
testLoader = torch.utils.data.DataLoader(CityscapesTestDataSet(
args.data_dir, args.test_data_list, mean=data['mean']),
batch_size=1,
shuffle=False,
num_workers=args.num_workers,
pin_memory=True)
if args.cuda:
cudnn.benchmark = True
print("=====> load pretrained model")
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
model.load_state_dict(convert_state_dict(checkpoint['model']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
# evaluate on test set
print("=====> beginning test")
print("test set length: ", len(testLoader))
test(args, testLoader, model)