{"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"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install torchsummary torchview","metadata":{"execution":{"iopub.status.busy":"2023-09-11T17:33:03.077011Z","iopub.execute_input":"2023-09-11T17:33:03.077508Z","iopub.status.idle":"2023-09-11T17:33:18.094452Z","shell.execute_reply.started":"2023-09-11T17:33:03.077472Z","shell.execute_reply":"2023-09-11T17:33:18.093317Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import re\nimport os\nimport numpy as np\nimport pandas as pd\nimport torch\nimport time\nimport shutil\n\nfrom torch import nn\nfrom imageio.v2 import imread\nfrom glob import glob\nfrom torch.nn import functional as F\nfrom torch.utils.data import Dataset, random_split, DataLoader\nfrom io import BytesIO\nfrom zipfile import ZipFile\nfrom urllib.request import urlopen\nfrom matplotlib import pyplot as plt\nfrom torchsummary import summary\nfrom torchview import draw_graph\nfrom tqdm.notebook import tqdm","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-09-11T17:33:50.015090Z","iopub.execute_input":"2023-09-11T17:33:50.015496Z","iopub.status.idle":"2023-09-11T17:33:53.628494Z","shell.execute_reply.started":"2023-09-11T17:33:50.015453Z","shell.execute_reply":"2023-09-11T17:33:53.627368Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n#This URL has a copy of the dataset, make sure you go to Kaggle though and make an account and follow the license rules. \ndata_url_zip = \"https://github.com/kamalkraj/DATA-SCIENCE-BOWL-2018/blob/master/data/stage1_train.zip?raw=true\"\n\nif not os.path.isdir('./data'):\n    os.mkdir(\"./data\")\n\n\n#Download this dataset if we have not already done so\nif not os.path.isdir('./data/stage1_train'):\n    resp = urlopen(data_url_zip)\n    os.mkdir(\"./data/stage1_train\")\n    zipfile = ZipFile(BytesIO(resp.read()))\n    zipfile.extractall(path = './data/stage1_train')","metadata":{"execution":{"iopub.status.busy":"2023-09-11T17:34:04.837164Z","iopub.execute_input":"2023-09-11T17:34:04.837698Z","iopub.status.idle":"2023-09-11T17:34:19.516388Z","shell.execute_reply.started":"2023-09-11T17:34:04.837669Z","shell.execute_reply":"2023-09-11T17:34:19.515381Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## DataLoader","metadata":{}},{"cell_type":"code","source":"#Grab all the paths to the images we just unziped \n\npaths = glob(\"./data/stage1_train/*\")","metadata":{"execution":{"iopub.status.busy":"2023-09-11T17:34:27.049502Z","iopub.execute_input":"2023-09-11T17:34:27.049862Z","iopub.status.idle":"2023-09-11T17:34:27.059236Z","shell.execute_reply.started":"2023-09-11T17:34:27.049834Z","shell.execute_reply":"2023-09-11T17:34:27.058104Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_imgs = {}\n\nfor i, path in enumerate(paths):\n    img = glob(path + \"/images/*\")[0]\n    mask_list = glob(path + \"/masks/*\")\n    \n    train_imgs[i] = {\"img\": img,\n                     \"mask\": mask_list}","metadata":{"execution":{"iopub.status.busy":"2023-09-11T17:34:32.015037Z","iopub.execute_input":"2023-09-11T17:34:32.015399Z","iopub.status.idle":"2023-09-11T17:34:32.182341Z","shell.execute_reply.started":"2023-09-11T17:34:32.015369Z","shell.execute_reply":"2023-09-11T17:34:32.181344Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class DSB2018(Dataset):\n    \"\"\"Dataset class for the 2018 Data science bowl\"\"\"\n\n    def __init__(self, paths):\n        \"\"\"paths: a lis of paths to every image folder in the dataset\"\"\"\n        self.paths = paths\n\n    def __len__(self):\n        return len(self.paths)\n\n    def __getitem__(self, idx):\n        img_path = self.paths[idx]['img']        \n        mask_imgs = self.paths[idx]['mask']\n        \n        img = imread(img_path)[:,:,0:3]\n        img = np.moveaxis(img, -1, 0)\n        img = img/255.0\n\n        masks = [imread(f)/255.0 for f in mask_imgs]\n        final_mask = np.zeros(masks[0].shape)\n        \n        for m in masks:\n            final_mask = np.logical_or(final_mask, m)\n            \n        final_mask = final_mask.astype(np.float32)    \n        del masks\n\n        img, final_mask = torch.tensor(img), torch.tensor(final_mask).unsqueeze(0) \n        img = F.interpolate(img.unsqueeze(0), (128, 128))\n        final_mask = F.interpolate(final_mask.unsqueeze(0), (128, 128))\n\n        return img.type(torch.FloatTensor)[0], final_mask.type(torch.FloatTensor)[0]","metadata":{"execution":{"iopub.status.busy":"2023-09-11T17:34:36.697735Z","iopub.execute_input":"2023-09-11T17:34:36.698091Z","iopub.status.idle":"2023-09-11T17:34:36.710183Z","shell.execute_reply.started":"2023-09-11T17:34:36.698063Z","shell.execute_reply":"2023-09-11T17:34:36.709111Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"dsb_data = DSB2018(train_imgs)","metadata":{"execution":{"iopub.status.busy":"2023-09-11T17:34:41.568173Z","iopub.execute_input":"2023-09-11T17:34:41.568623Z","iopub.status.idle":"2023-09-11T17:34:41.573378Z","shell.execute_reply.started":"2023-09-11T17:34:41.568588Z","shell.execute_reply":"2023-09-11T17:34:41.572500Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.figure(figsize=(12, 6))\nplt.subplot(1, 2, 1)\nplt.imshow(dsb_data[0][0].permute(1, 2, 0).numpy())\nplt.subplot(1, 2, 2)\nplt.imshow(dsb_data[0][1].numpy()[0, :], cmap='gray')","metadata":{"execution":{"iopub.status.busy":"2023-09-11T17:34:45.758139Z","iopub.execute_input":"2023-09-11T17:34:45.758846Z","iopub.status.idle":"2023-09-11T17:34:46.525850Z","shell.execute_reply.started":"2023-09-11T17:34:45.758812Z","shell.execute_reply":"2023-09-11T17:34:46.524915Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.figure(figsize=(12, 6))\nplt.subplot(1, 2, 1)\nplt.imshow(dsb_data[1][0].permute(1, 2, 0).numpy())\nplt.subplot(1, 2, 2)\nplt.imshow(dsb_data[1][1].numpy()[0, :], cmap='gray')","metadata":{"execution":{"iopub.status.busy":"2023-09-11T17:34:49.605156Z","iopub.execute_input":"2023-09-11T17:34:49.605547Z","iopub.status.idle":"2023-09-11T17:34:50.304291Z","shell.execute_reply.started":"2023-09-11T17:34:49.605516Z","shell.execute_reply":"2023-09-11T17:34:50.303387Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_split, test_split = random_split(dsb_data, [500, len(dsb_data) - 500])\ntrain_seg_loader = DataLoader(train_split, batch_size=16, shuffle=True)\ntest_seg_loader = DataLoader(test_split, batch_size=16)","metadata":{"execution":{"iopub.status.busy":"2023-09-11T17:38:54.179668Z","iopub.execute_input":"2023-09-11T17:38:54.180028Z","iopub.status.idle":"2023-09-11T17:38:54.190595Z","shell.execute_reply.started":"2023-09-11T17:38:54.179998Z","shell.execute_reply":"2023-09-11T17:38:54.189674Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Utils","metadata":{}},{"cell_type":"code","source":"def move_to(obj, device):\n    if isinstance(obj, list):\n        return [move_to(x, device) for x in obj]\n    \n    elif isinstance(obj, tuple):\n        return tuple(move_to(obj, device))\n    \n    elif isinstance(obj, set):\n        return set(move_to(obj, device))\n    \n    elif isinstance(obj, dict):\n        to_return = dict()\n        \n        for key, value in obj.items():\n            to_return[move_to(key, device)] = move_to(value, device)\n        \n        return to_return\n    \n    elif hasattr(obj, 'to'):\n        return obj.to(device)\n    \n    else:\n        return obj","metadata":{"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2023-09-11T17:38:59.558013Z","iopub.execute_input":"2023-09-11T17:38:59.558390Z","iopub.status.idle":"2023-09-11T17:38:59.566022Z","shell.execute_reply.started":"2023-09-11T17:38:59.558360Z","shell.execute_reply":"2023-09-11T17:38:59.565099Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def run_epoch(model, optimizer, data_loader,\n              loss_func, device, results,\n              score_funcs, prefix=\"\", desc=None, keep=False, epoch=None):\n    running_loss = []\n    y_true = []\n    y_pred = []\n    start = time.time()\n\n    for inputs, labels in tqdm(data_loader, desc=desc, leave=keep):\n        # Move the batch to the device we are using.\n        inputs = move_to(inputs, device)\n        labels = move_to(labels, device)\n\n        y_hat = model(inputs)  # this just computed f_Θ(x(i))\n        # Compute loss.\n        loss = loss_func(y_hat, labels)\n\n        if model.training:\n            loss.backward()\n            optimizer.step()\n            optimizer.zero_grad()\n\n        # Now we are just grabbing some information we would like to have\n        running_loss.append(loss.item())\n\n        if len(score_funcs) > 0 and isinstance(labels, torch.Tensor):\n            # moving labels & predictions back to CPU for computing / storing predictions\n            labels = labels.detach().cpu().numpy()\n            y_hat = y_hat.detach().cpu().numpy()\n            # add to predictions so far\n            y_true.extend(labels.tolist())\n            y_pred.extend(y_hat.tolist())\n\n    # end training epoch\n    end = time.time()\n\n    y_pred = np.asarray(y_pred)\n    if len(y_pred.shape) == 2 and y_pred.shape[1] > 1:  # We have a classification problem, convert to labels\n        y_pred = np.argmax(y_pred, axis=1)\n    # Else, we assume we are working on a regression problem\n\n    results[prefix + \" loss\"].append(np.mean(running_loss))\n    result_str = [f\"{prefix} loss: {np.mean(running_loss)}\"]\n\n    for name, score_func in score_funcs.items():\n        try:\n            score = score_func(y_true, y_pred)\n\n            results[prefix + \" \" + name].append(score)\n            result_str.append(f\"{prefix} {name}: {score}\")\n\n            if prefix == 'val' or prefix == 'test':\n                checkpointer(score, (epoch + 1), model, optimizer)\n\n        except:\n            results[prefix + \" \" + name].append(float(\"NaN\"))\n            result_str.append(f\"{prefix} {name}: {float('NaN')}\")\n\n    print(\" \".join(result_str))\n\n    return end - start","metadata":{"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2023-09-11T17:39:00.832894Z","iopub.execute_input":"2023-09-11T17:39:00.833728Z","iopub.status.idle":"2023-09-11T17:39:00.850637Z","shell.execute_reply.started":"2023-09-11T17:39:00.833686Z","shell.execute_reply":"2023-09-11T17:39:00.849485Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def train_loop(model, loss_func,\n               train_loader, test_loader=None, val_loader=None,\n               score_funcs=None, epochs=50, device=\"cpu\",\n               optimizer=None, lr_schedule=None, checkpoint_file=None, keep=True):\n    if score_funcs is None:\n        score_funcs = {}\n\n    to_track = [\"epoch\", \"total time\", \"train loss\", \"lr\"]\n\n    if val_loader is not None:\n        to_track.append(\"val loss\")\n\n    if test_loader is not None:\n        to_track.append(\"test loss\")\n\n    for eval_score in score_funcs:\n        to_track.append(\"train \" + eval_score)\n\n        if val_loader is not None:\n            to_track.append(\"val \" + eval_score)\n\n        if test_loader is not None:\n            to_track.append(\"test \" + eval_score)\n\n    total_train_time = 0  # How long have we spent in the training loop?\n    results = {}\n\n    # Initialize every item with an empty list\n    for item in to_track:\n        results[item] = []\n\n    if optimizer is None:\n        optimizer = torch.optim.AdamW(model.parameters(), lr=0.001)\n        del_opt = True\n\n    else:\n        del_opt = False\n\n    # Place the model on the correct computed resource (CPU or GPU)\n    model.to(device)\n\n    for epoch in tqdm(range(epochs), desc=\"Epoch\", leave=keep):\n        model = model.train()  # Put our model in training mode\n\n        run_epoch(model, optimizer, train_loader, loss_func, device, results, score_funcs,\n                  prefix=\"train\", desc=\"Training\")\n\n        results[\"total time\"].append(total_train_time)\n        results[\"epoch\"].append(epoch)\n        results[\"lr\"].append(optimizer.param_groups[0]['lr'])\n\n        if val_loader is not None:\n            model = model.eval()\n\n            with torch.no_grad():\n                run_epoch(model, optimizer, val_loader, loss_func, device, results, score_funcs, prefix=\"val\",\n                          desc=\"Validation\", epoch=epoch)\n\n        if test_loader is not None:\n            model = model.eval()\n\n            with torch.no_grad():\n                run_epoch(model, optimizer, test_loader, loss_func, device, results, score_funcs, prefix=\"test\",\n                          desc=\"Testing\", epoch=epoch)\n\n        if lr_schedule is not None:\n            if isinstance(lr_schedule, torch.optim.lr_scheduler.ReduceLROnPlateau):\n                lr_schedule.step(results[\"val loss\"][-1])\n\n            else:\n                lr_schedule.step()\n\n        if checkpoint_file is not None:\n            torch.save({\n                'epoch': epoch,\n                'model_state_dict': model.state_dict(),\n                'optimizer_state_dict': optimizer.state_dict(),\n                'results': results\n            }, checkpoint_file)\n\n    if del_opt:\n        del optimizer\n\n    return pd.DataFrame.from_dict(results)\n","metadata":{"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2023-09-11T17:39:01.882415Z","iopub.execute_input":"2023-09-11T17:39:01.882856Z","iopub.status.idle":"2023-09-11T17:39:01.902150Z","shell.execute_reply.started":"2023-09-11T17:39:01.882825Z","shell.execute_reply":"2023-09-11T17:39:01.900883Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def pred(model, img):    \n    with torch.no_grad():        \n        logits = model(img.unsqueeze(0).to(device))[0].cpu()\n\n        return torch.sigmoid(logits) >= 0.5","metadata":{"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2023-09-11T17:39:03.239964Z","iopub.execute_input":"2023-09-11T17:39:03.240886Z","iopub.status.idle":"2023-09-11T17:39:03.246491Z","shell.execute_reply.started":"2023-09-11T17:39:03.240842Z","shell.execute_reply":"2023-09-11T17:39:03.245473Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Configure fully convolutional model","metadata":{}},{"cell_type":"code","source":"def cnn_layer(in_filters, out_filters, kernel_size=3):\n    padding = kernel_size // 2\n\n    return nn.Sequential(nn.Conv2d(in_filters, out_filters, kernel_size, padding=padding),\n                         nn.BatchNorm2d(out_filters),\n                         nn.LeakyReLU()\n                         )\n\nclass BaseSegmentation:\n\n    def __init__(self):\n        self.C = 3\n        self.n_filters = 16\n        self.loss_func = nn.BCEWithLogitsLoss()\n\n    def __call__(self, *args, **kwargs):\n        return nn.Sequential(cnn_layer(self.C, self.n_filters),\n                             *[cnn_layer(self.n_filters, self.n_filters) for _ in range(5)],\n                             nn.Conv2d(self.n_filters, 1, (3,3), padding=1)\n                             )","metadata":{"execution":{"iopub.status.busy":"2023-09-11T17:39:06.430679Z","iopub.execute_input":"2023-09-11T17:39:06.431036Z","iopub.status.idle":"2023-09-11T17:39:06.438840Z","shell.execute_reply.started":"2023-09-11T17:39:06.431006Z","shell.execute_reply":"2023-09-11T17:39:06.437593Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# build the model\nbase_seg = BaseSegmentation()\nmodel = base_seg()","metadata":{"execution":{"iopub.status.busy":"2023-09-11T17:50:51.172383Z","iopub.execute_input":"2023-09-11T17:50:51.173581Z","iopub.status.idle":"2023-09-11T17:50:51.182770Z","shell.execute_reply.started":"2023-09-11T17:50:51.173526Z","shell.execute_reply":"2023-09-11T17:50:51.181684Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"summary(model.cuda(), input_size=(3, 128, 128), batch_size=16)","metadata":{"execution":{"iopub.status.busy":"2023-09-11T06:47:07.502171Z","iopub.execute_input":"2023-09-11T06:47:07.502565Z","iopub.status.idle":"2023-09-11T06:47:07.518976Z","shell.execute_reply.started":"2023-09-11T06:47:07.502534Z","shell.execute_reply":"2023-09-11T06:47:07.518047Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model_graph = draw_graph(model, input_size=(1, 3, 128, 128), expand_nested=True)\nmodel_graph.visual_graph","metadata":{"execution":{"iopub.status.busy":"2023-09-11T17:50:55.049505Z","iopub.execute_input":"2023-09-11T17:50:55.049917Z","iopub.status.idle":"2023-09-11T17:50:55.261384Z","shell.execute_reply.started":"2023-09-11T17:50:55.049885Z","shell.execute_reply":"2023-09-11T17:50:55.260374Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Training","metadata":{}},{"cell_type":"code","source":"seg_results = train_loop(model, \n                         base_seg.loss_func, \n                         train_seg_loader, \n                         epochs=10, \n                         device=device,\n                         val_loader=test_seg_loader\n                        )","metadata":{"execution":{"iopub.status.busy":"2023-09-11T06:09:09.688463Z","iopub.execute_input":"2023-09-11T06:09:09.688835Z","iopub.status.idle":"2023-09-11T06:18:41.657655Z","shell.execute_reply.started":"2023-09-11T06:09:09.688804Z","shell.execute_reply":"2023-09-11T06:18:41.656574Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"index = 6\nprediction = pred(model, test_split[index][0])\n\nplt.figure(figsize=(16,10))\nplt.subplot(1, 3, 1)\nplt.imshow(test_split[index][0].permute(1,2,0).numpy(), cmap='gray')\nplt.subplot(1, 3, 2)\nplt.imshow(test_split[index][1].numpy()[0,:], cmap='gray')\nplt.subplot(1, 3, 3)\nplt.imshow(prediction.numpy()[0,:], cmap='gray')","metadata":{"execution":{"iopub.status.busy":"2023-09-11T06:48:45.132072Z","iopub.execute_input":"2023-09-11T06:48:45.132992Z","iopub.status.idle":"2023-09-11T06:48:45.882564Z","shell.execute_reply.started":"2023-09-11T06:48:45.132943Z","shell.execute_reply":"2023-09-11T06:48:45.881556Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# clean up memmory\ndel model\ndel base_seg","metadata":{"execution":{"iopub.status.busy":"2023-09-11T06:57:30.548741Z","iopub.execute_input":"2023-09-11T06:57:30.549116Z","iopub.status.idle":"2023-09-11T06:57:30.553594Z","shell.execute_reply.started":"2023-09-11T06:57:30.549084Z","shell.execute_reply":"2023-09-11T06:57:30.552608Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Configure transpose convolutional model","metadata":{}},{"cell_type":"code","source":"class BaseTransCNNSegmentation:\n\n    def __init__(self):\n        self.C = 3\n        self.n_filters = 32\n        self.loss_func = nn.BCEWithLogitsLoss()\n\n    def __call__(self, *args, **kwargs):\n        return nn.Sequential(cnn_layer(self.C, self.n_filters),\n                             cnn_layer(self.n_filters, self.n_filters),\n                             nn.MaxPool2d(2),\n\n                             cnn_layer(self.n_filters, 2 * self.n_filters),\n                             cnn_layer(2 * self.n_filters, 2 * self.n_filters),\n                             cnn_layer(2 * self.n_filters, 2 * self.n_filters),\n\n                             nn.ConvTranspose2d(2 * self.n_filters, self.n_filters, kernel_size=3,\n                                                padding=1, output_padding=1, stride=2),\n                             nn.BatchNorm2d(self.n_filters),\n                             nn.LeakyReLU(),\n                             cnn_layer(self.n_filters, self.n_filters),\n\n                             nn.Conv2d(self.n_filters, 1, kernel_size=3, padding=1)\n                             )","metadata":{"execution":{"iopub.status.busy":"2023-09-11T17:51:22.272178Z","iopub.execute_input":"2023-09-11T17:51:22.272582Z","iopub.status.idle":"2023-09-11T17:51:22.282233Z","shell.execute_reply.started":"2023-09-11T17:51:22.272544Z","shell.execute_reply":"2023-09-11T17:51:22.281305Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# build the model\nbase_seg = BaseTransCNNSegmentation()\nmodel = base_seg()","metadata":{"execution":{"iopub.status.busy":"2023-09-11T17:51:25.001358Z","iopub.execute_input":"2023-09-11T17:51:25.002070Z","iopub.status.idle":"2023-09-11T17:51:25.013800Z","shell.execute_reply.started":"2023-09-11T17:51:25.002029Z","shell.execute_reply":"2023-09-11T17:51:25.012828Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"summary(model.cuda(), input_size=(3, 128, 128), batch_size=16)","metadata":{"execution":{"iopub.status.busy":"2023-09-11T07:00:57.359543Z","iopub.execute_input":"2023-09-11T07:00:57.359927Z","iopub.status.idle":"2023-09-11T07:00:57.382887Z","shell.execute_reply.started":"2023-09-11T07:00:57.359894Z","shell.execute_reply":"2023-09-11T07:00:57.381812Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model_graph = draw_graph(model, input_size=(1, 3, 128, 128), expand_nested=True)\nmodel_graph.visual_graph","metadata":{"execution":{"iopub.status.busy":"2023-09-11T17:51:29.352250Z","iopub.execute_input":"2023-09-11T17:51:29.352954Z","iopub.status.idle":"2023-09-11T17:51:29.414491Z","shell.execute_reply.started":"2023-09-11T17:51:29.352921Z","shell.execute_reply":"2023-09-11T17:51:29.413623Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Training","metadata":{}},{"cell_type":"code","source":"seg_results = train_loop(model, \n                         base_seg.loss_func, \n                         train_seg_loader, \n                         epochs=10, \n                         device=device,\n                         val_loader=test_seg_loader\n                        )","metadata":{"execution":{"iopub.status.busy":"2023-09-11T07:02:31.500321Z","iopub.execute_input":"2023-09-11T07:02:31.500708Z","iopub.status.idle":"2023-09-11T07:12:02.654471Z","shell.execute_reply.started":"2023-09-11T07:02:31.500676Z","shell.execute_reply":"2023-09-11T07:12:02.653310Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"index = 6\nprediction = pred(model, test_split[index][0])\n\nplt.figure(figsize=(16,10))\nplt.subplot(1, 3, 1)\nplt.imshow(test_split[index][0].permute(1,2,0).numpy(), cmap='gray')\nplt.subplot(1, 3, 2)\nplt.imshow(test_split[index][1].numpy()[0,:], cmap='gray')\nplt.subplot(1, 3, 3)\nplt.imshow(prediction.numpy()[0,:], cmap='gray')","metadata":{"execution":{"iopub.status.busy":"2023-09-11T07:12:21.069845Z","iopub.execute_input":"2023-09-11T07:12:21.070259Z","iopub.status.idle":"2023-09-11T07:12:21.815943Z","shell.execute_reply.started":"2023-09-11T07:12:21.070226Z","shell.execute_reply":"2023-09-11T07:12:21.815080Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# clean up memmory\n\ndel model\ndel base_seg","metadata":{"execution":{"iopub.status.busy":"2023-09-11T07:12:37.161658Z","iopub.execute_input":"2023-09-11T07:12:37.162035Z","iopub.status.idle":"2023-09-11T07:12:37.167030Z","shell.execute_reply.started":"2023-09-11T07:12:37.162001Z","shell.execute_reply":"2023-09-11T07:12:37.165733Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Configure Unet model","metadata":{}},{"cell_type":"code","source":"class UnetBlock(nn.Module):\n\n    def __init__(self, in_channels, mid_channels,\n                 out_channels=None, layers=1,\n                 sub_network=None, filter_size=3):\n        super().__init__()\n\n        in_layers = [cnn_layer(in_channels, mid_channels, filter_size)]\n\n        if sub_network is None:\n            inputs_to_outputs = 1\n\n        else:\n            inputs_to_outputs = 2\n\n        out_layers = [cnn_layer(mid_channels * inputs_to_outputs, mid_channels, filter_size)]\n\n        for _ in range(layers - 1):\n            in_layers.append(cnn_layer(mid_channels, mid_channels, filter_size))\n            out_layers.append(cnn_layer(mid_channels, mid_channels, filter_size))\n\n        if out_channels is not None:\n            out_layers.append(nn.Conv2d(mid_channels, out_channels, kernel_size=1, padding=0))\n\n        self.in_model = nn.Sequential(*in_layers)\n\n        if sub_network is not None:\n            self.bottleneck = nn.Sequential(nn.MaxPool2d(2),\n                                            sub_network,\n                                            nn.ConvTranspose2d(mid_channels, mid_channels, filter_size,\n                                                               padding=filter_size // 2, output_padding=1,\n                                                               stride=2)\n                                            )\n        else:\n            self.bottleneck = None\n\n        self.out_model = nn.Sequential(*out_layers)\n\n    def forward(self, x):\n        full_scale_result = self.in_model(x)\n\n        if self.bottleneck is not None:\n            bottle_result = self.bottleneck(full_scale_result)\n            full_scale_result = torch.cat([full_scale_result, bottle_result], dim=1)\n\n        return self.out_model(full_scale_result)\n    \n\nclass Unet:\n\n    def __init__(self):\n        self.C = 3\n        self.n_filters = 32\n        self.loss_func = nn.BCEWithLogitsLoss()\n\n    def __call__(self, *args, **kwargs):\n        return nn.Sequential(UnetBlock(self.C, self.n_filters, layers=2,\n                                       sub_network=UnetBlock(self.n_filters, 2 * self.n_filters,\n                                                             out_channels=self.n_filters, layers=2,\n                                                             sub_network=UnetBlock(2 * self.n_filters,\n                                                                                   4 * self.n_filters,\n                                                                                   out_channels=2 * self.n_filters,\n                                                                                   layers=2)),\n                                       ),\n                             nn.Conv2d(self.n_filters, 1, (3, 3), padding=1)\n                             )","metadata":{"execution":{"iopub.status.busy":"2023-09-11T17:39:28.803377Z","iopub.execute_input":"2023-09-11T17:39:28.804110Z","iopub.status.idle":"2023-09-11T17:39:28.818119Z","shell.execute_reply.started":"2023-09-11T17:39:28.804076Z","shell.execute_reply":"2023-09-11T17:39:28.817115Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# build the model\nbase_seg = Unet()\nmodel = base_seg()","metadata":{"execution":{"iopub.status.busy":"2023-09-11T17:39:32.449936Z","iopub.execute_input":"2023-09-11T17:39:32.450302Z","iopub.status.idle":"2023-09-11T17:39:32.476603Z","shell.execute_reply.started":"2023-09-11T17:39:32.450270Z","shell.execute_reply":"2023-09-11T17:39:32.475712Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"summary(model.cuda(), input_size=(3, 128, 128), batch_size=16)","metadata":{"execution":{"iopub.status.busy":"2023-09-11T17:39:35.116674Z","iopub.execute_input":"2023-09-11T17:39:35.117342Z","iopub.status.idle":"2023-09-11T17:39:42.814203Z","shell.execute_reply.started":"2023-09-11T17:39:35.117308Z","shell.execute_reply":"2023-09-11T17:39:42.813249Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model_graph = draw_graph(model, input_size=(1, 3, 128, 128), expand_nested=True)\nmodel_graph.visual_graph","metadata":{"execution":{"iopub.status.busy":"2023-09-11T07:17:00.131669Z","iopub.execute_input":"2023-09-11T07:17:00.132047Z","iopub.status.idle":"2023-09-11T07:17:00.363888Z","shell.execute_reply.started":"2023-09-11T07:17:00.132015Z","shell.execute_reply":"2023-09-11T07:17:00.362833Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Training","metadata":{}},{"cell_type":"code","source":"seg_results = train_loop(model, \n                         base_seg.loss_func, \n                         train_seg_loader, \n                         epochs=10, \n                         device=device,\n                         val_loader=test_seg_loader\n                        )","metadata":{"execution":{"iopub.status.busy":"2023-09-11T17:39:48.211403Z","iopub.execute_input":"2023-09-11T17:39:48.212151Z","iopub.status.idle":"2023-09-11T17:49:35.200200Z","shell.execute_reply.started":"2023-09-11T17:39:48.212109Z","shell.execute_reply":"2023-09-11T17:49:35.199269Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"index = 6\nprediction = pred(model, test_split[index][0])\n\nplt.figure(figsize=(16,10))\nplt.subplot(1, 3, 1)\nplt.imshow(test_split[index][0].permute(1,2,0).numpy(), cmap='gray')\nplt.subplot(1, 3, 2)\nplt.imshow(test_split[index][1].numpy()[0,:], cmap='gray')\nplt.subplot(1, 3, 3)\nplt.imshow(prediction.numpy()[0,:], cmap='gray')","metadata":{"execution":{"iopub.status.busy":"2023-09-11T17:49:51.341183Z","iopub.execute_input":"2023-09-11T17:49:51.341576Z","iopub.status.idle":"2023-09-11T17:49:52.153685Z","shell.execute_reply.started":"2023-09-11T17:49:51.341547Z","shell.execute_reply":"2023-09-11T17:49:52.152728Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# clean up memmory\n\ndel model\ndel base_seg","metadata":{"execution":{"iopub.status.busy":"2023-09-11T17:50:20.230986Z","iopub.execute_input":"2023-09-11T17:50:20.231363Z","iopub.status.idle":"2023-09-11T17:50:20.237276Z","shell.execute_reply.started":"2023-09-11T17:50:20.231332Z","shell.execute_reply":"2023-09-11T17:50:20.236336Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"shutil.rmtree(\"./data/\")","metadata":{"execution":{"iopub.status.busy":"2023-09-11T17:50:26.570969Z","iopub.execute_input":"2023-09-11T17:50:26.572005Z","iopub.status.idle":"2023-09-11T17:50:27.558068Z","shell.execute_reply.started":"2023-09-11T17:50:26.571962Z","shell.execute_reply":"2023-09-11T17:50:27.557035Z"},"trusted":true},"execution_count":null,"outputs":[]}],"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"}}