{"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 -q ../input/pycocotools/pycocotools-2.0-cp37-cp37m-linux_x86_64.whl","metadata":{"execution":{"iopub.status.busy":"2021-06-17T03:53:41.400491Z","iopub.execute_input":"2021-06-17T03:53:41.400812Z","iopub.status.idle":"2021-06-17T03:54:08.810976Z","shell.execute_reply.started":"2021-06-17T03:53:41.400735Z","shell.execute_reply":"2021-06-17T03:54:08.809963Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import sys\nsys.path = [\n    '../input/smp20210127/pytorch-image-models-master/pytorch-image-models-master',  # timm\n    '../input/hpapytorchzoozip/pytorch_zoo-master/',\n    '../input/hpa-seg/HPA-Cell-Segmentation/hpacellseg',\n    '../input/hpafinal'\n] + sys.path\nimport warnings\nwarnings.filterwarnings(\"ignore\")","metadata":{"execution":{"iopub.status.busy":"2021-06-17T03:54:44.870594Z","iopub.execute_input":"2021-06-17T03:54:44.87103Z","iopub.status.idle":"2021-06-17T03:54:44.87582Z","shell.execute_reply.started":"2021-06-17T03:54:44.870985Z","shell.execute_reply":"2021-06-17T03:54:44.87478Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport random\nimport numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\n\nfrom skimage import transform, util\nimport glob\nimport cv2\nfrom matplotlib.patches import Rectangle\nimport matplotlib.pyplot as plt\nimport base64\nfrom pycocotools import _mask as coco_mask\nimport typing as t\nimport zlib\nimport os.path\nimport urllib\nimport zipfile\nimport scipy.ndimage as ndi\nfrom skimage import filters, measure, segmentation\nfrom skimage.morphology import (binary_erosion, closing, disk,\n                                remove_small_holes, remove_small_objects)\n\nimport albumentations \nimport torch.nn.functional as F\nimport torch\nfrom torch import nn\nfrom torch.utils.data import DataLoader,Dataset\nfrom torch.utils.data import TensorDataset, DataLoader, Dataset\nimport timm\nfrom model import class_densenet121_dropout  # daishu\nimport PIL\nimport gc\nimport torch.cuda.amp as amp\n\ndevice = torch.device('cuda')","metadata":{"execution":{"iopub.status.busy":"2021-06-17T03:54:47.708255Z","iopub.execute_input":"2021-06-17T03:54:47.708776Z","iopub.status.idle":"2021-06-17T03:54:52.468961Z","shell.execute_reply.started":"2021-06-17T03:54:47.708736Z","shell.execute_reply":"2021-06-17T03:54:52.467982Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"seg_size = 512\nseg_bs = 8388608 // seg_size ** 2\nseg_TTA = 8\nsmall_th_dict = {\n    2048: 500,\n    1024: 125,\n    512 : 32,\n}\nsmall_th = small_th_dict[seg_size]\nmask_dir = 'test_mask_npz_fullsize_cell_mask'\n\nmodel_dirs = [\n    '../input/bo-hpa-models',  # bo\n    '../input/bo-hpa-models-3d256',  # bo\n    '../input/hpa-models',  # gary\n    '../input/hpa-models-qishen',\n]\n\nTTA = {\n    'orig': 2,\n    'masked': 3,\n    'cells_128': 8,\n    'cells_256': 6,\n    'center_cells': 3,\n}\nn_ch = 4\nnum_classes = 19\nimage_size = 512\n\norig_mean = [239.93038613, 246.05603962, 250.16871503, 250.50623682]\n\ndata_dir = '../input/hpa-single-cell-image-classification/test/'\ndf_sub = pd.read_csv('../input/hpa-single-cell-image-classification/sample_submission.csv')\ndf_sub = df_sub.head(10) if df_sub.shape[0] == 559 else df_sub\ndf_sub.shape","metadata":{"execution":{"iopub.status.busy":"2021-06-17T03:54:52.470548Z","iopub.execute_input":"2021-06-17T03:54:52.470922Z","iopub.status.idle":"2021-06-17T03:54:52.498656Z","shell.execute_reply.started":"2021-06-17T03:54:52.470886Z","shell.execute_reply":"2021-06-17T03:54:52.497889Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"os.makedirs(mask_dir, exist_ok=True)","metadata":{"execution":{"iopub.status.busy":"2021-06-17T03:54:55.323714Z","iopub.execute_input":"2021-06-17T03:54:55.32404Z","iopub.status.idle":"2021-06-17T03:54:55.330176Z","shell.execute_reply.started":"2021-06-17T03:54:55.32401Z","shell.execute_reply":"2021-06-17T03:54:55.329248Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"NORMALIZE = {\"mean\": [124 / 255, 117 / 255, 104 / 255], \"std\": [1 / (0.0167 * 255)] * 3}\n\n\ndef get_trans_seg(img, I, rev=False):\n    if I >= 4 and not rev:\n        img = img.transpose(2,3)\n    if I % 4 == 0:\n        pass\n    elif I % 4 == 1:\n        img = img.flip(2)\n    elif I % 4 == 2:\n        img = img.flip(3)\n    elif I % 4 == 3:\n        img = img.flip(2).flip(3)\n    if I >= 4 and rev:\n        img = img.transpose(2,3)\n    return img\n\n\nclass CellSegmentator(object):\n    \"\"\"Uses pretrained DPN-Unet models to segment cells from images.\"\"\"\n\n    def __init__(\n        self,\n        nuclei_model=\"../input/hpa-seg/dpn_unet_nuclei_v1.pth\",\n        cell_model=\"../input/hpa-seg/dpn_unet_cell_3ch_v1.pth\",\n        scale_factor=1.0,\n        device=\"cuda\",\n#         padding=False,\n        multi_channel_model=True,\n    ):\n        \"\"\"Class for segmenting nuclei and whole cells from confocal microscopy images.\n        It takes lists of images and returns the raw output from the\n        specified segmentation model. Models can be automatically\n        downloaded if they are not already available on the system.\n        When working with images from the Huan Protein Cell atlas, the\n        outputs from this class' methods are well combined with the\n        label functions in the utils module.\n        Note that for cell segmentation, there are two possible models\n        available. One that works with 2 channeled images and one that\n        takes 3 channels.\n        Keyword arguments:\n        nuclei_model -- A loaded torch nuclei segmentation model or the\n                        path to a file which contains such a model.\n                        If the argument is a path that points to a non-existant file,\n                        a pretrained nuclei_model is going to get downloaded to the\n                        specified path (default: './nuclei_model.pth').\n        cell_model -- A loaded torch cell segmentation model or the\n                      path to a file which contains such a model.\n                      The cell_model argument can be None if only nuclei\n                      are to be segmented (default: './cell_model.pth').\n        scale_factor -- How much to scale images before they are fed to\n                        segmentation models. Segmentations will be scaled back\n                        up by 1/scale_factor to match the original image\n                        (default: 0.25).\n        device -- The device on which to run the models.\n                  This should either be 'cpu' or 'cuda' or pointed cuda\n                  device like 'cuda:0' (default: 'cuda').\n        padding -- Whether to add padding to the images before feeding the\n                   images to the network. (default: False).\n        multi_channel_model -- Control whether to use the 3-channel cell model or not.\n                               If True, use the 3-channel model, otherwise use the\n                               2-channel version (default: True).\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        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                raise\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).eval()\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                raise\n\n            cell_model = torch.load(cell_model, map_location=torch.device(self.device))\n        self.cell_model = cell_model.to(self.device).eval()\n        self.scale_factor = scale_factor\n#         self.padding = padding\n\n    def _image_conversion(self, images):\n        \"\"\"Convert/Format images to RGB image arrays list for cell predictions.\n        Intended for internal use only.\n        Keyword arguments:\n        images -- list of lists of image paths/arrays. It should following the\n                 pattern if with er channel input,\n                 [\n                     [microtubule_path0/image_array0, microtubule_path1/image_array1, ...],\n                     [er_path0/image_array0, er_path1/image_array1, ...],\n                     [nuclei_path0/image_array0, nuclei_path1/image_array1, ...]\n                 ]\n                 or if without er input,\n                 [\n                     [microtubule_path0/image_array0, microtubule_path1/image_array1, ...],\n                     None,\n                     [nuclei_path0/image_array0, nuclei_path1/image_array1, ...]\n                 ]\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 pred_nuclei(self, images):\n        \"\"\"Predict the nuclei segmentation.\n        Keyword arguments:\n        images -- A list of image arrays or a list of paths to images.\n                  If as a list of image arrays, the images could be 2d images\n                  of nuclei data array only, or must have the nuclei data in\n                  the blue channel; If as a list of file paths, the images\n                  could be RGB image files or gray scale nuclei image file\n                  paths.\n        Returns:\n        predictions -- A list of predictions of nuclei segmentation for each nuclei image.\n        \"\"\"\n\n        def _preprocess(image):\n            self.target_shape = image.shape\n            if len(image.shape) == 2:\n                image = np.dstack((image, image, image))\n            image = transform.rescale(image, self.scale_factor, multichannel=True)\n            nuc_image = np.dstack((image[..., 2], image[..., 2], image[..., 2]))\n            nuc_image = nuc_image.transpose([2, 0, 1])\n            return nuc_image\n\n        def _segment_helper(imgs):\n            with torch.no_grad():\n                mean = torch.as_tensor(NORMALIZE[\"mean\"], device=self.device)\n                std = torch.as_tensor(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 = torch.stack([get_trans_seg(self.nuclei_model(get_trans_seg(imgs, I)), I, True).softmax(1) for I in range(1)], 0).mean(0)\n#                 imgs = self.nuclei_model(imgs).softmax(1)\n#                 imgs = F.softmax(imgs, dim=1)\n                return imgs\n\n        preprocessed_imgs = list(map(_preprocess, images))\n        bs = 24\n        predictions = []\n        for i in range(0, len(preprocessed_imgs), bs):\n            start = i\n            end = min(len(preprocessed_imgs), i+bs)\n            x = preprocessed_imgs[start:end]\n            pred = _segment_helper(x).cpu().numpy()\n            predictions.append(pred)\n        predictions = list(np.concatenate(predictions, axis=0))\n        predictions = map(util.img_as_ubyte, predictions)\n        predictions = list(map(self._restore_scaling_padding, predictions))\n        return predictions\n\n    def _restore_scaling_padding(self, n_prediction):\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 not self.scale_factor == 1:\n            n_prediction[..., 0] = 0\n            n_prediction = cv2.resize(\n                n_prediction,\n                (self.target_shape[0], self.target_shape[1]),\n                interpolation=cv2.INTER_AREA,\n            )\n        return n_prediction\n\n    def pred_cells(self, images, precombined=False):\n        \"\"\"Predict the cell segmentation for a list of images.\n        Keyword arguments:\n        images -- list of lists of image paths/arrays. It should following the\n                  pattern if with er channel input,\n                  [\n                      [microtubule_path0/image_array0, microtubule_path1/image_array1, ...],\n                      [er_path0/image_array0, er_path1/image_array1, ...],\n                      [nuclei_path0/image_array0, nuclei_path1/image_array1, ...]\n                  ]\n                  or if without er input,\n                  [\n                      [microtubule_path0/image_array0, microtubule_path1/image_array1, ...],\n                      None,\n                      [nuclei_path0/image_array0, nuclei_path1/image_array1, ...]\n                  ]\n                  The ER channel is required when multichannel is True\n                  and required to be None when multichannel is False.\n                  The images needs to be of the same size.\n        precombined -- If precombined is True, the list of images is instead supposed to be\n                       a list of RGB numpy arrays (default: False).\n        Returns:\n        predictions -- a list of predictions of cell segmentations.\n        \"\"\"\n\n        def _preprocess(image):\n            self.target_shape = image.shape\n            if not len(image.shape) == 3:\n                raise ValueError(\"image should has 3 channels\")\n            cell_image = transform.rescale(image, self.scale_factor, multichannel=True)\n            cell_image = cell_image.transpose([2, 0, 1])\n            return cell_image\n\n        def _segment_helper(imgs):\n            with torch.no_grad():\n                mean = torch.as_tensor(NORMALIZE[\"mean\"], device=self.device)\n                std = torch.as_tensor(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 = torch.stack([get_trans_seg(self.cell_model(get_trans_seg(imgs, I)), I, True).softmax(1) for I in range(seg_TTA)], 0).mean(0)\n#                 imgs = F.softmax(imgs, dim=1)\n                return imgs\n\n        if not precombined:\n            images = self._image_conversion(images)\n        preprocessed_imgs = list(map(_preprocess, images))\n        bs = 24\n        predictions = []\n        for i in range(0, len(preprocessed_imgs), bs):\n            start = i\n            end = min(len(preprocessed_imgs), i+bs)\n            x = preprocessed_imgs[start:end]\n            pred = _segment_helper(x).cpu().numpy()\n            predictions.append(pred)\n        ###\n#         return predictions\n        ###\n        predictions = list(np.concatenate(predictions, axis=0))\n        predictions = map(self._restore_scaling_padding, predictions)\n        predictions = list(map(util.img_as_ubyte, predictions))\n\n        return predictions\n\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\n\ndef label_cell(nuclei_pred, cell_pred):\n    \"\"\"Label the cells and the nuclei.\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    Returns:\n    A tuple containing:\n    nuclei-label -- A nuclei mask data array.\n    cell-label  -- A cell mask data array.\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    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.copy(mask_img)\n        m = seeds * border_img  # * dt\n        img_copy[m <= threshold + threshold_adjustment] = 0\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[mask_img <= threshold] = 0\n        mask_img[mask_img > threshold] = 1\n        mask_img = mask_img.astype(np.bool)\n        mask_img = remove_small_holes(mask_img, 63)\n        mask_img = remove_small_objects(mask_img, 1).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=small_th,\n    )\n\n    # for hpa_image, to remove the small pseduo nuclei\n    nuclei_label = remove_small_objects(nuclei_label, 157)\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    distance = np.clip(cell_pred[..., 2], 255 * threshold_value, cell_pred[..., 2])\n    cell_label = segmentation.watershed(-distance, nuclei_label, mask=sk)\n    cell_label = remove_small_objects(cell_label, 344).astype(np.uint8)\n    selem = disk(2)\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    cell_label = remove_small_objects(cell_label, 344)\n    cell_label = measure.label(cell_label)\n    cell_label = np.asarray(cell_label, dtype=np.uint16)\n    nuclei_label = np.multiply(cell_label > 0, nuclei_label) > 0\n    nuclei_label = measure.label(nuclei_label)\n    nuclei_label = remove_small_objects(nuclei_label, 157)\n    nuclei_label = np.multiply(cell_label, nuclei_label > 0)\n\n    return nuclei_label, cell_label","metadata":{"execution":{"iopub.status.busy":"2021-06-17T03:54:58.973359Z","iopub.execute_input":"2021-06-17T03:54:58.973721Z","iopub.status.idle":"2021-06-17T03:54:59.028306Z","shell.execute_reply.started":"2021-06-17T03:54:58.973688Z","shell.execute_reply":"2021-06-17T03:54:59.027357Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"## the cell seg model\ncellsegmentor = CellSegmentator()","metadata":{"execution":{"iopub.status.busy":"2021-06-17T03:55:32.503023Z","iopub.execute_input":"2021-06-17T03:55:32.503364Z","iopub.status.idle":"2021-06-17T03:55:53.457117Z","shell.execute_reply.started":"2021-06-17T03:55:32.503331Z","shell.execute_reply":"2021-06-17T03:55:53.456179Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class HPADatasetSeg(Dataset):\n    def __init__(self, df, root='../input/hpa-single-cell-image-classification/test/'):\n        self.df = df.reset_index(drop=True)\n        self.root = root\n\n    def __len__(self):\n        return len(self.df)\n        \n    def __getitem__(self, index):\n\n        row = self.df.loc[index]\n        r = os.path.join(self.root, f'{row.ID}_red.png')\n        y = os.path.join(self.root, f'{row.ID}_yellow.png')\n        b = os.path.join(self.root, f'{row.ID}_blue.png')\n        r = cv2.imread(r, 0)\n        y = cv2.imread(y, 0)\n        b = cv2.imread(b, 0)\n        target_shape = r.shape\n        gray_image = cv2.resize(b, (seg_size, seg_size))\n        rgb_image = cv2.resize(np.stack((r, y, b), axis=2), (seg_size, seg_size))\n\n        return gray_image, rgb_image, target_shape, row.ID\n    \n    \ndef collate_fn(batch):\n    gray = []\n    rgb_image = []\n    target_shape = []\n    IDs = []\n    for data_point in batch:\n        gray.append(data_point[0])\n        rgb_image.append(data_point[1])\n        target_shape.append(data_point[2])\n        IDs.append(data_point[3])\n    return gray, rgb_image, target_shape, IDs\n\n\ndataset_seg = HPADatasetSeg(df_sub)\nloader_seg = DataLoader(dataset_seg, batch_size=seg_bs, num_workers=2, collate_fn=collate_fn)","metadata":{"execution":{"iopub.status.busy":"2021-06-17T03:56:29.104352Z","iopub.execute_input":"2021-06-17T03:56:29.104774Z","iopub.status.idle":"2021-06-17T03:56:29.11578Z","shell.execute_reply.started":"2021-06-17T03:56:29.104736Z","shell.execute_reply":"2021-06-17T03:56:29.114825Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# cell_segmentations = cellsegmentor.pred_cells(rgb, precombined=True)\n# plt.imshow(cell_segmentations[0][0].transpose(1,2,0))","metadata":{"execution":{"iopub.status.busy":"2021-06-17T03:56:29.770149Z","iopub.execute_input":"2021-06-17T03:56:29.770495Z","iopub.status.idle":"2021-06-17T03:56:29.774244Z","shell.execute_reply.started":"2021-06-17T03:56:29.770461Z","shell.execute_reply":"2021-06-17T03:56:29.77323Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for gray, rgb, target_shapes, IDs in tqdm(loader_seg):\n    nuc_segmentations = cellsegmentor.pred_nuclei(gray)\n    cell_segmentations = cellsegmentor.pred_cells(rgb, precombined=True)\n    for data_id, target_shape, nuc_seg, cell_seg in zip(IDs, target_shapes, nuc_segmentations, cell_segmentations):\n        nuc, cell = label_cell(nuc_seg, cell_seg)\n        np.savez_compressed(f'./{mask_dir}/{data_id}', cell.astype(np.uint8))\nprint('---- finish mask write ----')","metadata":{"execution":{"iopub.status.busy":"2021-06-17T03:56:30.120034Z","iopub.execute_input":"2021-06-17T03:56:30.120331Z","iopub.status.idle":"2021-06-17T03:56:51.454565Z","shell.execute_reply.started":"2021-06-17T03:56:30.120301Z","shell.execute_reply":"2021-06-17T03:56:51.451828Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"del cellsegmentor\ngc.collect()\ntorch.cuda.empty_cache()","metadata":{"execution":{"iopub.status.busy":"2021-06-17T03:56:51.457566Z","iopub.execute_input":"2021-06-17T03:56:51.457849Z","iopub.status.idle":"2021-06-17T03:56:51.744669Z","shell.execute_reply.started":"2021-06-17T03:56:51.457821Z","shell.execute_reply":"2021-06-17T03:56:51.743537Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!nvidia-smi","metadata":{"execution":{"iopub.status.busy":"2021-06-17T03:56:51.748047Z","iopub.execute_input":"2021-06-17T03:56:51.748378Z","iopub.status.idle":"2021-06-17T03:56:52.450096Z","shell.execute_reply.started":"2021-06-17T03:56:51.74835Z","shell.execute_reply":"2021-06-17T03:56:52.449185Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def encode_binary_mask(mask: np.ndarray) -> t.Text:\n    \"\"\"Converts a binary mask into OID challenge encoding ascii text.\"\"\"\n\n    # check input mask --\n    if mask.dtype != np.bool:\n        raise ValueError(\n            \"encode_binary_mask expects a binary mask, received dtype == %s\" %\n            mask.dtype)\n\n    mask = np.squeeze(mask)\n    if len(mask.shape) != 2:\n        raise ValueError(\n            \"encode_binary_mask expects a 2d mask, received shape == %s\" %\n            mask.shape)\n\n    # convert input mask to expected COCO API input --\n    mask_to_encode = mask.reshape(mask.shape[0], mask.shape[1], 1)\n    mask_to_encode = mask_to_encode.astype(np.uint8)\n    mask_to_encode = np.asfortranarray(mask_to_encode)\n\n    # RLE encode mask --\n    encoded_mask = coco_mask.encode(mask_to_encode)[0][\"counts\"]\n\n    # compress and base64 encoding --\n    binary_str = zlib.compress(encoded_mask, zlib.Z_BEST_COMPRESSION)\n    base64_str = base64.b64encode(binary_str)\n    return base64_str.decode('ascii')\n\ndef binary_mask_to_ascii(mask, mask_val=1):\n    \"\"\"Converts a binary mask into OID challenge encoding ascii text.\"\"\"\n    mask = np.where(mask==mask_val, 1, 0).astype(np.bool)\n    \n    # check input mask --\n    if mask.dtype != np.bool:\n        raise ValueError(f\"encode_binary_mask expects a binary mask, received dtype == {mask.dtype}\")\n\n    mask = np.squeeze(mask)\n    if len(mask.shape) != 2:\n        raise ValueError(f\"encode_binary_mask expects a 2d mask, received shape == {mask.shape}\")\n\n    # convert input mask to expected COCO API input --\n    mask_to_encode = mask.reshape(mask.shape[0], mask.shape[1], 1)\n    mask_to_encode = mask_to_encode.astype(np.uint8)\n    mask_to_encode = np.asfortranarray(mask_to_encode)\n\n    # RLE encode mask --\n    encoded_mask = coco_mask.encode(mask_to_encode)[0][\"counts\"]\n\n    # compress and base64 encoding --\n    binary_str = zlib.compress(encoded_mask, zlib.Z_BEST_COMPRESSION)\n    base64_str = base64.b64encode(binary_str)\n    return base64_str.decode()","metadata":{"execution":{"iopub.status.busy":"2021-06-17T03:56:52.452119Z","iopub.execute_input":"2021-06-17T03:56:52.452521Z","iopub.status.idle":"2021-06-17T03:56:52.464121Z","shell.execute_reply.started":"2021-06-17T03:56:52.452478Z","shell.execute_reply":"2021-06-17T03:56:52.463105Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def read_img(image_id, color, train_or_test='test', image_size=None):\n    filename = f'../input/hpa-single-cell-image-classification/{train_or_test}/{image_id}_{color}.png'\n    img = cv2.imread(filename, 0)\n    return img\n\nclass HPADatasetTest(Dataset):\n    def __init__(self, image_ids, mode='test'):\n        self.image_ids = image_ids\n        self.mode = mode\n        \n    def __len__(self):\n        return len(self.image_ids)\n        \n    def __getitem__(self, index):\n        \n        try:\n            image_id = self.image_ids[index]\n            red = read_img(image_id, \"red\", self.mode, 0)\n            green = read_img(image_id, \"green\", self.mode, 0)\n            blue = read_img(image_id, \"blue\", self.mode, 0)\n            yellow = read_img(image_id, \"yellow\", self.mode, 0)\n            image = np.stack([blue, green, red, yellow], axis=-1)\n\n            image_512 = cv2.resize(image, (512, 512)).transpose(2,0,1).astype(np.float32)\n            image_768 = cv2.resize(image, (768, 768)).transpose(2,0,1).astype(np.float32)\n            cell_mask = np.load(f'{mask_dir}/{image_id}.npz')['arr_0']\n            ### for debug\n#             cell_mask = np.zeros(cell_mask.shape).astype(int) if random.random() < 0.5 else cell_mask\n            ###\n            cell_mask = cv2.resize(cell_mask, (image.shape[0], image.shape[1]), interpolation=cv2.INTER_NEAREST)\n\n            encs = ''\n            masked_images = []\n            cells_128 = []\n            cells_256 = []\n            center_cells = []\n            for cell_id in range(1, np.max(cell_mask)+1):\n                bmask = (cell_mask == cell_id).astype(np.uint8)\n                enc = encode_binary_mask(bmask==1)\n                x, y, w, h = cv2.boundingRect(bmask)\n\n                max_l = max(w, h)\n                cx = x + w // 2\n                cy = y + h // 2\n                x1 = max(0, cx - max_l // 2)\n                x1 = min(x1, image.shape[1] - max_l)\n                y1 = max(0, cy - max_l // 2)\n                y1 = min(y1, image.shape[0] - max_l)\n                tmp = image.copy()\n                tmp[bmask==0] = 0\n\n                cropped_cell_orig = tmp[y1:y1+max_l, x1:x1+max_l]\n                cropped_cell_128 = cv2.resize(cropped_cell_orig, (128, 128))\n                cells_128.append(cropped_cell_128)\n                cropped_cell_256 = cv2.resize(cropped_cell_orig, (256, 256))\n                cells_256.append(cropped_cell_256)\n                masked = cv2.resize(tmp, (image_size, image_size))\n                masked_images.append(masked)\n\n                ### daishu part\n                cropped_cell = cv2.resize(tmp[y:y+h, x:x+w], \n                                            (int(w / image.shape[1] * 768),\n                                             int(h / image.shape[0] * 768))\n                                         )\n                final_size = 512\n                new_h, new_w, _ = cropped_cell.shape\n                new_h = final_size if cropped_cell.shape[0] > final_size else new_h\n                new_w = final_size if cropped_cell.shape[1] > final_size else new_w\n                cropped_cell = cv2.resize(cropped_cell, (new_w, new_h))\n\n                center_cell = np.zeros((final_size, final_size, 4))\n                center = final_size // 2\n                h_start = max(0,center-cropped_cell.shape[0]//2)\n                h_end = min(final_size,h_start+cropped_cell.shape[0])\n                w_start = max(0,center-cropped_cell.shape[1]//2)\n                w_end = min(final_size,w_start+cropped_cell.shape[1])\n\n                center_cell[h_start:h_end,w_start:w_end,:] = cropped_cell\n                center_cells.append(center_cell)\n                ###\n\n                if encs == '':\n                    encs += enc\n                else:\n                    encs = encs + ' ' + enc\n\n            if len(masked_images) > 0:\n                masked_images = np.stack(masked_images).transpose(0, 3, 1, 2).astype(np.float32)\n                cells_128 = np.stack(cells_128).transpose(0, 3, 1, 2).astype(np.float32)\n                cells_256 = np.stack(cells_256).transpose(0, 3, 1, 2).astype(np.float32)\n                center_cells = np.stack(center_cells).transpose(0, 3, 1, 2).astype(np.float32)\n            else:\n                masked_images = np.zeros((4, 4, image_size, image_size))\n                cells_128 = np.zeros((4, 4, 128, 128))\n                cells_256 = np.zeros((4, 4, 256, 256))\n\n            for ch in range(4):\n                image_512[ch] /= orig_mean[ch]\n                image_768[ch] /= orig_mean[ch]\n                masked_images[:, ch] /= orig_mean[ch]\n                cells_128[:, ch] /= orig_mean[ch]\n                cells_256[:, ch] /= orig_mean[ch]\n                center_cells[:, ch] /= orig_mean[ch]\n\n        except:\n            image_id = ''\n            encs = ''\n            image_512 = np.zeros((4, 512, 512))\n            image_768 = np.zeros((4, 768, 768))\n            masked_images = np.zeros((5, 4, image_size, image_size))\n            cells_128 = np.zeros((5, 4, 128, 128))\n            cells_256 = np.zeros((5, 4, 256, 256))\n            center_cells = np.zeros((5, 4, 512, 512))\n\n        return image_id, encs, {\n            '512': torch.tensor(image_512),\n            '768': torch.tensor(image_768),\n            'masked': torch.tensor(masked_images),\n            'cells_128': torch.tensor(cells_128),\n            'cells_256': torch.tensor(cells_256),\n            'center_cells': torch.tensor(center_cells)\n        }\n","metadata":{"execution":{"iopub.status.busy":"2021-06-17T03:56:52.465719Z","iopub.execute_input":"2021-06-17T03:56:52.466067Z","iopub.status.idle":"2021-06-17T03:56:52.498048Z","shell.execute_reply.started":"2021-06-17T03:56:52.466034Z","shell.execute_reply":"2021-06-17T03:56:52.497157Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"dataset = HPADatasetTest(df_sub.ID.values, mode='test')\ndataloader = DataLoader(dataset, batch_size=1, num_workers=2)","metadata":{"execution":{"iopub.status.busy":"2021-06-17T03:56:52.499259Z","iopub.execute_input":"2021-06-17T03:56:52.499659Z","iopub.status.idle":"2021-06-17T03:56:52.511621Z","shell.execute_reply.started":"2021-06-17T03:56:52.499623Z","shell.execute_reply":"2021-06-17T03:56:52.510803Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if df_sub.shape[0] <= 559:\n    from pylab import rcParams\n    rcParams['figure.figsize'] = 20,15\n\n    f, axarr = plt.subplots(4,5)\n    ID, enc, imgs = dataset[0]\n    print(imgs['masked'].shape)\n    for p in range(5):\n        axarr[0, p].imshow(imgs['512'][:3].transpose(0, 1).transpose(1,2))\n        axarr[1, p].imshow(imgs['masked'][p, :3].transpose(0, 1).transpose(1,2))\n        axarr[2, p].imshow(imgs['cells_256'][p, :3].transpose(0, 1).transpose(1,2))\n        axarr[3, p].imshow(imgs['center_cells'][p, [2,1,0]].transpose(0, 1).transpose(1,2))","metadata":{"execution":{"iopub.status.busy":"2021-06-17T03:56:52.51337Z","iopub.execute_input":"2021-06-17T03:56:52.513683Z","iopub.status.idle":"2021-06-17T03:56:57.809801Z","shell.execute_reply.started":"2021-06-17T03:56:52.513649Z","shell.execute_reply":"2021-06-17T03:56:57.808992Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model","metadata":{}},{"cell_type":"code","source":"class enetv2(nn.Module):\n    def __init__(self, enet_type, out_dim=num_classes):\n        super(enetv2, self).__init__()\n        self.enet = timm.create_model(enet_type, False)\n        if ('efficientnet' in enet_type) or ('mixnet' in enet_type):\n            self.enet.conv_stem.weight = nn.Parameter(self.enet.conv_stem.weight.repeat(1,n_ch//3+1,1,1)[:, :n_ch])\n            self.myfc = nn.Linear(self.enet.classifier.in_features, out_dim)\n            self.enet.classifier = nn.Identity()\n        elif ('resnet' in enet_type or 'resnest' in enet_type) and 'vit' not in enet_type:\n            self.enet.conv1[0].weight = nn.Parameter(self.enet.conv1[0].weight.repeat(1,n_ch//3+1,1,1)[:, :n_ch])\n            self.myfc = nn.Linear(self.enet.fc.in_features, out_dim)\n            self.enet.fc = nn.Identity()\n        elif 'rexnet' in enet_type or 'regnety' in enet_type or 'nf_regnet' in enet_type:\n            self.enet.stem.conv.weight = nn.Parameter(self.enet.stem.conv.weight.repeat(1,n_ch//3+1,1,1)[:, :n_ch])\n            self.myfc = nn.Linear(self.enet.head.fc.in_features, out_dim)\n            self.enet.head.fc = nn.Identity()\n        elif 'resnext' in enet_type:\n            self.enet.conv1.weight = nn.Parameter(self.enet.conv1.weight.repeat(1,n_ch//3+1,1,1)[:, :n_ch])\n            self.myfc = nn.Linear(self.enet.fc.in_features, out_dim)\n            self.enet.fc = nn.Identity()\n        elif 'hrnet_w32' in enet_type:\n            self.enet.conv1.weight = nn.Parameter(self.enet.conv1.weight.repeat(1,n_ch//3+1,1,1)[:, :n_ch])\n            self.myfc = nn.Linear(self.enet.classifier.in_features, out_dim)\n            self.enet.classifier = nn.Identity()\n        elif 'densenet' in enet_type:\n            self.enet.features.conv0.weight = nn.Parameter(self.enet.features.conv0.weight.repeat(1,n_ch//3+1,1,1)[:, :n_ch])\n            self.myfc = nn.Linear(self.enet.classifier.in_features, out_dim)\n            self.enet.classifier = nn.Identity()\n        elif 'ese_vovnet39b' in enet_type or 'xception41' in enet_type:\n            self.enet.stem[0].conv.weight = nn.Parameter(self.enet.stem[0].conv.weight.repeat(1,n_ch//3+1,1,1)[:, :n_ch])\n            self.myfc = nn.Linear(self.enet.head.fc.in_features, out_dim)\n            self.enet.head.fc = nn.Identity()\n        elif 'dpn' in enet_type:\n            self.enet.features.conv1_1.conv.weight = nn.Parameter(self.enet.features.conv1_1.conv.weight.repeat(1,n_ch//3+1,1,1)[:, :n_ch])\n            self.myfc = nn.Linear(self.enet.classifier.in_channels, out_dim)\n            self.enet.classifier = nn.Identity()\n        elif 'inception' in enet_type:\n            self.enet.features[0].conv.weight = nn.Parameter(self.enet.features[0].conv.weight.repeat(1,n_ch//3+1,1,1)[:, :n_ch])\n            self.myfc = nn.Linear(self.enet.last_linear.in_features, out_dim)\n            self.enet.last_linear = nn.Identity()\n        elif 'vit_base_resnet50' in enet_type:\n            self.enet.patch_embed.backbone.stem.conv.weight = nn.Parameter(self.enet.patch_embed.backbone.stem.conv.weight.repeat(1,n_ch//3+1,1,1)[:, :n_ch])\n            self.myfc = nn.Linear(self.enet.head.in_features, out_dim)\n            self.enet.head = nn.Identity()\n        else:\n            raise\n    \n    def forward(self, x):\n        x = self.enet(x)\n        h = self.myfc(x)\n        return h\n    ","metadata":{"execution":{"iopub.status.busy":"2021-06-17T03:56:57.812105Z","iopub.execute_input":"2021-06-17T03:56:57.812495Z","iopub.status.idle":"2021-06-17T03:56:57.836154Z","shell.execute_reply.started":"2021-06-17T03:56:57.812458Z","shell.execute_reply":"2021-06-17T03:56:57.835325Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!ls ../input/bo-hpa-models-3d256","metadata":{"execution":{"iopub.status.busy":"2021-06-17T03:56:57.838081Z","iopub.execute_input":"2021-06-17T03:56:57.838573Z","iopub.status.idle":"2021-06-17T03:56:58.537466Z","shell.execute_reply.started":"2021-06-17T03:56:57.838536Z","shell.execute_reply":"2021-06-17T03:56:58.536378Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"kernel_types = {\n    'resnet50d_512_multilabel_8flips_ss22rot45_co2_lr1e4_bs32_focal_ext_15epo': {\n        'model_class': 'enetv2',\n        'folds': [1],\n        'enet_type': 'resnet50d',\n        'input_type': ['512', 'masked'],\n    },\n    'rex150_512_multilabel_8flips_ss22rot45_co7_lr3e4_bs32_ext_cellpseudo2full_15epo': {\n        'model_class': 'enetv2',\n        'folds': [0],\n        'enet_type': 'rexnet_150',\n        'input_type': ['512', 'masked'],\n    },\n    'densenet121_512_multilabel_8flips_ss22rot45_co2_lr1e4_bs32_focal_ext_15epo': {\n        'model_class': 'enetv2',\n        'folds': [2],\n        'enet_type': 'densenet121',\n        'input_type': ['512', 'masked'],\n    },\n    'b0_512_multilabel_8flips_ss22rot45_co7_lr1e4_bs32_focal_ext_15epo': {\n        'model_class': 'enetv2',\n        'folds': [3],\n        'enet_type': 'tf_efficientnet_b0_ns',\n        'input_type': ['512', 'masked'],\n    },\n    'resnet101d_512_multilabel_8flips_ss22rot45_co7_lr1e4_bs32_focal_ext_15epo': {\n        'model_class': 'enetv2',\n        'folds': [4],\n        'enet_type': 'resnet101d',\n        'input_type': ['512', 'masked'],\n    },\n    'dpn68b_512_multilabel_8flips_ss22rot45_co7_lr1e4_bs32_focal_ext_15epo': {\n        'model_class': 'enetv2',\n        'folds': [0],\n        'enet_type': 'dpn68b',\n        'input_type': ['512', 'masked'],\n    },\n    'densenet169_512_multilabel_8flips_ss22rot45_co2_lr1e4_bs32_focal_ext_15epo': {\n        'model_class': 'enetv2',\n        'folds': [1],\n        'enet_type': 'densenet169',\n        'input_type': ['512', 'masked'],\n    },\n    ### 2.5d\n    'b0_3d128_multilabel_lw41_8flips_ss22rot45_lr1e4_bs32cell16_ext_2019_15epo': {\n        'model_class': 'enetv2',\n        'folds': [2],\n        'enet_type': 'tf_efficientnet_b0_ns',\n        'input_type': ['cells_128'],\n    },\n    'resnet50d_3d128_multilabel_lw41_8flips_ss22rot45_lr1e4_bs32cell16_ext_15epo': {\n        'model_class': 'enetv2',\n        'folds': [0],\n        'enet_type': 'resnet50d',\n        'input_type': ['cells_128'],\n    },\n    'mixnet_m_3d128_multilabel_lw41_8flips_ss22rot45_lr1e4_bs32cell16_ext_15epo': {\n        'model_class': 'enetv2',\n        'folds': [0],\n        'enet_type': 'mixnet_m',\n        'input_type': ['cells_128'],\n    },\n    'densenet121_3d128_multilabel_lw41_8flips_ss22rot45_lr1e4_bs32cell16_ext_2019_15epo': {\n        'model_class': 'enetv2',\n        'folds': [3],\n        'enet_type': 'densenet121',\n        'input_type': ['cells_128'],\n    },\n    ### 2.5d 256\n    'b0_3d256_multilabel_lw41_8flips_ss22rot45_lr1e4_bs32cell16_ext_2019_15epo': {\n        'model_class': 'enetv2',\n        'folds': [0],\n        'enet_type': 'tf_efficientnet_b0_ns',\n        'input_type': ['cells_256'],\n    },\n    'b1_3d256_multilabel_lw41_8flips_ss22rot45_lr1e4_bs32cell16_ext_2019_15epo': {\n        'model_class': 'enetv2',\n        'folds': [3],\n        'enet_type': 'tf_efficientnet_b1_ns',\n        'input_type': ['cells_256'],\n    },\n    'densenet121_3d256_multilabel_lw41_8flips_ss22rot45_lr1e4_bs32cell16_ext_2019_15epo': {\n        'model_class': 'enetv2',\n        'folds': [2],\n        'enet_type': 'densenet121',\n        'input_type': ['cells_256'],\n    },\n    'dpn68b_3d256_multilabel_lw41_8flips_ss22rot45_lr1e4_bs32cell16_ext_2019_15epo': {\n        'model_class': 'enetv2',\n        'folds': [4],\n        'enet_type': 'dpn68b',\n        'input_type': ['cells_256'],\n    },\n    'mixnet_m_3d256_multilabel_lw41_8flips_ss22rot45_lr1e4_bs32cell16_ext_2019_15epo': {\n        'model_class': 'enetv2',\n        'folds': [2],\n        'enet_type': 'mixnet_m',\n        'input_type': ['cells_256'],\n    },\n    'resnet50d_3d256_multilabel_lw41_8flips_ss22rot45_lr1e4_bs32cell16_ext_2019_15epo': {\n        'model_class': 'enetv2',\n        'folds': [1],\n        'enet_type': 'resnet50d',\n        'input_type': ['cells_256'],\n    },\n}","metadata":{"execution":{"iopub.status.busy":"2021-06-17T03:56:58.539312Z","iopub.execute_input":"2021-06-17T03:56:58.539708Z","iopub.status.idle":"2021-06-17T03:56:58.553767Z","shell.execute_reply.started":"2021-06-17T03:56:58.539667Z","shell.execute_reply":"2021-06-17T03:56:58.552694Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def load_state_dict(model, model_file):\n    for folder in model_dirs:\n        model_path = os.path.join(folder, model_file)\n        if os.path.exists(model_path):\n            state_dict = torch.load(model_path)\n            state_dict = {k[7:] if k.startswith('module.') else k: state_dict[k] for k in state_dict.keys()}\n            model.load_state_dict(state_dict, strict=True)\n            model.eval()\n            return model\n    raise\n\nmodels = []\ninput_types = []\nfor key in kernel_types.keys():\n    for fold in kernel_types[key]['folds']:\n        model = eval(kernel_types[key]['model_class'])(\n            kernel_types[key]['enet_type'],\n        )\n        model = model.to(device)\n        model_file = f'{key}_final_fold{fold}.pth'\n        print(f'loading {model_file} ...')\n        model = load_state_dict(model, model_file)\n        models.append(model)\n\n        input_types.append(kernel_types[key]['input_type'])\n\nn_models = len(models)\nprint('done!')\nprint('model count:', n_models)","metadata":{"execution":{"iopub.status.busy":"2021-06-17T03:56:58.555466Z","iopub.execute_input":"2021-06-17T03:56:58.555907Z","iopub.status.idle":"2021-06-17T03:57:18.969475Z","shell.execute_reply.started":"2021-06-17T03:56:58.555872Z","shell.execute_reply":"2021-06-17T03:57:18.967897Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def load_model(model_name,path):\n    if model_name == 'densenet121':\n        state_dict = torch.load(path, torch.device('cuda') )\n        model = class_densenet121_dropout(num_classes=19,in_channels=4,pretrained_file=None)\n        model.cuda()\n        model.load_state_dict(state_dict)\n#         model = amp.initialize(model, opt_level=\"O1\")\n        model.eval()\n    return model","metadata":{"execution":{"iopub.status.busy":"2021-06-17T03:57:18.970807Z","iopub.execute_input":"2021-06-17T03:57:18.971152Z","iopub.status.idle":"2021-06-17T03:57:18.977438Z","shell.execute_reply.started":"2021-06-17T03:57:18.971114Z","shell.execute_reply":"2021-06-17T03:57:18.97621Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"folds = [0,1,2,3,4]\nmodel_dic = {'densenet121':'../input/hpafinal/output/run_nn_20210504_000509/'}","metadata":{"execution":{"iopub.status.busy":"2021-06-17T03:57:18.979218Z","iopub.execute_input":"2021-06-17T03:57:18.979699Z","iopub.status.idle":"2021-06-17T03:57:18.988733Z","shell.execute_reply.started":"2021-06-17T03:57:18.979636Z","shell.execute_reply":"2021-06-17T03:57:18.987785Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"rgby_models = []\nfor model_name in model_dic:\n    path = model_dic[model_name]\n    for fold in folds:\n        if os.path.exists(path+'fold%s.ckpt'%fold):\n            model = load_model(model_name,path+'fold%s.ckpt'%fold)\n            rgby_models.append(model)\nprint('daishu model count:', len(rgby_models))","metadata":{"execution":{"iopub.status.busy":"2021-06-17T03:57:18.992292Z","iopub.execute_input":"2021-06-17T03:57:18.99264Z","iopub.status.idle":"2021-06-17T03:57:25.332174Z","shell.execute_reply.started":"2021-06-17T03:57:18.99261Z","shell.execute_reply":"2021-06-17T03:57:25.330618Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_trans(img, I, mode='bgry'):\n    if mode == 'rgby':\n        img = img[:, [2,1,0,3]]\n    if I >= 4:\n        img = img.transpose(2,3)\n    if I % 4 == 0:\n        return img\n    elif I % 4 == 1:\n        return img.flip(2)\n    elif I % 4 == 2:\n        return img.flip(3)\n    elif I % 4 == 3:\n        return img.flip(2).flip(3)\n    \n\n# def get_trans_daishu(img, I, mode='bgry'):\n#     if mode == 'rgby':\n#         img = img[:, [2,1,0,3]]\n\n#     if I == 0:\n#         return img[:, :, :512, :512]\n#     if I == 1:\n#         return img[:, :, :512, 256:]\n#     if I == 2:\n#         return img[:, :, 256:, :512]\n#     if I == 3:\n#         return img[:, :, 256:, 256:]\n#     if I == 4:\n#         return img[:, :, 128:640, 128:640]\n#     raise\ndef get_trans_daishu(img, I, mode='bgry'):\n    if mode == 'rgby':\n        img = img[:, [2,1,0,3]]\n    \n    if I == 0:\n        img = img[:, :, 64:704, 64:704]\n    if I == 1:\n        img = img[:, :, :640, :640].flip(2)\n    if I == 2:\n        img = img[:, :, :640, 128:].flip(3)\n    if I == 3:\n        img = img[:, :, 128:, 128:].flip(2).flip(3)\n    if I == 4:\n        img = img[:, :, 128:, :640].transpose(2,3)\n    if I == 5:\n        img = img[:, :, 32:672, 96:736].transpose(2,3).flip(2)\n    if I == 6:\n        img = img[:, :, 96:736, 32:672].transpose(2,3).flip(3)\n    img = F.interpolate(img, size=[512, 512], mode=\"bilinear\")\n    return img","metadata":{"execution":{"iopub.status.busy":"2021-06-17T03:57:25.333501Z","iopub.execute_input":"2021-06-17T03:57:25.333844Z","iopub.status.idle":"2021-06-17T03:57:25.345517Z","shell.execute_reply.started":"2021-06-17T03:57:25.333808Z","shell.execute_reply":"2021-06-17T03:57:25.34448Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print('TTA', TTA)","metadata":{"execution":{"iopub.status.busy":"2021-06-17T03:57:25.346836Z","iopub.execute_input":"2021-06-17T03:57:25.347187Z","iopub.status.idle":"2021-06-17T03:57:25.360231Z","shell.execute_reply.started":"2021-06-17T03:57:25.347152Z","shell.execute_reply":"2021-06-17T03:57:25.359329Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"IDs = []\nencs = []\nPRED_FINAL = []\nlittle_bs = 16\nwith torch.no_grad():\n    for ID, enc, images in tqdm(dataloader):\n        try:\n            if len(enc[0]) > 0:\n                with amp.autocast():\n                    for k in images.keys():\n                        images[k] = images[k].cuda()\n                        if images[k].ndim == 5:\n                            images[k] = images[k].squeeze(0)\n\n                    preds = {\n                        'orig': [],\n                        'cells': [],\n                    }\n                    # orig 全图\n#                     for m, inp_types in zip(models, input_types):\n#                         for t in inp_types:\n#                             if t in ['512', '768']:\n#                                 for I in np.random.choice(range(8), TTA['orig'], replace=False):\n#                                     preds['orig'].append(m(get_trans(images[t], I)).sigmoid())\n\n                    # 遮盖 & 单cell\n                    for m, inp_types in zip(models, input_types):\n                        for t in inp_types:\n                            if t in ['masked', 'cells_128', 'cells_256']:\n                                for I in np.random.choice(range(8), TTA[t], replace=False):\n                                    this_pred = torch.cat([\n                                        m(get_trans(images[t][b:b+little_bs], I)).sigmoid() \\\n                                            for b in range(0, images[t].shape[0], little_bs)\n                                    ])\n                                    preds['cells'].append(this_pred)\n\n                    # daishu\n                    for m in rgby_models:\n                        # 全图\n#                         for I in np.random.choice(range(8), TTA['orig']-1, replace=False):\n#                             preds['orig'].append(m(get_trans_daishu(images['768'], I, 'rgby'))[1].sigmoid())\n#                         preds['orig'].append(m(images['512'][:, [2,1,0,3]])[1].sigmoid())\n                        # 单cell\n                        for I in np.random.choice(range(8), TTA['center_cells'], replace=False):\n                            this_pred = torch.cat([\n                                m(get_trans(images['center_cells'][b:b+little_bs], I, 'rgby'))[1].sigmoid() \\\n                                    for b in range(0, images['center_cells'].shape[0], little_bs)\n                            ])\n                            preds['cells'].append(this_pred)\n\n                    for k in preds.keys():\n                        if len(preds[k]) > 0:\n                            preds[k] = torch.stack(preds[k], 0).mean(0)\n                        else:\n                            preds[k] = 0\n\n                    pred_final = preds['cells']\n\n                    PRED_FINAL.append(pred_final.cpu())\n                    IDs += [ID[0]] * images['cells_128'].shape[0]\n                    encs += enc[0].split(' ')\n\n        except:\n            print('error')\n            pass\n\nPRED_FINAL = torch.cat(PRED_FINAL).float()","metadata":{"execution":{"iopub.status.busy":"2021-06-17T03:57:25.361652Z","iopub.execute_input":"2021-06-17T03:57:25.362063Z","iopub.status.idle":"2021-06-17T03:59:42.345115Z","shell.execute_reply.started":"2021-06-17T03:57:25.362025Z","shell.execute_reply":"2021-06-17T03:59:42.3439Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(PRED_FINAL.shape, PRED_FINAL.max(), PRED_FINAL.min(), PRED_FINAL.mean())","metadata":{"execution":{"iopub.status.busy":"2021-06-17T04:00:19.212795Z","iopub.execute_input":"2021-06-17T04:00:19.213146Z","iopub.status.idle":"2021-06-17T04:00:19.245922Z","shell.execute_reply.started":"2021-06-17T04:00:19.213111Z","shell.execute_reply":"2021-06-17T04:00:19.244909Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"PredictionString = []\nfor i in tqdm(range(PRED_FINAL.shape[0])):\n    enc = encs[i]\n    prob = PRED_FINAL[i]\n    sub_string = []\n    for cid, p in enumerate(prob):\n        sub_string.append(' '.join([str(cid), f'{p:.5f}', enc]))\n    sub_string = ' '.join(sub_string)\n    PredictionString.append(sub_string)","metadata":{"execution":{"iopub.status.busy":"2021-06-17T04:00:23.424883Z","iopub.execute_input":"2021-06-17T04:00:23.425213Z","iopub.status.idle":"2021-06-17T04:00:23.46522Z","shell.execute_reply.started":"2021-06-17T04:00:23.425175Z","shell.execute_reply":"2021-06-17T04:00:23.464205Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_pred = pd.DataFrame({\n    'ID': IDs,\n    'PredictionString': PredictionString\n})\ndf_pred = df_pred.groupby(['ID'])['PredictionString'].apply(lambda x: ' '.join(x)).reset_index()","metadata":{"execution":{"iopub.status.busy":"2021-06-17T04:00:23.900731Z","iopub.execute_input":"2021-06-17T04:00:23.901021Z","iopub.status.idle":"2021-06-17T04:00:23.918206Z","shell.execute_reply.started":"2021-06-17T04:00:23.900994Z","shell.execute_reply":"2021-06-17T04:00:23.917252Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_sub = df_sub[['ID', 'ImageWidth', 'ImageHeight']].merge(df_pred, on='ID', how=\"left\")\ndf_sub.fillna('', inplace=True)\ndf_sub.to_csv('submission.csv', index=False)","metadata":{"execution":{"iopub.status.busy":"2021-06-17T04:00:25.578211Z","iopub.execute_input":"2021-06-17T04:00:25.578563Z","iopub.status.idle":"2021-06-17T04:00:25.904798Z","shell.execute_reply.started":"2021-06-17T04:00:25.578529Z","shell.execute_reply":"2021-06-17T04:00:25.903973Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_sub.shape","metadata":{"execution":{"iopub.status.busy":"2021-06-17T04:00:42.126898Z","iopub.execute_input":"2021-06-17T04:00:42.127239Z","iopub.status.idle":"2021-06-17T04:00:42.135006Z","shell.execute_reply.started":"2021-06-17T04:00:42.127206Z","shell.execute_reply":"2021-06-17T04:00:42.133995Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!rm -rf {mask_dir}","metadata":{"execution":{"iopub.status.busy":"2021-06-17T04:00:43.675282Z","iopub.execute_input":"2021-06-17T04:00:43.675652Z","iopub.status.idle":"2021-06-17T04:00:44.491468Z","shell.execute_reply.started":"2021-06-17T04:00:43.67562Z","shell.execute_reply":"2021-06-17T04:00:44.490282Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}