💅 Improved API to return a result object.

Author: alexjc

src/texturize/api.py

@@ -1,6 +1,7 @@
 # neural-texturize — Copyright (c) 2020, Novelty Factory KG.  See LICENSE for details.
 
 import os
+import collections
 import progressbar
 
 import torch
@@ -77,18 +78,20 @@ class NotebookLog:
     class ProgressBar:
         def __init__(self, max_iter):
             import ipywidgets
+
             self.bar = ipywidgets.IntProgress(
                 value=0,
                 min=0,
                 max=max_iter,
                 step=1,
-                description='',
-                bar_style='',
-                orientation='horizontal',
+                description="",
+                bar_style="",
+                orientation="horizontal",
                 layout=ipywidgets.Layout(width="100%"),
             )
 
             from IPython.display import display
+
             display(self.bar)
 
         def update(self, value, **keywords):
@@ -100,10 +103,17 @@ def finish(self):
     def create_progress_bar(self, iterations):
         return NotebookLog.ProgressBar(iterations)
 
-    def debug(self, *args): pass
-    def notice(self, *args): pass
-    def info(self, *args): pass
-    def warn(self, *args): pass
+    def debug(self, *args):
+        pass
+
+    def notice(self, *args):
+        pass
+
+    def info(self, *args):
+        pass
+
+    def warn(self, *args):
+        pass
 
 
 def get_default_log():
@@ -114,6 +124,11 @@ def get_default_log():
         return EmptyLog()
 
 
