import os
import argparse
import numpy as np
from tqdm import tqdm
from collections import OrderedDict
import torch
import torch.nn.functional as F
from torch.utils import data
import torchvision.transforms as transform
from torch.nn.parallel.scatter_gather import gather
import encoding.utils as utils
from encoding.nn import SegmentationLosses, SyncBatchNorm
from encoding.parallel import DataParallelModel, DataParallelCriterion
from encoding.datasets import test_batchify_fn
from encoding.models.sseg import BaseNet
from modules.lseg_module import LSegModule
from utils import Resize
import cv2
import math
import types
import functools
import torchvision.transforms as torch_transforms
import copy
import itertools
from PIL import Image
import matplotlib.pyplot as plt
import clip
import matplotlib as mpl
import matplotlib.colors as mplc
import matplotlib.figure as mplfigure
import matplotlib.patches as mpatches
from matplotlib.backends.backend_agg import FigureCanvasAgg
from data import get_dataset
from additional_utils.encoding_models import MultiEvalModule as LSeg_MultiEvalModule
import torchvision.transforms as transforms
class Options:
def __init__(self):
parser = argparse.ArgumentParser(description="PyTorch Segmentation")
# model and dataset
parser.add_argument(
"--model", type=str, default="encnet", help="model name (default: encnet)"
)
parser.add_argument(
"--backbone",
type=str,
default="clip_vitl16_384",
help="backbone name (default: resnet50)",
)
parser.add_argument(
"--dataset",
type=str,
default="ade20k",
help="dataset name (default: pascal12)",
)
parser.add_argument(
"--workers", type=int, default=16, metavar="N", help="dataloader threads"
)
parser.add_argument(
"--base-size", type=int, default=520, help="base image size"
)
parser.add_argument(
"--crop-size", type=int, default=480, help="crop image size"
)
parser.add_argument(
"--train-split",
type=str,
default="train",
help="dataset train split (default: train)",
)
# training hyper params
parser.add_argument(
"--aux", action="store_true", default=False, help="Auxilary Loss"
)
parser.add_argument(
"--se-loss",
action="store_true",
default=False,
help="Semantic Encoding Loss SE-loss",
)
parser.add_argument(
"--se-weight", type=float, default=0.2, help="SE-loss weight (default: 0.2)"
)
parser.add_argument(
"--batch-size",
type=int,
default=16,
metavar="N",
help="input batch size for \
training (default: auto)",
)
parser.add_argument(
"--test-batch-size",
type=int,
default=16,
metavar="N",
help="input batch size for \
testing (default: same as batch size)",
)
# cuda, seed and logging
parser.add_argument(
"--no-cuda",
action="store_true",
default=False,
help="disables CUDA training",
)
parser.add_argument(
"--seed", type=int, default=1, metavar="S", help="random seed (default: 1)"
)
parser.add_argument(
"--weights", type=str, default=None, help="checkpoint to test"
)
parser.add_argument(
"--eval", action="store_true", default=False, help="evaluating mIoU"
)
parser.add_argument(
"--export",
type=str,
default=None,
help="put the path to resuming file if needed",
)
parser.add_argument(
"--acc-bn",
action="store_true",
default=False,
help="Re-accumulate BN statistics",
)
parser.add_argument(
"--test-val",
action="store_true",
default=False,
help="generate masks on val set",
)
parser.add_argument(
"--no-val",
action="store_true",
default=False,
help="skip validation during training",
)
parser.add_argument(
"--module",
default='lseg',
help="select model definition",
)
# test option
parser.add_argument(
"--data-path", type=str, default=None, help="path to test image folder"
)
parser.add_argument(
"--no-scaleinv",
dest="scale_inv",
default=True,
action="store_false",
help="turn off scaleinv layers",
)
parser.add_argument(
"--widehead", default=False, action="store_true", help="wider output head"
)
parser.add_argument(
"--widehead_hr",
default=False,
action="store_true",
help="wider output head",
)
parser.add_argument(
"--ignore_index",
type=int,
default=-1,
help="numeric value of ignore label in gt",
)
parser.add_argument(
"--label_src",
type=str,
default="default",
help="how to get the labels",
)
parser.add_argument(
"--jobname",
type=str,
default="default",
help="select which dataset",
)
parser.add_argument(
"--no-strict",
dest="strict",
default=True,
action="store_false",
help="no-strict copy the model",
)
parser.add_argument(
"--arch_option",
type=int,
default=0,
help="which kind of architecture to be used",
)
parser.add_argument(
"--block_depth",
type=int,
default=0,
help="how many blocks should be used",
)
parser.add_argument(
"--activation",
choices=['lrelu', 'tanh'],
default="lrelu",
help="use which activation to activate the block",
)
self.parser = parser
def parse(self):
args = self.parser.parse_args()
args.cuda = not args.no_cuda and torch.cuda.is_available()
print(args)
return args
def test(args):
module = LSegModule.load_from_checkpoint(
checkpoint_path=args.weights,
data_path=args.data_path,
dataset=args.dataset,
backbone=args.backbone,
aux=args.aux,
num_features=256,
aux_weight=0,
se_loss=False,
se_weight=0,
base_lr=0,
batch_size=1,
max_epochs=0,
ignore_index=args.ignore_index,
dropout=0.0,
scale_inv=args.scale_inv,
augment=False,
no_batchnorm=False,
widehead=args.widehead,
widehead_hr=args.widehead_hr,
map_locatin="cpu",
arch_option=args.arch_option,
strict=args.strict,
block_depth=args.block_depth,
activation=args.activation,
)
input_transform = module.val_transform
num_classes = module.num_classes
# dataset
testset = get_dataset(
args.dataset,
root=args.data_path,
split="val",
mode="testval",
transform=input_transform,
)
# dataloader
loader_kwargs = (
{"num_workers": args.workers, "pin_memory": True} if args.cuda else {}
)
test_data = data.DataLoader(
testset,
batch_size=args.test_batch_size,
drop_last=False,
shuffle=False,
collate_fn=test_batchify_fn,
**loader_kwargs
)
if isinstance(module.net, BaseNet):
model = module.net
else:
model = module
model = model.eval()
model = model.cpu()
print(model)
if args.acc_bn:
from encoding.utils.precise_bn import update_bn_stats
data_kwargs = {
"transform": input_transform,
"base_size": args.base_size,
"crop_size": args.crop_size,
}
trainset = get_dataset(
args.dataset, split=args.train_split, mode="train", **data_kwargs
)
trainloader = data.DataLoader(
ReturnFirstClosure(trainset),
root=args.data_path,
batch_size=args.batch_size,
drop_last=True,
shuffle=True,
**loader_kwargs
)
print("Reseting BN statistics")
model.cuda()
update_bn_stats(model, trainloader)
if args.export:
torch.save(model.state_dict(), args.export + ".pth")
return
scales = (
[0.75, 1.0, 1.25, 1.5, 1.75, 2.0, 2.25]
if args.dataset == "citys"
else [0.5, 0.75, 1.0, 1.25, 1.5, 1.75]
)
evaluator = LSeg_MultiEvalModule(
model, num_classes, scales=scales, flip=True
).cuda()
evaluator.eval()
metric = utils.SegmentationMetric(testset.num_class)
tbar = tqdm(test_data)
f = open("logs/log_test_{}_{}.txt".format(args.jobname, args.dataset), "a+")
per_class_iou = np.zeros(testset.num_class)
cnt = 0
for i, (image, dst) in enumerate(tbar):
if args.eval:
with torch.no_grad():
if False:
sample = {"image": image[0].cpu().permute(1, 2, 0).numpy()}
out = torch.zeros(
1, testset.num_class, image[0].shape[1], image[0].shape[2]
).cuda()
H, W = image[0].shape[1], image[0].shape[2]
for scale in scales:
long_size = int(math.ceil(520 * scale))
if H > W:
height = long_size
width = int(1.0 * W * long_size / H + 0.5)
short_size = width
else:
width = long_size
height = int(1.0 * H * long_size / W + 0.5)
short_size = height
rs = Resize(
width,
height,
resize_target=False,
keep_aspect_ratio=True,
ensure_multiple_of=32,
resize_method="minimal",
image_interpolation_method=cv2.INTER_AREA,
)
inf_image = (
torch.from_numpy(rs(sample)["image"])
.cuda()
.permute(2, 0, 1)
.unsqueeze(0)
)
inf_image = torch.cat((inf_image, torch.fliplr(inf_image)), 0)
try:
pred = model(inf_image)
except:
print(H, W, sz, i)
exit()
pred0 = F.softmax(pred[0], dim=1)
pred1 = F.softmax(pred[1], dim=1)
pred = pred0 + 0.2 * pred1
out += F.interpolate(
pred.sum(0, keepdim=True),
(out.shape[2], out.shape[3]),
mode="bilinear",
align_corners=True,
)
predicts = [out]
else:
predicts = evaluator.parallel_forward(image)
metric.update(dst, predicts)
pixAcc, mIoU = metric.get()
_, _, total_inter, total_union = metric.get_all()
per_class_iou += 1.0 * total_inter / (np.spacing(1) + total_union)
cnt+=1
tbar.set_description("pixAcc: %.4f, mIoU: %.4f" % (pixAcc, mIoU))
else:
with torch.no_grad():
outputs = evaluator.parallel_forward(image)
predicts = [
testset.make_pred(torch.max(output, 1)[1].cpu().numpy())
for output in outputs
]
# output folder
outdir = "outdir_ours"
if not os.path.exists(outdir):
os.makedirs(outdir)
for predict, impath in zip(predicts, dst):
mask = utils.get_mask_pallete(predict, args.dataset)
outname = os.path.splitext(impath)[0] + ".png"
mask.save(os.path.join(outdir, outname))
if args.eval:
each_classes_iou = per_class_iou/cnt
print("pixAcc: %.4f, mIoU: %.4f" % (pixAcc, mIoU))
print(each_classes_iou)
f.write("dataset {} ==> pixAcc: {:.4f}, mIoU: {:.4f}\n".format(args.dataset, pixAcc, mIoU))
for per_iou in each_classes_iou: f.write('{:.4f}, '.format(per_iou))
f.write('\n')
class ReturnFirstClosure(object):
def __init__(self, data):
self._data = data
def __len__(self):
return len(self._data)
def __getitem__(self, idx):
outputs = self._data[idx]
return outputs[0]
if __name__ == "__main__":
args = Options().parse()
torch.manual_seed(args.seed)
args.test_batch_size = torch.cuda.device_count()
test(args)