{"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":7503518,"sourceType":"datasetVersion","datasetId":4369572},{"sourceId":7511414,"sourceType":"datasetVersion","datasetId":4374867},{"sourceId":7556753,"sourceType":"datasetVersion","datasetId":4400741},{"sourceId":7560445,"sourceType":"datasetVersion","datasetId":4402426}],"dockerImageVersionId":30635,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# SenNet + HOA - Exploring 3D Human Vasculature Segmentation\n\nThis code is a simple implementation of the CNN model train and infer for the competition.\nAlthough **we cannot expect accuracy**, we hope that it will be helpful for those who are participating in a competition for the **first time on how to implement train and infer**.","metadata":{}},{"cell_type":"markdown","source":"Load the required libraries.","metadata":{}},{"cell_type":"code","source":"import os\nfrom glob import glob\nimport numpy as np\nimport pandas as pd\nimport cv2\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import DataLoader, Dataset, TensorDataset","metadata":{"execution":{"iopub.status.busy":"2024-02-05T10:18:40.621987Z","iopub.execute_input":"2024-02-05T10:18:40.622480Z","iopub.status.idle":"2024-02-05T10:18:40.627734Z","shell.execute_reply.started":"2024-02-05T10:18:40.622446Z","shell.execute_reply":"2024-02-05T10:18:40.626861Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"torch.cuda.is_available()","metadata":{"execution":{"iopub.status.busy":"2024-02-05T10:18:36.318977Z","iopub.execute_input":"2024-02-05T10:18:36.319376Z","iopub.status.idle":"2024-02-05T10:18:36.326352Z","shell.execute_reply.started":"2024-02-05T10:18:36.319345Z","shell.execute_reply":"2024-02-05T10:18:36.325402Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Set constants.\n\nNote that the shape of the image matrix is `(B, C, H, W)`.","metadata":{}},{"cell_type":"code","source":"def device_name():\n    if torch.cuda.is_available():\n        return 'cuda'\n    else:\n        return 'cpu'\n\n\n# ref.: https://www.kaggle.com/stainsby/fast-tested-rle\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    return ' '.join(str(x) for x in runs)\n\n\ndef to_csv(names, images):\n    rles = []\n\n    for image in images:\n        rle = rle_encode(image)\n        rles.append(rle)\n\n    pd.DataFrame({'id': names, 'rle': rles}).to_csv(\n        'submission.csv', index=False)\n\n\ndef to_img(names, images):\n    for i in range(images.shape[0]):\n        image = images[i][0]\n        image = image * 255\n        cv2.imwrite(names[i] + \".tif\", image)\n\n\nclass CheckPoint():\n    def __init__(self, path):\n        self.path = path\n\n    def save(self, model, optimizer, epoch, loss):\n        check_point = {\n            'model_state_dict':     model.state_dict(),\n            'optimizer_state_dict': optimizer.state_dict(),\n            'epoch':                epoch,\n            'loss':                 loss\n        }\n\n        torch.save(check_point, \"./model.pt\")\n\n    def load(self, model, optimizer, epoch, loss):\n        check_point = None\n\n        if os.path.exists(self.path):\n            print('loading... check point')\n            check_point = torch.load(self.path)\n\n        if check_point:\n            model.load_state_dict(check_point['model_state_dict'])\n            optimizer.load_state_dict(check_point['optimizer_state_dict'])\n            epoch = check_point['epoch']\n            loss = check_point['loss']\n\n        return model, optimizer, epoch, loss\n\n\nclass Model(nn.Module):\n    def __init__(self, config):\n        super().__init__()\n\n        self.config = config\n\n        self.conv = nn.Conv2d(1, 1, kernel_size=3, padding=1)\n        self.maxpool = nn.MaxPool2d(kernel_size=2)\n        self.relu = nn.ReLU()\n        # self.sigmoid = nn.Sigmoid()\n\n    def forward(self, x):\n        x = self.conv(x)\n        x = self.relu(x)\n        x = self.maxpool(x)\n        x = torch.nn.functional.interpolate(x, size=(self.config.H, self.config.W), mode='bilinear', align_corners=False)\n        # x = self.sigmoid(x)\n\n        return x\n\n\nclass Config():\n    def __init__(self, check_point_path=\"./model.pt\"):\n        self.check_point_path = check_point_path\n\n        self.B = 1\n        self.C = 1\n        self.H = 1303\n        self.W = 912\n\n        # self.RESIZE_IMAGE_H = 143\n        # self.RESIZE_IMAGE_W = 100\n\n        self.THRESHOLD = 0.5\n\n        self.TRAIN_BATCH_SIZE = 1\n        self.INFER_BATCH_SIZE = 1\n        self.VALID_BATCH_SIZE = 1\n\n        self.POS_WEIGHT = 1000\n\n        self.EPOCH_SIZE = 10\n\n        self.LR = 0.001\n\n        self.TRAIN_IMAGES_RANGE_S = 1000\n        self.TRAIN_IMAGES_RANGE_E = 1100\n        self.VALID_IMAGES_RANGE_S = 1000\n        self.VALID_IMAGES_RANGE_E = 1010\n        self.INFER_IMAGES_RANGE_S = None\n        self.INFER_IMAGES_RANGE_E = None\n\n\ndef load_image(path):\n    # print(\"loading...\", path)\n    image = cv2.imread(path, cv2.IMREAD_GRAYSCALE)\n\n    # image = cv2.resize(image, (self.config.RESIZE_IMAGE_W, self.config.RESIZE_IMAGE_H))\n    # add channel dimension\n    image = np.expand_dims(image, axis=0)\n\n    image = normalize_image(image)\n\n    return image\n\n\ndef load_label(path):\n    # print(\"loading...\", path)\n    image = cv2.imread(path, cv2.IMREAD_UNCHANGED)\n\n    # image = cv2.resize(image, (self.config.RESIZE_IMAGE_W, self.config.RESIZE_IMAGE_H))\n    # add channel dimension\n    image = np.expand_dims(image, axis=0)\n\n    image = normalize_image(image)\n\n    return image\n\n\ndef normalize_image(image):\n    # Min-max normalization\n    min_val = image.min()\n    max_val = image.max()\n    image = (image - min_val) / (max_val - min_val)\n    return image\n\n\nclass TrainDataset(Dataset):\n    def __init__(self, names, image_paths, label_paths):\n        self.names = names\n        self.image_paths = image_paths\n        self.label_paths = label_paths\n\n    def __len__(self):\n        return len(self.names)\n\n    def __getitem__(self, idx):\n        name = self.names[idx]\n        image_path = self.image_paths[idx]\n        label_path = self.label_paths[idx]\n\n        image = load_image(image_path)\n        label = load_label(label_path)\n\n        return name, image, label\n\n\nclass InferDataset(Dataset):\n    def __init__(self, names, image_paths):\n        self.names = names\n        self.image_paths = image_paths\n\n    def __len__(self):\n        return len(self.names)\n\n    def __getitem__(self, idx):\n        name = self.names[idx]\n        image_path = self.image_paths[idx]\n\n        image = load_image(image_path)\n\n        return name, image\n\n\nclass Processer():\n    def __init__(self, config=None) -> None:\n        if config == None:\n            self.config = Config()\n        else:\n            self.config = config\n\n        self.check_point = CheckPoint(self.config.check_point_path)\n\n    def image_path_to_name(self, image_path):\n        return image_path.split(\"/\")[-3] + \"_\" + image_path.split(\"/\")[-1].split(\".\")[0]\n\n    def gen_infer_data(self, source_image_path, range_s=None, range_e=None):\n        source_image_paths = glob(source_image_path)\n\n        source_image_paths.sort()\n\n        if range_s == None and range_e == None:\n            image_paths = source_image_paths\n        else:\n            image_paths = source_image_paths[range_s:range_e]\n\n        names = [self.image_path_to_name(image_path) for image_path in image_paths]\n\n        infer_dataset = InferDataset(names=names, image_paths=image_paths)\n\n        return infer_dataset\n\n    def gen_train_data(self, source_image_path, source_label_path, range_s=None, range_e=None):\n        source_image_paths = glob(source_image_path)\n        source_label_paths = glob(source_label_path)\n\n        source_image_paths.sort()\n        source_label_paths.sort()\n\n        if range_s == None and range_e == None:\n            image_paths = source_image_paths\n            label_paths = source_label_paths\n        else:\n            image_paths = source_image_paths[range_s:range_e]\n            label_paths = source_label_paths[range_s:range_e]\n\n        names = [self.image_path_to_name(image_path) for image_path in image_paths]\n\n        train_dataset = TrainDataset(names=names, image_paths=image_paths, label_paths=label_paths)\n\n        return train_dataset\n\n    def train(self, image_paths, label_paths):\n        model = Model(config=self.config)\n        model.to(device_name())\n        model.train()\n\n        # criterion = nn.BCELoss()\n        pos_weight = torch.full((self.config.H, self.config.W), self.config.POS_WEIGHT).to(device_name())\n        criterion = nn.BCEWithLogitsLoss(pos_weight=pos_weight)\n        optimizer = optim.Adam(model.parameters(), lr=self.config.LR)\n\n        model, optimizer, epoch, min_loss = self.check_point.load(model, optimizer, 0, float('inf'))\n\n        epoch_size = self.config.EPOCH_SIZE\n\n        if 0 < epoch:\n            epoch_size -= epoch + 1\n\n        dataset = self.gen_train_data(image_paths, label_paths, self.config.TRAIN_IMAGES_RANGE_S, self.config.TRAIN_IMAGES_RANGE_E)\n\n        data_loader = DataLoader(dataset, batch_size=self.config.TRAIN_BATCH_SIZE)\n\n        for epoch in range(epoch_size):\n            print(\"epoch: \", epoch)\n\n            for (names, images, labels) in data_loader:\n                # the tensor shape is B, C, H, W\n                images_tensor = torch.Tensor(np.array(images)).to(device_name())\n                labels_tensor = torch.Tensor(np.array(labels)).to(device_name())\n\n                optimizer.zero_grad()\n\n                x = model(images_tensor)\n                y = labels_tensor\n\n                loss = criterion(x, y)\n\n                loss.backward()\n\n                optimizer.step()\n\n            print(loss.item())\n\n            if loss.item() < min_loss:\n                min_loss = loss\n                self.check_point.save(model, optimizer, epoch, loss.item())\n\n    def infer(self, image_paths):\n        model = Model(config=self.config)\n        model.to(device_name())\n        model.eval()\n\n        optimizer = optim.Adam(model.parameters(), lr=self.config.LR)\n\n        model, optimizer, epoch, min_loss = self.check_point.load(model, optimizer, 0, float('inf'))\n\n        dataset = self.gen_infer_data(image_paths)\n\n        data_loader = DataLoader(dataset, batch_size=self.config.INFER_BATCH_SIZE)\n\n        with torch.no_grad():\n            data_records = np.ndarray((0, self.config.C, self.config.H, self.config.W))\n            name_records = []\n\n            for (names, images) in data_loader:\n                images_tensor = torch.Tensor(np.array(images)).to(device_name())\n\n                optimizer.zero_grad()\n\n                x = model(images_tensor)\n\n                x = x.to('cpu')\n                r = x.numpy()\n\n                r = (r > self.config.THRESHOLD).astype(np.uint8)\n\n                data_records = np.vstack((data_records, r))\n                name_records.extend(list(names))\n\n            name_records = np.array(name_records)\n\n            to_csv(name_records, data_records)\n\n    def valid(self, image_paths, label_paths):\n        model = Model(config=self.config)\n        model.to(device_name())\n        model.train()\n\n        # criterion = nn.BCELoss()\n        pos_weight = torch.full((self.config.H, self.config.W), self.config.POS_WEIGHT).to(device_name())\n        criterion = nn.BCEWithLogitsLoss(pos_weight=pos_weight)\n        optimizer = optim.Adam(model.parameters(), lr=self.config.LR)\n\n        model, optimizer, epoch, min_loss = self.check_point.load(model, optimizer, 0, float('inf'))\n\n        epoch_size = self.config.EPOCH_SIZE\n\n        if 0 < epoch:\n            epoch_size -= epoch + 1\n\n        dataset = self.gen_train_data(image_paths, label_paths, self.config.VALID_IMAGES_RANGE_S, self.config.VALID_IMAGES_RANGE_E)\n\n        data_loader = DataLoader(dataset, batch_size=self.config.VALID_BATCH_SIZE)\n\n        with torch.no_grad():\n            data_records = np.ndarray((0, self.config.C, self.config.H, self.config.W))\n            name_records = []\n\n            for (names, images, labels) in data_loader:\n                # the tensor shape is B, C, H, W\n                images_tensor = torch.Tensor(np.array(images)).to(device_name())\n                labels_tensor = torch.Tensor(np.array(labels)).to(device_name())\n\n                optimizer.zero_grad()\n\n                x = model(images_tensor)\n                y = labels_tensor\n\n                loss = criterion(x, y)\n\n                x = x.to('cpu')\n                r = x.numpy()\n\n                r = (r > self.config.THRESHOLD).astype(np.uint8)\n\n                data_records = np.vstack((data_records, r))\n                name_records.extend(list(names))\n\n            print(loss.item())\n\n            name_records = np.array(name_records)\n\n            to_csv(name_records, data_records)\n            to_img(name_records, data_records)\n","metadata":{"execution":{"iopub.status.busy":"2024-02-05T10:22:24.707897Z","iopub.execute_input":"2024-02-05T10:22:24.708250Z","iopub.status.idle":"2024-02-05T10:22:24.761969Z","shell.execute_reply.started":"2024-02-05T10:22:24.708222Z","shell.execute_reply":"2024-02-05T10:22:24.760617Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"You can call `train` or `infer`.\n\nThe model contents when running train are saved in `model.pt`.","metadata":{}},{"cell_type":"markdown","source":"Running `train`\n\n**Please delete it when submitting.**","metadata":{}},{"cell_type":"code","source":"!rm -rf ./model.pt\n!rm -rf ./submission.csv\n!rm -rf ./*.tif","metadata":{"execution":{"iopub.status.busy":"2024-02-05T10:23:06.792053Z","iopub.execute_input":"2024-02-05T10:23:06.792864Z","iopub.status.idle":"2024-02-05T10:23:09.734419Z","shell.execute_reply.started":"2024-02-05T10:23:06.792822Z","shell.execute_reply":"2024-02-05T10:23:09.732772Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# image_paths = \"/kaggle/input/blood-vessel-segmentation/train/*/images/*.tif\"\n# label_paths = \"/kaggle/input/blood-vessel-segmentation/train/*/labels/*.tif\"\n\n# processer = Processer()\n# processer.train(image_paths, label_paths)","metadata":{"execution":{"iopub.status.busy":"2024-02-05T10:17:44.720762Z","iopub.execute_input":"2024-02-05T10:17:44.721236Z","iopub.status.idle":"2024-02-05T10:17:44.725741Z","shell.execute_reply.started":"2024-02-05T10:17:44.721201Z","shell.execute_reply":"2024-02-05T10:17:44.724714Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Running `valid`\n\n**Please delete it when submitting.**","metadata":{}},{"cell_type":"code","source":"# image_paths = \"/kaggle/input/blood-vessel-segmentation/train/*/images/*.tif\"\n# label_paths = \"/kaggle/input/blood-vessel-segmentation/train/*/labels/*.tif\"\n\n# processer = Processer()\n# processer.valid(image_paths, label_paths)","metadata":{"execution":{"iopub.status.busy":"2024-02-05T10:17:37.044753Z","iopub.execute_input":"2024-02-05T10:17:37.045094Z","iopub.status.idle":"2024-02-05T10:17:37.049690Z","shell.execute_reply.started":"2024-02-05T10:17:37.045067Z","shell.execute_reply":"2024-02-05T10:17:37.048537Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Running `infer`\n\nPlease **change** `check_point_path`","metadata":{}},{"cell_type":"code","source":"image_paths = \"/kaggle/input/blood-vessel-segmentation/test/*/images/*.tif\"\ncheck_point_path = \"/kaggle/input/jtakahashi64-202402051900/model.pt\"\nconfig = Config(check_point_path=check_point_path)\n\nprocesser = Processer(config=config)\nprocesser.infer(image_paths)","metadata":{"execution":{"iopub.status.busy":"2024-02-05T10:22:30.875084Z","iopub.execute_input":"2024-02-05T10:22:30.875512Z","iopub.status.idle":"2024-02-05T10:22:31.071158Z","shell.execute_reply.started":"2024-02-05T10:22:30.875476Z","shell.execute_reply":"2024-02-05T10:22:31.070135Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# image_paths = \"/kaggle/input/blood-vessel-segmentation/test/*/images/*.tif\"\n\n# processer = Processer()\n# processer.infer(image_paths)","metadata":{"execution":{"iopub.status.busy":"2024-02-05T10:17:04.403233Z","iopub.execute_input":"2024-02-05T10:17:04.403983Z","iopub.status.idle":"2024-02-05T10:17:04.407771Z","shell.execute_reply.started":"2024-02-05T10:17:04.403948Z","shell.execute_reply":"2024-02-05T10:17:04.406837Z"},"trusted":true},"execution_count":null,"outputs":[]}]}