import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from skimage.segmentation._slic import _enforce_label_connectivity_cython
def initWave(nPeriodic):
buf = []
for i in range(nPeriodic // 4+1):
v = 0.5 + i / float(nPeriodic//4+1e-10)
buf += [0, v, v, 0]
buf += [0, -v, v, 0] #so from other quadrants as well..
buf = buf[:2*nPeriodic]
awave = np.array(buf, dtype=np.float32) * np.pi
awave = torch.FloatTensor(awave).unsqueeze(-1).unsqueeze(-1).unsqueeze(0)
return awave
class SPADEGenerator(nn.Module):
def __init__(self, hidden_dim):
super().__init__()
nf = hidden_dim // 16
self.head_0 = SPADEResnetBlock(16 * nf, 16 * nf)
self.G_middle_0 = SPADEResnetBlock(16 * nf, 16 * nf)
self.G_middle_1 = SPADEResnetBlock(16 * nf, 16 * nf)
self.up_0 = SPADEResnetBlock(16 * nf, 8 * nf)
self.up_1 = SPADEResnetBlock(8 * nf, 4 * nf)
self.up_2 = SPADEResnetBlock(4 * nf, nf)
#self.up_3 = SPADEResnetBlock(2 * nf, 1 * nf)
final_nc = nf
self.conv_img = nn.Conv2d(final_nc, 3, 3, padding=1)
self.up = nn.Upsample(scale_factor=2)
def forward(self, x, input):
seg = input
x = self.head_0(x, seg)
x = self.up(x)
x = self.G_middle_0(x, seg)
x = self.G_middle_1(x, seg)
x = self.up(x)
x = self.up_0(x, seg)
x = self.up(x)
x = self.up_1(x, seg)
x = self.up(x)
x = self.up_2(x, seg)
#x = self.up(x)
#x = self.up_3(x, seg)
x = self.conv_img(F.leaky_relu(x, 2e-1))
return x
class SPADE(nn.Module):
def __init__(self, norm_nc, label_nc):
super().__init__()
ks = 3
self.param_free_norm = nn.InstanceNorm2d(norm_nc, affine=False)
# The dimension of the intermediate embedding space. Yes, hardcoded.
nhidden = 128
pw = ks // 2
self.mlp_shared = nn.Sequential(
nn.Conv2d(label_nc, nhidden, kernel_size=ks, padding=pw),
nn.ReLU()
)
self.mlp_gamma = nn.Conv2d(nhidden, norm_nc, kernel_size=ks, padding=pw)
self.mlp_beta = nn.Conv2d(nhidden, norm_nc, kernel_size=ks, padding=pw)
def forward(self, x, segmap):
# Part 1. generate parameter-free normalized activations
normalized = self.param_free_norm(x)
# Part 2. produce scaling and bias conditioned on semantic map
#segmap = F.interpolate(segmap, size=x.size()[2:], mode='nearest')
segmap = F.interpolate(segmap, size=x.size()[2:], mode='bilinear', align_corners = False)
actv = self.mlp_shared(segmap)
gamma = self.mlp_gamma(actv)
beta = self.mlp_beta(actv)
# apply scale and bias
out = normalized * (1 + gamma) + beta
return out
class SPADEResnetBlock(nn.Module):
def __init__(self, fin, fout):
super().__init__()
# Attributes
self.learned_shortcut = (fin != fout)
fmiddle = min(fin, fout)
# create conv layers
self.conv_0 = nn.Conv2d(fin, fmiddle, kernel_size=3, padding=1)
self.conv_1 = nn.Conv2d(fmiddle, fout, kernel_size=3, padding=1)
if self.learned_shortcut:
self.conv_s = nn.Conv2d(fin, fout, kernel_size=1, bias=False)
# define normalization layers
self.norm_0 = SPADE(fin, 256)
self.norm_1 = SPADE(fmiddle, 256)
if self.learned_shortcut:
self.norm_s = SPADE(fin, 256)
# note the resnet block with SPADE also takes in |seg|,
# the semantic segmentation map as input
def forward(self, x, seg):
x_s = self.shortcut(x, seg)
dx = self.conv_0(self.actvn(self.norm_0(x, seg)))
dx = self.conv_1(self.actvn(self.norm_1(dx, seg)))
out = x_s + dx
return out
def shortcut(self, x, seg):
if self.learned_shortcut:
x_s = self.conv_s(self.norm_s(x, seg))
else:
x_s = x
return x_s
def actvn(self, x):
return F.leaky_relu(x, 2e-1)
def get_edges(sp_label, sp_num):
# This function returns a (hw) * (hw) matrix N.
# If Nij = 1, then superpixel i and j are neighbors
# Otherwise Nij = 0.
top = sp_label[:, :, :-1, :] - sp_label[:, :, 1:, :]
left = sp_label[:, :, :, :-1] - sp_label[:, :, :, 1:]
top_left = sp_label[:, :, :-1, :-1] - sp_label[:, :, 1:, 1:]
top_right = sp_label[:, :, :-1, 1:] - sp_label[:, :, 1:, :-1]
n_affs = []
edge_indices = []
for i in range(sp_label.shape[0]):
# change to torch.ones below to include self-loop in graph
n_aff = torch.zeros(sp_num, sp_num).unsqueeze(0).to(sp_label.device)
# top/bottom
top_i = top[i].squeeze()
x, y = torch.nonzero(top_i, as_tuple = True)
sp1 = sp_label[i, :, x, y].squeeze().long()
sp2 = sp_label[i, :, x+1, y].squeeze().long()
n_aff[:, sp1, sp2] = 1
n_aff[:, sp2, sp1] = 1
# left/right
left_i = left[i].squeeze()
try:
x, y = torch.nonzero(left_i, as_tuple = True)
except:
import pdb; pdb.set_trace()
sp1 = sp_label[i, :, x, y].squeeze().long()
sp2 = sp_label[i, :, x, y+1].squeeze().long()
n_aff[:, sp1, sp2] = 1
n_aff[:, sp2, sp1] = 1
# top left
top_left_i = top_left[i].squeeze()
x, y = torch.nonzero(top_left_i, as_tuple = True)
sp1 = sp_label[i, :, x, y].squeeze().long()
sp2 = sp_label[i, :, x+1, y+1].squeeze().long()
n_aff[:, sp1, sp2] = 1
n_aff[:, sp2, sp1] = 1
# top right
top_right_i = top_right[i].squeeze()
x, y = torch.nonzero(top_right_i, as_tuple = True)
sp1 = sp_label[i, :, x, y+1].squeeze().long()
sp2 = sp_label[i, :, x+1, y].squeeze().long()
n_aff[:, sp1, sp2] = 1
n_aff[:, sp2, sp1] = 1
n_affs.append(n_aff)
edge_index = torch.stack(torch.nonzero(n_aff.squeeze(), as_tuple=True))
edge_indices.append(edge_index.to(sp_label.device))
return edge_indices, torch.cat(n_affs)
def enforce_connectivity(segs, H, W, sp_num = 196, min_size = None, max_size = None):
rets = []
for i in range(segs.shape[0]):
seg = segs[i]
seg = seg.squeeze().cpu().numpy()
segment_size = H * W / sp_num
if min_size is None:
min_size = int(0.1 * segment_size)
if max_size is None:
max_size = int(1000.0 * segment_size)
seg = _enforce_label_connectivity_cython(seg[None], min_size, max_size)[0]
seg = torch.from_numpy(seg).unsqueeze(0).unsqueeze(0)
rets.append(seg)
return torch.cat(rets)