{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77a1e4ec",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "box_size:  tensor([[-0.9058,  1.2208],\n",
      "        [-0.9058,  1.2208],\n",
      "        [-0.9058,  1.2208]], dtype=torch.float64) box_shape [200 200 200]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.10/dist-packages/torch/autograd/function.py:573: UserWarning: Using torch.cross without specifying the dim arg is deprecated.\n",
      "Please either pass the dim explicitly or simply use torch.linalg.cross.\n",
      "The default value of dim will change to agree with that of linalg.cross in a future release. (Triggered internally at /opt/pytorch/pytorch/aten/src/ATen/native/Cross.cpp:62.)\n",
      "  return super().apply(*args, **kwargs)  # type: ignore[misc]\n"
     ]
    }
   ],
   "source": [
    "from deltamic import Gaussian_psf, compute_spatial_frequency_grid, normalize_tensor, render_image_from_ftmesh, Fourier3dMesh, compute_box_size, generate_gaussian_psf\n",
    "import trimesh\n",
    "import numpy as np\n",
    "import torch\n",
    "import skimage.io as io\n",
    "from skimage import img_as_ubyte\n",
    "\n",
    "\"\"\"\n",
    "Creation of\n",
    "initial\n",
    "configurations: \n",
    "\"\"\"\n",
    "\n",
    "assert torch.cuda.is_available()\n",
    "\n",
    "device = 'cuda:0'\n",
    "box_shape = np.array([200]*3)\n",
    "\n",
    "filename = \"../data/mesh/cow.obj\"\n",
    "\n",
    "Mesh_gt = trimesh.load(filename)\n",
    "faces = np.array(Mesh_gt.faces)\n",
    "verts = np.array(Mesh_gt.vertices)\n",
    "\n",
    "box_size = compute_box_size(verts, offset = 0.2)\n",
    "Verts = torch.tensor(verts, dtype = torch.float, device = device,requires_grad = True)\n",
    "Faces = torch.tensor(faces, dtype = torch.long, device = device)\n",
    "Faces_coeff = torch.ones(len(Faces),dtype = torch.float, device = device)\n",
    "box_size = torch.tensor(box_size)\n",
    "\n",
    "print(\"box_size: \",box_size, \"box_shape\",box_shape)\n",
    "\n",
    "narrowband_thresh = torch.tensor(0,dtype = torch.float, device = device)\n",
    "meshFT = Fourier3dMesh(box_size,box_shape,device=device, dtype = torch.float32)\n",
    "\n",
    "sigma_matrix = (1e-3*torch.eye(3,device=device))\n",
    "PSF = generate_gaussian_psf(sigma_matrix,meshFT.xi0,meshFT.xi1,meshFT.xi2).to(device)\n",
    "OTF = torch.abs(torch.fft.fftshift(torch.fft.fftn(PSF)))\n",
    "\n",
    "\"\"\"\n",
    "Image creation:\n",
    "\"\"\"\n",
    "ftmesh = meshFT(Verts,Faces, Faces_coeff)\n",
    "\n",
    "image_sp=normalize_tensor(render_image_from_ftmesh(ftmesh, OTF, box_shape))\n",
    "\n",
    "io.imsave(\"../data/image/cow.tif\",img_as_ubyte(image_sp.detach().cpu().numpy()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "54b68406",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Vertices shape: torch.Size([3225, 3]) tensor(1.0490, device='cuda:0', grad_fn=<MaxBackward1>) tensor(-0.7368, device='cuda:0', grad_fn=<MinBackward1>)\n",
      "tensor([[ 0.3488, -0.3350, -0.0832],\n",
      "        [ 0.3488, -0.3350, -0.0832],\n",
      "        [ 0.3131, -0.3991,  0.8812],\n",
      "        [ 0.3131, -0.3991,  0.8812],\n",
      "        [ 0.2668,  0.1816,  0.1227]], device='cuda:0',\n",
      "       grad_fn=<SliceBackward0>)\n",
      "Faces shape: torch.Size([5856, 3]) tensor(3224, device='cuda:0') tensor(0, device='cuda:0')\n",
      "tensor([[889, 885, 886],\n",
      "        [267, 886, 885],\n",
      "        [272, 888, 887],\n",
      "        [889, 887, 888],\n",
      "        [269, 892, 891]], device='cuda:0')\n"
     ]
    }
   ],
   "source": [
    "# 3D coordinates of Vertices\n",
    "print(\"Vertices shape:\", Verts.shape, torch.max(Verts), torch.min(Verts))\n",
    "print(Verts[:5, :])\n",
    "\n",
    "# Faces of the mesh - 3 vertices per triangular face - their respective indices\n",
    "print(\"Faces shape:\", Faces.shape, torch.max(Faces), torch.min(Faces))\n",
    "print(Faces[:5, :])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f024de62",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "box_size:  tensor([[-0.9058,  1.2208],\n",
      "        [-0.9058,  1.2208],\n",
      "        [-0.9058,  1.2208]], dtype=torch.float64) Box_ft_shape [200 200 200]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Current_Loss: [0.006070143077522516]: 100%|██████████| 11/11 [01:15<00:00,  6.84s/it]\n"
     ]
    }
   ],
   "source": [
    "from deltamic import normalize_tensor,render_image_from_ftmesh,Fourier3dMesh,Gaussian_psf, center_verts_in_box\n",
    "from tqdm import tqdm\n",
    "import trimesh\n",
    "import numpy as np\n",
    "import os\n",
    "import torch\n",
    "import skimage.io as io \n",
    "from largesteps.optimize import AdamUniform\n",
    "from largesteps.geometry import compute_matrix\n",
    "from largesteps.solvers import CholeskySolver\n",
    "from time import time\n",
    "\n",
    "\n",
    "folder_result = \"Results_cow/\"\n",
    "def create_dir(folder_name):\n",
    "    try: os.mkdir(folder_name)\n",
    "    except: pass\n",
    "  \n",
    "create_dir(folder_result)\n",
    "create_dir(folder_result+\"Images\")\n",
    "create_dir(folder_result+\"Meshes\")\n",
    "\n",
    "\"\"\"\n",
    "Parameters:\n",
    "\"\"\"\n",
    "\n",
    "device = 'cuda:0'\n",
    "name_micim = \"../data/image/cow.tif\"\n",
    "\n",
    "\"\"\"\n",
    "Creation of\n",
    "initial \n",
    "configurations: \n",
    "\"\"\"\n",
    "\n",
    "micim = normalize_tensor(torch.tensor(io.imread(name_micim).astype(np.float32),dtype=torch.float)).to(device)\n",
    "\n",
    "box_size = np.array([[-0.9058,  1.2208],\n",
    "                     [-0.9058,  1.2208],\n",
    "                     [-0.9058,  1.2208]])\n",
    "mesh = trimesh.primitives.Sphere(subdivisions = 4, radius=0.5)\n",
    "verts, faces = np.array(mesh.vertices)*2,np.array(mesh.faces)\n",
    "verts = center_verts_in_box(verts, box_size, offset=.4)\n",
    "Verts = torch.tensor(verts, dtype = torch.float, device = device)\n",
    "Faces = torch.tensor(faces, dtype = torch.long, device = device)\n",
    "Faces_coeff = torch.ones(len(Faces),dtype = torch.float, device = device)\n",
    "box_size = torch.tensor(box_size)\n",
    "box_shape= np.array([200, 200, 200])\n",
    "\n",
    "\n",
    "meshFT = Fourier3dMesh(box_size, box_shape, device=device, dtype = torch.float32,narrowband_thresh=1e-2)\n",
    "\n",
    "#Gaussian PSF\n",
    "model = Gaussian_psf(box_shape, box_size[:,1], sigma = 1e3, device = device)\n",
    "\n",
    "print(\"box_size: \",box_size, \"Box_ft_shape\",box_shape)\n",
    "\n",
    "\"\"\"\n",
    "Large \n",
    "Steps \n",
    "Optimization\n",
    "\"\"\"\n",
    "\n",
    "lr_base = 0.01\n",
    "lambda_=50.0\n",
    "alpha = 1.0\n",
    "Verts.requires_grad = True \n",
    "optimizer_geometry = AdamUniform([{'params': Verts}], lr=lr_base)\n",
    "optimizer_psf = torch.optim.Adam(model.parameters(), lr = lr_base)\n",
    "\n",
    "\n",
    "losses_mse = []\n",
    "losses_mse_wt = []\n",
    "\n",
    "M = compute_matrix(Verts, Faces, lambda_, device=device)\n",
    "solver = CholeskySolver(M@M)\n",
    "\n",
    "for k in (pbar:=tqdm(range(501))):#10000))):\n",
    "\n",
    "    #Loop of the optimization of the position of the vertices once we got a good approximation of the PSF\n",
    "    ###############################\n",
    "    optimizer_geometry.zero_grad()\n",
    "    ###############################\n",
    "    \n",
    "    PSF = model.forward()\n",
    "    OTF = torch.abs(torch.fft.fftshift(torch.fft.fftn(PSF)))\n",
    "    meshFT.OTF = OTF\n",
    "    \n",
    "    ftmesh = meshFT(Verts, Faces, Faces_coeff)\n",
    "    image_sp=normalize_tensor(render_image_from_ftmesh(ftmesh, OTF, box_shape))\n",
    "\n",
    "    loss_mse = torch.mean((image_sp-micim)**2)\n",
    "    loss_mse_wt = torch.mean(((image_sp-micim)**2)*(micim**alpha))\n",
    "\n",
    "    loss = loss_mse_wt\n",
    "    loss.backward()\n",
    "\n",
    "    pbar.set_description(\"Current_Loss: \"+str([loss.item()]))\n",
    "\n",
    "    with torch.no_grad():\n",
    "        Verts.grad = solver.solve(Verts.grad)\n",
    "\n",
    "    if k%100==0\n",
    "        np.savez(folder_result+'Meshes/'+str(k)+\".npz\", Verts=Verts.detach().cpu().numpy(), Faces=Faces.detach().cpu().numpy())\n",
    "        io.imsave(folder_result+'Images/'+str(k)+\".tif\", image_iter.detach().cpu().numpy()*255.0)\n",
    "    ###############################\n",
    "    optimizer_geometry.step()\n",
    "    optimizer_psf.step()\n",
    "    ###############################\n",
    "\n",
    "    losses_mse.append(loss_mse.item())\n",
    "    losses_mse_wt.append(loss_mse_wt.item())\n",
    "    np.save(folder_result+'Meshes/loss_mse.npy',losses_mse)\n",
    "    np.save(folder_result+'Meshes/loss_mse_wt.npy',losses_mse_wt)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c927ecc8",
   "metadata": {},
   "source": [
    "# Checking Rendered Initial Image against GT Image"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bbd98090",
   "metadata": {},
   "source": [
    "### Visualise PSF and OTF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47419f8c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Displaying interactive visualization...\n",
      "PSF shape: (200, 200, 200) \n",
      "OTF shape: (200, 200, 200)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 22604 characters]",
      "text/plain": [
       "<Figure size 1000x500 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "print(\"Displaying interactive visualization...\")\n",
    "\n",
    "PSF_np = PSF.detach().cpu().numpy()\n",
    "OTF_np = OTF.detach().cpu().numpy()\n",
    "print(\"PSF shape:\", PSF_np.shape, \"\\nOTF shape:\", OTF_np.shape)\n",
    "\n",
    "fig, axes = plt.subplots(1, 2, figsize=(10, 5))\n",
    "plt.subplots_adjust(bottom=0.2)\n",
    "\n",
    "z = PSF_np.shape[0] // 2\n",
    "im1 = axes[0].imshow(PSF_np[z], cmap=\"gray\")\n",
    "axes[0].set_title(\"PSF\")\n",
    "\n",
    "im2 = axes[1].imshow(OTF_np[z], cmap=\"gray\")\n",
    "axes[1].set_title(\"OTF\")\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83a884c5",
   "metadata": {},
   "source": [
    "# Visualise the Middle Slice of the Images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0fbe07d7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Displaying interactive visualization...\n",
      "Image shape: (200, 200, 200) \n",
      "Ground truth shape: (200, 200, 200)\n",
      "z slice:  100\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 148228 characters]",
      "text/plain": [
       "<Figure size 1500x500 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "print(\"Displaying interactive visualization...\")\n",
    "\n",
    "image_sp_np = image_sp.detach().cpu().numpy()\n",
    "micim_np = micim.detach().cpu().numpy()\n",
    "image_sp_final = io.imread('Results_cow/Images/500.tif')\n",
    "image_sp_final /= 255.0\n",
    "\n",
    "print(\"Image shape:\", image_sp_np.shape, \"\\nGround truth shape:\", micim_np.shape)\n",
    "\n",
    "fig, axes = plt.subplots(1, 3, figsize=(15, 5))\n",
    "plt.subplots_adjust(bottom=0.2)\n",
    "\n",
    "z = 100\n",
    "print(\"z slice: \", z)\n",
    "im1 = axes[0].imshow(micim_np[z], cmap=\"gray\")\n",
    "axes[0].set_title(\"Ground Truth\")\n",
    "\n",
    "im2 = axes[1].imshow(image_sp_np[z], cmap=\"gray\")\n",
    "axes[1].set_title(\"Rendered Initial Mesh\")\n",
    "\n",
    "im3 = axes[2].imshow(image_sp_final[z], cmap=\"gray\")\n",
    "axes[2].set_title(\"Rendered Final Mesh\")\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0eceb39d",
   "metadata": {},
   "source": [
    "\n",
    "# Visualise DeltaMic on Fluorescenc Microscopy Images of Embryos"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "id": "6610b2a9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 244264 characters]",
      "text/plain": [
       "<Figure size 1500x1000 with 9 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "celegan_gt = io.imread(\"../data/image/Sample01_007.tif\")\n",
    "celegan_1k = io.imread(\"Results_Sample01_007/Images/100.tif\")\n",
    "celegan_10k = io.imread(\"Results_Sample01_007/Images/10000.tif\")\n",
    "\n",
    "fig, axes = plt.subplots(3, 3, figsize=(15, 10))\n",
    "plt.subplots_adjust(bottom=0.2)\n",
    "\n",
    "x,y,z = (122, 122, 87)\n",
    "\n",
    "# Raw Data\n",
    "im1 = axes[0,0].imshow(celegan_gt[:, :, z], cmap=\"gray\")\n",
    "im1 = axes[0,1].imshow(celegan_gt[:, y, :], cmap=\"gray\")\n",
    "im1 = axes[0,2].imshow(celegan_gt[x], cmap=\"gray\")\n",
    "axes[0,0].set_title(\"Z Section\")\n",
    "\n",
    "# Reconstruction after 100 iterations\n",
    "im2 = axes[1,0].imshow(celegan_1k[:, :, z], cmap=\"gray\")\n",
    "im2 = axes[1,1].imshow(celegan_1k[:, y, :], cmap=\"gray\")\n",
    "im2 = axes[1,2].imshow(celegan_1k[x], cmap=\"gray\")\n",
    "axes[0,1].set_title(\"Y Section\")\n",
    "\n",
    "# Reconstruction after 10000 iterations\n",
    "im3 = axes[2,0].imshow(celegan_10k[:, :, z], cmap=\"gray\")\n",
    "im3 = axes[2,1].imshow(celegan_10k[:, y, :], cmap=\"gray\")\n",
    "im3 = axes[2,2].imshow(celegan_10k[x], cmap=\"gray\")\n",
    "axes[0,2].set_title(\"X Section\")\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1043cd1",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
