def main():
net = AFENet(classes=19, pretrained_model_path=None).cuda()
net.load_state_dict(
torch.load(os.path.join(args['model_save_path'], args['snapshot'])))
mean_std = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
input_transform = standard_transforms.Compose([
standard_transforms.ToTensor(),
standard_transforms.Normalize(*mean_std)
])
target_transform = extended_transforms.MaskToTensor()
restore_transform = standard_transforms.Compose([
extended_transforms.DeNormalize(*mean_std),
standard_transforms.ToPILImage()
])
dataset_path = args['dataset_path']
test_set = cityscapes.CityScapes(dataset_path,
'fine',
'test',
transform=input_transform,
target_transform=target_transform,
val_scale=args['scale'])
test_loader = DataLoader(test_set,
batch_size=1,
num_workers=1,
shuffle=False)
test(test_loader, net, input_transform, restore_transform, args['scale'])
def eval(epoch):
val_dataset = Cityscapes(root='.',
subset='val',
transform=MyTransform(augment=False))
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=1,
shuffle=False)
erfnet = torch.load('./model/erfnet' + str(epoch) + '.pth')
erfnet.cuda()
erfnet.eval()
iou_meter = meter.AverageValueMeter()
for i, (images, labels) in enumerate(val_loader):
images = Variable(images.cuda(), volatile=True).cuda()
labels = Variable(labels.cuda()).cuda()
outputs = erfnet(images)
iou = IoU(outputs.max(1)[1].data, labels.data)
iou_meter.add(iou)
draw(i, outputs[0].max(0)[1].byte().cpu().data)
iou_avg = iou_meter.value()[0]
print("Val_IoU : {:.5f}".format(iou_avg))
def train():
train_dataset = Cityscapes(root='.',
subset='train',
transform=MyTransform())
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=1,
shuffle=True)
erfnet = ERFNet(num_classes)
optimizer = Adam(erfnet.parameters(),
5e-4, (0.9, 0.999),
eps=1e-08,
weight_decay=1e-4)
lambda1 = lambda epoch: pow((1 - ((epoch - 1) / num_epochs)), 0.9)
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda1)
criterion = nn.NLLLoss2d(weight)
iou_meter = meter.AverageValueMeter()
erfnet.cuda()
erfnet.train()
for epoch in range(num_epochs):
print("Epoch {} : ".format(epoch + 1))
total_loss = []
iou_meter.reset()
scheduler.step(epoch)
for i, (images, labels) in enumerate(train_loader):
images = Variable(images.cuda(), requires_grad=True).cuda()
labels = Variable(labels.cuda()).cuda()
outputs = erfnet(images)
optimizer.zero_grad()
loss = criterion(torch.nn.functional.log_softmax(outputs), labels)
loss.backward()
optimizer.step()
iou = IoU(outputs.max(1)[1].data, labels.data)
iou_meter.add(iou)
total_loss.append(loss.data[0] * batch_size)
iou_avg = iou_meter.value()[0]
loss_avg = sum(total_loss) / len(total_loss)
scheduler.step(loss_avg)
print("IoU : {:.5f}".format(iou_avg))
print("Loss : {:.5f}".format(loss_avg))
torch.save(erfnet, './model/erfnet' + str(epoch) + '.pth')
eval(epoch)
def main():
net = AFENet(classes=19, pretrained_model_path=None).cuda()
net.load_state_dict(
torch.load(os.path.join(args['model_save_path'], args['snapshot'])))
net.eval()
mean_std = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
input_transform = standard_transforms.Compose([
standard_transforms.ToTensor(),
standard_transforms.Normalize(*mean_std)
])
img_path = args['image_path']
result_save_path = args['test_result_save_path']
scales = args['scale']
img = Image.open(img_path).convert('RGB')
crop_size = 768
prediction = np.zeros((19, 1024, 2048), dtype=float)
origin_w, origin_h = img.size
for scale in scales:
new_w = int(origin_w * scale)
new_h = int(origin_h * scale)
if scale == 1.0:
scale_img = img
else:
scale_img = img.resize((new_w, new_h), Image.BILINEAR)
prediction += scale_process(net, scale_img, origin_w, origin_h,
crop_size, new_w, new_h, input_transform)
prediction /= len(scales)
prediction = np.argmax(prediction, axis=0)
prediction = prediction.reshape([1024, 2048])
prediction = cityscapes.colorize_mask(prediction)
prediction.save(os.path.join(result_save_path))
print("Completed.")
torchvision.transforms.RandomHorizontalFlip(),
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
train_target_transform = torchvision.transforms.Compose([
torchvision.transforms.RandomHorizontalFlip(),
])
val_image_transform = torchvision.transforms.Compose([
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
# get cityscapes dataset from google drive link
train_dataset = Cityscapes(root=dataset_path, split='train', transform=train_image_transform, target_transform=train_target_transform)
val_dataset = Cityscapes(root=dataset_path, split='val', transform=val_image_transform)
# get train and val loaders for the corresponding datasets
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=bs_train, shuffle=True, num_workers=16)
val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=bs_val, shuffle=False, num_workers=16)
# model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model_map = {
'fcn_resnet18': network.fcn_resnet18,
'fcn_resnet50': network.fcn_resnet50,
'deeplabv3_resnet18': network.deeplabv3_resnet18,
'deeplabv3_resnet50': network.deeplabv3_resnet50,
'deeplabv3_resnet101': network.deeplabv3_resnet101,
'deeplabv3plus_resnet18': network.deeplabv3plus_resnet18,
optimizer.load_state_dict(ckpt['optimizer_state_dict'])
print('Optimizer recovered from {}.'.format(ckpt_file))
saver.load()
else:
epoch = 0
# Metrics
metrics = [
'SEMANTIC_LOSS', 'IOU_0', 'IOU_1', 'IOU_2', 'IOU_3', 'IOU_4', 'IOU_5',
'IOU_6', 'MEAN_IOU', 'PIX_ACC', 'DEPTH_LOSS', 'ABS_ERR', 'REL_ERR'
]
saver.add_metrics(metrics)
# Create datasets
dataset_path = opt.dataroot
city_train_set = Cityscapes(root=dataset_path, train=True)
city_test_set = Cityscapes(root=dataset_path, train=False)
batch_size = opt.batch_size
city_train_loader = torch.utils.data.DataLoader(dataset=city_train_set,
batch_size=batch_size,
shuffle=True)
city_test_loader = torch.utils.data.DataLoader(dataset=city_test_set,
batch_size=batch_size,
shuffle=False)
# Define parameters
total_epoch = opt.total_epoch
nb_train_batches = len(city_train_loader)
nb_test_batches = len(city_test_loader)
# you can add other transformations in this list
transforms.Resize([img_size,]),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
target_transform = transforms.Compose([
transforms.Resize([img_size,], interpolation=0),#0:InterpolationMode.NEAREST
targetToTensor
])
# Load Cityscapes train and test datasets
def load_dataset(batch_size, num_workers, split='train'):
"""
A util function for loading dataset and dataloader
Args:
batch_size: batch size (hyperparameters)
num_workers: num_workers (hyperparameters)
split: Takes input as string. Can be any split allowed by the dataloader.
Returns:
data_loader: An iterable element of the dataset
data_set: Loaded with processed dataset
"""
data_set = Cityscapes(root='cityscapes', split=split, mode='fine', target_type='semantic_basic', transform=image_transform, target_transform=target_transform)
data_loader = torch.utils.data.DataLoader(data_set, batch_size=batch_size, num_workers=num_workers, shuffle=True, pin_memory=True)
return data_loader, data_set
parser.add_argument('--backbone_type', default='resnet50', type=str, choices=['resnet50', 'resnet101'],
help='Backbone type')
parser.add_argument('--crop_h', default=512, type=int, help='Crop height for training images')
parser.add_argument('--crop_w', default=512, type=int, help='Crop width for training images')
parser.add_argument('--batch_size', default=4, type=int, help='Number of data for each batch to train')
parser.add_argument('--epochs', default=60, type=int, help='Number of sweeps over the dataset to train')
parser.add_argument('--save_path', default='results', type=str, help='Save path for results')
# args parse
args = parser.parse_args()
backbone_type, crop_h, crop_w = args.backbone_type, args.crop_h, args.crop_w
batch_size, epochs, save_path = args.batch_size, args.epochs, args.save_path
# dataset, model setup, optimizer config and loss definition
train_data = Cityscapes( split='train', crop_size=(crop_h, crop_w))
val_data = Cityscapes( split='val')
train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True, num_workers=min(4, batch_size),
drop_last=True)
val_loader = DataLoader(val_data, batch_size=batch_size, shuffle=False, num_workers=min(4, batch_size))
model = GatedSCNN(backbone_type=backbone_type, num_classes=6).cuda()
optimizer = SGD(model.parameters(), lr=1e-2, momentum=0.9, weight_decay=1e-4)
scheduler = LambdaLR(optimizer, lr_lambda=lambda eiter: math.pow(1 - eiter / epochs, 1.0))
semantic_criterion = nn.CrossEntropyLoss(ignore_index=255)
edge_criterion = BoundaryBCELoss(ignore_index=255)
task_criterion = DualTaskLoss(threshold=0.8, ignore_index=255)
results = {'train_loss': [], 'val_loss': [], 'train_PA': [], 'val_PA': [], 'train_mPA': [], 'val_mPA': [],
'train_class_mIOU': [], 'val_class_mIOU': [], 'train_category_mIOU': [], 'val_category_mIOU': []}
best_mIOU = 0.0
discriminator = ConditionalDiscriminator().to(device)
# optimizers
g_optimizer = torch.optim.Adam(generator.parameters(),
lr=args.lr,
betas=(0.5, 0.999))
d_optimizer = torch.optim.Adam(discriminator.parameters(),
lr=args.lr,
betas=(0.5, 0.999))
# loss functions
g_criterion = GeneratorLoss(alpha=100)
d_criterion = DiscriminatorLoss()
# dataset
print(f'Downloading "{args.dataset.upper()}" dataset!')
if args.dataset == 'cityscapes':
dataset = Cityscapes(root='.',
transform=transforms,
download=True,
mode='train')
elif args.dataset == 'maps':
dataset = Maps(root='.', transform=transforms, download=True, mode='train')
else:
dataset = Facades(root='.',
transform=transforms,
download=True,
mode='train')
dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True)
print('Start of training process!')
logger = Logger(filename=args.dataset)
for epoch in range(args.epochs):
ge_loss = 0.
de_loss = 0.
start = time.time()
def main():
net = AFENet(classes=19,
pretrained_model_path=args['pretrained_model_path']).cuda()
net_ori = [net.layer0, net.layer1, net.layer2, net.layer3, net.layer4]
net_new = [
net.ppm, net.cls, net.aux, net.ppm_reduce, net.aff1, net.aff2, net.aff3
]
mean_std = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
input_transform = standard_transforms.Compose([
standard_transforms.ToTensor(),
standard_transforms.Normalize(*mean_std)
])
target_transform = extended_transforms.MaskToTensor()
restore_transform = standard_transforms.Compose([
extended_transforms.DeNormalize(*mean_std),
standard_transforms.ToPILImage()
])
dataset_path = args['dataset_path']
# Loading dataset
train_set = cityscapes.CityScapes(dataset_path,
'fine',
'train',
transform=input_transform,
target_transform=target_transform)
train_loader = DataLoader(train_set,
batch_size=args['train_batch_size'],
num_workers=2,
shuffle=True)
val_set = cityscapes.CityScapes(dataset_path,
'fine',
'val',
transform=input_transform,
target_transform=target_transform)
val_loader = DataLoader(val_set,
batch_size=args['val_batch_size'],
num_workers=2,
shuffle=False)
if len(args['snapshot']) == 0:
curr_epoch = 1
args['best_record'] = {
'epoch': 0,
'val_loss': 1e10,
'acc': 0,
'acc_cls': 0,
'mean_iu': 0,
'fwavacc': 0
}
else:
print('training resumes from ' + args['snapshot'])
net.load_state_dict(
torch.load(os.path.join(args['model_save_path'],
args['snapshot'])))
split_snapshot = args['snapshot'].split('_')
curr_epoch = int(split_snapshot[1]) + 1
args['best_record'] = {
'epoch': int(split_snapshot[1]),
'val_loss': float(split_snapshot[3]),
'acc': float(split_snapshot[5]),
'acc_cls': float(split_snapshot[7]),
'mean_iu': float(split_snapshot[9]),
'fwavacc': float(split_snapshot[11])
}
params_list = []
for module in net_ori:
params_list.append(dict(params=module.parameters(), lr=args['lr']))
for module in net_new:
params_list.append(dict(params=module.parameters(),
lr=args['lr'] * 10))
args['index_split'] = 5
criterion = torch.nn.CrossEntropyLoss(ignore_index=cityscapes.ignore_label)
optimizer = torch.optim.SGD(params_list,
lr=args['lr'],
momentum=args['momentum'],
weight_decay=args['weight_decay'])
if len(args['snapshot']) > 0:
optimizer.load_state_dict(
torch.load(
os.path.join(args['model_save_path'],
'opt_' + args['snapshot'])))
check_makedirs(args['model_save_path'])
all_iter = args['epoch_num'] * len(train_loader)
for epoch in range(curr_epoch, args['epoch_num'] + 1):
train(train_loader, net, optimizer, epoch, all_iter)
validate(val_loader, net, criterion, optimizer, epoch,
restore_transform)
bn_eps = 1e-5
bn_momentum = 0.1
#specify the path of the Cityscape dataset, and construct the dataset and dataloader for training
train_path = [
'/opt/Cityscapes/leftImg8bit/train',
'/opt/Cityscapes/gtFine_trainvaltest/gtFine/train'
]
val_path = [
'/opt/Cityscapes/leftImg8bit/val',
'/opt/Cityscapes/gtFine_trainvaltest/gtFine/val'
]
use_val = False
DATASET = Cityscapes(train_path,
val_path,
train_size=train_size,
use_val=use_val)
print('the length of training set is:', len(DATASET),
',please check if consistent with your setting')
time.sleep(10)
num_classes = 19
num_workers = 4
train_loader = data.DataLoader(DATASET,
batch_size=batch_size,
num_workers=num_workers,
drop_last=True,
shuffle=True,
pin_memory=True,
sampler=None)
help='Number of steps to save predicted results')
parser.add_argument('--epochs',
default=100,
type=int,
help='Number of sweeps over the dataset to train')
# args parse
args = parser.parse_args()
data_path, crop_h, crop_w = args.data_path, args.crop_h, args.crop_w
batch_size, save_step, epochs = args.batch_size, args.save_step, args.epochs
if not os.path.exists('results'):
os.mkdir('results')
# dataset, model setup and optimizer config
train_data = Cityscapes(root=data_path,
split='train',
crop_size=(crop_h, crop_w))
val_data = Cityscapes(root=data_path, split='val')
train_loader = DataLoader(train_data,
batch_size=batch_size,
shuffle=True,
num_workers=4)
val_loader = DataLoader(val_data,
batch_size=batch_size,
shuffle=False,
num_workers=4)
model = FastSCNN(in_channels=3, num_classes=19).cuda()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
# model profile and loss definition
flops, params = profile(model,
def validate(t_val_loader, net, input_transform, restore, scales, count):
net.eval()
inputs_all, gts_all, predictions_all = [], [], []
crop_size = 768
for vi, (inputs, targets) in enumerate(t_val_loader):
prediction = np.zeros((19, 1024, 2048), dtype=float)
img = inputs.data[0]
img = restore(img)
inputs_all.append(img)
origin_w, origin_h = img.size
for scale in scales:
new_w = int(origin_w * scale)
new_h = int(origin_h * scale)
if scale == 1.0:
scale_img = img
else:
scale_img = img.resize((new_w, new_h), Image.BILINEAR)
prediction += scale_process(net, scale_img, origin_w, origin_h,
crop_size, new_w, new_h,
input_transform)
prediction /= len(scales)
prediction = np.argmax(prediction, axis=0)
prediction = prediction.reshape([1, 1024, 2048])
gts_all.append(targets.data.cpu().numpy())
predictions_all.append(prediction)
if (vi + 1) % 20 == 0:
localtime = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
print('%s Completed %d / %d' % (localtime, vi + 1, count))
gts_all = np.concatenate(gts_all)
predictions_all = np.concatenate(predictions_all)
print(gts_all.shape)
print(predictions_all.shape)
acc, acc_cls, mean_iu, fwavacc = evaluate(predictions_all, gts_all,
cityscapes.num_classes)
if args['val_save_or_not']:
check_mkdir(args['val_exp_save_path'])
for idx, data in enumerate(zip(inputs_all, gts_all, predictions_all)):
if data[0] is None:
continue
input_pil = data[0]
gt_pil = cityscapes.colorize_mask(data[1])
predictions_pil = cityscapes.colorize_mask(data[2])
input_pil.save(
os.path.join(args['val_exp_save_path'], '%d_input.png' % idx))
predictions_pil.save(
os.path.join(args['val_exp_save_path'],
'%d_prediction.png' % idx))
gt_pil.save(
os.path.join(args['val_exp_save_path'], '%d_gt.png' % idx))
print(
'-----------------------------------------------------------------------------------------------------------'
)
print('[acc %.5f], [acc_cls %.5f], [mean_iu %.5f], [fwavacc %.5f]' %
(acc, acc_cls, mean_iu, fwavacc))
print(
'-----------------------------------------------------------------------------------------------------------'
)
