{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.14","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":84969,"databundleVersionId":10033515,"sourceType":"competition"}],"dockerImageVersionId":30786,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install zarr &> null\n!pip install git+https://github.com/rostepifanov/voxelmentations &> null","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2024-11-30T18:33:10.624770Z","iopub.execute_input":"2024-11-30T18:33:10.625150Z","iopub.status.idle":"2024-11-30T18:33:39.234967Z","shell.execute_reply.started":"2024-11-30T18:33:10.625116Z","shell.execute_reply":"2024-11-30T18:33:39.233435Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import zarr\nimport json\nimport numpy as np\n\nfrom collections import OrderedDict\nfrom torch.utils.data import Dataset\n\nclass cryoetIMDataset(Dataset):\n    \"\"\"Dataset for loading CT images and points\n    \"\"\"\n    def __init__(self, datapath, names, *, augs=None):\n        \"\"\"\n            :args:\n                datapath: pathlib.Path\n                    directiry path with data\n                names: list of str\n                    names of data to load\n                augs: dict like\n                    option to augs\n        \"\"\"\n        self.augs = augs\n\n        self.imgs = OrderedDict()\n        self.masks = dict()\n        self.points = dict()\n        self.types = dict()\n\n        self.keys = list()\n        self.shapes = list()\n\n        self.mtypes = ['apo-ferritin', 'beta-galactosidase', 'ribosome', 'thyroglobulin', 'virus-like-particle']\n\n        particle_radius = {\n            'apo-ferritin': 60,\n            'beta-amylase': 65,\n            'beta-galactosidase': 90,\n            'ribosome': 150,\n            'thyroglobulin': 130,\n            'virus-like-particle': 135,\n        }\n        \n        for name in names:\n            group = zarr.open(datapath / 'static/ExperimentRuns' / name / 'VoxelSpacing10.000/denoised.zarr', mode='r')\n            voxel = group[0].astype(np.float32)\n            voxel = np.transpose(voxel, (1, 2, 0)) # set orientation as xyz\n            \n            scale = group.attrs['multiscales'][0]['datasets'][0]['coordinateTransformations'][0]['scale'] # read correct scale not just 10\n            \n            mask = np.zeros_like(voxel, dtype=np.uint8)\n\n            points = []\n            types = []\n            \n            x, y, z = np.indices(mask.shape)\n\n            for tdx, type_ in enumerate(self.mtypes):\n                with open(datapath / 'overlay/ExperimentRuns' / name / 'Picks' / (type_ + '.json')) as f:\n                    data = json.load(f)\n\n                    for entry in data['points']:\n                        location = entry['location']\n                        point = np.array([location['y'], location['x'], location['z']]) / scale\n\n                        ### some code to draw point on mask\n                        xc, yc, zc = point\n\n                        distance_sq = (x - xc)**2 + (y - yc)**2 + (z - zc)**2\n\n                        radius = 15 * particle_radius[type_] / 150\n                        mask[distance_sq <= radius**2] = tdx + 1\n                        ###\n\n                        points.append(np.array([*point, 1.], dtype=np.float32))\n                        types.append(tdx+1)\n\n            points = np.array(points, dtype=np.float32)\n            types = np.array(types, dtype=np.uint8)\n\n            self.imgs[name] = voxel\n            self.masks[name] = mask\n            self.points[name] = points\n            self.types[name] = types\n\n            self.keys.append(name)\n            self.shapes.append(voxel.shape)\n\n    def __len__(self):\n        return len(self.imgs)\n\n    def __getitem__(self, idx):\n        key, selector = idx\n\n        voxel = self.imgs[key][selector]\n        voxel = voxel[:, :, :, None]\n\n        mask = self.masks[key][selector]\n\n        if self.augs:\n            transformed = augs(voxel=voxel, mask=mask)\n\n            voxel = transformed['voxel']\n            mask = transformed['mask']\n\n        voxel = np.moveaxis(voxel, -1, 0)\n\n        return voxel, mask, key, selector","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-30T18:33:39.237717Z","iopub.execute_input":"2024-11-30T18:33:39.238085Z","iopub.status.idle":"2024-11-30T18:33:42.436081Z","shell.execute_reply.started":"2024-11-30T18:33:39.238052Z","shell.execute_reply":"2024-11-30T18:33:42.434929Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import voxelmentations as V\n\nfrom pathlib import Path\n\naugs = V.Sequential([\n    V.AxialPlaneFlip(p=0.5),\n    V.AxialPlaneTranpose(p=0.5),\n])\n\ngpath = Path('/kaggle/input/czii-cryo-et-object-identification')\n\ndataset = cryoetIMDataset(gpath / 'train', ['TS_5_4'], augs=augs)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-30T18:33:42.437372Z","iopub.execute_input":"2024-11-30T18:33:42.437959Z","iopub.status.idle":"2024-11-30T18:36:21.707020Z","shell.execute_reply.started":"2024-11-30T18:33:42.437925Z","shell.execute_reply":"2024-11-30T18:36:21.705651Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import matplotlib.pyplot as plt\n\nfig, ax = plt.subplots(1, 2, figsize=(10, 20))\n\nselector = (slice(0, None), slice(0, None), slice(0, 24))\ndata = dataset[('TS_5_4', selector)]\n\nax[0].set_title('Selected patch of image\\nwith the shape of 600x600x24')\nax[0].imshow(data[0][0, :, :, 12])\nax[0].axis('off')\n\nax[1].set_title('Selected patch of mask\\nwith the shape of 600x600x24')\nax[1].imshow(data[1][:, :, 12])\nax[1].axis('off')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-30T18:36:21.709451Z","iopub.execute_input":"2024-11-30T18:36:21.710324Z","iopub.status.idle":"2024-11-30T18:36:22.210035Z","shell.execute_reply.started":"2024-11-30T18:36:21.710269Z","shell.execute_reply":"2024-11-30T18:36:22.208670Z"}},"outputs":[],"execution_count":null}]}