def main(dataset, image_path, model_path, cuda, crf):
CONFIG = {
'voc12': {
'path_pytorch_model':
'data/models/deeplab_resnet101/voc12/deeplabv2_resnet101_VOC2012.pth',
'label_list': 'data/datasets/voc12/labels.txt',
'n_classes': 21,
'n_blocks': [3, 4, 23, 3],
'pyramids': [6, 3, 2, 1],
'image': {
'size': {
'train': 473,
'test': 473,
},
'mean': {
'R': 122.675,
'G': 116.669,
'B': 104.008,
}
},
},
'cocostuff': {
'path_pytorch_model':
'data/models/deeplab_resnet101/cocostuff_rgb/deeplabv2_resnet101_VOC2012.pth',
'label_list': 'data/datasets/voc12/labels.txt',
'n_classes': 183,
'n_blocks': [3, 4, 23, 3],
'pyramids': [6, 3, 2, 1],
'image': {
'size': {
'train': 321,
'test': 513,
},
'mean': {
'R': 122.675,
'G': 116.669,
'B': 104.008,
}
},
},
}.get(dataset)
cuda = cuda and torch.cuda.is_available()
# Label list
with open(CONFIG['label_list']) as f:
classes = {}
for label in f:
label = label.rstrip().split('\t')
classes[int(label[0])] = label[1].split(',')[0]
# Load a model
if model_path is None:
state_dict = torch.load(CONFIG['path_pytorch_model'])
else:
state_dict = torch.load(model_path)
# Model
model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG['n_classes'])
model.load_state_dict(state_dict)
model.eval()
if cuda:
model.cuda()
image_size = (CONFIG['image']['size']['test'],
CONFIG['image']['size']['test'])
# Image preprocessing
image = cv2.imread(image_path, cv2.IMREAD_COLOR).astype(float)
image = cv2.resize(image, image_size)
image_original = image.astype(np.uint8)
image -= np.array([
float(CONFIG['image']['mean']['B']),
float(CONFIG['image']['mean']['G']),
float(CONFIG['image']['mean']['R'])
])
image = torch.from_numpy(image.transpose(2, 0, 1)).float().unsqueeze(0)
image = image.cuda() if cuda else image
# Inference
output = model(Variable(image, volatile=True))
output = F.upsample(output, size=image_size, mode='bilinear')
output = F.softmax(output, dim=1)
output = output.data.cpu().numpy()[0]
if crf:
output = dense_crf(image_original, output)
labelmap = np.argmax(output.transpose(1, 2, 0), axis=2)
labels = np.unique(labelmap)
# Show results
rows = np.floor(np.sqrt(len(labels) + 1))
cols = np.ceil((len(labels) + 1) / rows)
plt.figure(figsize=(10, 10))
ax = plt.subplot(rows, cols, 1)
ax.set_title('Input image')
ax.imshow(image_original[:, :, ::-1])
ax.set_xticks([])
ax.set_yticks([])
for i, label in enumerate(labels):
print '{0:3d}: {1}'.format(label, classes[label])
mask = labelmap == label
ax = plt.subplot(rows, cols, i + 2)
ax.set_title(classes[label])
ax.imshow(np.dstack((mask, ) * 3) * image_original[:, :, ::-1])
ax.set_xticks([])
ax.set_yticks([])
plt.show()
def main(dataset):
WHITELIST = ["kernel_size", "stride", "padding", "dilation", "eps", "momentum"]
CONFIG = {
"voc12": {
"path_caffe_model": "data/models/deeplab_resnet101/voc12/train2_iter_20000.caffemodel",
"path_pytorch_model": "data/models/deeplab_resnet101/voc12/deeplabv2_resnet101_VOC2012.pth",
"n_classes": 21,
},
"init": {
"path_caffe_model": "data/models/deeplab_resnet101/coco_init/init.caffemodel",
"path_pytorch_model": "data/models/deeplab_resnet101/coco_init/deeplabv2_resnet101_COCO_init.pth",
"n_classes": 91,
},
}.get(dataset)
params = parse_caffemodel(CONFIG["path_caffe_model"])
model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG["n_classes"])
model.eval()
own_state = model.state_dict()
rel_tol = 1e-7
state_dict = OrderedDict()
for layer_name, layer_dict in params.items():
for param_name, values in layer_dict.items():
if param_name in WHITELIST and dataset != "coco_init" and dataset != "init":
attribute = translate_layer_name(layer_name)
attribute = eval("model." + attribute + "." + param_name)
if isinstance(attribute, tuple):
assert (
attribute[0] == values
), "Inconsistent values: {}@{}, {}@{}".format(
attribute[0],
translate_layer_name(layer_name) + "." + param_name,
values,
layer_name,
)
else:
assert (
abs(attribute - values) < rel_tol
), "Inconsistent values: {}@{}, {}@{}".format(
attribute,
translate_layer_name(layer_name) + "." + param_name,
values,
layer_name,
)
print(
layer_name.ljust(20),
"->",
param_name,
attribute,
values,
": Checked!",
)
continue
param_name = translate_layer_name(layer_name) + "." + param_name
if param_name in own_state:
values = torch.FloatTensor(values)
values = values.view_as(own_state[param_name])
state_dict[param_name] = values
print(layer_name.ljust(20), "->", param_name, ": Copied!")
try:
print("\033[32mVerify the converted model\033[00m")
model.load_state_dict(state_dict)
except:
import traceback
traceback.print_exc()
print("\033[32mVerify with ignoring ASPP (strict=False)\033[00m")
model.load_state_dict(state_dict, strict=False)
print("Saving to", CONFIG["path_pytorch_model"])
torch.save(state_dict, CONFIG["path_pytorch_model"])
def main(config, cuda, gpu):
# Configuration
CONFIG = Dict(yaml.load(open(config)))
# CUDA check
cuda = cuda and torch.cuda.is_available()
if cuda:
gpu_ids = [int(string) for string in gpu.split(',')]
current_device = torch.cuda.current_device()
print('Running on', torch.cuda.get_device_name(current_device),
gpu_ids)
# Dataset
dataset = CocoStuff10k(
root=CONFIG.ROOT,
split='train',
image_size=513,
crop_size=CONFIG.IMAGE.SIZE.TRAIN,
scale=True,
flip=True,
)
# DataLoader
loader = torch.utils.data.DataLoader(
dataset=dataset,
batch_size=CONFIG.BATCH_SIZE,
num_workers=CONFIG.NUM_WORKERS,
shuffle=True,
)
loader_iter = iter(loader)
# Model
model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG.N_CLASSES)
state_dict = torch.load(CONFIG.INIT_MODEL)
model.load_state_dict(state_dict, strict=False) # Skip "aspp" layer
model = nn.DataParallel(model, device_ids=gpu_ids)
if cuda:
model.cuda()
# Optimizer
optimizer = {
'sgd':
torch.optim.SGD(
# cf lr_mult and decay_mult in train.prototxt
params=[{
'params': get_lr_params(model.module, key='1x'),
'lr': CONFIG.LR,
'weight_decay': CONFIG.WEIGHT_DECAY
}, {
'params': get_lr_params(model.module, key='10x'),
'lr': 10 * CONFIG.LR,
'weight_decay': CONFIG.WEIGHT_DECAY
}, {
'params': get_lr_params(model.module, key='20x'),
'lr': 20 * CONFIG.LR,
'weight_decay': 0.0
}],
momentum=CONFIG.MOMENTUM,
),
}.get(CONFIG.OPTIMIZER)
# Loss definition
criterion = CrossEntropyLoss2d(ignore_index=CONFIG.IGNORE_LABEL)
if cuda:
criterion.cuda()
# TensorBoard Logger
writer = SummaryWriter(CONFIG.LOG_DIR)
loss_meter = MovingAverageValueMeter(20)
model.train()
model.module.scale.freeze_bn()
for iteration in tqdm(
range(1, CONFIG.ITER_MAX + 1),
total=CONFIG.ITER_MAX,
leave=False,
dynamic_ncols=True,
):
# Set a learning rate
poly_lr_scheduler(
optimizer=optimizer,
init_lr=CONFIG.LR,
iter=iteration - 1,
lr_decay_iter=CONFIG.LR_DECAY,
max_iter=CONFIG.ITER_MAX,
power=CONFIG.POLY_POWER,
)
# Clear gradients (ready to accumulate)
optimizer.zero_grad()
iter_loss = 0
for i in range(1, CONFIG.ITER_SIZE + 1):
data, target = next(loader_iter)
# Image
data = data.cuda() if cuda else data
data = Variable(data)
# Propagate forward
outputs = model(data)
# Loss
loss = 0
for output in outputs:
# Resize target for {100%, 75%, 50%, Max} outputs
target_ = resize_target(target, output.size(2))
target_ = target_.cuda() if cuda else target_
target_ = Variable(target_)
# Compute crossentropy loss
loss += criterion(output, target_)
# Backpropagate (just compute gradients wrt the loss)
loss /= float(CONFIG.ITER_SIZE)
loss.backward()
iter_loss += loss.data[0]
# Reload dataloader
if ((iteration - 1) * CONFIG.ITER_SIZE + i) % len(loader) == 0:
loader_iter = iter(loader)
loss_meter.add(iter_loss)
# Update weights with accumulated gradients
optimizer.step()
# TensorBoard
if iteration % CONFIG.ITER_TF == 0:
writer.add_scalar('train_loss', loss_meter.value()[0], iteration)
for i, o in enumerate(optimizer.param_groups):
writer.add_scalar('train_lr_group{}'.format(i), o['lr'],
iteration)
if iteration % 1000 != 0:
continue
for name, param in model.named_parameters():
name = name.replace('.', '/')
writer.add_histogram(name, param, iteration, bins="auto")
if param.requires_grad:
writer.add_histogram(name + '/grad',
param.grad,
iteration,
bins="auto")
# Save a model
if iteration % CONFIG.ITER_SNAP == 0:
torch.save(
model.module.state_dict(),
osp.join(CONFIG.SAVE_DIR,
'checkpoint_{}.pth'.format(iteration)),
)
# Save a model
if iteration % 100 == 0:
torch.save(
model.module.state_dict(),
osp.join(CONFIG.SAVE_DIR, 'checkpoint_current.pth'),
)
torch.save(
model.module.state_dict(),
osp.join(CONFIG.SAVE_DIR, 'checkpoint_final.pth'),
)
def main(config, embedding, model_path, run, imagedataset, local_rank, resnet,
bkg):
rank, world_size, device_id, device = setup(local_rank)
print("Local rank: {} Rank: {} World Size: {} Device_id: {} Device: {}".
format(local_rank, rank, world_size, device_id, device))
pth_extn = '.pth.tar'
# Configuration
CONFIG = Dict(yaml.load(open(config)))
datadir = os.path.join('data/datasets', imagedataset)
print("Split dir: ", datadir)
savedir = osp.dirname(model_path)
epoch = re.findall("checkpoint_(.*)\." + pth_extn[1:],
osp.basename(model_path))[-1]
if run == 'zlss' or run == 'flss':
val = np.load(datadir + '/split/test_list.npy')
visible_classes = np.load(datadir + '/split/novel_cls.npy')
if bkg:
visible_classes = np.asarray(np.concatenate(
[np.array([0]), visible_classes]),
dtype=int)
elif run == 'gzlss' or run == 'gflss':
val = np.load(datadir + '/split/test_list.npy')
vals_cls = np.asarray(np.concatenate([
np.load(datadir + '/split/seen_cls.npy'),
np.load(datadir + '/split/val_cls.npy')
]),
dtype=int)
if bkg:
vals_cls = np.asarray(np.concatenate([np.array([0]), vals_cls]),
dtype=int)
valu_cls = np.load(datadir + '/split/novel_cls.npy')
visible_classes = np.concatenate([vals_cls, valu_cls])
else:
print("invalid run ", run)
sys.exit()
cls_map = np.array([255] * 256)
for i, n in enumerate(visible_classes):
cls_map[n] = i
if run == 'gzlss' or run == 'gflss':
novel_cls_map = np.array([255] * 256)
for i, n in enumerate(list(valu_cls)):
novel_cls_map[cls_map[n]] = i
seen_cls_map = np.array([255] * 256)
for i, n in enumerate(list(vals_cls)):
seen_cls_map[cls_map[n]] = i
all_labels = np.genfromtxt(datadir + '/labels_2.txt',
delimiter='\t',
usecols=1,
dtype='str')
print("Visible Classes: ", visible_classes)
# Dataset
dataset = get_dataset(CONFIG.DATASET)(
train=None,
test=val,
root=CONFIG.ROOT,
split=CONFIG.SPLIT.TEST,
base_size=CONFIG.IMAGE.SIZE.TEST,
mean=(CONFIG.IMAGE.MEAN.B, CONFIG.IMAGE.MEAN.G, CONFIG.IMAGE.MEAN.R),
warp=CONFIG.WARP_IMAGE,
scale=None,
flip=False,
)
random.seed(42)
if embedding == 'word2vec':
class_emb = pickle.load(
open(datadir + '/word_vectors/word2vec.pkl', "rb"))
elif embedding == 'fasttext':
class_emb = pickle.load(
open(datadir + '/word_vectors/fasttext.pkl', "rb"))
elif embedding == 'fastnvec':
class_emb = np.concatenate([
pickle.load(open(datadir + '/word_vectors/fasttext.pkl', "rb")),
pickle.load(open(datadir + '/word_vectors/word2vec.pkl', "rb"))
],
axis=1)
else:
print("invalid emb ", embedding)
sys.exit()
class_emb = class_emb[visible_classes]
class_emb = F.normalize(torch.tensor(class_emb), p=2, dim=1).cuda()
print("Embedding dim: ", class_emb.shape[1])
print("# Visible Classes: ", class_emb.shape[0])
# DataLoader
loader = torch.utils.data.DataLoader(dataset=dataset,
batch_size=CONFIG.BATCH_SIZE.TEST,
num_workers=CONFIG.NUM_WORKERS,
shuffle=False,
sampler=DistributedSampler(
dataset,
num_replicas=world_size,
rank=rank,
shuffle=False),
pin_memory=True,
drop_last=True)
torch.set_grad_enabled(False)
# Model
model = DeepLabV2_ResNet101_MSC(class_emb.shape[1],
class_emb,
resnet=resnet)
state_dict = torch.load(model_path, map_location='cpu')
model = DistributedDataParallel(model.to(device), device_ids=[rank])
new_state_dict = OrderedDict()
if resnet == 'spnet':
for k, v in state_dict['state_dict'].items():
name = k.replace("scale", "base") # 'scale'->base
name = name.replace("stages.", "")
new_state_dict[name] = v
else:
new_state_dict = state_dict['state_dict']
model.load_state_dict(new_state_dict)
del state_dict
model.eval()
targets, outputs = [], []
loader_iter = iter(loader)
iterations = len(loader_iter)
print("Iterations: {}".format(iterations))
pbar = tqdm(loader,
total=iterations,
leave=False,
dynamic_ncols=True,
position=rank)
for iteration in pbar:
data, target, img_id = next(loader_iter)
# Image
data = data.to(device)
# Forward propagation
output = model(data)
output = F.interpolate(output,
size=data.shape[2:],
mode="bilinear",
align_corners=False)
output = F.softmax(output, dim=1)
target = cls_map[target.numpy()]
remote_target = torch.tensor(target).to(device)
if rank == 0:
remote_target = torch.zeros_like(remote_target).to(device)
output = torch.argmax(output, dim=1).cpu().numpy()
remote_output = torch.tensor(output).to(device)
if rank == 0:
remote_output = torch.zeros_like(remote_output).to(device)
for o, t in zip(output, target):
outputs.append(o)
targets.append(t)
torch.distributed.reduce(remote_output, dst=0)
torch.distributed.reduce(remote_target, dst=0)
torch.distributed.barrier()
if rank == 0:
remote_output = remote_output.cpu().numpy()
remote_target = remote_target.cpu().numpy()
for o, t in zip(remote_output, remote_target):
outputs.append(o)
targets.append(t)
if rank == 0:
if run == 'gzlss' or run == 'gflss':
score, class_iou = scores_gzsl(targets,
outputs,
n_class=len(visible_classes),
seen_cls=cls_map[vals_cls],
unseen_cls=cls_map[valu_cls])
else:
score, class_iou = scores(targets,
outputs,
n_class=len(visible_classes))
for k, v in score.items():
print(k, v)
score["Class IoU"] = {}
for i in range(len(visible_classes)):
score["Class IoU"][all_labels[visible_classes[i]]] = class_iou[i]
name = ""
name = model_path.replace(pth_extn, "_" + run + ".json")
if bkg == True:
with open(name.replace('.json', '_bkg.json'), "w") as f:
json.dump(score, f, indent=4, sort_keys=True)
else:
with open(name, "w") as f:
json.dump(score, f, indent=4, sort_keys=True)
print(score["Class IoU"])
return