File size: 1,087 Bytes
482ab8a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
#!/usr/bin/env python
# coding: utf-8
#
# Author: Kazuto Nakashima
# URL: https://kazuto1011.github.io
# Date: 09 January 2019
import numpy as np
import pydensecrf.densecrf as dcrf
import pydensecrf.utils as utils
class DenseCRF(object):
def __init__(self, iter_max, pos_w, pos_xy_std, bi_w, bi_xy_std, bi_rgb_std):
self.iter_max = iter_max
self.pos_w = pos_w
self.pos_xy_std = pos_xy_std
self.bi_w = bi_w
self.bi_xy_std = bi_xy_std
self.bi_rgb_std = bi_rgb_std
def __call__(self, image, probmap):
C, H, W = probmap.shape
U = utils.unary_from_softmax(probmap)
U = np.ascontiguousarray(U)
image = np.ascontiguousarray(image)
d = dcrf.DenseCRF2D(W, H, C)
d.setUnaryEnergy(U)
d.addPairwiseGaussian(sxy=self.pos_xy_std, compat=self.pos_w)
d.addPairwiseBilateral(
sxy=self.bi_xy_std, srgb=self.bi_rgb_std, rgbim=image, compat=self.bi_w
)
Q = d.inference(self.iter_max)
Q = np.array(Q).reshape((C, H, W))
return Q
|