From 0363a06963a3a542c1d97526c0c7ca9d0134d6af Mon Sep 17 00:00:00 2001
From: blessedcoolant <54517381+blessedcoolant@users.noreply.github.com>
Date: Sat, 23 Sep 2023 02:02:27 +0530
Subject: [PATCH] feat: Add Color Map Preprocessor

---
 invokeai/app/invocations/baseinvocation.py    |  1 +
 .../controlnet_image_processors.py            | 27 +++++++++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/invokeai/app/invocations/baseinvocation.py b/invokeai/app/invocations/baseinvocation.py
index 97c6e5fc21..af7a343274 100644
--- a/invokeai/app/invocations/baseinvocation.py
+++ b/invokeai/app/invocations/baseinvocation.py
@@ -90,6 +90,7 @@ class FieldDescriptions:
     mask = "The mask to use for the operation"
     board = "The board to save the image to"
     image = "The image to process"
+    tile_size = "Tile size"
 
 
 class Input(str, Enum):
diff --git a/invokeai/app/invocations/controlnet_image_processors.py b/invokeai/app/invocations/controlnet_image_processors.py
index 5696f43f8d..8655a42558 100644
--- a/invokeai/app/invocations/controlnet_image_processors.py
+++ b/invokeai/app/invocations/controlnet_image_processors.py
@@ -559,3 +559,30 @@ class SamDetectorReproducibleColors(SamDetector):
             img[:, :] = ann_color
             final_img.paste(Image.fromarray(img, mode="RGB"), (0, 0), Image.fromarray(np.uint8(m * 255)))
         return np.array(final_img, dtype=np.uint8)
+
+
+@invocation(
+    "color_map_image_processor",
+    title="Color Map Processor",
+    tags=["controlnet"],
+    category="controlnet",
+    version="1.0.0",
+)
+class ColorMapImageProcessorInvocation(ImageProcessorInvocation):
+    """Generates a color map from the provided image"""
+
+    color_map_tile_size: int = InputField(default=64, ge=0, description=FieldDescriptions.tile_size)
+
+    def run_processor(self, image: Image.Image):
+        image = image.convert("RGB")
+        image = np.array(image, dtype=np.uint8)
+        height, width = image.shape[:2]
+
+        color_map = cv2.resize(
+            image,
+            (width // self.color_map_tile_size, height // self.color_map_tile_size),
+            interpolation=cv2.INTER_CUBIC,
+        )
+        color_map = cv2.resize(color_map, (width, height), interpolation=cv2.INTER_NEAREST)
+        color_map = Image.fromarray(color_map)        
+        return color_map