training.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import os\n",
    "#os.environ[\"CUDA_VISIBLE_DEVICES\"] = ''  # restrict GPU usage\n",
    "import matplotlib.pyplot as plt\n",
    "import tensorflow as tf\n",
    "from zstitch import zstitch, get_z_step_mm\n",
    "from tqdm.notebook import tqdm\n",
    "import scipy.io\n",
    "from mcam_loading_scripts import load_xyz"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Choose MCAM datasets to train on\n",
    "Specify the paths to the MCAM datasets and calibration files as well as the hyperparameters for the sample you wish to train on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cam_slice0=(0, 9); cam_slice1=(0, 6)  # by default, use all cameras\n",
    "camera_dims_during_acquisition = (9, 6)  # shape used during acquisition\n",
    "filters_list = np.array([32, 32, 32, 64, 64, 64])  # CNN architecture (list of filter numbers)\n",
    "\n",
    "# pick a sample:\n",
    "sample_id = 'chair_painting'  # 'chair_painting', 'PCB', 'BGA', or 'PGA'\n",
    "\n",
    "if sample_id == 'chair_painting':\n",
    "    # chair painting sample:\n",
    "    directory = '/data/20220213_chair_painting/'\n",
    "    restore_path = '/data/20220213_green_noise_target/flat_ref_optimized_params.mat'\n",
    "    restore_path_single_cam = '/data/20220213_green_noise_target/flat_ref_optimized_params_single_cam.mat'\n",
    "    filters_list = np.array([32, 32, 32, 64, 64, 64])\n",
    "    weighted_sharpness_thresholds = [1, 1.5]\n",
    "elif sample_id == 'PCB':\n",
    "    # PCB sample:\n",
    "    directory = '/home/kevin/data/20211219_PCB_160B_rechunked'\n",
    "    restore_path = '/home/kevin/data/20211219_noise_target/flat_ref_optimized_params.mat'\n",
    "    restore_path_single_cam = '/home/kevin/data/20211219_noise_target/flat_ref_optimized_params_single_cam.mat'\n",
    "    filters_list = np.array([32, 32, 32, 64, 64])\n",
    "    weighted_sharpness_thresholds = [1.5, 2]\n",
    "    cam_slice0=(0, 8)\n",
    "elif sample_id == 'BGA':\n",
    "    # BGA chips sample:\n",
    "    directory = '/data/20220207_bga_chips/'\n",
    "    restore_path = '/data/20220207_green_noise_target/flat_ref_optimized_params.mat'\n",
    "    restore_path_single_cam = '/data/20220207_green_noise_target/flat_ref_optimized_params_single_cam.mat'\n",
    "    filters_list = np.array([32, 32, 32, 64, 64, 64])\n",
    "    weighted_sharpness_thresholds = [1.5, 2]\n",
    "    cam_slice0=(0, 8)\n",
    "elif sample_id == 'PGA':\n",
    "    # pin array sample:\n",
    "    directory = '/data/20220217_pin_array/'\n",
    "    restore_path = '/data/20220217_green_noise_target/flat_ref_optimized_params.mat'\n",
    "    restore_path_single_cam = '/data/20220217_green_noise_target/flat_ref_optimized_params_single_cam.mat'\n",
    "    filters_list = np.array([32, 32, 32, 64, 64, 64])\n",
    "    weighted_sharpness_thresholds = [1, 1.5]\n",
    "    cam_slice0=(0, 8)\n",
    "    cam_slice1=(0, 3)\n",
    "    camera_dims_during_acquisition = (8, 3)\n",
    "else:\n",
    "    raise Exception('invalid sample_id')\n",
    "\n",
    "blur_sigma = 8  # gaussian blur radius for calculating sharpness\n",
    "num_patch = 2  # number of patches per batch\n",
    "patch_size = 576  # size of square patch\n",
    "skip_list = [0]*len(filters_list)  # no skip connections"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Physics self-supervised training (dynamically loading from storage)\n",
    "Train a CNN to map from z-stacks to 3D height using focus cues and stereo, across the entire (up to) 2.1-TB datasets. Since the dataset is too large to load into computer RAM, we dynamically load patches randomly from the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_model_and_dataset(im_stack_inds_keep, recon_shape, ul_offset, downsamp, nominal_z_slices, \n",
    "                               nominal_z_slices_global,  # needed for training from disk, but not per camera\n",
    "                               filters_list, skip_list, variable_initial_values, z_step_mm,\n",
    "                               num_patch=num_patch, patch_size=patch_size,  # for making dataset\n",
    "                               preferred_camera=(0, 0)\n",
    "                              ):\n",
    "    # Although when training in this mode (i.e., streaming from disk), we don't need to load full per-camera\n",
    "    # datasets into memory -- load just one (supplying im_stack_inds_keep) so that we can monitor performance\n",
    "    # during optimization.\n",
    "    \n",
    "    # create visitation_log at lower scale for all cameras:\n",
    "    a = zstitch(stack=im_stack_inds_keep, \n",
    "                ul_coords=np.zeros((len(im_stack_inds_keep), 2)),  # ul_coords will be replaced\n",
    "                recon_shape=recon_shape, \n",
    "                ul_offset=ul_offset,\n",
    "                scale=.01*downsamp,\n",
    "                batch_size=None,\n",
    "                momentum=None,\n",
    "                report_error_map=False,\n",
    "                sigma=blur_sigma,\n",
    "                truncate=2,\n",
    "                z_step_mm=z_step_mm,\n",
    "                camera_dims=camera_dims_during_acquisition,\n",
    "               )\n",
    "    a.nominal_z_slices = tf.constant(nominal_z_slices)  # tensorarray.write doesn't accept np arrays\n",
    "    a.use_camera_calibration = True\n",
    "    a.effective_focal_length_mm = 25.05486\n",
    "    a.magnification_j = 0.8448\n",
    "    a.weighted_sharpness_loss = True \n",
    "    a.weighted_sharpness_thresholds = weighted_sharpness_thresholds\n",
    "    a.create_variables(deformation_model='camera_parameters_perspective_to_orthographic',\n",
    "        learning_rates={'camera_focal_length': -1e-3, 'camera_height': 1e-3, 'ground_surface_normal': 1e-3,\n",
    "                       'camera_in_plane_angle': 1e-3, 'rc': .1, 'gain': -1e-3, 'bias': -1e-3, 'ego_height': 1e-3,\n",
    "                       'radial_camera_distortion': 1e-3},\n",
    "                       variable_initial_values=variable_initial_values,\n",
    "                       remove_global_transform=True, antialiasing_filter=False)\n",
    "    stack_downsamp, rc_downsamp = a.generate_dataset()\n",
    "    \n",
    "    loss_i, recon, normalize, error_map, tracked = a.gradient_update(stack_downsamp, rc_downsamp,\n",
    "                                                                     return_tracked_tensors=True, \n",
    "                                                                     update_gradient=False)\n",
    "    a.recompute_CNN = True  # only if you need to save memory\n",
    "    a.unet_scale = .0001\n",
    "    a.define_network_and_camera_params(tracked['vanish_warp'], tracked['camera_to_vanish_point_xyz'], \n",
    "                                       num_channels_rgb=im_stack_inds_keep.shape[-1],\n",
    "                                       architecture='fcnn', filters_list=filters_list, \n",
    "                                       skip_list=skip_list,\n",
    "                                       learning_rate=1e-3/10)\n",
    "    \n",
    "    # generate visitation log:\n",
    "    with tf.device('/CPU:0'):\n",
    "        if camera_dims_during_acquisition == (9,6):\n",
    "            # If you choose to slice a fraction of the cameras, you don't need to indicate that here if you \n",
    "            # chose to slice the cameras AFTER acquisition.\n",
    "            # e.g., if you acquired all 9x6 cameras, but wish to only use 8x5 (specified via cam_slice0, \n",
    "            # cam_slice1), you don't need to indicate that here.\n",
    "            visitation_log_vars = a.generate_visitation_log_for_all_cameras(restore_path, reuse_log=False,  \n",
    "                                                                            preferred_camera=preferred_camera\n",
    "                                                                           )\n",
    "        else:\n",
    "            # however, if you acquired fewer than 9x6 cameras, then here you need to specify the crop \n",
    "            # here via cam_slice0/1.\n",
    "            visitation_log_vars = a.generate_visitation_log_for_all_cameras(restore_path, reuse_log=False,  \n",
    "                                                                            preferred_camera=preferred_camera,\n",
    "                                                                            cam_slice0=cam_slice0, \n",
    "                                                                            cam_slice1=cam_slice1\n",
    "                                                                           )\n",
    "        \n",
    "    # generate dataset\n",
    "    a.visitation_log_scale = a.scale\n",
    "    dataset = a.generate_patched_dataset_from_disk(directory, num_patch, patch_size, nominal_z_slices_global, \n",
    "                                                   patch_size*2, prefetch=-1, sample_margin=.2, \n",
    "                                                   cam_slice0=cam_slice0, cam_slice1=cam_slice1)\n",
    "        \n",
    "    return a, dataset, visitation_log_vars\n",
    "\n",
    "def generate_full_recon_on_cpu(a):\n",
    "    # compute reconstruction for the data loaded for a single camera, for monitoring progress\n",
    "    with tf.device('/CPU:0'):\n",
    "        recon_full, _ = a.generate_full_recon(margin=2)\n",
    "        recon_full = recon_full.numpy()\n",
    "        plt.figure(figsize=(10, 10))\n",
    "        plt.imshow(recon_full[:, :, :-1].astype(np.uint8))\n",
    "        plt.show()\n",
    "        plt.figure(figsize=(10, 10))\n",
    "        plt.imshow(recon_full[:, :, -1], cmap='turbo')\n",
    "        clims = np.percentile(recon_full[::4,::4, -1][recon_full[::4,::4, -1]!=0], [0.5, 99.99])\n",
    "        plt.clim(clims)\n",
    "        plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "downsamp = 1\n",
    "z_step_ratio = 1  # ratio of z step size of flat reference to sample of interest\n",
    "camera_dims = camera_dims_during_acquisition\n",
    "xy_scans = (8, 8)\n",
    "array_dims = (camera_dims[0]*xy_scans[0], camera_dims[1]*xy_scans[1])\n",
    "num_cameras = np.prod(camera_dims)  # e.g., 54\n",
    "num_images = np.prod(array_dims)  # e.g., 3456\n",
    "image_inds = np.arange(num_images).reshape(array_dims)  # 0, ..., 3455 shaped as (72, 48)\n",
    "# ^ used to pick out from variable_initial_values\n",
    "\n",
    "ckpt_path = os.path.join(directory, 'CNN_ckpts')\n",
    "print('CNN checkpoint path: ' + ckpt_path)\n",
    "train_from_scratch = True  # if not, then load from above ckpt\n",
    "\n",
    "# get precalibrated values:\n",
    "restored = scipy.io.loadmat(restore_path)\n",
    "if 'z_step_mm__' in restored:\n",
    "    z_step_mm = restored['z_step_mm__']\n",
    "else:\n",
    "    z_step_mm = get_z_step_mm(directory)  # z step in mm\n",
    "    \n",
    "nominal_z_slices_global = restored['nominal_z_slices__'].flatten().astype(np.float32)\n",
    "z_mean = nominal_z_slices_global.mean()\n",
    "nominal_z_slices_global = np.float32((nominal_z_slices_global - z_mean)*z_step_ratio + z_mean)\n",
    "\n",
    "if 'pre_downsample_factor__' in restored:\n",
    "    pre_downsample_factor = restored['pre_downsample_factor__'].flatten()\n",
    "else:\n",
    "    pre_downsample_factor = 1\n",
    "variable_initial_values_global = {key:restored[key].squeeze() \n",
    "                           for key in restored if '__' not in key}\n",
    "variable_initial_values_global['rc'] = variable_initial_values_global['rc'] / downsamp * pre_downsample_factor\n",
    "\n",
    "restored_single_cam = scipy.io.loadmat(restore_path_single_cam)\n",
    "recon_shape = restored_single_cam['recon_shape__'].flatten()\n",
    "ul_offset = restored_single_cam['ul_offset__'].flatten()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# load a single xyz stack from one camera, for monitoring:\n",
    "rr = 3  # pick a random camera (row and column)\n",
    "cc = 2\n",
    "\n",
    "im_stack_rgb = load_xyz(directory, 'full', cam_slice0=[rr, rr+1], cam_slice1=[cc, cc+1], keep_green_only=True)\n",
    "im_stack_rgb = im_stack_rgb.transpose(1, 2, 3, 4, 5, 0)  # move z-stack dim to end\n",
    "im_stack = im_stack_rgb.squeeze()  # just use green channel; converting to grayscale takes lot of memory\n",
    "im_stack = im_stack.reshape([-1] + list(im_stack.shape[2:]))  # flatten\n",
    "\n",
    "if downsamp > 1:\n",
    "    im_stack = im_stack[:, ::downsamp, ::downsamp, :]\n",
    "    recon_shape = recon_shape // downsamp\n",
    "    ul_offset = ul_offset // downsamp\n",
    "\n",
    "# pick out variable_initial_values from variable_initial_values_global:\n",
    "variable_inds = image_inds[rr*xy_scans[0]:(rr+1)*xy_scans[0], cc*xy_scans[1]:(cc+1)*xy_scans[1]].flatten()\n",
    "variable_initial_values = variable_initial_values_global.copy()\n",
    "for var_name in variable_initial_values:\n",
    "    variable = variable_initial_values[var_name]\n",
    "    if len(variable.shape) > 0 and variable.shape[0] == num_images:\n",
    "        variable_initial_values[var_name] = variable[variable_inds]\n",
    "nominal_z_slices = nominal_z_slices_global[variable_inds]\n",
    "\n",
    "with tf.device('/CPU:0'):\n",
    "    a, dataset, visitation_log_vars = generate_model_and_dataset(im_stack, recon_shape, ul_offset, \n",
    "                                                                 downsamp, nominal_z_slices, \n",
    "                                                                 nominal_z_slices_global,\n",
    "                                                                 filters_list, skip_list,\n",
    "                                                                 variable_initial_values,\n",
    "                                                                 z_step_mm=z_step_mm,\n",
    "                                                                 num_patch=num_patch, patch_size=patch_size,\n",
    "                                                                 preferred_camera=(rr, cc))\n",
    "\n",
    "plt.imshow((visitation_log_vars[0][rr, cc].numpy()>0).sum(2))\n",
    "plt.title('visitation log')\n",
    "plt.show()\n",
    "\n",
    "# restore CNN:\n",
    "if train_from_scratch:\n",
    "    print('did not restore CNN from ckpt')\n",
    "else:\n",
    "    a.ckpt = None\n",
    "    a.checkpoint_all_variables(ckpt_path, skip_saving=True)\n",
    "    a.restore_all_variables(ckpt_no=0)\n",
    "    a.ckpt = None\n",
    "    print('restored CNN from ckpt')\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "num_iter = 100001\n",
    "\n",
    "ii = 0\n",
    "losses = list()\n",
    "\n",
    "for batch in tqdm(dataset):\n",
    "    if any([len(batch[0][i]) < 2 for i in range(num_patch)]):\n",
    "        # there needs to be at least two patches to register!\n",
    "        continue\n",
    "\n",
    "    loss_i, recon_i, normalize_i, grads_i, norm_i = a.gradient_update_patch(batch, height_map_reg_coef=.5,\n",
    "                                                                            return_gradients=True,\n",
    "                                                                            clip_gradient_norm=10,\n",
    "                                                                            dither_coords=True,\n",
    "                                                                            downsample_factor=1,\n",
    "                                                                            sharpness_reg_coef=5000,\n",
    "                                                                            stitch_loss_coef=None,\n",
    "                                                                            argmax_loss_coef=1,\n",
    "                                                                            use_hpf_for_MSE_loss=False,\n",
    "                                                                            orthorectify=False,\n",
    "                                                                           )\n",
    "    losses.append([l.numpy() for l in loss_i])\n",
    "\n",
    "    if len(losses) % 100 == 0:\n",
    "        print(len(losses), 'Loss terms: ' + str(losses[-1]))\n",
    "\n",
    "    if len(losses) % 1000 == 0 or len(losses) == 1:\n",
    "        plt.plot(losses)\n",
    "        plt.title('loss history')\n",
    "        plt.show()\n",
    "\n",
    "        if len(losses) > 51:\n",
    "            if type(loss_i) is list:\n",
    "                for i in range(len(loss_i)):\n",
    "                    plt.plot(np.convolve(np.array([loss[i] for loss in losses][50:]), np.ones(200)/200, 'valid'))\n",
    "                    plt.show()\n",
    "            else:\n",
    "                plt.plot(np.convolve(np.array(losses[50:]).squeeze(), np.ones(200)/200, 'valid'))\n",
    "                plt.title('blurred loss')\n",
    "                plt.show()\n",
    "\n",
    "        plt.subplot(121)\n",
    "        plt.imshow(np.uint8(recon_i[:,:,0, :-1]))\n",
    "        plt.title('patch green channel')\n",
    "        plt.subplot(122)\n",
    "        plt.imshow(recon_i[:,:,0, -1], cmap='jet')\n",
    "        plt.title('patch height')\n",
    "        plt.show()\n",
    "\n",
    "    if ii % 10000 == 0 and ii != 0:\n",
    "        # check progress on single-camera data:\n",
    "        generate_full_recon_on_cpu(a)\n",
    "\n",
    "    if ii == num_iter:\n",
    "        break\n",
    "    else:\n",
    "        ii+=1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save CNN in a unique directory:\n",
    "CNN_ckpt_path = os.path.join(directory, 'CNN_ckpts')\n",
    "\n",
    "a.ckpt = None\n",
    "a.checkpoint_all_variables(path=CNN_ckpt_path)\n",
    "print(ckpt_path)\n",
    "print('checkpointed network')"
   ]
  }
 ],
 "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}