def __init__(self,
layer='r31',
alpha=0.05,
prob=0.1,
pseudo1=True,
Noise=False,
std=1.,
mean=0.):
self.net = StyleAugmentation(layer, alpha, prob, pseudo1, Noise, std,
mean).cuda() # CUDA is available
for name, child in self.net.named_children():
for param in child.parameters():
param.requires_grad = False
class Stylization(nn.Module):
def __init__(self, layer='r41', alpha=[0.5], remove_stochastic=True, prob=1.0, pseudo1=True, Noise=False, std=1.,mean=0., idx=0):
super(Stylization, self).__init__()
self.net = StyleAugmentation(layer, alpha, remove_stochastic, prob, pseudo1, Noise, std, mean, idx).cuda() # CUDA is available
for name, child in self.net.named_children():
for param in child.parameters():
param.requires_grad = False
def forward(self, img, idx, alpha):
img=F.interpolate(img, 256, mode='bilinear', align_corners=True)
self.net.idx = idx
self.net.alpha = alpha
with torch.no_grad():
img = self.net(img).detach()
return img
class Stylization(object):
def __init__(self,
layer='r31',
alpha=0.05,
prob=0.1,
pseudo1=True,
Noise=False,
std=1.,
mean=0.):
self.net = StyleAugmentation(layer, alpha, prob, pseudo1, Noise, std,
mean).cuda() # CUDA is available
for name, child in self.net.named_children():
for param in child.parameters():
param.requires_grad = False
def __call__(self, x):
with torch.no_grad():
return self.net(x[None, :].cuda())[0].cpu()
def __init__(self, layer='r41', alpha=[0.5], remove_stochastic=True, prob=1.0, pseudo1=True, Noise=False, std=1.,mean=0., idx=0):
super(Stylization, self).__init__()
self.net = StyleAugmentation(layer, alpha, remove_stochastic, prob, pseudo1, Noise, std, mean, idx).cuda() # CUDA is available
for name, child in self.net.named_children():
for param in child.parameters():
param.requires_grad = False
split='train',
transform=transforms.Compose([
transforms.Resize(sz, interpolation=2),
transforms.ToTensor(),
]),
target_transform=None,
download=True)
content_loader = torch.utils.data.DataLoader(train__setting,
batch_size=batch_size,
shuffle=False,
num_workers=1)
Stylenet = StyleAugmentation(layer="r41",
alpha=[alpha],
prob=1.0,
pseudo1=True,
Noise=False,
std=1.,
mean=0.).cuda()
output_img = np.zeros(
[batch_size * sz, (num_examples + 1) * sz, 3], dtype=np.uint8)
for it, (content, _) in enumerate(content_loader):
Image = np.uint8(
content.permute(0, 2, 3, 1).cpu().detach().numpy() * 255)
for n in range(batch_size):
output_img[n * sz:(n + 1) * sz, 0:sz] = Image[n]
if print_idx:
output_img[n * sz:(n + 1) * sz, 0:sz] = cv2.putText(
output_img[n * sz:(n + 1) * sz,
0:sz], 'Image', (0, 20),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1)