geo-tell / eanet Goto Github PK

    filt2 = torch.ones(self.n_classes, 1, ks, ks)
    filt2[:, :, self.radius:2*self.radius, self.radius:2*self.radius] = -8
    filt2.requires_grad = False
    filt2 = filt2.cuda()
    # get edge map of prediction
    prededge = F.conv2d(prediction.float(), filt2, bias=None,
                        stride=1, padding=self.radius, groups=self.n_classes)
    
    # squash C-dimensional (C is n_classes) edge map to single edge map
    norm = torch.sum(torch.pow(prededge,2), 1).unsqueeze(1)
    prededge = norm/(norm + self.alpha)
    
    
    # mask = lbedge.float()
    # num_positive = torch.sum((mask==1).float()).float()
    # num_negative = torch.sum((mask==0).float()).float()

    # mask[mask == 1] = 1.0 * num_negative / (num_positive + num_negative)
    # mask[mask == 0] = 1.5 * num_positive / (num_positive + num_negative)

    # cost = torch.nn.functional.binary_cross_entropy(
        # prededge.float(),lbedge.float(), weight=mask, reduce=False)
    # return torch.mean(cost)

    
    # visualize edge maps
    edge_maps.append(prededge.float())
    edge_maps.append(lbedge.float())
    edge_maps = torch.cat(edge_maps, dim=0)
    edge_maps_grid = torchvision.utils.make_grid(edge_maps, nrow=prededge.size(0))
    plt.imshow(edge_maps_grid.permute(1, 2, 0).detach().cpu(), cmap='gray')
    plt.title('Edge maps of 1-channel pred and gt')
    plt.show()
    
    return BinaryDiceLoss()(prededge.float(),lbedge.float())`