From de6b1267e6a4231b48c0027fb129b8d68c7bac0e Mon Sep 17 00:00:00 2001
From: lllyasviel <lyuminzhang@outlook.com>
Date: Mon, 9 Oct 2023 03:59:45 -0700
Subject: [PATCH] improve cn (#597)

improve cn
---
 fooocus_extras/preprocessors.py | 82 +++++++++++++++++++++------------
 fooocus_version.py              |  2 +-
 2 files changed, 53 insertions(+), 31 deletions(-)

diff --git a/fooocus_extras/preprocessors.py b/fooocus_extras/preprocessors.py
index ddcc363c..dde6929b 100644
--- a/fooocus_extras/preprocessors.py
+++ b/fooocus_extras/preprocessors.py
@@ -2,37 +2,68 @@ import cv2
 import numpy as np
 
 
-def canny_k(x, k=0.5):
-    import cv2
+def centered_canny(x: np.ndarray):
+    assert isinstance(x, np.ndarray)
+    assert x.ndim == 2 and x.dtype == np.uint8
+
+    y = cv2.Canny(x, 100, 200)
+    y = y.astype(np.float32) / 255.0
+    return y
+
+
+def centered_canny_color(x: np.ndarray):
+    assert isinstance(x, np.ndarray)
+    assert x.ndim == 3 and x.shape[2] == 3
+
+    result = [centered_canny(x[..., i]) for i in range(3)]
+    result = np.stack(result, axis=2)
+    return result
+
+
+def pyramid_canny_color(x: np.ndarray):
+    assert isinstance(x, np.ndarray)
+    assert x.ndim == 3 and x.shape[2] == 3
+
     H, W, C = x.shape
-    Hs, Ws = int(H * k), int(W * k)
-    small = cv2.resize(x, (Ws, Hs), interpolation=cv2.INTER_AREA)
-    return cv2.Canny(small, 100, 200).astype(np.float32) / 255.0
+    acc_edge = None
+
+    for k in [0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]:
+        Hs, Ws = int(H * k), int(W * k)
+        small = cv2.resize(x, (Ws, Hs), interpolation=cv2.INTER_AREA)
+        edge = centered_canny_color(small)
+        if acc_edge is None:
+            acc_edge = edge
+        else:
+            acc_edge = cv2.resize(acc_edge, (edge.shape[1], edge.shape[0]), interpolation=cv2.INTER_LINEAR)
+            acc_edge = acc_edge * 0.75 + edge * 0.25
+
+    return acc_edge
+
+
+def norm255(x, low=4, high=96):
+    assert isinstance(x, np.ndarray)
+    assert x.ndim == 2 and x.dtype == np.float32
+
+    v_min = np.percentile(x, low)
+    v_max = np.percentile(x, high)
+
+    x -= v_min
+    x /= v_max - v_min
+
+    return x * 255.0
 
 
 def canny_pyramid(x):
     # For some reasons, SAI's Control-lora Canny seems to be trained on canny maps with non-standard resolutions.
     # Then we use pyramid to use all resolutions to avoid missing any structure in specific resolutions.
 
-    ks = [0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
-    cs = [canny_k(x, k) for k in ks]
-    cur = None
+    color_canny = pyramid_canny_color(x)
+    result = np.sum(color_canny, axis=2)
 
-    for c in cs:
-        if cur is None:
-            cur = c
-        else:
-            H, W = c.shape
-            cur = cv2.resize(cur, (W, H), interpolation=cv2.INTER_LINEAR)
-            cur = cur * 0.75 + c * 0.25
-
-    cur *= 400.0
-
-    return cur.clip(0, 255).astype(np.uint8)
+    return norm255(result, low=1, high=99).clip(0, 255).astype(np.uint8)
 
 
 def cpds(x):
-    import cv2
     # cv2.decolor is not "decolor", it is Cewu Lu's method
     # See http://www.cse.cuhk.edu.hk/leojia/projects/color2gray/index.html
     # See https://docs.opencv.org/3.0-beta/modules/photo/doc/decolor.html
@@ -45,15 +76,6 @@ def cpds(x):
     boost = boost.astype(np.float32)
 
     offset = np.sum((raw - boost) ** 2.0, axis=2) ** 0.5
-
     result = density + offset
 
-    v_min = np.percentile(result, 4)
-    v_max = np.percentile(result, 96)
-
-    result -= v_min
-    result /= v_max - v_min
-
-    result *= 255.0
-
-    return result.clip(0, 255).astype(np.uint8)
+    return norm255(result, low=4, high=96).clip(0, 255).astype(np.uint8)
diff --git a/fooocus_version.py b/fooocus_version.py
index 5bec51d5..44ea5c50 100644
--- a/fooocus_version.py
+++ b/fooocus_version.py
@@ -1 +1 @@
-version = '2.1.25'
+version = '2.1.26'