{"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":"markdown","source":"# Overview\n\nThis notebook visualizes particle data in interactive 3d plots.\n\n* Version 6: Add annimatable visualization of particles and image slices.","metadata":{}},{"cell_type":"code","source":"!pip install zarr","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-13T07:49:09.699531Z","iopub.execute_input":"2024-11-13T07:49:09.699875Z","iopub.status.idle":"2024-11-13T07:49:22.172703Z","shell.execute_reply.started":"2024-11-13T07:49:09.699838Z","shell.execute_reply":"2024-11-13T07:49:22.171212Z"},"_kg_hide-output":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Import dependencies\nfrom glob import glob\nimport math\nimport os\n\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport pandas as pd\nimport plotly\nimport plotly.express as px\nimport plotly.graph_objects as go\nimport plotly.io as pio\nimport zarr\n\npio.renderers.default = 'iframe'","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-13T07:49:22.174983Z","iopub.execute_input":"2024-11-13T07:49:22.175570Z","iopub.status.idle":"2024-11-13T07:49:23.308332Z","shell.execute_reply.started":"2024-11-13T07:49:22.175508Z","shell.execute_reply":"2024-11-13T07:49:23.306771Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Get experiment runs\nruns = sorted(glob('/kaggle/input/czii-cryo-et-object-identification/train/overlay/ExperimentRuns/*'))\nruns = [os.path.basename(x) for x in runs]\nruns","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-13T07:49:23.311086Z","iopub.execute_input":"2024-11-13T07:49:23.311597Z","iopub.status.idle":"2024-11-13T07:49:23.322462Z","shell.execute_reply.started":"2024-11-13T07:49:23.311542Z","shell.execute_reply":"2024-11-13T07:49:23.321376Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Functions for data loading / visualization\n\ndef read_run(run: str) -> pd.DataFrame:\n    \"\"\"Read a experiment run.\"\"\"\n    # Read all types of particle data\n    paths = glob(\n        f\"/kaggle/input/czii-cryo-et-object-identification/train/overlay/ExperimentRuns/{run}/Picks/*.json\"\n    )\n    df = pd.concat([pd.read_json(x) for x in paths]).reset_index(drop=True)\n\n    # Append point information columns\n    for axis in \"x\", \"y\", \"z\":\n        df[axis] = df.points.apply(lambda x: x[\"location\"][axis])\n    for key in \"transformation_\", \"instance_id\":\n        df[key] = df.points.apply(lambda x: x[key])\n    return df\n\ndef plot_particles(\n    df: pd.DataFrame, scale: float = 1.0, marker_size: float = 2.0\n) -> plotly.graph_objs._figure.Figure:\n    \"\"\"Plot 3D scatter plot of particles.\"\"\"\n    df = df.copy()\n    df[[\"x\", \"y\", \"z\"]] *= scale\n    fig = px.scatter_3d(df, x=\"x\", y=\"y\", z=\"z\", color=\"pickable_object_name\")\n    fig.update_traces(marker=dict(size=marker_size))\n    fig.update_layout(\n        title=df.run_name.iloc[0],\n        scene=dict(\n            yaxis=dict(autorange=\"reversed\"),\n            camera=dict(eye=dict(x=1.25, y=-1.25, z=1.25)),\n        ),\n        width=800,\n        height=800,\n        template=\"plotly_dark\",\n    )\n    return fig\n\ndef read_zarr(run: str) -> zarr.hierarchy.Group:\n    \"\"\"Read a zarr data (denoised.zarr).\"\"\"\n    return zarr.open(\n        f\"/kaggle/input/czii-cryo-et-object-identification/train/static/ExperimentRuns/{run}/VoxelSpacing10.000/denoised.zarr\",\n        mode=\"r\",\n    )\n\ndef plot_zarr_images(arr: zarr.core.Array, ncols: int = 6, axsize: float = 2.0):\n    \"\"\"Plot zarr images.\"\"\"\n    nslices = len(arr)\n    nrows = math.ceil(nslices / ncols)\n\n    fig = plt.figure(figsize=(axsize * ncols, axsize * nrows))\n    for i in range(nslices):\n        ax = plt.subplot(nrows, ncols, i + 1)\n        ax.set_xticks([])\n        ax.set_yticks([])\n        ax.imshow(arr[i])\n    plt.subplots_adjust(left=0, bottom=0, right=1, top=1, wspace=0.05, hspace=0.05)\n    plt.show()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-13T07:49:23.324497Z","iopub.execute_input":"2024-11-13T07:49:23.324932Z","iopub.status.idle":"2024-11-13T07:49:23.341413Z","shell.execute_reply.started":"2024-11-13T07:49:23.324871Z","shell.execute_reply":"2024-11-13T07:49:23.339867Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Visualize All Samples","metadata":{}},{"cell_type":"markdown","source":"## TS_5_4","metadata":{}},{"cell_type":"code","source":"df = read_run(runs[0])\ndf","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-13T07:49:23.343128Z","iopub.execute_input":"2024-11-13T07:49:23.344295Z","iopub.status.idle":"2024-11-13T07:49:23.469957Z","shell.execute_reply.started":"2024-11-13T07:49:23.344239Z","shell.execute_reply":"2024-11-13T07:49:23.468718Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"* `user_id` is all \"curation\"\n* `session_id` is all 0\n* `run_name` is all \"TS_5_4\"\n* `voxel_spacing` is all `NaN`\n* `unit` is all \"angstrom\"\n* `trust_orientation` is all `True`\n* `transformation_` is all identity matrix\n* `instance_id` is all 0","metadata":{}},{"cell_type":"code","source":"# Particles\nplot_particles(df)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-13T07:49:23.471348Z","iopub.execute_input":"2024-11-13T07:49:23.471687Z","iopub.status.idle":"2024-11-13T07:49:24.220378Z","shell.execute_reply.started":"2024-11-13T07:49:23.471650Z","shell.execute_reply":"2024-11-13T07:49:24.218831Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Images\nzarr_group = read_zarr(runs[0])\nplot_zarr_images(zarr_group[2])  # index 2: low resolution array","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-13T07:49:24.222190Z","iopub.execute_input":"2024-11-13T07:49:24.222698Z","iopub.status.idle":"2024-11-13T07:49:29.117365Z","shell.execute_reply.started":"2024-11-13T07:49:24.222652Z","shell.execute_reply":"2024-11-13T07:49:29.115423Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## TS_69_2","metadata":{}},{"cell_type":"code","source":"# Particles\nplot_particles(read_run(runs[1]))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-13T07:49:29.124082Z","iopub.execute_input":"2024-11-13T07:49:29.124735Z","iopub.status.idle":"2024-11-13T07:49:29.325685Z","shell.execute_reply.started":"2024-11-13T07:49:29.124661Z","shell.execute_reply":"2024-11-13T07:49:29.324408Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Images\nzarr_group = read_zarr(runs[1])\nplot_zarr_images(zarr_group[2])  # index 2: low resolution array","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-13T07:49:29.327196Z","iopub.execute_input":"2024-11-13T07:49:29.327592Z","iopub.status.idle":"2024-11-13T07:49:34.476909Z","shell.execute_reply.started":"2024-11-13T07:49:29.327549Z","shell.execute_reply":"2024-11-13T07:49:34.474939Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## TS_6_4","metadata":{}},{"cell_type":"code","source":"# Particles\nplot_particles(read_run(runs[2]))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-13T07:49:34.478558Z","iopub.execute_input":"2024-11-13T07:49:34.478970Z","iopub.status.idle":"2024-11-13T07:49:34.676811Z","shell.execute_reply.started":"2024-11-13T07:49:34.478927Z","shell.execute_reply":"2024-11-13T07:49:34.674933Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Images\nzarr_group = read_zarr(runs[2])\nplot_zarr_images(zarr_group[2])  # index 2: low resolution array","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-13T07:49:34.678498Z","iopub.execute_input":"2024-11-13T07:49:34.679116Z","iopub.status.idle":"2024-11-13T07:49:39.444493Z","shell.execute_reply.started":"2024-11-13T07:49:34.679063Z","shell.execute_reply":"2024-11-13T07:49:39.442641Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## TS_6_6","metadata":{}},{"cell_type":"code","source":"# Particles\nplot_particles(read_run(runs[3]))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-13T07:49:39.445857Z","iopub.execute_input":"2024-11-13T07:49:39.446243Z","iopub.status.idle":"2024-11-13T07:49:39.842776Z","shell.execute_reply.started":"2024-11-13T07:49:39.446204Z","shell.execute_reply":"2024-11-13T07:49:39.841468Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Images\nzarr_group = read_zarr(runs[3])\nplot_zarr_images(zarr_group[2])  # index 2: low resolution array","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-13T07:49:39.844263Z","iopub.execute_input":"2024-11-13T07:49:39.844646Z","iopub.status.idle":"2024-11-13T07:49:44.614524Z","shell.execute_reply.started":"2024-11-13T07:49:39.844605Z","shell.execute_reply":"2024-11-13T07:49:44.612299Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## TS_73_6","metadata":{}},{"cell_type":"code","source":"# Particles\nplot_particles(read_run(runs[4]))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-13T07:49:44.616164Z","iopub.execute_input":"2024-11-13T07:49:44.616601Z","iopub.status.idle":"2024-11-13T07:49:44.810943Z","shell.execute_reply.started":"2024-11-13T07:49:44.616563Z","shell.execute_reply":"2024-11-13T07:49:44.809692Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Images\nzarr_group = read_zarr(runs[4])\nplot_zarr_images(zarr_group[2])  # index 2: low resolution array","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-13T07:49:44.812733Z","iopub.execute_input":"2024-11-13T07:49:44.813304Z","iopub.status.idle":"2024-11-13T07:49:49.568555Z","shell.execute_reply.started":"2024-11-13T07:49:44.813241Z","shell.execute_reply":"2024-11-13T07:49:49.566653Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## TS_86_3","metadata":{}},{"cell_type":"code","source":"# Particles\nplot_particles(read_run(runs[5]))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-13T07:49:49.570619Z","iopub.execute_input":"2024-11-13T07:49:49.571230Z","iopub.status.idle":"2024-11-13T07:49:49.776632Z","shell.execute_reply.started":"2024-11-13T07:49:49.571169Z","shell.execute_reply":"2024-11-13T07:49:49.775460Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Images\nzarr_group = read_zarr(runs[5])\nplot_zarr_images(zarr_group[2])  # index 2: low resolution array","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-13T07:49:49.778478Z","iopub.execute_input":"2024-11-13T07:49:49.779514Z","iopub.status.idle":"2024-11-13T07:49:54.604860Z","shell.execute_reply.started":"2024-11-13T07:49:49.779456Z","shell.execute_reply":"2024-11-13T07:49:54.602747Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## TS_99_9","metadata":{}},{"cell_type":"code","source":"# Particles\nplot_particles(read_run(runs[6]))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-13T07:49:54.606370Z","iopub.execute_input":"2024-11-13T07:49:54.606738Z","iopub.status.idle":"2024-11-13T07:49:54.813214Z","shell.execute_reply.started":"2024-11-13T07:49:54.606702Z","shell.execute_reply":"2024-11-13T07:49:54.811692Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Images\nzarr_group = read_zarr(runs[6])\nplot_zarr_images(zarr_group[2])  # index 2: low resolution array","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-13T07:49:54.814923Z","iopub.execute_input":"2024-11-13T07:49:54.815483Z","iopub.status.idle":"2024-11-13T07:49:59.546657Z","shell.execute_reply.started":"2024-11-13T07:49:54.815424Z","shell.execute_reply":"2024-11-13T07:49:59.543021Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Visualize Particles with Annimatable Slices\n\nThis computation is heavy.","metadata":{}},{"cell_type":"code","source":"def _normalize_array(zarr_array: zarr.core.Array) -> np.ndarray:\n    \"\"\"Normalize each slice of the array.\"\"\"\n    arr = np.array(zarr_array)\n    mins = arr.min(axis=(1, 2), keepdims=True)\n    maxs = arr.max(axis=(1, 2), keepdims=True)\n    arr = ((arr - mins) / (maxs - mins) * 255).astype(np.uint8)\n    return arr\n\ndef _plot_slice_surface(array: zarr.core.Array, z_index: int):\n    \"\"\"Plot surface plot of specified z-slice of array.\"\"\"\n    return go.Surface(\n        z=z_index * np.ones((array.shape[1], array.shape[2])),\n        surfacecolor=array[z_index],\n        colorscale=\"gray\",\n        cmin=0,\n        cmax=255,\n        showscale=False,\n    )\n\ndef plot_animatable_slices(\n    arr: zarr.core.Array,\n    step: int = 10,\n    init_z: int = 0,\n    title: str = \"\",\n    fig: plotly.graph_objs._figure.Figure | None = None,\n) -> plotly.graph_objs._figure.Figure:\n    \"\"\"Plot animatable slices.\"\"\"\n    base_traces = list(fig.data) if fig is not None else []\n\n    arr = _normalize_array(arr)\n    z_dim = len(arr)\n\n    # Z-indices to annimate\n    z_indices = range(0, z_dim, step)\n\n    # Initial plot\n    fig = go.Figure(\n        data=base_traces + [_plot_slice_surface(arr, init_z)],\n        layout=go.Layout(\n            title=title,\n            scene=dict(\n                yaxis=dict(autorange=\"reversed\"),\n                zaxis=dict(range=[0, z_dim], autorange=False),\n                aspectratio=dict(x=1, y=1, z=1),\n                camera=dict(eye=dict(x=1.25, y=-1.25, z=1.25)),\n            ),\n            width=800,\n            height=800,\n            template=\"plotly_dark\",\n            updatemenus=[\n                dict(\n                    type=\"buttons\",\n                    buttons=[\n                        dict(\n                            label=\"▶\",\n                            method=\"animate\",\n                            args=[\n                                None,\n                                dict(\n                                    frame=dict(duration=500, redraw=True),\n                                    fromcurrent=True,\n                                ),\n                            ],\n                        )\n                    ],\n                    font=dict(color=\"black\"),\n                )\n            ],\n        ),\n    )\n\n    # Set annimation frames\n    frames = []\n    for i in z_indices:\n        frame = go.Frame(data=base_traces + [_plot_slice_surface(arr, i)], name=str(i))\n        frames.append(frame)\n    fig.frames = frames\n\n    # Setup slider\n    sliders = [\n        dict(\n            active=0,\n            currentvalue=dict(prefix=\"z-index: \"),\n            pad=dict(t=50),\n            steps=[\n                dict(\n                    label=str(i),\n                    method=\"animate\",\n                    args=[\n                        [str(i)],\n                        dict(\n                            mode=\"immediate\",\n                            frame=dict(duration=200, redraw=True),\n                            transition=dict(duration=0),\n                        ),\n                    ],\n                )\n                for i in z_indices\n            ],\n        )\n    ]\n    fig.update_layout(sliders=sliders)\n    return fig","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-13T07:49:59.548402Z","iopub.execute_input":"2024-11-13T07:49:59.548823Z","iopub.status.idle":"2024-11-13T07:49:59.567043Z","shell.execute_reply.started":"2024-11-13T07:49:59.548779Z","shell.execute_reply":"2024-11-13T07:49:59.565355Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"run = runs[0]\ndf = read_run(run)\nfig = plot_particles(df, scale=0.1)\nzarr_group = read_zarr(run)\nplot_animatable_slices(zarr_group[0], title=df.run_name.iloc[0], fig=fig)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-13T07:49:59.568680Z","iopub.execute_input":"2024-11-13T07:49:59.569180Z","iopub.status.idle":"2024-11-13T07:50:05.118564Z","shell.execute_reply.started":"2024-11-13T07:49:59.569135Z","shell.execute_reply":"2024-11-13T07:50:05.117376Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"run = runs[1]\ndf = read_run(run)\nfig = plot_particles(df, scale=0.1)\nzarr_group = read_zarr(run)\nplot_animatable_slices(zarr_group[0], title=df.run_name.iloc[0], fig=fig)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-13T07:50:05.119963Z","iopub.execute_input":"2024-11-13T07:50:05.120424Z","iopub.status.idle":"2024-11-13T07:50:09.047056Z","shell.execute_reply.started":"2024-11-13T07:50:05.120380Z","shell.execute_reply":"2024-11-13T07:50:09.046047Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"run = runs[2]\ndf = read_run(run)\nfig = plot_particles(df, scale=0.1)\nzarr_group = read_zarr(run)\nplot_animatable_slices(zarr_group[0], title=df.run_name.iloc[0], fig=fig)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-13T07:50:09.048695Z","iopub.execute_input":"2024-11-13T07:50:09.049255Z","iopub.status.idle":"2024-11-13T07:50:14.554610Z","shell.execute_reply.started":"2024-11-13T07:50:09.049200Z","shell.execute_reply":"2024-11-13T07:50:14.553447Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"run = runs[3]\ndf = read_run(run)\nfig = plot_particles(df, scale=0.1)\nzarr_group = read_zarr(run)\nplot_animatable_slices(zarr_group[0], title=df.run_name.iloc[0], fig=fig)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-13T07:50:14.556274Z","iopub.execute_input":"2024-11-13T07:50:14.556739Z","iopub.status.idle":"2024-11-13T07:50:19.972235Z","shell.execute_reply.started":"2024-11-13T07:50:14.556695Z","shell.execute_reply":"2024-11-13T07:50:19.970923Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"run = runs[4]\ndf = read_run(run)\nfig = plot_particles(df, scale=0.1)\nzarr_group = read_zarr(run)\nplot_animatable_slices(zarr_group[0], title=df.run_name.iloc[0], fig=fig)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-13T07:50:19.973579Z","iopub.execute_input":"2024-11-13T07:50:19.973950Z","iopub.status.idle":"2024-11-13T07:50:25.534735Z","shell.execute_reply.started":"2024-11-13T07:50:19.973912Z","shell.execute_reply":"2024-11-13T07:50:25.533492Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"run = runs[5]\ndf = read_run(run)\nfig = plot_particles(df, scale=0.1)\nzarr_group = read_zarr(run)\nplot_animatable_slices(zarr_group[0], title=df.run_name.iloc[0], fig=fig)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-13T07:50:25.536519Z","iopub.execute_input":"2024-11-13T07:50:25.537084Z","iopub.status.idle":"2024-11-13T07:50:29.518217Z","shell.execute_reply.started":"2024-11-13T07:50:25.537023Z","shell.execute_reply":"2024-11-13T07:50:29.516925Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"run = runs[6]\ndf = read_run(run)\nfig = plot_particles(df, scale=0.1)\nzarr_group = read_zarr(run)\nplot_animatable_slices(zarr_group[0], title=df.run_name.iloc[0], fig=fig)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-13T07:50:29.520021Z","iopub.execute_input":"2024-11-13T07:50:29.520545Z","iopub.status.idle":"2024-11-13T07:50:33.583543Z","shell.execute_reply.started":"2024-11-13T07:50:29.520489Z","shell.execute_reply":"2024-11-13T07:50:33.582352Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Hope this notebook helps.","metadata":{"execution":{"iopub.status.busy":"2024-11-13T07:14:21.883427Z","iopub.execute_input":"2024-11-13T07:14:21.884095Z","iopub.status.idle":"2024-11-13T07:14:21.894144Z","shell.execute_reply.started":"2024-11-13T07:14:21.884040Z","shell.execute_reply":"2024-11-13T07:14:21.891994Z"}}}]}