pix2pix_inference.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import matplotlib.pyplot as plt\n",
    "from pix2pix.model.pix2pix import PIX2PIX\n",
    "from pix2pix.data.data_generator import DataGenerator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "validation_path = '/home/ubuntu/pix2pix_data'\n",
    "path = \"/home/ubuntu/pix2pix_tf/gen_weights/\"\n",
    "\n",
    "val_gen = DataGenerator(\n",
    "    mode=\"valid\",\n",
    "    data_path=validation_path,\n",
    "    batch_size=1,\n",
    "    resize_dim=(1024, 1024),\n",
    ")\n",
    "\n",
    "val_dataset = tf.data.Dataset.from_generator(\n",
    "    val_gen,\n",
    "    output_types=(tf.float32, tf.float32),\n",
    "    output_shapes=(\n",
    "        [1024, 1024, 1],\n",
    "        [1024, 1024, 3],\n",
    "    ),\n",
    ")\n",
    "val_dataset = val_dataset.batch(1)\n",
    "val_dataset = val_dataset.repeat()\n",
    "\n",
    "model = PIX2PIX(weights = path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_show_image(dataset, model, test_image_height=1024,test_image_width=1024):\n",
    "    for images in dataset.take(1):\n",
    "        for i in range(1):\n",
    "            gray_image = images[0][i]\n",
    "            colored_image = images[1][i]\n",
    "            fig, ax = plt.subplots(1, 2, figsize=(40,40))\n",
    "            ax[0].imshow(tf.cast(gray_image*255.0,tf.uint8),cmap='gray')\n",
    "            ax[0].set_title(\"Gray Image\")\n",
    "            predicted_colored_image = model(tf.expand_dims(gray_image,0))\n",
    "            predicted_colored_image = tf.squeeze(predicted_colored_image)\n",
    "            ax[1].imshow(tf.cast((predicted_colored_image*255.0),tf.uint8))\n",
    "            ax[1].set_title(\"Predicted Colored Image\")\n",
    "        plt.show()    \n",
    "\n",
    "test_show_image(val_dataset, model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "vents_count",
   "language": "python",
   "name": "vents_count"
  },
  "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.9.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}