{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":61446,"databundleVersionId":6962461,"sourceType":"competition"},{"sourceId":7103273,"sourceType":"datasetVersion","datasetId":4011732},{"sourceId":151009453,"sourceType":"kernelVersion"}],"dockerImageVersionId":30587,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!python -m pip install --no-index --find-links=/kaggle/input/pip-download-for-segmentation-models-pytorch segmentation-models-pytorch","metadata":{"execution":{"iopub.status.busy":"2023-12-02T04:04:37.929116Z","iopub.execute_input":"2023-12-02T04:04:37.929751Z","iopub.status.idle":"2023-12-02T04:04:49.962114Z","shell.execute_reply.started":"2023-12-02T04:04:37.929717Z","shell.execute_reply":"2023-12-02T04:04:49.960849Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from more_itertools import locate\nimport os\nfrom os.path import join as opj\n\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nimport cv2\nimport torch\nfrom PIL import Image\nfrom torch.utils.data import Dataset, DataLoader\nimport torch.nn as nn\nfrom typing import Optional, List\nfrom tqdm import tqdm\nimport rasterio\nfrom functools import lru_cache\nimport numpy as np\nimport pandas as pd\nimport segmentation_models_pytorch as smp\n\n\nimport warnings\nwarnings.filterwarnings(\"ignore\", category=rasterio.errors.NotGeoreferencedWarning)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-12-02T04:04:49.965242Z","iopub.execute_input":"2023-12-02T04:04:49.965627Z","iopub.status.idle":"2023-12-02T04:04:49.974893Z","shell.execute_reply.started":"2023-12-02T04:04:49.965591Z","shell.execute_reply":"2023-12-02T04:04:49.973989Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CFG:\n    seed          = 42\n    debug         = False # set debug=False for Full Training\n    exp_name      = 'unet_timm-resnest50d_1'\n    comment       = 'unet_timm-resnest50d-512x512'\n    model_name    = 'Unet'\n    backbone      = 'timm-resnest50d'\n    ckpt_path     = '/kaggle/input/blood-vessel-segmentation-weight/best_epoch_28_11_2023_09_10.bin'\n    valid_bs      = 32\n    img_size      = [800, 800]\n    num_classes   = 1\n    thresh        = 0.3\n    min_size      = 12\n    device        = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    multiple_gpu  = True\n    \n    data_transforms = {\n        \"valid\": A.Compose([\n            A.ToRGB(),\n            ToTensorV2()\n        ], p=1.0)\n    }","metadata":{"execution":{"iopub.status.busy":"2023-12-02T04:04:49.975937Z","iopub.execute_input":"2023-12-02T04:04:49.976248Z","iopub.status.idle":"2023-12-02T04:04:49.993851Z","shell.execute_reply.started":"2023-12-02T04:04:49.976223Z","shell.execute_reply":"2023-12-02T04:04:49.992984Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def rle_decode(mask_rle: str, img_shape: tuple = None) -> np.ndarray:\n    seq = mask_rle.split()\n    starts = np.array(list(map(int, seq[0::2])))\n    lengths = np.array(list(map(int, seq[1::2])))\n    assert len(starts) == len(lengths)\n    ends = starts + lengths\n    img = np.zeros((np.product(img_shape),), dtype=np.uint8)\n    for begin, end in zip(starts, ends):\n        img[begin:end] = 1\n    # https://stackoverflow.com/a/46574906/4521646\n    img.shape = img_shape\n    return img\n\ndef rle_encode(img):\n    '''\n    img: numpy array, 1 - mask, 0 - background\n    Returns run length as string formated\n    '''\n    pixels = img.flatten()\n    pixels = np.concatenate([[0], pixels, [0]])\n    runs = np.where(pixels[1:] != pixels[:-1])[0] + 1\n    runs[1::2] -= runs[::2]\n    rle = ' '.join(str(x) for x in runs)\n    if rle == '':\n        rle = '1 0'\n    return rle\n\n\ndef remove_small_objects(img, min_size):\n    # Find all connected components (labels)\n    num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(img, connectivity=8)\n\n    # Create a mask where small objects are removed\n    new_img = np.zeros_like(img)\n    for label in range(1, num_labels):\n        if stats[label, cv2.CC_STAT_AREA] >= min_size:\n            new_img[labels == label] = 1\n    return new_img\n\n\ndef build_model(backbone, num_classes):\n    model = smp.Unet(\n        encoder_name=backbone,      # choose encoder, e.g. mobilenet_v2 or efficientnet-b7\n        encoder_weights=None,     # use `imagenet` pre-trained weights for encoder initialization\n        in_channels=3,                  # model input channels (1 for gray-scale images, 3 for RGB, etc.)\n        classes=num_classes,        # model output channels (number of classes in your dataset)\n        activation='sigmoid',\n    )\n    return model\n\n\ndef load_model(backbone, num_classes, path):\n    model = build_model(backbone, num_classes)\n    model.load_state_dict(torch.load(path))\n    model.eval()\n    return model\n\n\n@lru_cache(maxsize=64)\ndef ropen(img_fpath):\n    return rasterio.open(img_fpath)","metadata":{"execution":{"iopub.status.busy":"2023-12-02T04:04:49.995131Z","iopub.execute_input":"2023-12-02T04:04:49.995416Z","iopub.status.idle":"2023-12-02T04:04:50.015926Z","shell.execute_reply.started":"2023-12-02T04:04:49.995392Z","shell.execute_reply":"2023-12-02T04:04:50.014864Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# In-Memory Tiling\n\nOne technique I've observed folks using is to resize the input image down to a constant scale. This introduces distortions and issues with the number of pixels-on-target disappearing (e.g., a 10 pixel target downsampled by 1/10 becomes a 1 pixel target, while a 5 pixel target becomes a 0 pixel target).\n\nRather than resize the imagery, I've adopted a technique that's common in geospatial processing- tiling. Tiling preserves the original target sizes and avoids distortions by simply applying a sliding window to the data. The window can be square or unsquare, avoiding aspect ratio issues. The window also has a stride associated with it, allowing the network multiple opportunities to \"detect\" a target if it misses. Critically, by tiling, this increases the percieved target scale as observed by the model (e.g., a 10 pixel target looks smaller in a 1000x1000 image than a 100x100 image). Finally, tiling the inflates the number of samples available from training and increases sample diversity","metadata":{}},{"cell_type":"code","source":"class TiledDataset(Dataset):\n\n    def __init__(\n            self,\n            path_img_dir: str,\n            transforms=None,\n            tile_size: Optional[List[int]] = None,\n            overlap_pct: float = 0.2,\n            cache_dir: str = None\n    ):\n        \"\"\"\n        Generates an in-memory tiled dataset using a sliding window approach (common in geospatial processing communities)\n\n        Args:\n\n            path_img_dir (str): Path to root directory containing data.\n            transforms (torchvision.transforms.Transform): Albumentation transforms.\n            tile_size (Optional[List[int]]): List describing the tile size. If None, tiling is not performed and the\n                full scene is utilized.\n            overlap_pct (float): Percentage of tile size to use as overlap between tiles. This effectively\n                controls the \"stride\" of the window as it's slid across the full scene.\n            cache_dir (str): Cache dir to read/write to. If None, no cache directory is utilized.\n        \"\"\"\n        self.path_img_dir = path_img_dir\n        self.transforms = transforms\n        self.samples = []\n        self.tile_size = np.array(tile_size) if tile_size is not None else None\n        self.overlap_pct = overlap_pct\n        self.cache_dir = cache_dir\n        path_images = []\n        for folder in os.listdir(self.path_img_dir):\n            # if folder != 'kidney_1_dense' and folder != 'kidney_2':\n            #     continue\n            # print(folder)\n            if not os.path.isdir(opj(self.path_img_dir, folder)):\n                continue\n            images_dir = opj(self.path_img_dir, folder, 'images')\n            for image in os.listdir(images_dir):\n                if not os.path.isfile(opj(images_dir, image)):\n                    continue\n                path_images.append(opj(images_dir, image))\n        # for _, row in self.data.iterrows():\n        #     p_img = os.path.join(self.path_img_dir, row[\"dataset\"], \"images\", f'{row[\"slice\"]}.tif')\n        for p_img in tqdm(path_images):\n            with rasterio.open(p_img) as reader:\n                width, height = reader.width, reader.height\n                img = reader.read()\n                px_max, px_min = img.max(), img.min()\n\n            if not os.path.isfile(p_img):\n                continue\n            self.samples.append((p_img, [], [], 1, [px_min, px_max], [width, height]))\n\n        if self.tile_size is not None:\n            empty = 0\n            nonempty = 0\n\n            empty_tiles = []\n            populated_tiles = []\n\n            for file_path, _, _, _, px_stats, img_dims in tqdm(self.samples, total=len(self.samples), desc='Generating tiles'):\n                width, height = img_dims\n\n                min_overlap = float(overlap_pct) * 0.01\n                max_stride = self.tile_size * (1.0 - min_overlap)\n                num_patches = np.ceil(np.array([height, width]) / max_stride).astype(np.int64)\n                # compute the cutoff points for the x-y dimensions\n                starts = [np.int64(np.linspace(0, width - self.tile_size[1], num_patches[1])),\n                          np.int64(np.linspace(0, height - self.tile_size[0], num_patches[0]))]\n                stops = [starts[0] + self.tile_size[0], starts[1] + self.tile_size[1]]\n                for y1, y2 in zip(starts[1], stops[1]):\n                    for x1, x2 in zip(starts[0], stops[0]):\n\n                        populated_tiles.append((file_path, [x1, y1, x2 - x1, y2 - y1], [height, width], 1, px_stats, img_dims))\n\n\n            pos_idxs_to_sample = range(len(populated_tiles))\n\n\n            self.samples = list(map(populated_tiles.__getitem__, pos_idxs_to_sample))\n\n    def __getitem__(self, idx: int) -> tuple:\n        \"\"\"\n        Returns a sample of data\n\n        \"\"\"\n        # Grab the sample from the sample list\n        img_fpath, bbox, original_img_size, target, px_stats, img_dims = self.samples[idx]\n\n\n        img = ropen(img_fpath).read(1, window=rasterio.windows.Window(*bbox) if len(bbox) > 0 else None)\n\n        if len(original_img_size) == 0:\n            original_img_size = img.shape\n\n        if img.ndim == 3:\n            img = np.mean(img, axis=2)\n\n        # If we read the window from the full scene, compress the dynamic range to UINT8\n        # TODO: Save full scene statistics and compress relative to those, rather than tile level statistics\n        img = (img - px_stats[0]) / (px_stats[1] - px_stats[0])\n        img *= 255.0\n        img = img.astype(np.uint8)\n\n        # Transform the image and mask using Albumentations\n        if self.transforms:\n            transformed = self.transforms(image=img)\n            img = transformed['image']\n\n        # The augmentations may scale the mask to the range 0-1, with 1's becoming ~0.0039 (1/255).\n        # Make sure we overwrite this in any case\n\n        # Return the sample\n        return img, img_fpath, bbox, original_img_size, target\n\n    def __len__(self) -> int:\n        return len(self.samples)","metadata":{"execution":{"iopub.status.busy":"2023-12-02T04:04:50.019239Z","iopub.execute_input":"2023-12-02T04:04:50.019589Z","iopub.status.idle":"2023-12-02T04:04:50.044318Z","shell.execute_reply.started":"2023-12-02T04:04:50.019563Z","shell.execute_reply":"2023-12-02T04:04:50.043254Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Construct a dataset and dataloader with some Augmentations\n\nThe in-memory tiling process isn't free- it takes some time at startup. The time required to tile the imagery scales inversely proportional with tile size (i.e., smaller tiles = longer wait time).","metadata":{}},{"cell_type":"code","source":"import albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\nDATASET_FOLDER = '/kaggle/input/blood-vessel-segmentation'\nIMG_PATH = DATASET_FOLDER + '/test'\n\n\ndataset = TiledDataset(\n                            path_img_dir=IMG_PATH,\n                            tile_size=CFG.img_size,\n                            overlap_pct=20,\n                            transforms=CFG.data_transforms['valid']\n)","metadata":{"execution":{"iopub.status.busy":"2023-12-02T04:04:50.045540Z","iopub.execute_input":"2023-12-02T04:04:50.045822Z","iopub.status.idle":"2023-12-02T04:04:50.118194Z","shell.execute_reply.started":"2023-12-02T04:04:50.045798Z","shell.execute_reply":"2023-12-02T04:04:50.117252Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_loader = DataLoader(dataset, batch_size=CFG.valid_bs, num_workers=0, shuffle=False, pin_memory=True)","metadata":{"execution":{"iopub.status.busy":"2023-12-02T04:04:50.119538Z","iopub.execute_input":"2023-12-02T04:04:50.120216Z","iopub.status.idle":"2023-12-02T04:04:50.125681Z","shell.execute_reply.started":"2023-12-02T04:04:50.120184Z","shell.execute_reply":"2023-12-02T04:04:50.124714Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = load_model(CFG.backbone, CFG.num_classes, CFG.ckpt_path)\nif CFG.multiple_gpu:\n    model = nn.DataParallel(model)\nmodel.to(CFG.device)","metadata":{"execution":{"iopub.status.busy":"2023-12-02T04:04:50.126698Z","iopub.execute_input":"2023-12-02T04:04:50.127093Z","iopub.status.idle":"2023-12-02T04:04:51.018748Z","shell.execute_reply.started":"2023-12-02T04:04:50.127066Z","shell.execute_reply":"2023-12-02T04:04:51.017812Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"count = 0\nposition = []\nimgs = []\nimg_fpaths = []\nbboxes = []\nshapes = []\npbar = tqdm(enumerate(test_loader), total=len(test_loader), desc='Inference')\n# for i, data in tqdm(enumerate(dataset)):\nfor step, (data) in pbar:\n\n    img, img_fpath, bbox, original_img_size, target = data\n    #print(img.shape)\n    img = img.to(CFG.device, dtype=torch.float)\n    with torch.no_grad():\n        preds = model(img)\n        preds = (preds>CFG.thresh).float()\n    img_fpaths.extend(img_fpath)\n    for i in range(CFG.valid_bs):\n        if i >= preds.shape[0]:\n            break\n        imgs.append(rle_encode((preds[i]).cpu().numpy().astype(np.uint8)))\n        \n        bboxes.append([int(bbox[0][i]), int(bbox[1][i]), int(bbox[2][i]), int(bbox[3][i])])\n        shapes.append([int(original_img_size[0][i]), int(original_img_size[1][i])])\n        position.append(CFG.valid_bs*step+i)\n    #position.append(i)\n    #img_fpaths.append(img_fpath)\n    #imgs.append(preds)","metadata":{"execution":{"iopub.status.busy":"2023-12-02T04:04:51.019799Z","iopub.execute_input":"2023-12-02T04:04:51.020064Z","iopub.status.idle":"2023-12-02T04:04:52.900465Z","shell.execute_reply.started":"2023-12-02T04:04:51.020041Z","shell.execute_reply":"2023-12-02T04:04:52.899533Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"del model, test_loader, dataset","metadata":{"execution":{"iopub.status.busy":"2023-12-02T04:04:52.901883Z","iopub.execute_input":"2023-12-02T04:04:52.902282Z","iopub.status.idle":"2023-12-02T04:04:52.907187Z","shell.execute_reply.started":"2023-12-02T04:04:52.902245Z","shell.execute_reply":"2023-12-02T04:04:52.906307Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"unique_fpath = list(set(img_fpaths))","metadata":{"execution":{"iopub.status.busy":"2023-12-02T04:04:52.908513Z","iopub.execute_input":"2023-12-02T04:04:52.908920Z","iopub.status.idle":"2023-12-02T04:04:52.918861Z","shell.execute_reply.started":"2023-12-02T04:04:52.908893Z","shell.execute_reply":"2023-12-02T04:04:52.918098Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# from more_itertools import locate\n# a_list = [1,2,3,4,1,2,1,2,3,4]\n#item_to_find=1\n# indices = locate(a_list, lambda x: x == item_to_find)\n# print(list(indices))","metadata":{"execution":{"iopub.status.busy":"2023-12-02T04:04:52.920148Z","iopub.execute_input":"2023-12-02T04:04:52.920433Z","iopub.status.idle":"2023-12-02T04:04:52.930890Z","shell.execute_reply.started":"2023-12-02T04:04:52.920408Z","shell.execute_reply":"2023-12-02T04:04:52.929942Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"position_index = dict()\nfor fpath in unique_fpath:\n    # position_index[fpath] = list(locate(img_fpaths, lambda x: x == fpath))\n    position_index[fpath] = [idx for idx, value in enumerate(img_fpaths) if value == fpath]\n    \n    ","metadata":{"execution":{"iopub.status.busy":"2023-12-02T04:04:52.932312Z","iopub.execute_input":"2023-12-02T04:04:52.932995Z","iopub.status.idle":"2023-12-02T04:04:52.942473Z","shell.execute_reply.started":"2023-12-02T04:04:52.932958Z","shell.execute_reply":"2023-12-02T04:04:52.941510Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def check_size(bbox, canvas):\n    x, y, w, h = bbox\n    if x + w > canvas.shape[1]:\n        bbox[2] = canvas.shape[1] - x\n    if y + h > canvas.shape[0]:\n        bbox[3] = canvas.shape[0] - y\n    return bbox","metadata":{"execution":{"iopub.status.busy":"2023-12-02T04:04:52.948130Z","iopub.execute_input":"2023-12-02T04:04:52.948569Z","iopub.status.idle":"2023-12-02T04:04:52.954683Z","shell.execute_reply.started":"2023-12-02T04:04:52.948538Z","shell.execute_reply":"2023-12-02T04:04:52.953705Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"rles = []\nfor key in position_index:\n    shape = shapes[position_index[key][0]].copy()\n    canvas = np.zeros(shape).astype(np.uint8)\n    for index in position_index[key]:\n        img =  rle_decode(imgs[index], tuple(CFG.img_size)).astype(np.uint8)\n        # img = img.cpu().numpy().astype(np.uint8)\n        bbox_ = bboxes[index]\n        bbox = check_size(bbox_, canvas)\n        x, y, w, h = bbox\n        canvas[y:y+h, x:x+w] = np.logical_or(img[0:h, 0:w], canvas[y:y+h, x:x+w])\n    # for index in position_index[key]\n    rle = rle_encode(remove_small_objects(canvas, CFG.min_size))\n    rles.append(rle)","metadata":{"execution":{"iopub.status.busy":"2023-12-02T04:04:52.955839Z","iopub.execute_input":"2023-12-02T04:04:52.956186Z","iopub.status.idle":"2023-12-02T04:04:53.026695Z","shell.execute_reply.started":"2023-12-02T04:04:52.956161Z","shell.execute_reply":"2023-12-02T04:04:53.025854Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"ids = []\nfor p_img in tqdm(position_index):\n    path_ = p_img.split(os.path.sep)\n    # parse the submission ID\n    dataset = path_[-3]\n    slice_id, _ = os.path.splitext(path_[-1])\n    ids.append(f\"{dataset}_{slice_id}\")","metadata":{"execution":{"iopub.status.busy":"2023-12-02T04:04:53.027910Z","iopub.execute_input":"2023-12-02T04:04:53.028239Z","iopub.status.idle":"2023-12-02T04:04:53.037243Z","shell.execute_reply.started":"2023-12-02T04:04:53.028210Z","shell.execute_reply":"2023-12-02T04:04:53.036295Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission = pd.DataFrame.from_dict({\n    \"id\": ids,\n    \"rle\": rles\n})\nsubmission.to_csv(\"submission.csv\", index=False)","metadata":{"execution":{"iopub.status.busy":"2023-12-02T04:04:53.038762Z","iopub.execute_input":"2023-12-02T04:04:53.039106Z","iopub.status.idle":"2023-12-02T04:04:53.046624Z","shell.execute_reply.started":"2023-12-02T04:04:53.039078Z","shell.execute_reply":"2023-12-02T04:04:53.045813Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission","metadata":{"execution":{"iopub.status.busy":"2023-12-02T04:04:53.047791Z","iopub.execute_input":"2023-12-02T04:04:53.048176Z","iopub.status.idle":"2023-12-02T04:04:53.063169Z","shell.execute_reply.started":"2023-12-02T04:04:53.048149Z","shell.execute_reply":"2023-12-02T04:04:53.062015Z"},"trusted":true},"execution_count":null,"outputs":[]}]}