{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aeb428d0-0817-462c-b5d8-455a0615d305",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from PIL import Image\n",
    "import numpy as np\n",
    "import cv2\n",
    "\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",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6140d4b7-8238-431c-848e-6f6ae27652f5",
   "metadata": {},
   "outputs": [],
   "source": [
    " # Load the model image.\n",
    "model_image = Image.open(\"/home/ryan/src/InvokeAI/invokeai/backend/vto_workflow/dress.jpeg\")\n",
    "\n",
    "# Load the pattern image.\n",
    "pattern_image = Image.open(\"/home/ryan/src/InvokeAI/invokeai/backend/vto_workflow/pattern1.jpg\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb7186ba-dc0c-4520-ac30-49073a65601a",
   "metadata": {},
   "outputs": [],
   "source": [
    "mask = generate_dress_mask(model_image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b935de4-94c5-4be5-bf8e-a5a6e445c811",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize mask\n",
    "model_image_np = np.array(model_image)\n",
    "masked_model_image = (model_image_np * np.expand_dims(mask, -1).astype(np.float32)).astype(np.uint8)\n",
    "mask_image = Image.fromarray(masked_model_image)\n",
    "mask_image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e51bb545",
   "metadata": {},
   "outputs": [],
   "source": [
    "shadows = extract_channel(np.array(model_image), ImageChannel.LAB_L)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec43de4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize masked shadows\n",
    "masked_shadows = (shadows * mask).astype(np.uint8)\n",
    "masked_shadows_image = Image.fromarray(masked_shadows)\n",
    "masked_shadows_image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbb53794",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}