def get_model(self, dataset):
model = PSPNet(dataset.num_classes, dataset.img_size).cuda()
start_epoch = 0
if self.args.resume:
setup.load_save(model, self.args)
start_epoch = self.args.resume_epoch
self.args.start_epoch = start_epoch
model.train()
return model
def build_network(snapshot, backend, in_channels, late_fusion):
epoch = 0
backend = backend.lower()
net = PSPNet(sizes=(1, 2, 3, 6),
psp_size=2048,
deep_features_size=1024,
backend='resnet50',
in_channels=in_channels,
pretrained=False,
late_fusion=late_fusion)
# net = nn.DataParallel(net)
if snapshot is not None:
_, epoch = os.path.basename(snapshot).split('_')
epoch = int(epoch)
net.load_state_dict(torch.load(snapshot))
net = net.cuda()
return net
def get_model():
if settings.MODEL_NAME == 'deeplab':
model = DeepLab(output_stride=settings.STRIDE,
num_classes=settings.NCLASS)
elif settings.MODEL_NAME == 'pspnet':
model = PSPNet(output_stride=settings.STRIDE,
num_classes=settings.NCLASS)
elif settings.MODEL_NAME == 'ann':
model = ANNNet(output_stride=settings.STRIDE,
num_classes=settings.NCLASS)
elif settings.MODEL_NAME == 'ocnet':
model = OCNet(output_stride=settings.STRIDE,
num_classes=settings.NCLASS)
elif settings.MODEL_NAME == 'danet':
model = DANet(output_stride=settings.STRIDE,
num_classes=settings.NCLASS)
elif settings.MODEL_NAME == 'ocrnet':
model = OCRNet(output_stride=settings.STRIDE,
num_classes=settings.NCLASS)
return model
def get_full_model(net, res, n_class, input_ch, is_data_parallel=True):
if net == "fcn":
from models.fcn import ResFCN
model = ResFCN(n_class, res, input_ch)
elif net == "fcnvgg":
from models.vgg_fcn import FCN8s
# TODO suport input_ch
model = FCN8s(n_class)
elif "drn" in net:
from models.dilated_fcn import DRNSeg
assert net in [
"drn_c_26", "drn_c_42", "drn_c_58", "drn_d_22", "drn_d_38",
"drn_d_54", "drn_d_105"
]
model = DRNSeg(net, n_class, input_ch=input_ch)
elif net == "psp":
from models.pspnet import PSPNet
model = PSPNet(n_class, layer=res, input_ch=input_ch)
elif net == "unet":
from models.unet import UNet
model = UNet(n_classes=n_class, input_ch=input_ch)
elif net == "fusenet":
from models.FuseNet import FuseNet
model = FuseNet(n_class=n_class, input_ch=input_ch)
else:
raise NotImplementedError(
"Only FCN, SegNet, PSPNet, DRNet, UNet are supported!")
if is_data_parallel:
return torch.nn.DataParallel(model)
return model
input_width = 256
output_height = 256
output_width = 256
# --- Predict data ---
image_dir = "/Users/Akash_Sengupta/Documents/4th_year_project_datasets/up-s31/trial/images3/train"
label_dir = "/Users/Akash_Sengupta/Documents/4th_year_project_datasets/up-s31/trial/masks3/train"
eval_transforms = transforms.Compose([
PadToSquare(),
Resize(input_height, input_width, output_height, output_width),
ToTensor()
])
eval_dataset = UPS31Dataset(image_dir=image_dir,
label_dir=label_dir,
transform=eval_transforms)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print('Eval images found: {}'.format(len(eval_dataset)))
print('Device: {}'.format(device))
# --- Model and Loss---
model = PSPNet(num_classes)
saved_model_path = "./saved_models/pspnet_test.tar"
criterion = nn.CrossEntropyLoss(weight=None)
# --- Predict ---
evaluate(model, saved_model_path, eval_dataset, criterion, batch_size, device,
num_classes)
shuffle=False,
num_workers=4,
# pin_memory=True,
)
if args.model == "unet":
model = UNet(input_channels=NUM_INPUT_CHANNELS,
output_channels=NUM_OUTPUT_CHANNELS)
elif args.model == "segnet":
model = SegNet(input_channels=NUM_INPUT_CHANNELS,
output_hannels=NUM_OUTPUT_CHANNELS)
else:
model = PSPNet(
layers=50,
bins=(1, 2, 3, 6),
dropout=0.1,
classes=NUM_OUTPUT_CHANNELS,
use_ppm=True,
pretrained=True,
)
# class_weights = 1.0 / train_dataset.get_class_probability()
# criterion = torch.nn.CrossEntropyLoss(weight=class_weights)
criterion = torch.nn.CrossEntropyLoss()
if CUDA:
model = model.cuda(device=GPU_ID)
# class_weights = class_weights.cuda(GPU_ID)
criterion = criterion.cuda(device=GPU_ID)
if args.checkpoint:
metrics = {
'cls_acc': cls_prec,
'head_acc': heading_prec,
'size_acc': size_prec,
'IoU_2D': iou2d_mean,
'IoU_3D': iou3d_mean,
'IoU_' + str(cfg.IOU_THRESH): iou3d_gt_mean
}
return losses, metrics
psp_models = {
'resnet18':
lambda: PSPNet(sizes=(1, 2, 3, 6),
psp_size=512,
deep_features_size=256,
backend='resnet18'),
'resnet34':
lambda: PSPNet(sizes=(1, 2, 3, 6),
psp_size=512,
deep_features_size=256,
backend='resnet34'),
'resnet50':
lambda: PSPNet(sizes=(1, 2, 3, 6),
psp_size=2048,
deep_features_size=1024,
backend='resnet50'),
'resnet101':
lambda: PSPNet(sizes=(1, 2, 3, 6),
psp_size=2048,
deep_features_size=1024,
from datasets.carpart import CarPart
from albumentations import *
from models.pspnet import PSPNet
#basic setting
num_epochs = 30
print_iter = 25
batch_size = 2
vis_result = True
validation_ratio = 0.05
startlr = 1e-4
#model load
model = model = PSPNet(n_classes=30,
n_blocks=[3, 4, 6, 3],
pyramids=[6, 3, 2, 1]).cuda()
#model = torch.load("./ckpt/25.pth")
model_name = model.__class__.__name__
#dataset load
aug = Compose([
HorizontalFlip(0.5),
OneOf([
MotionBlur(p=0.2),
MedianBlur(blur_limit=3, p=0.1),
ISONoise(p=0.3),
Blur(blur_limit=3, p=0.1),
],
p=0.5),
RandomBrightnessContrast(brightness_limit=(-0.3, 0.4), p=0.5),
import torch.optim as optim
from torch.autograd import Variable
from torch.utils.data import DataLoader
from torchvision.transforms import functional
from models.pspnet import PSPNet
from datasets.voc_loader import pascalVOCLoader
from utils.metrics import runningScore
import click
import numpy as np
import scipy.misc as misc
import PIL.Image as Image
import time
models = {
'squeezenet': lambda: PSPNet(sizes=(1, 2, 3, 6), psp_size=512, deep_features_size=256, backend='squeezenet'),
'densenet': lambda: PSPNet(sizes=(1, 2, 3, 6), psp_size=1024, deep_features_size=512, backend='densenet'),
'resnet18': lambda: PSPNet(sizes=(1, 2, 3, 6), psp_size=512, deep_features_size=256, backend='resnet18'),
'resnet34': lambda: PSPNet(sizes=(1, 2, 3, 6), psp_size=512, deep_features_size=256, backend='resnet34'),
'resnet50': lambda: PSPNet(sizes=(1, 2, 3, 6), psp_size=2048, deep_features_size=1024, backend='resnet50'),
'resnet101': lambda: PSPNet(sizes=(1, 2, 3, 6), psp_size=2048, deep_features_size=1024, backend='resnet101'),
'resnet152': lambda: PSPNet(sizes=(1, 2, 3, 6), psp_size=2048, deep_features_size=1024, backend='resnet152')
}
def build_network(snapshot, backend):
epoch = 0
backend = backend.lower()
net = models[backend]()
net = nn.DataParallel(net)
net = net.cuda()
save_path = os.path.join(args.paths["project_path"],
"saves{}".format(args.eval["version"]),
"{}-{}.pth".format(args.model, trained_epochs))
dataset = TestDataset()
if args.model == "GCN":
model = GCN(cityscapes.num_classes, args.img_size, k=args.K).cuda()
elif args.model == "GCN_DENSENET":
model = GCN_DENSENET(cityscapes.num_classes, args.img_size,
k=args.K).cuda()
elif args.model == "GCN_DECONV":
model = GCN_DECONV(cityscapes.num_classes, args.img_size,
k=args.K).cuda()
elif args.model == "GCN_PSP":
model = GCN_PSP(cityscapes.num_classes, args.img_size, k=args.K).cuda()
elif args.model == "GCN_COMB":
model = GCN_COMBINED(cityscapes.num_classes, args.img_size).cuda()
elif args.model == "GCN_RESNEXT":
model = GCN_RESNEXT(dataset.num_classes, k=self.args.K).cuda()
elif args.model == "UNet":
model = UNet(cityscapes.num_classes).cuda()
elif args.model == "PSPNet":
model = PSPNet(cityscapes.num_classes, args.img_size).cuda()
else:
raise ValueError("Invalid model arg.")
model.load_state_dict(torch.load(save_path))
evaluator = Evaluator(dataset)
evaluator.eval(model)
if __name__ == '__main__':
if torch.cuda.is_available():
if args.cuda:
device = 'cuda'
torch.cuda.empty_cache()
else:
device = 'cpu'
else:
device = 'cpu'
test_data = get_files(os.path.join(args.dataset_dir, args.test_folder),
extension_filter='.png')
#initialize the model and load the checkpoint of best model.
num_classes = len(color_encoding)
model = PSPNet(num_classes)
checkpoint = torch.load(args.model_path)
model = model.load_state_dict(checkpoint['state_dict'])
model.eval()
output_img_dir = args.output_img_dir
if not os.path.exists(output_img_dir):
os.makedirs(output_img_dir)
for data_path in test_data:
try:
image_name = data_path.split('/')[-1].split('_')[:3]
new_image_name = ('_').join(image_name) + '*.png'
data = Image.open(data_path)
h, w, _ = np.array(data).shape
image = image_transform(data)
def __init__(self):
super(TNet, self).__init__()
self.backbone = PSPNet()
if __name__ == '__main__':
dataset = {
phase: Salt_dataset('./data/train', 'images', 'masks',
phase == 'train', VAL_RATIO, transform)
for phase in ['train', 'val']
}
dataloader = {
phase: torch.utils.data.DataLoader(dataset[phase],
batch_size=BATCH_SIZE,
shuffle=True,
num_workers=8)
for phase in ['train', 'val']
}
# net = Unet(1,1,16).cuda()
net = PSPNet(1).cuda()
net = nn.DataParallel(net)
criterion = nn.BCELoss()
optimizer = optim.Adam(net.parameters(), lr=LEARNING_RATE)
scheduler = lr_scheduler.CosineAnnealingLR(optimizer,
T_max=40,
eta_min=1e-7)
# scheduler = lr_scheduler.ReduceLROnPlateau(optimizer)
min_iou = 0
print("train start")
for epoch in range(EPOCHS):
print('epoch: {}'.format(epoch))
train_running_loss = 0
train_running_dice_loss = 0
val_dataset = UPS31Dataset(image_dir=val_image_dir,
label_dir=val_label_dir,
transform=dataset_transforms['val'])
monitor_dataset = ImageFolder(monitor_image_dir,
transform=dataset_transforms['val'])
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print('Training images found: {}'.format(len(train_dataset)))
print('Validation images found: {}'.format(len(val_dataset)))
# print('Monitoring images found: {}'.format(len(monitor_dataset)))
print('Device: {}'.format(device))
# --- Model, loss and optimiser ---
model = PSPNet(num_classes)
# class_weights = enet_style_bounded_log_weighting(label_dir, output_height, output_width,
# num_classes)
# class_weights = median_frequency_balancing(label_dir, output_height, output_width,
# num_classes)
# class_weights = simple_bg_down_weight(0.1, num_classes)
# class_weights = class_weights.to(device)
# criterion = nn.CrossEntropyLoss(weight=None)
criterion = cross_entropy_with_aux_loss_pspnet(aux_weight=0.4,
class_weights=None)
optimiser = optim.Adam(model.parameters())
trained_model = train_model(model,
train_dataset,
val_dataset,
monitor_dataset,
def train(train_loader, val_loader, class_weights, class_encoding, visdom):
print("\nTraining...\n")
num_classes = len(class_encoding)
# Intialize ENet
# model = ENet(num_classes)
model = PSPNet(num_classes)
# Check if the network architecture is correct
print(model)
if torch.cuda.is_available():
if args.cuda:
device = 'cuda'
if torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model)
torch.cuda.empty_cache()
else:
device = 'cpu'
else:
device = 'cpu'
# We are going to use the CrossEntropyLoss loss function as it's most
# frequentely used in classification problems with multiple classes which
# fits the problem. This criterion combines LogSoftMax and NLLLoss.
criterion_seg = nn.CrossEntropyLoss(weight=class_weights)
#criterion_cls = nn.BCEWithLogitsLoss(weight=class_weights)
criterion_cls = nn.KLDivLoss(reduction='sum')
criterion = [criterion_seg, criterion_cls]
#authors used Adam as the optimizer
optimizer = optim.Adam(
model.parameters(),
lr=args.learning_rate,
weight_decay=args.weight_decay)
# Learning rate decay scheduler
lr_updater = lr_scheduler.StepLR(optimizer, args.lr_decay_epochs,
args.lr_decay)
# Evaluation metric
if args.ignore_unlabeled:
ignore_index = list(class_encoding).index('unlabeled')
else:
ignore_index = None
metric = IoU(num_classes, ignore_index=ignore_index)
# Optionally resume from a checkpoint
if args.resume:
model, optimizer, start_epoch, best_miou = utils.load_checkpoint(
model, optimizer, args.save_dir, args.name)
print("Resuming from model: Start epoch = {0} "
"| Best mean IoU = {1:.4f}".format(start_epoch, best_miou))
else:
start_epoch = 0
best_miou = 0
# Start Training
train = Train(model, train_loader, optimizer, criterion, metric, device)
val = Test(model, val_loader, criterion, metric, device)
for epoch in range(start_epoch, args.epochs):
print(">>>> [Epoch: {0:d}] Training".format(epoch))
lr_updater.step()
epoch_loss, (iou, miou) = train.run_epoch(args.print_step)
visdom.plot('train_loss', epoch_loss)
visdom.plot('train_miou', miou)
print(">>>> [Epoch: {0:d}] Avg. loss: {1:.4f} | Mean IoU: {2:.4f}".
format(epoch, epoch_loss, miou))
if (epoch + 1) % 10 == 0 or epoch + 1 == args.epochs:
print(">>>> [Epoch: {0:d}] Validation".format(epoch))
loss, (iou, miou) = val.run_epoch(args.print_step)
visdom.plot('val_loss', epoch_loss)
visdom.plot('val_miou', miou)
print(">>>> [Epoch: {0:d}] Avg. loss: {1:.4f} | Mean IoU: {2:.4f}".
format(epoch, loss, miou))
# Print per class IoU on last epoch or if best iou
if epoch + 1 == args.epochs or miou > best_miou:
for key, class_iou in zip(class_encoding.keys(), iou):
print("{0}: {1:.4f}".format(key, class_iou))
# Save the model if it's the best thus far
if miou > best_miou:
print("\nBest model thus far. Saving...\n")
best_miou = miou
utils.save_checkpoint(model, optimizer, epoch + 1, best_miou,
args)
return model
def create_model(cfg, name, backbone):
"""
freeze_encoder:
encoder_weights:
name: 'bisenetv1', 'bisenetv2', 'deeplabv3plus', 'fpn', 'hrnet_w18', 'hrnet_w32', 'hrnet_w48', 'lednet',
'linknet', 'nestnet', 'pspnet', 'u2net', 'unet', 'unet2plus', 'unet2s', 'unet3plus'
classes:
size:
batch_size:
channel:
backbone: 'vgg16', 'vgg19', 'resnet18', 'resnet34', 'resnet50', 'resnet101', 'resnet152', 'resnext50',
'resnext101', 'inceptionv3', 'inceptionresnet', 'densenet121', 'densenet169', 'densenet201',
'efficientnetb0', 'efficientnetb1', 'efficientnetb2', 'efficientnetb3', 'efficientnetb4',
'efficientnetb5', 'efficientnetb6', 'efficientnetb7', 'efficientnetl2', 'seresnext50', 'seresnext101',
'seresnet50', 'seresnet101', 'seresnet152', 'senet154'
activation:
"""
classes = cfg.DATASET.N_CLASSES
size = cfg.DATASET.SIZE
batch_size = cfg.TRAIN.BATCH_SIZE
channel = cfg.DATASET.CHANNELS
encoder_weights = cfg.TRAIN.ENCODER_WEIGHTS
activation = cfg.TRAIN.ACTIVATION
freeze_encoder = cfg.TRAIN.FREEZE_ENCODER
if name == 'bisenetv1':
model = BiseNetv1(cfg,
input_shape=(size[0], size[1], channel),
backbone_name=backbone,
encoder_weights=encoder_weights,
activation=activation)
elif name == 'bisenetv2':
model = BiseNetV2(classes,
size,
channel=channel,
activation=activation)
elif name == 'deeplabv3plus':
model = Deeplabv3(input_shape=(size[0], size[1], channel),
classes=classes,
backbone=backbone,
activation=activation)
elif name == 'fpn':
model = FPN(backbone,
input_shape=(size[0], size[1], channel),
classes=classes,
activation=activation,
encoder_weights=encoder_weights,
encoder_freeze=freeze_encoder)
elif name.startswith('hrnet'):
model = seg_hrnet(batch_size,
size[0],
size[1],
channel,
classes,
activation=activation,
hrnet_type=name)
elif name == 'lednet':
model = LEDNet(classes,
size[0],
size[1],
channel,
activation=activation)
elif name == 'linknet':
model = Linknet(backbone,
input_shape=(size[0], size[1], channel),
encoder_weights=encoder_weights,
freeze_encoder=freeze_encoder,
classes=classes,
activation=activation)
elif name == 'nestnet':
model = Nestnet(backbone,
input_shape=(size[0], size[1], channel),
encoder_weights=encoder_weights,
freeze_encoder=freeze_encoder,
classes=classes,
activation=activation)
elif name == 'pspnet':
model = PSPNet(backbone, (size[0], size[1], channel),
classes,
activation=activation,
encoder_weights=encoder_weights,
freeze_encoder=freeze_encoder)
elif name == 'u2net':
model = U2net((size[0], size[1], channel),
classes,
activation=activation)
elif name == 'unet':
model = Unet(backbone,
input_shape=(size[0], size[1], channel),
encoder_weights=encoder_weights,
freeze_encoder=freeze_encoder,
classes=classes,
activation=activation)
elif name == 'unet2plus':
model = Xnet(backbone, (size[0], size[1], channel),
classes=classes,
encoder_weights=encoder_weights,
freeze_encoder=freeze_encoder,
activation=activation)
elif name == 'unet2s':
model = Xnet_s(backbone, (size[0], size[1], channel),
classes=classes,
encoder_weights=encoder_weights,
activation=activation,
freeze_encoder=freeze_encoder)
elif name == 'unet3plus':
model = Unet3plus(name, (size[0], size[1], channel),
classes=classes,
activation=activation,
encoder_weights=encoder_weights,
freeze_encoder=freeze_encoder)
else:
raise TypeError('Unsupported model type')
return model