Shadow_Removal_Using_Diffusion.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "machine_shape": "hm"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "metadata": {
        "id": "3V8jhqa-1W-r"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%cd /content/drive/MyDrive/CV/Shadowremoval"
      ],
      "metadata": {
        "id": "cF3Soefa3InF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# %cd /content/drive/MyDrive/CV/Shadowremoval/ISTD_Dataset/train/train_A/\n",
        "# !ls\n",
        "# %cd /content/drive/MyDrive/CV/Shadowremoval/ISTD_Dataset/train/train_C_fixed_ours\n",
        "# !ls\n",
        "# %cd /content/drive/MyDrive/CV/Shadowremoval/ISTD_Dataset/test/test_A\n",
        "# !ls\n",
        "# %cd /content/drive/MyDrive/CV/Shadowremoval/ISTD_Dataset/test/test_C_fixed_official\n",
        "# !ls"
      ],
      "metadata": {
        "id": "WMDvW_cX4ZAt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# !mkdir images\n",
        "# !mv /content/drive/MyDrive/CV/Shadowremoval/ISTD_Dataset/test/test_A/*.png /content/drive/MyDrive/CV/Shadowremoval/ISTD_Dataset/test/test_A/images/\n",
        "# !mv /content/drive/MyDrive/CV/Shadowremoval/ISTD_Dataset/train/train_A/*.png /content/drive/MyDrive/CV/Shadowremoval/ISTD_Dataset/train/train_A/images/\n",
        "# !mv /content/drive/MyDrive/CV/Shadowremoval/ISTD_Dataset/train/train_C_fixed_ours/*.png /content/drive/MyDrive/CV/Shadowremoval/ISTD_Dataset/train/train_C_fixed_ours/images/\n",
        "# !mv /content/drive/MyDrive/CV/Shadowremoval/ISTD_Dataset/test/test_C_fixed_official/*.png /content/drive/MyDrive/CV/Shadowremoval/ISTD_Dataset/test/test_C_fixed_official/images/"
      ],
      "metadata": {
        "id": "bet4L3bA7-gn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torchvision\n",
        "import matplotlib.pyplot as plt\n",
        "from torchvision import datasets, transforms\n",
        "from torch.utils.data import DataLoader\n",
        "import numpy as np\n",
        "import torch.nn.functional as F"
      ],
      "metadata": {
        "id": "_HVfeg6f2mfF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "PATH = {\n",
        "    'train_shadow' : './ISTD_Dataset/train/train_A', \n",
        "    'train_out' : './ISTD_Dataset/train/train_C_fixed_ours', \n",
        "    'test_shadow' : './ISTD_Dataset/test', \n",
        "    'test_out' : './ISTD_Dataset/test/test_C_fixed_official', \n",
        "    }\n",
        "OUT_PATH = './output/'\n",
        "IMG_SIZE = 64"
      ],
      "metadata": {
        "id": "-HckygY42qXM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "data_transform = transforms.Compose([\n",
        "        transforms.Resize((IMG_SIZE, IMG_SIZE)),\n",
        "        transforms.ToTensor(), # Scales data into [0,1] \n",
        "        transforms.Lambda(lambda t: (t * 2) - 1) # Scale between [-1, 1] \n",
        "    ])"
      ],
      "metadata": {
        "id": "mE1tBM1r3xTD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "dataset_train_shadow = datasets.ImageFolder(PATH['train_shadow'], transform=data_transform)\n",
        "dataset_train_out = datasets.ImageFolder(PATH['train_out'], transform=data_transform)\n",
        "dataset_test_shadow = datasets.ImageFolder(PATH['test_shadow'], transform=data_transform)\n",
        "dataset_test_out = datasets.ImageFolder(PATH['test_out'], transform=data_transform)"
      ],
      "metadata": {
        "id": "fP-0ozKP2p-I"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "reverse_transform = transforms.Compose([\n",
        "        transforms.Lambda(lambda t: (t + 1) / 2),\n",
        "        transforms.Lambda(lambda t: t.permute(1, 2, 0)), # CHW to HWC\n",
        "        transforms.Lambda(lambda t: t * 255.),\n",
        "        transforms.Lambda(lambda t: t.numpy().astype(np.uint8)),\n",
        "        transforms.ToPILImage(),\n",
        "    ])"
      ],
      "metadata": {
        "id": "Glgh26BL7I7Z"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "INDEX = 0\n",
        "plt.figure(figsize=(10,10)) \n",
        "plt.subplot(2,2,1)\n",
        "plt.imshow(reverse_transform(dataset_train_shadow[INDEX][0]))\n",
        "plt.subplot(2,2,2)\n",
        "plt.imshow(reverse_transform(dataset_train_out[INDEX][0]))\n",
        "plt.subplot(2,2,3)\n",
        "plt.imshow(reverse_transform(dataset_test_shadow[INDEX][0]))\n",
        "plt.subplot(2,2,4)\n",
        "plt.imshow(reverse_transform(dataset_test_out[INDEX][0]))"
      ],
      "metadata": {
        "id": "vcddz56m2pSR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(dataset_test_out[0][0].shape)"
      ],
      "metadata": {
        "id": "8u_bkApApibW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch.nn.functional as F\n",
        "\n",
        "def linear_beta_schedule(timesteps, start=0.0001, end=0.02):\n",
        "    return torch.linspace(start, end, timesteps)\n",
        "\n",
        "def get_index_from_list(vals, t, x_shape):\n",
        "    \"\"\" \n",
        "    Returns a specific index t of a passed list of values vals\n",
        "    while considering the batch dimension.\n",
        "    \"\"\"\n",
        "    batch_size = t.shape[0]\n",
        "    out = vals.gather(-1, t.cpu())\n",
        "    return out.reshape(batch_size, *((1,) * (len(x_shape) - 1))).to(t.device)\n",
        "\n",
        "def forward_diffusion_sample(x_0, t, device=\"cpu\"):\n",
        "    \"\"\" \n",
        "    Takes an image and a timestep as input and \n",
        "    returns the noisy version of it\n",
        "    \"\"\"\n",
        "    noise = torch.randn_like(x_0)\n",
        "    sqrt_alphas_cumprod_t = get_index_from_list(sqrt_alphas_cumprod, t, x_0.shape)\n",
        "    sqrt_one_minus_alphas_cumprod_t = get_index_from_list(\n",
        "        sqrt_one_minus_alphas_cumprod, t, x_0.shape\n",
        "    )\n",
        "    # mean + variance\n",
        "    # print(sqrt_alphas_cumprod_t.shape)\n",
        "    # print(x_0.shape)\n",
        "    # print(sqrt_one_minus_alphas_cumprod_t.shape)\n",
        "    # print(noise.shape)\n",
        "    return sqrt_alphas_cumprod_t.to(device) * x_0.to(device) \\\n",
        "    + sqrt_one_minus_alphas_cumprod_t.to(device) * noise.to(device), noise.to(device)\n",
        "\n",
        "\n",
        "# Define beta schedule\n",
        "T = 300\n",
        "betas = linear_beta_schedule(timesteps=T)\n",
        "\n",
        "# Pre-calculate different terms for closed form\n",
        "alphas = 1. - betas\n",
        "alphas_cumprod = torch.cumprod(alphas, axis=0)\n",
        "alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value=1.0)\n",
        "sqrt_recip_alphas = torch.sqrt(1.0 / alphas)\n",
        "sqrt_alphas_cumprod = torch.sqrt(alphas_cumprod)\n",
        "sqrt_one_minus_alphas_cumprod = torch.sqrt(1. - alphas_cumprod)\n",
        "posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod)"
      ],
      "metadata": {
        "id": "qWw50ui9IZ5q"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# dataset = datasets.ImageFolder('path/to/data', transform=transform)"
      ],
      "metadata": {
        "id": "ho2_rRzW1Ujk"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from torchvision import transforms \n",
        "from torch.utils.data import DataLoader\n",
        "import numpy as np\n",
        "\n",
        "IMG_SIZE = 64\n",
        "BATCH_SIZE = 128\n",
        "\n",
        "# def load_transformed_dataset():\n",
        "#     data_transforms = [\n",
        "#         transforms.Resize((IMG_SIZE, IMG_SIZE)),\n",
        "#         transforms.RandomHorizontalFlip(),\n",
        "#         transforms.ToTensor(), # Scales data into [0,1] \n",
        "#         transforms.Lambda(lambda t: (t * 2) - 1) # Scale between [-1, 1] \n",
        "#     ]\n",
        "#     data_transform = transforms.Compose(data_transforms)\n",
        "\n",
        "#     train = torchvision.datasets.StanfordCars(root=\".\", download=True, \n",
        "#                                          transform=data_transform)\n",
        "\n",
        "#     test = torchvision.datasets.StanfordCars(root=\".\", download=True, \n",
        "#                                          transform=data_transform, split='test')\n",
        "#     return torch.utils.data.ConcatDataset([train, test])\n",
        "def show_tensor_image(image):\n",
        "    reverse_transforms = transforms.Compose([\n",
        "        transforms.Lambda(lambda t: (t + 1) / 2),\n",
        "        transforms.Lambda(lambda t: t.permute(1, 2, 0)), # CHW to HWC\n",
        "        transforms.Lambda(lambda t: t * 255.),\n",
        "        transforms.Lambda(lambda t: t.numpy().astype(np.uint8)),\n",
        "        transforms.ToPILImage(),\n",
        "    ])\n",
        "\n",
        "    # Take first image of batch\n",
        "    if len(image.shape) == 4:\n",
        "        image = image[0, :, :, :] \n",
        "    plt.imshow(reverse_transforms(image))\n",
        "\n",
        "# data = load_transformed_dataset()\n",
        "dataloader_train_shadow = DataLoader(dataset_train_shadow, batch_size=BATCH_SIZE, drop_last=True)\n",
        "dataloader_train_out = DataLoader(dataset_train_out, batch_size=BATCH_SIZE, drop_last=True)\n",
        "dataloader_test_shadow = DataLoader(dataset_test_shadow, batch_size=BATCH_SIZE, drop_last=True)\n",
        "dataloader_test_out = DataLoader(dataset_test_out, batch_size=BATCH_SIZE, drop_last=True)"
      ],
      "metadata": {
        "id": "uuckjpW_k1LN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Simulate forward diffusion\n",
        "image = next(iter(dataloader_train_shadow))[0]\n",
        "\n",
        "plt.figure(figsize=(15,15))\n",
        "plt.axis('off')\n",
        "num_images = 10\n",
        "stepsize = int(T/num_images)\n",
        "\n",
        "for idx in range(0, T, stepsize):\n",
        "    t = torch.Tensor([idx]).type(torch.int64)\n",
        "    plt.subplot(1, num_images+1, (idx//stepsize) + 1)\n",
        "    image, noise = forward_diffusion_sample(image, t)\n",
        "    show_tensor_image(image)"
      ],
      "metadata": {
        "id": "2fUPyJghdoUA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from torch import nn\n",
        "import torch as th\n",
        "import math\n",
        "\n",
        "\n",
        "class Block(nn.Module):\n",
        "    def __init__(self, in_ch, out_ch, time_emb_dim, up=False):\n",
        "        super().__init__()\n",
        "        self.time_mlp =  nn.Linear(time_emb_dim, out_ch)\n",
        "        if up:\n",
        "            self.conv1 = nn.Conv2d(2*in_ch, out_ch, 3, padding=1)\n",
        "            self.transform = nn.ConvTranspose2d(out_ch, out_ch, 4, 2, 1)\n",
        "        else:\n",
        "            self.conv1 = nn.Conv2d(in_ch, out_ch, 3, padding=1)\n",
        "            self.transform = nn.Conv2d(out_ch, out_ch, 4, 2, 1)\n",
        "        self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1)\n",
        "        self.bnorm1 = nn.BatchNorm2d(out_ch)\n",
        "        self.bnorm2 = nn.BatchNorm2d(out_ch)\n",
        "        self.relu  = nn.ReLU()\n",
        "        \n",
        "    def forward(self, x, t, ):\n",
        "        # First Conv\n",
        "        h = self.bnorm1(self.relu(self.conv1(x)))\n",
        "        # Time embedding\n",
        "        time_emb = self.relu(self.time_mlp(t))\n",
        "        # Extend last 2 dimensions\n",
        "        time_emb = time_emb[(..., ) + (None, ) * 2]\n",
        "        # Add time channel\n",
        "        h = h + time_emb\n",
        "        # Second Conv\n",
        "        h = self.bnorm2(self.relu(self.conv2(h)))\n",
        "        # Down or Upsample\n",
        "        return self.transform(h)\n",
        "\n",
        "\n",
        "class SinusoidalPositionEmbeddings(nn.Module):\n",
        "    def __init__(self, dim):\n",
        "        super().__init__()\n",
        "        self.dim = dim\n",
        "\n",
        "    def forward(self, time):\n",
        "        device = time.device\n",
        "        half_dim = self.dim // 2\n",
        "        embeddings = math.log(10000) / (half_dim - 1)\n",
        "        embeddings = torch.exp(torch.arange(half_dim, device=device) * -embeddings)\n",
        "        embeddings = time[:, None] * embeddings[None, :]\n",
        "        embeddings = torch.cat((embeddings.sin(), embeddings.cos()), dim=-1)\n",
        "        # TODO: Double check the ordering here\n",
        "        return embeddings\n",
        "\n",
        "\n",
        "class SimpleUnet(nn.Module):\n",
        "    \"\"\"\n",
        "    A simplified variant of the Unet architecture.\n",
        "    \"\"\"\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        image_channels = 3 * 2\n",
        "        down_channels = (64, 128, 256, 512, 1024)\n",
        "        up_channels = (1024, 512, 256, 128, 64)\n",
        "        out_dim = 1 \n",
        "        time_emb_dim = 32\n",
        "\n",
        "        # Time embedding\n",
        "        self.time_mlp = nn.Sequential(\n",
        "                SinusoidalPositionEmbeddings(time_emb_dim),\n",
        "                nn.Linear(time_emb_dim, time_emb_dim),\n",
        "                nn.ReLU()\n",
        "            )\n",
        "        \n",
        "        # Initial projection\n",
        "        self.conv0 = nn.Conv2d(image_channels, down_channels[0], 3, padding=1)\n",
        "\n",
        "        # Downsample\n",
        "        self.downs = nn.ModuleList([Block(down_channels[i], down_channels[i+1], \\\n",
        "                                    time_emb_dim) \\\n",
        "                    for i in range(len(down_channels)-1)])\n",
        "        # Upsample\n",
        "        self.ups = nn.ModuleList([Block(up_channels[i], up_channels[i+1], \\\n",
        "                                        time_emb_dim, up=True) \\\n",
        "                    for i in range(len(up_channels)-1)])\n",
        "\n",
        "        self.output = nn.Conv2d(up_channels[-1], 3, out_dim)\n",
        "\n",
        "    def forward(self, x, x1, timestep):\n",
        "        \n",
        "        x = th.cat((x.to(device), x1.to(device)), axis = 1).to(device)\n",
        "        # Embedd time\n",
        "        t = self.time_mlp(timestep)\n",
        "        \n",
        "\n",
        "        # Initial conv\n",
        "        x = self.conv0(x)\n",
        "        # Unet\n",
        "        residual_inputs = []\n",
        "        for down in self.downs:\n",
        "            x = down(x, t)\n",
        "            residual_inputs.append(x)\n",
        "        for up in self.ups:\n",
        "            residual_x = residual_inputs.pop()\n",
        "            # Add residual x as additional channels\n",
        "            x = torch.cat((x, residual_x), dim=1)           \n",
        "            x = up(x, t)\n",
        "        return self.output(x)\n",
        "\n",
        "model = SimpleUnet()\n",
        "print(\"Num params: \", sum(p.numel() for p in model.parameters()))\n",
        "model"
      ],
      "metadata": {
        "id": "KOYPSxPf_LL7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class ChromaticityConsistencyLoss(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(ChromaticityConsistencyLoss, self).__init__()\n",
        "\n",
        "    def forward(self, input_image, output_image):\n",
        "        # Convert images to chromaticity coordinates\n",
        "        input_chromaticity = input_image[0:3] / (input_image[0] + input_image[1] + input_image[2] + 1e-8)\n",
        "        output_chromaticity = output_image[0:3] / (output_image[0] + output_image[1] + output_image[2] + 1e-8)\n",
        "\n",
        "        # Calculate L2 loss between chromaticity coordinates of input and output images\n",
        "        chromaticity_consistency_loss = torch.mean((input_chromaticity - output_chromaticity) ** 2)\n",
        "\n",
        "        return chromaticity_consistency_loss\n",
        "\n",
        "\n",
        "# Create instance of ChromaticityConsistencyLoss\n",
        "chromaticity_consistency_loss = ChromaticityConsistencyLoss()\n",
        "\n",
        "# Calculate chromaticity consistency loss\n",
        "# loss = chromaticity_consistency_loss(input_image, output_image)\n",
        "\n",
        "\n",
        "class StructurePreservationLoss(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(StructurePreservationLoss, self).__init__()\n",
        "    \n",
        "    def forward(self, input_image, output_image):\n",
        "        # Calculate L1 loss between input and output images\n",
        "        structure_preservation_loss = torch.mean(torch.abs(input_image - output_image))\n",
        "\n",
        "        return structure_preservation_loss\n",
        "\n",
        "# Create instance of StructurePreservationLoss\n",
        "structure_preservation_loss = StructurePreservationLoss()\n",
        "\n",
        "# Calculate structure preservation loss\n",
        "# loss = structure_preservation_loss(input_image, output_image)"
      ],
      "metadata": {
        "id": "cetgkEXXNtO9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def get_loss(model, x_0, x_masked, t):\n",
        "    x_noisy, noise = forward_diffusion_sample(x_0, t, device)\n",
        "    noise_pred = model(x_noisy, x_masked, t)\n",
        "    return F.mse_loss(noise, noise_pred)"
      ],
      "metadata": {
        "id": "Ed12NNXPtDon"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "@torch.no_grad()\n",
        "def sample_timestep(x, t):\n",
        "    \"\"\"\n",
        "    Calls the model to predict the noise in the image and returns \n",
        "    the denoised image. \n",
        "    Applies noise to this image, if we are not in the last step yet.\n",
        "    \"\"\"\n",
        "    betas_t = get_index_from_list(betas, t, x.shape)\n",
        "    sqrt_one_minus_alphas_cumprod_t = get_index_from_list(\n",
        "        sqrt_one_minus_alphas_cumprod, t, x.shape\n",
        "    )\n",
        "    sqrt_recip_alphas_t = get_index_from_list(sqrt_recip_alphas, t, x.shape)\n",
        "    \n",
        "    # Call model (current image - noise prediction)\n",
        "    model_mean = sqrt_recip_alphas_t * (\n",
        "        x - betas_t * model(x, x, t) / sqrt_one_minus_alphas_cumprod_t\n",
        "    )\n",
        "    posterior_variance_t = get_index_from_list(posterior_variance, t, x.shape)\n",
        "    \n",
        "    if t == 0:\n",
        "        return model_mean\n",
        "    else:\n",
        "        noise = torch.randn_like(x)\n",
        "        return model_mean + torch.sqrt(posterior_variance_t) * noise \n",
        "\n",
        "@torch.no_grad()\n",
        "def sample_plot_image(fname):\n",
        "    # Sample noise\n",
        "    img_size = IMG_SIZE\n",
        "    img = torch.randn((1, 3, img_size, img_size), device=device)\n",
        "    plt.figure(figsize=(15,15))\n",
        "    plt.axis('off')\n",
        "    num_images = 10\n",
        "    stepsize = int(T/num_images)\n",
        "\n",
        "    for i in range(0,T)[::-1]:\n",
        "        t = torch.full((1,), i, device=device, dtype=torch.long)\n",
        "        img = sample_timestep(img, t)\n",
        "        if i % stepsize == 0:\n",
        "            plt.subplot(1, num_images, i//stepsize+1)\n",
        "            show_tensor_image(img.detach().cpu())\n",
        "    plt.savefig(fname)\n",
        "    plt.show()            "
      ],
      "metadata": {
        "id": "k13hj2mciCHA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from torch.optim import Adam\n",
        "\n",
        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "model.to(device)\n",
        "optimizer = Adam(model.parameters(), lr=0.001)\n",
        "epochs = 20 # Try more!\n",
        "\n",
        "for epoch in range(epochs):\n",
        "    for ((step, batch), (step_s, batch_s)) in zip(enumerate(dataloader_train_out), enumerate(dataloader_train_shadow)) :\n",
        "      optimizer.zero_grad()\n",
        "\n",
        "      t = torch.randint(0, T, (BATCH_SIZE,), device=device).long()\n",
        "      loss = get_loss(model, batch[0], batch_s[0], t)\n",
        "\n",
        "      loss.backward()\n",
        "      optimizer.step()\n",
        "\n",
        "      if epoch % 1 == 0 and step == 0:\n",
        "        print(f\"Epoch {epoch} | step {step:03d} Loss: {loss.item()} \")\n",
        "        fname = OUT_PATH + \"Epoch\" + str(epoch) + \".jpg\"\n",
        "        sample_plot_image(fname)"
      ],
      "metadata": {
        "id": "bpN_LKYwuLx0"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "torch.save(model.state_dict(), './model.pth')"
      ],
      "metadata": {
        "id": "QEwOXsxCdJtU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model.load_state_dict(torch.load('./model.pth'))"
      ],
      "metadata": {
        "id": "60MwJBx-BShG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Simulate forward diffusion\n",
        "image, shadow = next(iter(dataloader_train_shadow))[0], next(iter(dataloader_train_shadow))[0]\n",
        "\n",
        "plt.figure(figsize=(15,15))\n",
        "plt.axis('off')\n",
        "num_images = 10\n",
        "stepsize = int(T/num_images)\n",
        "\n",
        "for idx in range(0, T, stepsize):\n",
        "    t = torch.Tensor([idx]).type(torch.int64)\n",
        "    plt.subplot(1, num_images+1, (idx//stepsize) + 1)\n",
        "    image, noise = forward_diffusion_sample(image, t)\n",
        "    show_tensor_image(image)"
      ],
      "metadata": {
        "id": "j9C_4MY2XrHV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(image.shape)"
      ],
      "metadata": {
        "id": "1tWXNI6VSboE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "t = torch.randint(0, 5000, (BATCH_SIZE,), device=device).long()\n",
        "show_tensor_image(model(image, shadow, t).to('cpu').detach())"
      ],
      "metadata": {
        "id": "7ft_BS3ASdBt"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}