empatches_for_effi-seg_model_infer.ipynb

{
  "metadata": {
    "kernelspec": {
      "name": "python",
      "display_name": "Python (Pyodide)",
      "language": "python"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "python",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.8"
    }
  },
  "nbformat_minor": 4,
  "nbformat": 4,
  "cells": [
    {
      "cell_type": "markdown",
      "source": "# Patch-Based Efficient Segmentation Inference using Empatches Library\n\nWhile working with the Empatches library, we observed that dealing with large image sizes (e.g., 1500 x 1500 x 3) and generating patches with dimensions like 100 x 1500 x 1500 x 3 leads to memory overflow, especially when running on edge devices during model deployment. To address this issue, we propose a memory-efficient inference scheme. Instead of loading the entire patch tensor (100 x 1500 x 1500 x 3) into memory every time, we store the patches in temporary memory storage. Each patch is then fed individually into the model, and after processing, we use the indices to reconstruct the full-scale segmentation mask. This approach significantly reduces memory usage, making running on memory-constrained edge devices feasible while maintaining accurate segmentation. ",
      "metadata": {}
    },
    {
      "cell_type": "markdown",
      "source": "## Imports and Dependencies",
      "metadata": {}
    },
    {
      "cell_type": "code",
      "source": "import numpy as np\nimport math\nimport tempfile\nimport shutil\nimport os\nimport cv2",
      "metadata": {
        "trusted": true
      },
      "outputs": [],
      "execution_count": null
    },
    {
      "cell_type": "markdown",
      "source": "## EMPatches_Effi_Seg_Inference Code ",
      "metadata": {}
    },
    {
      "cell_type": "code",
      "source": "class EMPatches_Effi_Seg_Inference(object):\n    \n    def __init__(self):\n        self.temp_dir = None\n        self.temp_dir_path = None\n\n    def cleanup(self):\n        if self.temp_dir:\n            shutil.rmtree(self.temp_dir)\n            self.temp_dir = None\n\n    def extract_patches(self, data, patchsize, overlap=None, stride=None, vox=False ,base_temp_dir=None):\n        \n        if base_temp_dir is not None:\n            self.temp_dir = tempfile.mkdtemp(dir=base_temp_dir)\n        else:\n            self.temp_dir = tempfile.mkdtemp()\n\n        '''\n        Parameters\n        ----------\n        data : array to extract patches from; it can be 1D, 2D or 3D [W, H, D]. H: Height, W: Width, D: Depth,\n               3D data includes images (RGB, RGBA, etc) or Voxel data.\n        patchsize :  size of the patch to extract from the image only square patches can be\n                    extracted for now.\n        overlap (Optional): overlap between patched in percentage a float between [0, 1].\n        stride (Optional): Step size between patches\n        vox (Optional): Whether data is volumetric or not if set to the true array will be cropped in last dimension too.\n        \n        base_temp_dir (Optional) : temporary storage to save the patches and can be deleted later using  _ = emp.cleanup()\n\n        Returns\n        -------\n        temp_dir : Paths where the patches are saved into temporary memory.\n        indices : a list containing indices of patches in order, which can be used \n                at later stage for reconstruction.\n                \n        Original Dimensions:  (height, width,depth) of original shape of data\n\n        '''\n\n        height = data.shape[0]\n        width = data.shape[1]\n        depth = data.shape[2]\n\n        maxWindowSize = patchsize\n        windowSizeX = maxWindowSize\n        windowSizeY = maxWindowSize\n        windowSizeZ = maxWindowSize\n\n        windowSizeX = min(windowSizeX, width)\n        windowSizeY = min(windowSizeY, height)\n        windowSizeZ = min(windowSizeZ, depth)\n            \n        if stride is not None:\n                stepSizeX = stride\n                stepSizeY = stride\n                stepSizeZ = stride\n                        \n        elif overlap is not None:\n            overlapPercent = overlap\n\n            windowSizeX = maxWindowSize\n            windowSizeY = maxWindowSize\n            windowSizeZ = maxWindowSize\n            \n            # If the input data is smaller than the specified window size,\n            # clip the window size to the input size on both dimensions\n            windowSizeX = min(windowSizeX, width)\n            windowSizeY = min(windowSizeY, height)\n            windowSizeZ = min(windowSizeZ, depth)\n\n            # Compute the window overlap and step size\n            windowOverlapX = int(math.floor(windowSizeX * overlapPercent))\n            windowOverlapY = int(math.floor(windowSizeY * overlapPercent))\n            windowOverlapZ = int(math.floor(windowSizeZ * overlapPercent))\n\n            stepSizeX = windowSizeX - windowOverlapX\n            stepSizeY = windowSizeY - windowOverlapY                \n            stepSizeZ = windowSizeZ - windowOverlapZ                \n\n        else:\n            stepSizeX = 1\n            stepSizeY = 1\n            stepSizeZ = 1\n         \n        # Determine how many windows we will need in order to cover the input data\n        lastX = width - windowSizeX\n        lastY = height - windowSizeY\n        lastZ = depth - windowSizeZ\n        \n        xOffsets = list(range(0, lastX+1, stepSizeX))\n        yOffsets = list(range(0, lastY+1, stepSizeY))\n        zOffsets = list(range(0, lastZ+1, stepSizeZ))\n        \n        # Unless the input data dimensions are exact multiples of the step size,\n        # we will need one additional row and column of windows to get 100% coverage\n        if len(xOffsets) == 0 or xOffsets[-1] != lastX:\n            xOffsets.append(lastX)\n        if len(yOffsets) == 0 or yOffsets[-1] != lastY:\n            yOffsets.append(lastY)\n        if len(zOffsets) == 0 or zOffsets[-1] != lastZ:\n            zOffsets.append(lastZ)\n        indices = []\n        \n        patch_index = 0\n\n        for xOffset in xOffsets:\n            for yOffset in yOffsets:\n                #if len(data.shape) >= 3:\n                patch_path = os.path.join(self.temp_dir, f\"patch_{patch_index}.png\")\n                cv2.imwrite(patch_path, data[(slice(yOffset, yOffset+windowSizeY, None),\n                                        slice(xOffset, xOffset+windowSizeX, None))])\n                patch_index += 1    \n                indices.append((yOffset, yOffset+windowSizeY, xOffset, xOffset+windowSizeX))\n\n        return self.temp_dir, indices , (height,width,depth)",
      "metadata": {
        "trusted": true
      },
      "outputs": [],
      "execution_count": null
    },
    {
      "cell_type": "markdown",
      "source": "## Calling the EMPatches_Effi_Seg_Inference and Set paths",
      "metadata": {}
    },
    {
      "cell_type": "code",
      "source": "emp = EMPatches_Effi_Seg_Inference() \n\noutput_path = \"path to save segmentation mask\" + \"reconstructed_seg_mask.png\"\nimage_path = \"path to the input RGB image\"\ntemp_dir = \"Optional to save temporary patches\"",
      "metadata": {
        "trusted": true
      },
      "outputs": [],
      "execution_count": null
    },
    {
      "cell_type": "markdown",
      "source": "## Generate the patches and get the paths , indices and orignal image shape ",
      "metadata": {}
    },
    {
      "cell_type": "code",
      "source": "image = cv2.imread(image_path)\npatches_path, indices , org_shape = emp.extract_patches(image,patchsize=224, overlap=0.0 , base_temp_dir=temp_dir)",
      "metadata": {
        "trusted": true
      },
      "outputs": [],
      "execution_count": null
    },
    {
      "cell_type": "markdown",
      "source": " ## Simple Segmentation Model which converts RGB Image to Grayscale ",
      "metadata": {}
    },
    {
      "cell_type": "code",
      "source": "from torch import nn\nclass Seg_Model(nn.Module):\n    def __init__(self):\n        super().__init__()\n    def forward(self, x):\n        return cv2.cvtColor(x, cv2.COLOR_RGB2GRAY)\nmodel = Seg_Model()  ",
      "metadata": {
        "trusted": true
      },
      "outputs": [],
      "execution_count": null
    },
    {
      "cell_type": "markdown",
      "source": " ## Run the inference by loading each patch individually into memory",
      "metadata": {}
    },
    {
      "cell_type": "code",
      "source": "reconstructed_seg_mask= np.zeros((org_shape), dtype=np.uint8)\n# Iterate through patch files\nfor i, patch_file in enumerate(sorted(os.listdir(patches_path))):\n    patch_path = os.path.join(patches_path, patch_file)\n    patch = cv2.imread(patch_path)\n    y = model(patch)\n    y = np.stack((y,)*3, axis=-1)\n    y_start, y_end, x_start, x_end = indices[i]\n    reconstructed_seg_mask[y_start:y_end, x_start:x_end] = y\n    del patch,y\n\ncv2.imwrite(output_path, reconstructed_seg_mask)\nprint(f\"Reconstructed segmentaiton  saved to: {output_path}\")\n_ = emp.cleanup()   ## to delete the patches from memory ",
      "metadata": {
        "trusted": true
      },
      "outputs": [],
      "execution_count": null
    }
  ]
}