def main(config, model_path, cuda, crf, camera_id):
# Configuration
CONFIG = Dict(yaml.load(open(config)))
cuda = cuda and torch.cuda.is_available()
if cuda:
current_device = torch.cuda.current_device()
print('Running on', torch.cuda.get_device_name(current_device))
# Label list
with open(CONFIG.LABELS) as f:
classes = {}
for label in f:
label = label.rstrip().split('\t')
classes[int(label[0])] = label[1].split(',')[0]
# Load a model
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, ) * 2
cap = cv2.VideoCapture(camera_id)
cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*'YUYV'))
while True:
# Image preprocessing
ret, frame = cap.read()
image = cv2.resize(frame.astype(float), image_size)
raw_image = 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(raw_image, output)
labelmap = np.argmax(output.transpose(1, 2, 0), axis=2)
labelmap = labelmap.astype(float) / CONFIG.N_CLASSES
labelmap = cm.jet_r(labelmap)[..., :-1] * 255.0
cv2.addWeighted(np.uint8(labelmap), 0.5, raw_image, 0.5, 0.0,
raw_image)
cv2.imshow('DeepLabV2', raw_image)
cv2.waitKey(50)
def makePrediction(config_obj, image_path, model, cuda=True, crf=False):
torch.set_grad_enabled(False)
image_id = extractIdFromPath(image_path)
# Image preprocessing
image = cv2.imread(image_path, cv2.IMREAD_COLOR).astype(float)
image = preProcessImage(image, config_obj, cuda)
# Inference
output = model(image)
output = F.interpolate(output,
size=image.shape[2:],
mode="bilinear",
align_corners=True)
output = F.softmax(output, dim=1)
output = output.data.cpu().numpy()[0]
if crf:
output = dense_crf(image_original, output)
labelmap = np.argmax(output, axis=0)
cocoResFormat = cocostuff.segmentationToCocoResult(labelmap, image_id)
return cocoResFormat
def process_image(model, config, cuda, image_size, image_path):
image = cv2.imread(image_path, cv2.IMREAD_COLOR).astype(float)
image = cv2.resize(image, image_size)
image_original = image.astype(np.uint8)
image = image[..., ::-1] - np.array(
[config.IMAGE.MEAN.R, config.IMAGE.MEAN.G, config.IMAGE.MEAN.B])
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[0].cpu().data.numpy()
output = dense_crf(image_original, output)
# At this point output is label - x - y
prob_map = output.transpose(1, 2, 0)
label_map = np.argmax(prob_map, axis=2)
return prob_map, label_map
def main(config, model_path, cuda, crf):
# Configuration
CONFIG = Dict(yaml.load(open(config)))
cuda = cuda and torch.cuda.is_available()
if cuda:
current_device = torch.cuda.current_device()
print('Running on', torch.cuda.get_device_name(current_device))
image_size = (
CONFIG.IMAGE.SIZE.TEST,
CONFIG.IMAGE.SIZE.TEST,
)
# Dataset
dataset = CocoStuff10k(
root=CONFIG.ROOT,
split='test',
image_size=image_size,
scale=False,
flip=False,
preload=False,
)
# DataLoader
loader = torch.utils.data.DataLoader(
dataset=dataset,
batch_size=CONFIG.BATCH_SIZE,
num_workers=CONFIG.NUM_WORKERS,
shuffle=False,
)
state_dict = torch.load(model_path,
map_location=lambda storage, loc: storage)
# Model
model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG.N_CLASSES)
model.load_state_dict(state_dict)
model = nn.DataParallel(model)
model.eval()
if cuda:
model.cuda()
targets, outputs = [], []
for data, target in tqdm(
loader,
total=len(loader),
leave=False,
dynamic_ncols=True,
):
# Image
data = data.cuda() if cuda else data
data = Variable(data, volatile=True)
# Forward propagation
output = model(data)
output = F.upsample(output, size=image_size, mode='bilinear')
output = F.softmax(output, dim=1)
output = output.data.cpu().numpy()
# Postprocessing
if crf:
crf_output = np.zeros(output.shape)
images = data.data.cpu().numpy().astype(np.uint8)
for i, (image, prob_map) in enumerate(zip(images, output)):
image = image.transpose(1, 2, 0)
crf_output[i] = dense_crf(image, prob_map)
output = crf_output
output = np.argmax(output, axis=1)
target = target.numpy()
for o, t in zip(output, target):
outputs.append(o)
targets.append(t)
score, class_iou = scores(targets, outputs, n_class=CONFIG.N_CLASSES)
for k, v in score.items():
print(k, v)
score['Class IoU'] = {}
for i in range(CONFIG.N_CLASSES):
score['Class IoU'][i] = class_iou[i]
with open(model_path.replace('.pth', '.json'), 'w') as f:
json.dump(score, f, indent=4, sort_keys=True)
def main(config, image_path, cuda, crf):
CONFIG = Dict(yaml.load(open(config)))
cuda = cuda and torch.cuda.is_available()
# Label list
with open(CONFIG.LABELS) as f:
classes = {}
for label in f:
label = label.rstrip().split("\t")
classes[int(label[0])] = label[1].split(",")[0]
# Load a model
state_dict = torch.load(CONFIG.PYTORCH_MODEL)
# Model
model = PSPNet(n_classes=CONFIG.N_CLASSES,
n_blocks=CONFIG.N_BLOCKS,
pyramids=CONFIG.PYRAMIDS)
model.load_state_dict(state_dict)
model.eval()
if cuda:
model.cuda()
image_size = (CONFIG.IMAGE.SIZE.TEST, ) * 2
# 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 = image[..., ::-1] - np.array(
[CONFIG.IMAGE.MEAN.R, CONFIG.IMAGE.MEAN.G, CONFIG.IMAGE.MEAN.B])
image = torch.from_numpy(image.transpose(2, 0, 1)).float().unsqueeze(0)
image = image.cuda() if cuda else image
# Inference
with torch.no_grad():
output = model(Variable(image, volatile=True))
output = F.upsample(output, size=image_size, mode="bilinear")
output = F.softmax(output, dim=1)
output = output[0].cpu().data.numpy()
if crf:
output = dense_crf(image_original, output)
labelmap = np.argmax(output.transpose(1, 2, 0), axis=2)
labels = np.unique(labelmap)
rows = np.floor(np.sqrt(len(labels) + 1))
cols = np.ceil((len(labels) + 1) / rows)
plt.figure(figsize=(10, 10))
ax = plt.subplot(1, 2, 1)
ax.set_title("Input image")
ax.imshow(image_original[:, :, ::-1])
ax.set_xticks([])
ax.set_yticks([])
for i, label in enumerate(labels):
if classes[label] == 'sky':
ax = plt.subplot(1, 2, 2)
ax.set_title(classes[label])
print("{0:3d}: {1}".format(label, classes[label]))
mask = labelmap != label
mask = mask.astype(np.uint8) * 255
mask_temp = np.zeros(mask.shape, mask.dtype)
mask = np.expand_dims(mask, axis=2)
mask_temp = np.expand_dims(mask_temp, axis=2)
mask = np.concatenate((mask, mask_temp, mask_temp), axis=-1)
mask = cv2.cvtColor(mask, cv2.COLOR_RGB2GRAY)
origin = image_original[:, :, ::-1]
right_image = cv2.add(origin,
np.zeros(np.shape(origin), dtype=np.uint8),
mask=mask)
ax.imshow(right_image)
cv2.imwrite('demo_sky_masked.jpg',
cv2.cvtColor(right_image, cv2.COLOR_BGR2RGB))
# ax.imshow(mask.astype(np.float32), alpha=0.5, cmap="viridis")
ax.set_xticks([])
ax.set_yticks([])
plt.tight_layout()
plt.show()
def main(config, image_path, cuda, crf):
CONFIG = Dict(yaml.load(open(config)))
cuda = cuda and torch.cuda.is_available()
# Label list
with open(CONFIG.LABELS) as f:
classes = {}
for label in f:
label = label.rstrip().split('\t')
classes[int(label[0])] = label[1].split(',')[0]
# Load a model
state_dict = torch.load(CONFIG.PYTORCH_MODEL)
# Model
model = PSPNet(n_classes=CONFIG.N_CLASSES,
n_blocks=CONFIG.N_BLOCKS,
pyramids=CONFIG.PYRAMIDS)
model.load_state_dict(state_dict)
model.eval()
if cuda:
model.cuda()
image_size = (CONFIG.IMAGE.SIZE.TEST, ) * 2
# 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 = image[..., ::-1] - np.array([
CONFIG.IMAGE.MEAN.R,
CONFIG.IMAGE.MEAN.G,
CONFIG.IMAGE.MEAN.B,
])
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[0].cpu().data.numpy()
if crf:
output = dense_crf(image_original, output)
labelmap = np.argmax(output.transpose(1, 2, 0), axis=2)
labels = np.unique(labelmap)
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(image_original[:, :, ::-1])
ax.imshow(mask.astype(np.float32), alpha=0.5, cmap='viridis')
ax.set_xticks([])
ax.set_yticks([])
plt.tight_layout()
plt.show()
def semseg(rootpath, config, image_path, cuda, crf, out_class_figure):
CONFIG = Dict(yaml.load(open(config)))
cuda = cuda and torch.cuda.is_available()
# Label list
with open(os.path.join(rootpath, CONFIG.LABELS)) as f:
classes = {}
for label in f:
label = label.rstrip().split("\t")
classes[int(label[0])] = label[1].split(",")[0]
# Load a model
state_dict = torch.load(os.path.join(rootpath, CONFIG.PYTORCH_MODEL))
# Model
model = PSPNet(n_classes=CONFIG.N_CLASSES,
n_blocks=CONFIG.N_BLOCKS,
pyramids=CONFIG.PYRAMIDS)
model.load_state_dict(state_dict)
model.eval()
if cuda:
model.cuda()
image_size = (CONFIG.IMAGE.SIZE.TEST, ) * 2
# image_size = (512, 288)
# Image preprocessing
# MES change for grayscale
# for color image input, ignoring alpha channel: use cv2.IMREAD_COLOR
# for color image input, keeping alpha channel: use cv2.IMREAD_UNCHANGED
# for grayscale image input, use cv2.IMREAD_GRAYSCALE
img = cv2.imread(image_path, cv2.IMREAD_COLOR)
# img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
image = cv2.resize(img, image_size).astype(float)
image_original = image.astype(np.uint8)
image = image[..., ::-1] - np.array(
[CONFIG.IMAGE.MEAN.R, CONFIG.IMAGE.MEAN.G, CONFIG.IMAGE.MEAN.B])
image = torch.from_numpy(image.transpose(2, 0, 1)).float().unsqueeze(0)
image = image.cuda() if cuda else image
# Inference
# MES change to silence 'volatile' deprecation message
with torch.no_grad():
# output = model(Variable(image, volatile=True))
output = model(Variable(image))
# MES change to silence 'upsample' deprecation message
# MES change to silence 'align_corners' default message
# output = F.upsample(output, size=image_size, mode="bilinear")
output = F.interpolate(output,
size=image_size,
mode="bilinear",
align_corners=False)
output = F.softmax(output, dim=1)
output = output[0].cpu().data.numpy()
if crf:
output = dense_crf(image_original, output)
#
#
labelmap = np.argmax(output.transpose(1, 2, 0), axis=2)
labels = np.unique(labelmap)
#
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(image_original[:, :, ::-1])
ax.imshow(mask.astype(np.float32), alpha=0.5, cmap="viridis")
ax.set_xticks([])
ax.set_yticks([])
# MES change to save the sky class image as a separate image
if classes[label] == 'sky':
mask_invert = labelmap != label # preserve non-sky pixels
masked_image = cv2.bitwise_and(image_original,
image_original,
mask=mask_invert.astype(np.uint8))
# plt.imsave(out_masked_sky_image, masked_image)
plt.tight_layout()
# MES changes to save output
# plt.show()
plt.savefig(out_class_figure)
#
# return the masked sky image
return masked_image
def main(config, model_path, cuda, crf, camera_id):
cuda = cuda and torch.cuda.is_available()
device = torch.device("cuda" if cuda else "cpu")
if cuda:
current_device = torch.cuda.current_device()
print("Running on", torch.cuda.get_device_name(current_device))
else:
print("Running on CPU")
# Configuration
CONFIG = Dict(yaml.load(open(config)))
# Label list
with open(CONFIG.LABELS) as f:
classes = {}
for label in f:
label = label.rstrip().split("\t")
classes[int(label[0])] = label[1].split(",")[0]
torch.set_grad_enabled(False)
# Model
model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG.N_CLASSES)
model.load_state_dict(torch.load(model_path))
model.eval()
model.to(device)
cap = cv2.VideoCapture(camera_id)
cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*"YUYV"))
while True:
# Image preprocessing
ret, frame = cap.read()
h, w, c = frame.shape
image = frame.astype(np.float32)
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.to(device)
# Inference
output = model(image)
output = F.upsample(output,
size=(h, w),
mode="bilinear",
align_corners=False)
output = F.softmax(output, dim=1)
output = output.data.cpu().numpy()[0]
if crf:
output = dense_crf(raw_image, output)
labelmap = np.argmax(output.transpose(1, 2, 0), axis=2)
labelmap = labelmap.astype(float) / CONFIG.N_CLASSES
labelmap = cm.jet_r(labelmap)[..., :-1] * 255.0
cv2.addWeighted(np.uint8(labelmap), 0.5, frame, 0.5, 0.0, frame)
cv2.imshow("DeepLabV2", frame)
cv2.waitKey(10)
def main(config, model_path, cuda, crf):
cuda = cuda and torch.cuda.is_available()
device = torch.device("cuda" if cuda else "cpu")
if cuda:
current_device = torch.cuda.current_device()
print("Running on", torch.cuda.get_device_name(current_device))
else:
print("Running on CPU")
# Configuration
CONFIG = Dict(yaml.load(open(config)))
# Dataset 10k or 164k
dataset = get_dataset(CONFIG.DATASET)(
root=CONFIG.ROOT,
split=CONFIG.SPLIT.VAL,
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,
)
# DataLoader
loader = torch.utils.data.DataLoader(
dataset=dataset,
batch_size=CONFIG.BATCH_SIZE.TEST,
num_workers=CONFIG.NUM_WORKERS,
shuffle=False,
)
torch.set_grad_enabled(False)
# Model
model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG.N_CLASSES)
state_dict = torch.load(model_path,
map_location=lambda storage, loc: storage)
model.load_state_dict(state_dict)
model = nn.DataParallel(model)
model.eval()
model.to(device)
targets, outputs = [], []
for data, target in tqdm(loader,
total=len(loader),
leave=False,
dynamic_ncols=True):
# Image
data = data.to(device)
# Forward propagation
output = model(data)
output = F.interpolate(output, size=data.shape[2:], mode="bilinear")
output = F.softmax(output, dim=1)
output = output.data.cpu().numpy()
# Postprocessing
if crf:
crf_output = np.zeros(output.shape)
images = data.data.cpu().numpy().astype(np.uint8)
for i, (image, prob_map) in enumerate(zip(images, output)):
image = image.transpose(1, 2, 0)
crf_output[i] = dense_crf(image, prob_map)
output = crf_output
output = np.argmax(output, axis=1)
target = target.numpy()
for o, t in zip(output, target):
outputs.append(o)
targets.append(t)
score, class_iou = scores(targets, outputs, n_class=CONFIG.N_CLASSES)
for k, v in score.items():
print(k, v)
score["Class IoU"] = {}
for i in range(CONFIG.N_CLASSES):
score["Class IoU"][i] = class_iou[i]
with open(model_path.replace(".pth", ".json"), "w") as f:
json.dump(score, f, indent=4, sort_keys=True)
def main(config, image_path, model_path, cuda, crf):
# Configuration
CONFIG = Dict(yaml.load(open(config)))
cuda = cuda and torch.cuda.is_available()
if cuda:
current_device = torch.cuda.current_device()
print('Running on', torch.cuda.get_device_name(current_device))
# Label list
with open(CONFIG.LABELS) as f:
classes = {}
for label in f:
label = label.rstrip().split('\t')
classes[int(label[0])] = label[1].split(',')[0]
# Load a model
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, ) * 2
# 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(image_original[:, :, ::-1])
ax.imshow(mask.astype(np.float32), alpha=0.5, cmap='viridis')
ax.set_xticks([])
ax.set_yticks([])
plt.tight_layout()
plt.savefig('demo.jpg')
plt.show()
def main(config, image_path, model_path, cuda, crf):
cuda = cuda and torch.cuda.is_available()
device = torch.device("cuda" if cuda else "cpu")
if cuda:
current_device = torch.cuda.current_device()
print("Running on", torch.cuda.get_device_name(current_device))
else:
print("Running on CPU")
# Configuration
CONFIG = Dict(yaml.load(open(config)))
# Label list
with open(CONFIG.LABELS) as f:
classes = {}
for label in f:
label = label.rstrip().split("\t")
classes[int(label[0])] = label[1].split(",")[0]
torch.set_grad_enabled(False)
# Model
model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG.N_CLASSES)
state_dict = torch.load(model_path,
map_location=lambda storage, loc: storage)
model.load_state_dict(state_dict)
model.eval()
model.to(device)
# Image preprocessing
image = cv2.imread(image_path, cv2.IMREAD_COLOR).astype(float)
scale = CONFIG.IMAGE.SIZE.TEST / max(image.shape[:2])
image = cv2.resize(image, dsize=None, fx=scale, fy=scale)
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.to(device)
# Inference
output = model(image)
output = F.interpolate(output,
size=image.shape[2:],
mode="bilinear",
align_corners=True)
output = F.softmax(output, dim=1)
output = output.data.cpu().numpy()[0]
if crf:
output = dense_crf(image_original, output)
labelmap = np.argmax(output, axis=0)
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(image_original[..., ::-1])
ax.imshow(mask.astype(np.float32), alpha=0.5, cmap="viridis")
ax.set_xticks([])
ax.set_yticks([])
plt.tight_layout()
plt.show()
def main(config, excludeval, embedding, model_path, run, cuda, crf, redo,
imagedataset, threshold):
pth_extn = '.pth.tar'
if osp.isfile(model_path.replace(
pth_extn, "_" + run + ".json")) and not threshold and not redo:
print("Already Done!")
with open(model_path.replace(pth_extn,
"_" + run + ".json")) as json_file:
data = json.load(json_file)
for key, value in data.items():
if not key == "Class IoU":
print(key, value)
sys.exit()
cuda = cuda and torch.cuda.is_available()
device = torch.device("cuda" if cuda else "cpu")
if cuda:
current_device = torch.cuda.current_device()
print("Running on", torch.cuda.get_device_name(current_device))
else:
print("Running on CPU")
# 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]
val = None
visible_classes = None
if run == 'zlss' or run == 'flss':
val = np.load(datadir + '/split/test_list.npy')
visible_classes = np.load(datadir + '/split/novel_cls.npy')
elif run == 'gzlss' or run == 'gflss':
val = np.load(datadir + '/split/test_list.npy')
if excludeval:
vals_cls = np.asarray(np.load(datadir + '/split/seen_cls.npy'),
dtype=int)
else:
vals_cls = np.asarray(np.concatenate([
np.load(datadir + '/split/seen_cls.npy'),
np.load(datadir + '/split/val_cls.npy')
]),
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()
if threshold is not None and run != 'gzlss':
print("invalid run for threshold", run)
sys.exit()
cls_map = np.array([255] * 256)
for i, n in enumerate(visible_classes):
cls_map[n] = i
if threshold is not None:
savedir = osp.join(savedir, str(threshold))
if crf is not None:
savedir = savedir + '-crf'
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
if threshold is not None:
thresholdv = np.asarray(np.zeros((visible_classes.shape[0], 1)),
dtype=np.float)
thresholdv[np.in1d(visible_classes, vals_cls), 0] = threshold
thresholdv = torch.tensor(thresholdv).float().cuda()
visible_classesp = np.concatenate([visible_classes, [255]])
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,
)
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,
)
torch.set_grad_enabled(False)
# Model
model = DeepLabV2_ResNet101_MSC(class_emb.shape[1], class_emb)
sdir = osp.join(savedir, model_path.replace(pth_extn, ""), str(epoch), run)
state_dict = torch.load(model_path,
map_location=lambda storage, loc: storage)
model = nn.DataParallel(model)
model.load_state_dict(state_dict['state_dict'])
model.eval()
model.to(device)
imgfeat = []
targets, outputs = [], []
for data, target, img_id in tqdm(loader,
total=len(loader),
leave=False,
dynamic_ncols=True):
# 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)
if threshold is not None:
output = output - thresholdv.view(1, -1, 1, 1)
target = cls_map[target.numpy()]
# Postprocessing
if crf:
output = output.data.cpu().numpy()
crf_output = np.zeros(output.shape)
images = data.data.cpu().numpy().astype(np.uint8)
for i, (image, prob_map) in enumerate(zip(images, output)):
image = image.transpose(1, 2, 0)
crf_output[i] = dense_crf(image, prob_map)
output = crf_output
output = np.argmax(output, axis=1)
else:
output = torch.argmax(output, dim=1).cpu().numpy()
for o, t in zip(output, target):
outputs.append(o)
targets.append(t)
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]
if threshold is not None:
with open(
model_path.replace(pth_extn, "_" + run + '_T' +
str(threshold) + ".json"), "w") as f:
json.dump(score, f, indent=4, sort_keys=True)
else:
with open(model_path.replace(pth_extn, "_" + run + ".json"), "w") as f:
json.dump(score, f, indent=4, sort_keys=True)
print(score["Class IoU"])
def main(config, model_path, cuda):
# Configuration
with open(config) as f:
CONFIG = yaml.load(f)
cuda = cuda and torch.cuda.is_available()
image_size = (CONFIG['IMAGE']['SIZE']['TEST'],
CONFIG['IMAGE']['SIZE']['TEST'])
n_classes = CONFIG['N_CLASSES']
# Dataset
dataset = get_dataset(CONFIG['DATASET'])(root=CONFIG['ROOT'],
split='test',
image_size=image_size,
scale=False,
flip=False,
preload=False)
# DataLoader
loader = torch.utils.data.DataLoader(dataset=dataset,
batch_size=CONFIG['BATCH_SIZE'],
num_workers=CONFIG['NUM_WORKERS'],
shuffle=False)
state_dict = torch.load(model_path,
map_location=lambda storage, loc: storage)
# Model
model = DeepLabV2_ResNet101_MSC(n_classes=n_classes)
model.load_state_dict(state_dict)
model.eval()
if cuda:
model.cuda()
targets, outputs = [], []
for data, target in tqdm(loader,
total=len(loader),
leave=False,
dynamic_ncols=True):
# Image
data = data.cuda() if cuda else data
data = Variable(data, volatile=True)
# Forward propagation
output = model(data)
output = F.upsample(output, size=image_size, mode='bilinear')
output = F.softmax(output, dim=1)
output = output.data.cpu().numpy()
crf_output = np.zeros(output.shape)
images = data.data.cpu().numpy().astype(np.uint8)
for i, (image, prob_map) in enumerate(zip(images, output)):
image = image.transpose(1, 2, 0)
crf_output[i] = dense_crf(image, prob_map)
output = crf_output
output = np.argmax(output, axis=1)
target = target.numpy()
for o, t in zip(output, target):
outputs.append(o)
targets.append(t)
score, class_iou = scores(targets, outputs, n_class=n_classes)
for k, v in score.items():
print k, v
score['Class IoU'] = {}
for i in range(n_classes):
score['Class IoU'][i] = class_iou[i]
with open('results.json', 'w') as f:
json.dump(score, f)
def dense_crf_wrapper(args):
return dense_crf(args[0], args[1])
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()
