{"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":"markdown","source":"# Background Reading\n\nFor information on onmipose, check these\n\ncode: https://github.com/kevinjohncutler/omnipose  \ndoc: https://omnipose.readthedocs.io/   \npaper: https://www.nature.com/articles/s41592-022-01639-4    \n\"Omnipose: a high-precision morphology-independent solution for bacterial cell segmentation\" - Kevin J. Cutler \n\n# Basic Idea\n\n![https://i.ibb.co/QpGsRN2/Selection-999-2452.png](https://i.ibb.co/QpGsRN2/Selection-999-2452.png)\n\nif you cannot see the diagram please, use this link: https://ibb.co/S5pq8nh \n\n\nInstead of just predicting the usual semantic mask, we predict addition information like distance transform, flow Lx and Ly. We can then recover the instance masks via post processsing.\n\nBut how do we know what additional information to predict? how to design the post processing? Luckily, there is a libray for that. It is called \"Omnipose\". You can refer to the paper for more information. \n","metadata":{}},{"cell_type":"markdown","source":"In part1 notebook, we show:\n1) how to generate addition information as target (e.g. in the dataset class)  \n2) how to post process, i.e.recover instance segmentation using the additional targets.  \n  \nsince we are using the ground truth target, **we expect 100% instance masks recontsruction**.  \n\nIn future part2 notebook, we show predictions from a learned Unet model. We verify that although these predictions are not perfect,we cam still recover the instance masks with very high accuracy (e.g. better than mask-rcnn)","metadata":{}},{"cell_type":"code","source":"try:\n    import omnipose\nexcept:\n    !pip install torchvf\n    !pip install mgen\n    !pip install edt\n    !pip install ncolor\n    !pip install fastremap\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-07-02T12:33:45.930681Z","iopub.execute_input":"2023-07-02T12:33:45.931185Z","iopub.status.idle":"2023-07-02T12:33:45.941803Z","shell.execute_reply.started":"2023-07-02T12:33:45.931150Z","shell.execute_reply":"2023-07-02T12:33:45.940583Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import sys\nsys.path.append('/kaggle/input/onmi-pose-demo')\nfrom omnipose.core import compute_masks, masks_to_flows\n\nimport numpy as np  \nimport pandas as pd  \nimport cv2\n\nimport matplotlib\nimport matplotlib.pyplot as plt\nimport ncolor\n\nfrom omnipose.utils import normalize99\n\nimport torch\n\nprint('IMPORT OK!!!') \n\n#------\n\n#helper function\n\ndef to_instance(m):\n    # merge sure and unsure into one instance mask\n    m1 = m[...,1].astype(np.int32)\n    m2 = m[...,2].astype(np.int32)\n    m2[m1>0] = m1[m1>0]+256\n    _, c = np.unique(m2, return_inverse=True, )\n    c = c.reshape(m.shape[:2]) \n    return c\n\n\n\n#draw functions\ndef mask_to_inner_contour(mask):\n    mask = mask>0.5\n    pad = np.lib.pad(mask, ((1, 1), (1, 1)), 'reflect')\n    contour = mask & (\n            (pad[1:-1,1:-1] != pad[:-2,1:-1]) \\\n            | (pad[1:-1,1:-1] != pad[2:,1:-1]) \\\n            | (pad[1:-1,1:-1] != pad[1:-1,:-2]) \\\n            | (pad[1:-1,1:-1] != pad[1:-1,2:])\n    )\n    return contour\n\ndef draw_contour_overlay(image, mask, color=(0,0,255), thickness=1):\n    contour =  mask_to_inner_contour(mask)\n    if thickness==1:\n        image[contour] = color\n    else:\n        r = max(1,thickness//2)\n        for y,x in np.stack(np.where(contour)).T:\n            cv2.circle(image, (x,y), r, color, lineType=cv2.LINE_4 )\n    return image\n\n\ndef flow_to_hsv(flow, norm=True):\n    mag = np.sqrt(np.sum(flow ** 2, axis=0))\n    if norm:\n        mag = np.clip(normalize99(mag), 0, 1.)\n\n    angles = np.arctan2(flow[1], flow[0]) + np.pi\n\n    a = 2\n    r = ((np.cos(angles) + 1) / a)\n    g = ((np.cos(angles + 2 * np.pi / 3) + 1) / a)\n    b = ((np.cos(angles + 4 * np.pi / 3) + 1) / a)\n\n    hsv = np.stack((r * mag, g * mag, b * mag), axis=-1)\n    hsv = (np.clip(hsv, 0, 1) * 255).astype(np.uint8)\n    return hsv\n\ndef draw_instance_overlay1(instance):\n    h,w = instance.shape[:2]\n    _, c = np.unique(instance, return_inverse=True, )\n    c = c.reshape((h,w))\n\n    cinst = ncolor.label(instance)\n    cinst = cv2.applyColorMap(cinst*50, cv2.COLORMAP_JET)\n\n    #overlay = np.zeros((h,w,3))\n    #overlay[c>0]=255\n\n    overlay = cinst\n    overlay[c==0]=0\n    num_instance = c.max()\n    for i in range(num_instance):\n        draw_contour_overlay(overlay, (c==(i+1)).astype(np.float32),(0,0,255),2)\n    return overlay\n\n\ndef draw_flow_to_overlay(flow, norm=True):\n    mag = np.sqrt(np.sum(flow ** 2, axis=0))\n    if norm:\n        mag = np.clip(normalize99(mag), 0, 1.)\n\n    angles = np.arctan2(flow[1], flow[0]) + np.pi\n\n    a = 2\n    r = ((np.cos(angles) + 1) / a)\n    g = ((np.cos(angles + 2 * np.pi / 3) + 1) / a)\n    b = ((np.cos(angles + 4 * np.pi / 3) + 1) / a)\n\n    hsv = np.stack((r * mag, g * mag, b * mag), axis=-1)\n    hsv = (np.clip(hsv, 0, 1) * 255).astype(np.uint8)\n    return hsv\n\n#############################################################################\n#start here !!!\n\nimage_id =[ \n '2a4cc81cc5d6', '3378fe495259', '4ca084aec87b', '50dc42c72b45',\n '611599949a53', '67a7395eef7d', '8d24ea45c6b6', 'a373ae26f4f0',\n 'e1f6c8a7873e', 'f86347534ec1'\n]\ninstance_dir = \\\n    '/kaggle/input/onmi-pose-demo/example-instance-mask-ground-truth'\nimage_dir = \\\n    '/kaggle/input/hubmap-hacking-the-human-vasculature/train'\n\nPAD = 16\nDEVICE = 'cpu' #'cuda' \n    \n    \nfor i, id in enumerate(image_id):\n    print(i, id)\n\n    #1. read data\n    image = cv2.imread(f'{image_dir}/{id}.tif',cv2.IMREAD_COLOR)\n    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n    instance = to_instance(cv2.imread(f'{instance_dir}/{id}.png',cv2.IMREAD_COLOR))\n    mask = ((instance>0)*255).astype(np.uint8) #semantic\n  \n    instance_overlay1 = draw_instance_overlay1(instance)\n    semantic_overlay1 = cv2.cvtColor(mask, cv2.COLOR_GRAY2RGB)\n    \n    print(['image', 'instance mask', 'semantic mask'])\n    plt.figure(figsize=(18,3))\n    plt.imshow(np.hstack([\n        image, \n        instance_overlay1,\n        semantic_overlay1//2+image//2\n    ]))\n    plt.show()\n    \n    #2. generate additaional target\n    with torch.no_grad(): \n        pad_instance = np.pad(instance,[[PAD,PAD],[PAD,PAD]]) \n        mask, dist, boundary, T, mu = \\\n            masks_to_flows(\n                pad_instance,\n                affinity_graph=None,\n                dists=None,\n                coords=None,\n                links=None,\n                use_gpu=DEVICE=='cuda',\n                device=None,\n                omni=True,\n                dim=2,\n                smooth=False,\n                normalize=False,\n                n_iter=None,\n                verbose=False\n            )\n        #rescale the values\n        target_dist=T\n        target_dist[T <= 0] = -5\n        target_flow = mu*5\n        target_bd = boundary\n        target_bd[boundary<=0]=-5\n        \n        target_dist = target_dist.data.cpu().numpy() #keep as torch tensor if you are using it for training\n        target_flow = target_flow.data.cpu().numpy()     \n        \n    # normalise for visualisation \n    print(['target distance transform', 'flow Lx,Ly', 'boundary'])\n    plt.figure(figsize=(12,3))\n    plt.imshow(np.hstack([\n        normalize99(target_dist), \n        normalize99(target_flow[0]),\n        normalize99(target_flow[1]),\n        normalize99(target_bd),\n    ]))\n    plt.show()\n        \n    #3. post process\n    # assume unet results are perfect\n    predict_flow = target_flow\n    predict_dist = target_dist\n    predict_bd = target_bd\n\n    pad_post = compute_masks(\n        predict_flow,  #dP[:, i],\n        predict_dist,   #cellprob[i],\n        predict_bd,    #boundaries,\n        niter=200,\n        rescale=1.0,\n        resize=None,\n        min_size=20,\n        mask_threshold=0,\n        diam_threshold=12,\n        flow_threshold=0.4,\n        flow_factor=6, #5\n        interp=True,\n        cluster=False,\n        boundary_seg=False,\n        affinity_seg=False,\n        calc_trace=False,\n        verbose=False,\n        use_gpu=DEVICE=='cuda',\n        device=None,\n        nclasses=4, #3 if you use \"dist+flow\"; 4 for \"dist+flow+boundary\"\n        dim=2)\n   \n    post_instance = pad_post[0][...,PAD:-PAD,PAD:-PAD] \n    #print(post_instance.shape,pad_post[0].shape) \n    #print(post_flow.shape,pad_post[1].shape) \n    \n    post_instance_overlay1 = draw_instance_overlay1(post_instance)  \n\n    print(['reconstructed instance mask', 'diff',])\n    plt.figure(figsize=(6,3))\n    plt.imshow(np.hstack([\n        post_instance_overlay1, \n        post_instance_overlay1-instance_overlay1\n    ]))\n    plt.show()\n    print('------------------------------------------------------------------')\n","metadata":{"execution":{"iopub.status.busy":"2023-07-02T12:33:45.949138Z","iopub.execute_input":"2023-07-02T12:33:45.949557Z","iopub.status.idle":"2023-07-02T12:34:06.103847Z","shell.execute_reply.started":"2023-07-02T12:33:45.949515Z","shell.execute_reply":"2023-07-02T12:34:06.102663Z"},"trusted":true},"execution_count":null,"outputs":[]}]}