import CRF

params = CRF.FrankWolfeParams(scheme='fixed', # constant stepsize
            stepsize=1.0,
            regularizer='l2',
            lambda_=1.0, # regularization weight
            lambda_learnable=False,
            x0_weight=0.5, # useful for training, set to 0 if inference only
            x0_weight_learnable=False)

crf = CRF.DenseGaussianCRF(classes=21,
                alpha=160,
                beta=0.05,
                gamma=3.0,
                spatial_weight=1.0,
                bilateral_weight=1.0,
                compatibility=1.0,
                init='potts',
                solver='fw',
                iterations=5,
                params=params)

import torch

class CNNCRF(torch.nn.Module):
    """
    Simple CNN-CRF model
    """
    def __init__(self, cnn, crf):
        super().__init__()
        self.cnn = cnn
        self.crf = crf

    def forward(self, x):
        """
        x is a batch of input images
        """
        logits = self.cnn(x)
        logits = self.crf(x, logits)
        return logits

# Create a CNN-CRF model from given `cnn` and `crf`
# This is a PyTorch module that can be used in a usual way
model = CNNCRF(cnn, crf)