{"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"},{"sourceId":9906449,"sourceType":"datasetVersion","datasetId":6086262},{"sourceId":166497,"sourceType":"modelInstanceVersion","modelInstanceId":141676,"modelId":164257}],"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!cp -r /kaggle/input/zarr-wheels/ /kaggle/working\n!pip install /kaggle/working/zarr-wheels/asciitree-0.3.3/asciitree-0.3.3/\n!pip install --no-index --find-links=/kaggle/working/zarr-wheels zarr","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-11-14T14:22:53.889468Z","iopub.execute_input":"2024-11-14T14:22:53.890413Z","iopub.status.idle":"2024-11-14T14:23:43.509070Z","shell.execute_reply.started":"2024-11-14T14:22:53.890311Z","shell.execute_reply":"2024-11-14T14:23:43.508052Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import zarr\nimport glob\nimport os\nimport numpy as np\nimport json\nfrom tensorflow.keras.layers import Input\nfrom tensorflow.keras.models import clone_model\nimport tensorflow as tf\nfrom tensorflow.keras.layers import Conv2DTranspose\nfrom scipy.ndimage import label, distance_transform_edt\nfrom skimage.feature import peak_local_max\nfrom skimage.segmentation import watershed\n\nparticle_sizes = {'apo-ferritin': 60.0, 'beta-amylase': 65.0, 'beta-galactosidase': 90.0, 'ribosome': 150.0,\n                  'thyroglobulin': 130.0, 'virus-like-particle': 135.0}\nparticle_types = list(particle_sizes.keys())\napix = 10.012444196428572\nDEGENERACY = 4  # 1 = apply model to volumes once; 2 = apply twice, to original and to 90 rotated version; 3 = 0, 90, 180 deg.; 4 = 0, 90, 180, 270; 5-8: as before, but with horizontal flips for each angle (currently not implemented - training data also preserves chirality.)\nTHRESHOLD = 0.5\nVOLUME_FRACTION = 0.2\nSPACING_RADIUS_MULTIPLIER = 1.0\n\n\nclass CustomConv2DTranspose(Conv2DTranspose):\n    def __init__(self, *args, **kwargs):\n        kwargs.pop('groups', None)  # To work around a keras version compatibility issue\n        super().__init__(*args, **kwargs)\n\n    @classmethod\n    def from_config(cls, config):\n        config.pop('groups', None)\n        return super().from_config(config)\n\n\ntrained_model = tf.keras.models.load_model(\"/kaggle/input/ais_vggnet_l/keras/v1/1/ModelCheckpoint.h5\", compile=False,\n                                           custom_objects={'Conv2DTranspose': CustomConv2DTranspose})\nnew_input = Input(shape=(None, None, 1))\ninference_model = clone_model(trained_model, input_tensors=new_input)\ninference_model.set_weights(trained_model.get_weights())\n\n\ndef segment(volume, model):\n    vol_out = np.zeros((*volume.shape, len(particle_types)), dtype=np.float32)\n    for j in range(vol_out.shape[0]):\n        print(f\"{j + 1}/{vol_out.shape[0]}\")\n        slice_j = volume[j, :, :]\n        for k in range(DEGENERACY):\n            rotated_slice = np.rot90(slice_j, k=k, axes=(0, 1))\n            rotated_segmented_slice = np.squeeze(model.predict(rotated_slice[np.newaxis, :, :], verbose=0))\n            segmented_slice = np.rot90(rotated_segmented_slice, k=-k, axes=(0, 1))\n            vol_out[j, :, :, :] += segmented_slice\n    vol_out /= DEGENERACY\n    return vol_out\n\n\nclass Blob:\n    def __init__(self):\n        self.x = list()\n        self.y = list()\n        self.z = list()\n        self.v = list()\n\n    def get_centroid(self, scale=1):\n        return np.mean(self.x) * scale, np.mean(self.y) * scale, np.mean(self.z) * scale\n\n    def get_center_of_mass(self):\n        mx = np.sum(np.array(self.x) * np.array(self.v))\n        my = np.sum(np.array(self.y) * np.array(self.v))\n        mz = np.sum(np.array(self.z) * np.array(self.v))\n        m = self.get_weight()\n        return mx / m, my / m, mz / m\n\n    def get_volume(self):\n        return len(self.x)\n\n    def get_weight(self):\n        return np.sum(self.v)\n\n\ndef pick_particles(volume, threshold=0.5, min_spacing=1, min_volume=0):\n    binary_volume = volume > threshold\n    distance = distance_transform_edt(binary_volume)\n    maxima = peak_local_max(distance, min_distance=min_spacing)\n    mask = np.zeros(distance.shape, dtype=bool)\n    mask[tuple(maxima.T)] = True\n    markers, _ = label(mask)\n    labels = watershed(-distance, markers, mask=binary_volume)\n\n    blobs = dict()\n    Z, Y, X = np.nonzero(labels)\n    for i in range(len(X)):\n        z = Z[i]\n        y = Y[i]\n        x = X[i]\n\n        l = labels[z, y, x]\n        if l not in blobs:\n            blobs[l] = Blob()\n\n        blobs[l].x.append(x)\n        blobs[l].y.append(y)\n        blobs[l].z.append(z)\n        blobs[l].v.append(volume[z, y, x])\n\n    blobs = [blobs[k] for k in blobs if blobs[k].get_volume() > min_volume]\n    coordinates = [b.get_centroid(scale=1) for b in blobs]\n\n    remove = list()\n    j = 0\n    while j < len(coordinates):\n        for k in range(0, j):\n            if j in remove:\n                continue\n            p = np.array(coordinates[j])\n            q = np.array(coordinates[k])\n            d = np.sum((p - q) ** 2) ** 0.5\n            if d < min_spacing:\n                remove.append(k)\n        j += 1\n    remove.sort()\n    for j in reversed(remove):\n        coordinates.pop(j)\n\n    return coordinates\n\n\ndatasets = [os.path.basename(f) for f in\n            glob.glob(\"/kaggle/input/czii-cryo-et-object-identification/test/static/ExperimentRuns/*\")]\n\nTHRESHOLD = 0.5\n\na3topx3 = (10/apix)**3\nmin_volumes = {'apo-ferritin': 200.0 * a3topx3,\n               'beta-amylase': 250.0 * a3topx3,\n               'beta-galactosidase': 350.0 * a3topx3,\n               'ribosome': 1430.0 * a3topx3,\n               'thyroglobulin': 120.0 * a3topx3,\n               'virus-like-particle': 1430.0 * a3topx3}\n\nSPACING_RADIUS_MULTIPLIER = 1.0\n\ncsv_lines = list()\ncsv_lines.append(\"id,experiment,particle_type,x,y,z\\n\")\n\nparticle_id = 0\nfor tag in datasets:\n    print(f\"Processing {tag}\")\n    volume_path = f\"/kaggle/input/czii-cryo-et-object-identification/test/static/ExperimentRuns/{tag}/VoxelSpacing10.000/denoised.zarr\"\n    if not os.path.exists(volume_path):\n        continue\n    volume = np.array(zarr.open(volume_path, mode='r')['0']).astype(np.float32)\n    volume -= np.mean(volume)\n    volume /= np.std(volume)\n    volume = volume[:, :624, :624]\n\n    segmentation = segment(volume, inference_model)\n\n    for j, p in enumerate(particle_types):\n        print(f\"Picking {p}\")\n        s = segmentation[:, :, :, j]\n\n        coordinates = pick_particles(s,\n                                     threshold=THRESHOLD,\n                                     min_spacing=int((particle_sizes[p] / apix) * SPACING_RADIUS_MULTIPLIER),\n                                     min_volume=min_volumes[p])\n        print(f\"\\t found {len(coordinates)} particles\")\n        for c in coordinates:\n            csv_lines.append(f\"{particle_id},{tag},{p},{c[0] * apix},{c[1] * apix},{c[2] * apix}\\n\")\n            particle_id += 1\n\nwith open(\"submission.csv\", 'w') as f:\n    f.writelines(csv_lines)","metadata":{"execution":{"iopub.status.busy":"2024-11-14T13:44:15.316116Z","iopub.execute_input":"2024-11-14T13:44:15.316493Z","iopub.status.idle":"2024-11-14T13:59:24.529383Z","shell.execute_reply.started":"2024-11-14T13:44:15.316449Z","shell.execute_reply":"2024-11-14T13:59:24.528441Z"},"trusted":true},"execution_count":null,"outputs":[]}]}