Data management more efficient in jupyter notebooks

Author: FS-CodeBase

pain_emnist_tf.ipynb

@@ -11,7 +11,12 @@
     "**Update:** August 21, 2024\n",
     "\n",
     "**Implementation:** PAIN architecture presented in:\n",
-    "**[IEEE publication](https://ieeexplore.ieee.org/abstract/document/8682767)**"
+    "**[IEEE publication](https://ieeexplore.ieee.org/abstract/document/8682767)**\n",
+    "\n",
+    "***NOTE:*** Code is setup to reconstruct 7x7 compressed noisy images. Training parameters may need to be adjusted for 4x4, 14x14, and 28x28 cases.\n",
+    "\n",
+    "Python Version: 3.11.5  \n",
+    "Jupyer Notebook Version: 6.5.4"
    ]
   },
   {
@@ -55,10 +60,16 @@
     "\n",
     "# Dimension of the compressed/noisy images (width=height)  \n",
     "# cmp_dim = 4:(for 4x4), 7:(7x7), 14:(14x14), or 28:(28x284)\n",
-    "cmp_dim = 14\n",
+    "cmp_dim = 7\n",
     "\n",
     "# Dimension of output, original are 28 x 28\n",
-    "out_dim  = 28 "
+    "out_dim  = 28 \n",
+    "\n",
+    "# Number of training epochs\n",
+    "num_epochs = 25\n",
+    "\n",
+    "# Size of training batch sizes\n",
+    "batch_size = 250"
    ]
   },
   {
@@ -122,35 +133,24 @@
     "# Create compressed noisy data\n",
     "def create_training_test_data(clean_train, clean_test, cmp_dim, out_dim):\n",
     "    \n",
-    "    # Step 0: Create Training and Validation Sets\n",
-    "    clean_test, clean_valid = split_dataset_rnd(clean_test)\n",
-    "    \n",
-    "    # Step 1: Compress images using median in sliding window\n",
+    "    # Compress images using median in sliding window\n",
     "    cmp_train = down_sample_list(clean_train, cmp_dim)\n",
     "    cmp_test  = down_sample_list(clean_test, cmp_dim)\n",
-    "    cmp_valid = down_sample_list(clean_valid, cmp_dim)\n",
     "\n",
-    "    # Step 2: Add Poisson noise to compressed images\n",
+    "    # Add Poisson noise to compressed images\n",
     "    noisy_train_ = np.random.poisson(lam=cmp_train)\n",
     "    noisy_test_  = np.random.poisson(lam=cmp_test)\n",
-    "    noisy_valid_ = np.random.poisson(lam=cmp_valid)\n",
     "    \n",
+    "    # Scale to [0,255]\n",
     "    noisy_train  = np.clip(noisy_train_,0,255)\n",
     "    noisy_test   = np.clip(noisy_test_,0,255)\n",
-    "    noisy_valid  = np.clip(noisy_valid_,0,255)\n",
     "    \n",
-    "    # Step 3: Reshape Input Arrays and scale from [0, 255] to [0,1]\n",
+    "    # Reshape Compressed/Noisy Arrays and scale to [0,1]\n",
     "    noisy_train = np.array([matrix.reshape(cmp_dim**2,) for matrix in noisy_train/255])\n",
     "    noisy_test  = np.array([matrix.reshape(cmp_dim**2,) for matrix in noisy_test/255])\n",
-    "    noisy_valid = np.array([matrix.reshape(cmp_dim**2,) for matrix in noisy_valid/255])\n",
-    "    \n",
-    "    # Step 4: Reshape Expected Ouput Arraus and scale from [0, 255] to [0,1]\n",
-    "    clean_train = np.array([matrix.reshape(out_dim**2,) for matrix in clean_train/255])\n",
-    "    clean_test  = np.array([matrix.reshape(out_dim**2,) for matrix in clean_test/255])\n",
-    "    clean_valid = np.array([matrix.reshape(out_dim**2,) for matrix in clean_valid/255])\n",
     "    \n",
-    "    # Step 5: Return training data\n",
-    "    return [clean_train, clean_test, clean_valid, noisy_train, noisy_test, noisy_valid]\n",
+    "    # Return training data\n",
+    "    return noisy_train, noisy_test\n",
     "\n",
     "# PAIN architecture builder function\n",
     "def build_PAIN(in_dim, out_dim, enc_dim = 256):\n",
@@ -188,57 +188,43 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "# Check compression dimention variable:\n",
+    "if not (cmp_dim in [4, 7, 14, 28]):\n",
+    "    raise ValueError(f'cmp_dim = {cmp_dim}, not handled. The cmp_dim value must be 4, 7, 14, or 28.')\n",
+    "\n",
     "directory = 'training_data'\n",
     "\n",
     "# If the directory does not exist, create it\n",
     "if not os.path.exists(directory):\n",
     "    os.makedirs(directory)\n",
     "\n",
     "# Load the MNIST Dataset\n",
-    "file = f'mnist_training_data_cmp' # Do not modify\n",
-    "\n",
-    "# If the training data does not exist, create it\n",
-    "if not os.path.exists(f'{directory}/{file}_{cmp_dim}x{cmp_dim}.h5'):\n",
-    "    # Load MNIST\n",
-    "    (clean_train, _), (clean_test, _) = tf.keras.datasets.mnist.load_data()\n",
+    "dat_file = f'{directory}/mnist_{cmp_dim}x{cmp_dim}_train.h5'\n",
     "    \n",
-    "    # Create compressed/noisy data, test set, and validation set\n",
-    "    train_data = create_training_test_data(clean_train, clean_test, cmp_dim, out_dim)\n",
-    "    \n",
-    "    # MNIST data\n",
-    "    clean_train = train_data[0]\n",
-    "    clean_test  = train_data[1]\n",
-    "    clean_valid = train_data[2]\n",
-    "    \n",
-    "    # Compressed noisy signals\n",
-    "    noisy_train = train_data[3]\n",
-    "    noisy_test  = train_data[4]\n",
-    "    noisy_valid = train_data[5]\n",
-    "    \n",
-    "    # Save with compression\n",
-    "    with h5py.File(f'{directory}/{file}_{cmp_dim}x{cmp_dim}.h5', 'w') as f:\n",
-    "        # Save original images\n",
-    "        f.create_dataset('clean_train', data=clean_train, compression='gzip')\n",
-    "        f.create_dataset('clean_test', data=clean_test, compression='gzip')\n",
-    "        f.create_dataset('clean_valid', data=clean_valid, compression='gzip')\n",
+    "# Load MNIST\n",
+    "(clean_train, _), (clean_test, _) = tf.keras.datasets.mnist.load_data()\n",
+    "\n",
+    "# If the compressed/noisy data does not exist, create it\n",
+    "if not os.path.exists(dat_file):\n",
+    "    # Create compressed/noisy training/test data\n",
+    "    noisy_train, noisy_test  = create_training_test_data(clean_train, clean_test, cmp_dim, out_dim)\n",
     "        \n",
+    "    # Save compressed/noisy training/test data\n",
+    "    with h5py.File(dat_file, 'w') as f:\n",
     "        # Save compressed noisy images\n",
     "        f.create_dataset('noisy_train', data=noisy_train, compression='gzip')\n",
     "        f.create_dataset('noisy_test', data=noisy_test, compression='gzip')\n",
-    "        f.create_dataset('noisy_valid', data=noisy_valid, compression='gzip')\n",
     "        \n",
-    "else: # If data exists, load it\n",
-    "    with h5py.File(f'{directory}/{file}_{cmp_dim}x{cmp_dim}.h5', 'r') as dat_file:\n",
-    "\n",
-    "        # Load original MNIST images\n",
-    "        clean_train = dat_file['clean_train'][:]\n",
-    "        clean_test  = dat_file['clean_test'][:]\n",
-    "        clean_valid = dat_file['clean_valid'][:]\n",
-    "\n",
-    "        # Load compressed noisy images\n",
+    "else: \n",
+    "    # Load compressed/noisy training data\n",
+    "    with h5py.File(dat_file, 'r') as dat_file:\n",
+    "        # Load compressed noisy training/test images\n",
     "        noisy_train = dat_file['noisy_train'][:]\n",
     "        noisy_test  = dat_file['noisy_test'][:]\n",
-    "        noisy_valid = dat_file['noisy_valid'][:]"
+    "        \n",
+    "# Prepare original mnist data for model training\n",
+    "clean_train = np.array([matrix.reshape(out_dim**2,) for matrix in clean_train/255])\n",
+    "clean_test  = np.array([matrix.reshape(out_dim**2,) for matrix in clean_test/255])"
    ]
   },
   {
@@ -254,9 +240,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Display: training set\n",
-    "#######################\n",
-    "\n",
     "# Create a 2 by 4 subplot handle\n",
     "fig, axes = plt.subplots(2, 4, figsize=(9, 4.5))\n",
     "axes      = axes.flatten()\n",
@@ -311,14 +294,14 @@
     "# Create & Compile the PAIN model\n",
     "PAIN = build_PAIN(in_dim=cmp_dim, out_dim=out_dim)\n",
     "\n",
-    "# Create an RMSProp optimizer with a specific learning rate\n",
-    "RMSp = tf.keras.optimizers.RMSprop(learning_rate=0.05)\n",
+    "# Create an optimizer with a specific learning rate\n",
+    "tf_opt = tf.keras.optimizers.Adam(learning_rate=0.01)\n",
     "\n",
     "# Compile the model\n",
-    "PAIN.compile(optimizer=RMSp, loss='mean_squared_error')\n",
+    "PAIN.compile(optimizer=tf_opt, loss='mean_squared_error')\n",
     "\n",
     "# Train model and saving fitting history\n",
-    "fit_history = PAIN.fit(noisy_train, clean_train, epochs=120, batch_size=250, validation_data=(noisy_test, clean_test))"
+    "fit_history = PAIN.fit(noisy_train, clean_train, epochs=num_epochs, batch_size=batch_size, validation_split=0.2)"
    ]
   },
   {
@@ -403,7 +386,7 @@
     "axes[8].set_ylabel(f'Original\\n{out_dim} x {out_dim}\\nMNIST',fontdict={'fontsize': 12, 'fontfamily': 'serif'})\n",
     "\n",
     "# Adjust layout to decrease padding between subplots\n",
-    "plt.subplots_adjust(wspace=0.1, hspace=0.00)\n",
+    "plt.subplots_adjust(wspace=0.1, hspace=0)\n",
     "\n",
     "# Save results \n",
     "plt.savefig(f'{imdir}/{imfile}_{cmp_dim}x{cmp_dim}_train_PAIN.png')\n",
@@ -416,7 +399,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Apply PAIN to Validation Set and Display Output"
+    "## Apply PAIN to Test Set and Display Output"
    ]
   },
   {
@@ -425,8 +408,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Apply PAIN to all validation data inputs\n",
-    "pred_valid_out = PAIN.predict(noisy_valid[0:101])"
+    "# Apply PAIN to all test data inputs\n",
+    "pred_test_out = PAIN.predict(noisy_test[0:101])"
    ]
   },
   {
@@ -440,26 +423,26 @@
     "axes      = axes.flatten()\n",
     "\n",
     "# Add title\n",
-    "fig.suptitle('Application of PAIN Architecture\\n(MNIST Validation Set)',fontsize=20,fontweight='bold', fontfamily='serif')\n",
+    "fig.suptitle('Application of PAIN Architecture\\n(MNIST Test Set)',fontsize=20,fontweight='bold', fontfamily='serif')\n",
     "\n",
-    "# Shift window through validation dataset \n",
+    "# Shift window through test dataset \n",
     "sft_idx = 0\n",
     "\n",
     "# Plot noisy and decompressed realizations in each subplot\n",
     "for idx in range(4):\n",
     "    # Plot noisy\n",
-    "    axes[idx].imshow(noisy_valid[idx+sft_idx].reshape(cmp_dim,cmp_dim),cmap='gray')\n",
+    "    axes[idx].imshow(noisy_test[idx+sft_idx].reshape(cmp_dim,cmp_dim),cmap='gray')\n",
     "    axes[idx].set_xticks([]) # Remove xticklabels\n",
     "    axes[idx].set_yticks([]) # Remove yticklabels\n",
     "    axes[idx].set_xlabel('⇩',fontdict={'fontsize': 25, 'fontweight': 'bold', 'fontfamily': 'serif', 'color':'blue'})\n",
     "    \n",
     "    # Plot decompressed with PAIN\n",
-    "    axes[idx+4].imshow(pred_valid_out[idx+sft_idx].reshape(28,28),cmap='gray')\n",
+    "    axes[idx+4].imshow(pred_test_out[idx+sft_idx].reshape(28,28),cmap='gray')\n",
     "    axes[idx+4].set_xticks([]) # Remove xticklabels\n",
     "    axes[idx+4].set_yticks([]) # Remove yticklabels\n",
     "    \n",
     "    # Plot original\n",
-    "    axes[idx+8].imshow(clean_valid[idx+sft_idx].reshape(28,28),cmap='gray')\n",
+    "    axes[idx+8].imshow(clean_test[idx+sft_idx].reshape(28,28),cmap='gray')\n",
     "    axes[idx+8].set_xticks([]) # Remove xticklabels\n",
     "    axes[idx+8].set_yticks([]) # Remove yticklabels\n",
     "    \n",
@@ -469,10 +452,10 @@
     "axes[8].set_ylabel(f'(Original)\\n{out_dim} x {out_dim}\\nMNIST',fontdict={'fontsize': 12, 'fontfamily': 'serif'})\n",
     "\n",
     "# Adjust layout to decrease padding between subplots\n",
-    "plt.subplots_adjust(wspace=0.1, hspace=0.25)\n",
+    "plt.subplots_adjust(wspace=0.1, hspace=0)\n",
     "\n",
     "# Save results \n",
-    "plt.savefig(f'{imdir}/{imfile}_{cmp_dim}x{cmp_dim}_valid_PAIN.png')\n",
+    "plt.savefig(f'{imdir}/{imfile}_{cmp_dim}x{cmp_dim}_test_PAIN.png')\n",
     "\n",
     "# Display the figure\n",
     "plt.show()"