{"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":"## Efficient image loading with zarr\n### Brett Olsen, March 2023\n\nThis notebook should provide a simple, efficient tool for loading and accessing image data wrapped up in a single, persistent object.  We will use the zarr framework to store the compressed data and allow easy access to chunks of data at a time.\n\nIn particular, this allows fast-access persistent storage in your working directory that will enable loading of a full fragment dataset in milliseconds rather than minutes.\n\nFirst we will install the zarr package, import necessary packages, and define some path constants to appropriate folders.","metadata":{}},{"cell_type":"code","source":"!pip install zarr","metadata":{"execution":{"iopub.status.busy":"2023-03-27T17:10:05.362802Z","iopub.execute_input":"2023-03-27T17:10:05.363209Z","iopub.status.idle":"2023-03-27T17:10:18.974171Z","shell.execute_reply.started":"2023-03-27T17:10:05.363174Z","shell.execute_reply":"2023-03-27T17:10:18.972738Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport shutil\nimport time\nimport zarr\nimport numpy as np\nimport PIL.Image as Image\n\nINPUT_FOLDER = \"/kaggle/input/vesuvius-challenge-ink-detection\"\nWORKING_FOLDER = \"/kaggle/working/\"\nTEMP_FOLDER = \"kaggle/temp/\"","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-03-27T19:14:51.362460Z","iopub.execute_input":"2023-03-27T19:14:51.363547Z","iopub.status.idle":"2023-03-27T19:14:51.368687Z","shell.execute_reply.started":"2023-03-27T19:14:51.363476Z","shell.execute_reply":"2023-03-27T19:14:51.367790Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"This code is a simple context manager to make it simple to get timings for our load operations.","metadata":{}},{"cell_type":"code","source":"class TimerError(Exception):\n    pass\n\nclass Timer():\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()","metadata":{"execution":{"iopub.status.busy":"2023-03-27T19:32:59.569388Z","iopub.execute_input":"2023-03-27T19:32:59.569827Z","iopub.status.idle":"2023-03-27T19:32:59.579822Z","shell.execute_reply.started":"2023-03-27T19:32:59.569797Z","shell.execute_reply":"2023-03-27T19:32:59.578690Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"The FragmentImageData class is used to store the compressed data on disk.\nIt takes three parameters:  the sample type (\"test\" or \"train\"), the sample index in the folder, and a boolean that determines whether the data is stored in the (persistent) working directory or the temporary directory.  If stored persistently, then once generated, the zarr data format can be quickly loaded in other notebooks.\n\nIf the zarr file does not already exist, it will be generated by parsing the individual image files in the corresponding input directory.  Otherwise, it will quickly load from the\nzarr file.\n\nThe image data can be accessed as attributes of the object, largely (with the exception of fancy indexing) treated as numpy arrays:\n* `surface_volume`: the 3D X-ray tomography data\n* `mask`: the 2D boolean mask describing where data exists\n* `truth`: (Training data only) the 2D boolean mask with the ink truth set\n* `infrared`: (Training data only) the 2D infrared image of the parchment","metadata":{}},{"cell_type":"code","source":"class FragmentImageException(Exception):\n    pass\n\nclass FragmentImageData:\n    \"\"\"A general class that uses persistent zarr objects to store the surface volume data,\n    binary data mask, and for training sets, the truth data and infrared image of a papyrus\n    fragment, in a compressed and efficient way.\n    \"\"\"\n    def __init__(self, sample_type: str, sample_index: str, working: bool = True):\n        if sample_type not in (\"test, train\"):\n            raise FragmentImageException(\n                f\"Invalid sample type f{sample_type}, must be one of 'test' or 'train'\"\n            )\n        zarrpath = self._zarr_path(sample_type, sample_index, working)\n        if os.path.exists(zarrpath):\n            self.zarr = self.load_from_zarr(zarrpath)\n        else:\n            dirpath = os.path.join(INPUT_FOLDER, sample_type, sample_index)\n            if not os.path.exists(dirpath):\n                raise FragmentImageException(\n                    f\"No input data found at f{zarrpath} or f{dirpath}\"\n                )\n            self.zarr = self.load_from_directory(dirpath, zarrpath)\n    \n    @property\n    def surface_volume(self):\n        return self.zarr.surface_volume\n    \n    @property\n    def mask(self):\n        return self.zarr.mask\n    \n    @property\n    def truth(self):\n        return self.zarr.truth\n    \n    @property\n    def infrared(self):\n        return self.zarr.infrared\n    \n    @staticmethod\n    def _zarr_path(sample_type: str, sample_index: str, working: bool = True):\n        filename = f\"{sample_type}-{sample_index}.zarr\"\n        if working:\n            return os.path.join(WORKING_FOLDER, filename)\n        else:\n            return os.path.join(TEMP_FOLDER, filename)\n    \n    @staticmethod\n    def clean_zarr(sample_type: str, sample_index: str, working: bool = True):\n        zarrpath = FragmentImageData._zarr_path(sample_type, sample_index, working)\n        if os.path.exists(zarrpath):\n            shutil.rmtree(zarrpath)\n    \n    @staticmethod\n    def load_from_zarr(filepath):\n        with Timer(\"Loading from existing zarr\"):\n            return zarr.open(filepath, mode=\"r\")\n    \n    @staticmethod\n    def load_from_directory(dirpath, zarrpath):\n        if os.path.exists(zarrpath):\n            raise FragmentImageException(\n                f\"Trying to overwrite existing zarr at f{zarrpath}\"\n            )\n        # Initialize the root zarr group and write the file\n        root = zarr.open_group(zarrpath, mode=\"w\")\n        # Load in the surface volume tif files\n        with Timer(\"Surface volume loading\"):\n            init = True\n            imgfiles = sorted([\n                imgfile for imgfile in\n                os.listdir(os.path.join(dirpath, \"surface_volume\"))\n            ])\n            for imgfile in imgfiles:\n                print(f\"Loading file {imgfile}\", end=\"\\r\")\n                img_data = np.array(\n                    Image.open(os.path.join(dirpath, \"surface_volume\", imgfile))\n                )\n                if init:\n                    surface_volume = root.zeros(\n                        name=\"surface_volume\",\n                        shape=(img_data.shape[0], img_data.shape[1], len(imgfiles)),\n                        chunks=(1000, 1000, 4),\n                        dtype=img_data.dtype,\n                        write_empty_chunks=False,\n                    )\n                    init = False\n                z_index = int(imgfile.split(\".\")[0])\n                surface_volume[:,:,z_index] = img_data\n        # Load in the mask\n        with Timer(\"Mask loading\"):\n            img_data = np.array(Image.open(os.path.join(dirpath, \"mask.png\")), dtype=bool)\n            mask = root.array(\n                name=\"mask\",\n                data=img_data,\n                shape=img_data.shape,\n                chunks=(1000, 1000),\n                dtype=img_data.dtype,\n                write_empty_chunks=False,\n            )\n        # Load in the truth set (if it exists)\n        with Timer(\"Truth set loading\"):\n            truthfile = os.path.join(dirpath, \"inklabels.png\")\n            if os.path.exists(truthfile):\n                img_data = np.array(Image.open(truthfile), dtype=bool)\n                truth = root.array(\n                    name=\"truth\",\n                    data=img_data,\n                    shape=img_data.shape,\n                    chunks=(1000, 1000),\n                    dtype=img_data.dtype,\n                    write_empty_chunks=False,\n                )\n        # Load in the infrared image (if it exists)\n        with Timer(\"Infrared image loading\"):\n            irfile = os.path.join(dirpath, \"ir.png\")\n            if os.path.exists(irfile):\n                img_data = np.array(Image.open(irfile))\n                infrared = root.array(\n                    name = \"infrared\",\n                    data = img_data,\n                    shape = img_data.shape,\n                    chunks = (1000, 1000),\n                    dtype=img_data.dtype,\n                    write_empty_chunks=False,\n                )\n        return root        ","metadata":{"execution":{"iopub.status.busy":"2023-03-27T19:46:19.355390Z","iopub.execute_input":"2023-03-27T19:46:19.355923Z","iopub.status.idle":"2023-03-27T19:46:19.382095Z","shell.execute_reply.started":"2023-03-27T19:46:19.355880Z","shell.execute_reply":"2023-03-27T19:46:19.380542Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Let's see how long it takes to generate a new zarr file from scratch on the first training set data.  Note we need to clean up any pre-existing zarr data first, otherwise it will load directly from there instead of reading the input images.","metadata":{}},{"cell_type":"code","source":"FragmentImageData.clean_zarr(\"train\", 1)\ndata = FragmentImageData(\"train\", 1)","metadata":{"execution":{"iopub.status.busy":"2023-03-27T19:46:21.576974Z","iopub.execute_input":"2023-03-27T19:46:21.577398Z","iopub.status.idle":"2023-03-27T19:47:13.500624Z","shell.execute_reply.started":"2023-03-27T19:46:21.577356Z","shell.execute_reply":"2023-03-27T19:47:13.499494Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"We can see it took on the order of a minute to load all the images for this training dataset.  However, now that we have the data on disk in our working directory, we can reload the data from that zarr much faster:","metadata":{}},{"cell_type":"code","source":"data = FragmentImageData(\"train\", 1)","metadata":{"execution":{"iopub.status.busy":"2023-03-27T19:47:44.597458Z","iopub.execute_input":"2023-03-27T19:47:44.597907Z","iopub.status.idle":"2023-03-27T19:47:44.605214Z","shell.execute_reply.started":"2023-03-27T19:47:44.597874Z","shell.execute_reply":"2023-03-27T19:47:44.603405Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"We can look at the `info` attributes of the data to see some details about how the data is stored:","metadata":{}},{"cell_type":"code","source":"print(data.surface_volume.info)\nprint(data.mask.info)\nprint(data.truth.info)\nprint(data.infrared.info)","metadata":{"execution":{"iopub.status.busy":"2023-03-27T19:49:14.490730Z","iopub.execute_input":"2023-03-27T19:49:14.491689Z","iopub.status.idle":"2023-03-27T19:49:14.515098Z","shell.execute_reply.started":"2023-03-27T19:49:14.491653Z","shell.execute_reply":"2023-03-27T19:49:14.513606Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"We can get some useful information about how this works from this.\nThe data is \"chunked\" into groups of 1K by 1K by 4 pixels.  When slicing into the data arrays, only the chunks needed for that slice are loaded from disk.\nBecause the data is compressed, we use significantly less space than the raw images (about half).\nDepending on how one wants to use the data, tweaking the chunk sizes can improve performance by loading fewer chunks on the data you plan on working with.\n\nLet's see how to use the data.  First let's just look at a couple of standard slices.","metadata":{}},{"cell_type":"code","source":"import matplotlib.pyplot as plt","metadata":{"execution":{"iopub.status.busy":"2023-03-27T17:53:39.976443Z","iopub.execute_input":"2023-03-27T17:53:39.976888Z","iopub.status.idle":"2023-03-27T17:53:39.982569Z","shell.execute_reply.started":"2023-03-27T17:53:39.976853Z","shell.execute_reply":"2023-03-27T17:53:39.981346Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with Timer():\n    plt.imshow(data.infrared, cmap=\"gray\")","metadata":{"execution":{"iopub.status.busy":"2023-03-27T19:56:20.120627Z","iopub.execute_input":"2023-03-27T19:56:20.121954Z","iopub.status.idle":"2023-03-27T19:56:22.228532Z","shell.execute_reply.started":"2023-03-27T19:56:20.121894Z","shell.execute_reply":"2023-03-27T19:56:22.227351Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with Timer():\n    plt.imshow(data.surface_volume[:,:,20], cmap=\"gray\")","metadata":{"execution":{"iopub.status.busy":"2023-03-27T20:01:46.266604Z","iopub.execute_input":"2023-03-27T20:01:46.266973Z","iopub.status.idle":"2023-03-27T20:01:50.204256Z","shell.execute_reply.started":"2023-03-27T20:01:46.266942Z","shell.execute_reply":"2023-03-27T20:01:50.203235Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with Timer():\n    plt.imshow(data.surface_volume[2000:3000,2000:3000,20], cmap=\"gray\")","metadata":{"execution":{"iopub.status.busy":"2023-03-27T19:57:14.394711Z","iopub.execute_input":"2023-03-27T19:57:14.395395Z","iopub.status.idle":"2023-03-27T19:57:14.677279Z","shell.execute_reply.started":"2023-03-27T19:57:14.395318Z","shell.execute_reply":"2023-03-27T19:57:14.675645Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Note that showing an image of this smaller region was about 10 times faster, because only the chunk with that data in it was loaded, rather than all of the chunks.  \n\nOr we could look at a vertical slice:","metadata":{}},{"cell_type":"code","source":"with Timer():\n    plt.figure(figsize=(10, 1))\n    plt.imshow(data.surface_volume[2000,:,:].T, cmap=\"gray\", aspect=\"auto\")","metadata":{"execution":{"iopub.status.busy":"2023-03-27T20:02:47.081787Z","iopub.execute_input":"2023-03-27T20:02:47.082894Z","iopub.status.idle":"2023-03-27T20:02:47.537975Z","shell.execute_reply.started":"2023-03-27T20:02:47.082844Z","shell.execute_reply":"2023-03-27T20:02:47.536945Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Cool.  We can see the edges of the empty mask at either end, but we can also see the shape of the papyrus through the X-ray tomography.","metadata":{}}]}