Example usage Derrick.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-04-25 13:29:40.739298: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
      "2024-04-25 13:29:40.771789: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
      "2024-04-25 13:29:40.771816: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
      "2024-04-25 13:29:40.772746: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
      "2024-04-25 13:29:40.778582: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2024-04-25 13:29:41.426389: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from tensorflow.keras import layers, models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "connected components I"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import tensorflow as tf\n",
    "\n",
    "# def connected_components_tf(image, num_iterations=100):\n",
    "#     if not tf.is_tensor(image):\n",
    "#         raise TypeError(f\"Input type is not a tf.Tensor. Got: {type(image)}\")\n",
    "    \n",
    "#     if not isinstance(num_iterations, int) or num_iterations < 1:\n",
    "#         raise TypeError(\"Input num_iterations must be a positive integer.\")\n",
    "    \n",
    "#     if len(image.shape) < 3 or image.shape[-3] != 1:\n",
    "#         raise ValueError(f\"Input image shape must be (*, 1, H, W). Got: {image.shape}\")\n",
    "    \n",
    "#     # Reshape image to 2D if it's more than 3D\n",
    "#     shape = tf.shape(image)\n",
    "#     image = tf.reshape(image, [-1, shape[-2], shape[-3]])\n",
    "    \n",
    "#     # Create initial labels\n",
    "#     B, H, W = image.shape\n",
    "#     labels = tf.reshape(tf.range(B * H * W, dtype=image.dtype), [B, H, W])\n",
    "#     mask = tf.equal(image, 1)\n",
    "    \n",
    "#     # Initialize labels as zero where mask is False\n",
    "#     labels = tf.where(mask, labels, tf.zeros_like(labels))\n",
    "    \n",
    "#     for _ in range(num_iterations):\n",
    "#         # Max pool current labels to simulate the dilation effect\n",
    "#         labels = tf.expand_dims(labels, axis=3)  # Add a channel dimension\n",
    "#         pooled_labels = tf.nn.max_pool2d(labels, ksize=3, strides=1, padding='SAME')\n",
    "#         labels = tf.squeeze(pooled_labels, axis=[-1])  # Remove the channel dimension\n",
    "#         labels = tf.where(mask, labels, tf.zeros_like(labels))\n",
    "    \n",
    "#     return tf.reshape(labels, shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import tensorflow as tf\n",
    "# import matplotlib.pyplot as plt\n",
    "\n",
    "# # Create a sample binary image\n",
    "# H, W = 10, 10  # Image dimensions\n",
    "# image = tf.zeros((1, 1, H, W), dtype=tf.float32)  # Create a black image\n",
    "# image = image.numpy()  # Convert to numpy array for manipulation\n",
    "# image[0, 0, 2:4, 2:4] = 1  # Add a white square\n",
    "# image[0, 0, 6:8, 7:9] = 1  # Add another white square\n",
    "# image[0, 0, 2:4, 7:9] = 1  # Add a third white square\n",
    "# image = tf.convert_to_tensor(image)  # Convert back to tensor\n",
    "\n",
    "# # Visualize the input image\n",
    "# plt.figure(figsize=(12, 6))\n",
    "# plt.subplot(1, 2, 1)\n",
    "# plt.title(\"Original Binary Image\")\n",
    "# plt.imshow(image.numpy().squeeze(), cmap='gray')\n",
    "# plt.colorbar()\n",
    "\n",
    "# # Compute connected components\n",
    "# num_iterations = 100\n",
    "# labeled_image = connected_components_tf(image, num_iterations)\n",
    "\n",
    "# # Visualize the labeled image\n",
    "# plt.subplot(1, 2, 2)\n",
    "# plt.title(\"Labeled Image\")\n",
    "# plt.imshow(labeled_image.numpy().squeeze(), cmap='jet')\n",
    "# plt.colorbar()\n",
    "# plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "connected components II with example usage"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import tensorflow as tf\n",
    "# import numpy as np\n",
    "# import matplotlib.pyplot as plt\n",
    "\n",
    "# # Function definition (from the previous adaptation)\n",
    "# def connected_components_tf(image, num_iterations=100):\n",
    "#     if not tf.is_tensor(image):\n",
    "#         raise TypeError(f\"Input type is not a TensorFlow Tensor. Got: {type(image)}\")\n",
    "#     if not isinstance(num_iterations, int) or num_iterations < 1:\n",
    "#         raise TypeError(\"Input num_iterations must be a positive integer.\")\n",
    "#     if len(image.shape) < 3 or image.shape[-1] != 1:\n",
    "#         raise ValueError(f\"Input image shape must be (*, H, W, 1). Got: {image.shape}\")\n",
    "    \n",
    "#     original_shape = tf.shape(image)\n",
    "#     flat_image = tf.reshape(image, [-1, original_shape[-3], original_shape[-2], 1])\n",
    "    \n",
    "#     B, H, W, _ = tf.unstack(tf.shape(flat_image))\n",
    "#     labels = tf.range(B * H * W, dtype=tf.float32)\n",
    "#     labels = tf.reshape(labels, [B, H, W, 1])\n",
    "#     mask = tf.cast(flat_image == 1, tf.float32)\n",
    "    \n",
    "#     for _ in range(num_iterations):\n",
    "#         labels = mask * tf.nn.max_pool2d(labels, ksize=3, strides=1, padding='SAME') + (1 - mask) * labels\n",
    "\n",
    "#     return tf.reshape(labels, original_shape)\n",
    "\n",
    "# # Create a simple binary image\n",
    "# image_data = np.array([\n",
    "#     [0, 0, 1, 1, 0],\n",
    "#     [0, 0, 1, 1, 0],\n",
    "#     [1, 1, 0, 0, 0],\n",
    "#     [1, 1, 0, 0, 0],\n",
    "#     [0, 0, 0, 0, 1]\n",
    "# ], dtype=np.float32)\n",
    "\n",
    "# image_tensor = tf.constant(image_data)[tf.newaxis, ..., tf.newaxis]  # Shape: (1, 5, 5, 1)\n",
    "\n",
    "# # Apply the connected components function\n",
    "# labels = connected_components_tf(image_tensor, num_iterations=10)\n",
    "\n",
    "# # Display the result\n",
    "# plt.figure(figsize=(12, 6))\n",
    "# plt.subplot(1, 2, 1)\n",
    "# plt.title(\"Original Image\")\n",
    "# plt.imshow(image_data, cmap='gray')\n",
    "# plt.colorbar()\n",
    "# plt.subplot(1, 2, 2)\n",
    "# plt.title(\"Labeled Image\")\n",
    "# plt.imshow(labels[0, ..., 0], cmap='nipy_spectral')\n",
    "# plt.colorbar()\n",
    "# plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import tensorflow as tf\n",
    "# import torch\n",
    "# def stitch_windows_tf(windows, k, cropx, cropy):\n",
    "#     # Ensure windows is a TensorFlow tensor\n",
    "#     windows = tf.convert_to_tensor(windows)\n",
    "\n",
    "#     # Stitch the first row of windows\n",
    "#     row0 = tf.concat([windows[0, 0][:-k, :-k]] +\n",
    "#                      [win[:-k, k:-k] for win in windows[0, 1:-1]] +\n",
    "#                      [windows[0, -1][:-k, k:]], axis=1)\n",
    "\n",
    "#     rows = []\n",
    "\n",
    "#     # Stitch the middle rows of windows\n",
    "#     for r in range(1, windows.shape[0] - 1):\n",
    "#         row_r = tf.concat([windows[r, 0][k:-k, :-k]] +\n",
    "#                           [win[k:-k, k:-k] for win in windows[r, 1:-1]] +\n",
    "#                           [windows[r, -1][k:-k, k:]], axis=1)\n",
    "#         rows.append(row_r)\n",
    "\n",
    "#     # Stitch the last row of windows\n",
    "#     row_last = tf.concat([windows[-1, 0][k:, :-k]] +\n",
    "#                          [win[k:, k:-k] for win in windows[-1, 1:-1]] +\n",
    "#                          [windows[-1, -1][k:, k:]], axis=1)\n",
    "\n",
    "#     # Combine all rows\n",
    "#     final = tf.concat([row0] + rows + [row_last], axis=0)\n",
    "    \n",
    "#     # Crop to the original size if necessary\n",
    "#     final = final[:cropx, :cropy]\n",
    "    \n",
    "#     return final\n",
    "\n",
    "\n",
    "# def connected_components(image: torch.Tensor, num_iterations: int = 100) -> torch.Tensor:\n",
    "#     r\"\"\"Computes the Connected-component labelling (CCL) algorithm.\n",
    "\n",
    "#     .. image:: https://github.com/kornia/data/raw/main/cells_segmented.png\n",
    "\n",
    "#     The implementation is an adaptation of the following repository:\n",
    "\n",
    "#     https://gist.github.com/efirdc/5d8bd66859e574c683a504a4690ae8bc\n",
    "\n",
    "#     .. warning::\n",
    "#         This is an experimental API subject to changes and optimization improvements.\n",
    "\n",
    "#     .. note::\n",
    "#        See a working example `here <https://kornia-tutorials.readthedocs.io/en/latest/\n",
    "#        connected_components.html>`__.\n",
    "\n",
    "#     Args:\n",
    "#         image: the binarized input image with shape :math:`(*, 1, H, W)`.\n",
    "#           The image must be in floating point with range [0, 1].\n",
    "#         num_iterations: the number of iterations to make the algorithm to converge.\n",
    "\n",
    "#     Return:\n",
    "#         The labels image with the same shape of the input image.\n",
    "\n",
    "#     Example:\n",
    "#         >>> img = torch.rand(2, 1, 4, 5)\n",
    "#         >>> img_labels = connected_components(img, num_iterations=100)\n",
    "#     \"\"\"\n",
    "#     if not isinstance(image, torch.Tensor):\n",
    "#         raise TypeError(f\"Input imagetype is not a torch.Tensor. Got: {type(image)}\")\n",
    "\n",
    "#     if not isinstance(num_iterations, int) or num_iterations < 1:\n",
    "#         raise TypeError(\"Input num_iterations must be a positive integer.\")\n",
    "\n",
    "#     if len(image.shape) < 3 or image.shape[-3] != 1:\n",
    "#         raise ValueError(f\"Input image shape must be (*,1,H,W). Got: {image.shape}\")\n",
    "\n",
    "#     H, W = image.shape[-2:]\n",
    "#     image_view = image.view(-1, 1, H, W)\n",
    "\n",
    "#     # precompute a mask with the valid values\n",
    "#     mask = image_view == 1\n",
    "\n",
    "#     # allocate the output tensors for labels\n",
    "#     B, _, _, _ = image_view.shape\n",
    "#     out = torch.arange(B * H * W, device=image.device, dtype=image.dtype).view((-1, 1, H, W))\n",
    "#     out[~mask] = 0\n",
    "\n",
    "#     for _ in range(num_iterations):\n",
    "#         out[mask] = F.max_pool2d(out, kernel_size=3, stride=1, padding=1)[mask]\n",
    "\n",
    "#     return out.view_as(image)\n",
    "\n",
    "# class LocatorTF:\n",
    "#     def __init__(self, fastrcnn_model, process_stride=64, method='max', dark_threshold=20, locating_model=None,\n",
    "#                  mode='static', **kwargs):\n",
    "#         self.fastrcnn_model = fastrcnn_model\n",
    "#         self.mode = mode\n",
    "#         self.process_stride = process_stride\n",
    "#         self.method = method\n",
    "#         self.locating_model = locating_model\n",
    "#         self.dark_threshold = dark_threshold\n",
    "#         self.p_list = kwargs.get('p_list', [8, 6, 1.5, 1, 50])\n",
    "#         self.meanADU = kwargs.get('meanADU', 241.0)\n",
    "#         self.dynamic_thres = kwargs.get('dynamic_thres', True)\n",
    "#         self.pretune_thresholding = kwargs.get('pretune_thresholding')\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "#     def model_tune(self, arr):\n",
    "#         meanADU = self.meanADU * 4  # mean ADU * upsample_factor^2\n",
    "#         offset = 0\n",
    "#         limit = int(tf.reduce_sum(arr) / meanADU + offset)\n",
    "#         arr_t = tf.cast(arr[None, None, ...] > 30, tf.float32)\n",
    "        \n",
    "#         # Assuming a custom implementation for connected components that returns a count\n",
    "#         # since TensorFlow doesn't have a direct equivalent.\n",
    "#         limit_cca = connected_components(arr_t, num_iterations=10)\n",
    "#         limit = max(limit_cca, limit)\n",
    "#         limit = max(limit, 1)\n",
    "\n",
    "#         self.fastrcnn_model.rpn._pre_nms_top_n = {'training': limit * self.p_list[0], 'testing': limit * self.p_list[0]}\n",
    "#         self.fastrcnn_model.rpn._post_nms_top_n = {'training': limit * self.p_list[1],\n",
    "#                                                    'testing': limit * self.p_list[1]}\n",
    "#         self.fastrcnn_model.roi_heads.detections_per_img = int(limit * self.p_list[2])\n",
    "#         # self.fastrcnn_model.roi_heads.score_thresh = self.p_list[3] / limit if limit < self.p_list[4] else 0\n",
    "#         self.fastrcnn_model.roi_heads.score_thresh = self.p_list[3] / limit\n",
    "\n",
    "#         self.fastrcnn_model.roi_heads.nms_thresh = 0.02  # smaller, delete more detections\n",
    "\n",
    "#         if limit > (0.005 * arr.shape[0] * arr.shape[1]) and self.dynamic_thres:  # 0.002 is minimum for model13\n",
    "#             self.dark_threshold = 0  # for image that not quite sparse, lift the pre-thresholding.\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "#     def images_to_window_lists(self, inputs):\n",
    "#         outputs = []\n",
    "#         maxs = []\n",
    "#         mins = []\n",
    "#         h, w = inputs.shape[1], inputs.shape[2]\n",
    "\n",
    "#         # TensorFlow's way of defining the upsample layer\n",
    "#         upsample = tf.keras.layers.UpSampling2D(size=(2, 2), interpolation='nearest')\n",
    "\n",
    "#         if self.process_stride is None:\n",
    "#             for image in inputs:\n",
    "#                 image = tf.expand_dims(image, 0)  # Add batch dimension\n",
    "#                 windows = upsample(image)\n",
    "#                 outputs.extend(tf.reshape(windows, [-1, *windows.shape[2:]]))\n",
    "#                 maxs.extend([tf.reduce_max(image)] * (windows.shape[1] * windows.shape[2]))\n",
    "#                 mins.extend([tf.reduce_min(image)] * (windows.shape[1] * windows.shape[2]))\n",
    "#         else:\n",
    "#             # Asserts in TensorFlow work differently, and dynamic checks like this are less common,\n",
    "#             # but you can still perform a check using tf.debugging.assert_positive, for example.\n",
    "            \n",
    "#             for image in inputs:\n",
    "#                 # Pad the image if necessary\n",
    "#                 pad_height = ((h + self.process_stride - 1) // self.process_stride) * self.process_stride\n",
    "#                 pad_width = ((w + self.process_stride - 1) // self.process_stride) * self.process_stride\n",
    "\n",
    "#                 padding_height = pad_height - h\n",
    "#                 padding_width = pad_width - w\n",
    "\n",
    "#                 # The padding has the shape [n, 2], where 'n' is the rank of 'inputs'\n",
    "#                 # And we add padding to the height and width dimensions (1 and 2).\n",
    "#                 paddings = tf.constant([[0, 0], [padding_height // 2, padding_height - padding_height // 2],\n",
    "#                                         [padding_width // 2, padding_width - padding_width // 2], [0, 0]])\n",
    "\n",
    "#                 image_padded = tf.pad(image, paddings, mode='REFLECT')\n",
    "\n",
    "#                 # Splitting and upsampling the windows\n",
    "#                 # TensorFlow doesn't have an unfold function, but we can achieve something similar with tf.image.extract_patches\n",
    "#                 sizes = [1, self.process_stride, self.process_stride, 1]\n",
    "#                 strides = [1, self.process_stride - 6, self.process_stride - 6, 1]\n",
    "#                 rates = [1, 1, 1, 1]  # For dilation, not used here\n",
    "#                 windows = tf.image.extract_patches(images=tf.expand_dims(image_padded, 0),\n",
    "#                                                 sizes=sizes,\n",
    "#                                                 strides=strides,\n",
    "#                                                 rates=rates,\n",
    "#                                                 padding='VALID')\n",
    "#                 windows = tf.reshape(windows, [-1, self.process_stride, self.process_stride, inputs.shape[-1]])\n",
    "#                 windows = upsample(windows)\n",
    "#                 outputs.extend(tf.reshape(windows, [-1, *windows.shape[2:]]))\n",
    "#                 maxs.extend([tf.reduce_max(image)] * (windows.shape[0] * windows.shape[1]))\n",
    "#                 mins.extend([tf.reduce_min(image)] * (windows.shape[0] * windows.shape[1]))\n",
    "\n",
    "#         # Note: Returning `windows.shape` directly after the loop doesn't make sense in this context,\n",
    "#         # because `windows` would only refer to the last processed batch. You might want to collect\n",
    "#         # these shapes in a list if needed for each image/window.\n",
    "#         return outputs, maxs, mins\n",
    "\n",
    "#     # Assuming `images_to_window_lists` and `model_tune` are adapted to TensorFlow\n",
    "#     def predict_sequence(self, inputs):\n",
    "#         counted_list = []\n",
    "#         eventsize_all = []\n",
    "#         inputs = tf.cast(inputs, tf.float32)\n",
    "#         counted_images = tf.zeros_like(inputs)\n",
    "\n",
    "#         image_cell_list, windowshape, maxs, mins = self.images_to_window_lists(inputs)\n",
    "#         for i, image_cell in enumerate(image_cell_list):\n",
    "\n",
    "#             if self.mode == 'dynamic_window':\n",
    "#                 self.model_tune(image_cell)\n",
    "#             elif self.mode == 'dynamic_frame':\n",
    "#                 # Similar logic adapted for TensorFlow\n",
    "#                 pass  # Fill in based on adapted `model_tune`\n",
    "#             elif self.mode == 'static':\n",
    "#                 # Ensure process_stride compatibility\n",
    "#                 pass\n",
    "#             else:\n",
    "#                 raise ValueError(\"Use mode = 'dynamic_window', 'dynamic_frame', or 'static'.\")\n",
    "\n",
    "#             # Image cell thresholding and normalization\n",
    "#             image_cell = tf.where(image_cell < self.dark_threshold, 0, image_cell)\n",
    "#             image_cell_normalized = (image_cell - mins[i]) / (maxs[i] - mins[i])\n",
    "\n",
    "#             # Model inference (adapt based on your TensorFlow model)\n",
    "#             boxes = self.fastrcnn_model(image_cell_normalized[None, ...])[0]['boxes']\n",
    "\n",
    "#             select = []\n",
    "#             for row, value in enumerate(boxes):\n",
    "#                 if 0 < (value[2] - value[0]) < 30 and 0 < (value[3] - value[1]) < 30:\n",
    "#                     select.append(row)\n",
    "#             select = torch.as_tensor(select, dtype=torch.int, device=self.device)\n",
    "#             filtered_boxes = torch.index_select(boxes, 0, select)\n",
    "#             filtered_boxes = filtered_boxes / 2.0\n",
    "#             image_cell_ori = F.interpolate(image_cell_ori[None, None, ...], scale_factor=0.5, mode='nearest')[0, 0]\n",
    "#             filtered, _, eventsize = self.locate(image_cell_ori, filtered_boxes)\n",
    "#             counted_list.append(filtered[None, ...])  # [1,w,h]\n",
    "#             eventsize_all = eventsize_all + eventsize\n",
    "\n",
    "#         image_num = int(len(counted_list) / windowshape[0] / windowshape[1])\n",
    "#         for index in range(image_num):\n",
    "#             counted_cells = counted_list[\n",
    "#                             index * windowshape[0] * windowshape[1]:(index + 1) * windowshape[0] * windowshape[1]]\n",
    "#             counted_cells = torch.cat(counted_cells)\n",
    "#             counted_cells = counted_cells.reshape(windowshape[0], windowshape[1], int(windowshape[2] / 2),\n",
    "#                                                   int(windowshape[3] / 2))\n",
    "#             counted_images[index] = stitch_windows_tf(counted_cells, k=3, cropx=inputs.shape[1], cropy=inputs.shape[2])\n",
    "\n",
    "#         return counted_images, eventsize_all\n",
    "\n",
    "        \n",
    "    \n",
    "\n",
    "#     def locate(self, image_array, boxes):\n",
    "#         width = 10\n",
    "#         filtered = tf.zeros_like(image_array)\n",
    "#         boxes = tf.cast(tf.round(boxes), tf.int32)\n",
    "#         coor = []\n",
    "#         eventsize = []\n",
    "\n",
    "#         for box in boxes:\n",
    "#             xarea = image_array[box[1]:(box[3] + 1), box[0]:(box[2] + 1)]\n",
    "\n",
    "#             # Padding logic adapted for TensorFlow\n",
    "#             if xarea.shape[0] > (width + 1) or xarea.shape[1] > (width + 1):\n",
    "#                 paddings = tf.constant([[1, width], [1, width]])\n",
    "#                 patch = tf.pad(xarea, paddings, mode='CONSTANT')[:width + 2, :width + 2]\n",
    "#             else:\n",
    "#                 paddings = tf.constant([[1, width - xarea.shape[1] + 1], [1, width - xarea.shape[0] + 1]])\n",
    "#                 patch = tf.pad(xarea, paddings, mode='CONSTANT')\n",
    "\n",
    "#             if self.method == 'max':\n",
    "#                 model_x, model_y = tf.unravel_index(tf.argmax(patch, axis=None), patch.shape)\n",
    "#             elif self.method == 'binary_com':\n",
    "#                 patch = tf.where(patch < 30, 0, 1)\n",
    "#                 x = tf.range(0, patch.shape[0], dtype=tf.float32)\n",
    "#                 y = tf.range(0, patch.shape[1], dtype=tf.float32)\n",
    "#                 weights_x, weights_y = tf.meshgrid(x, y, indexing='ij')\n",
    "#                 model_x = tf.math.round(tf.reduce_sum(patch * weights_x) / tf.reduce_sum(patch))\n",
    "#                 model_y = tf.math.round(tf.reduce_sum(patch * weights_y) / tf.reduce_sum(patch))\n",
    "#             elif self.method == 'com':\n",
    "#                 x = tf.range(0, patch.shape[0], dtype=tf.float32)\n",
    "#                 y = tf.range(0, patch.shape[1], dtype=tf.float32)\n",
    "#                 weights_x, weights_y = tf.meshgrid(x, y, indexing='ij')\n",
    "#                 model_x = tf.math.round(tf.reduce_sum(patch * weights_x) / tf.reduce_sum(patch))\n",
    "#                 model_y = tf.math.round(tf.reduce_sum(patch * weights_y) / tf.reduce_sum(patch))\n",
    "#             else:\n",
    "#                 raise ValueError(\"Use 'max', 'com', or 'binary_com' to locate the entry position.\")\n",
    "\n",
    "#             cx = model_x + box[1] - 1\n",
    "#             cy = model_y + box[0] - 1\n",
    "\n",
    "#             if cx >= image_array.shape[0] or cy >= image_array.shape[1] or cx < 0 or cy < 0:\n",
    "#                 continue\n",
    "\n",
    "#             coor.append([cx, cy])\n",
    "#             eventsize.append(tf.reduce_sum(tf.cast(patch > 20, tf.int32)).numpy())\n",
    "\n",
    "#         coords = tf.cast(coor, tf.int32)\n",
    "#         for point in coords:\n",
    "#             filtered = tf.tensor_scatter_nd_add(filtered, [point], [1])\n",
    "\n",
    "#         return filtered, coords, eventsize"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "def stitch_windows_tf(windows, k, cropx, cropy):\n",
    "    # Example of adapting the stich_windows for TensorFlow\n",
    "    # Assumes 'windows' is a TensorFlow tensor\n",
    "    row0 = tf.concat([windows[0, 0][:-k, :-k]] +\n",
    "                     [win[:-k, k:-k] for win in windows[0, 1:-1]] +\n",
    "                     [windows[0, -1][:-k, k:]], axis=1)\n",
    "    rows = [row0]\n",
    "    for r in range(1, windows.shape[0] - 1):\n",
    "        rows.append(tf.concat([windows[r, 0][k:-k, :-k]] +\n",
    "                              [win[k:-k, k:-k] for win in windows[r, 1:-1]] +\n",
    "                              [windows[r, -1][k:-k, k:]], axis=1))\n",
    "    row_last = tf.concat([windows[-1, 0][k:, :-k]] +\n",
    "                         [win[k:, k:-k] for win in windows[-1, 1:-1]] +\n",
    "                         [windows[-1, -1][k:, k:]], axis=1)\n",
    "    final = tf.concat([row0] + rows[1:] + [row_last], axis=0)\n",
    "    final = final[:cropx, :cropy]\n",
    "    return final\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Stitched Image Shape: (6, 6)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-04-25 13:29:47.198929: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1929] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 13213 MB memory:  -> device: 0, name: NVIDIA GeForce RTX 3090, pci bus id: 0000:68:00.0, compute capability: 8.6\n"
     ]
    }
   ],
   "source": [
    "\n",
    "\n",
    "# Simulated image patches (overlapping windows)\n",
    "windows = tf.random.uniform((2, 3, 4, 4), minval=0, maxval=255, dtype=tf.float32)\n",
    "k = 1  # Overlap size\n",
    "cropx, cropy = 6, 6  # Dimensions of the final cropped image\n",
    "\n",
    "# Stitch the windows together\n",
    "stitched_image = stitch_windows_tf(windows, k, cropx, cropy)\n",
    "print(\"Stitched Image Shape:\", stitched_image.shape)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import tensorflow as tf\n",
    "# from tensorflow.keras import layers, models\n",
    "\n",
    "# class CustomModel(tf.keras.Model):\n",
    "#     def __init__(self):\n",
    "#         super(CustomModel, self).__init__()\n",
    "#         self.conv1 = layers.Conv2D(64, 3, padding='same', activation='relu')\n",
    "#         self.flatten = layers.Flatten()\n",
    "#         self.d1 = layers.Dense(10, activation='sigmoid')  # Dummy classification layer\n",
    "\n",
    "#     def call(self, x):\n",
    "#         x = self.conv1(x)\n",
    "#         x = self.flatten(x)\n",
    "#         return self.d1(x)\n",
    "\n",
    "#     def update_thresholds(self, thresholds):\n",
    "#         # This method would normally update operational parameters\n",
    "#         print(\"Updated model thresholds:\")\n",
    "#         for k, v in thresholds.items():\n",
    "#             print(f\"{k}: {v}\")\n",
    "#         # Store or use the thresholds as needed for operations\n",
    "\n",
    "# # Instantiate the custom model\n",
    "# model = CustomModel()\n",
    "\n",
    "\n",
    "# class LocatorTF:\n",
    "#     def __init__(self, model, **kwargs):\n",
    "#         self.model = model\n",
    "#         self.meanADU = kwargs.get('meanADU', 241.0)\n",
    "\n",
    "#     def model_tune(self, arr):\n",
    "#         meanADU = self.meanADU * 4\n",
    "#         limit = int(tf.reduce_sum(arr) / meanADU)\n",
    "#         # Simplified logic for the demo\n",
    "#         limit = max(limit, 1)\n",
    "#         thresholds = {\n",
    "#             'score_thresh': 0.5 / limit,  # Simplified calculation\n",
    "#             'nms_thresh': 0.02,\n",
    "#             'detections_per_img': limit * 10  # Simplified logic\n",
    "#         }\n",
    "#         self.model.update_thresholds(thresholds)\n",
    "\n",
    "# # Create an instance of LocatorTF\n",
    "# locator = LocatorTF(model=model)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Create a test input (e.g., a single 64x64 grayscale image)\n",
    "# test_input = tf.random.normal([1, 64, 64, 1])\n",
    "\n",
    "# # Call model_tune with the test input\n",
    "# locator.model_tune(test_input)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "\n",
    "# import kornia\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "\n",
    "\n",
    "def map01(mat):\n",
    "    return (mat - mat.min()) / (mat.max() - mat.min())\n",
    "\n",
    "\n",
    "def unravel_index(index, shape):\n",
    "    out = []\n",
    "    for dim in reversed(shape):\n",
    "        out.append(index % dim)\n",
    "        index = torch.div(index, dim, rounding_mode='floor')\n",
    "\n",
    "    return tuple(reversed(out))\n",
    "\n",
    "\n",
    "def stich_windows(windows, k, cropx, cropy):\n",
    "    if not torch.is_tensor(windows):\n",
    "        windows = torch.as_tensor(windows)\n",
    "    row0 = torch.cat([windows[0, 0][:-k, :-k]] +\n",
    "                     [win[:-k, k:-k] for win in windows[0, 1:-1]] +\n",
    "                     [windows[0, -1][:-k, k:]], dim=1)\n",
    "    rows = []\n",
    "\n",
    "    for r in range(windows.shape[0] - 1):\n",
    "        r = r + 1\n",
    "        rows.append(torch.cat([windows[r, 0][k:-k, :-k]] +\n",
    "                              [win[k:-k, k:-k] for win in windows[r, 1:-1]] +\n",
    "                              [windows[r, -1][k:-k, k:]], dim=1)\n",
    "                    )\n",
    "\n",
    "    row_last = torch.cat([windows[-1, 0][k:, :-k]] +\n",
    "                         [win[k:, k:-k] for win in windows[-1, 1:-1]] +\n",
    "                         [windows[-1, -1][k:, k:]], dim=1)\n",
    "\n",
    "    final = torch.cat([row0] + rows + [row_last], dim=0)\n",
    "    final = final[:cropx, :cropy]\n",
    "    return final\n",
    "\n",
    "\n",
    "def connected_components(image: torch.Tensor, num_iterations: int = 100) -> torch.Tensor:\n",
    "    r\"\"\"Computes the Connected-component labelling (CCL) algorithm.\n",
    "\n",
    "    .. image:: https://github.com/kornia/data/raw/main/cells_segmented.png\n",
    "    The implementation is an adaptation of the following repository:\n",
    "\n",
    "    https://gist.github.com/efirdc/5d8bd66859e574c683a504a4690ae8bc\n",
    "\n",
    "    .. warning::\n",
    "        This is an experimental API subject to changes and optimization improvements.\n",
    "\n",
    "    .. note::\n",
    "       See a working example `here <https://kornia-tutorials.readthedocs.io/en/latest/\n",
    "       connected_components.html>`__.\n",
    "\n",
    "    Args:\n",
    "        image: the binarized input image with shape :math:`(*, 1, H, W)`.\n",
    "          The image must be in floating point with range [0, 1].\n",
    "        num_iterations: the number of iterations to make the algorithm to converge.\n",
    "\n",
    "    Return:\n",
    "        The labels image with the same shape of the input image.\n",
    "\n",
    "    Example:\n",
    "        >>> img = torch.rand(2, 1, 4, 5)\n",
    "        >>> img_labels = connected_components(img, num_iterations=100)\n",
    "    \"\"\"\n",
    "    if not isinstance(image, torch.Tensor):\n",
    "        raise TypeError(f\"Input imagetype is not a torch.Tensor. Got: {type(image)}\")\n",
    "\n",
    "    if not isinstance(num_iterations, int) or num_iterations < 1:\n",
    "        raise TypeError(\"Input num_iterations must be a positive integer.\")\n",
    "\n",
    "    if len(image.shape) < 3 or image.shape[-3] != 1:\n",
    "        raise ValueError(f\"Input image shape must be (*,1,H,W). Got: {image.shape}\")\n",
    "\n",
    "    H, W = image.shape[-2:]\n",
    "    image_view = image.view(-1, 1, H, W)\n",
    "\n",
    "    # precompute a mask with the valid values\n",
    "    mask = image_view == 1\n",
    "\n",
    "    # allocate the output tensors for labels\n",
    "    B, _, _, _ = image_view.shape\n",
    "    out = torch.arange(B * H * W, device=image.device, dtype=image.dtype).view((-1, 1, H, W))\n",
    "    out[~mask] = 0\n",
    "\n",
    "    for _ in range(num_iterations):\n",
    "        out[mask] = F.max_pool2d(out, kernel_size=3, stride=1, padding=1)[mask]\n",
    "\n",
    "    return out.view_as(image)\n",
    "\n",
    "\n",
    "class Locator:\n",
    "    \"\"\"\n",
    "    Implements Faster R-CNN on a single image to detect boxes for electron events,\n",
    "    then use finding maximum or pre-trained FCN to assign the entry positions\n",
    "    The grid_predict and predict function works for image stack, but the locate is set to work for single image.\n",
    "    Returns boxes as xmin, ymin, xmax,  ymax, x means horizontal and y means vertical.\n",
    "\n",
    "    Args:\n",
    "        fastrcnn_model: the loaded fast rcnn model\n",
    "        device: torch.device('cpu') or torch.device('cuda')\n",
    "        process_stride: divide the image into pieces when applying the fast rcnn, default 64 for static mode.\n",
    "        None means no spliting and work on whole frames. Maximum stride 128 for mscdata gpu.\n",
    "        method: 'max' or 'fcn'\n",
    "        dark_threshold: the intensity threshold for remove dark noise for image patches with density < 0.01.\n",
    "        locating_model: the loaded fcn model for assigning entry position\n",
    "        mode: dynamic mode, whether apply model tune for images with different electron density.\n",
    "        'static': static parameters, same threshold, no model tune\n",
    "        'dynamic_window': apply model tune separately for each window\n",
    "        'dynamic_frame': apply model tune equally the whole frame based on whole frame intensity.\n",
    "        p_list: optional list of five multiplier for model tune, if none, will use default numbers: [6, 6, 1.3, 1.5, 23]\n",
    "        meanADU: optional float for mean intensity per electron (ADU), if none, will use default 241 for 200kV.\n",
    "        dynamic_thres: optional bool for wheather lift the threshold above some density within modeltune.\n",
    "        pretune_thresholding: optional threshold for reset noise pixels before model tune.\n",
    "    Example::\n",
    "\n",
    "        >>>from CountingNN.locator import Locator\n",
    "        >>>counting = Locator(model_object, device, process_stride, method,\n",
    "        >>>dark_threshold, locating_model, dynamic_param, p_list = p_list, meanADU=meanADU)\n",
    "        >>>filtered = counting.predict_sequence(inputs)  # inputs as the image tensor in shape [N,h,w]\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, fastrcnn_model, device, process_stride=64, method='max', dark_threshold=20, locating_model=None,\n",
    "                 mode='static', **kwargs):\n",
    "        super().__init__()\n",
    "        self.fastrcnn_model = fastrcnn_model\n",
    "        self.device = device\n",
    "        self.mode = mode\n",
    "        self.process_stride = process_stride\n",
    "        self.method = method\n",
    "        self.locating_model = locating_model\n",
    "        self.dark_threshold = dark_threshold\n",
    "        self.p_list = kwargs.get('p_list')\n",
    "        if self.p_list is None:\n",
    "            self.p_list = [8, 6, 1.5, 1, 50]\n",
    "        self.meanADU = kwargs.get('meanADU')\n",
    "        if self.meanADU is None:\n",
    "            self.meanADU = 241.0\n",
    "        self.dynamic_thres = kwargs.get('dynamic_thres')\n",
    "        self.pretune_thresholding = kwargs.get('pretune_thresholding')\n",
    "        if self.dynamic_thres is None:\n",
    "            self.dynamic_thres = True\n",
    "        self.fastrcnn_model = self.fastrcnn_model.to(self.device)\n",
    "\n",
    "    def model_tune(self, arr):\n",
    "        \"\"\"\n",
    "        Change the detection limits and thresholds of Fast R-CNN model by estimating the image sparsity\n",
    "        \"\"\"\n",
    "        meanADU = self.meanADU * 4  # mean ADU * upsample_factor^2\n",
    "        offset = 0\n",
    "        # fit from 200kV Validation data, between a 64x64\n",
    "        # up-sampled-by-2 image cell ans its original ground truth.\n",
    "        limit = int(arr.sum() / meanADU + offset)\n",
    "        arr_t = torch.as_tensor(arr[None, None, ...] > 30, dtype=torch.float32)\n",
    "        limit_cca = connected_components(arr_t, num_iterations=10)\n",
    "        limit = max(torch.unique(limit_cca).shape[0], limit)\n",
    "        limit = max(limit, 1)\n",
    "\n",
    "        # if limit > (0.035 * 2 * arr.shape[0] * arr.shape[1]):  # no pre-tune thresholding if density beyond 3%\n",
    "        #     self.pretune_thresholding = None\n",
    "        #\n",
    "        # if self.pretune_thresholding is not None:  # recalculate the limit after thresholding\n",
    "        #     arr[arr < self.pretune_thresholding] = 0\n",
    "        #     limit = int(arr.sum() / meanADU + offset)\n",
    "        #     limit = max(torch.unique(limit_cca).shape[0], limit)\n",
    "        #     limit = max(limit, 1)\n",
    "\n",
    "        self.fastrcnn_model.rpn._pre_nms_top_n = {'training': limit * self.p_list[0], 'testing': limit * self.p_list[0]}\n",
    "        self.fastrcnn_model.rpn._post_nms_top_n = {'training': limit * self.p_list[1],\n",
    "                                                   'testing': limit * self.p_list[1]}\n",
    "        self.fastrcnn_model.roi_heads.detections_per_img = int(limit * self.p_list[2])\n",
    "        # self.fastrcnn_model.roi_heads.score_thresh = self.p_list[3] / limit if limit < self.p_list[4] else 0\n",
    "        self.fastrcnn_model.roi_heads.score_thresh = self.p_list[3] / limit\n",
    "\n",
    "        self.fastrcnn_model.roi_heads.nms_thresh = 0.02  # smaller, delete more detections\n",
    "\n",
    "        if limit > (0.005 * arr.shape[0] * arr.shape[1]) and self.dynamic_thres:  # 0.002 is minimum for model13\n",
    "            self.dark_threshold = 0  # for image that not quite sparse, lift the pre-thresholding.\n",
    "\n",
    "    def images_to_window_lists(self, inputs: torch.tensor) -> List[torch.tensor]:\n",
    "        \"\"\"\n",
    "        transform a batch of images(3D tensor) into windows of the images, with up-sampling by 2.\n",
    "        if stride = None, not spliting and return whole images.\n",
    "        \"\"\"\n",
    "\n",
    "        inputs = inputs.to(self.device)\n",
    "        outputs = []\n",
    "        maxs = []\n",
    "        mins = []\n",
    "        h, w = inputs.shape[1:]\n",
    "\n",
    "        if self.process_stride is None:\n",
    "            for image in inputs:\n",
    "                windows = image[None, None, ...]\n",
    "                windows = torch.nn.Upsample(scale_factor=2, mode='nearest')(windows)\n",
    "                outputs = outputs + list(torch.flatten(windows, start_dim=0, end_dim=1))\n",
    "                maxs = maxs + [image.max()] * (windows.shape[0] * windows.shape[1])\n",
    "                mins = mins + [image.min()] * (windows.shape[0] * windows.shape[1])\n",
    "        else:\n",
    "            torch._assert((torch.as_tensor(inputs.shape[1:]) > self.process_stride).all(),\n",
    "                          f\"Your image dimension is {inputs.shape[1:]}, which is not larger than process stride, \"\n",
    "                          f\"please use process_stride<{min(inputs.shape[1:])}\"\n",
    "                          )\n",
    "            for image in inputs:\n",
    "                pad = [torch.div(image.shape[0], (self.process_stride - 6), rounding_mode='floor') * (\n",
    "                        self.process_stride - 6) + self.process_stride,\n",
    "                       torch.div(image.shape[1], (self.process_stride - 6), rounding_mode='floor') * (\n",
    "                               self.process_stride - 6) + self.process_stride]  # int works as floor for positive number\n",
    "                image = F.pad(image,\n",
    "                              (0, pad[1] - image.shape[1], 0, pad[0] - image.shape[0]))  # left, right, top, bottom !!!\n",
    "\n",
    "                # the zero pad area make dim counting results due to dynamic modell tune, so fill with edge values\n",
    "                image[h:, :w] = image[(2 * h - pad[0]):h, :w]\n",
    "                image[:h, w:] = image[:h, (2 * w - pad[1]):w]\n",
    "                image[h:, w:] = image[(2 * h - pad[0]):h, w:]\n",
    "\n",
    "                windows = image.unfold(0, self.process_stride, self.process_stride - 6)\n",
    "                windows = windows.unfold(1, self.process_stride, self.process_stride - 6)\n",
    "                # up-sampling the windows\n",
    "                windows = torch.nn.Upsample(scale_factor=2, mode='nearest')(windows)\n",
    "                outputs = outputs + list(torch.flatten(windows, start_dim=0, end_dim=1))\n",
    "                maxs = maxs + [image.max()] * (windows.shape[0] * windows.shape[1])\n",
    "                mins = mins + [image.min()] * (windows.shape[0] * windows.shape[1])\n",
    "        return outputs, windows.shape, maxs, mins\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def predict_sequence(self, inputs: torch.tensor):\n",
    "        \"\"\"\n",
    "        apply model on image one patch after another. The patches size equals the image size in training data, so that\n",
    "        no need to tune the model detection limits for different limit sizes.\n",
    "        \"\"\"\n",
    "        self.fastrcnn_model.transform.crop_max = max(inputs.shape[1], inputs.shape[2]) * 2\n",
    "        # make size_divisible equals process_stride here to avoid inconsistent padding issue.\n",
    "        # self.fastrcnn_model.transform.size_divisible = self.process_stride * 2\n",
    "        self.fastrcnn_model.eval()\n",
    "\n",
    "        counted_list = []\n",
    "        eventsize_all = []\n",
    "        inputs = torch.as_tensor(inputs, dtype=torch.float32)\n",
    "        counted_images = torch.zeros_like(inputs)\n",
    "\n",
    "        image_cell_list, windowshape, maxs, mins = self.images_to_window_lists(inputs)\n",
    "        for i, image_cell in enumerate(image_cell_list):\n",
    "\n",
    "            if self.mode == 'dynamic_window':\n",
    "                self.model_tune(image_cell)\n",
    "            elif self.mode == 'dynamic_frame':  # incorrect\n",
    "                image_i = torch.div(i, windowshape[0] * windowshape[1], rounding_mode='floor')\n",
    "                self.model_tune(torch.nn.Upsample(scale_factor=2, mode='nearest')(inputs[image_i][None, None, ...]))\n",
    "            elif self.mode == 'static':\n",
    "                torch._assert(self.process_stride == 64,\n",
    "                              f\"please use process_stride=64 for static mode.\"\n",
    "                              )\n",
    "                pass\n",
    "            else:\n",
    "                raise ValueError(\"Use mode = 'dynamic_window', dynamic_frame or 'static'. \")\n",
    "\n",
    "            # thresholding to remove dark noise before applying the model\n",
    "            image_cell[image_cell < self.dark_threshold] = 0\n",
    "\n",
    "            image_cell_ori = image_cell\n",
    "\n",
    "            image_cell = (image_cell - mins[i]) / (maxs[i] - mins[i])  # norm the image cells equally\n",
    "            # boxes = self.fastrcnn_model([image_cell[None, ...]])[0]['boxes']\n",
    "            boxes = self.fastrcnn_model(image_cell[None, None, ...])[0]['boxes']  # model direct input [N, C, H, W]\n",
    "\n",
    "            select = []\n",
    "            for row, value in enumerate(boxes):\n",
    "                if 0 < (value[2] - value[0]) < 30 and 0 < (value[3] - value[1]) < 30:\n",
    "                    select.append(row)\n",
    "            select = torch.as_tensor(select, dtype=torch.int, device=self.device)\n",
    "            filtered_boxes = torch.index_select(boxes, 0, select)\n",
    "            filtered_boxes = filtered_boxes / 2.0\n",
    "            image_cell_ori = F.interpolate(image_cell_ori[None, None, ...], scale_factor=0.5, mode='nearest')[0, 0]\n",
    "            filtered, _, eventsize = self.locate(image_cell_ori, filtered_boxes)\n",
    "            counted_list.append(filtered[None, ...])  # [1,w,h]\n",
    "            eventsize_all = eventsize_all + eventsize\n",
    "\n",
    "        image_num = int(len(counted_list) / windowshape[0] / windowshape[1])\n",
    "        for index in range(image_num):\n",
    "            counted_cells = counted_list[\n",
    "                            index * windowshape[0] * windowshape[1]:(index + 1) * windowshape[0] * windowshape[1]]\n",
    "            counted_cells = torch.cat(counted_cells)\n",
    "            counted_cells = counted_cells.reshape(windowshape[0], windowshape[1], int(windowshape[2] / 2),\n",
    "                                                  int(windowshape[3] / 2))\n",
    "            counted_images[index] = stich_windows(counted_cells, k=3, cropx=inputs.shape[1], cropy=inputs.shape[2])\n",
    "\n",
    "        return counted_images, eventsize_all\n",
    "\n",
    "    def locate(self, image_array, boxes):\n",
    "\n",
    "        width = 10\n",
    "        filtered = torch.zeros_like(image_array)\n",
    "        boxes = boxes.round().int()\n",
    "        coor = []\n",
    "        eventsize = []\n",
    "\n",
    "        # if torch.cuda.is_available() and self.device == torch.device('cuda'):\n",
    "        #     boxes = boxes.cpu()\n",
    "\n",
    "        for box in boxes:\n",
    "            xarea = image_array[box[1]:(box[3] + 1), box[0]:(box[2] + 1)]\n",
    "\n",
    "            # # if the box is just 1~2 pxs, take the intensity value at four line edge instead of padding zero\n",
    "            # if (xarea.shape[0]+xarea.shape[1]) <= 3 and self.ext_small:\n",
    "            #     xarea_ext = image_array[box[1]-1:(box[3] + 2), box[0]-1:(box[2] + 2)]\n",
    "            #     patch = np.pad(xarea_ext, ((0, width - xarea_ext.shape[0] + 2), (0, width - xarea_ext.shape[1] + 2)))\n",
    "\n",
    "            # one more row and column added at four edges.\n",
    "            if xarea.shape[0] > (width + 1) or xarea.shape[1] > (width + 1):\n",
    "                patch = F.pad(xarea, (1, width, 1, width))\n",
    "                patch = patch[:(width + 2), :(width + 2)]\n",
    "            else:\n",
    "                patch = F.pad(xarea, (1, (width - xarea.shape[1] + 1), 1, (width - xarea.shape[0] + 1)))\n",
    "\n",
    "            if self.method == 'max':\n",
    "\n",
    "                (model_x, model_y) = unravel_index(torch.argmax(patch), shape=(width + 2, width + 2))\n",
    "            elif self.method == 'binary_com':\n",
    "                patch[patch < 30] = 0\n",
    "                patch[patch >= 30] = 1\n",
    "                x = torch.linspace(0, patch.shape[0] - 1, patch.shape[0])\n",
    "                y = torch.linspace(0, patch.shape[1] - 1, patch.shape[1])\n",
    "                weights_x, weights_y = torch.meshgrid(x, y)\n",
    "                model_x = (patch * weights_x).sum() / patch.sum()\n",
    "                model_y = (patch * weights_y).sum() / patch.sum()\n",
    "                model_x = int(torch.round(model_x))\n",
    "                model_y = int(torch.round(model_y))\n",
    "            elif self.method == 'com':\n",
    "                x = torch.linspace(0, patch.shape[0] - 1, patch.shape[0])\n",
    "                y = torch.linspace(0, patch.shape[1] - 1, patch.shape[1])\n",
    "                weights_x, weights_y = torch.meshgrid(x, y)\n",
    "                model_x = (patch * weights_x).sum() / patch.sum()\n",
    "                model_y = (patch * weights_y).sum() / patch.sum()\n",
    "                model_x = int(torch.round(model_x))\n",
    "                model_y = int(torch.round(model_y))\n",
    "            else:\n",
    "                raise ValueError(\"Use 'max','com,'binary_com' to locate the entry position. \")\n",
    "\n",
    "            cx = model_x + box[1] - 1\n",
    "            cy = model_y + box[0] - 1\n",
    "            if cx > (image_array.shape[0] - 1) or cy > (image_array.shape[1] - 1) or cx < 0 or cy < 0:\n",
    "                continue\n",
    "            coor.append((cx, cy))\n",
    "            eventsize.append((torch.sum(patch > 20)).item())\n",
    "\n",
    "        coords = torch.as_tensor(coor, dtype=torch.long).to(self.device)\n",
    "        # eventsize = torch.as_tensor(eventsize, dtype=torch.long).to(self.device)\n",
    "        if coords.shape[0]:\n",
    "            # filtered[coords[:, 0], coords[:, 1]] = 1 # this does not account for >1 e on single pixel.\n",
    "            for point in coords:\n",
    "                filtered[point[0], point[1]] = filtered[point[0], point[1]] + 1\n",
    "\n",
    "        return filtered, coords, eventsize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_shape = (64,64,1)\n",
    "num_classes = 280\n",
    "num_coordinates = 4\n",
    "\n",
    "x_input = layers.Input(shape=input_shape)\n",
    "#Layer 1\n",
    "x = layers.Conv2D(64, kernel_size=3, padding='same', activation='relu')(x_input)\n",
    "x = layers.MaxPool2D()(x)\n",
    "x = layers.BatchNormalization()(x) \n",
    "x = layers.Conv2D(64, kernel_size=3, padding='same', activation='relu')(x)\n",
    "#Layer 2\n",
    "x = layers.Conv2D(128, kernel_size=3, padding='same', activation='relu')(x)\n",
    "x = layers.Conv2D(128, kernel_size=3, padding='same', activation='relu')(x)\n",
    "#Layer 3\n",
    "x = layers.Conv2D(256, kernel_size=3, padding='same', activation='relu')(x)\n",
    "x = layers.Conv2D(256, kernel_size=3, padding='same', activation='relu')(x)\n",
    "#Layer 4\n",
    "x = layers.Conv2D(512, kernel_size=3, padding='same', activation='relu')(x)\n",
    "x = layers.MaxPool2D()(x)\n",
    "x = layers.Conv2D(512, kernel_size=3, padding='same', activation='relu')(x)\n",
    "x = layers.MaxPool2D()(x)\n",
    "x = layers.Conv2D(512, kernel_size=3, padding='same', activation='relu')(x)\n",
    "x = layers.MaxPool2D()(x)\n",
    "#Layer 5\n",
    "x = layers.Conv2D(256, kernel_size=5, padding='same', activation='relu')(x)\n",
    "x = layers.MaxPool2D()(x)\n",
    "x = layers.BatchNormalization()(x) \n",
    "\n",
    "\n",
    "x = layers.Flatten()(x)\n",
    "# Probability output\n",
    "x_prob = layers.Dense(num_classes, activation='sigmoid', name='x_prob')(x)\n",
    "x_prob_reshape = layers.Reshape((-1, num_classes, 1), name='x_prob_reshape')(x_prob)\n",
    "\n",
    "# Bounding box output\n",
    "x_boxes = layers.Dense(num_classes * num_coordinates, activation='sigmoid', name='x_boxes')(x)\n",
    "x_boxes_reshape = layers.Reshape((-1, num_classes, num_coordinates), name='x_boxes_reshape')(x_boxes)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "model = tf.keras.models.Model(x_input, [x_prob_reshape, x_boxes_reshape])\n",
    "optimizer = tf.keras.optimizers.Adam(learning_rate=3e-5) \n",
    "model.compile(optimizer= optimizer, loss= {'x_prob_reshape': tf.keras.losses.BinaryCrossentropy(), 'x_boxes_reshape':tf.keras.losses.MeanSquaredError()}, metrics=['accuracy'])    \n",
    "# num_epochs = 3000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "metadata": {},
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "Input image shape must be (*, 1, H, W). Got: (1, 1, 1, 64, 64, 1)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[104], line 48\u001b[0m\n\u001b[1;32m     46\u001b[0m locator \u001b[38;5;241m=\u001b[39m LocatorTF(model\u001b[38;5;241m=\u001b[39mmodel)\n\u001b[1;32m     47\u001b[0m sample_image_tensor \u001b[38;5;241m=\u001b[39m tf\u001b[38;5;241m.\u001b[39mrandom\u001b[38;5;241m.\u001b[39mnormal([\u001b[38;5;241m1\u001b[39m, \u001b[38;5;241m64\u001b[39m, \u001b[38;5;241m64\u001b[39m, \u001b[38;5;241m1\u001b[39m])  \u001b[38;5;66;03m# A sample tensor for demonstration\u001b[39;00m\n\u001b[0;32m---> 48\u001b[0m \u001b[43mlocator\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmodel_tune\u001b[49m\u001b[43m(\u001b[49m\u001b[43msample_image_tensor\u001b[49m\u001b[43m)\u001b[49m\n",
      "Cell \u001b[0;32mIn[104], line 19\u001b[0m, in \u001b[0;36mLocatorTF.model_tune\u001b[0;34m(self, arr)\u001b[0m\n\u001b[1;32m     17\u001b[0m \u001b[38;5;66;03m# Assuming a function connected_components_tf exists that returns tensor of components\u001b[39;00m\n\u001b[1;32m     18\u001b[0m arr_t \u001b[38;5;241m=\u001b[39m tf\u001b[38;5;241m.\u001b[39mcast(arr \u001b[38;5;241m>\u001b[39m \u001b[38;5;241m30\u001b[39m, tf\u001b[38;5;241m.\u001b[39mfloat32)\n\u001b[0;32m---> 19\u001b[0m limit_cca \u001b[38;5;241m=\u001b[39m \u001b[43mconnected_components_tf\u001b[49m\u001b[43m(\u001b[49m\u001b[43marr_t\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mnum_iterations\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m10\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m     20\u001b[0m unique_labels \u001b[38;5;241m=\u001b[39m tf\u001b[38;5;241m.\u001b[39mreduce_max(limit_cca) \u001b[38;5;241m+\u001b[39m \u001b[38;5;241m1\u001b[39m  \u001b[38;5;66;03m# Assuming labels are zero-indexed\u001b[39;00m\n\u001b[1;32m     21\u001b[0m limit \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mmax\u001b[39m(unique_labels, limit)\n",
      "Cell \u001b[0;32mIn[92], line 11\u001b[0m, in \u001b[0;36mconnected_components_tf\u001b[0;34m(image, num_iterations)\u001b[0m\n\u001b[1;32m      8\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mInput num_iterations must be a positive integer.\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m     10\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(image\u001b[38;5;241m.\u001b[39mshape) \u001b[38;5;241m<\u001b[39m \u001b[38;5;241m3\u001b[39m \u001b[38;5;129;01mor\u001b[39;00m image\u001b[38;5;241m.\u001b[39mshape[\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m3\u001b[39m] \u001b[38;5;241m!=\u001b[39m \u001b[38;5;241m1\u001b[39m:\n\u001b[0;32m---> 11\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mInput image shape must be (*, 1, H, W). Got: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mimage\u001b[38;5;241m.\u001b[39mshape\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m     13\u001b[0m \u001b[38;5;66;03m# Reshape image to 2D if it's more than 3D\u001b[39;00m\n\u001b[1;32m     14\u001b[0m shape \u001b[38;5;241m=\u001b[39m tf\u001b[38;5;241m.\u001b[39mshape(image)\n",
      "\u001b[0;31mValueError\u001b[0m: Input image shape must be (*, 1, H, W). Got: (1, 1, 1, 64, 64, 1)"
     ]
    }
   ],
   "source": [
    "# import tensorflow as tf\n",
    "\n",
    "# class LocatorTF:\n",
    "#     def __init__(self, model, **kwargs):\n",
    "#         self.model = model\n",
    "#         self.meanADU = kwargs.get('meanADU', 241.0)\n",
    "#         self.p_list = kwargs.get('p_list', [8, 6, 1.5, 1, 50])\n",
    "#         self.dynamic_thres = kwargs.get('dynamic_thres', True)\n",
    "\n",
    "#     def model_tune(self, arr):\n",
    "#         \"\"\"\n",
    "#         Dynamically adjust model thresholds based on the estimated image sparsity.\n",
    "#         \"\"\"\n",
    "#         meanADU = self.meanADU * 4  # mean ADU * upsample_factor^2\n",
    "#         limit = int(tf.reduce_sum(arr) / meanADU)\n",
    "\n",
    "#         # Assuming a function connected_components_tf exists that returns tensor of components\n",
    "#         arr_t = tf.cast(arr > 30, tf.float32)\n",
    "#         limit_cca = connected_components_tf(arr_t[None, None, ...], num_iterations=10)\n",
    "#         unique_labels = tf.reduce_max(limit_cca) + 1  # Assuming labels are zero-indexed\n",
    "#         limit = max(unique_labels, limit)\n",
    "#         limit = max(limit, 1)\n",
    "\n",
    "#         # Dynamic adjustment of model parameters based on limit\n",
    "#         detections_per_img = int(limit * self.p_list[2])\n",
    "#         score_thresh = self.p_list[3] / limit\n",
    "#         nms_thresh = 0.02  # Considered a constant in this context\n",
    "\n",
    "#         # Example method to update model parameters\n",
    "#         self.update_model_params(detections_per_img, score_thresh, nms_thresh)\n",
    "\n",
    "#     def update_model_params(self, detections_per_img, score_thresh, nms_thresh):\n",
    "#         \"\"\"\n",
    "#         Stub function to illustrate how one might update model parameters.\n",
    "#         Actual implementation would depend on model architecture.\n",
    "#         \"\"\"\n",
    "#         print(f\"Updating model parameters: Detections per img: {detections_per_img}, Score threshold: {score_thresh}, NMS threshold: {nms_thresh}\")\n",
    "#         # This method should interface with your model to update its operational parameters.\n",
    "#         # For example:\n",
    "#         self.model.detections_per_img = detections_per_img\n",
    "#         self.model.score_thresh = score_thresh\n",
    "#         self.model.nms_thresh = nms_thresh\n",
    "\n",
    "# # Assume the model is already defined and instantiated\n",
    "\n",
    "# locator = LocatorTF(model=model)\n",
    "# sample_image_tensor = tf.random.normal([1, 64, 64, 1])  # A sample tensor for demonstration\n",
    "# locator.model_tune(sample_image_tensor)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-04-24 15:02:55.393563: W tensorflow/core/framework/op_kernel.cc:1839] OP_REQUIRES failed at strided_slice_op.cc:117 : INVALID_ARGUMENT: slice index 4 of dimension 0 out of bounds.\n"
     ]
    },
    {
     "ename": "InvalidArgumentError",
     "evalue": "{{function_node __wrapped__StridedSlice_device_/job:localhost/replica:0/task:0/device:GPU:0}} slice index 4 of dimension 0 out of bounds. [Op:StridedSlice] name: strided_slice/",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mInvalidArgumentError\u001b[0m                      Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[103], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43msample_image_tensor\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m4\u001b[39;49m\u001b[43m]\u001b[49m\n",
      "File \u001b[0;32m~/anaconda3/envs/tiny_yolo/lib/python3.10/site-packages/tensorflow/python/util/traceback_utils.py:153\u001b[0m, in \u001b[0;36mfilter_traceback.<locals>.error_handler\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    151\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[1;32m    152\u001b[0m   filtered_tb \u001b[38;5;241m=\u001b[39m _process_traceback_frames(e\u001b[38;5;241m.\u001b[39m__traceback__)\n\u001b[0;32m--> 153\u001b[0m   \u001b[38;5;28;01mraise\u001b[39;00m e\u001b[38;5;241m.\u001b[39mwith_traceback(filtered_tb) \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m    154\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[1;32m    155\u001b[0m   \u001b[38;5;28;01mdel\u001b[39;00m filtered_tb\n",
      "File \u001b[0;32m~/anaconda3/envs/tiny_yolo/lib/python3.10/site-packages/tensorflow/python/framework/ops.py:5883\u001b[0m, in \u001b[0;36mraise_from_not_ok_status\u001b[0;34m(e, name)\u001b[0m\n\u001b[1;32m   5881\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mraise_from_not_ok_status\u001b[39m(e, name) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m NoReturn:\n\u001b[1;32m   5882\u001b[0m   e\u001b[38;5;241m.\u001b[39mmessage \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m (\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m name: \u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;241m+\u001b[39m \u001b[38;5;28mstr\u001b[39m(name \u001b[38;5;28;01mif\u001b[39;00m name \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m\"\u001b[39m))\n\u001b[0;32m-> 5883\u001b[0m   \u001b[38;5;28;01mraise\u001b[39;00m core\u001b[38;5;241m.\u001b[39m_status_to_exception(e) \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "\u001b[0;31mInvalidArgumentError\u001b[0m: {{function_node __wrapped__StridedSlice_device_/job:localhost/replica:0/task:0/device:GPU:0}} slice index 4 of dimension 0 out of bounds. [Op:StridedSlice] name: strided_slice/"
     ]
    }
   ],
   "source": [
    "sample_image_tensor"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "electron",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}