From bead83a8bc5b73c7a62f7515219f7bdb520fe021 Mon Sep 17 00:00:00 2001
From: Ryan Dick <ryanjdick3@gmail.com>
Date: Thu, 25 Jul 2024 15:31:46 -0400
Subject: [PATCH] Add image multiplication.

---
 clothing_workflow.ipynb                       | 51 ++++++++++++++-----
 .../backend/vto_workflow/overlay_pattern.py   | 24 +++++++++
 2 files changed, 61 insertions(+), 14 deletions(-)

diff --git a/clothing_workflow.ipynb b/clothing_workflow.ipynb
index 5b0e00b0d9..7907ad48a2 100644
--- a/clothing_workflow.ipynb
+++ b/clothing_workflow.ipynb
@@ -12,7 +12,7 @@
     "import numpy as np\n",
     "import cv2\n",
     "\n",
-    "from invokeai.backend.vto_workflow.overlay_pattern import generate_dress_mask\n",
+    "from invokeai.backend.vto_workflow.overlay_pattern import generate_dress_mask, multiply_images\n",
     "from invokeai.backend.vto_workflow.extract_channel import extract_channel, ImageChannel\n",
     "from invokeai.backend.vto_workflow.seamless_mapping import map_seamless_tiles\n",
     "\n",
@@ -87,19 +87,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "expanded_pattern = map_seamless_tiles(seamless_tile=pattern_image, target_hw=(model_image.height, model_image.width), num_repeats_h=5.0)\n",
-    "expanded_pattern"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "2a12c166",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "expanded_pattern = map_seamless_tiles(seamless_tile=pattern_image, target_hw=(model_image.height, model_image.width), num_repeats_h=1.0)\n",
-    "expanded_pattern"
+    "# Tile the pattern.\n",
+    "expanded_pattern = map_seamless_tiles(seamless_tile=pattern_image, target_hw=(model_image.height, model_image.width), num_repeats_h=10.0)\n"
    ]
   },
   {
@@ -108,6 +97,40 @@
    "id": "f4f22d02",
    "metadata": {},
    "outputs": [],
+   "source": [
+    "# Multiply the pattern by the shadows.\n",
+    "pattern_with_shadows = multiply_images(expanded_pattern, shadows)\n",
+    "pattern_with_shadows"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "97db42b0",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "de32f7e3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Merge the pattern with the model image.\n",
+    "pattern_with_shadows_np = np.array(pattern_with_shadows)\n",
+    "merged_image = np.where(mask[:, :, None],  pattern_with_shadows_np,model_image_np)\n",
+    "merged_image = Image.fromarray(merged_image)\n",
+    "merged_image"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ff1d4044",
+   "metadata": {},
+   "outputs": [],
    "source": []
   }
  ],
diff --git a/invokeai/backend/vto_workflow/overlay_pattern.py b/invokeai/backend/vto_workflow/overlay_pattern.py
index d40af66bee..74b4a65b6f 100644
--- a/invokeai/backend/vto_workflow/overlay_pattern.py
+++ b/invokeai/backend/vto_workflow/overlay_pattern.py
@@ -29,6 +29,30 @@ def generate_dress_mask(model_image):
     return binary_mask
 
 
+def multiply_images(image_1: Image.Image, image_2: Image.Image) -> Image.Image:
+    """Multiply two images together.
+
+    Args:
+        image_1 (Image.Image): The first image.
+        image_2 (Image.Image): The second image.
+
+    Returns:
+        Image.Image: The product of the two images.
+    """
+    image_1_np = np.array(image_1, dtype=np.float32)
+    if image_1_np.ndim == 2:
+        # If the image is greyscale, add a channel dimension.
+        image_1_np = np.expand_dims(image_1_np, axis=-1)
+    image_2_np = np.array(image_2, dtype=np.float32)
+    if image_2_np.ndim == 2:
+        # If the image is greyscale, add a channel dimension.
+        image_2_np = np.expand_dims(image_2_np, axis=-1)
+    product_np = image_1_np * image_2_np // 255
+    product_np = np.clip(product_np, 0, 255).astype(np.uint8)
+    product = Image.fromarray(product_np)
+    return product
+
+
 @torch.inference_mode()
 def main():
     # Load the model image.