{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Lab 6 — A U-Net for segmentation and denoising\n",
        "\n",
        "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/kleinric/cv-labs/blob/main/lab-06.ipynb)\n",
        "\n",
        "**COMS4036A / COMS7050A Computer Vision**\n",
        "\n",
        "Groups of up to three; one member submits the group's notebook (`.ipynb`) on Moodle. Every member must be able to explain every cell.\n",
        "\n",
        "Companion reading is [Chapter 6 of the course book](https://courses.ms.wits.ac.za/~richard/cv/book/chapters/06-segmentation.html). Labs 4 and 5 asked where the objects were; this one asks which pixels they occupy. You build a U-Net from the chapter's description, train it to separate pet from background on the same Oxford-IIIT Pet photographs Lab 3 classified, on a network of under two million parameters; after eight epochs the reference solution reaches a mean per-image foreground IoU of about 0.7, with trimap label 3 treated as background. Runs differ by a few hundredths; two reference runs gave 0.69 and 0.72. Then you point the same network at a different target and retrain it, changing only the input, the target, the number of output channels and the loss.\n",
        "\n",
        "This lab's rung of the skills ladder is **checkpointing and resume**: you save the model, optimiser and random-number states after every epoch, store each checkpoint as a W&B artefact, and resume training from it in a new run after deliberately stopping the first. Long runs die: the Colab session times out, the queue evicts you, someone trips over the cable. A run that cannot be resumed is a run you have to start again.\n",
        "\n",
        "## 0. GPU, W&B, and the data\n",
        "\n",
        "Open this notebook in [Colab](https://colab.research.google.com) with the badge above, or [download it from the book](https://courses.ms.wits.ac.za/~richard/cv/book/labs/lab-06.ipynb) and use **File ▸ Upload notebook**; then **File ▸ Save a copy in Drive**. Set **Runtime ▸ Change runtime type** to a GPU, and log in to W&B; this lab logs to a project called `cv-lab6`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# Group members — fill in before submitting.\n",
        "MEMBERS = [\n",
        "    # (\"Student name\", \"Student number\"),\n",
        "]\n",
        "for name, number in MEMBERS:\n",
        "    print(f\"{number}  {name}\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "!nvidia-smi --query-gpu=name,memory.total --format=csv"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "import wandb\n",
        "wandb.login()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The data is Lab 3's download, both archives again: the images you have already classified, plus the `annotations/` directory you only read a text file from. This lab uses the trimaps in `annotations/trimaps/`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "!wget -q -nc https://thor.robots.ox.ac.uk/pets/images.tar.gz\n",
        "!wget -q -nc https://thor.robots.ox.ac.uk/pets/annotations.tar.gz\n",
        "!tar -xzf images.tar.gz && tar -xzf annotations.tar.gz\n",
        "!ls annotations/trimaps | head -3"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "import math\n",
        "import os\n",
        "import random\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from PIL import Image\n",
        "from torch.utils.data import DataLoader, Dataset, Subset\n",
        "from torchvision import transforms\n",
        "from torchvision.transforms import functional as TF\n",
        "\n",
        "DEV = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "SIZE = 128\n",
        "\n",
        "\n",
        "def seed_everything(seed=0):\n",
        "    random.seed(seed)\n",
        "    np.random.seed(seed)\n",
        "    torch.manual_seed(seed)\n",
        "    torch.cuda.manual_seed_all(seed)\n",
        "\n",
        "\n",
        "seed_everything(0)\n",
        "print(\"device:\", DEV)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 1. Data and paired image-mask transforms\n",
        "\n",
        "`annotations/trimaps/Abyssinian_100.png` is a single-channel PNG the same size as `images/Abyssinian_100.jpg`, and every pixel of it holds one of three integers: **1** for the animal, **2** for the background, and **3** for ambiguous or unclassified regions, mostly a band along the animal's outline but also accessories such as collars. There are 7,390 of them, one per photograph in `images/`, and they were drawn by hand (Parkhi et al., 2012).\n",
        "\n",
        "The target this lab predicts is binary: pet or not-pet. That leaves a decision about label 3, which you make once, write down, and keep.\n",
        "\n",
        "The photograph and its mask go through one pipeline together, and it has to treat them consistently. Whatever geometry you apply to the photograph, a crop or a flip or a rotation, the mask gets the *same* geometry, chosen once per item and not sampled twice. Resampling is where the two part company: a photograph interpolates bilinearly, and a mask must not. An average of label 1 and label 3 is label 2, which gives the pixel a class it does not have.\n",
        "\n",
        "**Q1.1.** Write `PetSegDataset(split, size=128, train=False)`, which is Lab 3's `PetsDataset` with the trimap attached. `__getitem__(i)` opens photograph $i$ and its trimap and returns `(image, mask, trimap)`: the image as a float tensor in $[0,1]$ (`TF.to_tensor`, without normalisation), the binary pet mask as a float tensor of shape $(1, s, s)$, and the raw three-label trimap as a `long` tensor of shape $(s, s)$ for Q1.2 and Q1.4 to count.\n",
        "\n",
        "- when `train=True`: resize both to `int(1.15 * s)` pixels in each dimension, take one `transforms.RandomCrop.get_params` window and apply it to both with `TF.crop`, then flip both or neither;\n",
        "- when `train=False`: resize both to $s \\times s$ and stop, so evaluation is deterministic;\n",
        "- images resize with `TF.InterpolationMode.BILINEAR`, trimaps with `NEAREST`.\n",
        "\n",
        "Build `train_set` (from `trainval`, `train=True`), `test_set` (from `test`, `train=False`), a training loader with `batch_size=32, shuffle=True, num_workers=2, drop_last=True`, and a test loader with `batch_size=32`. Print the size of each split."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "class PetSegDataset(Dataset):\n",
        "    # YOUR CODE HERE\n",
        "    ..."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "**Q1.2.** Display eight training pairs as two rows: the photograph, and the photograph with its mask drawn over it (`plt.contour(mask, [0.5])` or a translucent overlay; anything you can read). Run the cell twice so you see two different crops of the same eight animals. Then report, over 200 training images, the fraction of pixels carrying each of the three trimap labels. Does every mask sit on its animal?"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# YOUR CODE HERE"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "*Answer:*\n",
        "\n",
        "**Q1.3.** Break it on purpose. Copy your `__getitem__` into a throwaway function that draws the crop window for the image and the crop window for the mask from two *separate* calls to `get_params`, and display eight pairs from it. Then say what you see, and more usefully what a training run on this dataset would look like from the outside, given that nothing here raises an exception and the loss still falls."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# YOUR CODE HERE"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "*Answer:*\n",
        "\n",
        "**Q1.4.** Resample one trimap to $128 \\times 128$ twice, once with `NEAREST` and once with `BILINEAR`, and report the percentage of pixels that end up with a different label. Separately, take the *binary* pet mask as a float tensor at full resolution, resize it bilinearly to $128 \\times 128$, and print how many distinct values it holds. Then answer: what should a mask be resampled with, and what did you do with trimap label 3?"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# YOUR CODE HERE"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "*Answer:*\n",
        "\n",
        "## 2. Pixel accuracy and IoU\n",
        "\n",
        "Two numbers describe a predicted binary mask against its target. **Pixel accuracy** is the fraction of pixels whose predicted label matches. **Intersection over union** is $|P \\cap T| / |P \\cup T|$ over the foreground: the pixels both call pet, divided by the pixels either calls pet.\n",
        "\n",
        "Throughout this lab, \"mean IoU\" means foreground IoU computed per image and then averaged over images. This differs from the benchmark mIoU of Chapter 6, which accumulates intersections and unions for each class across the whole dataset and then averages over classes.\n",
        "\n",
        "Implement both metrics and evaluate a predictor that outputs background everywhere before you evaluate your own network.\n",
        "\n",
        "**Q2.1.** Implement `pixel_accuracy(pred, target)` and `iou(pred, target)` for batches of shape $(N, 1, H, W)$ holding zeros and ones, each returning a tensor of $N$ per-image values. Check both on a case you compute first on paper: two $4 \\times 4$ masks, one covering the top two rows and the other the middle two rows."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "def pixel_accuracy(pred, target):\n",
        "    # YOUR CODE HERE\n",
        "    ...\n",
        "\n",
        "\n",
        "def iou(pred, target, eps=1e-7):\n",
        "    # YOUR CODE HERE\n",
        "    ..."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "**Q2.2.** Run the predictor that outputs background everywhere over the whole test split, and report its mean pixel accuracy and mean IoU. Report the mean fraction of pixels that are pet as well. What is pixel accuracy worth as the headline number for this task?"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# YOUR CODE HERE"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "*Answer:*\n",
        "\n",
        "**Q2.3.** The union in the denominator can be zero. Count the test images whose mask holds no pet pixel at all at 128 px, display them next to their trimaps, and check whether their *full-resolution* trimaps have any pet pixels either. A mask can be empty because the animal is small and resampling lost it, or because nobody labelled it. Report what your `iou` returns for such an image; the `eps` above is a convention, and it is making a claim."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# YOUR CODE HERE"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "*Answer:*\n",
        "\n",
        "## 3. The U-Net\n",
        "\n",
        "This is the U-Net of Chapter 6 at lab scale, with *same* padding rather than the original paper's unpadded convolutions, so that every map keeps its size through each block and the skips need no cropping. An encoder of three blocks at 32, 64 and 128 channels, each a `DoubleConv` (two $3 \\times 3$ convolutions, each followed by batch norm and ReLU). A bottleneck `DoubleConv` at 256 channels, and a $2 \\times 2$ max-pool in front of every block after the first, the bottleneck included. A decoder of three blocks that mirror the encoder: a $2 \\times 2$ transposed convolution that doubles the resolution and halves the channels, then a **concatenation** of the encoder tensor at that resolution, then a `DoubleConv` back down to the encoder's channel count. Finally a $1 \\times 1$ convolution to the output channels.\n",
        "\n",
        "Use `padding=1` and `bias=False` for both convolutions in each `DoubleConv`, and `kernel_size=2, stride=2` for each transposed convolution, keeping the default biases there and in the final convolution. Return the final convolution's raw outputs: `F.binary_cross_entropy_with_logits` applies the sigmoid itself.\n",
        "\n",
        "The concatenation is also where the most common shape error occurs: the `DoubleConv` after `up3` receives 256 input channels, not 128, because 128 come up the decoder and 128 come across the skip.\n",
        "\n",
        "**Q3.1.** Implement `DoubleConv` and `UNet(in_ch=3, out_ch=1, base=32)`. Give `forward` a `trace=False` argument that, when true, prints the shape of every encoder, bottleneck and decoder tensor, so that you can check the intermediate shapes."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "class DoubleConv(nn.Module):\n",
        "    # YOUR CODE HERE\n",
        "    ...\n",
        "\n",
        "\n",
        "class UNet(nn.Module):\n",
        "    # YOUR CODE HERE\n",
        "    ..."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "**Q3.2.** Before you run it: write down the shape you expect at each of the seven traced tensors for a $1 \\times 3 \\times 128 \\times 128$ input, and compute the bottleneck block's parameter count by hand ($C_{\\text{out}} \\times C_{\\text{in}} \\times 9$ per convolution, no bias because batch norm follows, plus $2C$ per batch-norm layer). Then run the trace and `sum(p.numel() ...)` and check both. Report the total."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# YOUR CODE HERE"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "*Answer:*\n",
        "\n",
        "**Q3.3.** The ritual, from Lab 2 and expected before every real run since: take one batch of 16 training pairs and drive a fresh seeded `UNet` to memorise it with Adam at $10^{-3}$ for 250 steps, `F.binary_cross_entropy_with_logits` against the mask. Log it to W&B under a name that makes clear it is the ritual. Print the loss and the batch IoU every 50 steps, and report the first step at which the batch IoU passes 0.99."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# YOUR CODE HERE"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 4. Training, checkpointing and resume\n",
        "\n",
        "A checkpoint is not the weights. Weights alone will restore what the network predicts and nothing about how it was moving: Adam carries a first and second moment for every parameter, built up over every step so far, and a run that resumes without them starts with neither. Save, every epoch: the model's `state_dict`, the optimiser's `state_dict`, the epoch number, and the random-number state. Then `wandb.log_artifact` it, so the file is attached to the run rather than to a virtual machine that is about to disappear. (W&B's own API calls the object an `Artifact`.)\n",
        "\n",
        "Logging a checkpoint as a W&B **artefact** takes three lines:\n",
        "\n",
        "```python\n",
        "art = wandb.Artifact(f\"unet-seg-{wandb.run.id}\", type=\"model\")\n",
        "art.add_file(\"ckpt.pt\")\n",
        "wandb.log_artifact(art, aliases=[f\"epoch-{epoch}\"])\n",
        "```\n",
        "\n",
        "and getting it back in another run, in another session or on another machine, takes two:\n",
        "\n",
        "```python\n",
        "path = wandb.run.use_artifact(f\"unet-seg-{run_id}:epoch-4\", type=\"model\").download()\n",
        "ckpt = torch.load(os.path.join(path, \"ckpt.pt\"), weights_only=False)\n",
        "```\n",
        "\n",
        "**Q4.1.** Write `save_ckpt(path, model, opt, epoch)` and `load_ckpt(path, model, opt=None)`. `save_ckpt` stores the four things named above; `load_ckpt` restores the model, restores the optimiser when one is passed, restores the random-number state, and returns the epoch. Write `evaluate(model, loader)` returning mean pixel accuracy and mean IoU over a loader (in evaluation mode, without gradients), and `run_epoch(model, opt, loader)` returning the mean training loss and the list of per-step losses."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "def save_ckpt(path, model, opt, epoch):\n",
        "    # YOUR CODE HERE\n",
        "    ...\n",
        "\n",
        "\n",
        "def load_ckpt(path, model, opt=None):\n",
        "    # YOUR CODE HERE\n",
        "    ...\n",
        "\n",
        "\n",
        "def evaluate(model, loader):\n",
        "    # YOUR CODE HERE\n",
        "    ...\n",
        "\n",
        "\n",
        "def run_epoch(model, opt, loader):\n",
        "    # YOUR CODE HERE\n",
        "    ..."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "**Q4.2.** Seed, build a fresh `UNet`, and train for 8 epochs with Adam at $10^{-3}$ and batch size 32. Every epoch: evaluate on the test split, log `train/loss`, `test/acc` and `test/iou` to W&B, save a checkpoint, and log it as an artefact with an `epoch-N` alias. This lab has no separate validation split, so these curves monitor the test split for illustration only: keep the prescribed settings fixed, and do not use the curves to choose them. The config must carry the architecture's base width, the image size, the optimiser and learning rate, the batch size, the epoch count, the parameter count and the seed. Report the final pixel accuracy and IoU beside Q2.2's trivial predictor, and plot both against epoch."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# YOUR CODE HERE"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "**Q4.3.** Now the rung. Pretend the run above died after epoch 4. Start a new W&B run for each resume, build a *fresh* `UNet` and a *fresh* Adam, pull the **epoch-4** artefact from Q4.2's run, load it, and train two more epochs, epochs 5 and 6, which Q4.2 also ran, so you have something to compare against. Do the whole thing twice: once passing the optimiser to `load_ckpt`, once not passing it. Three things to report.\n",
        "\n",
        "1. The pixel accuracy and IoU straight after loading, before any training. Compare them against what Q4.2 logged at epoch 4.\n",
        "2. The per-step training loss over the first 20 steps of each resume, plotted together. Both resumes see the same batches in the same order, because the random-number state came out of the checkpoint too, so the two curves are comparable step by step. If you evaluate straight after loading, save the random-number state before that evaluation and restore it afterwards, so the check does not shift the batch order. The same order does not guarantee bit-identical results on a GPU.\n",
        "3. The mean loss over those 20 steps, and the mean loss over each of the two resumed epochs.\n",
        "\n",
        "Then answer: what did the optimiser state buy, how long did the cost of dropping it last, and which of the three reported quantities would have caught the mistake if you had not been looking for it?"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# YOUR CODE HERE"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "*Answer:*\n",
        "\n",
        "**Q4.4.** Nothing in your `UNet` names a resolution. Build test loaders at 64, 128 and 256 px, the same photographs and the same trimaps resized differently, and evaluate the Q4.2 weights on all three, unchanged. Report accuracy and IoU at each. Then say what limits the network at 64 and what limits it at 256, given that the convolutions have not moved."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# YOUR CODE HERE"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "*Answer:*\n",
        "\n",
        "## 5. Comparison with a pretrained Mask R-CNN\n",
        "\n",
        "`maskrcnn_resnet50_fpn_v2` is torchvision's instance segmenter, trained on COCO: 118,000 images, 80 classes, of which two are `cat` (label 17) and `dog` (label 18). We use the COCO-pretrained weights without fine-tuning on Oxford-IIIT Pet. It returns, per image, a set of boxes, labels, scores and soft masks of shape $(N, 1, H, W)$."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "from torchvision.models.detection import (MaskRCNN_ResNet50_FPN_V2_Weights,\n",
        "                                          maskrcnn_resnet50_fpn_v2)\n",
        "\n",
        "mrcnn = maskrcnn_resnet50_fpn_v2(\n",
        "    weights=MaskRCNN_ResNet50_FPN_V2_Weights.COCO_V1).to(DEV).eval()\n",
        "print(f\"{sum(p.numel() for p in mrcnn.parameters()):,} parameters\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "**Q5.1.** Take a fixed seeded subset of 300 test images. Evaluate your Q4.2 network on it. Then run Mask R-CNN on the *same* 128 px tensors; it takes a list of images in $[0,1]$ and resizes internally. Turn its output into a binary pet mask: keep the detections labelled cat or dog with score above 0.5, threshold their soft masks at 0.5, and take the union. Report mean IoU for both, and the number of images in which Mask R-CNN found no cat and no dog at all."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# YOUR CODE HERE"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "**Q5.2.** Find the six images where the two disagree most, ranking by the absolute difference of their per-image IoU, and display each as a row: photograph, ground-truth mask, your mask, Mask R-CNN's mask. Describe any recurring failure mode you can see for each model, using examples from these images."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# YOUR CODE HERE"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "*Answer:*\n",
        "\n",
        "**Q5.3.** Mask R-CNN has 46 million parameters and was fitted on 118,000 annotated images of eighty everyday classes; yours has under two million and saw 3,680 photographs of cats and dogs for eight epochs. Compare training this U-Net from scratch with using Mask R-CNN without fine-tuning, on the two IoU numbers you measured. Which properties of this dataset may explain the result? Explain why this experiment does not establish whether fine-tuning a pretrained model would do better.\n",
        "\n",
        "*Answer:*\n",
        "\n",
        "## 6. Retargeting the U-Net to denoising\n",
        "\n",
        "Nothing in the U-Net knows about pets. It is a function from an image to an image of the same size, and what it computes is set by the pair you train it on and the loss you score it with. Keep the encoder, decoder and skip connections. Change the output head from one channel to three, use noisy images as inputs and clean images as targets, and replace binary cross-entropy with mean squared error.\n",
        "\n",
        "The new pair is (noisy photograph, clean photograph). Add Gaussian noise of standard deviation $\\sigma$ to an image in $[0,1]$, draw $\\sigma$ per image from $\\mathcal{U}(0.05, 0.4)$ so the network cannot specialise to one level, and clamp back to $[0,1]$. The target is the clean image. The loss is mean squared error. Quality is reported as **PSNR** (peak signal-to-noise ratio), $10 \\log_{10}(1/\\text{MSE})$ for images in $[0,1]$, in decibels, higher being better; an increase of about 3 dB corresponds to halving the mean squared error.\n",
        "\n",
        "When you evaluate, use fixed seeded noise, so that every model sees the same noisy images, and clamp the network's outputs to $[0,1]$ for evaluation only. Compute the MSE over all pixels and channels of the full evaluation split, then the PSNR from that MSE.\n",
        "\n",
        "**Q6.1.** Write `add_noise(x, sigma)` and display four training triples: clean, noisy, and the difference, at $\\sigma$ = 0.05, 0.15, 0.25 and 0.4, with the input PSNR in each title. Look at 0.4: that is what the network is being asked to undo."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# YOUR CODE HERE"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "**Q6.2.** Call `seed_everything(0)`, then build `UNet(in_ch=3, out_ch=3)` and train for 6 epochs with Adam at $10^{-3}$ on the same loader, with noise drawn fresh per batch and MSE against the clean image. Log to W&B with a complete config, checkpoint each epoch, and log the final checkpoint as an artefact; a later lab uses these weights. Report the parameter count beside Section 4's network and the test MSE at $\\sigma = 0.2$ per epoch."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# YOUR CODE HERE"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "**Q6.3.** For $\\sigma \\in \\{0.1, 0.2, 0.4\\}$ on the test split, report the PSNR of the noisy input and the PSNR of the network's output, and display four (noisy, output, clean) triples at $\\sigma = 0.2$."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# YOUR CODE HERE"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "**Q6.4.** Now $\\sigma = 0.6$, which is outside the range it trained on. Report input and output PSNR and display four triples. Compare the result with the in-range noise levels, and describe what happens to the residual noise and to the image detail."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# YOUR CODE HERE"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "*Answer:*\n",
        "\n",
        "**Q6.5.** List what changed between Section 4 and Section 6 and what did not, down to the line and including the parameter count. A later chapter returns to this network.\n",
        "\n",
        "*Answer:*\n",
        "\n",
        "## 7. W&B runs\n",
        "\n",
        "Paste links to your W&B runs below. At minimum: the one-batch ritual, the eight-epoch segmentation run with its per-epoch checkpoint artefacts, both resumes (one run each), and the second-target run. Each needs a meaningful name and a complete config: base width, image size, optimiser, learning rate, batch size, epochs, parameter count and seed. The artefact aliases must be visible on the segmentation run.\n",
        "\n",
        "*W&B run links:*\n",
        "\n",
        "## 8. Before you submit\n",
        "\n",
        "- [ ] **Runtime ▸ Restart session and run all** on a GPU runtime, then read every output. The full re-run downloads roughly 800 MB of dataset archives plus pretrained weights, trains three networks and performs two resumes; how long it takes depends on the GPU you are allocated.\n",
        "- [ ] Fill in the group members; check that every member can explain every cell.\n",
        "- [ ] Complete every *Answer:* cell, and paste the W&B links into Section 7 with the checkpoint artefacts attached.\n",
        "- [ ] **File ▸ Download ▸ Download .ipynb**; one member submits it on Moodle by the deadline shown there."
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 4
}