import torch
import torch.nn.functional as F
from torch.optim import Adam
from torch import nn

import numpy as np
import matplotlib.pyplot as plt

from demo_helpers.demo_utils import pil, pil_grid, dummy_triangles
import pytorch.lovasz_losses as L


class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.bias = nn.Parameter(torch.Tensor([0]))

    def forward(self, x):
        return x + self.bias

def main():
    np.random.seed(18)
    labels_ = [dummy_triangles(200, [0, 255, 1]) for b in range(10)]
    labels = torch.stack([torch.from_numpy(a) for a in map(np.array, labels_)]).long()
    pil_grid(labels_, 5, margin=1).show()

    np.random.seed(57)
    feats = labels.float()
    feats[labels == 0] = feats[labels == 0].normal_(-1, 2.5)
    feats[labels == 1] = feats[labels == 1].normal_(1, 2.5)
    feats[labels == 255] = feats[labels == 255].normal_(0, 5)

    feats = feats.cuda()
    labels = labels.cuda()
    m = Model()
    m = m.cuda()

    preds = (m(feats) > 0).long().data
    print("Initial image-IoU:", L.iou_binary(preds, labels, ignore=255, per_image=False))

    m.bias.data.fill_(0)  # reset b

    optimizer = Adam(m.parameters(), lr=0.05)

    loss_plot = []
    iou_plot = []

    for iter in range(500):
        optimizer.zero_grad()
        out = m(feats)
        loss = L.lovasz_hinge(out, labels, ignore=255)
        loss.backward()
        optimizer.step()

        preds = (out.data > 0).long()
        loss_plot.append(loss.item())
        iou_plot.append(L.iou_binary(preds, labels, ignore=255, per_image=True))


if __name__ == '__main__':
    main()
import torch
import torch.nn.functional as F
from torch.optim import Adam
from torch import nn

import numpy as np
import matplotlib.pyplot as plt

from demo_helpers.demo_utils import pil, pil_grid, dummy_triangles
import pytorch.lovasz_losses as L


class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.bias = nn.Parameter(torch.Tensor([0]))

    def forward(self, x):
        return x + self.bias

def main():
    np.random.seed(18)
    labels_ = [dummy_triangles(200, [0, 255, 1]) for b in range(10)]
    labels = torch.stack([torch.from_numpy(a) for a in map(np.array, labels_)]).long()
    pil_grid(labels_, 5, margin=1).show()

    np.random.seed(57)
    feats = labels.float()
    feats[labels == 0] = feats[labels == 0].normal_(-1, 2.5)
    feats[labels == 1] = feats[labels == 1].normal_(1, 2.5)
    feats[labels == 255] = feats[labels == 255].normal_(0, 5)

    feats = feats.cuda()
    labels = labels.cuda()
    m = Model()
    m = m.cuda()

    preds = (m(feats) > 0).long().data
    print("Initial image-IoU:", L.iou_binary(preds, labels, ignore=255, per_image=False))

    m.bias.data.fill_(0)  # reset b

    optimizer = Adam(m.parameters(), lr=0.05)

    loss_plot = []
    iou_plot = []

    for iter in range(500):
        optimizer.zero_grad()
        out = m(feats)
        loss = L.lovasz_hinge(out, labels, ignore=255)
        loss.backward()
        optimizer.step()

        preds = (out.data > 0).long()
        loss_plot.append(loss.item())
        iou_plot.append(L.iou_binary(preds, labels, ignore=255, per_image=True))


if __name__ == '__main__':
    main()