From b24b1eb96190ee969f9aa6f9c8320a9c949b312b Mon Sep 17 00:00:00 2001
From: Sean Sube <seansube@gmail.com>
Date: Sat, 14 Jan 2023 15:19:41 -0600
Subject: [PATCH] feat(gui): produce noise based on source image histogram

---
 api/onnx_web/image.py | 64 +++++++++++++++++++++++++++++++++++++++++++
 api/onnx_web/serve.py | 42 ++--------------------------
 2 files changed, 67 insertions(+), 39 deletions(-)
 create mode 100644 api/onnx_web/image.py

diff --git a/api/onnx_web/image.py b/api/onnx_web/image.py
new file mode 100644
index 00000000..c8ea228d
--- /dev/null
+++ b/api/onnx_web/image.py
@@ -0,0 +1,64 @@
+from numpy import random
+from PIL import Image, ImageStat
+from typing import Tuple
+
+import numpy as np
+
+def blend_mask_inverse_source(source: Tuple[int, int, int], mask: Tuple[int, int, int], noise: int) -> Tuple[int, int, int]:
+    m = float(noise) / 256
+    n = 1.0 - m
+
+    return (
+        int((source[0] * n) + (mask[0] * m)),
+        int((source[1] * n) + (mask[1] * m)),
+        int((source[2] * n) + (mask[2] * m)),
+    )
+
+
+def blend_source_histogram(source_image: Image, dims: Tuple[int, int], sigma = 200) -> Tuple[float, float, float]:
+    r, g, b = source_image.split()
+    width, height = dims
+
+    hist_r = r.histogram()
+    hist_g = g.histogram()
+    hist_b = b.histogram()
+
+    rng_r = random.choice(256, p=hist_r)
+    rng_g = random.choice(256, p=hist_g)
+    rng_b = random.choice(256, p=hist_b)
+
+    noise_r = rng_r.integers(0, size=width * height)
+    noise_g = rng_g.integers(0, size=width * height)
+    noise_b = rng_b.integers(0, size=width * height)
+
+    noise = Image.fromarray(zip(noise_r, noise_g, noise_b))
+
+    return noise
+
+
+
+# based on https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/master/scripts/outpainting_mk_2.py#L175-L232
+def expand_image(source_image: Image, mask_image: Image, dims: Tuple[int, int, int, int], fill = 'white', blend_source=blend_source_histogram, blend_op=blend_mask_inverse_source):
+    left, right, top, bottom = dims
+
+    full_width = left + source_image.width + right
+    full_height = top + source_image.height + bottom
+
+    full_source = Image.new('RGB', (full_width, full_height), fill)
+    full_source.paste(source_image, (left, top))
+
+    full_mask = Image.new('RGB', (full_width, full_height), fill)
+    full_mask.paste(mask_image, (left, top))
+
+    full_noise = blend_source(source_image, (full_width, full_height))
+
+    for x in range(full_source.width):
+        for y in range(full_source.height):
+            mask_color = full_mask.getpixel((x, y))
+            noise_color = full_noise.getpixel((x, y))
+            source_color = full_source.getpixel((x, y))
+
+            if mask_color[0] > 0:
+                full_source.putpixel((x, y), blend_op(source_color, mask_color, noise_color))
+
+    return (full_source, full_mask, (full_width, full_height))
diff --git a/api/onnx_web/serve.py b/api/onnx_web/serve.py
index 20d95a98..1d7afa3d 100644
--- a/api/onnx_web/serve.py
+++ b/api/onnx_web/serve.py
@@ -24,11 +24,13 @@ from flask_cors import CORS
 from flask_executor import Executor
 from hashlib import sha256
 from io import BytesIO
-from PIL import Image, ImageDraw
+from PIL import Image
 from struct import pack
 from os import environ, makedirs, path, scandir
 from typing import Any, Dict, Tuple, Union
 
+from .image import expand_image
+
 import json
 import numpy as np
 
@@ -102,44 +104,6 @@ def get_latents_from_seed(seed: int, width: int, height: int) -> np.ndarray:
     return image_latents
 
 
-def blend_pixel(source: Tuple[int, int, int], mask: Tuple[int, int, int], noise: int) -> Tuple[int, int, int]:
-    m = float(noise) / 256
-    n = 1.0 - m
-
-    return (
-        int((source[0] * n) + (mask[0] * m)),
-        int((source[1] * n) + (mask[1] * m)),
-        int((source[2] * n) + (mask[2] * m)),
-    )
-
-
-# based on https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/master/scripts/outpainting_mk_2.py#L175-L232
-def expand_image(source_image: Image, mask_image: Image, dims: Tuple[int, int, int, int]):
-    (left, right, top, bottom) = dims
-
-    full_width = left + source_image.width + right
-    full_height = top + source_image.height + bottom
-
-    full_source = Image.new('RGB', (full_width, full_height), 'white')
-    full_source.paste(source_image, (left, top))
-
-    full_mask = Image.new('RGB', (full_width, full_height), 'white')
-    full_mask.paste(mask_image, (left, top))
-
-    full_noise = Image.effect_noise((full_width, full_height), 200)
-
-    for x in range(full_source.width):
-        for y in range(full_source.height):
-            mask_color = full_mask.getpixel((x, y))
-            noise_color = full_noise.getpixel((x, y))
-            source_color = full_source.getpixel((x, y))
-
-            if mask_color[0] > 0:
-                full_source.putpixel((x, y), blend_pixel(source_color, mask_color, noise_color))
-
-    return (full_source, full_mask, (full_width, full_height))
-
-
 def load_pipeline(pipeline: DiffusionPipeline, model: str, provider: str, scheduler):
     global last_pipeline_instance
     global last_pipeline_scheduler