def main():
chainer.config.train = False
parser = argparse.ArgumentParser()
parser.add_argument('--gpu', type=int, default=-1)
parser.add_argument('--pretrained_model', default='camvid')
parser.add_argument('image')
args = parser.parse_args()
model = SegNetBasic(
n_class=len(camvid_label_names),
pretrained_model=args.pretrained_model)
if args.gpu >= 0:
chainer.cuda.get_device_from_id(args.gpu).use()
model.to_gpu()
img = utils.read_image(args.image, color=True)
labels = model.predict([img])
label = labels[0]
fig = plot.figure()
ax1 = fig.add_subplot(1, 2, 1)
vis_image(img, ax=ax1)
ax2 = fig.add_subplot(1, 2, 2)
vis_label(label, camvid_label_names, camvid_label_colors, ax=ax2)
plot.show()
def test_vis_label_mismatch_names_and_colors(self):
label = np.random.randint(-1, 2, size=(48, 64)).astype(np.int32)
if optional_modules:
with self.assertRaises(ValueError):
vis_label(label,
label_names=('class0', 'class1', 'class2'),
label_colors=((255, 0, 0), (0, 255, 0)))
def test_vis_label(self):
if optional_modules:
ax, legend_handles = vis_label(self.label,
label_names=self.label_names,
label_colors=self.label_colors)
self.assertIsInstance(ax, matplotlib.axes.Axes)
for handle in legend_handles:
self.assertIsInstance(handle, matplotlib.patches.Patch)
def test_vis_label_exceed_value(self):
label = np.random.randint(10, 20, size=(48, 64)).astype(np.int32)
if optional_modules:
with self.assertRaises(ValueError):
vis_label(label, label_names=('class0', 'class1', 'class2'))
from chainercv.visualizations import vis_image
from chainercv.visualizations import vis_label
import matplotlib.pyplot as plot
fig = plot.figure(figsize=(26, 10))
ax1 = fig.add_subplot(1, 2, 1)
plot.axis('off')
ax2 = fig.add_subplot(1, 2, 2)
plot.axis('off')
dataset = VOCDetectionDataset()
img, bbox, label = dataset[310]
vis_bbox(img, bbox, label,
label_names=voc_detection_label_names,
ax=ax1)
dataset = VOCSemanticSegmentationDataset()
img, label = dataset[30]
vis_image(img, ax=ax2)
_, legend_handles = vis_label(
label,
label_names=voc_semantic_segmentation_label_names,
label_colors=voc_semantic_segmentation_label_colors,
alpha=0.9, ax=ax2)
# ax2.legend(handles=legend_handles, bbox_to_anchor=(1, 1), loc=2)
plot.tight_layout()
plot.savefig('../images/vis_visualization.png')
plot.show()
chainer.config.train = False
if args.model == 'voc2012':
model = PSPNet(pretrained_model='voc2012')
labels = voc_semantic_segmentation_label_names
colors = voc_semantic_segmentation_label_colors
elif args.model == 'cityscapes':
model = PSPNet(pretrained_model='cityscapes')
labels = cityscapes_label_names
colors = cityscapes_label_colors
elif args.model == 'ade20k':
model = PSPNet(pretrained_model='ade20k')
labels = ade20k_label_names
colors = ade20k_label_colors
if args.gpu >= 0:
chainer.cuda.get_device_from_id(args.gpu).use()
model.to_gpu(args.gpu)
img = read_image(args.img_fn)
pred = model.predict([img])[0]
# Save the result image
ax = vis_image(img)
_, legend_handles = vis_label(pred, labels, colors, alpha=1.0, ax=ax)
ax.legend(handles=legend_handles, bbox_to_anchor=(1.05, 1), loc=2,
borderaxespad=0.)
base = os.path.splitext(os.path.basename(args.img_fn))[0]
plot.savefig('predict_{}.png'.format(base), bbox_inches='tight', dpi=400)