Eval_Graphene_model.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Gomb - Net\n",
    "### Performance test on the Graphene dataset\n",
    "### And analysis of Graphene experimental data\n",
    "\n",
    "Austin Houston"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[![OpenInColab](https://colab.research.google.com/assets/colab-badge.svg)](\n",
    "    https://colab.research.google.com/github/ahoust17/Gomb-Net/blob/main/Eval_Graphene_model.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# basics\n",
    "import os\n",
    "import sys\n",
    "import numpy as np\n",
    "\n",
    "# plotting\n",
    "import matplotlib.pylab as plt\n",
    "import matplotlib.colors as mcolors\n",
    "from matplotlib.colors import Normalize\n",
    "from matplotlib import cm\n",
    "\n",
    "# colab interactive plots and drive\n",
    "drive = False\n",
    "if 'google.colab' in sys.modules:\n",
    "    from  google.colab import drive \n",
    "    from google.colab import output\n",
    "    drive.mount('/content/drive')\n",
    "    output.enable_custom_widget_manager()\n",
    "    drive = True\n",
    "else:\n",
    "    %matplotlib widget\n",
    "\n",
    "# other imports\n",
    "from scipy.ndimage import label, center_of_mass, gaussian_filter, zoom, uniform_filter\n",
    "from scipy.spatial import KDTree\n",
    "from scipy.interpolate import griddata\n",
    "from scipy.stats import norm, gaussian_kde\n",
    "from skimage.filters import threshold_otsu\n",
    "from skimage.feature import blob_log\n",
    "\n",
    "# for cropping function\n",
    "if drive:\n",
    "    print('installing DataGenSTEM')\n",
    "    !pip install ase\n",
    "    !git clone https://github.com/ahoust17/DataGenSTEM.git\n",
    "    sys.path.append('./DataGenSTEM/DataGenSTEM')\n",
    "    import data_generator as dg\n",
    "\n",
    "# for Gomb-Net\n",
    "if drive:\n",
    "    print('installing Gomb-Net')\n",
    "    !git clone https://github.com/ahoust17/Gomb-Net.git\n",
    "    sys.path.append('./Gomb-Net/GombNet')    \n",
    "from GombNet.networks import *\n",
    "from GombNet.loss_func import GombinatorialLoss\n",
    "from GombNet.utils import *\n",
    "\n",
    "import torch\n",
    "# Check if CUDA is available\n",
    "if torch.cuda.is_available():\n",
    "    device = torch.device(\"cuda\")\n",
    "    print(\"CUDA is available. Using GPU.\")\n",
    "else:\n",
    "    device = torch.device(\"cpu\")\n",
    "    print(\"CUDA is not available. Using CPU.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set mps, just for my computer\n",
    "# device = torch.device('mps')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Now, you need to add the following shared drive to your google drive:\n",
    "*** WARNING: it is a big file.  Check before you download ***\n",
    "\n",
    "\n",
    "https://drive.google.com/file/d/1DyKtrmJ8wNYQg3YEJ8_iXjz6lB_DQfwy/view?usp=sharing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run the following cell after the download is complete"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# should be someting like 'content/drive/My Drive/Gomb-Net files'\n",
    "if drive:\n",
    "    shared_folder = 'drive/My Drive/Gomb-Net files'\n",
    "else:\n",
    "    shared_folder = '/Users/austin/Desktop/Gomb-Net aux files'\n",
    "\n",
    "print('available files & directories:')\n",
    "!ls '{shared_folder}'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Now, on to running the code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "let's look at the dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create dataloaders\n",
    "images_dir = str(shared_folder + '/Graphene_dataset/images')\n",
    "labels_dir = str(shared_folder + '/Graphene_dataset/labels')\n",
    "train_loader, val_loader, test_loader = get_dataloaders(images_dir, labels_dir, batch_size = 1, val_split=0.2, test_split=0.1, seed = 42)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_iter = 3\n",
    "test = test_loader.dataset[test_iter][0].unsqueeze(0)\n",
    "gt = test_loader.dataset[test_iter][1]\n",
    "\n",
    "fig, ax = plt.subplots(1, 3, figsize=(10, 5))\n",
    "ax[0].imshow(test[0, 0].cpu().numpy(), cmap='gray')\n",
    "ax[0].set_title('Input')\n",
    "\n",
    "titles = ['L1: C', 'L2: C']\n",
    "for i in range(2):\n",
    "    ax[i+1].imshow(gt[i].cpu().numpy(), cmap='gray')\n",
    "    ax[i+1].set_title(titles[i])\n",
    "for a in ax:\n",
    "    a.axis('off')\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "now let's look at the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize model\n",
    "input_channels = 1\n",
    "num_classes = 2\n",
    "num_filters = [32, 64, 128, 256]\n",
    "\n",
    "model = TwoLeggedUnet(input_channels, num_classes, num_filters, dropout = 0.2)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)\n",
    "loss = GombinatorialLoss(group_size = num_classes//2, loss = 'Dice', epsilon=1e-6, class_weights = None, alpha=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the number of trainable parameters\n",
    "def count_trainable_parameters(model):\n",
    "    return sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "num_trainable_params = count_trainable_parameters(model)\n",
    "print(f\"Number of trainable parameters: {num_trainable_params}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "visualize the training history for the pretrained model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_history = np.load(str(shared_folder + '/Pretrained_models/Graphene_model_loss_history.npz'))\n",
    "train_loss = loss_history['train_loss_history']\n",
    "val_loss = loss_history['val_loss_history']\n",
    "\n",
    "plt.figure(figsize = (6,4))\n",
    "plt.plot(train_loss, label='training', color = '#1f77b4')\n",
    "plt.plot(val_loss, label='validation', color = '#d62728')\n",
    "plt.xlabel('Epoch')\n",
    "plt.xlim(0,30)\n",
    "plt.legend(title='Losses')\n",
    "plt.tight_layout()  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "load in the pretrained weights onto our model 'skeleton'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_path = str(shared_folder + '/Pretrained_models/Graphene_model.pth')\n",
    "\n",
    "checkpoint = torch.load(model_path, map_location=device)\n",
    "model.load_state_dict(checkpoint['model_state_dict'])\n",
    "optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_iter = 0\n",
    "\n",
    "test = test_loader.dataset[test_iter][0].unsqueeze(0)\n",
    "gt = test_loader.dataset[test_iter][1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate prediction\n",
    "with torch.no_grad():\n",
    "    #test.to(device)\n",
    "    probability = model(test)\n",
    "    prediction = F.sigmoid(probability)#>0.50\n",
    "probability = probability.squeeze().cpu().numpy() \n",
    "prediction = prediction.squeeze().cpu().numpy()\n",
    "\n",
    "threshold = threshold_otsu(prediction)\n",
    "prediction = (prediction > threshold).astype(float)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.imshow(test.squeeze().cpu().numpy(), cmap='gray')\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plotting\n",
    "num_classes\n",
    "fig, axs = plt.subplots(3,num_classes,dpi = 300, sharex=True, sharey=True)\n",
    "\n",
    "for i in range(num_classes):\n",
    "    axs[0,i].imshow(gt[i], cmap='gray')\n",
    "\n",
    "for i in range(num_classes):\n",
    "    axs[1,i].imshow(prediction[i], cmap='gray')\n",
    "\n",
    "for i in range(num_classes)[:1]:\n",
    "    axs[2,i].imshow(probability[i], cmap='plasma')\n",
    "for i in range(num_classes)[1:]:\n",
    "    axs[2,i].imshow(probability[i], cmap='viridis')\n",
    "\n",
    "\n",
    "for ax in axs.ravel():\n",
    "    ax.axis('off')\n",
    "\n",
    "axs[0,0].set_ylabel('GrounTruth')\n",
    "axs[1,0].set_ylabel('Prediction')\n",
    "axs[2,0].set_ylabel('Probability')\n",
    "\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Comparison to blob-finder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_number = 0\n",
    "# try a regular blob finder for comparison\n",
    "blob_im = test.squeeze().cpu().numpy()\n",
    "\n",
    "blobs = blob_log(blob_im, min_sigma=1, max_sigma=20, num_sigma=5, threshold=0.1)\n",
    "blobs_com = [center_of_mass(blob_im, blobs[i, 0], blobs[i, 1]) for i in range(blobs.shape[0])]\n",
    "\n",
    "plt.figure()\n",
    "plt.imshow(blob_im, cmap='gray')\n",
    "plt.scatter(blobs[:, 1], blobs[:, 0], c='r', s=20)\n",
    "plt.axis('off')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The following code is marked out, because it takes a few minutes to run.  Just measuring network performance metrics across the test dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pwa_total = 0\n",
    "# dice_total = 0\n",
    "# IoU_total = 0\n",
    "# \n",
    "# def iou(pred, gt):\n",
    "#     intersection = np.logical_and(pred, gt).sum()\n",
    "#     union = np.logical_or(pred, gt).sum()\n",
    "#     return intersection / union\n",
    "# \n",
    "# def dice_coefficient(pred, gt):\n",
    "#     intersection = np.logical_and(pred, gt).sum()\n",
    "#     return 2 * intersection / (pred.sum() + gt.sum())\n",
    "# \n",
    "# # Calculate the accuracy\n",
    "# for i in range(len(test_loader)):\n",
    "#     test = test_loader.dataset[i][0].unsqueeze(0)\n",
    "#     gt = test_loader.dataset[i][1].numpy()  # Convert to numpy array\n",
    "#     \n",
    "#     # Switch ground truth layers\n",
    "#     gt_switched = np.flip(gt, axis=0)\n",
    "#     \n",
    "#     with torch.no_grad():\n",
    "#         probability = model(test)\n",
    "#         prediction = torch.sigmoid(probability)  # Use torch.sigmoid instead of F.sigmoid (deprecated)\n",
    "#     \n",
    "#     probability = probability.squeeze().cpu().numpy()\n",
    "#     prediction = prediction.squeeze().cpu().numpy()\n",
    "# \n",
    "#     threshold = threshold_otsu(prediction)\n",
    "#     prediction = (prediction > threshold).astype(float)\n",
    "#     \n",
    "#     # Calculate metrics for original and switched ground truths\n",
    "#     pwa_original = np.sum(prediction == gt) / np.prod(gt.shape)\n",
    "#     pwa_switched = np.sum(prediction == gt_switched) / np.prod(gt_switched.shape)\n",
    "#     \n",
    "#     dice_original = dice_coefficient(prediction, gt)\n",
    "#     dice_switched = dice_coefficient(prediction, gt_switched)\n",
    "#     \n",
    "#     iou_original = iou(prediction, gt)\n",
    "#     iou_switched = iou(prediction, gt_switched)\n",
    "#     \n",
    "#     # Take the highest value for each metric\n",
    "#     pwa_total += max(pwa_original, pwa_switched)\n",
    "#     dice_total += max(dice_original, dice_switched)\n",
    "#     IoU_total += max(iou_original, iou_switched)\n",
    "# \n",
    "# # Calculate the average for each metric\n",
    "# pwa_total /= len(test_loader)\n",
    "# dice_total /= len(test_loader)\n",
    "# IoU_total /= len(test_loader)\n",
    "# \n",
    "# print(f\"Pixel-wise Accuracy: {pwa_total}\")\n",
    "# print(f\"Mean Dice Coefficient: {dice_total}\")\n",
    "# print(f\"Mean IoU: {IoU_total}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Now, on Experimental data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "exp_data = np.load(str(shared_folder + '/Experimental_datasets/moire.npz'))\n",
    "im_array = exp_data['im_array']\n",
    "pixel_size = exp_data['pixel_size']\n",
    "\n",
    "\n",
    "# im_array = gaussian_filter(im_array, sigma=1)\n",
    "im_array = im_array - np.min(im_array)\n",
    "im_array = im_array / np.max(im_array)\n",
    "# zoom_factor = 0.7\n",
    "# \n",
    "# im_array = dg.resize_image(im_array, zoom_factor * 512)\n",
    "\n",
    "print(f\"Pixel size: {pixel_size.astype(float)} m/pix\")\n",
    "plt.figure()\n",
    "plt.imshow(im_array, cmap='gray')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "images = dg.shotgun_crop(im_array, crop_size=256, n_crops = 5, roi = None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(5, 5, figsize=(10, 10), dpi = 300)\n",
    "masks = np.zeros((5, 2, 256, 256))\n",
    "\n",
    "for i, im in enumerate(images):\n",
    "    # make nn prediction\n",
    "    im = im.astype(np.float32)\n",
    "    im = torch.tensor(im).unsqueeze(0).unsqueeze(0)\n",
    "    with torch.no_grad():\n",
    "        probability = model(im)\n",
    "        prediction = torch.sigmoid(probability)\n",
    "    probability = probability.squeeze().cpu().numpy()\n",
    "    prediction = prediction.squeeze().cpu().numpy()\n",
    "    threshold = threshold_otsu(prediction)\n",
    "    prediction = (prediction > threshold).astype(float)\n",
    "    masks[i] = prediction\n",
    "\n",
    "\n",
    "    # plot\n",
    "    ax[i, 0].imshow(im.squeeze().numpy(), cmap='gray')\n",
    "    ax[i, 1].imshow(prediction[0], cmap='gray')\n",
    "    ax[i, 2].imshow(prediction[1], cmap='gray')\n",
    "    ax[i, 3].imshow(probability[0], cmap='plasma')\n",
    "    ax[i, 4].imshow(probability[1], cmap='viridis')\n",
    "\n",
    "for a in ax.ravel():\n",
    "    a.axis('off')\n",
    "fig.tight_layout()\n",
    "\n",
    "#plt.savefig('moire_segmentation.png')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resize_factor = 4\n",
    "zoom_order = 3\n",
    "i = 0\n",
    "dist_hist = []\n",
    "colors = ['#d62728','#1f77b4']\n",
    "for image, mask in zip(images, masks):\n",
    "    # fig, ax = plt.subplots(1, 2, dpi=300, figsize=(5, 5), subplot_kw={'aspect': 'equal'})\n",
    "    for layer, color, a in zip(mask, colors, ax):\n",
    "        # resize layer\n",
    "        layer = zoom(layer, resize_factor, order=zoom_order)\n",
    "        layer = gaussian_filter(layer, sigma=1)\n",
    "        threshold = threshold_otsu(layer)\n",
    "        layer = (layer > threshold).astype(float)\n",
    "\n",
    "        plt.figure()\n",
    "        plt.imshow(layer, cmap='gray')\n",
    "\n",
    "        labeled_array, num_features = label(layer)\n",
    "        centroids = center_of_mass(layer, labeled_array, range(1, num_features + 1))\n",
    "        centroids = np.array(centroids)\n",
    "\n",
    "        # Calculate the distance between the centroids\n",
    "        tree = KDTree(centroids)\n",
    "        distances, indices = tree.query(centroids, k=3)\n",
    "        nearest_distances = distances[:, 1:] * float(pixel_size) # angstroms\n",
    "        dist_hist.append(nearest_distances.flatten())\n",
    "\n",
    "    #     a.scatter(centroids[:, 1], centroids[:, 0], s=38, c=color, edgecolors='k', linewidths=0.5)\n",
    "    #     a.axis('off')\n",
    "    #     a.set_xlim(0, layer.shape[1])\n",
    "    #     a.set_ylim(layer.shape[0], 0)  # Flip the y-axis to match image coordinates\n",
    "    #     \n",
    "    # plt.subplots_adjust(wspace=0, hspace=0)\n",
    "    #plt.savefig(f'atoms_{i}.png', transparent=True, bbox_inches='tight', pad_inches=0)\n",
    "    i += 1\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# make dist_hist a 1D array\n",
    "dist_hist = np.concatenate(dist_hist) * 1e10 / resize_factor\n",
    "\n",
    "avg_dist = np.mean(dist_hist)\n",
    "\n",
    "plt.figure(figsize = (8,8), dpi=300)\n",
    "plt.hist(dist_hist, bins=100, color='gray')\n",
    "\n",
    "plt.vlines(1.42, 0, 1000, color='k', label='C-C bond 1.42 Å')\n",
    "plt.vlines(avg_dist, 0, 1000, color='k', linestyle = '--', label=f'Avg. bond {avg_dist:.2f} Å')\n",
    "plt.xlabel('Distance (Å)', fontsize=24)\n",
    "plt.xlim(0.8,2)\n",
    "plt.ylim(0,850)\n",
    "plt.legend(loc='upper right', fontsize=24)\n",
    "# set xtick fontsize\n",
    "plt.xticks(fontsize=24)\n",
    "plt.yticks([])\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Average distance: {avg_dist:.2f} Å\")\n",
    "\n",
    "# FWHM\n",
    "half_max = np.max(np.histogram(dist_hist, bins=100)[0]) / 2\n",
    "hist, bin_edges = np.histogram(dist_hist, bins=100)\n",
    "bin_centers = (bin_edges[:-1] + bin_edges[1:]) / 2\n",
    "fwhm = bin_centers[hist > half_max]\n",
    "fwhm = fwhm[-1] - fwhm[0]\n",
    "print(f\"FWHM: {fwhm:.2f} Å\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pixel_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torch-nightly",
   "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}