{"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":"nvidiaTeslaT4","dataSources":[{"sourceId":84969,"databundleVersionId":10033515,"sourceType":"competition"}],"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# !pip install copick git+https://github.com/copick/copick-utils.git scikit-image cupy-cuda12x torch torchvision tqdm matplotlib","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2024-11-11T11:49:15.92033Z","iopub.execute_input":"2024-11-11T11:49:15.920985Z","iopub.status.idle":"2024-11-11T11:49:15.927779Z","shell.execute_reply.started":"2024-11-11T11:49:15.920948Z","shell.execute_reply":"2024-11-11T11:49:15.926904Z"}},"outputs":[],"execution_count":4},{"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\": \"/kaggle/working/test/overlay\",\n\n    \"overlay_fs_args\": {\n        \"auto_mkdir\": true\n    },\n\n    \"static_root\": \"/kaggle/input/czii-cryo-et-object-identification/test/static\"\n}\"\"\"\n\ncopick_config_path = \"/kaggle/working/copick.config\"\n\nwith open(copick_config_path, \"w\") as f:\n    f.write(config_blob)\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-11T11:49:15.929952Z","iopub.execute_input":"2024-11-11T11:49:15.930368Z","iopub.status.idle":"2024-11-11T11:49:15.937183Z","shell.execute_reply.started":"2024-11-11T11:49:15.930325Z","shell.execute_reply":"2024-11-11T11:49:15.936307Z"}},"outputs":[],"execution_count":5},{"cell_type":"code","source":"import matplotlib.pyplot as plt\nimport torch\nimport torch.nn.functional as F\nimport numpy as np\nfrom skimage.measure import regionprops\nfrom skimage.morphology import ball\nfrom skimage.segmentation import watershed\nfrom tqdm import tqdm\nimport scipy.ndimage as ndi\nimport time\nimport csv\nimport os\nimport copick\nimport zarr\n\nDEVICE = 'cuda'\nOUTPUT_CSV_PATH = \"submission.csv\"\ntomo_type = 'denoised'\nRESOLUTION_THRESHOLD = 16\n\ndef gaussian_kernel(size, sigma):\n    \"\"\"Generate a 3D Gaussian kernel.\"\"\"\n    kernel = np.fromfunction(\n        lambda x, y, z: (1/ (2 * np.pi * sigma**2)) * \n        np.exp(- ((x - (size[0] - 1) / 2) ** 2 + \n                   (y - (size[1] - 1) / 2) ** 2 + \n                   (z - (size[2] - 1) / 2) ** 2) / (2 * sigma ** 2)),\n        size\n    )\n    return torch.tensor(kernel).float().unsqueeze(0).unsqueeze(0).to(DEVICE)  # Add batch and channel dimensions\n\ndef create_hessian_particle_mask(tomogram, sigma):\n    \"\"\"\n    Generate a binary mask for dark, blob-like particles in a cryo-ET tomogram\n    using Hessian-based filtering with PyTorch.\n\n    Args:\n        tomogram (torch.Tensor): The input 3D tomogram (C, D, H, W).\n        sigma (float): The standard deviation for Gaussian smoothing.\n\n    Returns:\n        torch.Tensor: Binary mask highlighting dark blob-like areas in the tomogram.\n    \"\"\"\n    kernel_size = (5, 5, 5)\n    gaussian_k = gaussian_kernel(kernel_size, sigma)\n    \n    tomogram_smoothed = F.conv3d(tomogram.unsqueeze(0).unsqueeze(0), gaussian_k, padding=2).squeeze()\n\n    # Compute Hessian components\n    hessian_xx = F.conv3d(tomogram_smoothed.unsqueeze(0).unsqueeze(0), gaussian_k, padding=2)\n    hessian_yy = F.conv3d(tomogram_smoothed.unsqueeze(0).unsqueeze(0), gaussian_k, padding=2)\n    hessian_xy = F.conv3d(tomogram_smoothed.unsqueeze(0).unsqueeze(0), gaussian_k, padding=2)\n\n    hessian_response = hessian_xx + hessian_yy + hessian_xy  # Simplified combination\n    binary_mask = hessian_response < 0  # Adjust threshold based on your needs\n\n    return binary_mask.squeeze().byte()\n\ndef erode_dilate_mask(mask, radius):\n    \"\"\"\n    Perform binary erosion and dilation on a binary mask using a spherical structuring element.\n    \n    Args:\n        mask (torch.Tensor): Input binary mask\n        radius (int): Radius of the spherical structuring element\n        \n    Returns:\n        torch.Tensor: Dilated mask after erosion and dilation operations\n    \"\"\"\n    # Create a spherical structuring element\n    radius = int(radius)  # Ensure radius is an integer\n    struct_elem = ball(radius)\n    struct_elem_tensor = torch.tensor(struct_elem, dtype=torch.float32, device=DEVICE).unsqueeze(0).unsqueeze(0)\n\n    # Reshape mask for conv3d\n    mask_reshaped = mask.unsqueeze(0).unsqueeze(0).float()  # Shape (1, 1, D, H, W)\n    \n    # Calculate padding size - ensure it's an integer\n    pad_size = int(radius // 2)\n    \n    # Debug: Print shapes\n    print(f\"Mask shape for erosion: {mask_reshaped.shape}\")\n    print(f\"Structuring element shape: {struct_elem_tensor.shape}\")\n    print(f\"Padding size: {pad_size}\")\n    \n    # Erosion: Use a negative structuring element for max pooling\n    # Convert padding to the expected format (left, right, top, bottom, front, back)\n    # Ensure all values are integers\n    pad_3d = (int(pad_size), int(pad_size), \n              int(pad_size), int(pad_size), \n              int(pad_size), int(pad_size))\n    \n    mask_padded = F.pad(mask_reshaped, pad_3d, mode='constant', value=1)\n    eroded = -F.conv3d(\n        -mask_padded,\n        struct_elem_tensor,\n        stride=1,\n        padding=0,\n        dilation=1,\n        groups=1\n    )\n    eroded = (eroded >= struct_elem_tensor.sum()).squeeze().byte()\n\n    # Dilation\n    mask_padded = F.pad(eroded.unsqueeze(0).unsqueeze(0).float(), pad_3d, mode='constant', value=0)\n    dilated = F.conv3d(\n        mask_padded,\n        struct_elem_tensor,\n        stride=1,\n        padding=0,\n        dilation=1,\n        groups=1\n    )\n    dilated = (dilated > 0).squeeze().byte()\n    \n    return dilated\n\ndef distance_transform(mask):\n    \"\"\"\n    Compute the distance transform using a simple distance transform approach.\n    \n    Args:\n        mask (torch.Tensor): Binary mask tensor\n        \n    Returns:\n        torch.Tensor: Distance transform result\n    \"\"\"\n    # Ensure mask is boolean, then convert to float for distance calculation\n    mask = mask.bool()\n    # Invert the mask (using logical not instead of bitwise not)\n    inverted_mask = (~mask).float()\n    \n    # Add batch and channel dimensions\n    inverted_mask = inverted_mask.unsqueeze(0).unsqueeze(0)\n    \n    # Create kernel on the correct device\n    kernel = torch.ones(1, 1, 3, 3, 3, device=DEVICE)\n    \n    # Compute distance transform using convolution\n    distance = F.conv3d(inverted_mask, kernel, padding=1)\n    \n    return distance.squeeze()\n\ndef local_maxima(distance, radius):\n    \"\"\"\n    Detect local maxima in the distance transform.\n    \n    Args:\n        distance (torch.Tensor): Distance transform tensor\n        radius (int): Radius for local maxima detection\n        \n    Returns:\n        torch.Tensor: Binary mask of local maxima\n    \"\"\"\n    # Ensure radius is an integer\n    radius = int(radius)\n    \n    # Add batch dimension for max_pool3d\n    distance = distance.unsqueeze(0)\n    \n    # Create kernel size tuple (must be odd numbers)\n    kernel_size = (2 * radius + 1, 2 * radius + 1, 2 * radius + 1)\n    \n    # Compute local maxima\n    maxpool = F.max_pool3d(\n        distance,\n        kernel_size=kernel_size,\n        stride=1,\n        padding=radius\n    )\n    \n    # Compare with original distance to find local maxima\n    local_max = (distance == maxpool)\n    \n    return local_max.squeeze()\n\ndef get_tomogram_data(run, voxel_spacing, radius):\n    \"\"\"\n    Get tomogram data at appropriate resolution based on particle radius.\n    \n    Args:\n        run: Run object\n        voxel_spacing (float): Base voxel spacing\n        radius (float): Particle radius\n        \n    Returns:\n        tuple: (tomogram tensor, effective_voxel_spacing, scale_factor)\n    \"\"\"\n    tomogram_wrapper = run.get_voxel_spacing(voxel_spacing).get_tomogram(tomo_type)\n    z = zarr.open(store=tomogram_wrapper.zarr(), path=\"/\", mode=\"r\")\n    \n    if radius <= RESOLUTION_THRESHOLD:\n        # Use highest resolution\n        tomogram = z['0'][:]\n        effective_voxel_spacing = voxel_spacing\n        scale_factor = 1\n    else:\n        # Use medium resolution\n        tomogram = z['1'][:]\n        effective_voxel_spacing = voxel_spacing * 2  # Scale factor is 2 for level 1\n        scale_factor = 2\n        \n    return torch.tensor(tomogram).to(DEVICE), effective_voxel_spacing, scale_factor\n\ndef process_all_runs(root, session_id, user_id, voxel_spacing):\n    \"\"\"Process all runs and save results to CSV.\"\"\"\n    results = []\n    pick_id = 0\n    \n    for run in tqdm(root.runs):\n        start_time = time.time()\n        print(f\"\\nProcessing run: {run.meta.name}\")\n        \n        # Process each particle type separately since they might need different resolutions\n        for obj in root.pickable_objects:\n            if not obj.is_particle:\n                continue\n                \n            radius = obj.radius\n            print(f\"Processing {obj.name} with radius {radius}\")\n            \n            # Get appropriate resolution data\n            tomogram_tensor, effective_voxel_spacing, scale_factor = get_tomogram_data(\n                run, voxel_spacing, radius)\n\n            print(f\"Using scale factor {scale_factor} (effective voxel spacing: {effective_voxel_spacing})\")\n\n            # Create segmentation at appropriate scale\n            segmentation = create_hessian_particle_mask(tomogram_tensor, sigma=3)\n            \n            if torch.sum(segmentation) == 0:\n                print(f\"No particles detected in segmentation for {obj.name}\")\n                continue\n\n            # Adjust radius for effective voxel spacing\n            scaled_radius = radius / effective_voxel_spacing\n\n            # Erode and dilate the segmentation\n            dilated_mask = erode_dilate_mask(segmentation, scaled_radius)\n\n            # Distance transform and local maxima detection\n            distance = distance_transform(dilated_mask)\n            local_max = local_maxima(distance, scaled_radius)\n\n            # Convert tensors to numpy for watershed\n            local_max_np = local_max.cpu().numpy()\n            distance_np = distance.cpu().numpy()\n            dilated_mask_np = dilated_mask.cpu().numpy()\n\n            # Watershed segmentation\n            markers, _ = ndi.label(local_max_np)\n            watershed_labels = watershed(-distance_np, markers, mask=dilated_mask_np)\n\n            # Extract region properties and scale coordinates back to original space\n            centroids = []\n            for region in regionprops(watershed_labels):\n                # Scale the centroid coordinates back to original space\n                centroid = np.array(region.centroid) * scale_factor\n                centroids.append(centroid)  # ZYX order\n\n            # Save centroids as picks and add to results\n            if centroids:\n                pick_set = run.get_picks(obj.name)\n                if pick_set:\n                    pick_set = pick_set[0]\n                else:\n                    pick_set = run.new_picks(obj.name, session_id, user_id)\n                \n                for centroid in centroids:\n                    # Convert from ZYX to XYZ order and apply voxel spacing\n                    x = centroid[2] * voxel_spacing  # Z -> X\n                    y = centroid[1] * voxel_spacing  # Y -> Y\n                    z = centroid[0] * voxel_spacing  # X -> Z\n                    \n                    # Add to results list\n                    row = [pick_id, run.meta.name, obj.name, x, y, z]\n                    results.append(row)\n                    pick_id += 1\n                    \n                # Store pick set\n                pick_set.points = [{'x': c[2] * voxel_spacing,\n                                  'y': c[1] * voxel_spacing,\n                                  'z': c[0] * voxel_spacing}\n                                 for c in centroids]\n                pick_set.store()\n                print(f\"Saved {len(centroids)} centroids for {obj.name}\")\n            else:\n                print(f\"No valid centroids found for {obj.name}\")\n\n        # Print timing for this run\n        end_time = time.time()\n        print(f\"Run {run.meta.name} completed in {end_time - start_time:.2f} seconds\")\n\n    print(f\"\\nTotal picks found: {len(results)}\")\n\n    # Write results to CSV\n    with open(OUTPUT_CSV_PATH, mode='w', newline='') as file:\n        writer = csv.writer(file)\n        writer.writerow([\"id\", \"experiment\", \"particle_type\", \"x\", \"y\", \"z\"])\n        writer.writerows(results)\n    \n    print(f\"Results saved to {OUTPUT_CSV_PATH}\")\n    return results\n\n# Run the processing\nroot = copick.from_file(copick_config_path)\nresults = process_all_runs(\n    root=root,\n    session_id=\"0\",\n    user_id=\"blobDetector\",\n    voxel_spacing=10\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-11T11:49:16.027756Z","iopub.execute_input":"2024-11-11T11:49:16.02807Z","iopub.status.idle":"2024-11-11T11:49:16.13498Z","shell.execute_reply.started":"2024-11-11T11:49:16.028037Z","shell.execute_reply":"2024-11-11T11:49:16.133725Z"}},"outputs":[{"traceback":["\u001b[0;31m---------------------------------------------------------------------------\u001b[0m","\u001b[0;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)","Cell \u001b[0;32mIn[6], line 13\u001b[0m\n\u001b[1;32m     11\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mcsv\u001b[39;00m\n\u001b[1;32m     12\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mos\u001b[39;00m\n\u001b[0;32m---> 13\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mcopick\u001b[39;00m\n\u001b[1;32m     14\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mzarr\u001b[39;00m\n\u001b[1;32m     16\u001b[0m DEVICE \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mcuda\u001b[39m\u001b[38;5;124m'\u001b[39m\n","\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'copick'"],"ename":"ModuleNotFoundError","evalue":"No module named 'copick'","output_type":"error"}],"execution_count":6},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}