From 1b0d0d29ce8496a49f5aa3f8dc3ecf09c2c97251 Mon Sep 17 00:00:00 2001
From: TeaCrab <brentliu131@gmail.com>
Date: Fri, 18 Aug 2023 19:51:53 -0400
Subject: [PATCH] Changed how Equalize Node's parameters work

The way equalize_clahe() function's parameters being auto generated based on image size was resulting in errors.

Issue seems to be some kind of dimensional mismatch which was required to be the same by the kornia function.

Reverting back to a working method.

Plus some minor fixes.
---
 _func.py |  2 +-
 nodes.py | 16 +++++++++-------
 2 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/_func.py b/_func.py
index 4e9e6d2..55e524f 100644
--- a/_func.py
+++ b/_func.py
@@ -3,7 +3,7 @@ from math import ceil, floor, log
 
 debouncer = set()
 
-def po2(value, fill=False):
+def po2(value, fill=False) -> int:
     func = ceil if fill else floor
     return pow(2, func(log(value)/log(2)))
 
diff --git a/nodes.py b/nodes.py
index 4a1f917..b61991e 100644
--- a/nodes.py
+++ b/nodes.py
@@ -12,7 +12,7 @@ class EqualizeCLAHE:
     def INPUT_TYPES(s):
         return {
             "required": {
-                "images": ("IMAGE", ),
+                "image": ("IMAGE", ),
                 "clip_limit": ("FLOAT", {"default": 64, "min": 0.0, "max": 255, "step": 0.1}),
                 "grid_size": ("INT", {"default": 8, "min": 1, "max": 64, "step": 1}),
             },
@@ -28,12 +28,14 @@ class EqualizeCLAHE:
         _image = image.movedim(-1, 1)
         if size != (1024, 1024):
             grid_ratio = grid_size / clip_limit
-            size_ratio = min(size) / max(size)
-            clip_limit = po2(max(size), True)
-            grid_x = clip_limit * grid_ratio
-            grid_y = max(2, grid_x * size_ratio // 2 * 2)
-            if size[0] < size[1]: grid_x, grid_y = grid_y, grid_x
-        _image = equalize_clahe(_image, clip_limit, (grid_x, grid_y))
+            # size_ratio = min(size) / max(size)
+            clip_limit = int(max(8, po2(max(size)) * (clip_limit / 1024)))
+            grid_size = int(max(2, clip_limit * grid_ratio))
+            print(clip_limit, grid_size)
+            # grid_x = max(2, int(clip_limit * grid_ratio) // 2 * 2)
+            # grid_y = max(2, int(grid_x * size_ratio) // 2 * 2)
+            # if size[0] < size[1]: grid_x, grid_y = grid_y, grid_x
+        _image = equalize_clahe(_image, float(clip_limit), (grid_size, grid_size))
         result = _image.movedim(1, -1)
 
         return (result,)