{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## Simpler zarr .tif image loading\n### Brett Olsen, April 2023\n\nMy previous approaches to loading the tif image stack used some special handling and a custom class that had to be imported from a github repository.  I've since identified a much easier but just as efficient way of loading it, which I'll demonstrate here.  This just requires a single new package install.","metadata":{}},{"cell_type":"code","source":"!pip install zarr","metadata":{"execution":{"iopub.status.busy":"2023-04-27T22:52:28.698967Z","iopub.execute_input":"2023-04-27T22:52:28.699889Z","iopub.status.idle":"2023-04-27T22:52:40.158556Z","shell.execute_reply.started":"2023-04-27T22:52:28.699837Z","shell.execute_reply":"2023-04-27T22:52:40.156929Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import tifffile\nimport numpy \nimport matplotlib.pyplot as plt\nimport os\nimport time\nimport psutil\nimport zarr\n\nINPUT_FOLDER = \"/kaggle/input/vesuvius-challenge-ink-detection/\"","metadata":{"execution":{"iopub.status.busy":"2023-04-27T22:52:40.165724Z","iopub.execute_input":"2023-04-27T22:52:40.166055Z","iopub.status.idle":"2023-04-27T22:52:40.280520Z","shell.execute_reply.started":"2023-04-27T22:52:40.166021Z","shell.execute_reply":"2023-04-27T22:52:40.279451Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"The `tifffile` package includes methods for loading a directory of .tif files in as a `zarr` array without directly needing to use the `zarr` package. `zarr` arrays only load data into memory lazily, which means we can hold all these zarr objects in memory and only when we pull slices out from the array will we use memory.","metadata":{}},{"cell_type":"code","source":"class TimerError(Exception):\n    pass\n\nclass Timer():\n    \"\"\"This is a utility class that, when used as a context manager,\n    will report the time spent on code inside its block.\n    \"\"\"\n    def __init__(self, text=None):\n        if text is not None:\n            self.text = text + \": {:0.4f} seconds\"\n        else:\n            self.text = \"Elapsed time: {:0.4f} seconds\"\n        def logfunc(x):\n            print(x)\n        self.logger = logfunc\n        self._start_time = None\n\n    def start(self):\n        if self._start_time is not None:\n            raise TimerError(\"Timer is already running.  Use .stop() to stop it.\")\n        self._start_time = time.time()\n\n    def stop(self):\n        if self._start_time is None:\n            raise TimerError(\"Timer is not running.  Use .start() to start it.\")\n        elapsed_time = time.time() - self._start_time\n        self._start_time = None\n\n        if self.logger is not None:\n            self.logger(self.text.format(elapsed_time))\n\n        return elapsed_time\n\n    def __enter__(self):\n        self.start()\n        return self\n    \n    def __exit__(self, exc_type, exc_value, exc_traceback):\n        self.stop()\n\ndef show_mem_use():\n    process = psutil.Process()\n    mb_mem = process.memory_info().rss / 1e6\n    print(f\"{mb_mem:6.2f} MB used\")","metadata":{"execution":{"iopub.status.busy":"2023-04-27T22:52:40.281967Z","iopub.execute_input":"2023-04-27T22:52:40.282302Z","iopub.status.idle":"2023-04-27T22:52:40.294287Z","shell.execute_reply.started":"2023-04-27T22:52:40.282270Z","shell.execute_reply":"2023-04-27T22:52:40.292767Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"show_mem_use()\nwith Timer():\n    tif_folder = os.path.join(INPUT_FOLDER, \"train\", \"2\", \"surface_volume\", \"*.tif\")\n    frag2_store = tifffile.imread(tif_folder, aszarr=True)\n    frag2_stack = zarr.open(frag2_store, mode=\"r\")\nshow_mem_use()\nprint(frag2_stack)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-04-27T22:52:40.297620Z","iopub.execute_input":"2023-04-27T22:52:40.298049Z","iopub.status.idle":"2023-04-27T22:52:40.414863Z","shell.execute_reply.started":"2023-04-27T22:52:40.298012Z","shell.execute_reply":"2023-04-27T22:52:40.413514Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"You see that loading the whole .tif stack only took a small amount of time and used a couple hundred megabytes of RAM.\nThe z-levels are the first index and the x,y coordinates are the other indices.\n\nNow we can access any bit of data we want without using an enormous amount of memory:","metadata":{}},{"cell_type":"code","source":"show_mem_use()\nwith Timer():\n    frag_slice = frag2_stack[35,:,:]\nplt.imshow(frag_slice)\nshow_mem_use()","metadata":{"execution":{"iopub.status.busy":"2023-04-27T22:52:40.416554Z","iopub.execute_input":"2023-04-27T22:52:40.417548Z","iopub.status.idle":"2023-04-27T22:52:45.772114Z","shell.execute_reply.started":"2023-04-27T22:52:40.417497Z","shell.execute_reply":"2023-04-27T22:52:45.770704Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"show_mem_use()\nwith Timer():\n    frag_slice = frag2_stack[32,2000:12000,2000:8000]\nplt.imshow(frag_slice)\nshow_mem_use()","metadata":{"execution":{"iopub.status.busy":"2023-04-27T22:52:45.773738Z","iopub.execute_input":"2023-04-27T22:52:45.774322Z","iopub.status.idle":"2023-04-27T22:52:48.652586Z","shell.execute_reply.started":"2023-04-27T22:52:45.774267Z","shell.execute_reply":"2023-04-27T22:52:48.651678Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"We used 500 MB or so to load that single z-slice in and take a look at it, but as soon as we used that same variable to hold a different slice, we did _not_ use additional memory for the new slice - the old one got deallocated.\n\nYou can use the zarr array just like a numpy array for most purposes, except that the data's held on disk and loaded into RAM when you slice.","metadata":{}}]}