| |
| |
| import numpy as np |
| from caffe2.python.crf import CRFWithLoss |
| |
| |
| def crf_update_predictions(model, crf_with_loss, classes): |
| return apply_crf( |
| model.param_init_net, |
| model.net, |
| crf_with_loss.transitions, |
| classes, |
| crf_with_loss.num_classes, |
| ) |
| |
| |
| def apply_crf(init_net, net, transitions, predictions, num_classes): |
| padded_classes = CRFWithLoss.pad_predictions( |
| predictions, init_net, net, num_classes |
| ) |
| bestPath = net.ViterbiPath([padded_classes, transitions]) |
| new_padded_classes = net.SwapBestPath([padded_classes, bestPath]) |
| # Revert the effect of pad_predictions by removing the last two rows and |
| # the last two columns |
| new_classes = net.RemovePadding( |
| [new_padded_classes], padding_width=1, end_padding_width=1 |
| ) |
| slice_starts = np.array([0, 0]).astype(np.int32) |
| slice_ends = np.array([-1, -3]).astype(np.int32) |
| slice_starts = net.GivenTensorIntFill([], shape=[2], values=slice_starts) |
| slice_ends = net.GivenTensorIntFill([], shape=[2], values=slice_ends) |
| new_classes = net.Slice([new_classes, slice_starts, slice_ends]) |
| return new_classes |