def main(args):
dataset = load_config(args.dataset)
num_classes = len(dataset["common"]["classes"])
net = UNet(num_classes)
chkpt = torch.load(args.checkpoint, map_location="cpu")
net.load_state_dict(chkpt)
net = torch.nn.DataParallel(net)
# Todo: make input channels configurable, not hard-coded to three channels for RGB
batch = torch.autograd.Variable(torch.randn(1, 3, args.image_size, args.image_size))
torch.onnx.export(net, batch, args.model)
def main(args):
config = load_config(args.config)
if args.type == "onnx":
os.environ["CUDA_VISIBLE_DEVICES"] = ""
# Workaround: PyTorch ONNX, DataParallel with GPU issue, cf https://github.com/pytorch/pytorch/issues/5315
num_classes = len(config["classes"]["classes"])
num_channels = 0
for channel in config["channels"]:
num_channels += len(channel["bands"])
export_channels = num_channels if not args.export_channels else args.export_channels
assert num_channels >= export_channels, "Will be hard indeed, to export more channels than thoses dataset provide"
def map_location(storage, _):
return storage.cpu()
net = UNet(num_classes, num_channels=num_channels).to("cpu")
chkpt = torch.load(args.checkpoint, map_location=map_location)
net = torch.nn.DataParallel(net)
net.load_state_dict(chkpt["state_dict"])
if export_channels < num_channels:
weights = torch.zeros((64, export_channels, 7, 7))
weights.data = net.module.resnet.conv1.weight.data[:, :
export_channels, :, :]
net.module.resnet.conv1 = nn.Conv2d(num_channels,
64,
kernel_size=7,
stride=2,
padding=3,
bias=False)
net.module.resnet.conv1.weight = nn.Parameter(weights)
if args.type == "onnx":
batch = torch.autograd.Variable(
torch.randn(1, export_channels, args.image_size, args.image_size))
torch.onnx.export(net, batch, args.out)
elif args.type == "pth":
states = {
"epoch": chkpt["epoch"],
"state_dict": net.state_dict(),
"optimizer": chkpt["optimizer"]
}
torch.save(states, args.out)
def net_from_chkpt_(self):
def map_location(storage, _):
return storage.cuda() if self.cuda else storage.cpu()
# https://github.com/pytorch/pytorch/issues/7178
chkpt = torch.load(self.checkpoint, map_location=map_location)
num_classes = len(self.dataset['common']['classes'])
net = UNet(num_classes).to(self.device)
net = nn.DataParallel(net)
if self.cuda:
torch.backends.cudnn.benchmark = True
net.load_state_dict(chkpt)
net.eval()
return net
def net_from_chkpt_(self):
def map_location(storage, _):
return storage.cuda() if self.cuda else storage.cpu()
# https://github.com/pytorch/pytorch/issues/7178
chkpt = torch.load(self.checkpoint, map_location=map_location)
num_classes = len(self.dataset['common']['classes'])
net = UNet(num_classes).to(self.device)
net = nn.DataParallel(net)
if self.cuda:
torch.backends.cudnn.benchmark = True
net.load_state_dict(chkpt)
net.eval()
return net
def main(args):
model = load_config(args.model)
dataset = load_config(args.dataset)
device = torch.device("cuda" if model["common"]["cuda"] else "cpu")
if model["common"]["cuda"] and not torch.cuda.is_available():
sys.exit("Error: CUDA requested but not available")
os.makedirs(model["common"]["checkpoint"], exist_ok=True)
num_classes = len(dataset["common"]["classes"])
net = UNet(num_classes)
net = DataParallel(net)
net = net.to(device)
if model["common"]["cuda"]:
torch.backends.cudnn.benchmark = True
try:
weight = torch.Tensor(dataset["weights"]["values"])
except KeyError:
if model["opt"]["loss"] in ("CrossEntropy", "mIoU", "Focal"):
sys.exit(
"Error: The loss function used, need dataset weights values")
optimizer = Adam(net.parameters(),
lr=model["opt"]["lr"],
weight_decay=model["opt"]["decay"])
resume = 0
if args.checkpoint:
def map_location(storage, _):
return storage.cuda() if model["common"]["cuda"] else storage.cpu()
# https://github.com/pytorch/pytorch/issues/7178
chkpt = torch.load(args.checkpoint, map_location=map_location)
net.load_state_dict(chkpt["state_dict"])
if args.resume:
optimizer.load_state_dict(chkpt["optimizer"])
resume = chkpt["epoch"]
if model["opt"]["loss"] == "CrossEntropy":
criterion = CrossEntropyLoss2d(weight=weight).to(device)
elif model["opt"]["loss"] == "mIoU":
criterion = mIoULoss2d(weight=weight).to(device)
elif model["opt"]["loss"] == "Focal":
criterion = FocalLoss2d(weight=weight).to(device)
elif model["opt"]["loss"] == "Lovasz":
criterion = LovaszLoss2d().to(device)
else:
sys.exit("Error: Unknown [opt][loss] value !")
train_loader, val_loader = get_dataset_loaders(model, dataset,
args.workers)
num_epochs = model["opt"]["epochs"]
if resume >= num_epochs:
sys.exit(
"Error: Epoch {} set in {} already reached by the checkpoint provided"
.format(num_epochs, args.model))
history = collections.defaultdict(list)
log = Log(os.path.join(model["common"]["checkpoint"], "log"))
log.log("--- Hyper Parameters on Dataset: {} ---".format(
dataset["common"]["dataset"]))
log.log("Batch Size:\t {}".format(model["common"]["batch_size"]))
log.log("Image Size:\t {}".format(model["common"]["image_size"]))
log.log("Learning Rate:\t {}".format(model["opt"]["lr"]))
log.log("Weight Decay:\t {}".format(model["opt"]["decay"]))
log.log("Loss function:\t {}".format(model["opt"]["loss"]))
if "weight" in locals():
log.log("Weights :\t {}".format(dataset["weights"]["values"]))
log.log("---")
for epoch in range(resume, num_epochs):
log.log("Epoch: {}/{}".format(epoch + 1, num_epochs))
train_hist = train(train_loader, num_classes, device, net, optimizer,
criterion)
log.log(
"Train loss: {:.4f}, mIoU: {:.3f}, {} IoU: {:.3f}, MCC: {:.3f}".
format(
train_hist["loss"],
train_hist["miou"],
dataset["common"]["classes"][1],
train_hist["fg_iou"],
train_hist["mcc"],
))
for k, v in train_hist.items():
history["train " + k].append(v)
val_hist = validate(val_loader, num_classes, device, net, criterion)
log.log(
"Validate loss: {:.4f}, mIoU: {:.3f}, {} IoU: {:.3f}, MCC: {:.3f}".
format(val_hist["loss"], val_hist["miou"],
dataset["common"]["classes"][1], val_hist["fg_iou"],
val_hist["mcc"]))
for k, v in val_hist.items():
history["val " + k].append(v)
visual = "history-{:05d}-of-{:05d}.png".format(epoch + 1, num_epochs)
plot(os.path.join(model["common"]["checkpoint"], visual), history)
checkpoint = "checkpoint-{:05d}-of-{:05d}.pth".format(
epoch + 1, num_epochs)
states = {
"epoch": epoch + 1,
"state_dict": net.state_dict(),
"optimizer": optimizer.state_dict()
}
torch.save(states,
os.path.join(model["common"]["checkpoint"], checkpoint))
def main(args):
model = load_config(args.model)
dataset = load_config(args.dataset)
device = torch.device("cuda" if model["common"]["cuda"] else "cpu")
if model["common"]["cuda"] and not torch.cuda.is_available():
sys.exit("Error: CUDA requested but not available")
# if args.batch_size < 2:
# sys.exit('Error: PSPNet requires more than one image for BatchNorm in Pyramid Pooling')
os.makedirs(model["common"]["checkpoint"], exist_ok=True)
num_classes = len(dataset["common"]["classes"])
net = UNet(num_classes)
net = DataParallel(net)
net = net.to(device)
if model["common"]["cuda"]:
torch.backends.cudnn.benchmark = True
if args.checkpoint:
def map_location(storage, _):
return storage.cuda() if model["common"]["cuda"] else storage.cpu()
# https://github.com/pytorch/pytorch/issues/7178
chkpt = torch.load(args.checkpoint, map_location=map_location)
net.load_state_dict(chkpt)
optimizer = Adam(net.parameters(),
lr=model["opt"]["lr"],
weight_decay=model["opt"]["decay"])
weight = torch.Tensor(dataset["weights"]["values"])
criterion = CrossEntropyLoss2d(weight=weight).to(device)
# criterion = FocalLoss2d(weight=weight).to(device)
train_loader, val_loader = get_dataset_loaders(model, dataset)
num_epochs = model["opt"]["epochs"]
history = collections.defaultdict(list)
for epoch in range(num_epochs):
print("Epoch: {}/{}".format(epoch + 1, num_epochs))
train_hist = train(train_loader, num_classes, device, net, optimizer,
criterion)
print("Train loss: {:.4f}, mean IoU: {:.4f}".format(
train_hist["loss"], train_hist["iou"]))
for k, v in train_hist.items():
history["train " + k].append(v)
val_hist = validate(val_loader, num_classes, device, net, criterion)
print("Validate loss: {:.4f}, mean IoU: {:.4f}".format(
val_hist["loss"], val_hist["iou"]))
for k, v in val_hist.items():
history["val " + k].append(v)
visual = "history-{:05d}-of-{:05d}.png".format(epoch + 1, num_epochs)
plot(os.path.join(model["common"]["checkpoint"], visual), history)
checkpoint = "checkpoint-{:05d}-of-{:05d}.pth".format(
epoch + 1, num_epochs)
torch.save(net.state_dict(),
os.path.join(model["common"]["checkpoint"], checkpoint))
def main(args):
model = load_config(args.model)
dataset = load_config(args.dataset)
cuda = model["common"]["cuda"]
device = torch.device("cuda" if cuda else "cpu")
def map_location(storage, _):
return storage.cuda() if cuda else storage.cpu()
if cuda and not torch.cuda.is_available():
sys.exit("Error: CUDA requested but not available")
num_classes = len(dataset["common"]["classes"])
# https://github.com/pytorch/pytorch/issues/7178
chkpt = torch.load(args.checkpoint, map_location=map_location)
net = UNet(num_classes).to(device)
net = nn.DataParallel(net)
if cuda:
torch.backends.cudnn.benchmark = True
net.load_state_dict(chkpt["state_dict"])
net.eval()
mean, std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
transform = Compose([
ConvertImageMode(mode="RGB"),
ImageToTensor(),
Normalize(mean=mean, std=std)
])
directory = BufferedSlippyMapDirectory(args.tiles,
transform=transform,
size=args.tile_size,
overlap=args.overlap)
loader = DataLoader(directory,
batch_size=args.batch_size,
num_workers=args.workers)
# don't track tensors with autograd during prediction
with torch.no_grad():
for images, tiles in tqdm(loader,
desc="Eval",
unit="batch",
ascii=True):
images = images.to(device)
outputs = net(images)
# manually compute segmentation mask class probabilities per pixel
probs = nn.functional.softmax(outputs, dim=1).data.cpu().numpy()
for tile, prob in zip(tiles, probs):
x, y, z = list(map(int, tile))
# we predicted on buffered tiles; now get back probs for original image
prob = directory.unbuffer(prob)
# Quantize the floating point probabilities in [0,1] to [0,255] and store
# a single-channel `.png` file with a continuous color palette attached.
assert prob.shape[
0] == 2, "single channel requires binary model"
assert np.allclose(
np.sum(prob, axis=0), 1.
), "single channel requires probabilities to sum up to one"
foreground = prob[1:, :, :]
anchors = np.linspace(0, 1, 256)
quantized = np.digitize(foreground, anchors).astype(np.uint8)
palette = continuous_palette_for_color("pink", 256)
out = Image.fromarray(quantized.squeeze(), mode="P")
out.putpalette(palette)
os.makedirs(os.path.join(args.probs, str(z), str(x)),
exist_ok=True)
path = os.path.join(args.probs, str(z), str(x),
str(y) + ".png")
out.save(path, optimize=True)
def main(args):
model = load_config(args.model)
dataset = load_config(args.dataset)
cuda = model['common']['cuda']
device = torch.device('cuda' if cuda else 'cpu')
def map_location(storage, _):
return storage.cuda() if cuda else storage.cpu()
if cuda and not torch.cuda.is_available():
sys.exit('Error: CUDA requested but not available')
num_classes = len(dataset['common']['classes'])
# https://github.com/pytorch/pytorch/issues/7178
chkpt = torch.load(args.checkpoint, map_location=map_location)
net = UNet(num_classes).to(device)
net = nn.DataParallel(net)
if cuda:
torch.backends.cudnn.benchmark = True
net.load_state_dict(chkpt)
net.eval()
transform = Compose([
ConvertImageMode(mode='RGB'),
ImageToTensor(),
Normalize(mean=dataset['stats']['mean'], std=dataset['stats']['std'])
])
directory = BufferedSlippyMapDirectory(args.tiles, transform=transform, size=args.tile_size, overlap=args.overlap)
loader = DataLoader(directory, batch_size=args.batch_size)
# don't track tensors with autograd during prediction
with torch.no_grad():
for images, tiles in tqdm(loader, desc='Eval', unit='batch', ascii=True):
images = images.to(device)
outputs = net(images)
# manually compute segmentation mask class probabilities per pixel
probs = nn.functional.softmax(outputs, dim=1).data.cpu().numpy()
for tile, prob in zip(tiles, probs):
x, y, z = list(map(int, tile))
# we predicted on buffered tiles; now get back probs for original image
prob = directory.unbuffer(prob)
# Quantize the floating point probabilities in [0,1] to [0,255] and store
# a single-channel `.png` file with a continuous color palette attached.
assert prob.shape[0] == 2, 'single channel requires binary model'
assert np.allclose(np.sum(prob, axis=0), 1.), 'single channel requires probabilities to sum up to one'
foreground = prob[1:, :, :]
anchors = np.linspace(0, 1, 256)
quantized = np.digitize(foreground, anchors).astype(np.uint8)
palette = continuous_palette_for_color('pink', 256)
out = Image.fromarray(quantized.squeeze(), mode='P')
out.putpalette(palette)
os.makedirs(os.path.join(args.probs, str(z), str(x)), exist_ok=True)
path = os.path.join(args.probs, str(z), str(x), str(y) + '.png')
out.save(path, optimize=True)
def main(args):
model = load_config(args.model)
dataset = load_config(args.dataset)
device = torch.device('cuda' if model['common']['cuda'] else 'cpu')
if model['common']['cuda'] and not torch.cuda.is_available():
sys.exit('Error: CUDA requested but not available')
# if args.batch_size < 2:
# sys.exit('Error: PSPNet requires more than one image for BatchNorm in Pyramid Pooling')
os.makedirs(model['common']['checkpoint'], exist_ok=True)
num_classes = len(dataset['common']['classes'])
net = UNet(num_classes).to(device)
if args.resume:
path = os.path.join(model['common']['checkpoint'], args.resume)
cuda = model['common']['cuda']
def map_location(storage, _):
return storage.cuda() if cuda else storage.cpu()
chkpt = torch.load(path, map_location=map_location)
net.load_state_dict(chkpt)
resume_at_epoch = int(args.resume[11:16])
else:
resume_at_epoch = 0
if model['common']['cuda']:
torch.backends.cudnn.benchmark = True
net = DataParallel(net)
optimizer = SGD(net.parameters(),
lr=model['opt']['lr'],
momentum=model['opt']['momentum'])
scheduler = MultiStepLR(optimizer,
milestones=model['opt']['milestones'],
gamma=model['opt']['gamma'])
weight = torch.Tensor(dataset['weights']['values'])
for i in range(resume_at_epoch):
scheduler.step()
criterion = CrossEntropyLoss2d(weight=weight).to(device)
# criterion = FocalLoss2d(weight=weight).to(device)
train_loader, val_loader = get_dataset_loaders(model, dataset)
num_epochs = model['opt']['epochs']
history = collections.defaultdict(list)
for epoch in range(resume_at_epoch, num_epochs):
print('Epoch: {}/{}'.format(epoch + 1, num_epochs))
train_hist = train(train_loader, num_classes, device, net, optimizer,
scheduler, criterion)
print('Train loss: {:.4f}, mean IoU: {:.4f}'.format(
train_hist['loss'], train_hist['iou']))
for k, v in train_hist.items():
history['train ' + k].append(v)
val_hist = validate(val_loader, num_classes, device, net, criterion)
print('Validate loss: {:.4f}, mean IoU: {:.4f}'.format(
val_hist['loss'], val_hist['iou']))
for k, v in val_hist.items():
history['val ' + k].append(v)
visual = 'history-{:05d}-of-{:05d}.png'.format(epoch + 1, num_epochs)
plot(os.path.join(model['common']['checkpoint'], visual), history)
checkpoint = 'checkpoint-{:05d}-of-{:05d}.pth'.format(
epoch + 1, num_epochs)
torch.save(net.state_dict(),
os.path.join(model['common']['checkpoint'], checkpoint))
def main(args):
model = load_config(args.model)
dataset = load_config(args.dataset)
device = torch.device('cuda' if model['common']['cuda'] else 'cpu')
if model['common']['cuda'] and not torch.cuda.is_available():
sys.exit('Error: CUDA requested but not available')
# if args.batch_size < 2:
# sys.exit('Error: PSPNet requires more than one image for BatchNorm in Pyramid Pooling')
os.makedirs(model['common']['checkpoint'], exist_ok=True)
num_classes = len(dataset['common']['classes'])
net = UNet(num_classes).to(device)
if args.resume:
path = os.path.join(model['common']['checkpoint'], args.resume)
cuda = model['common']['cuda']
def map_location(storage, _):
return storage.cuda() if cuda else storage.cpu()
chkpt = torch.load(path, map_location=map_location)
net.load_state_dict(chkpt)
resume_at_epoch = int(args.resume[11:16])
else:
resume_at_epoch = 0
if model['common']['cuda']:
torch.backends.cudnn.benchmark = True
net = DataParallel(net)
optimizer = SGD(net.parameters(), lr=model['opt']['lr'], momentum=model['opt']['momentum'])
scheduler = MultiStepLR(optimizer, milestones=model['opt']['milestones'], gamma=model['opt']['gamma'])
weight = torch.Tensor(dataset['weights']['values'])
for i in range(resume_at_epoch):
scheduler.step()
criterion = CrossEntropyLoss2d(weight=weight).to(device)
# criterion = FocalLoss2d(weight=weight).to(device)
train_loader, val_loader = get_dataset_loaders(model, dataset)
num_epochs = model['opt']['epochs']
history = collections.defaultdict(list)
for epoch in range(resume_at_epoch, num_epochs):
print('Epoch: {}/{}'.format(epoch + 1, num_epochs))
train_hist = train(train_loader, num_classes, device, net, optimizer, scheduler, criterion)
print('Train loss: {:.4f}, mean IoU: {:.4f}'.format(train_hist['loss'], train_hist['iou']))
for k, v in train_hist.items():
history['train ' + k].append(v)
val_hist = validate(val_loader, num_classes, device, net, criterion)
print('Validate loss: {:.4f}, mean IoU: {:.4f}'.format(val_hist['loss'], val_hist['iou']))
for k, v in val_hist.items():
history['val ' + k].append(v)
visual = 'history-{:05d}-of-{:05d}.png'.format(epoch + 1, num_epochs)
plot(os.path.join(model['common']['checkpoint'], visual), history)
checkpoint = 'checkpoint-{:05d}-of-{:05d}.pth'.format(epoch + 1, num_epochs)
torch.save(net.state_dict(), os.path.join(model['common']['checkpoint'], checkpoint))
def main(args):
config = load_config(args.config)
num_classes = len(config["classes"]["titles"])
if torch.cuda.is_available():
device = torch.device("cuda")
torch.backends.cudnn.benchmark = True
else:
device = torch.device("cpu")
def map_location(storage, _):
return storage.cuda() if torch.cuda.is_available() else storage.cpu()
# https://github.com/pytorch/pytorch/issues/7178
chkpt = torch.load(args.checkpoint, map_location=map_location)
net = UNet(num_classes).to(device)
net = nn.DataParallel(net)
net.load_state_dict(chkpt["state_dict"])
net.eval()
mean, std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
transform = Compose([ImageToTensor(), Normalize(mean=mean, std=std)])
directory = BufferedSlippyMapDirectory(args.tiles,
transform=transform,
size=args.tile_size,
overlap=args.overlap)
loader = DataLoader(directory,
batch_size=args.batch_size,
num_workers=args.workers)
if args.masks_output:
palette = make_palette(config["classes"]["colors"][0],
config["classes"]["colors"][1])
else:
palette = continuous_palette_for_color("pink", 256)
# don't track tensors with autograd during prediction
with torch.no_grad():
for images, tiles in tqdm(loader,
desc="Eval",
unit="batch",
ascii=True):
images = images.to(device)
outputs = net(images)
# manually compute segmentation mask class probabilities per pixel
probs = nn.functional.softmax(outputs, dim=1).data.cpu().numpy()
for tile, prob in zip(tiles, probs):
x, y, z = list(map(int, tile))
# we predicted on buffered tiles; now get back probs for original image
prob = directory.unbuffer(prob)
assert prob.shape[
0] == 2, "single channel requires binary model"
assert np.allclose(
np.sum(prob, axis=0), 1.0
), "single channel requires probabilities to sum up to one"
if args.masks_output:
image = np.around(prob[1:, :, :]).astype(
np.uint8).squeeze()
else:
image = (prob[1:, :, :] * 255).astype(np.uint8).squeeze()
out = Image.fromarray(image, mode="P")
out.putpalette(palette)
os.makedirs(os.path.join(args.probs, str(z), str(x)),
exist_ok=True)
path = os.path.join(args.probs, str(z), str(x),
str(y) + ".png")
out.save(path, optimize=True)
if args.web_ui:
template = "leaflet.html" if not args.web_ui_template else args.web_ui_template
tiles = [tile for tile, _ in tiles_from_slippy_map(args.tiles)]
web_ui(args.probs, args.web_ui, tiles, tiles, "png", template)
