Hello, I'm looking for a library that trains a CRF model in Python (if Pytorch, that would be even better).

I am working on a semantic segmentation task where we are trying to segment curvilinear structures.

My requirements for the CRF training are a bit specific:

- In my case, the image pixels are not the graph nodes. Instead, since the dataset is curvilinear structures, for every image I have a set of edges (small pieces of the curvilinear structure). I now want to train the CRF on these edge-pieces, that is, the graph nodes will be these edge-pieces. Thus the trained CRF essentially does a binary classification for each of these edge-pieces (that is, whether this edge-piece should be part of the segmentation output or not).

- I need a library where I can specify the unary and pairwise potentials of these edge-pieces in order to train the CRF. As a simple example, the unary potential is the average likelihood of the edge-piece, and the pairwise potential is the angle between two edge-pieces.

- It is not a linear-chain CRF because edge-pieces could be connected to multiple other pieces.

- Currently, I have frozen a deep neural network (DNN) which generates the edge-pieces. If the CRF library is in PyTorch, I could train the DNN and the CRF end-to-end, but if the CRF library is in Python, I would only be able to train the CRF. At this stage, even a Python library would be very helpful.

Some of the existing libraries don't work for my requirement:

- PyDenseCRF : It does not have learnable parameters.

- python-crfsuite : It does not allow me to specify the unary and pairwise potentials.

- pytorch-crf : It does linear-chain CRF while I need a graph one.

- crfasrnn_pytorch : It by default assumes the image pixels as the graph nodes. I cannot specify the unary and pairwise potentials.

If I could get any leads, that would be immensely helpful, thank you.