pupilseg_ps.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pupil Segmentation CenterFix - Phase Shift\n",
    "\n",
    "> Author: [Zhengyi Zhan](mailto:zhanzy@zju.edu.cn)\n",
    "> \n",
    "> Date: Sept 15, 2024\n",
    "\n",
    "### Introduction\n",
    "\n",
    "This script uses a whole blazed grating on an SLM to perform pupil segmentation for wavefront sensing.\n",
    "\n",
    "It implements a 'CenterFix' pattern where one mask is fixed in the center and another moves in a rotational order.\n",
    "\n",
    "Key Features:\n",
    "- Utilizes HPC_SLM_Calib for SLM control and GxCam for camera operations\n",
    "- Generates rotational coordinates for grating mask placement\n",
    "- Applies phase shifts to the moving mask\n",
    "- Captures and saves image data for further analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initialize Camera & SLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from hpc_slm.hpc_slm import *\n",
    "from gxipy.MyGxCam import GxCam\n",
    "\n",
    "monitor_idx = 2             # Index of the monitor connected to SLM (via DVI)\n",
    "size = [1024, 1272]         # SLM resolution [height, width]\n",
    "\n",
    "cam_model_query = '1070'    # Camera model identifier\n",
    "\n",
    "# Initialize SLM & grating object\n",
    "slm = HPC_SLM_Calib(monitor_idx, 0, size)\n",
    "slm.display_img(np.zeros(size, dtype=np.uint8)) # Display a blank image on SLM to verify correct monitor\n",
    "\n",
    "# Initialize cam\n",
    "cam = GxCam(cam_model_query)\n",
    "cam.set_trigger_mode(True)      # Enable trigger mode for synchronized capture"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set Grating Cooridnates Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "center_x = 636\n",
    "center_y = 460\n",
    "\n",
    "beam_diameter = 310     # Diameter of the beam\n",
    "grating_size = 29       # Size of each grating mask\n",
    "d1 = 36                 # Distance from the center to the center of the first circle\n",
    "\n",
    "grating_type = 'blazed' # The type of the grating\n",
    "grating_step = 10       # Step size for the grating\n",
    "\n",
    "turns = 5               # Number of rotational turns\n",
    "# numbers = 9\n",
    "numbers = np.array([9, 18, 26, 30, 40]) # Custom numbers of points for each turn\n",
    "\n",
    "beam_diameter += grating_size   # Adjust beam diameter to include grating size\n",
    "\n",
    "from hpc_slm.gratings import *\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "g = Gratings(size)\n",
    "# Generate rotational coordinates and preview image\n",
    "cx, cy, img = g.gen_rotational_coors(grating_size, d1, beam_diameter, \n",
    "                                     center_x, center_y, turns, numbers, img_flag=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Preview Rotational Cooridnates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(9, 6))\n",
    "plt.imshow(img)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Acquisition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "calib_file = Path(r'D:\\.MINFLUX\\CalibData\\slm-calib\\yzt-642-100ms.h5')  # SLM calibration file\n",
    "flat_flag = True        # Apply flat field correction\n",
    "local_flag = True       # Use local LUT for calibration\n",
    "global_lut = 200        # Global LUT value (only used when local_flag is False)\n",
    "\n",
    "slm_delay = 0.1         # Delay between SLM updates (in seconds)\n",
    "frame_per_postion = 1   # Number of frames to capture per position\n",
    "\n",
    "phase_step = 4          # Number of phase steps (use 1 to disable phase shift)\n",
    "\n",
    "from tqdm import tqdm\n",
    "import time\n",
    "\n",
    "# Load SLM calibration file\n",
    "slm.load_calib_file(calib_file)\n",
    "\n",
    "# Generate a whole grating pattern\n",
    "whole_grating = g.gen_grating(grating_type, beam_diameter,\n",
    "                        grating_step = grating_step,\n",
    "                        circle_mask = True, phase_inverse = False)\n",
    "\n",
    "# Generate phase shift array\n",
    "if phase_step < 1:\n",
    "    phase_step = 1\n",
    "phase_shift_rad = np.linspace(0, 2 * np.pi, phase_step, endpoint = False)\n",
    "\n",
    "# Open camera stream\n",
    "cam.open_stream()\n",
    "time.sleep(0.1) # Short delay to ensure camera is ready\n",
    "\n",
    "# Place the whole grating into the specified beam center of canvas\n",
    "whole_grating_canvas = g.place_patterns(cx[0:1], cy[0:1], whole_grating)\n",
    "\n",
    "# Capture image with only fixed center grating to retrieve the center of interference\n",
    "grating_mask = g.gen_canvas_mask(np.array([grating_size]), cx[0:1], cy[0:1], circle_mask=True)\n",
    "grating_canvas = g.apply_canvas_mask(whole_grating_canvas, grating_mask)\n",
    "slm.phs_to_display(grating_canvas, flat_flag, local_flag, global_lut)\n",
    "time.sleep(slm_delay)\n",
    "img_center = cam.get_snapshot()\n",
    "\n",
    "# Initialize memory and set up image array\n",
    "imgs = np.zeros(img_center.shape + (phase_step, len(cx)-1,), dtype = img_center.dtype)\n",
    "\n",
    "# Main acquisition loop\n",
    "for i in tqdm(range(len(cx)-1)):\n",
    "    # Set up one fixed grating in the center and one moving grating\n",
    "    d_arr = np.array([grating_size, grating_size])\n",
    "    cx_arr = np.array([cx[0], cx[i+1]])\n",
    "    cy_arr = np.array([cy[0], cy[i+1]])\n",
    "\n",
    "    # Generate canvas mask for two circles and apply it\n",
    "    grating_mask = g.gen_canvas_mask(d_arr, cx_arr, cy_arr, circle_mask=True)\n",
    "    grating_canvas = g.apply_canvas_mask(whole_grating_canvas, grating_mask)\n",
    "\n",
    "    # Generate canvas mask for the moving circle (for phase shift only)\n",
    "    phase_shift_mask = g.gen_canvas_mask(np.array([grating_size]), cx[i+1:i+2], cy[i+1:i+2], circle_mask=True)\n",
    " \n",
    "    for j, value in enumerate(phase_shift_rad):\n",
    "        # Apply the phase shift to the moving grating\n",
    "        phase_shift_canvas = g.apply_canvas_mask(np.ones(size) * value, phase_shift_mask)\n",
    "\n",
    "        # Send pattern to SLM\n",
    "        slm.phs_to_display(grating_canvas+phase_shift_canvas, flat_flag, local_flag, global_lut)\n",
    "\n",
    "        time.sleep(slm_delay)   # Wait for SLM to ensure the pattern is changed\n",
    "\n",
    "        # Capture one frame from camera\n",
    "        imgs[:, :, j, i] = cam.get_snapshot()\n",
    "\n",
    "# Close camera stream\n",
    "cam.close_stream()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data Preview"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Preview the first 9 images\n",
    "plt.figure(figsize=(10, 10))\n",
    "for i in range(9):\n",
    "    plt.subplot(3,3,i+1)\n",
    "    plt.imshow(imgs[:,:,0,i])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Save Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.io import savemat\n",
    "from datetime import datetime\n",
    "\n",
    "# Generate filename with current timestamp\n",
    "current_time = datetime.now().strftime('%Y%m%d_%H%M%S')\n",
    "file_name = Path(f\"pupilseg_aps_{current_time}.mat\")\n",
    "final_path = Path(r\"D:\\BeamSense\\PyPupilSegAPS\") / file_name\n",
    "\n",
    "# Prepare the structs of metadata\n",
    "grating_para = {\n",
    "    'type': grating_type,\n",
    "    'size': grating_size,\n",
    "    'step': grating_step,\n",
    "    'beam_diameter': beam_diameter,\n",
    "    'd1': d1,\n",
    "    'turns': turns,\n",
    "    'numbers': numbers\n",
    "}\n",
    "\n",
    "slm_para = {\n",
    "    'local': local_flag,\n",
    "    'flat': flat_flag,\n",
    "    'calib_file': str(calib_file),\n",
    "    'delay': slm_delay\n",
    "}\n",
    "\n",
    "# Save data and metadata to '.mat'\n",
    "savemat(str(final_path), {'imgs': imgs, 'img_center': img_center,\n",
    "                          'cx': cx, 'cy': cy, \n",
    "                          'grating_para': grating_para, \n",
    "                          'slm_para': slm_para})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Release Camera & SLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "del cam\n",
    "del slm"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}