{"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":"import gc\nimport glob\nimport torch\nimport numpy as np\nimport torch.nn as nn\nfrom tqdm import tqdm\nimport PIL.Image as Image\nimport torch.optim as optim\nimport matplotlib.pyplot as plt\nimport torch.utils.data as data\nimport matplotlib.patches as patches\n\nPREFIX_1 = ['/kaggle/input/vesuvius-challenge-ink-detection/train/1/', '/kaggle/input/vesuvius-v4-files/train_1_BCE_layer_1.npy', '/kaggle/input/vesuvius-v4-files/train_1_BCE_layer_2.npy', [16,18,20,22,24,26,28,30,32,34]]\nPREFIX_2 = ['/kaggle/input/vesuvius-challenge-ink-detection/train/2/', '/kaggle/input/vesuvius-v4-files/train_2_BCE_layer_1.npy', '/kaggle/input/vesuvius-v4-files/train_2_BCE_layer_2.npy', [22,24,26,28,30,32,34,36,38,40]]\nPREFIX_3 = ['/kaggle/input/vesuvius-challenge-ink-detection/train/3/', '/kaggle/input/vesuvius-v4-files/train_3_BCE_layer_1.npy', '/kaggle/input/vesuvius-v4-files/train_3_BCE_layer_2.npy', [20,22,24,26,28,30,32,34,36,38]]\n\nBUFFER = 30  # Buffer size in x and y direction\nWINDOW = 1  # Window = pixels in label in x and y direction \nLEARNING_RATE = 0.0002\nBATCH_SIZE = 32\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-06-08T19:35:53.855600Z","iopub.execute_input":"2023-06-08T19:35:53.856585Z","iopub.status.idle":"2023-06-08T19:35:57.447379Z","shell.execute_reply.started":"2023-06-08T19:35:53.856539Z","shell.execute_reply":"2023-06-08T19:35:57.446032Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# This is the updated SubvolumeDataset code block with transforms.\nimport random as random\nimport torchvision.transforms.functional as TF\n\nclass SubvolumeDataset(data.Dataset):\n    def __init__(self, image_stack, label, pixels, DATASET_WINDOW, NEW_BUFFER, transform=False):\n        self.image_stack = image_stack\n        self.label = label\n        self.pixels = pixels\n        self.transform = transform\n        self.DATASET_WINDOW = DATASET_WINDOW\n        self.NEW_BUFFER = NEW_BUFFER\n    def __len__(self):\n        return len(self.pixels)\n    def __getitem__(self, index):\n        y, x = self.pixels[index]\n        subvolume = self.image_stack[:, y-self.NEW_BUFFER:y+self.NEW_BUFFER+1, x-self.NEW_BUFFER:x+self.NEW_BUFFER+1].view(1, len(self.image_stack), self.NEW_BUFFER*2+1, self.NEW_BUFFER*2+1)\n        inklabel = self.label[y-self.DATASET_WINDOW:y+self.DATASET_WINDOW+1, x-self.DATASET_WINDOW:x+self.DATASET_WINDOW+1]\n\n        subvolme = torch.tensor(subvolume)\n        inklabel = torch.tensor(inklabel)\n\n        if self.transform == True:\n            if random.random() > 0.5:\n                subvolume = TF.hflip(subvolume)\n                inklabel = TF.hflip(inklabel)\n            if random.random() > 0.5:\n                subvolume = TF.vflip(subvolume)\n                inklabel = TF.vflip(inklabel)\n        \n        inklabel = inklabel.reshape((self.DATASET_WINDOW*2+1)**2)             \n        return subvolume, inklabel\n    ","metadata":{"execution":{"iopub.status.busy":"2023-06-08T19:36:08.915053Z","iopub.execute_input":"2023-06-08T19:36:08.915802Z","iopub.status.idle":"2023-06-08T19:36:09.194408Z","shell.execute_reply.started":"2023-06-08T19:36:08.915764Z","shell.execute_reply":"2023-06-08T19:36:09.193211Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n# !!! This is the code to generate the layer_1 NN. !!!\n# !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!","metadata":{"execution":{"iopub.status.busy":"2023-05-16T01:56:01.056764Z","iopub.execute_input":"2023-05-16T01:56:01.057519Z","iopub.status.idle":"2023-05-16T01:56:01.084130Z","shell.execute_reply.started":"2023-05-16T01:56:01.057468Z","shell.execute_reply":"2023-05-16T01:56:01.082757Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# This looks outside the rectangle to select random pixels.\ndef train_dataset_gen(PREFIX):   \n    rect = (1,1,1,1)\n    image_selections = PREFIX[3]\n    Z_DIM = len(image_selections) # We add x to this if we're including prediction layer(s).\n    img_layers = sorted(glob.glob(PREFIX[0]+\"surface_volume/*.tif\"))\n\n    images = [np.array(Image.open(img_layers[i-1]), dtype=np.float32)/65535.0 for i in tqdm(image_selections)]\n    \n    # We send the image stack to the GPU, doesn't fit on CPU.\n    image_stack = torch.stack([torch.from_numpy(image) for image in images], dim=0).to(DEVICE)\n    \n    del images\n\n    mask = np.array(Image.open(PREFIX[0]+\"mask.png\").convert('1'))\n    # The label portion goes to the CPU where it's used to compute losses/accuracy.\n    label = torch.from_numpy(np.array(Image.open(PREFIX[0]+\"inklabels.png\"))).gt(0).float()#.to(DEVICE)\n    \n    print(\"Generating pixel lists...\")\n    # rect = (3000, 4500, 1000, 1000)\n    # rect = (1,1,1,1)\n    not_border = np.zeros(mask.shape, dtype=bool)\n    not_border[BUFFER:mask.shape[0]-BUFFER, BUFFER:mask.shape[1]-BUFFER] = True\n    arr_mask = np.array(mask) * not_border\n\n    outside_rect = np.ones(mask.shape, dtype=bool) * arr_mask\n    outside_rect[rect[1]:rect[1]+rect[3]+1, rect[0]:rect[0]+rect[2]+1] = False\n    pixels_outside_rect = np.argwhere(outside_rect)\n    del mask, not_border, arr_mask, outside_rect,\n\n    # Will have to add this back in for the second pass.\n    # image_stack = torch.concatenate((output.reshape(1,mask.shape[0],mask.shape[1]), image_stack))\n    \n    dataset = SubvolumeDataset(image_stack, label, pixels_outside_rect)\n    \n    del label, pixels_outside_rect, image_stack\n    gc.collect()\n    \n    return dataset\n","metadata":{"execution":{"iopub.status.busy":"2023-06-06T17:34:04.776092Z","iopub.execute_input":"2023-06-06T17:34:04.777102Z","iopub.status.idle":"2023-06-06T17:34:04.792062Z","shell.execute_reply.started":"2023-06-06T17:34:04.777045Z","shell.execute_reply":"2023-06-06T17:34:04.790932Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# This is to combine multiple datasets. We do 2 first beacause it's huge.\nprint(\"processing dataset 2...\")\ntrain_dataset_1 = train_dataset_gen(PREFIX_2)\nprint(\"processing dataset 1...\")\ntrain_dataset_2 = train_dataset_gen(PREFIX_1)\ntrain_dataset = data.ConcatDataset((train_dataset_1, train_dataset_2))\n\ndel train_dataset_1\ndel train_dataset_2\ngc.collect()\n\nprint(\"processing dataset 3...\")\ntrain_dataset_3 = train_dataset_gen(PREFIX_3)\ntrain_dataset = data.ConcatDataset((train_dataset, train_dataset_3))\n\ndel train_dataset_3\ngc.collect()\n\ntrain_loader = data.DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)\n\ndel train_dataset\ngc.collect()\n\n# Defining the model:\nmodel = nn.Sequential(\n    nn.Conv3d(1, 16, 3, 1, 1), nn.MaxPool3d(2, 2),\n    nn.Conv3d(16, 32, 3, 1, 1), nn.MaxPool3d(2, 2),\n    nn.Conv3d(32, 64, 3, 1, 1), nn.MaxPool3d(2, 2),\n    nn.Flatten(start_dim=1),\n    nn.LazyLinear(128), nn.ReLU(),\n    nn.LazyLinear(9), nn.Sigmoid()\n).to(DEVICE)\n\n# Training the model:\ncriterion = nn.BCELoss()\n\nTRAINING_STEPS = 200000\noptimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)\nscheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=LEARNING_RATE, total_steps=TRAINING_STEPS)\nmodel.train()\ngc.collect()\n\nprint(\"Training...\")\nrunning_loss = 0.0\nfor i, (subvolumes, inklabels) in tqdm(enumerate(train_loader), total=TRAINING_STEPS):\n    if i >= TRAINING_STEPS:\n        break\n    optimizer.zero_grad()\n    outputs = model(subvolumes.to(DEVICE))\n    loss = criterion(outputs, inklabels.to(DEVICE))\n    loss.backward()\n    optimizer.step()\n    scheduler.step()\n    running_loss += loss.item()\n    if i % (TRAINING_STEPS/10) == (TRAINING_STEPS/10)-1:\n        print(\"Loss:\", running_loss / (TRAINING_STEPS/10))\n        running_loss = 0.0     \n\ngc.collect()\n    \ntorch.save(model, 'bce_model_L1.pth')","metadata":{"execution":{"iopub.status.busy":"2023-06-06T17:34:07.582068Z","iopub.execute_input":"2023-06-06T17:34:07.582973Z","iopub.status.idle":"2023-06-06T18:45:35.545457Z","shell.execute_reply.started":"2023-06-06T17:34:07.582930Z","shell.execute_reply":"2023-06-06T18:45:35.544237Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# This is to combine multiple datasets. We do 2 first beacause it's huge. \n# The rect can be used to block pixels if images take up too much RAM.\nrect_1 = (1,1,1,1)\nrect_2 = (1,1,1,1)\nrect_3 = (1,1,1,1)\n","metadata":{"execution":{"iopub.status.busy":"2023-06-06T01:40:18.427178Z","iopub.execute_input":"2023-06-06T01:40:18.427880Z","iopub.status.idle":"2023-06-06T01:40:18.433827Z","shell.execute_reply.started":"2023-06-06T01:40:18.427839Z","shell.execute_reply":"2023-06-06T01:40:18.432611Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"j=21\nwhile j < 32:\n    surface_volumes = [j, j+1, j+2, j+3, j+4, j+5, j+6, j+7, j+8, j+9]\n\n    print(\"processing dataset...\")\n    train_dataset = train_dataset_gen(PREFIX_3, rect_3, surface_volumes)\n    train_loader = data.DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)\n\n    del train_dataset\n    gc.collect()\n\n    # Defining the model:\n    model = nn.Sequential(\n        nn.Conv3d(1, 16, 3, 1, 1), nn.MaxPool3d(2, 2),\n        nn.Conv3d(16, 32, 3, 1, 1), nn.MaxPool3d(2, 2),\n        nn.Conv3d(32, 64, 3, 1, 1), nn.MaxPool3d(2, 2),\n        nn.Flatten(start_dim=1),\n        nn.LazyLinear(128), nn.ReLU(),\n        nn.LazyLinear(9), nn.Sigmoid()\n    ).to(DEVICE)\n\n    # Training the model:\n    criterion = nn.BCELoss()\n\n    TRAINING_STEPS = 25000\n    optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)\n    scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=LEARNING_RATE, total_steps=TRAINING_STEPS)\n    model.train()\n\n    print(\"Training...\")\n\n    running_loss = 0.0\n    pbar = tqdm(enumerate(train_loader), total=TRAINING_STEPS)\n    for i, (subvolumes, inklabels) in pbar:\n        if i >= TRAINING_STEPS:\n            break\n        optimizer.zero_grad()\n        outputs = model(subvolumes.to(DEVICE))\n        loss = criterion(outputs, inklabels.to(DEVICE))\n        loss.backward()\n        optimizer.step()\n        scheduler.step()\n        running_loss += loss.item()\n        if i % (TRAINING_STEPS/10) == (TRAINING_STEPS/10)-1:\n            print(\"Loss:\", running_loss / (TRAINING_STEPS/10))\n            running_loss = 0.0    \n    \n    print(surface_volumes)\n    print(\"----------------------\")\n    j+=1\n    del train_loader, model, pbar, outputs, loss\n    gc.collect()\n\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n# !!! Now we're training the model again, but with a prediction output layer included as input. !!!\n# !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"OLD_WINDOW = 1\nDATASET_WINDOW = 2\nOLD_BUFFER = 30\nNEW_BUFFER = 60\n\ndef train_dataset_gen(PREFIX):    \n    image_selections = PREFIX[3]\n    Z_DIM = len(image_selections) + 1 # We add x to this if we're including prediction layer(s).\n    img_layers = sorted(glob.glob(PREFIX[0]+\"surface_volume/*.tif\"))\n\n    images = [np.array(Image.open(img_layers[i-1]), dtype=np.float32)/65535.0 for i in tqdm(image_selections)]\n    image_stack = torch.stack([torch.from_numpy(image) for image in images], dim=0).to(DEVICE)\n    del images\n    gc.collect()\n\n    output_1 = np.float32(np.load(PREFIX[1]))\n    mask = np.array(Image.open(PREFIX[0]+\"mask.png\").convert('1'))\n    label = torch.from_numpy(np.array(Image.open(PREFIX[0]+\"inklabels.png\"))).gt(0).float().to(DEVICE)\n    \n    not_border = np.zeros(mask.shape, dtype=bool)\n    not_border[OLD_BUFFER:mask.shape[0]-OLD_BUFFER, OLD_BUFFER:mask.shape[1]-OLD_BUFFER] = True\n\n    grid = np.zeros(mask.shape)\n    i = OLD_WINDOW\n    while i < grid.shape[0]:\n        j = OLD_WINDOW\n        while j < grid.shape[1]:\n            grid[i,j] = 1\n            j += 2*OLD_WINDOW+1\n        i += 2*OLD_WINDOW+1\n\n    arr_mask = np.array(mask) * not_border * grid    \n    full_rect = np.ones(mask.shape, dtype=bool) * arr_mask\n    all_pixels = np.argwhere(full_rect)\n    \n    for i in range(len(output_1)):\n        pixel_y = all_pixels[i][0]\n        pixel_x = all_pixels[i][1]\n        full_rect[pixel_y-OLD_WINDOW:pixel_y+OLD_WINDOW+1,pixel_x-OLD_WINDOW:pixel_x+OLD_WINDOW+1] = output_1[i].reshape(3,3)\n    \n    del output_1, all_pixels\n\n    image_stack = torch.concatenate((torch.tensor(full_rect, dtype=torch.float32).reshape(1,mask.shape[0],mask.shape[1]).to(DEVICE), image_stack))\n    del full_rect\n    gc.collect()\n    \n    print(\"Generating pixel lists...\")\n    not_border = np.zeros(mask.shape, dtype=bool)\n    not_border[NEW_BUFFER:mask.shape[0]-NEW_BUFFER, NEW_BUFFER:mask.shape[1]-NEW_BUFFER] = True\n    rect = (3000, 4500, 1000, 1000)\n    arr_mask = np.array(mask) * not_border\n\n    outside_rect = np.ones(mask.shape, dtype=bool) * arr_mask\n    outside_rect[rect[1]:rect[1]+rect[3]+1, rect[0]:rect[0]+rect[2]+1] = False\n    pixels_outside_rect = np.argwhere(outside_rect)\n    del mask, not_border, arr_mask, outside_rect, grid\n    gc.collect()\n    \n    dataset = SubvolumeDataset(image_stack, label, pixels_outside_rect, DATASET_WINDOW, NEW_BUFFER)\n    del label, pixels_outside_rect, image_stack\n    gc.collect()\n    \n    return dataset\n","metadata":{"execution":{"iopub.status.busy":"2023-06-08T19:36:15.985259Z","iopub.execute_input":"2023-06-08T19:36:15.985747Z","iopub.status.idle":"2023-06-08T19:36:16.010394Z","shell.execute_reply.started":"2023-06-08T19:36:15.985712Z","shell.execute_reply":"2023-06-08T19:36:16.009015Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# This is to combine multiple datasets. We do 2 first beacause it's huge.\nprint(\"processing dataset 2...\")\ntrain_dataset_1 = train_dataset_gen(PREFIX_2)\nprint(\"processing dataset 1...\")\ntrain_dataset_2 = train_dataset_gen(PREFIX_1)\ntrain_dataset = data.ConcatDataset((train_dataset_1, train_dataset_2))\n\ndel train_dataset_1\ndel train_dataset_2\ngc.collect()\n\nprint(\"processing dataset 3...\")\ntrain_dataset_3 = train_dataset_gen(PREFIX_3)\ntrain_dataset = data.ConcatDataset((train_dataset, train_dataset_3))\n\ndel train_dataset_3\ngc.collect()\n\ntrain_loader = data.DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)\n\ndel train_dataset\ngc.collect()\n\n# Defining the model:\nmodel = nn.Sequential(\n    nn.Conv3d(1, 16, 3, 1, 1), nn.MaxPool3d(2, 2),\n    nn.Conv3d(16, 32, 3, 1, 1), nn.MaxPool3d(2, 2),\n    nn.Conv3d(32, 64, 3, 1, 1), nn.MaxPool3d(2, 2),\n    nn.Flatten(start_dim=1),\n    nn.LazyLinear(128), nn.ReLU(),\n    nn.LazyLinear((DATASET_WINDOW*2+1)**2), nn.Sigmoid()\n).to(DEVICE)\n","metadata":{"execution":{"iopub.status.busy":"2023-06-08T19:36:19.273986Z","iopub.execute_input":"2023-06-08T19:36:19.275061Z","iopub.status.idle":"2023-06-08T19:39:47.074450Z","shell.execute_reply.started":"2023-06-08T19:36:19.275001Z","shell.execute_reply":"2023-06-08T19:39:47.073097Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Training the model:\ncriterion = nn.BCELoss()\n\nTRAINING_STEPS = 200000\noptimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)\nscheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=LEARNING_RATE, total_steps=TRAINING_STEPS)\nmodel.train()\ngc.collect()\n\nprint(\"Training...\")\nrunning_loss = 0.0\nfor i, (subvolumes, inklabels) in tqdm(enumerate(train_loader), total=TRAINING_STEPS):\n    if i >= TRAINING_STEPS:\n        break\n    optimizer.zero_grad()\n    outputs = model(subvolumes.to(DEVICE))\n    loss = criterion(outputs, inklabels.to(DEVICE))\n    loss.backward()\n    optimizer.step()\n    scheduler.step()\n    running_loss += loss.item()\n    if i % (TRAINING_STEPS/10) == (TRAINING_STEPS/10)-1:\n        print(\"Loss:\", running_loss / (TRAINING_STEPS/10))\n        running_loss = 0.0     \n\ngc.collect()\n    \ntorch.save(model, 'bce_model_L2.pth')","metadata":{"execution":{"iopub.status.busy":"2023-06-08T19:39:47.076700Z","iopub.execute_input":"2023-06-08T19:39:47.077106Z","iopub.status.idle":"2023-06-08T23:31:41.659897Z","shell.execute_reply.started":"2023-06-08T19:39:47.077058Z","shell.execute_reply":"2023-06-08T23:31:41.658478Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n# !!! This trains the layer-3 NN !!!n\n# !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def train_dataset_gen(PREFIX, rect):    \n    image_selections = PREFIX[3]\n    image_selections = image_selections[1:len(image_selections)-1]\n    Z_DIM = len(image_selections) + 2 # We add x to this if we're including prediction layer(s).\n    img_layers = sorted(glob.glob(PREFIX[0]+\"surface_volume/*.tif\"))\n\n    images = [np.array(Image.open(img_layers[i-1]), dtype=np.float32)/65535.0 for i in tqdm(image_selections)]\n    image_stack = torch.stack([torch.from_numpy(image) for image in images], dim=0).to(DEVICE)\n    del images\n    gc.collect()\n\n    # This adds the first output layer.\n    output_1 = np.float32(np.load(PREFIX[1]))\n    mask = np.array(Image.open(PREFIX[0]+\"mask.png\").convert('1'))\n    label = torch.from_numpy(np.array(Image.open(PREFIX[0]+\"inklabels.png\"))).gt(0).float()#.to(DEVICE)\n    \n    not_border = np.zeros(mask.shape, dtype=bool)\n    not_border[BUFFER:mask.shape[0]-BUFFER, BUFFER:mask.shape[1]-BUFFER] = True\n\n    grid = np.zeros(mask.shape)\n    i = WINDOW\n    while i < grid.shape[0]:\n        j = WINDOW\n        while j < grid.shape[1]:\n            grid[i,j] = 1\n            j += 2*WINDOW+1\n        i += 2*WINDOW+1\n\n    arr_mask = np.array(mask) * not_border * grid    \n    full_rect = np.ones(mask.shape, dtype=bool) * arr_mask\n    all_pixels = np.argwhere(full_rect)\n    \n    for i in range(len(output_1)):\n        pixel_y = all_pixels[i][0]\n        pixel_x = all_pixels[i][1]\n        full_rect[pixel_y-WINDOW:pixel_y+WINDOW+1,pixel_x-WINDOW:pixel_x+WINDOW+1] = output_1[i].reshape(3,3)\n    \n    del output_1, all_pixels\n\n    image_stack = torch.concatenate((torch.tensor(full_rect, dtype=torch.float32).reshape(1,mask.shape[0],mask.shape[1]).to(DEVICE), image_stack))\n    del full_rect\n    gc.collect()\n    \n    # This adds the second output layer.\n    output_2 = np.float32(np.load(PREFIX[2]))\n    mask = np.array(Image.open(PREFIX[0]+\"mask.png\").convert('1'))\n    label = torch.from_numpy(np.array(Image.open(PREFIX[0]+\"inklabels.png\"))).gt(0).float().to(DEVICE)\n    \n    not_border = np.zeros(mask.shape, dtype=bool)\n    not_border[BUFFER:mask.shape[0]-BUFFER, BUFFER:mask.shape[1]-BUFFER] = True\n\n    grid = np.zeros(mask.shape)\n    i = WINDOW\n    while i < grid.shape[0]:\n        j = WINDOW\n        while j < grid.shape[1]:\n            grid[i,j] = 1\n            j += 2*WINDOW+1\n        i += 2*WINDOW+1\n\n    arr_mask = np.array(mask) * not_border * grid    \n    full_rect = np.ones(mask.shape, dtype=bool) * arr_mask\n    all_pixels = np.argwhere(full_rect)\n    \n    for i in range(len(output_2)):\n        pixel_y = all_pixels[i][0]\n        pixel_x = all_pixels[i][1]\n        full_rect[pixel_y-WINDOW:pixel_y+WINDOW+1,pixel_x-WINDOW:pixel_x+WINDOW+1] = output_2[i].reshape(3,3)\n    \n    del output_2, all_pixels\n\n    image_stack = torch.concatenate((torch.tensor(full_rect, dtype=torch.float32).reshape(1,mask.shape[0],mask.shape[1]).to(DEVICE), image_stack))\n    del full_rect\n    gc.collect()   \n    \n    print(\"Generating pixel lists...\")\n    # rect = (3000, 4500, 1000, 1000)\n    arr_mask = np.array(mask) * not_border\n\n    outside_rect = np.ones(mask.shape, dtype=bool) * arr_mask\n    outside_rect[rect[1]:rect[1]+rect[3]+1, rect[0]:rect[0]+rect[2]+1] = False\n    pixels_outside_rect = np.argwhere(outside_rect)\n    del mask, not_border, arr_mask, outside_rect, grid\n    gc.collect()\n    \n    dataset = SubvolumeDataset(image_stack, label, pixels_outside_rect)\n    del label, pixels_outside_rect, image_stack\n    gc.collect()\n    \n    return dataset\n","metadata":{"execution":{"iopub.status.busy":"2023-06-07T19:34:47.961754Z","iopub.execute_input":"2023-06-07T19:34:47.962258Z","iopub.status.idle":"2023-06-07T19:34:48.005612Z","shell.execute_reply.started":"2023-06-07T19:34:47.962211Z","shell.execute_reply":"2023-06-07T19:34:48.004243Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# This is to combine multiple datasets. We do 2 first beacause it's huge.\nrect_1 = (1,1,1,1)\nrect_2 = (1,1,4000,4000)\nrect_3 = (1,1,1,1)\n\nprint(\"processing dataset 2...\")\ntrain_dataset_1 = train_dataset_gen(PREFIX_2, rect_2)\nprint(\"processing dataset 1...\")\ntrain_dataset_2 = train_dataset_gen(PREFIX_1, rect_1)\ntrain_dataset = data.ConcatDataset((train_dataset_1, train_dataset_2))\n\ndel train_dataset_1\ndel train_dataset_2\ngc.collect()\n\nprint(\"processing dataset 3...\")\ntrain_dataset_3 = train_dataset_gen(PREFIX_3, rect_3)\ntrain_dataset = data.ConcatDataset((train_dataset, train_dataset_3))\n\ndel train_dataset_3\ngc.collect()\n\ntrain_loader = data.DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)\n\ndel train_dataset\ngc.collect()\n\n# Defining the model:\nmodel = nn.Sequential(\n    nn.Conv3d(1, 16, 3, 1, 1), nn.MaxPool3d(2, 2),\n    nn.Conv3d(16, 32, 3, 1, 1), nn.MaxPool3d(2, 2),\n    nn.Conv3d(32, 64, 3, 1, 1), nn.MaxPool3d(2, 2),\n    nn.Flatten(start_dim=1),\n    nn.LazyLinear(128), nn.ReLU(),\n    nn.LazyLinear(9), nn.Sigmoid()\n).to(DEVICE)","metadata":{"execution":{"iopub.status.busy":"2023-06-07T19:34:51.553542Z","iopub.execute_input":"2023-06-07T19:34:51.554513Z","iopub.status.idle":"2023-06-07T19:39:34.858801Z","shell.execute_reply.started":"2023-06-07T19:34:51.554435Z","shell.execute_reply":"2023-06-07T19:39:34.857548Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Training the model:\ncriterion = nn.BCELoss()\n\nTRAINING_STEPS = 300000\noptimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)\nscheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=LEARNING_RATE, total_steps=TRAINING_STEPS)\nmodel.train()\ngc.collect()\n\nprint(\"Training...\")\nrunning_loss = 0.0\nfor i, (subvolumes, inklabels) in tqdm(enumerate(train_loader), total=TRAINING_STEPS):\n    if i >= TRAINING_STEPS:\n        break\n    optimizer.zero_grad()\n    outputs = model(subvolumes.to(DEVICE))\n    loss = criterion(outputs, inklabels.to(DEVICE))\n    loss.backward()\n    optimizer.step()\n    scheduler.step()\n    running_loss += loss.item()\n    if i % (TRAINING_STEPS/10) == (TRAINING_STEPS/10)-1:\n        print(\"Loss:\", running_loss / (TRAINING_STEPS/10))\n        running_loss = 0.0     \n\ngc.collect()\n    \ntorch.save(model, 'bce_model_L3.pth')","metadata":{"execution":{"iopub.status.busy":"2023-06-07T19:39:34.860866Z","iopub.execute_input":"2023-06-07T19:39:34.864118Z","iopub.status.idle":"2023-06-07T21:23:03.945144Z","shell.execute_reply.started":"2023-06-07T19:39:34.864085Z","shell.execute_reply":"2023-06-07T21:23:03.944019Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n# !!! This processes the output if you want to check it. !!!\n# !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n\nfrom pathlib import Path \n\nclass SubvolumeDataset_eval(data.Dataset):\n    def __init__(self, image_stack, pixels):\n        self.image_stack = image_stack\n        self.pixels = pixels\n    def __len__(self):\n        return len(self.pixels)\n    def __getitem__(self, index):\n        y, x = self.pixels[index]\n        subvolume = self.image_stack[:, y-BUFFER:y+BUFFER+1, x-BUFFER:x+BUFFER+1].view(1, len(self.image_stack), BUFFER*2+1, BUFFER*2+1)\n        return subvolume\n\n# This is detecting the locations of test fragments.\nbase_path = Path(\"/kaggle/input/vesuvius-challenge-ink-detection/test/\")\ntest_fragments = sorted([fragment_name for fragment_name in base_path.iterdir()])\n\n# This creates the output dictionary.\nsubmission_dict = {'ID': [], 'Predicted': []}\n\nfor PREFIX in test_fragments:\n    fragment_name = PREFIX.name\n    PREFIX = str(PREFIX)+ '/'\n    print(PREFIX)\n\n    # This is processing the layers for the first pass.\n    image_selections = [21,23,25,27,29,31,33,35,37,39]\n    # Z_DIM = len(image_selections) # We add x to this if we're including prediction layer(s).\n    img_layers = sorted(glob.glob(PREFIX+\"surface_volume/*.tif\"))\n\n    images = [np.array(Image.open(img_layers[i-1]), dtype=np.float32)/65535.0 for i in tqdm(image_selections)]\n    image_stack = torch.stack([torch.from_numpy(image) for image in images], dim=0).to(DEVICE)\n    del images\n\n    mask = np.array(Image.open(PREFIX+\"mask.png\").convert('1'))\n\n    not_border = np.zeros(mask.shape, dtype=bool)\n    not_border[BUFFER:mask.shape[0]-BUFFER, BUFFER:mask.shape[1]-BUFFER] = True\n\n    grid = np.zeros(mask.shape)\n    i = WINDOW\n    while i < grid.shape[0]:\n        j = WINDOW\n        while j < grid.shape[1]:\n            grid[i,j] = 1\n            j += 2*WINDOW+1\n        i += 2*WINDOW+1\n\n    arr_mask = np.array(mask) * not_border * grid    \n    full_rect = np.ones(mask.shape, dtype=bool) * arr_mask\n    all_pixels = np.argwhere(full_rect)\n\n    del not_border, arr_mask, grid\n\n    eval_dataset = SubvolumeDataset_eval(image_stack, all_pixels)\n\n    image_stack\n    gc.collect()\n\n    eval_loader = data.DataLoader(eval_dataset, batch_size=BATCH_SIZE, shuffle=False)\n\n#     # This is running the model on the first pass.\n#     if torch.cuda.is_available():\n#         model = torch.load('/kaggle/input/vesuvius-v4-files/bce4_model_L1.pth').to(DEVICE)\n#     else:\n#         model = torch.load('/kaggle/input/vesuvius-v4-files/bce4_model_L1.pth', map_location=torch.device('cpu')).to(DEVICE)\n    output_1=[]\n    model.eval()\n    with torch.no_grad():\n        for i, (subvolumes) in enumerate(tqdm(eval_loader)):\n            for j, value in enumerate(model(subvolumes.to(DEVICE))):\n                output_1.append(value.tolist())\n\n    # This is processing the layers for the second pass.\n    output_1 = np.array(output_1)\n    layer_2_input = full_rect\n\n    for i in range(len(output_1)):\n        pixel_y = all_pixels[i][0]\n        pixel_x = all_pixels[i][1]\n        layer_2_input[pixel_y-WINDOW:pixel_y+WINDOW+1,pixel_x-WINDOW:pixel_x+WINDOW+1] = output_1[i].reshape(3,3)\n\n    # This is checking to make sure it looks normal.\n    fig, (ax1) = plt.subplots(1, 1)\n    ax1.imshow(layer_2_input, cmap='gray')\n    plt.show()\n","metadata":{"execution":{"iopub.status.busy":"2023-05-16T17:31:11.838368Z","iopub.execute_input":"2023-05-16T17:31:11.838828Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# This was an experiment with transforms, which seemed to help with exploding gradients,\n# but I couldn't get it to work when predicting on 9 pixels at a time. It also doesn't\n# seem to be as useful when loading all three data sets.\n\n# import random as random\n# from torchvision import transforms\n# my_transforms = transforms.Compose([\n#             transforms.RandomHorizontalFlip(p=0.5),\n#             transforms.RandomVerticalFlip(p=0.5)\n#             ])\n\n# import torchvision.transforms.functional as TF\n\n# class SubvolumeDataset(data.Dataset):\n#     def __init__(self, image_stack, label, pixels, transform=False):\n#         self.image_stack = image_stack\n#         self.label = label\n#         self.pixels = pixels\n#         self.transform = transform\n#     def __len__(self):\n#         return len(self.pixels)\n#     def __getitem__(self, index):\n#         y, x = self.pixels[index]\n#         subvolume = self.image_stack[:, y-BUFFER:y+BUFFER+1, x-BUFFER:x+BUFFER+1].view(1, len(self.image_stack), BUFFER*2+1, BUFFER*2+1)\n#         inklabel = self.label[y-WINDOW:y+WINDOW+1, x-WINDOW:x+WINDOW+1].reshape(9)\n#         # inklabel = self.label[y, x].view(1)\n#         if self.transform == True:\n#             if random.random() > 0.5:\n#                 subvolume = torch.tensor(TF.hflip(subvolume))\n#                 inklabel = torch.tensor(TF.hflip(inklabel))\n#             if random.random() > 0.5:\n#                 subvolume = torch.tensor(TF.vflip(subvolume))\n#                 inklabel = torch.tensor(TF.vflip(inklabel))\n#             # subvolume = my_transforms(subvolume)\n#         return subvolume, inklabel\n","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# This produces a prediction dataset of all pixels within the rectangle. !! NO GRID, 100% of pixels !!\n\n# def eval_dataset_gen(PREFIX[0]):\n#     mask = np.array(Image.open(PREFIX[0]+\"mask.png\").convert('1'))\n#     label = torch.from_numpy(np.array(Image.open(PREFIX[0]+\"inklabels.png\"))).gt(0).float().to(DEVICE)\n#     images = [np.array(Image.open(filename), dtype=np.float32)/65535.0 for filename in tqdm(sorted(glob.glob(PREFIX[0]+\"surface_volume/*.tif\"))[Z_START:Z_START+Z_DIM])]\n\n#     print(\"Generating pixel lists...\")\n#     rect = (3000, 4500, 1000, 1000)\n#     not_border = np.zeros(mask.shape, dtype=bool)\n#     not_border[BUFFER:mask.shape[0]-BUFFER, BUFFER:mask.shape[1]-BUFFER] = True\n#     arr_mask = np.array(mask) * not_border\n\n#     inside_rect = np.zeros(mask.shape, dtype=bool) * arr_mask\n#     inside_rect[rect[1]:rect[1]+rect[3]+1, rect[0]:rect[0]+rect[2]+1] = True\n#     pixels_inside_rect = np.argwhere(inside_rect)\n#     del mask, not_border, arr_mask, inside_rect\n\n#     image_stack = torch.stack([torch.from_numpy(image) for image in images], dim=0).to(DEVICE)\n#     del images\n#     # Will have to add this back in for the second pass.\n#     # image_stack = torch.concatenate((output.reshape(1,mask.shape[0],mask.shape[1]), image_stack))\n    \n#     # Because this is generating a validation dataset, we pass 'pixels_inside_rect' to SubvolumeDataset\n#     dataset = SubvolumeDataset(image_stack, label, pixels_inside_rect)\n    \n#     del image_stack\n#     gc.collect()\n    \n#     return label, dataset, pixels_inside_rect, rect\n","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# This is to combine two datasets instead of three. Was used before I reduced memory requirements.\n\n# train_dataset_1 = train_dataset_gen(PREFIX_3)\n# train_dataset_2 = train_dataset_gen(PREFIX_2)\n# train_dataset = data.ConcatDataset((train_dataset_1, train_dataset_2))\n\n# del train_dataset_1\n# del train_dataset_2\n\n# train_loader = data.DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)\n\n# label, eval_dataset, pixels_inside_rect, rect = eval_dataset_gen(PREFIX_3)\n# eval_loader = data.DataLoader(eval_dataset, batch_size=BATCH_SIZE, shuffle=False)\n","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# This is the code from Namgil Lee: https://arxiv.org/abs/2104.01459\n# It makes f-beta differentiable so we can train on it:\n\n# def reduce_loss(loss: torch.Tensor, reduction: str = 'mean') -> torch.Tensor:\n#     if reduction == 'mean':\n#         loss = torch.mean(loss)\n#     elif reduction == 'sum':\n#         loss = torch.sum(loss)\n#     elif reduction == 'none':\n#         pass\n#     else:\n#         log.warning(f\"Unknown reduction {reduction}\")\n#     return loss\n\n\n# class SurrogateFBetaLoss(nn.Module):\n#     torch_eps = torch.tensor(1e-12)\n#     eps = 1e-12\n\n#     def __init__(self, beta: float = 0.99, reduction: str = 'sum') -> None:\n#         \"\"\"\n#         Surrogate F-beta loss for binary classification.\n\n#         :param beta: parameter beta for balancing between precision and recall.\n#         :param positive_rate: rate of positive samples with value in [0, 1].\n#         :param reduction: the type of reduction to apply to the return value.\n#         \"\"\"\n#         super(SurrogateFBetaLoss, self).__init__()\n#         self.beta = beta\n#         # self.p = positive_rate if positive_rate < (1 - self.eps) else (1 - self.eps)\n#         self.reduction = reduction\n\n#     def forward(self, predictions: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:\n#         \"\"\"\n#         :param predictions: Tensor of arbitrary shape with values between 0 and 1.\n#         :param labels: Tensor of the same shape as predictions.\n#         :return: the surrogate F-beta loss between the predictions and labels.\n#         \"\"\"\n#         positive_rate = torch.mean(labels)\n#         device = predictions.device\n#         torch_eps = self.torch_eps.to(device=device)\n#         predictions = torch.where(predictions > torch_eps, predictions, torch_eps)\n#         loss = -labels * torch.log(predictions) + \\\n#             (1 - labels) * torch.log((self.beta * self.beta) * positive_rate / (1 - positive_rate) + predictions)\n\n#         return reduce_loss(loss, self.reduction)\n","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# This is used to debug the dataloader:\n\n# import itertools\n# my_sample = next(itertools.islice(train_loader, 10, None))\n\n# enum = enumerate(train_loader)\n# print(len(train_loader))\n","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Here we're putting the training and prediction into a loop. We call each loop an 'epoch'.\n# This also includes the SurrogateFBetaLoss as a loss function option.\n\n# criterion = nn.BCELoss()\n# epoch = 1 # Start this at 1, not 0.\n\n# while epoch < 2:\n#     TRAINING_STEPS = 100000*epoch\n#     optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)\n#     scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=LEARNING_RATE, total_steps=TRAINING_STEPS)\n#     model.train()\n    \n#     print(\"Training...\")\n#     running_loss = 0.0\n#     for i, (subvolumes, inklabels) in tqdm(enumerate(train_loader), total=TRAINING_STEPS):\n#         if i >= TRAINING_STEPS:\n#             break\n#         optimizer.zero_grad()\n#         outputs = model(subvolumes.to(DEVICE))\n#         loss = criterion(outputs, inklabels.to(DEVICE))\n#         # loss = SurrogateFBetaLoss().forward(outputs, inklabels.to(DEVICE))\n#         loss.backward()\n#         optimizer.step()\n#         scheduler.step()\n#         running_loss += loss.item()\n#         if i % (TRAINING_STEPS/10) == (TRAINING_STEPS/10)-1:\n#             print(\"Loss:\", running_loss / (TRAINING_STEPS/10))\n#             running_loss = 0.0\n\n#     print(\"Predicting...\")\n#     output = torch.zeros_like(label).float()\n#     model.eval()\n#     with torch.no_grad():\n#         for i, (subvolumes, _) in enumerate(tqdm(eval_loader)):\n#             for j, value in enumerate(model(subvolumes.to(DEVICE))):\n#                 output[tuple(pixels_inside_rect[i*BATCH_SIZE+j])] = value\n\n#     fig, (ax1, ax2) = plt.subplots(1, 2)\n#     ax1.imshow(output.cpu(), cmap='gray')\n#     ax2.imshow(label.cpu())\n#     patch = patches.Rectangle((rect[0], rect[1]), rect[2], rect[3], linewidth=2, edgecolor='r', facecolor='none')\n#     ax2.add_patch(patch)\n#     plt.show()        \n    \n#     gc.collect()\n#     epoch += 1\n    \n# torch.save(model, 'bce_model_L1.pth')\n","metadata":{},"execution_count":null,"outputs":[]}]}