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)