{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"submission_mode=False # ommit preliminary data analysis, just run the model for submission\n\n# ensure reproducibility\nseed = 184618484\nimport random\nimport numpy as np\nimport numpy.random as rd\nimport tensorflow as tf\nrandom.seed(seed)\nnp.random.seed(seed)\nrd.seed(seed)\ntf.random.set_seed(seed)","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:36:45.373526Z","iopub.execute_input":"2022-01-17T12:36:45.374302Z","iopub.status.idle":"2022-01-17T12:36:49.233594Z","shell.execute_reply.started":"2022-01-17T12:36:45.374179Z","shell.execute_reply":"2022-01-17T12:36:49.232829Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# copy saved models -- turn off for submissions!\n#! cp -r /kaggle/input/2021-sartorius-neuronsegmentation/*.model /kaggle/input/2021-sartorius-neuronsegmentation/*.pkl /kaggle/working\n#! ls\n# adjust for submissions\nepochs_to_train = 1024\nearly_stopping_patience = 16\nsimulation_epochs = 8\nsimulation_dt = 4.0","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:44:49.607822Z","iopub.execute_input":"2022-01-17T12:44:49.608083Z","iopub.status.idle":"2022-01-17T12:44:49.612427Z","shell.execute_reply.started":"2022-01-17T12:44:49.608055Z","shell.execute_reply":"2022-01-17T12:44:49.611459Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nfrom scipy import signal\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\nimport pickle as pkl\nfrom tqdm import tqdm\nfrom copy import copy\n\nfrom io import StringIO\n\nimport matplotlib.pyplot as plt\nimport matplotlib.gridspec as gridspec\nfrom matplotlib import cm\nfrom matplotlib.patches import Rectangle\n\nfrom PIL import Image\n\nimport os\nimport glob\nprint(glob.glob('/kaggle/input/sartorius-cell-instance-segmentation/*'))\n\nimport tensorflow_addons as tfa\nimport tensorflow_probability as tfp\nfrom tensorflow import keras\nfrom tensorflow.keras import layers, models, callbacks\n# the line below causes the code to fail if running with GPU\n#tf.config.experimental.enable_mlir_graph_optimization()\n\nval_fraction = 0.15\n\n# decide which models to train during the current run of the notebook\ntask_list = {\n    'img_classifier':True, \n    'astro_recognitionNN':True,\n    'astro_segmentation':True,\n    'cort_recognitionNN':True,\n    'cort_segmentation':True,\n    'shsy5y_recognitionNN':True,\n    'shsy5y_segmentation':True\n}\n\n# handle model dependencies\nforce_retrain = False\ndependencies = {\n    'img_classifier':[], \n    'astro_recognitionNN':[],\n    'astro_segmentation':[],\n    'cort_recognitionNN':[],\n    'cort_segmentation':[],\n    'shsy5y_recognitionNN':[],\n    'shsy5y_segmentation':[]\n}\nfor task in dependencies.keys():\n    if task_list[task]:\n        for dependency in dependencies[task]:\n            if not task_list[dependency] and (not os.path.exists(dependency+'.model') or force_retrain):\n                task_list[dependency] = True\n    \nprint(task_list)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2022-01-17T13:55:02.476891Z","iopub.execute_input":"2022-01-17T13:55:02.477153Z","iopub.status.idle":"2022-01-17T13:55:18.811Z","shell.execute_reply.started":"2022-01-17T13:55:02.477124Z","shell.execute_reply":"2022-01-17T13:55:18.810249Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"First, let's plot a number of images from the test set, to see what we're dealing with..","metadata":{}},{"cell_type":"code","source":"img_filenames = glob.glob('/kaggle/input/sartorius-cell-instance-segmentation/train/*')\n\nif not submission_mode:\n    fig = plt.figure(figsize=(16,16))\n    gs = gridspec.GridSpec(4,4)\n    filenames_to_plot = np.random.choice(img_filenames, 16)\n    print(len(img_filenames))\n    for i in range(8):\n        plt.subplot(gs[2*(i%2), int(i/2)])\n        img = np.array(Image.open(filenames_to_plot[i]))\n        plt.imshow(img, cmap='inferno')\n        print(img.shape)\n        plt.subplot(gs[2*(i%2)+1, int(i/2)])\n        plt.hist(img.flatten(), 64)\n    plt.show()\n    plt.close()","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:37:08.108983Z","iopub.execute_input":"2022-01-17T12:37:08.109252Z","iopub.status.idle":"2022-01-17T12:37:08.150987Z","shell.execute_reply.started":"2022-01-17T12:37:08.109205Z","shell.execute_reply":"2022-01-17T12:37:08.150369Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Observations:\n - there are 606 train images total,\n - all images have the same shape (520x704) with a **single** color channel,\n - the color scale ranges from 0 to 255, with the vast majority of pixels being at the half of it (128),\n - the most noticeable cell features appear dark on the background.","metadata":{}},{"cell_type":"code","source":"# Let's confirm these observations on the full dataset:\nok = True\nif False: # already checked\n    for filename in tqdm(img_filenames):\n        img = np.array(Image.open(filenames_to_plot[i]))\n        if img.shape != (520, 704):\n            ok = False\n            break\n        img = img.flatten()\n        if np.min(img) < 0 or np.max(img) > 255:\n            ok = False\n            break\n        if np.abs(np.median(img) - 128) > 20.:\n            ok = False\n            break\nprint(ok)","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:39:52.249754Z","iopub.execute_input":"2022-01-17T12:39:52.250006Z","iopub.status.idle":"2022-01-17T12:39:52.256769Z","shell.execute_reply.started":"2022-01-17T12:39:52.24998Z","shell.execute_reply":"2022-01-17T12:39:52.256076Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"It may be easier to deal with the images (even visually) if we rescale them to have a mean of 0 and the color to vary between -1 and 1.","metadata":{}},{"cell_type":"code","source":"def img_color_rescale (img):\n    return (img-np.mean(img.flatten()))/128.","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:39:53.379049Z","iopub.execute_input":"2022-01-17T12:39:53.379678Z","iopub.status.idle":"2022-01-17T12:39:53.383886Z","shell.execute_reply.started":"2022-01-17T12:39:53.379643Z","shell.execute_reply":"2022-01-17T12:39:53.383085Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if not submission_mode:\n    fig = plt.figure(figsize=(16,16))\n    gs = gridspec.GridSpec(4,4)\n    filenames_to_plot = np.random.choice(img_filenames, 16)\n    for i in range(8):\n        plt.subplot(gs[2*(i%2), int(i/2)])\n        img = img_color_rescale(np.array(Image.open(filenames_to_plot[i])))\n        plt.imshow(img, cmap='seismic')\n        plt.subplot(gs[2*(i%2)+1, int(i/2)])\n        plt.hist(img.flatten(), 64)\n    plt.show()\n    plt.close()\n\n    fig = plt.figure(figsize=(16,16))\n    img = img_color_rescale(np.array(Image.open(filenames_to_plot[0])))\n    plt.imshow(img, cmap='seismic')\n    plt.show()\n    plt.close()","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:39:54.134506Z","iopub.execute_input":"2022-01-17T12:39:54.134755Z","iopub.status.idle":"2022-01-17T12:39:54.143262Z","shell.execute_reply.started":"2022-01-17T12:39:54.134729Z","shell.execute_reply":"2022-01-17T12:39:54.142545Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Ok, this looks much better. Now let's have a look at the label segmentation...","metadata":{}},{"cell_type":"code","source":"train_csv = pd.read_csv('/kaggle/input/sartorius-cell-instance-segmentation/train.csv', parse_dates=['sample_date'])\nfor time_col in ['plate_time', 'elapsed_timedelta']:\n    train_csv[time_col] = pd.to_timedelta(train_csv[time_col])","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:39:55.694822Z","iopub.execute_input":"2022-01-17T12:39:55.69551Z","iopub.status.idle":"2022-01-17T12:39:57.091755Z","shell.execute_reply.started":"2022-01-17T12:39:55.695473Z","shell.execute_reply":"2022-01-17T12:39:57.090908Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if not submission_mode:\n    print(train_csv.head())","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:39:57.093381Z","iopub.execute_input":"2022-01-17T12:39:57.093652Z","iopub.status.idle":"2022-01-17T12:39:57.098003Z","shell.execute_reply.started":"2022-01-17T12:39:57.093612Z","shell.execute_reply":"2022-01-17T12:39:57.097279Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if not submission_mode:\n    print(train_csv.dtypes)","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:39:57.270025Z","iopub.execute_input":"2022-01-17T12:39:57.270473Z","iopub.status.idle":"2022-01-17T12:39:57.27428Z","shell.execute_reply.started":"2022-01-17T12:39:57.270439Z","shell.execute_reply":"2022-01-17T12:39:57.273262Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"cell_types = train_csv.cell_type.unique()\nif not submission_mode:\n    print(cell_types)\n    train_csv.cell_type.hist()","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:39:58.187755Z","iopub.execute_input":"2022-01-17T12:39:58.18817Z","iopub.status.idle":"2022-01-17T12:39:58.204042Z","shell.execute_reply.started":"2022-01-17T12:39:58.188139Z","shell.execute_reply":"2022-01-17T12:39:58.203171Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if not submission_mode:\n    print(sorted(train_csv.plate_time.astype('timedelta64[h]').unique()/60.))\n    fig = plt.figure()\n    plt.hist(train_csv.plate_time.astype('timedelta64[h]'))\n    plt.title('timedelta64[h]')\n    plt.show()\n    plt.close()","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:39:59.014689Z","iopub.execute_input":"2022-01-17T12:39:59.014953Z","iopub.status.idle":"2022-01-17T12:39:59.021361Z","shell.execute_reply.started":"2022-01-17T12:39:59.014926Z","shell.execute_reply":"2022-01-17T12:39:59.018637Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if not submission_mode:\n    print(sorted(train_csv.elapsed_timedelta.astype('timedelta64[m]').unique()/60.))\n    fig = plt.figure()\n    plt.hist(train_csv.elapsed_timedelta.astype('timedelta64[m]'))\n    plt.title('timedelta64[m]')\n    plt.show()\n    plt.close()","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:39:59.859328Z","iopub.execute_input":"2022-01-17T12:39:59.859858Z","iopub.status.idle":"2022-01-17T12:39:59.864586Z","shell.execute_reply.started":"2022-01-17T12:39:59.859824Z","shell.execute_reply":"2022-01-17T12:39:59.86381Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if not submission_mode:\n    print(sorted((train_csv.plate_time.astype('timedelta64[s]')-train_csv.elapsed_timedelta.astype('timedelta64[s]')).unique()))\n    fig = plt.figure()\n    plt.hist(train_csv.plate_time.astype('timedelta64[s]')-train_csv.elapsed_timedelta.astype('timedelta64[s]'))\n    plt.title('plate time - timedelta [s]')\n    plt.show()\n    plt.close()","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:40:02.172484Z","iopub.execute_input":"2022-01-17T12:40:02.172739Z","iopub.status.idle":"2022-01-17T12:40:02.178046Z","shell.execute_reply.started":"2022-01-17T12:40:02.172712Z","shell.execute_reply":"2022-01-17T12:40:02.176844Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if not submission_mode:\n    print(sorted(train_csv.sample_date.unique()))\n    fig = plt.figure()\n    plt.hist(train_csv.sample_date)\n    plt.title(\"Sample date\")\n    plt.show()\n    plt.close()","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:40:02.31005Z","iopub.execute_input":"2022-01-17T12:40:02.310497Z","iopub.status.idle":"2022-01-17T12:40:02.315149Z","shell.execute_reply.started":"2022-01-17T12:40:02.310465Z","shell.execute_reply":"2022-01-17T12:40:02.314273Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if not submission_mode:\n    print('No of cells histogram')\n    train_csv.groupby('id').width.count().hist()","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:40:02.876975Z","iopub.execute_input":"2022-01-17T12:40:02.877611Z","iopub.status.idle":"2022-01-17T12:40:02.881613Z","shell.execute_reply.started":"2022-01-17T12:40:02.877573Z","shell.execute_reply":"2022-01-17T12:40:02.88062Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Observations:\n - there is a clear imbalance in the type of cells: shsy5y are overabundant by about a factor of 5 over astro and cort (which, in turn, have about equal populations),\n - the fields plate_time and elapsed_timedelta are identical,\n - Samples were taken in two rounds, about 2/3 of them in June of 2019 and 1/3 between Sept. and Nov. of 2020. Besides this bimodality, there is a very small number of distinct dates in the dataset, so it is likely that sample_date only designated the date the sample data was included in the database (not the date it was photographed)...\n - There is a large number of images with very small number of cells (<80), and a broader distribution centered around n ~ 400","metadata":{}},{"cell_type":"markdown","source":"With these in hand, let us now parse and plot the segmentation masks given in the training dataset. It would be memory-wasteful to store the masks as separate numpy arrays at all times, so let us just use a function to translate the run length encoded pixels from the csv into a numpy array **on the fly**. Fortunately, an efficient solution can be found on Stackoverflow (thanks, jdehesa and M. Innat!).","metadata":{}},{"cell_type":"code","source":"def rle_decode_tf(mask_rle, shape):\n    '''\n    Input [string]: Run-length encoded pixel mask\n    Output [tf.array of shape shape]: Segmentation mask\n    By: M. Innat, jdehesa\n    Source: https://stackoverflow.com/questions/58693261/decoding-rle-run-length-encoding-mask-with-tensorflow-datasets\n    '''\n    shape = tf.convert_to_tensor(shape, tf.int64)\n    size = tf.math.reduce_prod(shape)\n    # Split string\n    s = tf.strings.split(mask_rle)\n    s = tf.strings.to_number(s, tf.int64)\n    # Get starts and lengths\n    starts = s[::2] - 1\n    lens = s[1::2]\n    # Make ones to be scattered\n    total_ones = tf.reduce_sum(lens)\n    ones = tf.ones([total_ones], tf.uint8)\n    # Make scattering indices\n    r = tf.range(total_ones)\n    lens_cum = tf.math.cumsum(lens)\n    s = tf.searchsorted(lens_cum, r, 'right')\n    idx = r + tf.gather(starts - tf.pad(lens_cum[:-1], [(1, 0)]), s)\n    # Scatter ones into flattened mask\n    mask_flat = tf.scatter_nd(tf.expand_dims(idx, 1), ones, [size])\n    # Reshape into mask\n    return tf.reshape(mask_flat, shape)","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:40:07.893655Z","iopub.execute_input":"2022-01-17T12:40:07.894264Z","iopub.status.idle":"2022-01-17T12:40:07.904939Z","shell.execute_reply.started":"2022-01-17T12:40:07.894207Z","shell.execute_reply":"2022-01-17T12:40:07.904255Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Let's have a look at a single segmentation..","metadata":{}},{"cell_type":"code","source":"if not submission_mode:\n    print(filenames_to_plot[0])\n    img_id = filenames_to_plot[0].split('/')[-1].split('.')[0]\n    print(train_csv[train_csv.id == img_id])","metadata":{"execution":{"iopub.status.busy":"2021-12-30T10:28:37.906599Z","iopub.execute_input":"2021-12-30T10:28:37.907129Z","iopub.status.idle":"2021-12-30T10:28:37.948918Z","shell.execute_reply.started":"2021-12-30T10:28:37.907088Z","shell.execute_reply":"2021-12-30T10:28:37.948265Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"So there are only 5 segments on this image!","metadata":{}},{"cell_type":"code","source":"celltype_colors = {'shsy5y':'k', 'astro':'green', 'cort':'cyan'}\ndef plot_image (img_filename, ax=None):\n\n    # Let us plot an example of the segmentation\n    standalone = (ax == None)\n    if standalone:\n        fig = plt.figure(figsize=(16,16), clear=True)\n        ax = plt.gca()\n\n    # plot the image itself\n    img = img_color_rescale(np.array(Image.open(img_filename)))\n    ax.imshow(img, cmap='seismic', zorder=-3, rasterized=True)\n    shape = img.shape\n    del img\n\n    # overplot segmentation\n    img_id = img_filename.split('/')[-1].split('.')[0]\n    mask = np.zeros(shape).astype(np.float)\n    buffer = 0. * mask\n    for idx, row in tqdm(train_csv[train_csv.id == img_id].iterrows()):\n        # decode the mask\n        buffer = rle_decode_tf(row.annotation, shape).numpy().astype(np.float)\n        mask += buffer\n        # draw the bounding box\n        ypx, xpx = np.where(buffer == 1)\n        xmin, xmax = np.min(xpx), np.max(xpx)\n        ymin, ymax = np.min(ypx), np.max(ypx)\n        box = Rectangle(xy=(xmin,ymin), width=(xmax-xmin), height=(ymax-ymin), fill=False, ec=celltype_colors[row.cell_type], zorder=-1, rasterized=True)\n        ax.add_patch(box)\n        # add annotations to mark the cell details\n        ax.text(xmin,ymin, row.cell_type, color=celltype_colors[row.cell_type], va='bottom', weight='bold', zorder=-1, rasterized=True)\n    # draw the mask\n    mask[mask > 0] = 1\n    mask[mask == 0] = np.nan\n    cmap = cm.get_cmap('YlGn_r')\n    cmap.set_bad(alpha=0.0)\n    ax.imshow(mask, cmap=cmap, alpha=0.5, zorder=-2, rasterized=True)\n    del mask, buffer\n        \n    # rasterize to conserve RAM\n    ax.set_rasterization_zorder(0)\n        \n    if standalone:\n        plt.show()\n        fig.clear()\n        plt.close(fig)\n        del fig","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:40:27.860085Z","iopub.execute_input":"2022-01-17T12:40:27.860541Z","iopub.status.idle":"2022-01-17T12:40:27.874131Z","shell.execute_reply.started":"2022-01-17T12:40:27.860505Z","shell.execute_reply":"2022-01-17T12:40:27.873435Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if not submission_mode:\n    for filename in filenames_to_plot[:3]:\n        plot_image(filename)","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:40:34.481798Z","iopub.execute_input":"2022-01-17T12:40:34.482089Z","iopub.status.idle":"2022-01-17T12:40:34.488708Z","shell.execute_reply.started":"2022-01-17T12:40:34.482056Z","shell.execute_reply":"2022-01-17T12:40:34.488098Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":" - cort-type cells are tiny, circular, and only occupy small spaces on their cortesponding images\n - astro-type cells are HUGE, occupy significant fraction of the cortesponding images (both in surface area, as well as linear dimensions), and are elongated / brancing in shape\n - shsy5y are somewhat between these two types in both size and appearance\n - **the three types of cells are clearly found in different environments**\n  - cort are surrounded by a spiderweb-like network of threads (axions / dendrites, neural connections?)\n  - astro usually occur by themselves, with plain featureless background (given size of other cells present, we can be relatively confident that the scale is the same?)\n  - again, shsy5y is the transitional type","metadata":{}},{"cell_type":"markdown","source":"Given these observations, it may be helpful to work on two separate models -- one for cort-like cells and one for astro-like cells (a simple background-classifying model should be able to tell us which one we should use) -- and then merge them into a combined model that could handle the transitional shsy5y type as well.","metadata":{}},{"cell_type":"markdown","source":"First, let us find out the percentages of types for each picture..","metadata":{}},{"cell_type":"code","source":"img_topics = pd.DataFrame(train_csv['id'].unique(), columns=['id']).copy()\nfor cell_type in cell_types:\n    img_topics['perc_%s' % cell_type] = train_csv.groupby('id')['cell_type'].apply(lambda x : np.sum(x == cell_type)).reset_index()['cell_type'] * 1.0 / train_csv.groupby('id')['cell_type'].apply(len).reset_index()['cell_type']\nif not submission_mode:\n    img_topics.head()","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:40:56.044116Z","iopub.execute_input":"2022-01-17T12:40:56.04468Z","iopub.status.idle":"2022-01-17T12:40:56.464418Z","shell.execute_reply.started":"2022-01-17T12:40:56.044643Z","shell.execute_reply":"2022-01-17T12:40:56.46371Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if not submission_mode:\n    for cell_type in cell_types:\n        print('perc_%s\\n - unique values: ' % cell_type, img_topics['perc_%s' % cell_type].unique())\n        print(' - no of images: ', np.sum(img_topics['perc_%s' % cell_type] == 1))","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:41:10.120488Z","iopub.execute_input":"2022-01-17T12:41:10.120737Z","iopub.status.idle":"2022-01-17T12:41:10.125129Z","shell.execute_reply.started":"2022-01-17T12:41:10.120711Z","shell.execute_reply":"2022-01-17T12:41:10.124507Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":" - There are NO images with more than one type of cells (labeled) on them.\n - The number of images is not too large -- we **definitely** need to use augmentation to train a model to distinguish between cort and astro-like cells!\n \nThe problem of distinguishing these images does, however, seem relatively simple (they are immediately different even to my non-expert eye), so we probably don't need to resort to any state-of-the art pre-trained networks. Instead, let us use a simple CNN trained from scratch on the images.","metadata":{}},{"cell_type":"code","source":"# since our dataset is imbalanced, we will need sample weights to balance the training..\nsample_weights = img_topics.drop('id', axis=1).sum()\nsample_weights /= sample_weights.sum()\nsample_weights = sample_weights.to_dict()\nprint(sample_weights)\nimg_topics['sample_weight'] = img_topics.drop('id', axis=1).apply(\n    lambda row : np.sum([sample_weights[x] if row[x] > 0 else 0 for x in img_topics.drop('id', axis=1).columns]), axis=1)\nprint(img_topics.head())","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:41:26.496432Z","iopub.execute_input":"2022-01-17T12:41:26.496873Z","iopub.status.idle":"2022-01-17T12:41:26.779009Z","shell.execute_reply.started":"2022-01-17T12:41:26.496838Z","shell.execute_reply":"2022-01-17T12:41:26.77827Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# reload the dataset, just to be safe\ntrain_csv = pd.read_csv('/kaggle/input/sartorius-cell-instance-segmentation/train.csv', parse_dates=['sample_date'])\nfor time_col in ['plate_time', 'elapsed_timedelta']:\n    train_csv[time_col] = pd.to_timedelta(train_csv[time_col])\n\n# before we start, let's set apart the validation sample\nall_ids = train_csv.id.unique()\nval_ids = rd.choice(all_ids, int(val_fraction * len(all_ids)))\nval_csv = train_csv[train_csv.id.isin(val_ids)]\ntrain_csv = train_csv[~train_csv.id.isin(val_ids)]\n\nprint(len(train_csv), len(val_csv))","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:42:07.297022Z","iopub.execute_input":"2022-01-17T12:42:07.297291Z","iopub.status.idle":"2022-01-17T12:42:08.362455Z","shell.execute_reply.started":"2022-01-17T12:42:07.297262Z","shell.execute_reply":"2022-01-17T12:42:08.361729Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# split the sample weights\ntrain_weights = img_topics[~img_topics.id.isin(val_ids)]['sample_weight'].to_numpy()\nval_weights = img_topics[img_topics.id.isin(val_ids)]['sample_weight'].to_numpy()","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:42:09.411264Z","iopub.execute_input":"2022-01-17T12:42:09.4117Z","iopub.status.idle":"2022-01-17T12:42:09.419185Z","shell.execute_reply.started":"2022-01-17T12:42:09.411666Z","shell.execute_reply":"2022-01-17T12:42:09.418321Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# ... and finally let us prepare the labels\ntrain_labels = img_topics[~img_topics.id.isin(val_ids)].copy().drop(['id', 'sample_weight'], axis=1).to_numpy(dtype=np.int32)\nval_labels = img_topics[img_topics.id.isin(val_ids)].copy().drop(['id', 'sample_weight'], axis=1).to_numpy(dtype=np.int32)\nprint(train_labels.shape, val_labels.shape)","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:42:27.080642Z","iopub.execute_input":"2022-01-17T12:42:27.080883Z","iopub.status.idle":"2022-01-17T12:42:27.091727Z","shell.execute_reply.started":"2022-01-17T12:42:27.080857Z","shell.execute_reply":"2022-01-17T12:42:27.091049Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# preprocessing\ndef tf_img_load (img_id, directory='train'):\n    filestem = ('/kaggle/input/sartorius-cell-instance-segmentation/%s/' % directory)\n    img = tf.keras.utils.load_img(filestem+img_id+'.png', color_mode='grayscale')\n    return tf.keras.preprocessing.image.img_to_array(img)\ndef tf_img_color_rescale (img):\n    return (img-tf.reduce_mean(img))/128.","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:42:30.455904Z","iopub.execute_input":"2022-01-17T12:42:30.456204Z","iopub.status.idle":"2022-01-17T12:42:30.461615Z","shell.execute_reply.started":"2022-01-17T12:42:30.456173Z","shell.execute_reply":"2022-01-17T12:42:30.460329Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def tf_img_augment (model):\n    model.add(layers.RandomTranslation(\n        height_factor=(-1,1),\n        width_factor=(-1,1),\n        fill_mode='wrap'\n    ))\n    model.add(layers.RandomZoom(\n        height_factor=(-0.3, 0.1), # zoom-in 30% to zoom-out 10%\n        fill_mode='wrap')\n    )\n    model.add(layers.RandomRotation(\n        factor=(-1,1), # full rotation\n        fill_mode='wrap')\n    )\n    model.add(layers.RandomFlip())\n    return model","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:42:41.610503Z","iopub.execute_input":"2022-01-17T12:42:41.610905Z","iopub.status.idle":"2022-01-17T12:42:41.621549Z","shell.execute_reply.started":"2022-01-17T12:42:41.610865Z","shell.execute_reply":"2022-01-17T12:42:41.620407Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# in principle, we could do image loading as part of the model\n# , but, since there aren't that many of them,\n# it seems more sebsible to load them into our RAM memory all at once\nif task_list['img_classifier']:\n    train_imgs = []\n    for img_id in tqdm(train_csv.id.unique()):\n        train_imgs.append(\n            tf_img_color_rescale(\n                tf_img_load(img_id)\n            ).numpy()\n        )\n    train_imgs = np.array(train_imgs)\n\n    val_imgs = []\n    for img_id in tqdm(val_csv.id.unique()):\n        val_imgs.append(\n            tf_img_color_rescale(\n                tf_img_load(img_id)\n            ).numpy()\n        )\n    val_imgs = np.array(val_imgs)\n    print(train_imgs.shape, val_imgs.shape)","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:43:18.695083Z","iopub.execute_input":"2022-01-17T12:43:18.69534Z","iopub.status.idle":"2022-01-17T12:43:34.171731Z","shell.execute_reply.started":"2022-01-17T12:43:18.695313Z","shell.execute_reply":"2022-01-17T12:43:34.171028Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# build our model\nmodel_name = 'img_classifier'\nmodel = models.Sequential()\nmodel.add(keras.Input(shape=(520,704,1)))\n# first, augment the input image\n#model = tf_img_augment(model)\n# now some vanilla CNN layers\nmodel.add(layers.Conv2D(\n    filters=32, kernel_size=(5,5),\n    activation='relu')\n)\nmodel.add(layers.MaxPooling2D((2,2)))\nmodel.add(layers.Conv2D(\n    filters=32, kernel_size=(5,5),\n    activation='relu')\n)\nmodel.add(layers.MaxPooling2D((2,2)))\nmodel.add(layers.Conv2D(\n    filters=32, kernel_size=(5,5),\n    activation='relu')\n)\nmodel.add(layers.MaxPooling2D((2,2)))\n# and finish up with a perceptron\nmodel.add(layers.Flatten())\nmodel.add(layers.Dense(\n    units=128, activation='relu'\n))\nmodel.add(layers.Dropout(0.5))\nmodel.add(layers.Dense(\n    units=3, activation='relu'\n))\n\nmodel.compile(optimizer='adam',\n              loss=tf.keras.losses.CategoricalCrossentropy(\n                  from_logits=True\n              ),\n              metrics=['accuracy'])\n\nprint('input', model.input.shape)\nfor layer in model.layers:\n    print(layer.name, layer.output_shape)","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:43:36.331791Z","iopub.execute_input":"2022-01-17T12:43:36.332412Z","iopub.status.idle":"2022-01-17T12:43:36.751111Z","shell.execute_reply.started":"2022-01-17T12:43:36.332373Z","shell.execute_reply":"2022-01-17T12:43:36.750421Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if task_list[model_name]:\n    stop_early = tf.keras.callbacks.EarlyStopping(\n        monitor='val_loss',\n        patience=early_stopping_patience,\n        restore_best_weights=True\n    )\n    callbacks = [stop_early,]\n    history = model.fit(\n        train_imgs,\n        train_labels,\n        epochs=epochs_to_train,\n        validation_data=(val_imgs, val_labels),\n        sample_weight=train_weights,\n        callbacks=callbacks\n    )","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:46:51.535706Z","iopub.execute_input":"2022-01-17T12:46:51.536362Z","iopub.status.idle":"2022-01-17T12:48:35.270313Z","shell.execute_reply.started":"2022-01-17T12:46:51.536323Z","shell.execute_reply":"2022-01-17T12:48:35.269473Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if task_list[model_name] and not submission_mode:\n    plt.plot(history.history['loss'], label='loss')\n    plt.plot(history.history['val_loss'], label = 'val_loss')\n    plt.xlabel('Epoch')\n    plt.ylabel('Loss')\n    plt.legend(loc='best')\n    plt.show()\n    plt.close()","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:48:35.275512Z","iopub.execute_input":"2022-01-17T12:48:35.277623Z","iopub.status.idle":"2022-01-17T12:48:35.285195Z","shell.execute_reply.started":"2022-01-17T12:48:35.277583Z","shell.execute_reply":"2022-01-17T12:48:35.284574Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Hmm.. without data augmentation, the network clearly memorizes the training data (as expected, for a measly sample of 300 images). While the validation performance surprisingly isn't bad (at 98%!), it would be embarassing to use such an overfit model. Let's do this again, with image augmentation turned on:","metadata":{}},{"cell_type":"code","source":"# build our model\n# prepend the previous pre-trained model with an augmentation layer\nnew_model = models.Sequential()\nnew_model.add(keras.Input(shape=(520,704,1)))\n# first, augment the input image\nnew_model = tf_img_augment(new_model)\nfor layer in model.layers:\n    new_model.add(layer)\nmodel = new_model\n\nmodel.compile(optimizer='adam',\n              loss=tf.keras.losses.CategoricalCrossentropy(from_logits=True),\n              metrics=['accuracy'])\n\nfor layer in model.layers:\n    print(layer.name, layer.output_shape)","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:48:35.289974Z","iopub.execute_input":"2022-01-17T12:48:35.292347Z","iopub.status.idle":"2022-01-17T12:48:35.618299Z","shell.execute_reply.started":"2022-01-17T12:48:35.292313Z","shell.execute_reply":"2022-01-17T12:48:35.617604Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if task_list[model_name]:\n    if force_retrain or not os.path.exists(model_name+'.model'):\n        stop_early = tf.keras.callbacks.EarlyStopping(\n            monitor='val_loss',\n            patience=early_stopping_patience,\n            restore_best_weights=True\n        )\n        callbacks = [stop_early,]\n        history = model.fit(train_imgs, train_labels, epochs=epochs_to_train,\n                           validation_data=(val_imgs, val_labels),\n                           sample_weight=train_weights,\n                           callbacks=callbacks)\n        model.save(model_name + '.model')\n    else:\n        model = keras.models.load_model(\n            model_name + '.model'\n        )\n    img_classifier = model","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:48:35.620292Z","iopub.execute_input":"2022-01-17T12:48:35.620524Z","iopub.status.idle":"2022-01-17T12:53:02.651592Z","shell.execute_reply.started":"2022-01-17T12:48:35.62049Z","shell.execute_reply":"2022-01-17T12:53:02.650931Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if task_list['img_classifier'] and not submission_mode:\n    plt.plot(history.history['loss'], label='loss')\n    plt.plot(history.history['val_loss'], label = 'val_loss')\n    plt.xlabel('Epoch')\n    plt.ylabel('Loss')\n    plt.legend(loc='best')\n    plt.gca().set_yscale('log')\n    plt.show()\n    plt.close()","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:53:02.657269Z","iopub.execute_input":"2022-01-17T12:53:02.657466Z","iopub.status.idle":"2022-01-17T12:53:02.6623Z","shell.execute_reply.started":"2022-01-17T12:53:02.657442Z","shell.execute_reply":"2022-01-17T12:53:02.661358Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Well, this is probably more reasonable. This proves that we can recognize environments of our different cell classes in the images.\n\nThus, let us proceed to build three separate segmentation models, each specified for one type of neurons. We will then connect them to the model we've just trained, to build our solution to the competition problem.","metadata":{}},{"cell_type":"code","source":"# clean up after the img_classifier training\nif task_list['img_classifier']:\n    del train_imgs, val_imgs","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:53:02.663862Z","iopub.execute_input":"2022-01-17T12:53:02.664409Z","iopub.status.idle":"2022-01-17T12:53:02.675176Z","shell.execute_reply.started":"2022-01-17T12:53:02.664372Z","shell.execute_reply":"2022-01-17T12:53:02.674434Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"---------------\n\n**ASTRO-cell segmentation**\n\nLet's start with a segmentation model to recognize the astro-type cells. First, let's create a dataset containing only images with these types.","metadata":{}},{"cell_type":"code","source":"if not submission_mode:\n    img_topics.head()","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:53:02.676377Z","iopub.execute_input":"2022-01-17T12:53:02.67697Z","iopub.status.idle":"2022-01-17T12:53:02.684334Z","shell.execute_reply.started":"2022-01-17T12:53:02.676933Z","shell.execute_reply":"2022-01-17T12:53:02.683632Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"astro_ids = img_topics[img_topics.perc_astro == 1.0].id.to_numpy()\nprint(len(astro_ids))","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:53:02.685343Z","iopub.execute_input":"2022-01-17T12:53:02.685613Z","iopub.status.idle":"2022-01-17T12:53:02.69473Z","shell.execute_reply.started":"2022-01-17T12:53:02.685579Z","shell.execute_reply":"2022-01-17T12:53:02.69387Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"That's not a lot of images... Hopefully image segmentation will help us...\n\nLet's split into train and validation sample:","metadata":{}},{"cell_type":"code","source":"astro_csv = train_csv[train_csv.id.isin(astro_ids)]\nastro_val_ids = rd.choice(astro_ids, int(val_fraction * len(astro_ids)))\nastro_val_csv = astro_csv[astro_csv.id.isin(astro_val_ids)]\nastro_train_csv = astro_csv[~astro_csv.id.isin(astro_val_ids)]\nprint(len(astro_train_csv), len(astro_val_csv))","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:53:02.695844Z","iopub.execute_input":"2022-01-17T12:53:02.696114Z","iopub.status.idle":"2022-01-17T12:53:02.714746Z","shell.execute_reply.started":"2022-01-17T12:53:02.696076Z","shell.execute_reply":"2022-01-17T12:53:02.714104Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"At least the number of cells is pretty decent.","metadata":{}},{"cell_type":"code","source":"# in principle, we could do image loading as part of the model\n# , but, since there aren't that many of them,\n# it seems more sensible to load them into our RAM memory all at once\nshape = (520,704)\nsave_path = 'astro_train_data.pkl'\nif task_list['astro_recognitionNN'] or task_list['astro_segmentation']:\n    if force_retrain or not os.path.exists(save_path):\n        astro_train_imgs = []\n        astro_train_masks = []\n        for img_id in tqdm(astro_train_csv.id.unique()):\n            astro_train_imgs.append(\n                tf_img_color_rescale(\n                    tf_img_load(img_id)\n                ).numpy()\n            )\n            # could probably make this faster with a tensorflow map, but as-is, a dtype error is returned (can't map a string-type tensor to int)...\n            mask = np.zeros(shape)\n            for ann in astro_train_csv[astro_train_csv.id == img_id].annotation:\n                mask += rle_decode_tf(ann, shape=shape)\n            mask = tf.where(mask > 0, 1,0)\n            astro_train_masks.append(mask.numpy().astype(np.uint8))\n        astro_train_imgs = np.array(astro_train_imgs)\n        astro_train_masks = np.array(astro_train_masks)\n\n        astro_val_imgs = []\n        astro_val_masks = []\n        for img_id in tqdm(astro_val_csv.id.unique()):\n            astro_val_imgs.append(\n                tf_img_color_rescale(\n                    tf_img_load(img_id)\n                ).numpy()\n            )\n            mask = np.zeros(shape)\n            for ann in astro_val_csv[astro_val_csv.id == img_id].annotation:\n                mask += rle_decode_tf(ann, shape=shape)\n            mask = tf.where(mask > 0, 1,0)\n            astro_val_masks.append(mask.numpy().astype(np.uint8))\n        astro_val_imgs = np.array(astro_val_imgs)\n        astro_val_masks = np.array(astro_val_masks)\n        with open(save_path, 'wb') as f:\n            pkl.dump((astro_train_imgs, astro_train_masks, astro_val_imgs, astro_val_masks), f)\n    else:\n        with open(save_path, 'rb') as f:\n            astro_train_imgs, astro_train_masks, astro_val_imgs, astro_val_masks = pkl.load(f)\n    print(astro_train_imgs.shape, astro_val_imgs.shape)\n    print(astro_train_masks.shape, astro_val_masks.shape)","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:53:02.717718Z","iopub.execute_input":"2022-01-17T12:53:02.718127Z","iopub.status.idle":"2022-01-17T12:54:25.212696Z","shell.execute_reply.started":"2022-01-17T12:53:02.7181Z","shell.execute_reply":"2022-01-17T12:54:25.211895Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if not submission_mode:\n    plt.imshow(astro_val_masks[2])\n    plt.show()\n    plt.close()","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:54:25.213864Z","iopub.execute_input":"2022-01-17T12:54:25.214141Z","iopub.status.idle":"2022-01-17T12:54:25.218982Z","shell.execute_reply.started":"2022-01-17T12:54:25.214104Z","shell.execute_reply":"2022-01-17T12:54:25.218323Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Let's have a look at a couple of the images, just to refresh our memory..","metadata":{}},{"cell_type":"code","source":"if not submission_mode:\n    for filename in ['/kaggle/input/sartorius-cell-instance-segmentation/train/' + x + '.png' for x in astro_ids][:3]:\n        plot_image(filename)","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:54:25.220143Z","iopub.execute_input":"2022-01-17T12:54:25.220522Z","iopub.status.idle":"2022-01-17T12:54:25.228952Z","shell.execute_reply.started":"2022-01-17T12:54:25.220486Z","shell.execute_reply":"2022-01-17T12:54:25.228263Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":" - The cells cortespond to a large fraction of the image, so the segmentation network's footprint should be at least equal to ~25% of the image, to capture (ideally) whole cells\n - Cell segmentation maps contain fine structures (like threads, spikes, etc.), so lowering the image resolution to generate a segmentation map might not be the best idea...\n \n Given these observations, let us start with a \"brute-force\" idea, with a lossless convolutional network at a fixed size (no maxpool) mapping our image to the segmentation map..","metadata":{}},{"cell_type":"code","source":"# due to imbalance numbers of on-cell and off-cell pixels, let's use a weighing function\ndef balanced_loss (y_true, y_pred):\n    loss_falsepositive = tf.reduce_mean(tf.where(y_true > 0, (y_true-y_pred)**2, 0))\n    loss_falsenegative = tf.reduce_mean(tf.where(y_true < 1, (y_true-y_pred)**2, 0))\n    return loss_falsepositive + loss_falsenegative","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:54:25.230166Z","iopub.execute_input":"2022-01-17T12:54:25.230513Z","iopub.status.idle":"2022-01-17T12:54:25.237907Z","shell.execute_reply.started":"2022-01-17T12:54:25.230477Z","shell.execute_reply":"2022-01-17T12:54:25.23709Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# since we're working with segmentation now, \n# we need to apply the same augmentation\n# to both the input and the label masks\n# We will do this by carrying masks over with the images,\n# contained within an extra dimension.\ndef tf_img_augment_nomodel (X, y):\n    # concatenate along channels\n    # to process both the img and mask with the same operations\n    buffer = tf.concat([X,y], axis=3)\n    # build the augmentation model\n    buffer = layers.RandomTranslation(\n        height_factor=(-1,1),\n        width_factor=(-1,1),\n        fill_mode='wrap'\n    )(buffer)\n    buffer = layers.RandomZoom(\n        height_factor=(-0.3, 0.1), # zoom-in 30% to zoom-out 10%\n        fill_mode='wrap'\n    )(buffer)\n    buffer = layers.RandomRotation(\n        factor=(-1,1), # full rotation\n        fill_mode='wrap'\n    )(buffer)\n    buffer = layers.RandomFlip()(buffer)\n    # split up the results into img and mask\n    X = tf.expand_dims(buffer[:,:,:,0], axis=3)\n    y = tf.expand_dims(buffer[:,:,:,1], axis=3)\n    # clean up\n    del buffer\n    return X,y","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:54:25.239015Z","iopub.execute_input":"2022-01-17T12:54:25.239823Z","iopub.status.idle":"2022-01-17T12:54:25.250351Z","shell.execute_reply.started":"2022-01-17T12:54:25.239787Z","shell.execute_reply":"2022-01-17T12:54:25.249517Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# To work with our custom augmentation,\n# we will need to overwrite the fitting method\n# - adapted from https://keras.io/guides/customizing_what_happens_in_fit/\nclass AugmentedSequential (keras.Sequential):\n    def train_step(self, data):\n        # read the data passed to fit()\n        X, y = data\n        # augment\n        X, y = tf_img_augment_nomodel(X, y)\n        # perform the augmentation and turn into tensors\n        X, y = tf.constant(X.numpy()), tf.constant(y.numpy())\n        \n        # then proceed as usual\n        return super().train_step((X,y))","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:54:25.25336Z","iopub.execute_input":"2022-01-17T12:54:25.253602Z","iopub.status.idle":"2022-01-17T12:54:25.262296Z","shell.execute_reply.started":"2022-01-17T12:54:25.253577Z","shell.execute_reply":"2022-01-17T12:54:25.261552Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# build our astro_model\nmodel_name = 'astro_recognitionNN'\nastro_model = AugmentedSequential()\nastro_model.add(keras.Input(\n    shape=(520,704,1)\n))\n# now some vanilla CNN layers\nif True: # val_loss: 0.0659\n    astro_model.add(layers.Conv2D(\n        filters=16, kernel_size=(3,3),\n        activation='relu',\n        padding='same')\n    )\n    astro_model.add(layers.Conv2D(\n        filters=8, kernel_size=(5,5),\n        activation='relu',\n        padding='same')\n    )\n    astro_model.add(layers.Conv2D(\n        filters=8, kernel_size=(7,7),\n        activation='relu',\n        padding='same')\n    )\n# Apply a dense layer along the axis of the filters, keeping the image size the same\nastro_model.add(layers.Conv2D(\n    filters=8, kernel_size=(1,1),\n    activation='relu',\n    padding='same')\n)\n# OUPTUT -----------------------------\n#astro_model.add(layers.Dropout(0.5))\nastro_model.add(layers.Conv2D(\n    filters=1, kernel_size=(1,1),\n    activation='sigmoid',\n    padding='same')\n)\n\nastro_model.compile(optimizer='adam',\n              loss=balanced_loss,\n                   run_eagerly=True)  #tf.keras.losses.CategoricalCrossentropy(from_logits=False),\n              #metrics=['accuracy']\n\nastro_model.summary()","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:54:25.263708Z","iopub.execute_input":"2022-01-17T12:54:25.26424Z","iopub.status.idle":"2022-01-17T12:54:25.324584Z","shell.execute_reply.started":"2022-01-17T12:54:25.264185Z","shell.execute_reply":"2022-01-17T12:54:25.323897Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def plot_history (history):\n    plt.plot(history.history['loss'], label='loss')\n    plt.plot(history.history['val_loss'], label = 'val_loss')\n    plt.xlabel('Epoch')\n    plt.ylabel('Loss')\n    plt.legend(loc='best')\n    plt.gca().set_yscale('log')\n    plt.show()\n    plt.close()","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:54:25.325721Z","iopub.execute_input":"2022-01-17T12:54:25.326142Z","iopub.status.idle":"2022-01-17T12:54:25.332109Z","shell.execute_reply.started":"2022-01-17T12:54:25.326107Z","shell.execute_reply":"2022-01-17T12:54:25.331274Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if task_list[model_name]:\n    print(model_name)\n    if force_retrain or not os.path.exists(model_name+'.model'):\n        stop_early = tf.keras.callbacks.EarlyStopping(\n            monitor='val_loss',\n            patience=early_stopping_patience,\n            restore_best_weights=True\n        )\n        callbacks = [stop_early,]\n        history = astro_model.fit(\n            tf.cast(astro_train_imgs,tf.float32).numpy(),\n            tf.cast(tf.expand_dims(astro_train_masks, axis=3),tf.float32).numpy(),\n            validation_data=(\n                tf.cast(astro_val_imgs,tf.float32).numpy(),\n                tf.cast(tf.expand_dims(astro_val_masks, axis=3),tf.float32).numpy()),\n            callbacks=callbacks,\n            epochs=epochs_to_train)\n        if not submission_mode:\n            astro_model.save(model_name + '.model')\n            plot_history(history)\n    else:\n        astro_model = keras.models.load_model(\n            model_name+'.model',\n            custom_objects={'balanced_loss':balanced_loss}\n        )","metadata":{"execution":{"iopub.status.busy":"2022-01-17T12:54:25.333869Z","iopub.execute_input":"2022-01-17T12:54:25.334275Z","iopub.status.idle":"2022-01-17T13:04:20.655742Z","shell.execute_reply.started":"2022-01-17T12:54:25.334145Z","shell.execute_reply":"2022-01-17T13:04:20.655053Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# make sure augmentation works, check the predictions\ndef test_predictions (model):\n    for i in range(3):\n        X_val, y_val = tf_img_augment_nomodel(\n            astro_val_imgs, \n            tf.expand_dims(astro_val_masks, axis=3).numpy()\n        )\n\n        ex_img = X_val[4].numpy()\n        ex_truemask = y_val[4].numpy()\n\n        # show an example prediction\n        fig = plt.figure(figsize=(16,6), clear=True)\n\n        plt.subplot(131)\n        plt.imshow(ex_img, cmap='seismic')\n        plt.title('Original image')\n\n        plt.subplot(132)\n        ex_mask = model.predict(ex_img.reshape((1,520,704,1)))[0]\n        print(ex_mask.shape)\n        plt.imshow(ex_mask, cmap=\"YlGn_r\")\n        plt.title(\"Model prediction\")\n        plt.colorbar()\n\n        plt.subplot(133)\n        plt.imshow(ex_truemask, cmap=\"YlGn_r\")\n        plt.title(\"True mask\")\n\n        plt.show()\n        plt.close()\n    \nif task_list[model_name] and not submission_mode:\n    test_predictions(astro_model)","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:04:20.656859Z","iopub.execute_input":"2022-01-17T13:04:20.657099Z","iopub.status.idle":"2022-01-17T13:04:20.666031Z","shell.execute_reply.started":"2022-01-17T13:04:20.657055Z","shell.execute_reply":"2022-01-17T13:04:20.665347Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"This looks quite good. I have tried some other ideas, but this simple network performed best for me..\n\nLet us now proceed to map this detection mask into segmentation into separate neurons. First, we will need to threshold the map.","metadata":{}},{"cell_type":"code","source":"# The first step will probably be to binarize our output map with a threshold\n# it would be great for the threshold value to be trainable\n# Fortunately, this was already done by zeka0 (https://github.com/keras-team/keras/issues/6926)\nclass ThresholdLayer(keras.layers.Layer):\n    def __init__(self, **kwargs):\n        super(ThresholdLayer, self).__init__(**kwargs)\n\n    def build(self, input_shape):\n        self.kernel = self.add_weight(name=\"threshold\", shape=(1,), initializer=\"uniform\",\n                                      trainable=True)\n        super(ThresholdLayer, self).build(input_shape)\n\n    def call(self, x):\n        return keras.backend.sigmoid(100*(x-self.kernel))\n\n    def compute_output_shape(self, input_shape):\n        return input_shape","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:04:20.678793Z","iopub.execute_input":"2022-01-17T13:04:20.679234Z","iopub.status.idle":"2022-01-17T13:04:20.690414Z","shell.execute_reply.started":"2022-01-17T13:04:20.679177Z","shell.execute_reply":"2022-01-17T13:04:20.689694Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# First, let's improve the contrast of the network output, to make it easier to segment\nif not submission_mode:\n    img = astro_train_imgs[5]\n    mask = astro_train_masks[5]\n    pred = astro_model.predict(np.array([img,]))[0]\n    th_layer = ThresholdLayer()\n    th_layer.build(input_shape=pred.shape)\n    print(th_layer.get_weights())\n    th_layer.set_weights([np.array([0.6])])\n    th = th_layer(pred).numpy()\n\n    plt.figure(figsize=(20,6), clear=True)\n    plt.subplot(151)\n    plt.imshow(img, cmap='seismic')\n    plt.title('Original image')\n    plt.subplot(152)\n    plt.imshow(mask, cmap=\"YlGn_r\")\n    plt.title('True mask')\n    plt.subplot(153)\n    plt.imshow(pred, cmap=\"YlGn_r\")\n    plt.title('Predicted mask')\n    plt.subplot(154)\n    plt.imshow(pred > 0.6, cmap=\"YlGn_r\")\n    plt.title('Threshold 1')\n    plt.subplot(155)\n    plt.imshow(th, cmap=\"YlGn_r\")\n    plt.title('Threshold 2')\n    plt.show()\n    plt.close()\n\n    plt.hist(pred.flatten())\n    plt.title('Prediction value histogram')\n    plt.show()\n    plt.close()","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:04:20.691875Z","iopub.execute_input":"2022-01-17T13:04:20.69214Z","iopub.status.idle":"2022-01-17T13:04:20.703399Z","shell.execute_reply.started":"2022-01-17T13:04:20.692106Z","shell.execute_reply":"2022-01-17T13:04:20.702675Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Now, we need to segment the cortesponding binary map. First let's try to do the following:\n - translate the np array map into a list of positive pixel coordinates (~ point cloud)\n - perform clustering of the resulting dataset using classical ml methods (K-means?)\n \nIdeally, we would need a clustering algorithm that:\n - doesn't require knowing the number of clusters in advance,\n - does well on elongated structures, connected into complex networks.\n \nJudgning by [https://scikit-learn.org/stable/modules/clustering.html](https://scikit-learn.org/stable/modules/clustering.html), this would include:\n - DBSCAN,\n - OPTICS.\n \nLet's try all of them and see how they perform. Let's use the \"true\" segmentation mask, to split their performance from the performance of our segmentation map extraction network.","metadata":{}},{"cell_type":"code","source":"# Let's plot the true segmentation, so that we know what we're looking for...\nif not submission_mode:\n    myid = astro_train_csv.id.unique()[5]\n    astro_train_csv[astro_train_csv.id == myid].head()\n\n    pixels_true = []\n    cl = 0\n    for annotation in tqdm(tf.random.shuffle( # shuffle helps to make adjacent cells different color\n        astro_train_csv[astro_train_csv.id == myid].annotation)\n                          ):\n        mask_here = rle_decode_tf(annotation, shape=img.shape)\n        pixels_true_here = tf.where(tf.squeeze(mask_here) > 0)\n        pixels_true_here = tf.concat([pixels_true_here, cl*np.ones([pixels_true_here.shape[0],1])], axis=1)\n        cl += 1\n        pixels_true.append(pixels_true_here)\n    pixels_true = tf.concat(pixels_true, axis=0)\n\n    x = tf.transpose(pixels_true)[0]\n    y = tf.transpose(pixels_true)[1]\n    c = tf.transpose(pixels_true)[2]\n\n    plt.scatter(x, y, c=c, s=0.125, cmap='prism')\n    plt.show()\n    plt.close()","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:04:20.704732Z","iopub.execute_input":"2022-01-17T13:04:20.705087Z","iopub.status.idle":"2022-01-17T13:04:20.715168Z","shell.execute_reply.started":"2022-01-17T13:04:20.705042Z","shell.execute_reply":"2022-01-17T13:04:20.714488Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if not submission_mode:\n    pixels = tf.where(tf.squeeze(mask) > 0).numpy()\n\n    x, y = tf.transpose(pixels)\n    plt.scatter(x, y, s=0.125)\n    plt.show()\n    plt.close()","metadata":{"execution":{"iopub.status.busy":"2021-12-30T10:36:08.493335Z","iopub.execute_input":"2021-12-30T10:36:08.493763Z","iopub.status.idle":"2021-12-30T10:36:08.763382Z","shell.execute_reply.started":"2021-12-30T10:36:08.493727Z","shell.execute_reply":"2021-12-30T10:36:08.762491Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Now let's try classical clustering algorithms on this point cloud\nfrom sklearn.cluster import DBSCAN\nif not submission_mode:\n    clusterer = DBSCAN(eps=1, min_samples=1)\n\n    clusterer.fit(pixels)\n    print(len(np.unique(clusterer.labels_)))\n\n    plt.scatter(x, y, c=clusterer.labels_, s=0.125, cmap='prism')\n    plt.show()\n    plt.close()","metadata":{"execution":{"iopub.status.busy":"2021-12-30T10:36:08.768556Z","iopub.execute_input":"2021-12-30T10:36:08.7692Z","iopub.status.idle":"2021-12-30T10:36:10.560373Z","shell.execute_reply.started":"2021-12-30T10:36:08.769153Z","shell.execute_reply":"2021-12-30T10:36:10.559656Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# https://scikit-learn.org/stable/auto_examples/cluster/plot_optics.html#sphx-glr-auto-examples-cluster-plot-optics-py\nif True:\n    from sklearn.cluster import OPTICS\n    clusterer = OPTICS() #(min_samples=20, xi=1.0, min_cluster_size=5)\n\n    clusterer.fit(pixels)\n    print(len(np.unique(clusterer.labels_)))\n\n    plt.scatter(x, y, c=clusterer.labels_, s=0.125, cmap='prism')\n    plt.show()\n    plt.close()","metadata":{"execution":{"iopub.status.busy":"2021-12-30T10:36:10.5615Z","iopub.execute_input":"2021-12-30T10:36:10.561748Z","iopub.status.idle":"2021-12-30T10:36:10.568905Z","shell.execute_reply.started":"2021-12-30T10:36:10.561712Z","shell.execute_reply":"2021-12-30T10:36:10.567662Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"DBSCAN is probably a bit better (depending on parameters, I suspect), but still not ideal. Still, it is much faster, so we will stick with that model.\n\nSegmentation is still, however, far from perfect. The clustering model has trouble splitting cells that are connected by threads (axions?). We should be able to fix this by pre-processing the segmentation map to push the cells further apart. Since I am, by background, a physicist, an idea that comes to mind is the following:\n - treat the point cloud as a collection of test particles,\n - introduce a force that will make particles stick (be attracted) to each other at close distances and repel each other at larger distances,\n - evolve the segmentation map with a particle simulator that applies the introduced force,\n - the particles should split according to their groupings in the original map, splitting at narrow features,\n - the resulting evolved map will be much simpler to segment.","metadata":{}},{"cell_type":"markdown","source":"There are already great particle simulators for Python available. However, since the problem here is very simple, it may be best to use a custom solution (the available simulators will be rather multi-purpose, so may be a bit unwieldy for us here). Let us quickly write a simple particle simulator then.","metadata":{}},{"cell_type":"code","source":"from matplotlib.animation import FuncAnimation\nfrom IPython import display\nclass ParticleSimulator ():\n    \n    def __init__ (self, particles_init, force, force_radius=20., force_kwargs={}, store_history=False):\n        self.particles = tf.cast(tf.convert_to_tensor(particles_init), tf.float16) # particles stored in shape (...,2+) with the last dimension being (x_pos, y_pos, ...)\n        self.particles_shape = tf.shape(self.particles)\n        self.transposed = tf.transpose(self.particles) # useful for plotting\n        self.n_part = len(particles_init)\n        self.force = force # a function f(part1, part2) returning a force acting between the given two particles (represented by (x_pos, y_pos, _))\n        # make sure that force(part1,part1) = 0\n        self.force_radius = force_radius # force only applied to particles that are closer than that\n        self.force_kwargs = force_kwargs\n        self.time = 0 # progress of the simulation in timesteps\n        self.store_history = store_history\n        if self.store_history:\n            self.history = tf.expand_dims(copy(self.particles.numpy()), axis=0)\n        else:\n            self.history = []\n    \n    def all_forces (self, sparse=True, symmetric=True):\n        # calculate distances between particles\n        u = tf.repeat(tf.expand_dims(self.particles[:,:2], axis=0), tf.shape(self.particles)[0], axis=0)\n        v = tf.transpose(u, perm=[1,0,2])\n        # if the forces are symmetric, create upper triangular slices of the distance matrix, to save computation and memory\n        if symmetric:\n            distance_matrix = tf.transpose(\n                tf.linalg.band_part(\n                    tf.transpose(u-v, perm=[2,0,1]),\n                    0,-1\n                ),\n                perm=[1,2,0]\n            )\n        else:\n            distance_matrix = u-v\n        sqr_distance_matrix = tf.reduce_sum(distance_matrix**2, axis=2)\n        # limit to forces between nearby particles\n        mask = tf.logical_and(sqr_distance_matrix > 0.5, sqr_distance_matrix < tf.cast(self.force_radius**2, tf.float16))\n        indices = tf.cast(tf.where(mask), tf.int32)\n        indices_vec = tf.concat(\n            [tf.repeat(indices, tf.shape(distance_matrix)[2], axis=0),\n             tf.tile([[0,],[1,]], [tf.shape(indices)[0],1])],\n            axis=1\n        )\n        if sparse:\n            # turn the distance matrices into sparse tensors\n            distance_matrix = tf.SparseTensor(\n                indices=tf.cast(indices_vec, tf.int64),\n                values=tf.gather_nd(distance_matrix, indices_vec),\n                dense_shape=tf.cast(tf.shape(distance_matrix), tf.int64)\n            )\n            sqr_distance_matrix = tf.SparseTensor(\n                indices=tf.cast(indices, tf.int64),\n                values=tf.gather_nd(sqr_distance_matrix, indices),\n                dense_shape=tf.cast(tf.shape(sqr_distance_matrix), tf.int64)\n            )\n            # fill the forces matrix\n            forces = tf.SparseTensor(\n                indices=distance_matrix.indices,\n                values=self.force(distance_matrix,sqr_distance_matrix, self.force_radius), # THIS IS THE ACTUAL FORCE\n                dense_shape=distance_matrix.shape\n            )\n            # sum forces for each particle, including the anti-symmetric lower diagonal, if needed\n            if symmetric:\n                forces2 = tf.sparse.transpose(forces, perm=[1,0,2])\n            forces = tf.sparse.reduce_sum(forces, axis=1)\n            # forces are anti-symmetric, recreate the lower diagonal\n            if symmetric:\n                forces -= tf.sparse.reduce_sum(forces2, axis=1)\n        else: # dense matrices used\n            forces = tf.where(\n                tf.repeat(tf.expand_dims(mask,axis=2), tf.shape(distance_matrix)[2], axis=2),\n                0.1 * distance_matrix / tf.repeat(tf.expand_dims(sqr_distance_matrix, axis=2), tf.shape(distance_matrix)[2], axis=2),\n                tf.zeros(tf.shape(distance_matrix), dtype=tf.float16)\n            )\n            forces -= tf.transpose(forces, perm=[1,0,2])\n            forces = tf.reduce_sum(forces, axis=1)\n        return forces\n    \n    def integrator (self, part, force, dt=1.0): # evolve the particle by dt=1 given the force f\n        # the forward Euler integrator would be **terrible** for actual physics simulations, but for our use it is probably fine...\n        return part[:2] + dt*force\n            \n    def plot (self):\n        if len(self.transposed) > 2:\n            c = self.transposed[2]\n        else:\n            c = None\n        plt.scatter(self.transposed[0], self.transposed[1], c=c, s=0.125, cmap='prism')\n        plt.show()\n        plt.close()\n        \n    def plot_history_frame (self, i, ax):\n        xlim, ylim = ax.get_xlim(), ax.get_ylim()\n        ax.clear()\n        ax.set_xlim(xlim)\n        ax.set_ylim(ylim)\n        transposed = tf.transpose(self.history[i])\n        if tf.shape(transposed)[0] > 2:\n            c = transposed[2].numpy()\n        else:\n            c = self.time * tf.ones(tf.shape(transposed[0])[0])\n        return ax.scatter(transposed[0].numpy(), transposed[1].numpy(), c=c, s=0.125, cmap='prism')\n    \n    def animate_history (self):\n        fig, ax = plt.subplots()\n        ax.set_xlim(tf.reduce_min(self.history[:,:,0])-10, tf.reduce_max(self.history[:,:,0])+10)\n        ax.set_ylim(tf.reduce_min(self.history[:,:,1])-10, tf.reduce_max(self.history[:,:,1])+10)\n        self.anim = FuncAnimation(fig, lambda i : self.plot_history_frame(i, ax), frames=len(self.history), interval=200)\n        video = self.anim.to_html5_video()\n        html = display.HTML(video)\n        display.display(html)\n        plt.close()\n        \n    def evolve (self, n_steps=1, animate=False, dt=1.0, verbose=True):\n        if verbose:\n            print('Evolving the particles...', flush=True)\n            view = tqdm\n        else:\n            view = (lambda x : x)\n        for i in view(range(n_steps)):\n            \n            indices = tf.cast(tf.range(tf.shape(self.particles)[0]), tf.int32)\n            indices = tf.reshape(\n                tf.stack(\n                    [\n                        indices,\n                        tf.zeros(tf.shape(self.particles)[0], dtype=tf.int32),\n                        indices,\n                        tf.ones(tf.shape(self.particles)[0], dtype=tf.int32)\n                    ],\n                    axis=1\n                ),\n                [-1,2,2]\n            )\n            \n            self.particles = tf.tensor_scatter_nd_add(\n                self.particles,\n                indices,\n                dt * self.all_forces()\n            )\n            \n            if self.store_history:\n                self.history = tf.concat([self.history, tf.expand_dims(copy(self.particles.numpy()), axis=0)], axis=0)\n            self.time += dt\n        if animate:\n            self.animate_history()","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:04:20.716779Z","iopub.execute_input":"2022-01-17T13:04:20.717043Z","iopub.status.idle":"2022-01-17T13:04:47.835959Z","shell.execute_reply.started":"2022-01-17T13:04:20.717008Z","shell.execute_reply":"2022-01-17T13:04:47.834914Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# trying out some different force prescriptions here..\n\ndef force_axis (sparse_distance, sparse_sqr_distance, force_radius, norm=0.05, attractive=1.0, repulsive=1.0):\n    distance = tf.reshape(sparse_distance.values, [-1,2])\n    sqr_distance = sparse_sqr_distance.values\n    part0_indices = sparse_sqr_distance.indices[:,0]\n    # calculate the number of neighbours\n    n_neighbours = tf.SparseTensor(\n        indices = sparse_sqr_distance.indices, \n        values = tf.ones(tf.shape(sparse_sqr_distance.values), dtype=tf.float16),\n        dense_shape = sparse_sqr_distance.dense_shape\n    )\n    n_neighbours = tf.sparse.reduce_sum(n_neighbours, axis=1)\n    n_neighbours = tf.gather_nd(n_neighbours, tf.expand_dims(part0_indices,-1))\n    # calculate the local mean direction of neighbours\n    direction = distance / tf.expand_dims(tf.sqrt(tf.reduce_sum(distance**2,-1)),-1)\n    sparse_direction = tf.SparseTensor(\n        indices = sparse_distance.indices,\n        values = tf.reshape(direction, [-1]),\n        dense_shape = sparse_distance.dense_shape\n    )\n    axis = tf.sparse.reduce_sum(sparse_direction, axis=1)\n    axis = tf.reshape(axis, [-1,2])\n    axis = tf.gather_nd(axis, tf.expand_dims(part0_indices,-1))\n    axis /= tf.sqrt(tf.reduce_sum(axis**2))\n    # have the axis always point at positive x\n    axis = tf.where(\n        tf.expand_dims(axis[:,0] > 0,-1),\n        axis,\n        -axis\n    )\n    axis_perp = axis[:,::-1]\n    # the particles will feel the strongest force towards particles in the mean direction\n    force_proposed = (\n        #+ 0.5*attractive / ((1.0 + tf.abs(tf.reduce_sum(axis_perp*distance, axis=1))**2.0)**4 * sqr_distance)\n        + attractive / ((1.0 + tf.abs(tf.reduce_sum(axis_perp*distance, axis=1))**2)**2 * sqr_distance) / 25**(n_neighbours/force_radius**2)\n    )\n    # and will be repelled from particles perpendicular to the mean direction\n    force_proposed += (\n        - repulsive * tf.sin(sqr_distance*0.5*np.pi/force_radius**2)\n            * tf.abs(tf.reduce_sum(axis_perp*distance, axis=1))\n            / ((0.2*force_radius + tf.abs(tf.reduce_sum(axis_perp*distance, axis=1))**3) * sqr_distance)\n    )\n    force_proposed = distance * tf.expand_dims(norm*force_proposed, axis=-1)\n    # push the particles along the axis to avoid them being stuck in threads\n    #force_proposed += norm * tf.expand_dims(tf.abs(tf.reduce_sum(axis*distance, axis=-1)), axis=-1) * axis\n    force_value = tf.reduce_sum(force_proposed**2, axis=1)\n    border_force = tfp.stats.percentile(force_value,25)\n    force_proposed = tf.where(\n        tf.expand_dims(force_value > border_force, axis=-1),\n        force_proposed,\n        #tf.expand_dims(border_force,-1)*axis\n        1.0*axis\n    )\n    # apply limits\n    force_max = tf.sqrt(sqr_distance) / n_neighbours\n    force_value = tf.reduce_sum(force_proposed**2, axis=1)\n    force_direction = force_proposed / tf.expand_dims(force_value, axis=-1)\n    force_proposed = tf.where(\n        tf.expand_dims(tf.abs(force_value) < force_max, -1),\n        force_proposed,\n        norm * force_direction * tf.expand_dims(force_max,-1)\n    )\n    # reshape to sparse representation\n    force_proposed = tf.reshape(force_proposed, [-1])\n    return force_proposed","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:04:47.837367Z","iopub.execute_input":"2022-01-17T13:04:47.838623Z","iopub.status.idle":"2022-01-17T13:04:47.856539Z","shell.execute_reply.started":"2022-01-17T13:04:47.838585Z","shell.execute_reply":"2022-01-17T13:04:47.855745Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Well, let's try this\nif not submission_mode:\n    test_size = 20\n    simulator = ParticleSimulator(\n        tf.cast(\n            tf.concat([\n                tf.concat([np.expand_dims(np.linspace(0,test_size, test_size), axis=1),np.expand_dims(np.zeros([test_size,]),axis=1)], axis=1),\n                #tf.concat([np.expand_dims(np.linspace(0,20, 20), axis=1),np.expand_dims(np.ones([20,]),axis=1)], axis=1),\n                tf.concat([np.expand_dims(np.sin(np.pi*2./3.)*np.linspace(0,test_size, test_size), axis=1),np.expand_dims(np.cos(np.pi*2./3.)*np.linspace(0,test_size, test_size),axis=1)], axis=1),\n                tf.concat([np.expand_dims(np.sin(np.pi*4./3.)*np.linspace(0,test_size, test_size), axis=1),np.expand_dims(np.cos(np.pi*4./3.)*np.linspace(0,test_size, test_size),axis=1)], axis=1),\n            ], axis=0),\n            tf.float32\n        ),\n        force=force_axis,\n        store_history=True,\n        force_radius=10.\n    )\n    simulator.plot()\n    simulator.evolve(animate=True, n_steps=simulation_epochs, dt=simulation_dt)","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:04:47.859753Z","iopub.execute_input":"2022-01-17T13:04:47.85994Z","iopub.status.idle":"2022-01-17T13:04:47.871571Z","shell.execute_reply.started":"2022-01-17T13:04:47.859917Z","shell.execute_reply":"2022-01-17T13:04:47.870738Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"This looks pretty good, if not perfect. The pixels detach at network nodes, but otherwise stay together. Let's try on our dataset.\n\nThe full mask has a very large number of pixels (N~50000), which will not fit as a particle simulator object (with N^2 matrices) in the GPU memory. Thus, let's limit our test to only part of the map.","metadata":{}},{"cell_type":"code","source":"if not submission_mode:\n    plt.figure(figsize=(12,4))\n\n    # Let's plot the true segmentation, so that we know what we're looking for...\n    myid = astro_train_csv.id.unique()[5]\n    astro_train_csv[astro_train_csv.id == myid].head()\n\n    pixels_true = []\n    cl = 0\n    for annotation in tqdm(tf.random.shuffle( # shuffle helps to make adjacent cells different color\n        astro_train_csv[astro_train_csv.id == myid].annotation)\n                          ):\n        mask_here = rle_decode_tf(annotation, shape=img.shape)\n        pixels_true_here = tf.where(tf.squeeze(mask_here) > 0)\n        pixels_true_here = tf.concat([pixels_true_here, cl*np.ones([pixels_true_here.shape[0],1])], axis=1)\n        cl += 1\n        pixels_true.append(pixels_true_here)\n    pixels_true = tf.concat(pixels_true, axis=0)\n\n    x = tf.transpose(pixels_true)[0]\n    y = tf.transpose(pixels_true)[1]\n    c = tf.transpose(pixels_true)[2]\n\n    plt.subplot(121)\n    plt.title('Full mask')\n    plt.scatter(x, y, c=c, s=0.125, cmap='prism')\n\n    # limit to a test region\n    mask = tf.logical_and(x > 300, y < 400).numpy()\n    print(pixels_true.shape)\n    pixels_true = pixels_true.numpy()[mask,:]\n    print(pixels_true.shape)\n\n    x = tf.transpose(pixels_true)[0]\n    y = tf.transpose(pixels_true)[1]\n    c = tf.transpose(pixels_true)[2]\n\n    plt.subplot(122)\n    plt.title('Cropped fragment')\n    plt.scatter(x, y, c=c, s=0.125, cmap='prism')\n\n    plt.show()\n    plt.close()","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:04:47.877627Z","iopub.execute_input":"2022-01-17T13:04:47.877902Z","iopub.status.idle":"2022-01-17T13:04:47.88967Z","shell.execute_reply.started":"2022-01-17T13:04:47.877875Z","shell.execute_reply":"2022-01-17T13:04:47.88875Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if not submission_mode:\n    plt.scatter(x, y, s=0.125)\n    plt.show()\n    plt.close()","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:04:47.891136Z","iopub.execute_input":"2022-01-17T13:04:47.891671Z","iopub.status.idle":"2022-01-17T13:04:47.900905Z","shell.execute_reply.started":"2022-01-17T13:04:47.891633Z","shell.execute_reply":"2022-01-17T13:04:47.900258Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# some adjustments to avoid OOM errors\nphysical_devices = tf.config.list_physical_devices('GPU')\ntry:\n    tf.config.experimental.set_memory_growth(physical_devices[0], True)\nexcept:\n    # Invalid device or cannot modify virtual devices once initialized.\n    pass","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:04:47.902091Z","iopub.execute_input":"2022-01-17T13:04:47.902515Z","iopub.status.idle":"2022-01-17T13:04:47.908597Z","shell.execute_reply.started":"2022-01-17T13:04:47.902479Z","shell.execute_reply":"2022-01-17T13:04:47.907803Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# now evolve our cropped fragment\nif not submission_mode:\n    simulator = ParticleSimulator(\n        tf.cast(pixels_true, tf.float16),\n        force=force_axis,\n        store_history=True,\n        force_radius=25\n    )\n    simulator.plot()\n    simulator.evolve(animate=True, n_steps=simulation_epochs, dt=simulation_dt)","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:04:47.910041Z","iopub.execute_input":"2022-01-17T13:04:47.910373Z","iopub.status.idle":"2022-01-17T13:04:47.917877Z","shell.execute_reply.started":"2022-01-17T13:04:47.91034Z","shell.execute_reply":"2022-01-17T13:04:47.917188Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Now we can use this evolved picture to cluster our point cloud. Let us compare 3 cases:\n - using the original recognition map (the NN output) only,\n - using the evolved map (the final state from our particle simulator) only,\n - using both the original and the evolved map as features.","metadata":{}},{"cell_type":"code","source":"# limit the clusters to only those with at least 10 members\ndef limit_labels (label_list, min_members=10, verbose=True):\n    label_list = np.array(label_list)\n    classes = {c:tf.reduce_sum(tf.cast(label_list == c, tf.int32)).numpy() for c in np.unique(label_list)}\n    classes = {a:(classes[a] if classes[a] > 10 else -1) for a in classes.keys()}\n    labels = [classes[x] for x in label_list]\n    if verbose:\n        print(len(np.unique(list(classes.values()))))\n    return labels","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:04:47.919206Z","iopub.execute_input":"2022-01-17T13:04:47.919522Z","iopub.status.idle":"2022-01-17T13:04:47.928069Z","shell.execute_reply.started":"2022-01-17T13:04:47.919487Z","shell.execute_reply":"2022-01-17T13:04:47.92737Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Now let's try classical clustering algorithms on this point cloud\nfrom sklearn.cluster import DBSCAN\ndbscan_eps = 6\nif not submission_mode:\n\n    plt.figure(figsize=(12,16))\n\n    # target classification\n    plt.subplot(321)\n    plt.title('TARGET (%i)' % len(np.unique(c)))\n    plt.scatter(x, y, c=c, s=0.125, cmap='prism')\n\n    # classification without evolution\n    clusterer = DBSCAN(eps=1, min_samples=10)\n    clusterer.fit(pixels_true[:,:2])\n\n    # discard classes below 10 pixels\n    labels = limit_labels(clusterer.labels_)\n\n    plt.subplot(322)\n    plt.title('No evolution (%i)' % len(np.unique(clusterer.labels_)))\n    plt.scatter(x, y, c=labels, s=0.125, cmap='prism')\n\n    # just evolved\n\n    clusterer = DBSCAN(eps=dbscan_eps, min_samples=10)\n    clusterer.fit(simulator.particles[:,:2])\n\n    # discard classes below 10 pixels\n    labels = limit_labels(clusterer.labels_)\n\n    plt.subplot(323)\n    plt.title('Evolved only (%i)' % len(np.unique(clusterer.labels_)))\n    plt.scatter(x, y, c=labels, s=0.125, cmap='prism')\n\n    # for better clustering, combine the original positions and the evolved ones\n    features = tf.concat([pixels_true[:,:2], simulator.particles[:,:2]], axis=1)\n\n    clusterer = DBSCAN(eps=dbscan_eps, min_samples=10)\n    clusterer.fit(features)\n\n    # discard classes below 10 pixels\n    labels = limit_labels(clusterer.labels_)\n\n    plt.subplot(324)\n    plt.title('Original + evolved (%i)' % len(np.unique(labels)))\n    plt.scatter(x, y, c=labels, s=0.125, cmap='prism')\n\n    # let's try optics too, just out of curiosity\n    from sklearn.cluster import OPTICS\n    clusterer = OPTICS(xi=0.3) #(min_samples=20, xi=1.0, min_cluster_size=5)\n\n    clusterer.fit(features)\n\n    # discard classes below 10 pixels\n    labels = limit_labels(clusterer.labels_)\n\n    plt.subplot(325)\n    plt.title('Original + evolved (%i), OPTICS' % len(np.unique(labels)))\n    plt.scatter(x, y, c=labels, s=0.125, cmap='prism')\n\n    plt.show()\n    plt.close()","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:04:47.92985Z","iopub.execute_input":"2022-01-17T13:04:47.93021Z","iopub.status.idle":"2022-01-17T13:04:48.254101Z","shell.execute_reply.started":"2022-01-17T13:04:47.930178Z","shell.execute_reply":"2022-01-17T13:04:48.253397Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"The original+evolved DBSCAN model perfomes best here. The result is not perfect, but with some tuning, this could do great!\n\nSince the trick with particle simulation paid off, let us now apply it to the full map. Since we can't hold all the pixels in memory at the same time, let us use cropouts. We will proceed as follows:\n - split the full binary detection map into a number of cropouts, where each single one can be held in GPU memory at a time -- note that they should overlap each other!,\n - evolve the cropouts with the particle simulator in sequence, by a small time step at a time,\n - repeat until the cells are sufficiently separated.","metadata":{}},{"cell_type":"code","source":"# recalculate the full list of pixels\nif not submission_mode:\n    pixels_true = []\n    cl = 0\n    for annotation in tqdm(tf.random.shuffle( # shuffle helps to make adjacent cells different color\n        astro_train_csv[astro_train_csv.id == myid].annotation)\n                          ):\n        mask_here = rle_decode_tf(annotation, shape=img.shape)\n        pixels_true_here = tf.where(tf.squeeze(mask_here) > 0)\n        pixels_true_here = tf.concat([pixels_true_here, cl*np.ones([pixels_true_here.shape[0],1])], axis=1)\n        cl += 1\n        pixels_true.append(pixels_true_here)\n    pixels_true = tf.concat(pixels_true, axis=0)\n\n    x = tf.transpose(pixels_true)[0]\n    y = tf.transpose(pixels_true)[1]\n    c = tf.transpose(pixels_true)[2]","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:04:48.257182Z","iopub.execute_input":"2022-01-17T13:04:48.257435Z","iopub.status.idle":"2022-01-17T13:04:48.266778Z","shell.execute_reply.started":"2022-01-17T13:04:48.257409Z","shell.execute_reply":"2022-01-17T13:04:48.266001Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class ParticleSimulatorGroup ():\n    \n    def __init__ (\n        self, particles_init,\n        force, force_radius=20., force_kwargs={}, store_history=False, # ParticleSimulator arguments\n        max_px=10000, # the maximum number of particles per simulator instance\n        overlap=0.25, # fractional overlap between cropped images\n    ):\n        # read the properties of the particle list\n        self.time = 0.0\n        self.particles = tf.cast(particles_init, tf.float16)\n        self.n_particles = tf.shape(particles_init)[0].numpy()\n        transposed = tf.transpose(self.particles[:,:2])\n        self.x = transposed[0]\n        self.y = transposed[1]\n        self.xmin, self.xmax = tf.reduce_min(self.x).numpy(), tf.reduce_max(self.x).numpy()\n        self.xlen = self.xmax - self.xmin\n        self.ymin, self.ymax = tf.reduce_min(self.y).numpy(), tf.reduce_max(self.y).numpy()\n        self.ylen = self.ymax - self.ymin\n        # decide how to split the pixel list\n        if self.n_particles < max_px: # use a single simulator\n            self.n_simulators = 1\n            self.split_borders = [[self.xmin, self.xmax, self.ymin, self.ymax],]\n            self.split = [tf.range(self.n_particles),] # index list\n            self.overlap = 0.\n        else:\n            # split, ensuring significant overlap between cropped images\n            self.n_simulators = 2 * self.n_particles / max_px\n            # bring to the closest larger square of int to make it easier to split\n            sqrt_n_simulators = np.sqrt(self.n_simulators)\n            self.n_simulators = sqrt_n_simulators**2\n            # calculate the split boundaries\n            self.overlap = overlap\n            self.split_borders = [[\n                self.xmin + (i[0]  -self.overlap) * self.xlen/sqrt_n_simulators,\n                self.xmin + (i[0]+1+self.overlap) * self.xlen/sqrt_n_simulators,\n                self.ymin + (i[1]  -self.overlap) * self.ylen/sqrt_n_simulators,\n                self.ymin + (i[1]+1+self.overlap) * self.ylen/sqrt_n_simulators,\n            ] for i in np.reshape(np.transpose(np.meshgrid(np.arange(sqrt_n_simulators), np.arange(sqrt_n_simulators), indexing='ij')),[-1,2])]\n            # now, gather the indices in each border set\n            self.split = [\n                tf.where(\n                    tf.logical_and(\n                        tf.logical_and(\n                            self.particles[:,0] > b[0], self.particles[:,0] < b[1]\n                        ),\n                        tf.logical_and(\n                            self.particles[:,1] > b[2], self.particles[:,1] < b[3]\n                        )\n                    )\n                ) for b in self.split_borders\n            ]\n        # create particle simulators -- one for each crop-out\n        self.simulators = [\n            ParticleSimulator(\n                tf.gather_nd(self.particles, indices),\n                force=force, force_radius=force_radius, force_kwargs=force_kwargs, store_history=store_history\n            ) for indices in self.split\n        ]\n        # prepare for history storage, if needed\n        self.store_history = store_history\n        if self.store_history:\n            self.history = tf.expand_dims(copy(self.particles.numpy()), axis=0)\n        else:\n            self.history = []\n            \n    def plot (self):\n        if len(self.transposed) > 2:\n            c = self.transposed[2]\n        else:\n            c = None\n        plt.scatter(self.transposed[0], self.transposed[1], c=c, s=0.125, cmap='prism')\n        plt.show()\n        plt.close()\n        \n    def plot_history_frame (self, i, ax):\n        xlim, ylim = ax.get_xlim(), ax.get_ylim()\n        ax.clear()\n        ax.set_xlim(xlim)\n        ax.set_ylim(ylim)\n        transposed = tf.transpose(self.history[i])\n        if tf.shape(transposed)[0] > 2:\n            c = transposed[2].numpy()\n        else:\n            c = self.time * tf.ones(tf.shape(transposed[0])[0])\n        return ax.scatter(transposed[0].numpy(), transposed[1].numpy(), c=c, s=0.125, cmap='prism')\n    \n    def animate_history (self):\n        fig, ax = plt.subplots()\n        ax.set_xlim(tf.reduce_min(self.history[:,:,0])-10, tf.reduce_max(self.history[:,:,0])+10)\n        ax.set_ylim(tf.reduce_min(self.history[:,:,1])-10, tf.reduce_max(self.history[:,:,1])+10)\n        self.anim = FuncAnimation(fig, lambda i : self.plot_history_frame(i, ax), frames=len(self.history), interval=200)\n        video = self.anim.to_html5_video()\n        html = display.HTML(video)\n        display.display(html)\n        plt.close()\n        \n    def evolve (self, n_steps=1, animate=False, dt=1.0, verbose=True):\n        if verbose:\n            print('Evolving the particles...', flush=True)\n        for i in tqdm(range(n_steps)):\n            \n            # evolve each particle simulator by one time step, in order\n            for idx in range(tf.cast(self.n_simulators, tf.int32)):\n                simulator = self.simulators[idx]\n                indices = self.split[idx]\n                simulator.evolve(\n                    n_steps=1,\n                    animate=False,\n                    dt=dt,\n                    verbose=False\n                )\n                self.particles = tf.tensor_scatter_nd_update(\n                    self.particles,\n                    indices,\n                    simulator.particles\n                )\n            \n            # save history if needed\n            if self.store_history:\n                self.history = tf.concat([self.history, tf.expand_dims(copy(self.particles.numpy()), axis=0)], axis=0)\n            self.time += dt\n        if animate:\n            self.animate_history()","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:04:48.269513Z","iopub.execute_input":"2022-01-17T13:04:48.270006Z","iopub.status.idle":"2022-01-17T13:04:48.301521Z","shell.execute_reply.started":"2022-01-17T13:04:48.269964Z","shell.execute_reply":"2022-01-17T13:04:48.30054Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if not submission_mode:\n    test = ParticleSimulatorGroup(\n        pixels_true,\n        force=force_axis,\n        store_history=True,\n        force_radius=25\n    )\n    print(test.simulators[3].n_part)","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:04:48.30382Z","iopub.execute_input":"2022-01-17T13:04:48.304025Z","iopub.status.idle":"2022-01-17T13:04:48.312764Z","shell.execute_reply.started":"2022-01-17T13:04:48.304Z","shell.execute_reply":"2022-01-17T13:04:48.312081Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if not submission_mode:\n    test.evolve(n_steps=simulation_epochs, animate=True, dt=simulation_dt)","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:04:48.314175Z","iopub.execute_input":"2022-01-17T13:04:48.31446Z","iopub.status.idle":"2022-01-17T13:04:48.321017Z","shell.execute_reply.started":"2022-01-17T13:04:48.314424Z","shell.execute_reply":"2022-01-17T13:04:48.320275Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Now let's try classical clustering algorithms on this point cloud\nfrom sklearn.cluster import DBSCAN\ndbscan_eps = 6\nastro_clusterer = DBSCAN(eps=dbscan_eps, min_samples=10)\n\nif not submission_mode:\n    plt.figure(figsize=(12,16))\n\n    # target classification\n    plt.subplot(321)\n    plt.title('TARGET (%i)' % len(np.unique(c)))\n    plt.scatter(x, y, c=c, s=0.125, cmap='prism')\n\n    # classification without evolution\n    clusterer = DBSCAN(eps=1, min_samples=10)\n    clusterer.fit(pixels_true[:,:2])\n\n    # discard classes below 10 pixels\n    labels = limit_labels(clusterer.labels_)\n\n    plt.subplot(322)\n    plt.title('No evolution (%i)' % len(np.unique(clusterer.labels_)))\n    plt.scatter(x, y, c=labels, s=0.125, cmap='prism')\n\n    # just evolved\n\n    clusterer = DBSCAN(eps=dbscan_eps, min_samples=10)\n    clusterer.fit(test.particles[:,:2])\n\n    # discard classes below 10 pixels\n    labels = limit_labels(clusterer.labels_)\n\n    plt.subplot(323)\n    plt.title('Evolved only (%i)' % len(np.unique(clusterer.labels_)))\n    plt.scatter(x, y, c=labels, s=0.125, cmap='prism')\n\n    # for better clustering, combine the original positions and the evolved ones\n    features = tf.concat([pixels_true[:,:2], tf.cast(test.particles[:,:2], tf.int64)], axis=1)\n\n    astro_clusterer.fit(features)\n\n    # discard classes below 10 pixels\n    labels = limit_labels(astro_clusterer.labels_)\n\n    plt.subplot(324)\n    plt.title('Original + evolved (%i)' % len(np.unique(labels)))\n    plt.scatter(x, y, c=labels, s=0.125, cmap='prism')\n\n    # let's try optics too, just out of curiosity\n    from sklearn.cluster import OPTICS\n    clusterer = OPTICS(xi=0.3) #(min_samples=20, xi=1.0, min_cluster_size=5)\n\n    clusterer.fit(features)\n\n    # discard classes below 10 pixels\n    labels = limit_labels(clusterer.labels_)\n\n    plt.subplot(325)\n    plt.title('Original + evolved (%i), OPTICS' % len(np.unique(labels)))\n    plt.scatter(x, y, c=labels, s=0.125, cmap='prism')\n\n    plt.show()\n    plt.close()","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:47:05.793565Z","iopub.execute_input":"2022-01-17T13:47:05.793815Z","iopub.status.idle":"2022-01-17T13:47:05.807664Z","shell.execute_reply.started":"2022-01-17T13:47:05.793788Z","shell.execute_reply":"2022-01-17T13:47:05.806711Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"All right, let's put it all together into a classifier.\n - Input: raw image with astro cells\n - Processing:\n   - preprocessing (color adjustments, etc.),\n   - detection of cell structures with the NN trained above,\n   - transformation of the resulting mask with particle simulators,\n   - mask splitting into separate cells with DBSCAN,\n   - conversion into competition pixel format.\n - Output:\n   - imaged clustering map (side-by-side with the original image)\n   - the clustering map in the competition pixel format, as a string","metadata":{}},{"cell_type":"code","source":"def rle_encode (mask, classes=[1,], img_id='noid'):\n    '''Encode an image using run-length-encoding'''\n    mask_here = tf.squeeze(mask).numpy()\n    shape = mask_here.shape\n    \n    rle_encoded = ''\n    for cl in tqdm(classes):\n        j,i = np.where(mask_here == cl)\n        i = np.append(i, -1)\n\n        kernel = np.array([1,-1])\n        starting_idxs = np.where(\n            np.convolve(i,kernel, mode='valid') != 1\n        )[0] + 1\n        starting_idxs = np.insert(starting_idxs,0,0)\n        starting_px = j[starting_idxs[:-1]]*shape[1] + i[starting_idxs[:-1]] + 1\n        run_lengths = np.convolve(starting_idxs,kernel, mode='valid')\n\n        rle_encoded_here = np.transpose(\n            [starting_px,\n            run_lengths]\n        ).flatten()\n        rle_encoded += '\\n' + img_id + ',' + ' '.join(rle_encoded_here.astype(str))\n    \n    return rle_encoded[1:]\n\n# test that the rle encoding works properly\nif not submission_mode:\n    annotation = astro_train_csv[astro_train_csv.id == myid].annotation.iloc[0]\n    rle_encoded = rle_encode(rle_decode_tf(annotation, shape=img.shape))\n    rle_encoded[5:] == annotation","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:04:48.339378Z","iopub.execute_input":"2022-01-17T13:04:48.340136Z","iopub.status.idle":"2022-01-17T13:04:48.350014Z","shell.execute_reply.started":"2022-01-17T13:04:48.3401Z","shell.execute_reply":"2022-01-17T13:04:48.349342Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class AstroModel:\n    \n    def __init__ (\n        self,\n        preprocess=False,\n        detection_model=astro_model,\n        detection_threshold=0.6,\n        simulator_force=force_axis,\n        simulator_force_radius=25,\n        simulator_force_kwargs={},\n        simulator_steps=simulation_epochs,\n        simulator_dt=simulation_dt,\n        clusterer=DBSCAN(eps=dbscan_eps, min_samples=10)\n    ):\n        self.preprocess = preprocess\n        self.detection_model = detection_model\n        self.detection_threshold = detection_threshold\n        self.simulator_force = simulator_force\n        self.simulator_force_radius = simulator_force_radius\n        self.simulator_force_kwargs = simulator_force_kwargs\n        self.simulator_steps = simulator_steps\n        self.simulator_dt = simulator_dt\n        self.clusterer = astro_clusterer\n    \n    def predict (self, img, img_id='noid', mask_pixels=[], verbose=True):\n        # image preprocessing\n        if verbose:\n            print(' - predicting an astro-cell image segmentation:')\n            print('   - image preprocessing.. ', end='', flush=True)\n            fig = plt.figure(figsize=(12,6))\n            plt.subplot(121)\n            plt.imshow(1.*img, cmap=\"seismic\")\n        if self.preprocess:\n            img = img_color_rescale(img)\n        if verbose:\n            plt.subplot(122)\n            plt.imshow(img, cmap=\"seismic\")\n            plt.show()\n            plt.close()\n            print('done.', flush=True)\n        # use the NN to detect cell structures\n        if verbose:\n            print('   - detecting cell pixels with a NN.. ', end='', flush=True)\n        pred = self.detection_model.predict(np.array([img,]))[0]\n        pixels_true = tf.where(pred > self.detection_threshold)\n        x = tf.transpose(pixels_true)[0]\n        y = tf.transpose(pixels_true)[1]\n        if verbose:\n            print('done.', flush=True)\n            plt.figure(figsize=(12,6))\n            plt.subplot(121)\n            plt.imshow(pred)\n            plt.colorbar()\n            plt.subplot(122)\n            plt.scatter(y,x, s=0.125)\n            plt.show()\n            plt.close()\n        # transform with particle simulators\n        if verbose:\n            print('   - transforming with particle simulators.. ', end='', flush=True)\n        self.psg = ParticleSimulatorGroup(\n            pixels_true,\n            force=self.simulator_force,\n            store_history=verbose,\n            force_radius=self.simulator_force_radius,\n            force_kwargs=self.simulator_force_kwargs\n        )\n        self.psg.evolve(n_steps=self.simulator_steps, animate=verbose, dt=self.simulator_dt)\n        if verbose:\n            print('done.', flush=True)\n        # cluster with DBSCAN\n        if verbose:\n            print('   - clustering.. ', end='', flush=True)\n        features = tf.concat([pixels_true[:,:2], tf.cast(self.psg.particles[:,:2], tf.int64)], axis=1)\n        self.clusterer.fit(features)\n        labels = np.array(limit_labels(self.clusterer.labels_))\n        if verbose:\n            if len(mask_pixels) > 0:\n                fig = plt.figure(figsize=(12,6))\n                plt.subplot(121)\n                mask_pixels = tf.transpose(mask_pixels)\n                xm = mask_pixels[0]\n                ym = mask_pixels[1]\n                cmask = mask_pixels[2]\n                plt.scatter(ym,xm,c=cmask, s=0.125, cmap='prism')\n                plt.subplot(122)\n            else:\n                fig = plt.figure(figsize=(6,6))\n            plt.scatter(x, y, c=labels, s=0.125, cmap='prism')\n            plt.show()\n            plt.close()\n            print('done.', flush=True)\n        # convert to the competitioon pixel number format\n        if verbose:\n            print('   - converting to RLE string.. ', end='', flush=True)\n        mask_pred = np.squeeze(np.zeros(img.shape))\n        mask_pred[tuple(np.transpose(pixels_true[:,:2]))] = labels\n        rle_encoded = rle_encode(mask_pred, classes=np.unique(labels), img_id=img_id)\n        if verbose:\n            print('done.', flush=True)\n        return rle_encoded","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:04:48.35318Z","iopub.execute_input":"2022-01-17T13:04:48.353416Z","iopub.status.idle":"2022-01-17T13:04:48.375324Z","shell.execute_reply.started":"2022-01-17T13:04:48.353392Z","shell.execute_reply":"2022-01-17T13:04:48.374599Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if not submission_mode:\n    img = astro_train_imgs[5]\n    mask = astro_train_masks[5]\n\n    model = AstroModel()\n    encoded = model.predict(img)\n    print(encoded[:50])","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:04:48.376411Z","iopub.execute_input":"2022-01-17T13:04:48.377104Z","iopub.status.idle":"2022-01-17T13:04:48.385905Z","shell.execute_reply.started":"2022-01-17T13:04:48.37704Z","shell.execute_reply":"2022-01-17T13:04:48.385111Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"When improving the model later, it will be very helpful to have our own implementation of the competition score. Let us prepare one here.","metadata":{}},{"cell_type":"code","source":"def single_intersection (b1,e1,b2,e2):\n    result = min(e1,e2) - max(b1,b2)\n    return max(result, 0)\ndef rle_intersection_over_union (rle1, rle2):\n    '''Calculate IoU for rle-encoded data'''\n    start_px1, end_px1 = tuple(np.transpose(np.array(rle1.split()).astype(np.int).reshape([-1,2])))\n    end_px1 += start_px1\n    start_px2, end_px2 = tuple(np.transpose(np.array(rle2.split()).astype(np.int).reshape([-1,2])))\n    end_px2 += start_px2\n    intersections = [\n        single_intersection(\n            start_px1[i1],\n            end_px1[i1],\n            start_px2[i2],\n            end_px2[i2]\n        ) for i1,i2 in zip(range(len(start_px1)), range(len(start_px2)))\n    ]\n    intersection = np.sum(intersections)\n    union = np.sum(end_px1-start_px1) + np.sum(end_px2-start_px2) - intersection\n    return intersection / union\n\ndef comp_score (rle_pred, rle_true):\n    '''Calculate competition score, https://www.kaggle.com/c/sartorius-cell-instance-segmentation/overview/evaluation'''\n    # key annotations with their image ids\n    data_true = pd.DataFrame([[x.split(',')[0], ' '.join(x.split(',')[1].split())] for x in rle_true.split('\\n')], columns=['id','annotation'])\n    data_pred = pd.DataFrame([[x.split(',')[0], ' '.join(x.split(',')[1].split())] for x in rle_pred.split('\\n')], columns=['id','annotation'])\n    # loop over all images\n    global_precision = []\n    for img_id in tqdm(data_true.id.unique()):\n        # build an IoU array\n        ious = np.array([\n            [rle_intersection_over_union(rle1, rle2) \\\n            for rle2 in data_true[data_true.id == img_id].annotation] \\\n            for rle1 in np.array(data_pred[data_pred.id == img_id].annotation)\n        ])\n        # sweep over thresholds\n        precisions = []\n        for threshold in np.arange(0.5,0.99,0.05):\n            matches = (ious > threshold)\n            true_positives = np.sum(np.any(matches, axis=0))\n            false_positives = matches.shape[0] - true_positives\n            false_negatives = matches.shape[1] - np.sum(np.any(matches, axis=1))\n            precisions.append(\n                true_positives / (true_positives + false_positives + false_negatives)\n            )\n        global_precision.append(np.mean(precisions))\n    return np.mean(global_precision)\n\n# test that it works cortectly\nif not submission_mode:\n    print(\"Should be 1: \", comp_score(encoded, encoded))","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:04:48.389416Z","iopub.execute_input":"2022-01-17T13:04:48.389782Z","iopub.status.idle":"2022-01-17T13:04:48.406835Z","shell.execute_reply.started":"2022-01-17T13:04:48.389753Z","shell.execute_reply":"2022-01-17T13:04:48.406128Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# let's see how our astro model performs with respect to the competition score\nif not submission_mode:\n    \n    # choose the indices to test on\n    val_ids = np.random.choice(astro_val_csv.id.unique(), 1)\n    print(val_ids)\n\n    # parse the cortect answers\n    rows_true = astro_val_csv[astro_val_csv.id.isin(val_ids)]\n    encoded_true = ''\n    for idx, row in rows_true[['id', 'annotation']].iterrows():\n        encoded_true += '\\n' + row['id'] + ',' + row['annotation']\n    encoded_true = encoded_true[1:]\n\n    # calculate predictions\n    if True:\n        encoded = ''\n        for myid in val_ids:\n            print('--- Predicting image %s ---' % myid, flush=True)\n            img = tf_img_color_rescale(\n                tf_img_load(myid)\n            ).numpy()\n            # decode the target mask for comparison\n            pixels_true = []\n            cl = 0\n            for annotation in tf.random.shuffle( # shuffle helps to make adjacent cells different color\n                astro_val_csv[astro_val_csv.id == myid].annotation\n            ):\n                mask_here = rle_decode_tf(annotation, shape=img.shape)\n                pixels_true_here = tf.where(tf.squeeze(mask_here) > 0)\n                pixels_true_here = tf.concat([pixels_true_here, cl*np.ones([pixels_true_here.shape[0],1])], axis=1)\n                cl += 1\n                pixels_true.append(pixels_true_here)\n            pixels_true = tf.concat(pixels_true, axis=0)\n            # pass the image through our model\n            model = AstroModel(\n                preprocess=False,\n                detection_model=astro_model,\n                detection_threshold=0.6,\n                simulator_force=force_axis,\n                simulator_force_radius=25,\n                simulator_steps=32,\n                clusterer=DBSCAN(eps=dbscan_eps, min_samples=10)\n            )\n            encoded += '\\n' + model.predict(\n                img,\n                img_id=myid,\n                mask_pixels=pixels_true\n            )\n            print('--- Image %s predicted ---' % myid, flush=True)\n        encoded = encoded[1:]\n\n        score = comp_score(encoded, encoded_true)\n        print('Competition score: %.3f' % score)","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:04:48.408265Z","iopub.execute_input":"2022-01-17T13:04:48.408696Z","iopub.status.idle":"2022-01-17T13:04:48.420767Z","shell.execute_reply.started":"2022-01-17T13:04:48.408657Z","shell.execute_reply":"2022-01-17T13:04:48.420021Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Hmm.. not so great.. Still, let us obtain a full model first -- with the other types of cells included -- before we proceed to any tuning. This way we can compare with a submission score and see whether we calculate it cortectly (also, time is now starting to run out ;), and I can do only so much Kaggle work on weekends).","metadata":{}},{"cell_type":"markdown","source":"-----------\n**CORT-cell segmentation**","metadata":{}},{"cell_type":"code","source":"cort_ids = img_topics[img_topics.perc_cort == 1.0].id.to_numpy()\nprint(len(cort_ids))","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:04:48.42281Z","iopub.execute_input":"2022-01-17T13:04:48.423531Z","iopub.status.idle":"2022-01-17T13:04:48.433618Z","shell.execute_reply.started":"2022-01-17T13:04:48.423496Z","shell.execute_reply":"2022-01-17T13:04:48.432834Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"cort_csv = train_csv[train_csv.id.isin(cort_ids)]\ncort_val_ids = rd.choice(cort_ids, int(val_fraction * len(cort_ids)))\ncort_val_csv = cort_csv[cort_csv.id.isin(cort_val_ids)]\ncort_train_csv = cort_csv[~cort_csv.id.isin(cort_val_ids)]\nprint(len(cort_train_csv), len(cort_val_csv))","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:04:48.434772Z","iopub.execute_input":"2022-01-17T13:04:48.435865Z","iopub.status.idle":"2022-01-17T13:04:48.459518Z","shell.execute_reply.started":"2022-01-17T13:04:48.435767Z","shell.execute_reply":"2022-01-17T13:04:48.458868Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# in principle, we could do image loading as part of the model\n# , but, since there aren't that many of them,\n# it seems more sensible to load them into our RAM memory all at once\nshape = (520,704)\nsave_path = 'cort_train_data.pkl'\nif task_list['cort_recognitionNN'] or task_list['cort_segmentation']:\n    if force_retrain or not os.path.exists(save_path):\n        cort_train_imgs = []\n        cort_train_masks = []\n        for img_id in tqdm(cort_train_csv.id.unique()):\n            cort_train_imgs.append(\n                tf_img_color_rescale(\n                    tf_img_load(img_id)\n                ).numpy()\n            )\n            # could probably make this faster with a tensorflow map, but as-is, a dtype error is returned (can't map a string-type tensor to int)...\n            mask = np.zeros(shape)\n            for ann in cort_train_csv[cort_train_csv.id == img_id].annotation:\n                mask += rle_decode_tf(ann, shape=shape)\n            mask = tf.where(mask > 0, 1,0)\n            cort_train_masks.append(mask.numpy().astype(np.uint8))\n        cort_train_imgs = np.array(cort_train_imgs)\n        cort_train_masks = np.array(cort_train_masks)\n\n        cort_val_imgs = []\n        cort_val_masks = []\n        for img_id in tqdm(cort_val_csv.id.unique()):\n            cort_val_imgs.append(\n                tf_img_color_rescale(\n                    tf_img_load(img_id)\n                ).numpy()\n            )\n            mask = np.zeros(shape)\n            for ann in cort_val_csv[cort_val_csv.id == img_id].annotation:\n                mask += rle_decode_tf(ann, shape=shape)\n            mask = tf.where(mask > 0, 1,0)\n            cort_val_masks.append(mask.numpy().astype(np.uint8))\n        cort_val_imgs = np.array(cort_val_imgs)\n        cort_val_masks = np.array(cort_val_masks)\n        if not submission_mode:\n            with open(save_path, 'wb') as f:\n                pkl.dump((cort_train_imgs, cort_train_masks, cort_val_imgs, cort_val_masks), f)\n    else:\n        with open(save_path, 'rb') as f:\n            cort_train_imgs, cort_train_masks, cort_val_imgs, cort_val_masks = pkl.load(f)\n    print(cort_train_imgs.shape, cort_val_imgs.shape)\n    print(cort_train_masks.shape, cort_val_masks.shape)","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:04:48.461902Z","iopub.execute_input":"2022-01-17T13:04:48.462152Z","iopub.status.idle":"2022-01-17T13:06:19.93064Z","shell.execute_reply.started":"2022-01-17T13:04:48.462121Z","shell.execute_reply":"2022-01-17T13:06:19.929911Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if not submission_mode:\n    plt.imshow(cort_val_masks[2])\n    plt.show()\n    plt.close()","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:06:19.9321Z","iopub.execute_input":"2022-01-17T13:06:19.932373Z","iopub.status.idle":"2022-01-17T13:06:19.937003Z","shell.execute_reply.started":"2022-01-17T13:06:19.932338Z","shell.execute_reply":"2022-01-17T13:06:19.936077Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if not submission_mode:\n    for filename in ['/kaggle/input/sartorius-cell-instance-segmentation/train/' + x + '.png' for x in cort_ids][:3]:\n        plot_image(filename)","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:06:19.938405Z","iopub.execute_input":"2022-01-17T13:06:19.938945Z","iopub.status.idle":"2022-01-17T13:06:19.947509Z","shell.execute_reply.started":"2022-01-17T13:06:19.938907Z","shell.execute_reply":"2022-01-17T13:06:19.946627Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":" - The cells are very small, circular, and very well separated.\n - Within the images, there are lots of similar and dis-similar features in addition to cort-cells.\n \nGiven these observations, a CNN with small, but deep filters could prove useful. Let's try.","metadata":{}},{"cell_type":"code","source":"# build our cort_model\nmodel_name = 'cort_recognitionNN'\ncort_model = AugmentedSequential()\ncort_model.add(keras.Input(\n    shape=(520,704,1)\n))\n# now some vanilla CNN layers\nif False: # val_loss: 0.0142\n    cort_model.add(layers.Conv2D(\n        filters=16, kernel_size=(3,3),\n        activation='relu',\n        padding='same')\n    )\n    cort_model.add(layers.Conv2D(\n        filters=16, kernel_size=(3,3),\n        activation='relu',\n        padding='same')\n    )\n    cort_model.add(layers.Conv2D(\n        filters=16, kernel_size=(3,3),\n        activation='relu',\n        padding='same')\n    )\nelse: # val_loss: 0.0098\n    cort_model.add(layers.Conv2D(\n        filters=16, kernel_size=(3,3),\n        activation='relu',\n        padding='same')\n    )\n    cort_model.add(layers.Conv2D(\n        filters=8, kernel_size=(5,5),\n        activation='relu',\n        padding='same')\n    )\n    cort_model.add(layers.Conv2D(\n        filters=8, kernel_size=(7,7),\n        activation='relu',\n        padding='same')\n    )\n# Apply a dense layer along the axis of the filters, keeping the image size the same\ncort_model.add(layers.Conv2D(\n    filters=8, kernel_size=(1,1),\n    activation='relu',\n    padding='same')\n)\n# OUPTUT -----------------------------\n#cort_model.add(layers.Dropout(0.5))\ncort_model.add(layers.Conv2D(\n    filters=1, kernel_size=(1,1),\n    activation='sigmoid',\n    padding='same')\n)\n\ncort_model.compile(optimizer='adam',\n              loss=balanced_loss,\n                   run_eagerly=True)  #tf.keras.losses.CategoricalCrossentropy(from_logits=False),\n              #metrics=['accuracy']\n\ncort_model.summary()","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:06:19.948971Z","iopub.execute_input":"2022-01-17T13:06:19.949317Z","iopub.status.idle":"2022-01-17T13:06:20.009582Z","shell.execute_reply.started":"2022-01-17T13:06:19.949282Z","shell.execute_reply":"2022-01-17T13:06:20.007315Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if task_list[model_name]:\n    print(model_name)\n    if force_retrain or not os.path.exists(model_name+'.model'):\n        stop_early = tf.keras.callbacks.EarlyStopping(\n            monitor='val_loss',\n            patience=early_stopping_patience,\n            restore_best_weights=True\n        )\n        callbacks = [stop_early,]\n        history = cort_model.fit(\n            tf.cast(cort_train_imgs,tf.float32).numpy(),\n            tf.cast(tf.expand_dims(cort_train_masks, axis=3),tf.float32).numpy(),\n            validation_data=(\n                tf.cast(cort_val_imgs,tf.float32).numpy(),\n                tf.cast(tf.expand_dims(cort_val_masks, axis=3),tf.float32).numpy()),\n            callbacks=callbacks,\n            epochs=epochs_to_train)\n        if not submission_mode:\n            cort_model.save(model_name + '.model')\n            plot_history(history)\n    else:\n        cort_model = keras.models.load_model(\n            model_name+'.model',\n            custom_objects={'balanced_loss':balanced_loss}\n        )","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:06:20.010619Z","iopub.execute_input":"2022-01-17T13:06:20.01086Z","iopub.status.idle":"2022-01-17T13:29:13.175157Z","shell.execute_reply.started":"2022-01-17T13:06:20.010826Z","shell.execute_reply":"2022-01-17T13:29:13.174462Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# make sure augmentation works, check the predictions\ndef test_predictions (model):\n    for i in range(3):\n        X_val, y_val = tf_img_augment_nomodel(\n            cort_val_imgs, \n            tf.expand_dims(cort_val_masks, axis=3).numpy()\n        )\n\n        ex_img = X_val[4].numpy()\n        ex_truemask = y_val[4].numpy()\n\n        # show an example prediction\n        fig = plt.figure(figsize=(16,6), clear=True)\n\n        plt.subplot(131)\n        plt.imshow(ex_img, cmap='seismic')\n        plt.title('Original image')\n\n        plt.subplot(132)\n        ex_mask = model.predict(ex_img.reshape((1,520,704,1)))[0]\n        print(ex_mask.shape)\n        plt.imshow(ex_mask, cmap=\"YlGn_r\")\n        plt.title('Predicted mask')\n        plt.colorbar()\n\n        plt.subplot(133)\n        plt.imshow(ex_truemask, cmap=\"YlGn_r\")\n        plt.title(\"True mask\")\n\n        plt.show()\n        plt.close()\n    \nif task_list[model_name] and not submission_mode:\n    test_predictions(cort_model)","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:29:13.176479Z","iopub.execute_input":"2022-01-17T13:29:13.176706Z","iopub.status.idle":"2022-01-17T13:29:13.184514Z","shell.execute_reply.started":"2022-01-17T13:29:13.176671Z","shell.execute_reply":"2022-01-17T13:29:13.183632Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# First, let's improve the contrast of the network output, to make it easier to segment\nif not submission_mode:\n    img = cort_train_imgs[5]\n    mask = cort_train_masks[5]\n    pred = cort_model.predict(np.array([img,]))[0]\n\n    plt.figure(figsize=(20,6), clear=True)\n    plt.subplot(141)\n    plt.imshow(img, cmap='seismic')\n    plt.title('Original image')\n    plt.subplot(142)\n    plt.imshow(mask, cmap=\"YlGn_r\")\n    plt.title('True mask')\n    plt.subplot(143)\n    plt.imshow(pred, cmap=\"YlGn_r\")\n    plt.title('Prediction')\n    plt.subplot(144)\n    plt.imshow(pred > 0.2, cmap=\"YlGn_r\")\n    plt.title('Threshold')\n    plt.show()\n    plt.close()\n\n    plt.hist(pred.flatten())\n    plt.title('Prediction histogram')\n    plt.show()\n    plt.close()","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:29:13.186872Z","iopub.execute_input":"2022-01-17T13:29:13.187507Z","iopub.status.idle":"2022-01-17T13:29:13.196095Z","shell.execute_reply.started":"2022-01-17T13:29:13.187471Z","shell.execute_reply":"2022-01-17T13:29:13.195407Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"This shold be very easy to split into cells, so particle simulator is probably unnecessary here.","metadata":{}},{"cell_type":"code","source":"# Let's plot the true segmentation, so that we know what we're looking for...\nif not submission_mode:\n    myid = cort_train_csv.id.unique()[5]\n    cort_train_csv[cort_train_csv.id == myid].head()\n\n    pixels_true = []\n    cl = 0\n    for annotation in tqdm(tf.random.shuffle( # shuffle helps to make adjacent cells different color\n        cort_train_csv[cort_train_csv.id == myid].annotation)\n                          ):\n        mask_here = rle_decode_tf(annotation, shape=img.shape)\n        pixels_true_here = tf.where(tf.squeeze(mask_here) > 0)\n        pixels_true_here = tf.concat([pixels_true_here, cl*np.ones([pixels_true_here.shape[0],1])], axis=1)\n        cl += 1\n        pixels_true.append(pixels_true_here)\n    pixels_true = tf.concat(pixels_true, axis=0)\n\n    x = tf.transpose(pixels_true)[0]\n    y = tf.transpose(pixels_true)[1]\n    c = tf.transpose(pixels_true)[2]\n\n    plt.scatter(x, y, c=c, s=0.125, cmap='prism')\n    plt.show()\n    plt.close()","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:29:13.196867Z","iopub.execute_input":"2022-01-17T13:29:13.19707Z","iopub.status.idle":"2022-01-17T13:29:13.207929Z","shell.execute_reply.started":"2022-01-17T13:29:13.197046Z","shell.execute_reply":"2022-01-17T13:29:13.207192Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Now let's try classical clustering algorithms on this point cloud\nfrom sklearn.cluster import DBSCAN\n\nif not submission_mode:\n    pixels = tf.where(tf.squeeze(mask) > 0).numpy()\n    x = tf.transpose(pixels)[0]\n    y = tf.transpose(pixels)[1]\n\n    clusterer = DBSCAN(eps=1, min_samples=1)\n\n    clusterer.fit(pixels)\n    print(len(np.unique(clusterer.labels_)))\n\n    plt.scatter(x, y, c=clusterer.labels_, s=0.125, cmap='prism')\n    plt.show()\n    plt.close()","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:29:13.210181Z","iopub.execute_input":"2022-01-17T13:29:13.21094Z","iopub.status.idle":"2022-01-17T13:29:13.217962Z","shell.execute_reply.started":"2022-01-17T13:29:13.210902Z","shell.execute_reply":"2022-01-17T13:29:13.217209Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Looks good! This completes our model for cort-type cells.","metadata":{}},{"cell_type":"code","source":"class CortModel:\n    \n    def __init__ (\n        self,\n        preprocess=False,\n        detection_model=cort_model,\n        detection_threshold=0.2,\n        clusterer=DBSCAN(eps=1, min_samples=1)\n    ):\n        self.preprocess = preprocess\n        self.detection_model = detection_model\n        self.detection_threshold = detection_threshold\n        self.clusterer = clusterer\n    \n    def predict (self, img, img_id='noid', mask_pixels=[], verbose=True):\n        # image preprocessing\n        if verbose:\n            print(' - predicting an cort-cell image segmentation:')\n            print('   - image preprocessing.. ', end='', flush=True)\n            fig = plt.figure(figsize=(12,6))\n            plt.subplot(121)\n            plt.imshow(1.*img, cmap=\"seismic\")\n        if self.preprocess:\n            img = img_color_rescale(img)\n        if verbose:\n            plt.subplot(122)\n            plt.imshow(img, cmap=\"seismic\")\n            plt.show()\n            plt.close()\n            print('done.', flush=True)\n        # use the NN to detect cell structures\n        if verbose:\n            print('   - detecting cell pixels with a NN.. ', end='', flush=True)\n        pred = self.detection_model.predict(np.array([img,]))[0]\n        pixels_true = tf.where(pred > self.detection_threshold)\n        x = tf.transpose(pixels_true)[0]\n        y = tf.transpose(pixels_true)[1]\n        if verbose:\n            print('done.', flush=True)\n            plt.figure(figsize=(12,6))\n            plt.subplot(121)\n            plt.imshow(pred)\n            plt.colorbar()\n            plt.subplot(122)\n            plt.scatter(x,y, s=0.125)\n            plt.show()\n            plt.close()\n        # cluster with DBSCAN\n        if verbose:\n            print('   - clustering.. ', end='', flush=True)\n        features = pixels_true[:,:2]\n        self.clusterer.fit(features)\n        labels = np.array(limit_labels(self.clusterer.labels_))\n        if verbose:\n            if len(mask_pixels) > 0:\n                fig = plt.figure(figsize=(12,6))\n                plt.subplot(121)\n                mask_pixels = tf.transpose(mask_pixels)\n                xm = mask_pixels[0]\n                ym = mask_pixels[1]\n                cmask = mask_pixels[2]\n                plt.scatter(xm,ym,c=cmask, s=0.125, cmap='prism')\n                plt.subplot(122)\n            else:\n                fig = plt.figure(figsize=(6,6))\n            plt.scatter(x, y, c=labels, s=0.125, cmap='prism')\n            plt.show()\n            plt.close()\n            print('done.', flush=True)\n        # convert to the competitioon pixel number format\n        if verbose:\n            print('   - converting to RLE string.. ', end='', flush=True)\n        mask_pred = np.squeeze(np.zeros(img.shape))\n        mask_pred[tuple(np.transpose(pixels_true[:,:2]))] = labels\n        rle_encoded = rle_encode(mask_pred, classes=np.unique(labels), img_id=img_id)\n        if verbose:\n            print('done.', flush=True)\n        return rle_encoded","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:29:13.219305Z","iopub.execute_input":"2022-01-17T13:29:13.219609Z","iopub.status.idle":"2022-01-17T13:29:44.692838Z","shell.execute_reply.started":"2022-01-17T13:29:13.219577Z","shell.execute_reply":"2022-01-17T13:29:44.692061Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if not submission_mode:\n    img = cort_train_imgs[5]\n    mask = cort_train_masks[5]\n\n    model = CortModel()\n    encoded = model.predict(img)\n    print(encoded[:50])","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:29:44.694597Z","iopub.execute_input":"2022-01-17T13:29:44.694842Z","iopub.status.idle":"2022-01-17T13:29:44.705716Z","shell.execute_reply.started":"2022-01-17T13:29:44.694807Z","shell.execute_reply":"2022-01-17T13:29:44.705077Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# let's see how our cort model performs with respect to the competition score\nif not submission_mode:\n    \n    # choose the indices to test on\n    val_ids = np.random.choice(cort_val_csv.id.unique(), 1)\n    print(val_ids)\n\n    # parse the correct answers\n    rows_true = cort_val_csv[cort_val_csv.id.isin(val_ids)]\n    encoded_true = ''\n    for idx, row in rows_true[['id', 'annotation']].iterrows():\n        encoded_true += '\\n' + row['id'] + ',' + row['annotation']\n    encoded_true = encoded_true[1:]\n\n    # calculate predictions\n    if True:\n        encoded = ''\n        for myid in val_ids:\n            print('--- Predicting image %s ---' % myid, flush=True)\n            img = tf_img_color_rescale(\n                tf_img_load(myid)\n            ).numpy()\n            # decode the target mask for comparison\n            pixels_true = []\n            cl = 0\n            for annotation in tf.random.shuffle( # shuffle helps to make adjacent cells different color\n                cort_val_csv[cort_val_csv.id == myid].annotation\n            ):\n                mask_here = rle_decode_tf(annotation, shape=img.shape)\n                pixels_true_here = tf.where(tf.squeeze(mask_here) > 0)\n                pixels_true_here = tf.concat([pixels_true_here, cl*np.ones([pixels_true_here.shape[0],1])], axis=1)\n                cl += 1\n                pixels_true.append(pixels_true_here)\n            pixels_true = tf.concat(pixels_true, axis=0)\n            # pass the image through our model\n            model = CortModel(\n                preprocess=False,\n                detection_model=cort_model,\n                detection_threshold=0.2,\n                clusterer=DBSCAN(eps=1, min_samples=1)\n            )\n            encoded += '\\n' + model.predict(\n                img,\n                img_id=myid,\n                mask_pixels=pixels_true\n            )\n            print('--- Image %s predicted ---' % myid, flush=True)\n        encoded = encoded[1:]\n\n        score = comp_score(encoded, encoded_true)\n        print('Competition score: %.3f' % score)","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:29:44.708694Z","iopub.execute_input":"2022-01-17T13:29:44.708879Z","iopub.status.idle":"2022-01-17T13:29:44.720378Z","shell.execute_reply.started":"2022-01-17T13:29:44.708858Z","shell.execute_reply":"2022-01-17T13:29:44.719702Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"---------\n**shsy5y-cell segmentation**","metadata":{}},{"cell_type":"code","source":"shsy5y_ids = img_topics[img_topics.perc_shsy5y == 1.0].id.to_numpy()\nprint(len(shsy5y_ids))","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:29:44.723173Z","iopub.execute_input":"2022-01-17T13:29:44.723404Z","iopub.status.idle":"2022-01-17T13:29:44.733204Z","shell.execute_reply.started":"2022-01-17T13:29:44.723376Z","shell.execute_reply":"2022-01-17T13:29:44.732516Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"shsy5y_csv = train_csv[train_csv.id.isin(shsy5y_ids)]\nshsy5y_val_ids = rd.choice(shsy5y_ids, int(val_fraction * len(shsy5y_ids)))\nshsy5y_val_csv = shsy5y_csv[shsy5y_csv.id.isin(shsy5y_val_ids)]\nshsy5y_train_csv = shsy5y_csv[~shsy5y_csv.id.isin(shsy5y_val_ids)]\nprint(len(shsy5y_train_csv), len(shsy5y_val_csv))","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:29:44.735017Z","iopub.execute_input":"2022-01-17T13:29:44.735287Z","iopub.status.idle":"2022-01-17T13:29:44.772786Z","shell.execute_reply.started":"2022-01-17T13:29:44.735252Z","shell.execute_reply":"2022-01-17T13:29:44.771642Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# in principle, we could do image loading as part of the model\n# , but, since there aren't that many of them,\n# it seems more sensible to load them into our RAM memory all at once\nshape = (520,704)\nsave_path = 'shsy5y_train_data.pkl'\nif task_list['shsy5y_recognitionNN'] or task_list['shsy5y_segmentation']:\n    if force_retrain or not os.path.exists(save_path):\n        shsy5y_train_imgs = []\n        shsy5y_train_masks = []\n        for img_id in tqdm(shsy5y_train_csv.id.unique()):\n            shsy5y_train_imgs.append(\n                tf_img_color_rescale(\n                    tf_img_load(img_id)\n                ).numpy()\n            )\n            # could probably make this faster with a tensorflow map, but as-is, a dtype error is returned (can't map a string-type tensor to int)...\n            mask = np.zeros(shape)\n            for ann in shsy5y_train_csv[shsy5y_train_csv.id == img_id].annotation:\n                mask += rle_decode_tf(ann, shape=shape)\n            mask = tf.where(mask > 0, 1,0)\n            shsy5y_train_masks.append(mask.numpy().astype(np.uint8))\n        shsy5y_train_imgs = np.array(shsy5y_train_imgs)\n        shsy5y_train_masks = np.array(shsy5y_train_masks)\n\n        shsy5y_val_imgs = []\n        shsy5y_val_masks = []\n        for img_id in tqdm(shsy5y_val_csv.id.unique()):\n            shsy5y_val_imgs.append(\n                tf_img_color_rescale(\n                    tf_img_load(img_id)\n                ).numpy()\n            )\n            mask = np.zeros(shape)\n            for ann in shsy5y_val_csv[shsy5y_val_csv.id == img_id].annotation:\n                mask += rle_decode_tf(ann, shape=shape)\n            mask = tf.where(mask > 0, 1,0)\n            shsy5y_val_masks.append(mask.numpy().astype(np.uint8))\n        shsy5y_val_imgs = np.array(shsy5y_val_imgs)\n        shsy5y_val_masks = np.array(shsy5y_val_masks)\n        with open(save_path, 'wb') as f:\n            pkl.dump((shsy5y_train_imgs, shsy5y_train_masks, shsy5y_val_imgs, shsy5y_val_masks), f)\n    else:\n        with open(save_path, 'rb') as f:\n            shsy5y_train_imgs, shsy5y_train_masks, shsy5y_val_imgs, shsy5y_val_masks = pkl.load(f)\n    print(shsy5y_train_imgs.shape, shsy5y_val_imgs.shape)\n    print(shsy5y_train_masks.shape, shsy5y_val_masks.shape)","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:29:44.775091Z","iopub.execute_input":"2022-01-17T13:29:44.775469Z","iopub.status.idle":"2022-01-17T13:36:41.970523Z","shell.execute_reply.started":"2022-01-17T13:29:44.775433Z","shell.execute_reply":"2022-01-17T13:36:41.969703Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if not submission_mode:\n    plt.imshow(shsy5y_val_masks[2])\n    plt.show()\n    plt.close()","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:36:41.971783Z","iopub.execute_input":"2022-01-17T13:36:41.973375Z","iopub.status.idle":"2022-01-17T13:36:41.978065Z","shell.execute_reply.started":"2022-01-17T13:36:41.973334Z","shell.execute_reply":"2022-01-17T13:36:41.977261Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if not submission_mode:\n    for filename in ['/kaggle/input/sartorius-cell-instance-segmentation/train/' + x + '.png' for x in shsy5y_ids][:3]:\n        plot_image(filename)","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:36:41.979418Z","iopub.execute_input":"2022-01-17T13:36:41.979686Z","iopub.status.idle":"2022-01-17T13:36:41.992395Z","shell.execute_reply.started":"2022-01-17T13:36:41.97965Z","shell.execute_reply":"2022-01-17T13:36:41.99155Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"The previous approach worked well for both corr and astro, so it might fit here as well.","metadata":{}},{"cell_type":"code","source":"# build our shsy5y_model\nmodel_name = 'shsy5y_recognitionNN'\nshsy5y_model = AugmentedSequential()\nshsy5y_model.add(keras.Input(\n    shape=(520,704,1)\n))\n# now some vanilla CNN layers\nif True: # val_loss: \n    shsy5y_model.add(layers.Conv2D(\n        filters=16, kernel_size=(3,3),\n        activation='relu',\n        padding='same')\n    )\n    shsy5y_model.add(layers.Conv2D(\n        filters=8, kernel_size=(5,5),\n        activation='relu',\n        padding='same')\n    )\n    shsy5y_model.add(layers.Conv2D(\n        filters=8, kernel_size=(7,7),\n        activation='relu',\n        padding='same')\n    )\n# Apply a dense layer along the axis of the filters, keeping the image size the same\nshsy5y_model.add(layers.Conv2D(\n    filters=8, kernel_size=(1,1),\n    activation='relu',\n    padding='same')\n)\n# OUPTUT -----------------------------\n#shsy5y_model.add(layers.Dropout(0.5))\nshsy5y_model.add(layers.Conv2D(\n    filters=1, kernel_size=(1,1),\n    activation='sigmoid',\n    padding='same')\n)\n\nshsy5y_model.compile(optimizer='adam',\n              loss=balanced_loss,\n                   run_eagerly=True)  #tf.keras.losses.CategoricalCrossentropy(from_logits=False),\n              #metrics=['accuracy']\n\nshsy5y_model.summary()","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:36:41.993774Z","iopub.execute_input":"2022-01-17T13:36:41.994117Z","iopub.status.idle":"2022-01-17T13:36:42.054989Z","shell.execute_reply.started":"2022-01-17T13:36:41.994081Z","shell.execute_reply":"2022-01-17T13:36:42.054071Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if task_list[model_name]:\n    print(model_name)\n    if force_retrain or not os.path.exists(model_name+'.model'):\n        stop_early = tf.keras.callbacks.EarlyStopping(\n            monitor='val_loss',\n            patience=early_stopping_patience,\n            restore_best_weights=True\n        )\n        callbacks = [stop_early,]\n        history = shsy5y_model.fit(\n            tf.cast(shsy5y_train_imgs,tf.float32).numpy(),\n            tf.cast(tf.expand_dims(shsy5y_train_masks, axis=3),tf.float32).numpy(),\n            validation_data=(\n                tf.cast(shsy5y_val_imgs,tf.float32).numpy(),\n                tf.cast(tf.expand_dims(shsy5y_val_masks, axis=3),tf.float32).numpy()),\n            callbacks=callbacks,\n            epochs=epochs_to_train)\n        if not submission_mode:\n            shsy5y_model.save(model_name + '.model')\n            plot_history(history)\n    else:\n        shsy5y_model = keras.models.load_model(\n            model_name+'.model',\n            custom_objects={'balanced_loss':balanced_loss}\n        )","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:36:42.056622Z","iopub.execute_input":"2022-01-17T13:36:42.056896Z","iopub.status.idle":"2022-01-17T13:43:57.104718Z","shell.execute_reply.started":"2022-01-17T13:36:42.05686Z","shell.execute_reply":"2022-01-17T13:43:57.099858Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# make sure augmentation works, check the predictions\ndef test_predictions (model):\n    for i in range(3):\n        X_val, y_val = tf_img_augment_nomodel(\n            shsy5y_val_imgs, \n            tf.expand_dims(shsy5y_val_masks, axis=3).numpy()\n        )\n\n        ex_img = X_val[4].numpy()\n        ex_truemask = y_val[4].numpy()\n\n        # show an example prediction\n        fig = plt.figure(figsize=(16,6), clear=True)\n\n        plt.subplot(131)\n        plt.imshow(ex_img, cmap='seismic')\n        plt.title('Original image')\n\n        plt.subplot(132)\n        ex_mask = model.predict(ex_img.reshape((1,520,704,1)))[0]\n        print(ex_mask.shape)\n        plt.imshow(ex_mask, cmap=\"YlGn_r\")\n        plt.title('Predicted mask')\n        plt.colorbar()\n\n        plt.subplot(133)\n        plt.imshow(ex_truemask, cmap=\"YlGn_r\")\n        plt.title('True mask')\n\n        plt.show()\n        plt.close()\n    \nif task_list[model_name] and not submission_mode:\n    test_predictions(shsy5y_model)","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:43:57.10644Z","iopub.execute_input":"2022-01-17T13:43:57.106738Z","iopub.status.idle":"2022-01-17T13:43:57.126904Z","shell.execute_reply.started":"2022-01-17T13:43:57.1067Z","shell.execute_reply":"2022-01-17T13:43:57.125344Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# First, let's improve the contrast of the network output, to make it easier to segment\nif not submission_mode:\n    img = shsy5y_train_imgs[5]\n    mask = shsy5y_train_masks[5]\n    pred = shsy5y_model.predict(np.array([img,]))[0]\n\n    plt.figure(figsize=(20,6), clear=True)\n    plt.subplot(141)\n    plt.imshow(img, cmap='seismic')\n    plt.title('Original image')\n    plt.subplot(142)\n    plt.imshow(mask, cmap=\"YlGn_r\")\n    plt.title('True mask')\n    plt.subplot(143)\n    plt.imshow(pred, cmap=\"YlGn_r\")\n    plt.title('Prediction')\n    plt.subplot(144)\n    plt.imshow(pred > 0.6, cmap=\"YlGn_r\")\n    plt.title('Threshold')\n    plt.show()\n    plt.close()\n\n    plt.hist(pred.flatten())\n    plt.title('Prediction value histogram')\n    plt.show()\n    plt.close()","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:43:57.128557Z","iopub.execute_input":"2022-01-17T13:43:57.132652Z","iopub.status.idle":"2022-01-17T13:43:57.14309Z","shell.execute_reply.started":"2022-01-17T13:43:57.132612Z","shell.execute_reply":"2022-01-17T13:43:57.142323Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Let's plot the true segmentation, so that we know what we're looking for...\nif not submission_mode:\n    myid = shsy5y_train_csv.id.unique()[5]\n    shsy5y_train_csv[shsy5y_train_csv.id == myid].head()\n\n    pixels_true = []\n    cl = 0\n    for annotation in tqdm(tf.random.shuffle( # shuffle helps to make adjacent cells different color\n        shsy5y_train_csv[shsy5y_train_csv.id == myid].annotation)\n                          ):\n        mask_here = rle_decode_tf(annotation, shape=img.shape)\n        pixels_true_here = tf.where(tf.squeeze(mask_here) > 0)\n        pixels_true_here = tf.concat([pixels_true_here, cl*np.ones([pixels_true_here.shape[0],1])], axis=1)\n        cl += 1\n        pixels_true.append(pixels_true_here)\n    pixels_true = tf.concat(pixels_true, axis=0)\n\n    x = tf.transpose(pixels_true)[0]\n    y = tf.transpose(pixels_true)[1]\n    c = tf.transpose(pixels_true)[2]\n\n    plt.scatter(x, y, c=c, s=0.125, cmap='prism')\n    plt.show()\n    plt.close()","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:43:57.146881Z","iopub.execute_input":"2022-01-17T13:43:57.147145Z","iopub.status.idle":"2022-01-17T13:43:57.166593Z","shell.execute_reply.started":"2022-01-17T13:43:57.14712Z","shell.execute_reply":"2022-01-17T13:43:57.16517Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Hmm this might benefit from using a particle simulator, so the shsy5y model would be an adjusted version of the Astro model.","metadata":{}},{"cell_type":"code","source":"# trying out some different force prescriptions here..\n\ndef force_RA (sparse_distance, sparse_sqr_distance, force_radius, norm=0.5, sqr_scale1=1.5, sqr_scale3=2.):\n    distance = tf.reshape(sparse_distance.values, [-1,2])\n    sqr_distance = sparse_sqr_distance.values\n    part0_indices = sparse_sqr_distance.indices[:,0]\n    # calculate the number of neighbours\n    n_neighbours = tf.SparseTensor(\n        indices = sparse_sqr_distance.indices, \n        values = tf.ones(tf.shape(sparse_sqr_distance.values), dtype=tf.float16),\n        dense_shape = sparse_sqr_distance.dense_shape\n    )\n    n_neighbours = tf.sparse.reduce_sum(n_neighbours, axis=1)\n    n_neighbours = tf.gather_nd(n_neighbours, tf.expand_dims(part0_indices,-1))\n    # the actual force prescription\n    force_proposed = (\n        - norm * (sqr_scale1-sqr_distance) * tf.exp( - sqr_distance/sqr_scale3)\n    )\n    force_proposed = distance * tf.expand_dims(norm*force_proposed, axis=-1)\n    # apply limits\n    force_max = tf.sqrt(sqr_distance) / n_neighbours\n    force_value = tf.reduce_sum(force_proposed**2, axis=1)\n    force_direction = force_proposed / tf.expand_dims(force_value, axis=-1)\n    force_proposed = tf.where(\n        tf.expand_dims(tf.abs(force_value) < force_max, -1),\n        force_proposed,\n        norm * force_direction * tf.expand_dims(force_max,-1)\n    )\n    # reshape to sparse representation\n    force_proposed = tf.reshape(force_proposed, [-1])\n    return force_proposed","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:43:57.168045Z","iopub.execute_input":"2022-01-17T13:43:57.168466Z","iopub.status.idle":"2022-01-17T13:43:57.185827Z","shell.execute_reply.started":"2022-01-17T13:43:57.168427Z","shell.execute_reply":"2022-01-17T13:43:57.184627Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Shsy5yModel (AstroModel):\n    \n    def __init__ (\n        self,\n        preprocess=False,\n        detection_model=shsy5y_model,\n        detection_threshold=0.6,\n        simulator_force=force_RA,\n        simulator_force_radius=10,\n        simulator_force_kwargs={},\n        simulator_steps=simulation_epochs,\n        simulator_dt=simulation_dt,\n        clusterer=DBSCAN(eps=6, min_samples=10)\n    ):\n        super(Shsy5yModel, self).__init__(\n            preprocess=preprocess,\n            detection_model=detection_model,\n            detection_threshold=detection_threshold,\n            simulator_force=simulator_force,\n            simulator_force_radius=simulator_force_radius,\n            simulator_force_kwargs=simulator_force_kwargs,\n            simulator_steps=simulator_steps,\n            simulator_dt=simulator_dt,\n            clusterer=clusterer\n        )","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:43:57.190646Z","iopub.execute_input":"2022-01-17T13:43:57.191098Z","iopub.status.idle":"2022-01-17T13:43:57.20196Z","shell.execute_reply.started":"2022-01-17T13:43:57.190875Z","shell.execute_reply":"2022-01-17T13:43:57.201194Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if not submission_mode:\n    img = shsy5y_train_imgs[5]\n    mask = shsy5y_train_masks[5]\n\n    model = Shsy5yModel()\n    encoded = model.predict(img)\n    print(encoded[:50])","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:43:57.206683Z","iopub.execute_input":"2022-01-17T13:43:57.206944Z","iopub.status.idle":"2022-01-17T13:43:57.216048Z","shell.execute_reply.started":"2022-01-17T13:43:57.206911Z","shell.execute_reply":"2022-01-17T13:43:57.215171Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# let's see how our shsy5y model performs with respect to the competition score\nif not submission_mode:\n    \n    # choose the indices to test on\n    val_ids = np.random.choice(shsy5y_val_csv.id.unique(), 1)\n    print(val_ids)\n\n    # parse the cortect answers\n    rows_true = shsy5y_val_csv[shsy5y_val_csv.id.isin(val_ids)]\n    encoded_true = ''\n    for idx, row in rows_true[['id', 'annotation']].iterrows():\n        encoded_true += '\\n' + row['id'] + ',' + row['annotation']\n    encoded_true = encoded_true[1:]\n\n    # calculate predictions\n    if True:\n        encoded = ''\n        for myid in val_ids:\n            print('--- Predicting image %s ---' % myid, flush=True)\n            img = tf_img_color_rescale(\n                tf_img_load(myid)\n            ).numpy()\n            # decode the target mask for comparison\n            pixels_true = []\n            cl = 0\n            for annotation in tf.random.shuffle( # shuffle helps to make adjacent cells different color\n                shsy5y_val_csv[shsy5y_val_csv.id == myid].annotation\n            ):\n                mask_here = rle_decode_tf(annotation, shape=img.shape)\n                pixels_true_here = tf.where(tf.squeeze(mask_here) > 0)\n                pixels_true_here = tf.concat([pixels_true_here, cl*np.ones([pixels_true_here.shape[0],1])], axis=1)\n                cl += 1\n                pixels_true.append(pixels_true_here)\n            pixels_true = tf.concat(pixels_true, axis=0)\n            # pass the image through our model\n            model = Shsy5yModel()\n            encoded += '\\n' + model.predict(\n                img,\n                img_id=myid,\n                mask_pixels=pixels_true\n            )\n            print('--- Image %s predicted ---' % myid, flush=True)\n        encoded = encoded[1:]\n\n        score = comp_score(encoded, encoded_true)\n        print('Competition score: %.3f' % score)","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:43:57.221069Z","iopub.execute_input":"2022-01-17T13:43:57.221405Z","iopub.status.idle":"2022-01-17T13:43:57.257072Z","shell.execute_reply.started":"2022-01-17T13:43:57.221373Z","shell.execute_reply":"2022-01-17T13:43:57.256236Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"-----\n**Final Model**","metadata":{}},{"cell_type":"code","source":"def logit_to_celltype (logit):\n    return cell_types[np.argmax(logit)]\n\nclass FinalModel:\n    \n    def __init__ (\n        self,\n        preprocess=False,\n        classifier=img_classifier,\n        astro_kwargs={},\n        cort_kwargs={},\n        shsy5y_kwargs={}\n    ):\n        self.preprocess=preprocess\n        self.classifier=img_classifier\n        self.segmenters = {\n            'astro':AstroModel(preprocess=False, **astro_kwargs),\n            'cort':CortModel(preprocess=False, **cort_kwargs),\n            'shsy5y':Shsy5yModel(preprocess=False, **shsy5y_kwargs)\n        }\n        \n    def predict (\n        self,\n        img, img_id='noid',\n        mask_pixels=[],\n        verbose=True\n    ):\n        if verbose:\n            print(\"Processing image %s.\" % img_id, flush=True)\n        # preprocess if needed\n        if verbose:\n            print(' - image preprocessing.. ', end='', flush=True)\n            fig = plt.figure(figsize=(12,6))\n            plt.subplot(121)\n            plt.imshow(1.*img, cmap=\"seismic\")\n        if self.preprocess:\n            img = img_color_rescale(img)\n        if verbose:\n            plt.subplot(122)\n            plt.imshow(img, cmap=\"seismic\")\n            plt.show()\n            plt.close()\n            print('done.', flush=True)\n        # classify the image\n        if verbose:\n            print(' - image classification.. ', flush=True)\n        img_class = logit_to_celltype(\n            self.classifier.predict(tf.expand_dims(img, axis=0))\n        )\n        if verbose:\n            print('    - classified as ', img_class, ', done.')\n        # apply the appropriate type-specific segmentation model\n        result = self.segmenters[img_class].predict(\n            img, img_id=img_id,\n            mask_pixels=mask_pixels,\n            verbose=verbose\n        )\n        # cleanup and output\n        return result","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:49:42.285264Z","iopub.execute_input":"2022-01-17T13:49:42.285765Z","iopub.status.idle":"2022-01-17T13:49:42.297519Z","shell.execute_reply.started":"2022-01-17T13:49:42.28573Z","shell.execute_reply":"2022-01-17T13:49:42.296669Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Test the model on a couple of examples\nif not submission_mode:\n    for celltype in cell_types:\n        print('Testing cell type =', celltype, flush=True)\n        img_id = train_csv[train_csv.cell_type == celltype].id.iloc[7]\n        # load the image\n        img = tf_img_load(img_id, directory='train')\n        # parse the cortect answers\n        rows_true = train_csv[train_csv.id == img_id]\n        encoded_true = ''\n        for idx, row in rows_true[['id', 'annotation']].iterrows():\n            encoded_true += '\\n' + row['id'] + ',' + row['annotation']\n        encoded_true = encoded_true[1:]\n        # decode the target mask for comparison\n        pixels_true = []\n        cl = 0\n        for annotation in tf.random.shuffle( # shuffle helps to make adjacent cells different color\n            train_csv[train_csv.id == img_id].annotation\n        ):\n            mask_here = rle_decode_tf(annotation, shape=img.shape)\n            pixels_true_here = tf.where(tf.squeeze(mask_here) > 0)\n            pixels_true_here = tf.concat([\n                pixels_true_here, \n                cl*np.ones([pixels_true_here.shape[0],1])\n            ], axis=1)\n            cl += 1\n            pixels_true.append(pixels_true_here)\n        pixels_true = tf.concat(pixels_true, axis=0)\n        # run everything through the model\n        model = FinalModel(preprocess=True)\n        encoded = model.predict(\n            img, img_id,\n            mask_pixels=pixels_true\n        )\n        print(encoded[:50])\n        # calculate competition score\n        score = comp_score(encoded, encoded_true)\n        print('Competition score: %.3f' % score)","metadata":{"execution":{"iopub.status.busy":"2021-12-30T10:57:36.972574Z","iopub.execute_input":"2021-12-30T10:57:36.973049Z","iopub.status.idle":"2021-12-30T11:02:06.365696Z","shell.execute_reply.started":"2021-12-30T10:57:36.973005Z","shell.execute_reply":"2021-12-30T11:02:06.364909Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# process the submission\n\n# load the submission examples\ntest_ids = glob.glob('/kaggle/input/sartorius-cell-instance-segmentation/test/*.png')\ntest_ids = [x.split('/')[-1][:-4] for x in test_ids]\nprint(len(test_ids), test_ids)\n\n# run the model\nsubmission = 'id, predicted\\n'\nfor img_id in test_ids:\n    # load the image\n    img = tf_img_load(img_id, directory='test')\n    # run everything through the model\n    model = FinalModel(preprocess=True)\n    encoded = model.predict(\n        img, img_id, verbose=(not submission_mode)\n    )\n    print(encoded[:50])\n    submission += encoded + '\\n'\n# parse as csv and view to ensure correct format\nsubmission = pd.read_csv(StringIO(submission[:-1]))\nprint(submission.head())\n# save the submission file\nsubmission.to_csv('submission.csv', index=False)","metadata":{"execution":{"iopub.status.busy":"2022-01-17T13:55:18.812786Z","iopub.execute_input":"2022-01-17T13:55:18.813227Z","iopub.status.idle":"2022-01-17T13:58:23.718067Z","shell.execute_reply.started":"2022-01-17T13:55:18.813175Z","shell.execute_reply":"2022-01-17T13:58:23.716923Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}