{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":5441,"databundleVersionId":38425,"sourceType":"competition"}],"dockerImageVersionId":30732,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# 1. Simple Dataset downloading pipeline","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport torch, torch.nn as nn\nimport torch.nn.functional as F\nfrom matplotlib import pyplot as plt\nfrom PIL import Image\nimport datasets\nimport torchvision\nfrom tqdm import tqdm\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")","metadata":{"execution":{"iopub.status.busy":"2024-06-27T10:12:25.375390Z","iopub.execute_input":"2024-06-27T10:12:25.376034Z","iopub.status.idle":"2024-06-27T10:12:32.978381Z","shell.execute_reply.started":"2024-06-27T10:12:25.375998Z","shell.execute_reply":"2024-06-27T10:12:32.977249Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Custom_Dataset(torch.utils.data.Dataset):\n    def __init__(self, huggingface_dataset):\n        self.dataset = huggingface_dataset\n        self.transform =     transformation_list =  torchvision.transforms.Compose([\n            torchvision.transforms.Resize((224, 224)), # IMPORTANT SIZE OF IMAGE (224, 224) or (256,256)\n            torchvision.transforms.Grayscale(),        # TODO: Change to colored image\n            torchvision.transforms.ToTensor(),\n        ])\n\n    def __len__(self):\n        self.dataset_length = len(self.dataset)\n        return self.dataset_length\n\n    def __getitem__(self,index):\n        img   = self.dataset[index]['image']\n        label = self.dataset[index]['labels']\n        \n        img_transformed = self.transform(img).to(dtype=torch.float32)\n\n        return img_transformed,label\n    \ntraining_dataset                     = Custom_Dataset(huggingface_dataset = datasets.load_dataset(\"microsoft/cats_vs_dogs\", split=\"train\"))\ntraining_dataset, validation_dataset = torch.utils.data.random_split(training_dataset, [0.9, 0.1])\ntraining_dataloader   = torch.utils.data.DataLoader( dataset= training_dataset     , batch_size= 5, shuffle= True )\nvalidation_dataloader = torch.utils.data.DataLoader( dataset= validation_dataset   , batch_size= 5, shuffle= True )","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-06-27T10:12:37.906441Z","iopub.execute_input":"2024-06-27T10:12:37.907052Z","iopub.status.idle":"2024-06-27T10:13:13.832985Z","shell.execute_reply.started":"2024-06-27T10:12:37.907017Z","shell.execute_reply":"2024-06-27T10:13:13.831868Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"x, y = next(iter(training_dataloader))\nprint('Input shape:', x.shape)\nprint('Labels:', y)\nplt.imshow(torchvision.utils.make_grid(x)[0], cmap=\"Greys\"); # TODO: Change to colored image\n","metadata":{"execution":{"iopub.status.busy":"2024-06-27T10:13:27.934468Z","iopub.execute_input":"2024-06-27T10:13:27.934879Z","iopub.status.idle":"2024-06-27T10:13:28.274429Z","shell.execute_reply.started":"2024-06-27T10:13:27.934845Z","shell.execute_reply":"2024-06-27T10:13:28.273421Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.imshow(torchvision.transforms.ToPILImage()(x[0]), cmap=\"Greys\")","metadata":{"execution":{"iopub.status.busy":"2024-06-27T10:13:32.053088Z","iopub.execute_input":"2024-06-27T10:13:32.053505Z","iopub.status.idle":"2024-06-27T10:13:32.421115Z","shell.execute_reply.started":"2024-06-27T10:13:32.053474Z","shell.execute_reply":"2024-06-27T10:13:32.420090Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 1.5 Adding Noise to existing Data","metadata":{}},{"cell_type":"code","source":"def add_noise(input_image, amount): # AMOUNT IS A FLAG, to tell how important the noise is\n    \"\"\"Corrupt the input `x` by mixing it with noise according to `amount`\"\"\"\n    noise       = torch.rand_like(input_image)\n    amount      = amount.view(-1, 1, 1, 1) # Sort shape so broadcasting works\n    noisy_image = input_image * (1-amount) + (noise * amount)\n    return noisy_image\n\ndef plot_progressive_noise_addition(x): # Plotting the input data\n    # Plotting the input data\n    fig, axs = plt.subplots(2, 1, figsize=(12, 5))\n    axs[0].set_title('Input data')\n    axs[0].imshow(torchvision.utils.make_grid(x)[0], cmap='Greys')\n\n    # Adding noise\n    # AMOUNT between 0 to 1\n    amount = torch.linspace(0, 1, x.shape[0]) # Left to right -> more corruption\n    noised_x = add_noise(x, amount)\n\n    # Plotting the noised version\n    axs[1].set_title('Corrupted data (-- amount increases -->)')\n    axs[1].imshow(torchvision.utils.make_grid(noised_x)[0]);\n    \nplot_progressive_noise_addition(x)","metadata":{"execution":{"iopub.status.busy":"2024-06-27T10:14:00.190789Z","iopub.execute_input":"2024-06-27T10:14:00.191777Z","iopub.status.idle":"2024-06-27T10:14:00.947494Z","shell.execute_reply.started":"2024-06-27T10:14:00.191744Z","shell.execute_reply":"2024-06-27T10:14:00.946456Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 2. Simple Model","metadata":{}},{"cell_type":"code","source":"class BasicUNet(torch.nn.Module):\n    \"\"\"A minimal UNet implementation.\"\"\"\n    def __init__(self, in_channels=1, out_channels=1):\n        super().__init__()\n        self.down_layers = torch.nn.ModuleList([ \n            torch.nn.Conv2d(in_channels, 32, kernel_size=5, padding=2),\n            torch.nn.Conv2d(32, 64, kernel_size=5, padding=2),\n            torch.nn.Conv2d(64, 64, kernel_size=5, padding=2),\n        ])\n        self.up_layers = torch.nn.ModuleList([\n            torch.nn.Conv2d(64, 64, kernel_size=5, padding=2),\n            torch.nn.Conv2d(64, 32, kernel_size=5, padding=2),\n            torch.nn.Conv2d(32, out_channels, kernel_size=5, padding=2), \n        ])\n        self.act = nn.SiLU() # The activation function\n        self.downscale = nn.MaxPool2d(2)\n        self.upscale = nn.Upsample(scale_factor=2)\n\n    def forward(self, x):\n        h = []\n        for layer_number, layer in enumerate(self.down_layers):\n            x = self.act(layer(x)) # Through the layer and the activation function\n            if layer_number < 2: # For all but the third (final) down layer:\n              h.append(x) # Storing output for skip connection\n              x = self.downscale(x) # Downscale ready for the next layer\n              \n        for layer_number, layer in enumerate(self.up_layers):\n            if layer_number > 0: # For all except the first up layer\n              x = self.upscale(x) # Upscale\n              x += h.pop() # Fetching stored output (skip connection)\n            x = self.act(layer(x)) # Through the layer and the activation function\n            \n        return x\n    \nnet = BasicUNet(in_channels = 1, out_channels = 1)\nx   = torch.rand(5, 1, 224, 224)\nnet(x).shape","metadata":{"execution":{"iopub.status.busy":"2024-06-27T10:14:23.791881Z","iopub.execute_input":"2024-06-27T10:14:23.792290Z","iopub.status.idle":"2024-06-27T10:14:24.240326Z","shell.execute_reply.started":"2024-06-27T10:14:23.792257Z","shell.execute_reply":"2024-06-27T10:14:24.239403Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 3. Simple Model Training Pipeline","metadata":{}},{"cell_type":"code","source":"def TRAIN_MODEL(model, training_dataloader, validation_dataloader= None):\n    model.train(mode=True)\n    LOSS_FUNC  = nn.MSELoss()\n    OPTIMIZER  = torch.optim.Adam(net.parameters(), lr=1e-3)\n    for epoch in range(n_epochs):\n        for input_image, y_actual in (pbar := tqdm(training_dataloader)):\n            # Get some data and prepare the corrupted version\n            # INPUT_IMAGE -> NOISY_INPUT_IMAGE\n            input_image        = input_image.to(device) # Data on the GPU\n            noise_amount       = torch.rand(input_image.shape[0]).to(device) # Pick random noise amounts\n            input_image_noisy  = add_noise(input_image, noise_amount) # Create our noisy x\n\n            # Get the model prediction\n            predicted_denoised_image = model(input_image_noisy)\n\n            # Calculate the loss\n            loss = LOSS_FUNC(predicted_denoised_image, input_image) # How close is the output to the true 'clean' x?\n\n            # Backprop and update the params:\n            OPTIMIZER.zero_grad()\n            loss.backward()\n            OPTIMIZER.step()\n\n            # Store the loss for later\n            losses.append(loss.item())\n            pbar.set_description(f'loss={loss.item():4f}')\n\n        # Print our the average of the loss values for this epoch:\n        avg_loss = sum(losses[-len(training_dataloader):])/len(training_dataloader)\n        print(f'Finished epoch {epoch}. Average loss for this epoch: {avg_loss:05f}')\n    \n    # View the loss curve\n    plt.plot(losses)\n    plt.ylim(0, 0.1);\n    ","metadata":{"execution":{"iopub.status.busy":"2024-06-27T10:14:32.502343Z","iopub.execute_input":"2024-06-27T10:14:32.503242Z","iopub.status.idle":"2024-06-27T10:14:32.513103Z","shell.execute_reply.started":"2024-06-27T10:14:32.503204Z","shell.execute_reply":"2024-06-27T10:14:32.511979Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\n# Create the network\nnet = BasicUNet(in_channels = 1, out_channels = 1).to(device)\n \nlosses  = []\nn_epochs = 3\nTRAIN_MODEL(net, training_dataloader)","metadata":{"execution":{"iopub.status.busy":"2024-06-27T10:14:40.717376Z","iopub.execute_input":"2024-06-27T10:14:40.717799Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}