+Result = collections.namedtuple(
+    "Result", ["images", "loss", "octave", "scale", "iteration"]
+)
+
+
 @torch.no_grad()
 def process_octaves(
     source,
@@ -127,12 +142,12 @@ def process_octaves(
     device: str = None,
     precision: str = None,
 ):
-    # Setup the output and logging to use throughout the synthesis. 
+    # Setup the output and logging to use throughout the synthesis.
     log = log or get_default_log()
 
     # Determine which device and dtype to use by default, then set it up.
     device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
-    precision = getattr(torch, precision or "float32") 
+    precision = getattr(torch, precision or "float32")
 
     # Load the original image, always on the host device to save memory.
     texture_img = load_tensor_from_image(source, device="cpu").to(dtype=precision)
@@ -171,22 +186,18 @@ def process_octaves(
     ).to(dtype=precision)
 
     # Coarse-to-fine rendering, number of octaves specified by user.
-    for i, octave in enumerate(2 ** s for s in range(octaves - 1, -1, -1)):
-        if i == 5:
-            precision = torch.float16
-            encoder = encoder.half()
-
+    for octave, scale in enumerate(2 ** s for s in range(octaves - 1, -1, -1)):
         # Each octave we start a new optimization process.
         synth = TextureSynthesizer(
             device, encoder, lr=1.0, threshold=threshold, max_iter=iterations,
         )
-        log.info(f"\n OCTAVE #{i} ")
-        log.debug("<- scale:", f"1/{octave}")
+        log.info(f"\n OCTAVE #{octave} ")
+        log.debug("<- scale:", f"1/{scale}")
 
         # Create downscaled version of original texture to match this octave.
         texture_cur = F.interpolate(
             texture_img,
-            scale_factor=1.0 / octave,
+            scale_factor=1.0 / scale,
             mode="area",
             recompute_scale_factor=False,
         ).to(device=device, dtype=precision)
@@ -195,7 +206,7 @@ def process_octaves(
         del texture_cur
 
         # Compute the seed image for this octave, sprinkling a bit of gaussian noise.
-        result_size = size[1] // octave, size[0] // octave
+        result_size = size[1] // scale, size[0] // scale
         seed_img = F.interpolate(
             result_img, result_size, mode="bicubic", align_corners=False
         )
@@ -207,33 +218,41 @@ def process_octaves(
 
         # Now we can enable the automatic gradient computation to run the optimization.
         with torch.enable_grad():
-            for loss, result_img in synth.run(log, seed_img.to(dtype=precision), critics):
+            for iteration, (loss, result_img) in enumerate(
+                synth.run(log, seed_img.to(dtype=precision), critics)
+            ):
                 pass
         del synth
 
         output_img = F.interpolate(
             result_img, size=(size[1], size[0]), mode="nearest"
         ).cpu()
-        yield octave, loss, [
-            save_tensor_to_image(output_img[j : j + 1])
-            for j in range(output_img.shape[0])
-        ]
+        yield Result(
+            loss=loss,
+            octave=octave,
+            scale=scale,
+            iteration=iteration,
+            images=[
+                save_tensor_to_image(output_img[j : j + 1])
+                for j in range(output_img.shape[0])
+            ],
+        )
         del output_img
 
 
 def process_single_file(source, log: object, output: str = None, **config: dict):
-    for octave, _, result_img in process_octaves(
+    for result in process_octaves(
         load_image_from_file(source), log, **config
     ):
         filenames = []
-        for i, result in enumerate(result_img):
+        for i, image in enumerate(result.images):
             # Save the files for each octave to disk.
             filename = output.format(
-                octave=octave,
+                octave=result.octave,
                 source=os.path.splitext(os.path.basename(source))[0],
                 variation=i,
             )
-            result.save(filename)
+            image.save(filename)
             log.debug("\n=> output:", filename)
             filenames.append(filename)
 

src/texturize/io.py

@@ -1,5 +1,6 @@
 # neural-texturize — Copyright (c) 2020, Novelty Factory KG.  See LICENSE for details.
 
+import urllib
 import asyncio
 from io import BytesIO
 
@@ -20,6 +21,12 @@ def load_tensor_from_image(image, device):
     return V.to_tensor(image).unsqueeze(0).to(device)
 
 
+def load_image_from_url(url, mode="RGB"):
+    response = urllib.request.urlopen(url)
+    buffer = BytesIO(response.read())
+    return PIL.Image.open(buffer).convert(mode)
+
+
 def save_tensor_to_file(tensor, filename, mode="RGB"):
     img = save_tensor_to_image(tensor)
     img.save(filename)
@@ -38,10 +45,18 @@ def save_tensor_to_image(tensor, mode="RGB"):
     pass
 
 
-def show_result_in_notebook(images):
+def show_result_in_notebook(result, title="Generated Image"):
     clear_output()
-    for out in images:
-        html = ipywidgets.HTML(value="<h3>Octave #1</h3>")
+
+    for out in result.images:
+        html = ipywidgets.HTML(value=f"""
+        <h3>{title}</h3>
+        <ul style="font-size: 16px;">
+            <li>octave: {result.octave}</li>
+            <li>size: {out.size}</li>
+            <li>scale: 1/{result.scale}</li>
+            <li>loss: {result.loss:0.4f}</li>
+        </ul>""")
 
         buffer = io.BytesIO()
         out.save(buffer, format="webp", quality=80)

tests/app_gram.py

@@ -7,20 +7,18 @@
 
 
 def test_gram_single(image, size=(96, 88)):
-    for _, loss, images in process_octaves(
-        image(size), octaves=2, size=size, mode="gram"
-    ):
-        assert len(images) == 1
-        assert all(isinstance(img, PIL.Image.Image) for img in images)
-        assert all(img.size == size for img in images)
-        assert loss < 5e-2
+    for result in process_octaves(image(size), octaves=2, size=size, mode="gram"):
+        assert len(result.images) == 1
+        assert all(isinstance(img, PIL.Image.Image) for img in result.images)
+        assert all(img.size == size for img in result.images)
+        assert result.loss < 5e-2
 
 
 def test_gram_variations(image, size=(72, 64)):
-    for _, loss, images in process_octaves(
+    for result in process_octaves(
         image(size), variations=2, octaves=2, size=size, mode="gram"
     ):
-        assert len(images) == 2
-        assert all(isinstance(img, PIL.Image.Image) for img in images)
-        assert all(img.size == size for img in images)
-        assert loss < 5e-1
+        assert len(result.images) == 2
+        assert all(isinstance(img, PIL.Image.Image) for img in result.images)
+        assert all(img.size == size for img in result.images)
+        assert result.loss < 5e-1

tests/app_hist.py

@@ -7,10 +7,8 @@
 
 
 def test_hist_single(image, size=(32, 48)):
-    for _, loss, images in process_octaves(
-        image(size), octaves=2, size=size, mode="hist"
-    ):
-        assert len(images) == 1
-        assert all(isinstance(img, PIL.Image.Image) for img in images)
-        assert all(img.size == size for img in images)
-        assert loss < 1e-1
+    for result in process_octaves(image(size), octaves=2, size=size, mode="hist"):
+        assert len(result.images) == 1
+        assert all(isinstance(img, PIL.Image.Image) for img in result.images)
+        assert all(img.size == size for img in result.images)
+        assert result.loss < 1e-1

tests/app_patch.py

@@ -7,10 +7,10 @@
 
 
 def test_patch_single(image, size=(64, 48)):
-    for _, loss, images in process_octaves(
+    for result in process_octaves(
         image(size), octaves=2, size=size, mode="patch", threshold=1e-3
     ):
-        assert len(images) == 1
-        assert all(isinstance(img, PIL.Image.Image) for img in images)
-        assert all(img.size == size for img in images)
-        assert loss < 5.0
+        assert len(result.images) == 1
+        assert all(isinstance(img, PIL.Image.Image) for img in result.images)
+        assert all(img.size == size for img in result.images)
+        assert result.loss < 5.0