def run(self): """Main entry point for style transfer, operates coarse-to-fine as specified by the number of scales. """ for self.scale in range(0, self.args.scales): # Pre-process the input images so they have the expected size. factor = 2 ** (self.args.scales - self.scale - 1) content_imgs = [] for img in self.content_imgs: content_imgs.append(resize.DownscaleBuilder(factor, cuda=self.cuda).build(img)) style_imgs = [] for img in self.style_imgs: style_imgs.append(resize.DownscaleBuilder(factor, cuda=self.cuda).build(img)) # Determine the stating point for the optimizer, was there an output of previous scale? if self.seed_img is None: # a) Load an image from disk, this needs to be the exact right size. if self.args.seed is not None: seed_img = images.load_from_file(self.args.seed, self.device) #seed_img = resize.DownscaleBuilder(factor).build(self.seed_img) #print(seed_img.shape, content_img.shape) assert seed_img.shape == content_imgs[0].shape # b) Use completely random buffer from a normal distribution. else: seed_img = torch.empty_like(content_imgs[0]).normal_(std=0.5).clamp_(-2.0, +2.0) else: # c) There was a previous scale, so resize and add noise from normal distribution. seed_img = (resize.DownscaleBuilder(factor, cuda=self.cuda).build(self.seed_img) + torch.empty_like(content_imgs[0]).normal_(std=0.1)).clamp_(-2.0, +2.0) # Pre-compute the cross-correlation statistics for the style image layers (aka. gram matrices). self.style_gram = {} n = 0 for img in style_imgs: for i, f in self.model.extract(img, layers=self.args.style_layers): self.style_gram[n, i] = histogram.square_matrix(f - 1.0).detach() n = n + 1 # Pre-compute feature histograms for the style image layers specified. self.style_hist = {} n = 0 for img in style_imgs: for k, v in self.model.extract(img, layers=self.args.histogram_layers): self.style_hist[n, k] = histogram.extract_histograms(v, bins=5, min=torch.tensor(-1.0), max=torch.tensor(+4.0)) n = n + 1 # Prepare and store the content image activations for image layers too. self.content_feat = {} n = 0 for img in content_imgs: for i, f in self.model.extract(img, layers=self.args.content_layers): self.content_feat[n, i] = f.detach() n = n + 1 # Now run the optimization using L-BFGS starting from the seed image. output = self.optimize(seed_img, self.iterations[self.scale]) #, lr=0.2) # For the next scale, we'll reuse a biliniear interpolated version of this output. self.seed_img = resize.UpscaleBuilder(factor, mode='bilinear').build(output).detach() # Save the final image at the finest scale to disk. basename = os.path.splitext(os.path.basename(self.args.content or self.args.style))[0] images.save_to_file(self.image.clone().detach().cpu(), self.args.output or ('output/%s_final.png' % basename))
def test_match_offset(source, offset): h = histogram.extract_histograms(source) output = histogram.match_histograms(source + offset, h) assert pytest.approx(0.0, abs=1e-4) == torch.max(output - source)
def test_match_identity(source): h = histogram.extract_histograms(source) output = histogram.match_histograms(source, h) assert pytest.approx(0.0, abs=1e-6) == torch.max(output - source)
def test_extract_normalized(source): h = histogram.extract_histograms(source) assert pytest.approx(1.0, abs=1e-6) == torch.sum(h[0])
def test_extract_balanced(source): h = histogram.extract_histograms(source) assert pytest.approx(0.0, abs=1e-6) == torch.mean(h[0] - 1.0 / h[0].shape[2])
def test_extract_deterministic(source): h1 = histogram.extract_histograms(source) h2 = histogram.extract_histograms(source) assert (h1[0] == h2[0]).all()