{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Smoke test for the tile infill\"\"\"\n",
    "\n",
    "from pathlib import Path\n",
    "from typing import Optional\n",
    "from PIL import Image\n",
    "from invokeai.backend.image_util.infill_methods.tile import infill_tile\n",
    "\n",
    "images: list[tuple[str, Image.Image]] = []\n",
    "\n",
    "for i in sorted(Path(\"./test_images/\").glob(\"*.webp\")):\n",
    "    images.append((i.name, Image.open(i)))\n",
    "    images.append((i.name, Image.open(i).transpose(Image.FLIP_LEFT_RIGHT)))\n",
    "    images.append((i.name, Image.open(i).transpose(Image.FLIP_TOP_BOTTOM)))\n",
    "    images.append((i.name, Image.open(i).resize((512, 512))))\n",
    "    images.append((i.name, Image.open(i).resize((1234, 461))))\n",
    "\n",
    "outputs: list[tuple[str, Image.Image, Image.Image, Optional[Image.Image]]] = []\n",
    "\n",
    "for name, image in images:\n",
    "    try:\n",
    "        output = infill_tile(image, seed=0, tile_size=32)\n",
    "        outputs.append((name, image, output.infilled, output.tile_image))\n",
    "    except ValueError as e:\n",
    "        print(f\"Skipping image {name}: {e}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Display the images in jupyter notebook\n",
    "import matplotlib.pyplot as plt\n",
    "from PIL import ImageOps\n",
    "\n",
    "fig, axes = plt.subplots(len(outputs), 3, figsize=(10, 3 * len(outputs)))\n",
    "plt.subplots_adjust(hspace=0)\n",
    "\n",
    "for i, (name, original, infilled, tile_image) in enumerate(outputs):\n",
    "    # Add a border to each image, helps to see the edges\n",
    "    size = original.size\n",
    "    original = ImageOps.expand(original, border=5, fill=\"red\")\n",
    "    filled = ImageOps.expand(infilled, border=5, fill=\"red\")\n",
    "    if tile_image:\n",
    "        tile_image = ImageOps.expand(tile_image, border=5, fill=\"red\")\n",
    "\n",
    "    axes[i, 0].imshow(original)\n",
    "    axes[i, 0].axis(\"off\")\n",
    "    axes[i, 0].set_title(f\"Original ({name} - {size})\")\n",
    "\n",
    "    if tile_image:\n",
    "        axes[i, 1].imshow(tile_image)\n",
    "        axes[i, 1].axis(\"off\")\n",
    "        axes[i, 1].set_title(\"Tile Image\")\n",
    "    else:\n",
    "        axes[i, 1].axis(\"off\")\n",
    "        axes[i, 1].set_title(\"NO TILES GENERATED (NO TRANSPARENCY)\")\n",
    "\n",
    "    axes[i, 2].imshow(filled)\n",
    "    axes[i, 2].axis(\"off\")\n",
    "    axes[i, 2].set_title(\"Filled\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".invokeai",
   "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": 2
}