def forward(self, bottom, top):
probs = bottom[0].data
_, _, h, w = probs.shape
probs[probs < min_prob] = min_prob
unary = np.transpose(np.array(probs), [0, 2, 3, 1])
mean_pixel = np.array([104.0, 117.0, 123.0])
im = bottom[1].data[...]
im = zoom(im,
(1.0, 1.0, float(h) / im.shape[2], float(w) / im.shape[3]),
order=1)
im = np.transpose(im, [0, 2, 3, 1])
im = im + mean_pixel[None, None, None, :]
im = np.round(im)
N = unary.shape[0]
self.result = np.zeros(unary.shape)
for i in range(N):
self.result[i] = CRF(im[i], unary[i], scale_factor=12.0)
self.result = np.transpose(self.result, [0, 3, 1, 2])
self.result[self.result < min_prob] = min_prob
self.result = self.result / np.sum(self.result, axis=1, keepdims=True)
top[0].data[...] = np.log(self.result)
def run_infer(args):
args.batch_size = 1
args.pretrained = os.path.join(
args.snapshot, '{}-{:04d}.pth'.format(args.model, args.num_epoch - 1))
model = eval(args.model)(args)
loader = voc2012_loader(args)
pred_root = args.snapshot
# infer
model.eval()
for step, batch in enumerate(loader):
name, origin_image, image = batch
name = name[0]
origin_image = origin_image[0].numpy()
model.infer(image)
prob = model.output[0].detach().cpu().numpy().transpose(1, 2, 0)
d1, d2 = int(origin_image.shape[0]), int(origin_image.shape[1])
prob_exp = np.exp(prob - np.max(prob, axis=2, keepdims=True))
prob = prob_exp / np.sum(prob_exp, axis=2, keepdims=True)
prob = cv2.resize(prob, (d2, d1))
eps = 1e-5
prob[prob < eps] = eps
ans = np.argmax(prob, axis=2)
ans_crf = np.argmax(CRF(origin_image, np.log(prob), scale_factor=1.0),
axis=2)
imwrite(os.path.join(pred_root, 'pred', name + '.png'), ans)
imwrite(os.path.join(pred_root, 'pred_crf', name + '.png'), ans_crf)
info(None, 'Infer batch={}'.format(step))
def expand_seed_21(ini_seed, ini_prob, im, expand_bkg):
"""
expand seed when seed is of shape (41, 41, 21)
"""
loc = np.where(ini_seed > 0)
im_labels = np.unique(loc[2]).tolist()
im_non_labels = list(
set([k for k in range(21)]).difference(set(im_labels)))
res_unary = np.argmax(ini_prob, axis=2)
unary_zeros_ratio = len(np.where(res_unary == 0)[0]) / res_unary.size
ini_cue, _ = collapse_seed_to_cue(ini_seed)
if unary_zeros_ratio > 0.95:
# return np.expand_dims(ini_cue, axis=2)
return ini_seed
seed_unary = set_seed_unary_21(ini_seed)
seed_unary[seed_unary < eps] = eps
unary = seed_ratio * seed_unary + ini_prob
unary = unary / (np.sum(unary, axis=2, keepdims=True) + eps)
crf_prob = CRF(im, np.log(unary), scale_factor=12.0)
crf_prob[:, :, im_non_labels] = 0
# ############################ generate new seed #########################
# new_seed = np.zeros(shape=ini_seed.shape, dtype=np.float32)
# for cls in im_labels:
# thresh = bg_thresh if cls == 0 else fg_thresh
# cls_prob = crf_prob[:, :, cls]
# if cls != 0:
# new_seed[:, :, cls] = cls_prob > thresh * np.max(cls_prob)
# new_seed[:, :, 0] = ini_seed[:, :, 0]
# ############################ generate new seed #########################
res_crf = np.argmax(crf_prob, axis=2)
max_prob = np.max(crf_prob, axis=2)
res_crf[max_prob < prob_conf_thresh] = 21
new_seed = np.copy(ini_seed)
for v in sorted(im_labels):
if v == 21:
continue
if v == 0:
if not expand_bkg:
continue
ind = np.where(res_crf == v)
new_seed[ind[0], ind[1], v] = 1.0
# return np.expand_dims(new_seed, axis=2)
return new_seed
def expand_seed(ini_seed, ini_prob, im):
"""
expand seed when seed is of shape (41, 41), range [0, 21]
"""
assert ini_seed.shape[:2] == (41, 41)
assert ini_prob.shape == (41, 41, 21)
assert im.shape == (41, 41, 3)
im_labels = np.unique(ini_seed).tolist()
im_non_labels = list(
set([k for k in range(21)]).difference(set(im_labels)))
# set_trace()
################################ check if unary gives all bkg prediction ########################
res_unary = np.argmax(ini_prob, axis=2)
unary_zeros_ratio = len(np.where(res_unary == 0)[0]) / res_unary.size
# seed_non_zeros_ratio = len(np.where((ini_seed > 0) & (ini_seed < 21))[0]) / ini_seed.size
if unary_zeros_ratio > 0.95:
return np.expand_dims(ini_seed, axis=2)
################################ check if unary gives all bkg prediction ########################
################################ expand seed with CRF ########################
seed_unary = set_seed_unary(ini_seed, num_labels=21)
seed_unary[seed_unary < eps] = eps
unary = seed_ratio * seed_unary + ini_prob
unary = unary / (np.sum(unary, axis=2, keepdims=True) + eps)
crf_prob = CRF(im, np.log(unary), scale_factor=12.0)
crf_prob[:, :, im_non_labels] = 0
res_crf = np.argmax(crf_prob, axis=2)
max_prob = np.max(crf_prob, axis=2)
res_crf[max_prob < prob_conf_thresh] = 21
################################ expand seed with CRF ########################
################################ do not update bkg ########################
new_seed = np.copy(ini_seed)
unique_v = np.unique(ini_seed).tolist()
for v in unique_v:
if v == 21:
continue
# if v == 0: # do not update bkg
# continue
ind = np.where(res_crf == v)
new_seed[ind[0], ind[1]] = v
################################ do not update bkg ########################
return np.expand_dims(new_seed, axis=2)
def run_single(self, mask, img, flag_train): # flag_train is for the strange dif between train & test in org SEC
batch_size = mask.shape[0]
unary = np.transpose(mask, [0, 2, 3, 1])
if flag_train:
result = np.zeros(unary.shape)
for i in range(batch_size):
result[i] = CRF(np.round(resize(img[i]/255.0, self.mask_size, mode='constant')*255), unary[i], scale_factor=12.0)
result = np.transpose(result, [0, 3, 1, 2])
result[result < self.min_prob] = self.min_prob
result = result / np.sum(result, axis=1, keepdims=True)
return result
else:
result = np.zeros([batch_size, self.input_size[0], self.input_size[1]])
unary[unary < self.min_prob] = self.min_prob
for i in range(batch_size):
u_temp = resize(unary[i], self.input_size, mode='constant')
result[i, :, :] = np.argmax(CRF(img[i], np.log(u_temp), scale_factor=1.0), axis=2)
return result
def crf_layer(fc8_SEC, images, iternum):
unary = np.transpose(np.array(fc8_SEC.cpu().clone().data), [0, 2, 3, 1])
mean_pixel = np.array([104.0, 117.0, 123.0])
im = images.cpu().data
im = zoom(im, (1, 1, 41 / im.shape[2], 41 / im.shape[3]), order=1)
im = im + mean_pixel[None, :, None, None]
im = np.transpose(np.round(im), [0, 2, 3, 1])
N = unary.shape[0]
result = np.zeros(unary.shape)
for i in range(N):
result[i] = CRF(im[i], unary[i], maxiter=iternum, scale_factor=12.0)
result = np.transpose(result, [0, 3, 1, 2])
result[result < MIN_PROB] = MIN_PROB
result = result / np.sum(result, axis=1, keepdims=True)
return np.log(result)
def CRFLayer(fc8_SEC_Softmax, images, iternum):
# fc8-SEC-Softmax: (batch_size=15,21,41,41)
# images -> resize: (batch_size=15,21,41,41)
unary = np.transpose(np.array(fc8_SEC_Softmax.cpu().data), [0, 2, 3, 1])
mean_pixel = np.array([104.0, 117.0, 123.0])
im = np.array(images.cpu().data)
im = zoom(im, (1.0, 1.0, 41.0 / im.shape[2], 41.0 / im.shape[3]), order=1)
im = im + mean_pixel[None, :, None, None]
im = np.transpose(np.round(im), [0, 2, 3, 1])
N = unary.shape[0]
result = np.zeros(unary.shape)
for i in range(N):
result[i] = CRF(im[i], unary[i], maxiter=iternum, scale_factor=12.0)
result = np.transpose(result, [0, 3, 1, 2])
result[result < min_prob] = min_prob
result = result / np.sum(result, axis=1, keepdims=True)
CRF_result = np.log(result)
return CRF_result
def refinement(self, probs, im, scale_factor=12.0):
_, _, h, w = probs.shape
probs[probs < min_prob] = min_prob
unary = np.transpose(np.array(probs), [0, 2, 3, 1])
mean_pixel = np.array([104.0, 117.0, 123.0])
im = zoom(im,
(1.0, 1.0, float(h) / im.shape[2], float(w) / im.shape[3]),
order=1)
im = np.transpose(im, [0, 2, 3, 1])
im = im + mean_pixel[None, None, None, :]
im = np.round(im)
N = unary.shape[0]
result = np.zeros(unary.shape)
for i in range(N):
result[i] = CRF(im[i], unary[i], scale_factor=scale_factor)
result = np.transpose(result, [0, 3, 1, 2])
result[result < min_prob] = min_prob
result = result / np.sum(result, axis=1, keepdims=True)
return result