{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image, ImageEnhance\n",
    "import PIL.Image\n",
    "from editor.utils import random, get_colour_lut, apply_pixel_shader\n",
    "from editor.operations import add_noise, add_random_colour_spill\n",
    "from pathlib import Path\n",
    "\n",
    "\n",
    "PIL.Image.MAX_IMAGE_PIXELS = None\n",
    "\n",
    "\n",
    "DATA = Path('/mnt/wsl/PHYSICALDRIVE1/data/unsplash')\n",
    "EDIT_COUNT = 5\n",
    "TARGET_SIZE = 640, 640\n",
    "originals = sorted(DATA.glob('*.jpg'))\n",
    "originals_iter = iter(originals)\n",
    "\n",
    "\n",
    "def random_edit(img: Image) -> Image:\n",
    "    img = add_noise(img, random(0, 0.1))\n",
    "    img = ImageEnhance.Contrast(img).enhance(random(0.5, 2))\n",
    "    img = add_random_colour_spill(img, 1.5)\n",
    "    img = img.convert(\"HSV\")\n",
    "    saturation_lut = get_colour_lut(variance=0.2, count=5, type='linear')\n",
    "    brightness_lut = get_colour_lut(variance=0.2, count=5, type='cubic')\n",
    "    img = apply_pixel_shader(img, lambda h, s, v: (h, saturation_lut[s], brightness_lut[v]))\n",
    "    img = img.convert(\"RGB\")\n",
    "    img = ImageEnhance.Brightness(img).enhance(random(0.5, 2))\n",
    "    return img\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 978444 characters]",
      "text/plain": [
       "<Figure size 1200x800 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "example = next(originals_iter)\n",
    "img = Image.open(example)\n",
    "img.thumbnail(TARGET_SIZE, Image.Resampling.LANCZOS)\n",
    "\n",
    "edits = [\n",
    "    random_edit(img.copy())\n",
    "    for i in range(5)\n",
    "]\n",
    "\n",
    "images = [img] + edits\n",
    "titles = ['Original'] + [f'Edit {i}' for i in range(1, len(images))]\n",
    "\n",
    "fig, axes = plt.subplots(nrows=2, ncols=3, figsize=(12, 8))\n",
    "\n",
    "axes = axes.flatten()\n",
    "\n",
    "for i, (title, image) in enumerate(zip(titles, images)):\n",
    "    axes[i].imshow(image)\n",
    "    axes[i].axis('off')\n",
    "    axes[i].set_title(title)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import concurrent.futures\n",
    "from tqdm import tqdm\n",
    "\n",
    "\n",
    "def create_random_edits(file: Path):\n",
    "    try:\n",
    "        edits_folder = DATA / 'edited' / file.stem\n",
    "        edits_folder.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "        img = Image.open(file)\n",
    "        img.thumbnail(TARGET_SIZE, Image.Resampling.LANCZOS)\n",
    "        img.save(edits_folder / 'original.jpg')\n",
    "\n",
    "        for i in range(EDIT_COUNT):\n",
    "            clone = img.copy()\n",
    "            clone = random_edit(clone)\n",
    "            clone.save(edits_folder / f'{i}.jpg')\n",
    "    except Exception as e:\n",
    "        print(e, file)\n",
    "\n",
    "\n",
    "with tqdm(total=len(originals)) as progress_bar:\n",
    "    with concurrent.futures.ProcessPoolExecutor(32) as executor:\n",
    "        futures = [executor.submit(create_random_edits, file) for file in originals]\n",
    "        for future in concurrent.futures.as_completed(futures):\n",
    "            future.result()\n",
    "            progress_bar.update(1)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "bipolaroid",
   "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
