{"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":"part1 is here  \nhttps://www.kaggle.com/code/hengck23/unet-instance-segmentation-onmipose-part1\n","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    \n    !pip install timm\n\nprint('PIP OK!!!')","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-07-03T07:19:54.219653Z","iopub.execute_input":"2023-07-03T07:19:54.220067Z","iopub.status.idle":"2023-07-03T07:19:54.234319Z","shell.execute_reply.started":"2023-07-03T07:19:54.220037Z","shell.execute_reply":"2023-07-03T07:19:54.233127Z"},"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\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom timm.models.resnet import *\n\nprint('IMPORT OK!!!') \n\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'   ","metadata":{"execution":{"iopub.status.busy":"2023-07-03T07:19:54.236368Z","iopub.execute_input":"2023-07-03T07:19:54.236809Z","iopub.status.idle":"2023-07-03T07:19:54.251065Z","shell.execute_reply.started":"2023-07-03T07:19:54.236780Z","shell.execute_reply":"2023-07-03T07:19:54.250122Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#from segmentation_models_pytorch.decoders.unet.decoder import  DecoderBlock\nclass Conv2dReLU(nn.Sequential):\n    def __init__(\n        self,\n        in_channels,\n        out_channels,\n        kernel_size,\n        padding=0,\n        stride=1,\n    ):\n        conv = nn.Conv2d(\n            in_channels,\n            out_channels,\n            kernel_size,\n            stride=stride,\n            padding=padding,\n            bias=False,\n        )\n        relu = nn.ReLU(inplace=True)\n        bn = nn.BatchNorm2d(out_channels)\n\n        super(Conv2dReLU, self).__init__(conv, bn, relu)\n\nclass DecoderBlock(nn.Module):\n    def __init__(\n        self,\n        in_channels,\n        skip_channels,\n        out_channels,\n        use_batchnorm=True,\n        attention_type=None,\n    ):\n        super().__init__()\n        self.conv1 = Conv2dReLU(\n            in_channels + skip_channels,\n            out_channels,\n            kernel_size=3,\n            padding=1,\n        )\n        self.attention1 = nn.Identity()\n        self.conv2 = Conv2dReLU(\n            out_channels,\n            out_channels,\n            kernel_size=3,\n            padding=1,\n        )\n        self.attention2 = nn.Identity()\n\n    def forward(self, x, skip=None):\n        x = F.interpolate(x, scale_factor=2, mode=\"nearest\")\n        if skip is not None:\n            x = torch.cat([x, skip], dim=1)\n            x = self.attention1(x)\n        x = self.conv1(x)\n        x = self.conv2(x)\n        x = self.attention2(x)\n        return x\n\nclass SmpUnetDecoder(nn.Module):\n    def __init__(self,\n                 in_channel,\n                 skip_channel,\n                 out_channel,\n                 ):\n        super().__init__()\n        self.center = nn.Identity()\n\n        i_channel = [in_channel, ] + out_channel[:-1]\n        s_channel = skip_channel\n        o_channel = out_channel\n        block = [\n            DecoderBlock(i, s, o, use_batchnorm=True, attention_type=None)\n            for i, s, o in zip(i_channel, s_channel, o_channel)\n        ]\n        self.block = nn.ModuleList(block)\n\n    def forward(self, feature, skip):\n        d = self.center(feature)\n        decode = []\n        for i, block in enumerate(self.block):\n            # print(i, d.shape, skip[i].shape if skip[i] is not None else 'none')\n            # print(block.conv1[0])\n            # print('')\n\n            s = skip[i]\n            d = block(d, s)\n            decode.append(d)\n\n        last = d\n        return last, decode\n\n\n#######################################################################################\n# num_class = 5 # semantic, dist, fx fy, boundary\n\nclass Net(nn.Module):\n    def __init__(self, ):\n        super().__init__()\n        self.output_type = ['inference', 'loss']\n\n        # --------------------------------\n\n        encoder_dim = [64, 256, 512, 1024, 2048]\n        decoder_dim = [256, 128, 64, 32, 16]\n\n        self.encoder = seresnext101d_32x8d(pretrained=False, in_chans=3)\n        self.decoder = SmpUnetDecoder(\n            in_channel=encoder_dim[-1],\n            skip_channel=encoder_dim[:-1][::-1] + [0],\n            out_channel=decoder_dim,\n        )\n        self.all  = nn.Conv2d(decoder_dim[-1], 5, kernel_size=1)\n        self.boundary = nn.Conv2d(decoder_dim[-1], 1, kernel_size=1)\n        self.dist     = nn.Conv2d(decoder_dim[-1], 1, kernel_size=1)\n        self.flow     = nn.Conv2d(decoder_dim[-1], 2, kernel_size=1)\n\n\n    def forward(self, batch):\n        x = batch['image']\n        B,_,H,W = x.shape\n\n        # ---------------------------------\n        encoder = []\n        e = self.encoder\n\n        x = e.conv1(x)\n        x = e.bn1(x)\n        x = e.act1(x); encoder.append(x)\n        x = F.avg_pool2d(x, kernel_size=2, stride=2)\n        x = e.layer1(x); encoder.append(x)\n        x = e.layer2(x); encoder.append(x)\n        x = e.layer3(x); encoder.append(x)\n        x = e.layer4(x); encoder.append(x)\n        ##[print(f'encoder1_{i}',f.shape) for i,f in enumerate(encoder)]\n\n        #--------------------------\n        last, decoder = self.decoder(feature=encoder[-1], skip=encoder[:-1][::-1]+ [None])\n        ##print('last', last.shape)\n\n        # ---------------------------------\n        #last = F.dropout(last, p=0.5, training=self.training)\n        all      = self.all(last)\n        logit    = all[:,[0]]\n        boundary = all[:,[1]]\n        dist     = all[:,[2]]\n        flow     = all[:,[3,4]]\n        ##print('logit', logit.shape, H/logit.shape[2])\n\n        output = {}\n        if 'loss' in self.output_type:\n            output['mask_loss'] = F.binary_cross_entropy_with_logits(logit, batch['blood_vessel'])\n            output['boundary_loss'] = F.binary_cross_entropy_with_logits(boundary, batch['boundary'])\n            output['dist_loss'] = F.mse_loss(dist, batch['dist'])\n            output['flow_loss'] = F.mse_loss(flow, batch['flow'])\n            # <todo>: weighing MSE ???\n\n        if 'inference' in self.output_type:\n            if all.shape[2:]!=(H, W):\n                all = F.interpolate(all, size=(H, W), mode='bilinear', align_corners=False, antialias=True)\n                logit    = all[:,[0]]\n                boundary = all[:,[1]]\n                dist     = all[:,[2]]\n                flow     = all[:,[3,4]]\n\n            output['blood_vessel'] = torch.sigmoid(logit)\n            output['boundary'] = torch.sigmoid(boundary)\n            output['dist'] = dist\n            output['flow'] = flow\n\n        return output\n    \n#post process\n# larger flow_threshold values mores more masks\n# 0.4\ndef post_process_one(dist, flow, flow_threshold=1):\n    dist = np.pad(dist, [[PAD, PAD], [PAD, PAD]])\n    flow = np.pad(flow, [[0, 0], [PAD, PAD], [PAD, PAD]])\n\n    post = compute_masks(\n        flow,  #  #dP[:, i],\n        dist,   # cellprob[i],\n        None,  # 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=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!='cpu',\n        device=None,\n        nclasses=3,  # 3, #4\n        dim=2)\n\n\n    instance = post[0][..., PAD:-PAD, PAD:-PAD]\n    return instance\n\n\n\ncheckpoint = '/kaggle/input/onmi-pose-demo/00045125.model.pth'\n\nnet = Net()\nf = torch.load(checkpoint, map_location=lambda storage, loc: storage)\nstate_dict = f['state_dict']\nprint(net.load_state_dict(state_dict,strict=True))\n","metadata":{"execution":{"iopub.status.busy":"2023-07-03T07:19:54.252622Z","iopub.execute_input":"2023-07-03T07:19:54.253512Z","iopub.status.idle":"2023-07-03T07:19:56.347275Z","shell.execute_reply.started":"2023-07-03T07:19:54.253480Z","shell.execute_reply":"2023-07-03T07:19:56.345972Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# helper function for data\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\ndef read_as_one_batch(id):\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    \n    image = np.ascontiguousarray(np.transpose(image,(2,0,1,)))/255\n    image = torch.from_numpy(image).float()\n    instance = torch.from_numpy(instance).long()\n    \n    batch = {\n        'image' : image.reshape(1,3,512,512),\n        'instance' : instance.reshape(1,1,512,512), \n    }\n    return batch\n    \n\n##########################################################################################\n# draw result\ndef my_norm(x):\n    x = x-x.min()\n    x = x/(x.max()+0.0001)\n    return x\n\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\n\ndef instance_to_overlay(instance, contour_thickness=2, contour_color=(0,0,255),ncolor=False):\n    h,w = instance.shape[:2]\n    _, c = np.unique(instance, return_inverse=True, )\n    c = c.reshape((h,w))\n\n    if ncolor:\n        cinst = ncolor.label(instance)\n        cinst = cv2.applyColorMap(cinst*50, cv2.COLORMAP_JET)\n        overlay = cinst\n    else:\n        overlay = np.zeros((h,w,3))\n        overlay[c>0]=255\n\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),contour_color,thickness=contour_thickness)\n    return overlay\n\n\ndef draw_flow_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\ndef show_batch(batch, output, name='xxx'):\n    batch_size = len(batch['image'])\n\n    image = batch['image'].permute(0, 2, 3, 1)\n    image = image.float().data.cpu().numpy()\n    instance = batch['instance'].long().data.cpu().numpy()\n\n    semantic = output['blood_vessel'].float().data.cpu().numpy()\n    boundary = output['boundary'].float().data.cpu().numpy() #x5\n    dist = output['dist'].float().data.cpu().numpy()\n    flow = output['flow'].float().data.cpu().numpy()\n    \n    if 0: #debug with groud truth\n        semantic = batch['blood_vessel'].float().data.cpu().numpy()\n        boundary = batch['boundary'].float().data.cpu().numpy()  # x5\n        dist = batch['dist'].float().data.cpu().numpy()\n        flow = batch['flow'].float().data.cpu().numpy()\n\n    # print(blood_vessel.min(),blood_vessel.max())\n    for b in range(batch_size):\n        m = image[b]\n        t = instance_to_overlay(instance[b,0], contour_color=(255,0,0))/255\n        s = cv2.cvtColor(semantic[b,0],cv2.COLOR_GRAY2BGR)\n        d = cv2.cvtColor(my_norm(dist[b,0]),cv2.COLOR_GRAY2BGR)\n        r = cv2.cvtColor(boundary[b,0],cv2.COLOR_GRAY2BGR)\n        f = draw_flow_overlay(flow[b], norm=True) / 255\n\n        post = post_process_one(dist[b,0], flow[b])\n        post = instance_to_overlay(post, contour_color=(0,0,255))/255\n\n        s = m * 0.3 + s * 0.7\n        yx = np.where(np.all(t == [1,0,0], axis=-1))\n        s[yx] = [1,0,0]\n\n        diff = t-post\n        yx = np.where(np.all(diff < [0,0,0], axis=-1))\n        diff[yx] = [0.4,0.4,0.4]\n\n        overlay = np.hstack([\n            m,\n            t ,\n            post,\n            s,\n            d,\n            f,\n            diff,\n        ])\n        height,width=overlay.shape[:2]\n        for x in range(0,width,512):\n            overlay[:,x]=1\n\n        \n        #image_show_norm(f'{name}: [image,truth,predict, semantic,dist,flow, diff]', overlay, min=0, max=1, resize=0.5)\n        #cv2.waitKey(0)\n        print(f'{name}: [image,truth,predict, semantic,dist,flow, diff]')\n        \n        plt.figure(figsize=(21,3))\n        plt.imshow(overlay)\n        plt.show()\n    ","metadata":{"execution":{"iopub.status.busy":"2023-07-03T07:19:56.348562Z","iopub.status.idle":"2023-07-03T07:19:56.349130Z","shell.execute_reply.started":"2023-07-03T07:19:56.348946Z","shell.execute_reply":"2023-07-03T07:19:56.348965Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\nnet.eval()\nnet.output_type = ['inference',]\nif DEVICE!='cpu':\n    net.cuda()\n    \nfor i, id in enumerate(image_id):\n    print(i, id)\n\n    #1. read data\n    batch = read_as_one_batch(id)\n    if DEVICE!='cpu':\n        batch['image'] = batch['image'].cuda()\n    output = net(batch)\n    show_batch(batch, output, name='valid')\n    print('------------------------------------------------------------------')\n  ","metadata":{"execution":{"iopub.status.busy":"2023-07-03T07:19:56.350003Z","iopub.status.idle":"2023-07-03T07:19:56.350882Z","shell.execute_reply.started":"2023-07-03T07:19:56.350659Z","shell.execute_reply":"2023-07-03T07:19:56.350679Z"},"trusted":true},"execution_count":null,"outputs":[]}]}