Spaces:
Runtime error
Runtime error
File size: 1,987 Bytes
583456e |
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 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 |
# Copyright (c) Facebook, Inc. and its affiliates.
# Copyright (c) Meta Platforms, Inc. All Rights Reserved
import torch
from torch.nn import functional as F
import numpy as np
try:
import pydensecrf.densecrf as dcrf
from pydensecrf.utils import (
unary_from_softmax,
unary_from_labels,
create_pairwise_bilateral,
create_pairwise_gaussian,
)
except:
dcrf = None
def dense_crf_post_process(
logits,
image,
n_labels=None,
max_iters=5,
pos_xy_std=(3, 3),
pos_w=3,
bi_xy_std=(80, 80),
bi_rgb_std=(13, 13, 13),
bi_w=10,
):
"""
logits : [C,H,W]
image : [3,H,W]
"""
if dcrf is None:
raise FileNotFoundError(
"pydensecrf is required to perform dense crf inference."
)
if isinstance(logits, torch.Tensor):
logits = F.softmax(logits, dim=0).detach().cpu().numpy()
U = unary_from_softmax(logits)
n_labels = logits.shape[0]
elif logits.ndim == 3:
U = unary_from_softmax(logits)
n_labels = logits.shape[0]
else:
assert n_labels is not None
U = unary_from_labels(logits, n_labels, zero_unsure=False)
d = dcrf.DenseCRF2D(image.shape[1], image.shape[0], n_labels)
d.setUnaryEnergy(U)
# This adds the color-independent term, features are the locations only.
d.addPairwiseGaussian(
sxy=pos_xy_std,
compat=pos_w,
kernel=dcrf.DIAG_KERNEL,
normalization=dcrf.NORMALIZE_SYMMETRIC,
)
# This adds the color-dependent term, i.e. features are (x,y,r,g,b).
d.addPairwiseBilateral(
sxy=bi_xy_std,
srgb=bi_rgb_std,
rgbim=image,
compat=bi_w,
kernel=dcrf.DIAG_KERNEL,
normalization=dcrf.NORMALIZE_SYMMETRIC,
)
# Run five inference steps.
logits = d.inference(max_iters)
logits = np.asarray(logits).reshape((n_labels, image.shape[0], image.shape[1]))
return torch.from_numpy(logits)
|