{"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":"!pip install \"/kaggle/input/hpamisc/pytorch_zoo-master\"\n!pip install \"/kaggle/input/hpamisc/pycocotools-2.0-cp37-cp37m-linux_x86_64.whl\"\n!pip install \"/kaggle/input/hpamisc/faiss_gpu-1.7.0-cp37-cp37m-manylinux2014_x86_64.whl\"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import sys\nimport importlib\n\nMODULE_PATH = \"/kaggle/input/timmilya/pytorch-image-models-master/timm/__init__.py\"\nMODULE_NAME = \"timm\"\n\nspec = importlib.util.spec_from_file_location(MODULE_NAME, MODULE_PATH)\nmodule = importlib.util.module_from_spec(spec)\nsys.modules[spec.name] = module\nspec.loader.exec_module(module)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport cv2\nimport math\nimport timm\nimport zlib\nimport faiss\nimport torch\nimport base64\nimport pickle\nimport torch.nn\nimport numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\nimport albumentations as A\nimport scipy.ndimage as ndi\nimport matplotlib.pyplot as plt\nimport torch.nn.functional as F\nfrom scipy.sparse import csr_matrix\nfrom pycocotools import _mask as coco_mask\nfrom skimage import filters, measure, segmentation, transform, util\nfrom skimage.morphology import closing, disk, remove_small_holes, remove_small_objects\nfrom fastai.vision.all import cnn_learner, DataBlock, aug_transforms, Resize, Normalize, ImageBlock, MultiCategoryBlock, RandomSplitter, partial, PILImage, create_body, create_head, num_features_model","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"directory = '/kaggle/input/hpa-single-cell-image-classification'\n\nPUBLIC_ONLY = True\n\nSEGMENTATION_SCALE = 0.25\nARCFACE_THRESH = 0.9\n\nnum_classes = 19\ncrop_input_size = 168\nfull_input_size = 512","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_df = pd.read_csv(os.path.join(directory, 'sample_submission.csv'))\n\nif len(test_df) == 559:\n    test_df = test_df[:8][::2]\n    \nif PUBLIC_ONLY:\n    with open(\"/kaggle/input/hpamisc/public_image_ids.pickle\", 'rb') as f:\n        public_image_ids = pickle.load(f)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"crop_weights = ['dm_nfnet_f0_168_2021_04_23__11_13_04/epoch_0',\n                'dm_nfnet_f1_168_2021_04_26__14_52_16/epoch_0',\n                'dm_nfnet_f2_168_2021_04_27__15_13_50/epoch_0',\n                'dm_nfnet_f3_168_2021_04_30__18_30_44/epoch_0',\n                'ecaresnet50d_168_2021_04_24__08_48_09/epoch_0',\n                'ecaresnet50t_168_2021_05_01__20_25_39/epoch_0',\n                'seresnet152d_168_2021_04_23__17_30_06/epoch_0',\n                'efficientnet_v2s_168_2021_04_23__22_05_47/epoch_0',\n                'seresnext50_32x4d_168_2021_04_24__08_50_52/epoch_0',\n                'tf_efficientnet_b5_ns_168_2021_05_01__20_24_56/epoch_0']\n\ncrop_mdls = []\nfor w in crop_weights:\n    parts = w.split('_2021')[0].split('_')\n    crop_mdl = '_'.join(parts[:-1])\n    crop_mdls.append(crop_mdl)\n    print(crop_mdl)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"full_weights = ['eca_nfnet_l0_512_2021_04_29__16_56_08/epoch_9',\n                'eca_nfnet_l0_512_2021_04_29__22_06_03/epoch_9',\n                'eca_nfnet_l1_512_2021_04_30__00_35_08/epoch_9',\n                'eca_nfnet_l1_512_2021_04_30__00_36_12/epoch_9',\n                'ecaresnet50t_512_2021_04_30__15_18_41/epoch_9',\n                'ecaresnet50t_512_2021_04_30__15_18_52/epoch_9',\n                'efficientnet_v2s_512_2021_04_30__12_21_17/epoch_9',\n                'efficientnet_v2s_512_2021_04_30__15_14_15/epoch_9',\n                'tf_efficientnet_b5_ns_512_2021_05_01__14_58_14/epoch_9',\n                'tf_efficientnet_b5_ns_512_2021_05_01__14_58_27/epoch_9']\n\nfull_mdls = []\nfor w in full_weights:\n    parts = w.split('_2021')[0].split('_')\n    full_mdl = '_'.join(parts[:-1])\n    full_mdls.append(full_mdl)\n    print(full_mdl)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"arcface_weights = 'eca_nfnet_l0_512_2021_04_25__23_37_16/epoch_14'\n\nparts = arcface_weights.split('_2021')[0].split('_')\n\narcface_mdl = '_'.join(parts[:-1])\n\nprint(arcface_mdl)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with open(f\"/kaggle/input/hpaembeddings/embeddings_{arcface_weights.replace('/', '_')}/embeddings.pickle\", 'rb') as f:\n    embeddings = pickle.load(f)\n\nembeddings_labels = np.zeros((len(embeddings), num_classes), dtype=np.float32)\nembeddings_features = np.zeros((len(embeddings), 512), dtype=np.float32)\n\nfor i, (k,v) in enumerate(embeddings.items()):\n    embeddings_labels[i] = v['labels']\n    embeddings_features[i] = v['embeddings']\n\nembeddings_features /= np.linalg.norm(embeddings_features, axis=1, keepdims=True)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"gpu_index = faiss.IndexFlatIP(embeddings_features.shape[1])\ngpu_index = faiss.index_cpu_to_all_gpus(gpu_index)\ngpu_index.add(embeddings_features)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def read_img(image_id_path):\n    img = cv2.imread(image_id_path, 0)\n    return img\n\ndef load_RGBY_images(image_id_path):\n    \n    red_image = read_img(image_id_path+\"_red.png\")\n    green_image = read_img(image_id_path+\"_green.png\")\n    blue_image = read_img(image_id_path+\"_blue.png\")\n    yellow_image = read_img(image_id_path+\"_yellow.png\")\n    \n    return red_image, green_image, blue_image, yellow_image\n\ndef encode_binary_mask(mask):\n\n    if mask.dtype != np.bool:\n        raise ValueError(\"encode_binary_mask expects a binary mask, received dtype == %s\" % mask.dtype)\n\n    mask = np.squeeze(mask)\n    if len(mask.shape) != 2:\n        raise ValueError(\"encode_binary_mask expects a 2d mask, received shape == %s\" % mask.shape)\n    \n    mask = mask.astype(np.uint8)\n    \n    mask_to_encode = mask.reshape(mask.shape[0], mask.shape[1], 1)\n    mask_to_encode = np.asfortranarray(mask_to_encode)\n    \n    encoded_mask = coco_mask.encode(mask_to_encode)[0][\"counts\"]\n    \n    binary_str = zlib.compress(encoded_mask, zlib.Z_BEST_COMPRESSION)\n    base64_str = base64.b64encode(binary_str)\n    \n    return base64_str.decode()\n\ndef compute_M(data):\n    cols = np.arange(data.size)\n    return csr_matrix((cols, (data.ravel(), cols)), shape=(data.max() + 1, data.size))\n\ndef get_indices_sparse(data):\n    M = compute_M(data)\n    return [np.unravel_index(row.data, data.shape) for row in M]\n\nclass CellSegmentator(object):\n    \"\"\"Uses pretrained DPN-Unet models to segment cells from images.\"\"\"\n\n    NORMALIZE = {\"mean\": [124 / 255, 117 / 255, 104 / 255], \"std\": [1 / (0.0167 * 255)] * 3}\n    \n    def __init__(\n            self,\n            nuclei_model=\"./nuclei_model.pth\",\n            cell_model=\"./cell_model.pth\",\n            model_width_height=None,\n            device=\"cuda\",\n            multi_channel_model=True,\n            return_without_scale_restore=False,\n            scale_factor=0.25,\n            padding=False\n    ):\n\n        if device != \"cuda\" and device != \"cpu\" and \"cuda\" not in device:\n            raise ValueError(f\"{device} is not a valid device (cuda/cpu)\")\n        if device != \"cpu\":\n            try:\n                assert torch.cuda.is_available()\n            except AssertionError:\n                print(\"No GPU found, using CPU.\", file=sys.stderr)\n                device = \"cpu\"\n                \n        self.device = device\n\n        if isinstance(nuclei_model, str):\n            if not os.path.exists(nuclei_model):\n                print(\n                    f\"Could not find {nuclei_model}. Downloading it now\",\n                    file=sys.stderr,\n                )\n                download_with_url(NUCLEI_MODEL_URL, nuclei_model)\n            nuclei_model = torch.load(\n                nuclei_model, map_location=torch.device(self.device)\n            )\n        if isinstance(nuclei_model, torch.nn.DataParallel) and device == \"cpu\":\n            nuclei_model = nuclei_model.module\n\n        self.nuclei_model = nuclei_model.to(self.device)\n\n        self.multi_channel_model = multi_channel_model\n        if isinstance(cell_model, str):\n            if not os.path.exists(cell_model):\n                print(\n                    f\"Could not find {cell_model}. Downloading it now\", file=sys.stderr\n                )\n                if self.multi_channel_model:\n                    download_with_url(MULTI_CHANNEL_CELL_MODEL_URL, cell_model)\n                else:\n                    download_with_url(TWO_CHANNEL_CELL_MODEL_URL, cell_model)\n            cell_model = torch.load(cell_model, map_location=torch.device(self.device))\n        self.cell_model = cell_model.to(self.device)\n        self.model_width_height = model_width_height\n        self.return_without_scale_restore = return_without_scale_restore\n        self.scale_factor = scale_factor\n        self.padding = padding\n\n    def _image_conversion(self, images):\n\n        microtubule_imgs, er_imgs, nuclei_imgs = images\n        if self.multi_channel_model:\n            if not isinstance(er_imgs, list):\n                raise ValueError(\"Please speicify the image path(s) for er channels!\")\n        else:\n            if not er_imgs is None:\n                raise ValueError(\n                    \"second channel should be None for two channel model predition!\"\n                )\n\n        if not isinstance(microtubule_imgs, list):\n            raise ValueError(\"The microtubule images should be a list\")\n        if not isinstance(nuclei_imgs, list):\n            raise ValueError(\"The microtubule images should be a list\")\n\n        if er_imgs:\n            if not len(microtubule_imgs) == len(er_imgs) == len(nuclei_imgs):\n                raise ValueError(\"The lists of images needs to be the same length\")\n        else:\n            if not len(microtubule_imgs) == len(nuclei_imgs):\n                raise ValueError(\"The lists of images needs to be the same length\")\n\n        if not all(isinstance(item, np.ndarray) for item in microtubule_imgs):\n            microtubule_imgs = [\n                os.path.expanduser(item) for _, item in enumerate(microtubule_imgs)\n            ]\n            nuclei_imgs = [\n                os.path.expanduser(item) for _, item in enumerate(nuclei_imgs)\n            ]\n\n            microtubule_imgs = list(\n                map(lambda item: imageio.imread(item), microtubule_imgs)\n            )\n            nuclei_imgs = list(map(lambda item: imageio.imread(item), nuclei_imgs))\n            if er_imgs:\n                er_imgs = [os.path.expanduser(item) for _, item in enumerate(er_imgs)]\n                er_imgs = list(map(lambda item: imageio.imread(item), er_imgs))\n\n        if not er_imgs:\n            er_imgs = [\n                np.zeros(item.shape, dtype=item.dtype)\n                for _, item in enumerate(microtubule_imgs)\n            ]\n        cell_imgs = list(\n            map(\n                lambda item: np.dstack((item[0], item[1], item[2])),\n                list(zip(microtubule_imgs, er_imgs, nuclei_imgs)),\n            )\n        )\n\n        return cell_imgs\n\n    def _pad(self, image):\n        \n        rows, cols = image.shape[:2]\n        self.scaled_shape = rows, cols\n        img_pad= cv2.copyMakeBorder(\n                    image,\n                    32,\n                    (32 - rows % 32),\n                    32,\n                    (32 - cols % 32),\n                    cv2.BORDER_REFLECT,\n                )\n        \n        return img_pad\n\n    def pred_nuclei(self, images):\n\n        def _preprocess(images):\n            if isinstance(images[0], str):\n                raise NotImplementedError('Currently the model requires images as numpy arrays, not paths.')\n                # images = [imageio.imread(image_path) for image_path in images]\n            self.target_shapes = [image.shape for image in images]\n            #print(images.shape)\n            #resize like in original implementation with https://scikit-image.org/docs/dev/api/skimage.transform.html#skimage.transform.resize\n            if self.model_width_height:\n                images = np.array([transform.resize(image, (self.model_width_height,self.model_width_height)) \n                                  for image in images])\n            else:\n                images = [transform.rescale(image, self.scale_factor) for image in images]\n\n            if self.padding:\n                images = [self._pad(image) for image in images]\n\n            nuc_images = np.array([np.dstack((image[..., 2], image[..., 2], image[..., 2])) if len(image.shape) >= 3\n                                   else np.dstack((image, image, image)) for image in images])\n            \n            nuc_images = nuc_images.transpose([0, 3, 1, 2])\n            #print(\"nuc\", nuc_images.shape)\n\n            return nuc_images\n\n        def _segment_helper(imgs):\n            with torch.no_grad():\n                mean = torch.as_tensor(self.NORMALIZE[\"mean\"], device=self.device)\n                std = torch.as_tensor(self.NORMALIZE[\"std\"], device=self.device)\n                imgs = torch.tensor(imgs).float()\n                imgs = imgs.to(self.device)\n                imgs = imgs.sub_(mean[:, None, None]).div_(std[:, None, None])\n\n                imgs = self.nuclei_model(imgs)\n                imgs = F.softmax(imgs, dim=1)\n                return imgs\n\n        preprocessed_imgs = _preprocess(images)\n        predictions = _segment_helper(preprocessed_imgs)\n        predictions = predictions.to(\"cpu\").numpy()\n        #dont restore scaling, just save and scale later ...\n        predictions = [self._restore_scaling(util.img_as_ubyte(pred), target_shape)\n                       for pred, target_shape in zip(predictions, self.target_shapes)]\n        return predictions\n\n    def _restore_scaling(self, n_prediction, target_shape):\n        \"\"\"Restore an image from scaling and padding.\n        This method is intended for internal use.\n        It takes the output from the nuclei model as input.\n        \"\"\"\n        n_prediction = n_prediction.transpose([1, 2, 0])\n        if self.padding:\n            n_prediction = n_prediction[\n                32 : 32 + self.scaled_shape[0], 32 : 32 + self.scaled_shape[1], ...\n            ]\n        n_prediction[..., 0] = 0\n        if not self.return_without_scale_restore:\n            n_prediction = cv2.resize(\n                n_prediction,\n                (target_shape[0], target_shape[1]),\n                #try INTER_NEAREST_EXACT\n                interpolation=cv2.INTER_AREA,\n            )\n        return n_prediction\n\n    def pred_cells(self, images, precombined=False):\n\n        def _preprocess(images):\n            self.target_shapes = [image.shape for image in images]\n            for image in images:\n                if not len(image.shape) == 3:\n                    raise ValueError(\"image should has 3 channels\")\n            #resize like in original implementation with https://scikit-image.org/docs/dev/api/skimage.transform.html#skimage.transform.resize\n            if self.model_width_height:\n                images = np.array([transform.resize(image, (self.model_width_height,self.model_width_height)) \n                                  for image in images])\n            else:\n                images = np.array([transform.rescale(image, self.scale_factor, multichannel=True) for image in images])\n\n            if self.padding:\n                images = np.array([self._pad(image) for image in images])\n\n            cell_images = images.transpose([0, 3, 1, 2])\n\n            return cell_images\n\n        def _segment_helper(imgs):\n            with torch.no_grad():\n                mean = torch.as_tensor(self.NORMALIZE[\"mean\"], device=self.device)\n                std = torch.as_tensor(self.NORMALIZE[\"std\"], device=self.device)\n                imgs = torch.tensor(imgs).float()\n                imgs = imgs.to(self.device)\n                imgs = imgs.sub_(mean[:, None, None]).div_(std[:, None, None])\n                imgs = self.cell_model(imgs)\n                imgs = F.softmax(imgs, dim=1)\n                return imgs\n\n        if not precombined:\n            images = self._image_conversion(images)\n        preprocessed_imgs = _preprocess(images)\n        predictions = _segment_helper(preprocessed_imgs)\n        predictions = predictions.to(\"cpu\").numpy()\n        predictions = [self._restore_scaling(util.img_as_ubyte(pred), target_shape)\n                       for pred, target_shape in zip(predictions, self.target_shapes)]\n        \n        return predictions\n\ndef __fill_holes(image):\n    \"\"\"Fill_holes for labelled image, with a unique number.\"\"\"\n    boundaries = segmentation.find_boundaries(image)\n    image = np.multiply(image, np.invert(boundaries))\n    image = ndi.binary_fill_holes(image > 0)\n    image = ndi.label(image)[0]\n    return image\n\ndef label_cell_vlad(nuclei_pred, cell_pred, img_size=512, return_nuclei_label=True):\n    \"\"\"Label the cells and the nuclei.\n\n    Keyword arguments:\n    nuclei_pred -- a 3D numpy array of a prediction from a nuclei image.\n    cell_pred -- a 3D numpy array of a prediction from a cell image.\n\n    Returns:\n    A tuple containing:\n    nuclei-label -- A nuclei mask data array.\n    cell-label  -- A cell mask data array.\n\n    0's in the data arrays indicate background while a continous\n    strech of a specific number indicates the area for a specific\n    cell.\n    The same value in cell mask and nuclei mask refers to the identical cell.\n\n    NOTE: The nuclei labeling from this function will be sligthly\n    different from the values in :func:`label_nuclei` as this version\n    will use information from the cell-predictions to make better\n    estimates.\n    \"\"\"\n    def __wsh(\n        mask_img,\n        threshold,\n        border_img,\n        seeds,\n        threshold_adjustment=0.35,\n        small_object_size_cutoff=10,\n    ):\n        img_copy = np.zeros_like(mask_img)\n        m = seeds * border_img  # * dt\n        img_copy[m > threshold + threshold_adjustment] = 1\n        img_copy = img_copy.astype(np.bool)\n        img_copy = remove_small_objects(img_copy, small_object_size_cutoff).astype(\n            np.uint8\n        )\n\n        mask_img = np.where(mask_img <= threshold, 0, 1)\n        mask_img = mask_img.astype(np.bool)\n        \n        ### New segmentation ###\n        mask_img = remove_small_holes(mask_img, int(63 * (img_size / 512)**2))\n        ########################\n        \n        mask_img = remove_small_objects(mask_img, 8).astype(np.uint8)\n        markers = ndi.label(img_copy, output=np.uint32)[0]\n        labeled_array = segmentation.watershed(\n            mask_img, markers, mask=mask_img, watershed_line=True\n        )\n        return labeled_array\n\n    nuclei_label = __wsh(\n        nuclei_pred[..., 2] / 255.0,\n        0.4,\n        1 - (nuclei_pred[..., 1] + cell_pred[..., 1]) / 255.0 > 0.05,\n        nuclei_pred[..., 2] / 255,\n        threshold_adjustment=-0.25,\n        small_object_size_cutoff=32,\n    )\n\n    # for hpa_image, to remove the small pseduo nuclei\n    \n    ### New segmentation ###\n    nuclei_label = remove_small_objects(nuclei_label, int(157 * (img_size / 512)**2))\n    ########################\n    \n    nuclei_label = measure.label(nuclei_label)\n    # this is to remove the cell borders' signal from cell mask.\n    # could use np.logical_and with some revision, to replace this func.\n    # Tuned for segmentation hpa images\n    threshold_value = max(0.22, filters.threshold_otsu(cell_pred[..., 2] / 255) * 0.5)\n    # exclude the green area first\n    cell_region = np.multiply(\n        cell_pred[..., 2] / 255 > threshold_value,\n        np.invert(np.asarray(cell_pred[..., 1] / 255 > 0.05, dtype=np.int8)),\n    )\n    sk = np.asarray(cell_region, dtype=np.int8)\n    \n    ################### CHANGES HERE ###################################\n \n    ####################### V4 ####################\n    distance_old = np.clip(cell_pred[..., 2], 255 * threshold_value, cell_pred[..., 2])\n    cell_label_old = segmentation.watershed(-distance_old, nuclei_label, mask=sk)\n    \n    ### New segmentation ###\n    cell_label_old = remove_small_objects(cell_label_old, int(344 * (img_size / 512)**2)).astype(np.uint8)\n    ########################\n    \n    distance = distance_old.copy()\n    distance[distance<225] = 0\n    \n    cell_label = segmentation.watershed(-distance, nuclei_label, mask=distance)\n    \n    ### New segmentation ###\n    cell_label = remove_small_objects(cell_label, int(344 * (img_size / 512)**2)).astype(np.uint8)\n    ########################\n    \n    unqs = np.unique(cell_label)\n    if 0 in unqs:\n        unqs = unqs[1:]\n        \n    lst = [cv2.dilate((cell_label==unq).astype(np.uint8), kernel=np.ones((10, 10), np.uint8), iterations=1) for unq in unqs]\n    cell_label = np.zeros_like(cell_label)\n    \n    for i, l in enumerate(lst):\n        cell_label[l==True] = unqs[i]\n                \n    for unq in unqs:\n        smth = cell_label_old==unq\n        if True not in smth[0,:] and True not in smth[-1,:] and True not in smth[:,0] and True not in smth[:,-1]:\n            cell_label[cell_label==unq] = 0\n            cell_label[cell_label_old==unq] = unq\n            \n    ################### CHANGES HERE ###################################\n    \n#     selem = disk(max(1, int(6 * 2048 / img_size)))\n    selem = disk(6)\n    cell_label = closing(cell_label, selem)\n    cell_label = __fill_holes(cell_label)\n    # this part is to use green channel, and extend cell label to green channel\n    # benefit is to exclude cells clear on border but without nucleus\n    sk = np.asarray(\n        np.add(\n            np.asarray(cell_label > 0, dtype=np.int8),\n            np.asarray(cell_pred[..., 1] / 255 > 0.05, dtype=np.int8),\n        )\n        > 0,\n        dtype=np.int8,\n    )\n    cell_label = segmentation.watershed(-distance, cell_label, mask=sk)\n    cell_label = __fill_holes(cell_label)\n    cell_label = np.asarray(cell_label > 0, dtype=np.uint8)\n    cell_label = measure.label(cell_label)\n    \n    ### New segmentation ###\n    cell_label = remove_small_objects(cell_label, int(344 * (img_size / 512)**2))\n    ########################\n    \n    cell_label = measure.label(cell_label)\n    cell_label = np.asarray(cell_label, dtype=np.uint16)\n    if not return_nuclei_label:\n        return cell_label\n    nuclei_label = np.multiply(cell_label > 0, nuclei_label) > 0\n    nuclei_label = measure.label(nuclei_label)\n    \n    ### New segmentation ###\n    nuclei_label = remove_small_objects(nuclei_label, int(157 * (img_size / 512)**2))\n    ########################\n    \n    nuclei_label = np.multiply(cell_label, nuclei_label > 0)\n\n    return nuclei_label, cell_label\n\ndef crop_net(crop_mdl, pretrained=False, **kwargs):\n    model = timm.create_model(crop_mdl, pretrained=pretrained, num_classes=num_classes)\n    return model\n\ndef full_net(full_mdl, pretrained=False, **kwargs):\n    model = timm.create_model(full_mdl, pretrained=pretrained, num_classes=num_classes)\n    return model\n\ndef arcface_net(pretrained=False, **kwargs):\n    model = timm.create_model('eca_nfnet_l0', pretrained=pretrained, num_classes=num_classes)\n    return model\n\nclass ArcMarginProduct(torch.nn.Module):\n\n    def __init__(self, in_features=512, out_features=11582):\n        super().__init__()\n        self.weight = torch.nn.Parameter(torch.FloatTensor(out_features, in_features))\n        self.reset_parameters()\n        \n    def reset_parameters(self):\n        stdv = 1. / math.sqrt(self.weight.size(1))\n        self.weight.data.uniform_(-stdv, stdv)\n        \n    def forward(self, x):\n        cosine = torch.functional.F.linear(torch.functional.F.normalize(x), torch.functional.F.normalize(self.weight.cuda()))\n        return cosine  \n    \nclass Customhead(torch.nn.Module):\n\n    def __init__(self, in_features=512, out_features=11582):\n        \n        super(Customhead, self).__init__()\n        \n        body = create_body(arcface_net, n_in=3, pretrained=False)\n        nf = num_features_model(torch.nn.Sequential(*body.children()))\n        \n        self.head = create_head(nf, n_out=in_features, concat_pool=True)\n        self.arc_margin = ArcMarginProduct(in_features, out_features)\n\n    def forward(self, features):\n        x = self.head(features)\n        return x","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"NUC_MODEL = '/kaggle/input/hpamisc/HPA-Cell-Segmentation-weights/dpn_unet_nuclei_v1.pth'\nCELL_MODEL = '/kaggle/input/hpamisc/HPA-Cell-Segmentation-weights/dpn_unet_cell_3ch_v1.pth'\n\nsegmentator = CellSegmentator(\n    NUC_MODEL,\n    CELL_MODEL,\n    device=\"cuda\",\n    multi_channel_model=True,\n    scale_factor=SEGMENTATION_SCALE,\n    padding=True,\n    return_without_scale_restore=True\n)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df = pd.read_csv('/kaggle/input/hpa-single-cell-image-classification/train.csv')[:1]\ndf['ID'] = df['ID'].map(lambda x: f'../input/hpa-single-cell-image-classification/train/{x}_red.png')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"crop_batch_tfms = [Normalize.from_stats(mean=[0.135, 0.085, 0.100], std=[0.152, 0.099, 0.175])]\ncrop_resizer = Resize(crop_input_size, method='squish')\n\ncrop_dblock = DataBlock(blocks=(ImageBlock, MultiCategoryBlock(vocab=[str(i) for i in range(19)])),\n                        splitter=RandomSplitter(),\n                        get_x=lambda x: x[0],\n                        get_y=lambda x: x['Label'].split('|'),\n                        item_tfms=[],\n                        batch_tfms=crop_batch_tfms)\n\ncrop_dls = crop_dblock.dataloaders(df, bs=64, val_bs=64)\n\nt = PILImage.create(np.zeros((crop_input_size,crop_input_size,3), dtype=np.uint8)) # Needed for the first iteration to initialize model on GPU\n\ncrop_models = []\nfor crop_mdl, w in zip(crop_mdls, crop_weights):\n    crop_model = cnn_learner(crop_dls, partial(crop_net, crop_mdl), pretrained=False).load(f\"/kaggle/input/hpaweights/{w}\")\n    crop_model.get_preds(dl=crop_dls.test_dl([t]))[0] # Needed for the first iteration to initialize model on GPU\n    crop_models.append(crop_model)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"full_batch_tfms = [Normalize.from_stats(mean=[0.110, 0.063, 0.066], std=[0.157, 0.098, 0.160])]\nfull_resizer = Resize(full_input_size, method='squish')\n\nfull_dblock = DataBlock(blocks=(ImageBlock, MultiCategoryBlock(vocab=[str(i) for i in range(19)])),\n                        splitter=RandomSplitter(),\n                        get_x=lambda x: x[0],\n                        get_y=lambda x: x['Label'].split('|'),\n                        item_tfms=[],\n                        batch_tfms=full_batch_tfms)\n\nfull_dls = full_dblock.dataloaders(df, bs=16, val_bs=16)\n\nt = PILImage.create(np.zeros((full_input_size,full_input_size,3), dtype=np.uint8)) # Needed for the first iteration to initialize model on GPU\n\nfull_models = []\nfor full_mdl, w in zip(full_mdls, full_weights):\n    full_model = cnn_learner(full_dls, partial(full_net, full_mdl), pretrained=False).load(f\"/kaggle/input/hpaweights/{w}\")\n    full_model.get_preds(dl=full_dls.test_dl([t]))[0] # Needed for the first iteration to initialize model on GPU\n    full_models.append(full_model)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"arcface_model = cnn_learner(full_dls, arcface_net, custom_head=Customhead(512, 11582), pretrained=False).load(f\"/kaggle/input/hpaweights/{arcface_weights}\")\narcface_model.get_preds(dl=full_dls.test_dl([t]))[0] # Needed for the first iteration to initialize model on GPU\nprint()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def remove_small_nuclei_on_border(nuclei_mask, nuclei_area_thresh=0.5):\n    \n    h,w = nuclei_mask.shape[:2]\n    \n    touching_border = []\n    \n    nuclei_uniques = np.unique(nuclei_mask)\n    if 0 in nuclei_uniques:\n        nuclei_uniques = nuclei_uniques[1:]\n        \n    for unq in nuclei_uniques:\n        idxs = np.where(nuclei_mask==unq)\n        y_min, y_max = min(idxs[0]), max(idxs[0])\n        x_min, x_max = min(idxs[1]), max(idxs[1])\n        if x_min == 0 or y_min == 0 or x_max == (w - 1) or y_max == (h - 1):\n            touching_border.append(unq)\n    \n    nuclei_areas = np.array([np.count_nonzero(nuclei_mask==unq) for unq in nuclei_uniques])\n    not_touching_border_idxs = [i for i, unq in enumerate(nuclei_uniques) if unq not in touching_border]\n\n    ignore_nuclei_idxs = nuclei_uniques[(nuclei_areas < np.median(nuclei_areas[not_touching_border_idxs])*nuclei_area_thresh) & np.isin(nuclei_uniques, touching_border)]\n    \n    return ignore_nuclei_idxs\n\ndef do_segmentation(ryb_image, blue_image):\n    \n    img_size = blue_image.shape[0] * SEGMENTATION_SCALE\n\n    nuc_segmentation = segmentator.pred_nuclei([blue_image])\n    cell_segmentation = segmentator.pred_cells([ryb_image], precombined=True)\n    \n    nuclei_mask, cell_mask = label_cell_vlad(nuc_segmentation[0], cell_segmentation[0], img_size, return_nuclei_label=True)\n    \n    cell_mask = cell_mask.astype(np.uint8)\n\n    # Remove border cells with nuclei_area < median/2\n    small_border_nuclei_idxs = remove_small_nuclei_on_border(nuclei_mask)\n\n    # Remove cells without nuclei\n    cell_uniques = np.unique(cell_mask)\n    if 0 in cell_uniques:\n        cell_uniques = cell_uniques[1:]\n        \n    nuclei_uniques = np.unique(nuclei_mask)\n    if 0 in nuclei_uniques:\n        nuclei_uniques = nuclei_uniques[1:]\n            \n    cells_without_nuclei_idxs = np.setdiff1d(np.union1d(cell_uniques, nuclei_uniques), np.intersect1d(cell_uniques, nuclei_uniques))\n    \n    if len(small_border_nuclei_idxs):\n        small_border_cell_mask = np.array([cell_mask != i for i in small_border_nuclei_idxs]).prod(axis=0).astype(np.uint8)\n    else:\n        small_border_cell_mask = np.ones(cell_mask.shape, dtype=np.uint8)\n        \n    for ig in np.union1d(small_border_nuclei_idxs, cells_without_nuclei_idxs):\n        cell_mask[cell_mask == ig] = 0\n        \n    return cell_mask, small_border_cell_mask\n\ndef get_bboxes(cell_mask, scale_factor):\n    \n    cell_bboxes = {}\n\n    unqs = np.unique(cell_mask)\n    if 0 in unqs:\n        unqs = unqs[1:]\n    \n    bboxes = get_indices_sparse(cell_mask)\n    \n    for c in unqs:\n        w, h = bboxes[c]\n        x_0, x_1, y_0, y_1 = w.min(), w.max(), h.min(), h.max()\n        \n        bbox = [int(scale_factor*x_0), int((x_1+1)*scale_factor), int(y_0*scale_factor), int((y_1+1)*scale_factor)]\n        cell_bboxes[c] = bbox\n    \n    return cell_bboxes\n\ndef do_preprocessing(rgb_image, cell_bboxes, cell_mask, small_border_cell_mask):\n    \n    crops_batch = []\n    full_crops_batch = []\n    \n    arcface_image = full_resizer(PILImage.create(rgb_image))\n    full_image = arcface_image * cv2.resize(small_border_cell_mask, (full_input_size, full_input_size), interpolation=cv2.INTER_NEAREST)[..., None]\n\n    resized_mask = cv2.resize(cell_mask, (full_input_size, full_input_size), interpolation=cv2.INTER_NEAREST)\n    \n    for cell_idx, cell_bbox in cell_bboxes.items():\n        \n        temp_mask = (cell_mask == cell_idx).astype(np.uint8)\n        \n        # CROPS\n        temp_crop = (rgb_image * temp_mask[..., None])[cell_bbox[0]:cell_bbox[1], cell_bbox[2]:cell_bbox[3]]\n        crops_batch.append(crop_resizer(PILImage.create(temp_crop)))\n        \n        # FULL CROPS\n        masked_img_resized = full_image * (resized_mask == cell_idx).astype(np.uint8)[..., None] \n        full_crops_batch.append(masked_img_resized)\n    \n    return crops_batch, full_crops_batch, full_image, arcface_image","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"idx = 0\n\nimage_id, ImageWidth, ImageHeight, PredictionString = test_df.iloc[idx]\n\n# Load images\nred_image, green_image, blue_image, yellow_image = load_RGBY_images(f\"{directory}/test/{image_id}\")\n\nimage_size = red_image.shape[0]\nscale_factor = 1/SEGMENTATION_SCALE\n\n# Segmentation\nryb_image = np.transpose(np.array([red_image, yellow_image, blue_image]), (1,2,0))\ncell_mask, small_border_cell_mask = do_segmentation(ryb_image, blue_image)\n\n# Get bboxes\ncell_bboxes = get_bboxes(cell_mask, scale_factor)\n\n# RLE\ncell_mask = cv2.resize(cell_mask, (image_size, image_size), interpolation=cv2.INTER_NEAREST)\nsmall_border_cell_mask = cv2.resize(small_border_cell_mask, (image_size, image_size), interpolation=cv2.INTER_NEAREST)\nrles = [encode_binary_mask(cell_mask==cell_id) for cell_id in cell_bboxes.keys()]\n\n# Preprocessing\nrgb_image = np.transpose(np.array([red_image, green_image, blue_image]), (1,2,0))\ncrops_batch, full_crops_batch, full_image, arcface_image = do_preprocessing(rgb_image, cell_bboxes, cell_mask, small_border_cell_mask)\n\nwith torch.no_grad():\n\n    # Crops Classification\n    crop_dl = crop_dls.test_dl(crops_batch)\n\n    y_crops = []\n    for crop_model in crop_models:\n        preds = []\n        for batch in crop_dl:\n            preds.extend(crop_model.model(batch[0]))\n        preds = torch.sigmoid(torch.stack(preds))\n        y_crops.append(preds)\n\n    y_crops = torch.mean(torch.stack(y_crops), axis=0)\n\n    # Full Crops Classification\n    full_dl = full_dls.test_dl(full_crops_batch)\n\n    y_full_crops = []\n    for full_model in full_models:\n        preds = []\n        for batch in full_dl:\n            preds.extend(full_model.model(batch[0]))\n        preds = torch.sigmoid(torch.stack(preds))\n        y_full_crops.append(preds)\n\n    y_full_crops = torch.mean(torch.stack(y_full_crops), axis=0)\n\n    # Full Image Classification\n    full_image_dl = full_dls.test_dl([full_image])\n\n    y_full_image = []\n    for full_model in full_models:\n        preds = []\n        for batch in full_image_dl:\n            preds.extend(full_model.model(batch[0]))\n        preds = torch.sigmoid(torch.stack(preds))\n        y_full_image.append(preds)\n\n    y_full_image = torch.mean(torch.stack(y_full_image), axis=0).flatten()\n\n    # Arcface Prediction\n    arcface_image_dl = full_dls.test_dl([arcface_image])\n\n    preds = []\n    for batch in arcface_image_dl:\n        preds.extend(arcface_model.model(batch[0]))\n    y_arcface = torch.stack(preds).cpu().numpy()\n\n    y_arcface /= np.linalg.norm(y_arcface, axis=1, keepdims=True)\n\nall_dists, all_topk_idxs = gpu_index.search(x=y_arcface, k=1)\n\nif all_dists.item() >= ARCFACE_THRESH:\n    y_full_image = embeddings_labels[all_topk_idxs.item()]\n\nfor cell_id in range(len(cell_bboxes)):\n    for class_id in range(num_classes):\n        conf = y_crops[cell_id, class_id]/3 + y_full_crops[cell_id, class_id]/3 + y_full_image[class_id]/3","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.imshow(cell_mask)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"fig, axs = plt.subplots(figsize=(15, 13), sharex=True, sharey=True, ncols=5, nrows=len(crops_batch) // 5 + 1 if len(crops_batch) % 5 else len(crops_batch) // 5)\n\nfor i, (cr, k) in enumerate(zip(crops_batch, cell_bboxes.keys())):\n    axs.reshape(-1)[i].imshow(cr)\n    axs.reshape(-1)[i].set_title(f'Cell {k}')\n    \nif len(crops_batch) % 5:\n    for i in range(len(crops_batch), len(crops_batch) + 5 - len(crops_batch) % 5):\n        axs.reshape(-1)[i].set_axis_off()\n        \nplt.tight_layout()\nplt.show()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"fig, axs = plt.subplots(figsize=(15, 13), sharex=True, sharey=True, ncols=5, nrows=len(full_crops_batch) // 5 + 1 if len(full_crops_batch) % 5 else len(full_crops_batch) // 5)\n\nfor i, (cr, k) in enumerate(zip(full_crops_batch, cell_bboxes.keys())):\n    axs.reshape(-1)[i].imshow(cr)\n    axs.reshape(-1)[i].set_title(f'Cell {k}')\n    \nif len(full_crops_batch) % 5:\n    for i in range(len(full_crops_batch), len(full_crops_batch) + 5 - len(full_crops_batch) % 5):\n        axs.reshape(-1)[i].set_axis_off()\n        \nplt.tight_layout()\nplt.show()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.imshow(full_image)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.imshow(arcface_image)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with open('submission.csv', 'w') as outf:\n        \n    print('ID,ImageWidth,ImageHeight,PredictionString', file=outf)\n    \n    for idx in tqdm(range(len(test_df)), total=len(test_df)):\n        \n        image_id, ImageWidth, ImageHeight, PredictionString = test_df.iloc[idx]\n                \n        if (PUBLIC_ONLY and (image_id in public_image_ids)) or not PUBLIC_ONLY:\n            \n            # Load images\n            red_image, green_image, blue_image, yellow_image = load_RGBY_images(f\"{directory}/test/{image_id}\")\n            \n            image_size = red_image.shape[0]\n            scale_factor = 1/SEGMENTATION_SCALE\n            \n            # Segmentation\n            ryb_image = np.transpose(np.array([red_image, yellow_image, blue_image]), (1,2,0))\n            cell_mask, small_border_cell_mask = do_segmentation(ryb_image, blue_image)\n\n            # Get bboxes\n            cell_bboxes = get_bboxes(cell_mask, scale_factor)\n            \n            # RLE\n            cell_mask = cv2.resize(cell_mask, (image_size, image_size), interpolation=cv2.INTER_NEAREST)\n            small_border_cell_mask = cv2.resize(small_border_cell_mask, (image_size, image_size), interpolation=cv2.INTER_NEAREST)\n            rles = [encode_binary_mask(cell_mask==cell_id) for cell_id in cell_bboxes.keys()]\n            \n            # Preprocessing\n            rgb_image = np.transpose(np.array([red_image, green_image, blue_image]), (1,2,0))\n            crops_batch, full_crops_batch, full_image, arcface_image = do_preprocessing(rgb_image, cell_bboxes, cell_mask, small_border_cell_mask)\n\n            with torch.no_grad():\n                \n                # Crops Classification\n                crop_dl = crop_dls.test_dl(crops_batch)\n\n                y_crops = []\n                for crop_model in crop_models:\n                    preds = []\n                    for batch in crop_dl:\n                        preds.extend(crop_model.model(batch[0]))\n                    preds = torch.sigmoid(torch.stack(preds))\n                    y_crops.append(preds)\n\n                y_crops = torch.mean(torch.stack(y_crops), axis=0)\n\n                # Full Crops Classification\n                full_dl = full_dls.test_dl(full_crops_batch)\n\n                y_full_crops = []\n                for full_model in full_models:\n                    preds = []\n                    for batch in full_dl:\n                        preds.extend(full_model.model(batch[0]))\n                    preds = torch.sigmoid(torch.stack(preds))\n                    y_full_crops.append(preds)\n\n                y_full_crops = torch.mean(torch.stack(y_full_crops), axis=0)\n\n                # Arcface Prediction\n                arcface_image_dl = full_dls.test_dl([arcface_image])\n\n                preds = []\n                for batch in arcface_image_dl:\n                    preds.extend(arcface_model.model(batch[0]))\n                y_arcface = torch.stack(preds).cpu().numpy()\n\n                y_arcface /= np.linalg.norm(y_arcface, axis=1, keepdims=True)\n\n                all_dists, all_topk_idxs = gpu_index.search(x=y_arcface, k=1)\n\n                if all_dists.item() >= ARCFACE_THRESH:\n                    y_full_image = embeddings_labels[all_topk_idxs.item()]\n                else:\n                    # Full Image Classification\n                    full_image_dl = full_dls.test_dl([full_image])\n\n                    y_full_image = []\n                    for full_model in full_models:\n                        preds = []\n                        for batch in full_image_dl:\n                            preds.extend(full_model.model(batch[0]))\n                        preds = torch.sigmoid(torch.stack(preds))\n                        y_full_image.append(preds)\n\n                    y_full_image = torch.mean(torch.stack(y_full_image), axis=0).flatten()\n                            \n            # Submission\n            pred_strs = []\n        \n            for cell_id in range(len(cell_bboxes)):\n                \n                rle = rles[cell_id]\n                \n                for class_id in range(num_classes):\n                    \n                    conf = y_crops[cell_id, class_id]/3 + y_full_crops[cell_id, class_id]/3 + y_full_image[class_id]/3\n                    pred_strs.append(f\"{class_id} {conf} {rle}\")\n                    \n            PredictionString = ' '.join(pred_strs)\n                            \n        print(f\"{image_id},{ImageWidth},{ImageHeight},{PredictionString}\", file=outf)","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}