{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "metadata": {}
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-09-05 22:06:52,993 - INFO - Loading model from /home/andras/projects/bipolaroid/saved_models/run_53\n",
      "2024-09-05 22:06:52,994 - INFO - Hyperparameters: {'batch_size': 128, 'edit_count': 12, 'bin_count': 16, 'learning_rate': 0.002331030431920231, 'scheduler_gamma': 0.8396514875981009, 'num_epochs': 10, 'elu_alpha': 1.5997542544517163, 'leaky_relu_slope': 0.014905901603030717, 'dropout_prob': 0.0445357673555878, 'features': [64, 64, 64], 'use_residual': True, 'kernel_size': 5, 'model_type': 'HistogramNet', 'use_instance_norm': True, 'use_elu': False, 'leaky_relu_alpha': 0.04778301696189257}\n",
      "2024-09-05 22:06:54,564 - INFO - Parameter count: 5136129\n",
      "2024-09-05 22:06:54,605 - INFO - Loaded 14719 original images\n"
     ]
    }
   ],
   "source": [
    "from utils import set_up_logging\n",
    "\n",
    "set_up_logging()\n",
    "\n",
    "import torch\n",
    "from training import HistogramDataset\n",
    "from visualisation import display_images\n",
    "from models import load_model\n",
    "from config import MODELS_PATH, TEST_DATA\n",
    "from utils import get_device\n",
    "\n",
    "MODEL = \"run_53\"\n",
    "IMAGE_ID = 1\n",
    "\n",
    "device = get_device()\n",
    "model, hyperparameters = load_model(MODELS_PATH / MODEL, device)\n",
    "model.eval()\n",
    "dataset = HistogramDataset(\n",
    "    TEST_DATA,\n",
    "    edit_count=hyperparameters[\"edit_count\"],\n",
    "    bin_count=hyperparameters[\"bin_count\"],\n",
    "    target_size=(640, 640),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "img = dataset.get_original_image(IMAGE_ID)\n",
    "\n",
    "edits = {\"Original\": img}\n",
    "edits.update(\n",
    "    {f\"Edit {i}\": dataset.get_edited_image(IMAGE_ID, edit_idx=i) for i in range(1, 36)}\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 5634944 characters]",
      "text/plain": [
       "<Figure size 1650x2400 with 36 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display_images(edits, images_per_row=6, img_size_inches=4)\n",
    "None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy\n",
    "from utils import compute_histogram\n",
    "\n",
    "with torch.no_grad():\n",
    "    inputs = (\n",
    "        torch.tensor(\n",
    "            numpy.array(\n",
    "                [\n",
    "                    compute_histogram(\n",
    "                        image, bins=hyperparameters[\"bin_count\"], normalize=True\n",
    "                    )\n",
    "                    for image in edits.values()\n",
    "                ]\n",
    "            ),\n",
    "            dtype=torch.float,\n",
    "        ).unsqueeze(1)\n",
    "    ).to(device)\n",
    "\n",
    "    outputs = model(inputs).to(\"cpu\").detach().numpy().squeeze()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 5788964 characters]",
      "text/plain": [
       "<Figure size 1650x2400 with 36 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from histogram_transfer import apply_histogram\n",
    "from PIL import Image\n",
    "import concurrent.futures\n",
    "\n",
    "\n",
    "with concurrent.futures.ProcessPoolExecutor() as executor:\n",
    "    results = list(\n",
    "        executor.map(\n",
    "            apply_histogram,\n",
    "            edits.values(),\n",
    "            outputs,\n",
    "        )\n",
    "    )\n",
    "\n",
    "predicted = {\n",
    "    f\"Predicted {i}\": Image.fromarray(result) for i, result in enumerate(results)\n",
    "}\n",
    "display_images(predicted, images_per_row=6, img_size_inches=4)\n",
    "None"
   ]
  }
 ],
 "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
}
