{"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 copick ipywidgets","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-15T16:03:07.079588Z","iopub.execute_input":"2024-11-15T16:03:07.080071Z","iopub.status.idle":"2024-11-15T16:03:22.349495Z","shell.execute_reply.started":"2024-11-15T16:03:07.080015Z","shell.execute_reply":"2024-11-15T16:03:22.348243Z"},"_kg_hide-input":false},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport glob\nimport json\nimport numpy as np\n\nimport copick\nfrom copick.impl.filesystem import CopickRootFSSpec\nimport zarr\n\nimport ipywidgets as widgets\nimport matplotlib.pyplot as plt\nfrom IPython.display import display, clear_output\nimport skimage","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-15T16:02:06.858071Z","iopub.execute_input":"2024-11-15T16:02:06.858496Z","iopub.status.idle":"2024-11-15T16:02:06.865430Z","shell.execute_reply.started":"2024-11-15T16:02:06.858454Z","shell.execute_reply":"2024-11-15T16:02:06.864233Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"data_path='/kaggle/input/czii-cryo-et-object-identification/'\ntrain_path=data_path+'train/'\ntest_path=data_path+'test/'\nexperiment_runs=glob.glob(train_path+'static/ExperimentRuns/**')\nexperiment_runs=[os.path.split(e)[1] for e in experiment_runs]\nexperiment_runs","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-15T16:02:06.893607Z","iopub.execute_input":"2024-11-15T16:02:06.894040Z","iopub.status.idle":"2024-11-15T16:02:06.906838Z","shell.execute_reply.started":"2024-11-15T16:02:06.893971Z","shell.execute_reply":"2024-11-15T16:02:06.905890Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"config_blob = \"\"\"{\n    \"name\": \"czii_cryoet_mlchallenge_2024\",\n    \"description\": \"2024 CZII CryoET ML Challenge training data.\",\n    \"version\": \"1.0.0\",\n\n    \"pickable_objects\": [\n        {\n            \"name\": \"apo-ferritin\",\n            \"is_particle\": true,\n            \"pdb_id\": \"4V1W\",\n            \"label\": 1,\n            \"color\": [  0, 117, 220, 128],\n            \"radius\": 60,\n            \"map_threshold\": 0.0418\n        },\n        {\n            \"name\": \"beta-amylase\",\n            \"is_particle\": true,\n            \"pdb_id\": \"1FA2\",\n            \"label\": 2,\n            \"color\": [153,  63,   0, 128],\n            \"radius\": 65,\n            \"map_threshold\": 0.035\n        },\n        {\n            \"name\": \"beta-galactosidase\",\n            \"is_particle\": true,\n            \"pdb_id\": \"6X1Q\",\n            \"label\": 3,\n            \"color\": [ 76,   0,  92, 128],\n            \"radius\": 90,\n            \"map_threshold\": 0.0578\n        },\n        {\n            \"name\": \"ribosome\",\n            \"is_particle\": true,\n            \"pdb_id\": \"6EK0\",\n            \"label\": 4,\n            \"color\": [  0,  92,  49, 128],\n            \"radius\": 150,\n            \"map_threshold\": 0.0374\n        },\n        {\n            \"name\": \"thyroglobulin\",\n            \"is_particle\": true,\n            \"pdb_id\": \"6SCJ\",\n            \"label\": 5,\n            \"color\": [ 43, 206,  72, 128],\n            \"radius\": 130,\n            \"map_threshold\": 0.0278\n        },\n        {\n            \"name\": \"virus-like-particle\",\n            \"is_particle\": true,\n            \"pdb_id\": \"6N4V\",            \n            \"label\": 6,\n            \"color\": [255, 204, 153, 128],\n            \"radius\": 135,\n            \"map_threshold\": 0.201\n        }\n    ],\n\n    \"overlay_root\": \"XXXXXXXXXXtrain/overlay\",\n\n    \"overlay_fs_args\": {\n        \"auto_mkdir\": true\n    },\n\n    \"static_root\": \"XXXXXXXXXXtrain/static/\"\n}\"\"\"\n\nconfig_blob=config_blob.replace('XXXXXXXXXX',data_path)\ncopick_config=json.loads(config_blob)\n\ncopick_config_path = \"copick.config\"\nwith open(copick_config_path, \"w\") as f:\n    f.write(config_blob)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-15T16:02:06.908788Z","iopub.execute_input":"2024-11-15T16:02:06.909612Z","iopub.status.idle":"2024-11-15T16:02:06.918611Z","shell.execute_reply.started":"2024-11-15T16:02:06.909557Z","shell.execute_reply":"2024-11-15T16:02:06.917451Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"root = CopickRootFSSpec.from_file(copick_config_path)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-15T16:02:06.919905Z","iopub.execute_input":"2024-11-15T16:02:06.920300Z","iopub.status.idle":"2024-11-15T16:02:06.935298Z","shell.execute_reply.started":"2024-11-15T16:02:06.920254Z","shell.execute_reply":"2024-11-15T16:02:06.934265Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import json\n\ndef get_particles(run,scale=10):\n    path=train_path+'overlay/ExperimentRuns/'+run+'Picks/'\n    path=train_path+'/overlay/ExperimentRuns/'+run+'/Picks/'\n    json_files=glob.glob(path+'*.json')\n    particles={}\n    particles_dict={}\n    for j in json_files:\n        with open(j,'rt') as jf:\n            data=json.load(jf)\n            for j in data:\n                particles_dict[data['pickable_object_name']]=np.array([[c['location']['z']/scale, c['location']['x']/scale,c['location']['y']/scale] for c in data['points']])\n            particles[run]=particles_dict\n    return particles\n\ndef get_points(run):\n    path=train_path+'overlay/ExperimentRuns/'+run+'Picks/'\n    path=train_path+'/overlay/ExperimentRuns/'+run+'/Picks/'\n    json_files=glob.glob(path+'*.json')\n    particles={}\n    particles_dict={}\n    for j in json_files:\n        with open(j,'rt') as jf:\n            data=json.load(jf)\n            for j in data:\n                particles_dict[data['pickable_object_name']]=data['points']\n            particles[run]=particles_dict\n    return particles\n\ndef get_radius(particle_type, scale=1):\n    for po in copick_config['pickable_objects']:\n        if po['name']==particle_type:\n            return po['radius']/scale\n    return None","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-15T16:02:06.936841Z","iopub.execute_input":"2024-11-15T16:02:06.937224Z","iopub.status.idle":"2024-11-15T16:02:06.950114Z","shell.execute_reply.started":"2024-11-15T16:02:06.937167Z","shell.execute_reply":"2024-11-15T16:02:06.949032Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#Annexed from copick source\n#Creates 3D spherical annotations based on point data\ndef from_picks(points, \n               seg_volume,\n               radius: float = 10.0, \n               label_value: int = 1,\n               voxel_spacing: float = 10):\n    \"\"\"\n    Paints picks into a segmentation volume as spheres.\n\n    Parameters:\n    -----------\n    pick : copick.models.CopickPicks\n        Copick object containing `points`, where each point has a `location` attribute with `x`, `y`, `z` coordinates.\n    seg_volume : numpy.ndarray\n        3D segmentation volume (numpy array) where the spheres are painted. Shape should be (Z, Y, X).\n    radius : float, optional\n        The radius of the spheres to be inserted in physical units (not voxel units). Default is 10.0.\n    label_value : int, optional\n        The integer value used to label the sphere regions in the segmentation volume. Default is 1.\n    voxel_spacing : float, optional\n        The spacing of voxels in the segmentation volume, used to scale the radius of the spheres. Default is 10.\n\n    Returns:\n    --------\n    numpy.ndarray\n        The modified segmentation volume with spheres inserted at pick locations.\n    \"\"\"\n        \n    def create_sphere(shape, center, radius, val):\n        \"\"\"Creates a 3D sphere within the given shape, centered at the given coordinates.\"\"\"\n        zc, yc, xc = center\n        z, y, x = np.indices(shape)\n        \n        # Compute squared distance from the center\n        distance_sq = (x - xc)**2 + (y - yc)**2 + (z - zc)**2\n        \n        # Create a mask for points within the sphere\n        sphere = np.zeros(shape, dtype=np.float32)\n        sphere[distance_sq <= radius**2] = val\n        return sphere\n\n    def get_relative_target_coordinates(center, delta, shape):\n        \"\"\"\n        Calculate the low and high index bounds for placing a sphere within a 3D volume, \n        ensuring that the indices are clamped to the valid range of the volume dimensions.\n        \"\"\"\n\n        low = max(int(np.floor(center) - delta), 0)\n        high = min(int(np.ceil(center) + delta + 1), shape)\n\n        return low, high\n\n    # Adjust radius for voxel spacing\n    radius_voxel = radius / voxel_spacing\n    delta = int(np.ceil(radius_voxel))\n\n    # Get volume dimensions\n    vol_shape_x, vol_shape_y, vol_shape_z = seg_volume.shape\n\n    # Paint each pick as a sphere\n    for pick in points:\n        \n        # Adjust the pick's location for voxel spacing\n        cx, cy, cz = pick['location']['z'] / voxel_spacing, pick['location']['y'] / voxel_spacing, pick['location']['x'] / voxel_spacing\n\n        # Calculate subarray bounds, clamped to the valid volume dimensions\n        xLow, xHigh = get_relative_target_coordinates(cx, delta, vol_shape_x)\n        yLow, yHigh = get_relative_target_coordinates(cy, delta, vol_shape_y)\n        zLow, zHigh = get_relative_target_coordinates(cz, delta, vol_shape_z)\n\n        # Subarray shape\n        subarray_shape = (xHigh - xLow, yHigh - yLow, zHigh - zLow)\n\n        # Compute the local center of the sphere within the subarray\n        local_center = (cx - xLow, cy - yLow, cz - zLow)\n\n        # Create the sphere\n        sphere = create_sphere(subarray_shape, local_center, radius_voxel, label_value)\n\n        # Assign Sphere to Segmentation Target Volume\n        seg_volume[xLow:xHigh, yLow:yHigh, zLow:zHigh] = np.maximum(seg_volume[xLow:xHigh, yLow:yHigh, zLow:zHigh], sphere)\n\n    return seg_volume","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-15T16:02:06.951893Z","iopub.execute_input":"2024-11-15T16:02:06.952323Z","iopub.status.idle":"2024-11-15T16:02:06.971134Z","shell.execute_reply.started":"2024-11-15T16:02:06.952282Z","shell.execute_reply":"2024-11-15T16:02:06.969750Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def get_hex_color(c):#utility function to get hexadecimal HTML color code\n    return f'#{c[0]:02x}{c[1]:02x}{c[2]:02x}'","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-15T16:02:06.972979Z","iopub.execute_input":"2024-11-15T16:02:06.973484Z","iopub.status.idle":"2024-11-15T16:02:06.988195Z","shell.execute_reply.started":"2024-11-15T16:02:06.973429Z","shell.execute_reply":"2024-11-15T16:02:06.987063Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"axis='xy'\n\ndef process_run(): \n    #extract a tomogram and points from the selected expiriment run \n    global volume, particles, points, run_str, run\n    with output:\n        run_str=runs_dropdown.value\n        run = root.get_run(run_str)\n        tomogram = run.get_voxel_spacing(10).get_tomogram(\"denoised\")\n        group = zarr.open(tomogram.zarr())\n        arrays = list(group.arrays())\n        _, volume = arrays[0]\n        particles=get_particles(run_str)\n        points=get_points(run_str)\n\ndef on_run_change(change):\n    #called when expiriment run is changed\n    global axis\n    with output:\n        process_run()\n        process_volume_and_mask()\n        reset_widgets()\n        update()\n\ndef process_volume_and_mask():\n    #create volume from tomogram data, do some minimal processing\n    #create 3d mask data from points\n    global mask, volume, vol\n    with output:#necessary for debugging or else ipwidgets supress all exceptions\n        vol = np.clip(volume, np.percentile(volume, 0.5), np.percentile(volume, 99.5))\n        vol = (vol - vol.min()) / (vol.max() - vol.min())\n        \n        mask=np.zeros(vol.shape+(3,), dtype=np.uint8)\n        \n        for p in copick_config['pickable_objects']:#create colored mask\n            particle_type=p['name']\n            color=np.array(p['color'],dtype=np.uint8)[0:3].reshape(1,1,1,3)\n            m=np.zeros_like(volume, dtype=np.uint8)\n            m=from_picks(points=points[run_str][particle_type], \n                           seg_volume=m,\n                           radius=get_radius(particle_type), \n                           label_value = 1,\n                           voxel_spacing = 10)\n            m=np.tile(m[:,:,:,None], (1,1,1,3))\n            mask+=m*color\n    \n        #rotated views\n        if axis=='xz':\n            print(vol.shape, mask.shape)\n            vol=np.transpose(np.swapaxes(vol, 0,2),(0,2,1))\n            mask=np.transpose(np.swapaxes(mask, 0,2),(0,2,1,3))\n            print(vol.shape, mask.shape)\n    \n        if axis=='zy':\n            vol=np.swapaxes(vol, 0,1)\n            mask=np.swapaxes(mask, 0,1)\n\ndef process_image():\n    with output:\n        n=slice_slider.value\n        image=vol[n]\n        image=image[:,:,None]\n        image=np.tile(image,(1,1,3))\n        \n        a=alpha_slider.value\n        msk=mask[n]/255.0\n        image=image*(msk>0)*(1-a)+msk*a+image*(msk==0)#needed to overlay colorful mask over BW image without clipping\n        #image=image.clip(0,1)\n    return image\n    \ndef update():#this function is called when slice number is changed\n    global image\n    with output:\n        image=process_image()\n        redraw_plot()\n\ndef reset_widgets():#this function is called when slice slider max value needs to be adjusted\n    slice_slider.max=len(vol)-1\n\n\ndef on_value_change(change):#This is the ipwidgets callback\n    update()\n\ndef prev_slide(instance):\n    slice_slider.value=max(slice_slider.value-1,0)\n    update()\n\ndef next_slide(instance):\n    slice_slider.value=min(slice_slider.value+1,slice_slider.max-1)\n    update()\n\ndef on_view_change(change):\n    global axis\n    with output:\n        axis=view_radiobutton.value\n        process_volume_and_mask()\n        reset_widgets()\n        update()\n\ndef redraw_plot():#redraw the image\n    global fig, imshow, image\n    with output:\n        clear_output(True)\n        fig=plt.figure(figsize=(10,10))\n        imshow=plt.imshow(image)\n        plt.xticks([]), plt.yticks([])\n        plt.axis(\"off\")\n        plt.subplots_adjust(hspace=0, wspace=0)\n        plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-15T16:02:06.991909Z","iopub.execute_input":"2024-11-15T16:02:06.992389Z","iopub.status.idle":"2024-11-15T16:02:07.018438Z","shell.execute_reply.started":"2024-11-15T16:02:06.992331Z","shell.execute_reply":"2024-11-15T16:02:07.017417Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#slice number\nslice_slider=widgets.IntSlider(\n    value=0,\n    min=0,\n    max=1,\n    step=1,\n    disabled=False,\n    continuous_update=False,\n    orientation='horizontal',\n    readout=True,\n    readout_format='d')\nslice_slider.observe(on_value_change, names='value')\n\nbuttonPrev = widgets.Button(\n    description='Previous',\n    disabled=False,\n    button_style='', \n    tooltip='Previous',)\nbuttonPrev.on_click(prev_slide)\n\nbuttonNext = widgets.Button(\n    description='Next',\n    disabled=False,\n    button_style='', \n    tooltip='Next',)\nbuttonNext.on_click(next_slide)\n\n#mask transparency\nalpha_slider=widgets.FloatSlider(\n    value=0.5,\n    min=0.0,\n    max=1.0,\n    step=0.05,\n    description='Transparency',\n    disabled=False,\n    continuous_update=False,\n    orientation='horizontal',\n    readout=True,\n    readout_format='.2f',)\nalpha_slider.observe(on_value_change, names='value')\n\n#Horizontal radiobuttons requires some mad html slillz\n_style = widgets.HTML(\n    \"<style>.widget-radio-box {flex-direction: row !important;}.widget-radio-box\"\n    \" label{margin:5px !important;width: 120px !important;}</style>\",\n    layout=widgets.Layout(display=\"none\"),)\nview_radiobutton=widgets.RadioButtons(\n    options=['xy', 'xz', 'zy'],\n    description='View:',\n    disabled=False,)\nh_radiobutton=widgets.HBox([view_radiobutton, _style])\nview_radiobutton.observe(on_view_change, names='value')\n\n#labels to show color codings for particles\nlabels=[]\nfor (n,c) in [(chr(0x2B24)+p['name'],get_hex_color(p['color'][0:3])) for p in copick_config['pickable_objects']]:\n    l=widgets.Label(value=n, layout=widgets.Layout(height='50%'))\n    l.style.text_color=c\n    l.style.height=1\n    labels.append(l)\n\nlabels=widgets.HBox(labels)\n\n#to select experiment run\nruns_dropdown=widgets.Dropdown(options=[(er,er) for er in experiment_runs],description='Run:',)\nruns_dropdown.observe(on_run_change, names='value')\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-15T16:02:07.020049Z","iopub.execute_input":"2024-11-15T16:02:07.020425Z","iopub.status.idle":"2024-11-15T16:02:07.091199Z","shell.execute_reply.started":"2024-11-15T16:02:07.020386Z","shell.execute_reply":"2024-11-15T16:02:07.089753Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"output = widgets.Output()\nprocess_run()\nprocess_volume_and_mask()\nimage=process_image()\nupdate()\nreset_widgets()\n\nHBox1=widgets.HBox([buttonPrev, slice_slider,buttonNext])\nHBox2=widgets.HBox([h_radiobutton, alpha_slider])\nBox=widgets.VBox([runs_dropdown, HBox2, labels,HBox1])\ndisplay(Box, output)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-15T16:02:07.092702Z","iopub.execute_input":"2024-11-15T16:02:07.093111Z","iopub.status.idle":"2024-11-15T16:02:27.213764Z","shell.execute_reply.started":"2024-11-15T16:02:07.093069Z","shell.execute_reply":"2024-11-15T16:02:27.212617Z"}},"outputs":[],"execution_count":null}]}