CM_GAN_Jan5.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "authorship_tag": "ABX9TyPcTRx4Ut5mkhdIsmwBpSfx",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/fjadidi2001/Image_Inpaint/blob/main/CM_GAN_Jan5.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The architecture of CM-GAN focuses on image inpainting, specifically designed to fill in missing or corrupted regions of images with realistic content. While\n",
        "### **1. Generator Architecture**\n",
        "The generator is responsible for creating realistic inpainted images. CM-GAN uses **cascaded modulation** to process inputs. A common flow:\n",
        "\n",
        "#### Input:\n",
        "- An **image** with missing regions (masked image).\n",
        "- A **binary mask** representing the missing areas (1 for missing, 0 for existing pixels).\n",
        "\n",
        "#### Layers:\n",
        "1. **Convolutional Layers with Mask Concatenation**:\n",
        "   - Initial layers concatenate the image with the binary mask.\n",
        "   - Convolutions extract features from the masked regions.\n",
        "\n",
        "   **Purpose**: Learn the structure and surrounding context of the image.\n",
        "\n",
        "2. **Cascaded Modulation Block**:\n",
        "   - Combines **global modulation** (to understand overall image semantics) with **spatially adaptive modulation** (to handle local details).\n",
        "   - Global modulation uses a feature map that spans the entire image.\n",
        "   - Adaptive modulation applies location-specific adjustments.\n",
        "\n",
        "   **Purpose**: Balance global coherence and local realism.\n",
        "\n",
        "3. **Feature Propagation via Attention Mechanisms**:\n",
        "   - **Enhanced Attention** to propagate contextual information from known to unknown areas.\n",
        "\n",
        "   **Purpose**: Ensures accurate filling of missing regions based on surrounding context.\n",
        "\n",
        "4. **Output Layers**:\n",
        "   - A final set of convolutions or deconvolutions reconstructs the inpainted image.\n",
        "\n",
        "   **Purpose**: Generate the final high-quality inpainted output.\n",
        "\n",
        "---\n",
        "\n",
        "### **2. Discriminator Architecture**\n",
        "The discriminator evaluates the inpainted images for realism.\n",
        "\n",
        "1. **Input**:\n",
        "   - The inpainted image (from the generator).\n",
        "   - The corresponding ground truth image (actual image without missing areas).\n",
        "\n",
        "2. **Layers**:\n",
        "   - Convolutional layers extract features.\n",
        "   - Outputs a **realism score**, indicating how realistic the inpainted image is.\n",
        "\n",
        "3. **Loss Function**:\n",
        "   - Often uses an **adversarial loss** (e.g., Wasserstein or hinge loss) to train the generator and discriminator in a competitive manner.\n",
        "\n",
        "---\n",
        "\n",
        "### **Key Components of CM-GAN**\n",
        "1. **Object-Aware Training**:\n",
        "   - Focuses on challenging regions, like objects, using annotations (e.g., panoptic segmentation).\n",
        "   - Ensures that the generator fills object regions more realistically.\n",
        "\n",
        "2. **Mask-Aware Encoding**:\n",
        "   - Explicitly considers the mask during feature extraction.\n",
        "   - Helps the generator learn to handle varied mask sizes and shapes.\n",
        "\n",
        "3. **Enhanced Attention**:\n",
        "   - Propagates information from visible areas to missing areas.\n",
        "   - Improves inpainting quality for complex patterns.\n",
        "\n",
        "---\n",
        "\n",
        "### **How the Architecture Works**\n",
        "1. **Training**:\n",
        "   - The generator creates inpainted images.\n",
        "   - The discriminator evaluates their realism.\n",
        "   - Both networks are updated iteratively to improve their performance.\n",
        "\n",
        "2. **Inference**:\n",
        "   - Given an input image and a mask, the generator fills the missing regions.\n",
        "   - No discriminator is needed during inference.\n"
      ],
      "metadata": {
        "id": "qSeics-fhrmJ"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Step 1: Gather the Dataset\n"
      ],
      "metadata": {
        "id": "D5Wl0pJ4jriZ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!wget http://data.csail.mit.edu/places/places365/train_large_places365standard.tar\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "vA6ohVVspUws",
        "outputId": "42ff7276-b1f7-4886-c04f-61bfce1ffcb8"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "--2025-01-05 18:08:34--  http://data.csail.mit.edu/places/places365/train_large_places365standard.tar\n",
            "Resolving data.csail.mit.edu (data.csail.mit.edu)... 128.52.131.233\n",
            "Connecting to data.csail.mit.edu (data.csail.mit.edu)|128.52.131.233|:80... connected.\n",
            "HTTP request sent, awaiting response... 301 Moved Permanently\n",
            "Location: https://data.csail.mit.edu/places/places365/train_large_places365standard.tar [following]\n",
            "--2025-01-05 18:08:35--  https://data.csail.mit.edu/places/places365/train_large_places365standard.tar\n",
            "Connecting to data.csail.mit.edu (data.csail.mit.edu)|128.52.131.233|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 112598435840 (105G) [application/x-tar]\n",
            "Saving to: ‘train_large_places365standard.tar’\n",
            "\n",
            "d.tar                 2%[                    ]   2.10G  15.4MB/s    eta 1h 55m "
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!tar -xvf train_large_places365standard.tar"
      ],
      "metadata": {
        "id": "5faLa2eSqhYR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!ls"
      ],
      "metadata": {
        "id": "f9ZfAbSpqz_d"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Load and Use the Dataset\n"
      ],
      "metadata": {
        "id": "lT70RwhUq8O-"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "from tensorflow.keras.preprocessing.image import load_img, img_to_array\n",
        "from skimage.draw import random_shapes\n",
        "from sklearn.model_selection import train_test_split\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "# Parameters\n",
        "DATA_DIR = \"./places2\"  # Path to your dataset\n",
        "IMG_SIZE = (128, 128)  # Target image size\n",
        "BATCH_SIZE = 32\n",
        "MASK_TYPE = \"random\"  # Options: \"random\", \"rectangle\"\n",
        "\n",
        "# Generate a binary mask\n",
        "def generate_mask(img_size, mask_type=\"random\"):\n",
        "    if mask_type == \"random\":\n",
        "        # Generate random irregular mask\n",
        "        mask, _ = random_shapes(img_size, max_shapes=5, min_size=50, max_size=100, multichannel=False)\n",
        "        mask = (mask == 255).astype(np.float32)  # Convert to binary mask\n",
        "    elif mask_type == \"rectangle\":\n",
        "        # Generate rectangular mask\n",
        "        mask = np.ones(img_size, dtype=np.float32)\n",
        "        x1, y1 = np.random.randint(0, img_size[0] // 2), np.random.randint(0, img_size[1] // 2)\n",
        "        x2, y2 = np.random.randint(x1, img_size[0]), np.random.randint(y1, img_size[1])\n",
        "        mask[x1:x2, y1:y2] = 0\n",
        "    return mask\n",
        "\n",
        "# Load and preprocess a single image\n",
        "def load_and_preprocess_image(img_path, img_size):\n",
        "    img = load_img(img_path, target_size=img_size)\n",
        "    img = img_to_array(img) / 255.0  # Normalize to [0, 1]\n",
        "    return img\n",
        "\n",
        "# Create a TensorFlow Dataset\n",
        "def create_dataset(data_dir, img_size, batch_size, mask_type=\"random\"):\n",
        "    # Get list of image paths\n",
        "    image_paths = [os.path.join(data_dir, img_name) for img_name in os.listdir(data_dir)\n",
        "                   if img_name.endswith(('.jpg', '.png', '.jpeg'))]\n",
        "\n",
        "    if not image_paths:\n",
        "        raise ValueError(f\"No images found in {data_dir}. Please check the dataset path.\")\n",
        "\n",
        "    # Split into training and validation sets\n",
        "    train_paths, val_paths = train_test_split(image_paths, test_size=0.2, random_state=42)\n",
        "\n",
        "    # Function to load and preprocess images and masks\n",
        "    def process_image(img_path):\n",
        "        # Load and preprocess image\n",
        "        img = tf.numpy_function(load_and_preprocess_image, [img_path, img_size], tf.float32)\n",
        "        img.set_shape(img_size + (3,))  # Set shape explicitly\n",
        "\n",
        "        # Generate mask\n",
        "        mask = tf.numpy_function(generate_mask, [img_size, mask_type], tf.float32)\n",
        "        mask.set_shape(img_size)  # Set shape explicitly\n",
        "\n",
        "        # Apply mask to image\n",
        "        masked_img = img * tf.expand_dims(mask, axis=-1)\n",
        "\n",
        "        return masked_img, img, mask\n",
        "\n",
        "    # Create TensorFlow Dataset\n",
        "    train_dataset = tf.data.Dataset.from_tensor_slices(train_paths)\n",
        "    train_dataset = train_dataset.map(process_image, num_parallel_calls=tf.data.AUTOTUNE)\n",
        "    train_dataset = train_dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE)\n",
        "\n",
        "    val_dataset = tf.data.Dataset.from_tensor_slices(val_paths)\n",
        "    val_dataset = val_dataset.map(process_image, num_parallel_calls=tf.data.AUTOTUNE)\n",
        "    val_dataset = val_dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE)\n",
        "\n",
        "    return train_dataset, val_dataset\n",
        "\n",
        "# Create the dataset\n",
        "try:\n",
        "    train_dataset, val_dataset = create_dataset(DATA_DIR, IMG_SIZE, BATCH_SIZE, MASK_TYPE)\n",
        "    print(\"Dataset created successfully.\")\n",
        "except Exception as e:\n",
        "    print(f\"Error creating dataset: {e}\")\n",
        "\n",
        "# Visualize a sample batch\n",
        "def visualize_batch(dataset):\n",
        "    for masked_images, original_images, masks in dataset.take(1):\n",
        "        plt.figure(figsize=(10, 5))\n",
        "        for i in range(3):  # Display 3 samples\n",
        "            plt.subplot(3, 3, i * 3 + 1)\n",
        "            plt.title(\"Masked Image\")\n",
        "            plt.imshow(masked_images[i])\n",
        "            plt.axis(\"off\")\n",
        "\n",
        "            plt.subplot(3, 3, i * 3 + 2)\n",
        "            plt.title(\"Original Image\")\n",
        "            plt.imshow(original_images[i])\n",
        "            plt.axis(\"off\")\n",
        "\n",
        "            plt.subplot(3, 3, i * 3 + 3)\n",
        "            plt.title(\"Mask\")\n",
        "            plt.imshow(masks[i], cmap='gray')\n",
        "            plt.axis(\"off\")\n",
        "        plt.show()\n",
        "\n",
        "# Visualize a batch from the training dataset\n",
        "if 'train_dataset' in locals():\n",
        "    visualize_batch(train_dataset)\n",
        "else:\n",
        "    print(\"Dataset not available for visualization.\")"
      ],
      "metadata": {
        "id": "IB80Tsl8GEDY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Organize the Dataset"
      ],
      "metadata": {
        "id": "rrcW9k6rrFjp"
      }
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "K_AbgCKMr_Sz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Use the Dataset in a Machine Learning Framework"
      ],
      "metadata": {
        "id": "u70-TMFxrSfo"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Task"
      ],
      "metadata": {
        "id": "6fE1YaHCrvQA"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "---\n",
        "\n",
        "### **Step-by-Step Guide**\n",
        "\n",
        "---\n",
        "\n",
        "### **Step 1: Prepare Your Dataset**\n",
        "1. **Organize Your Dataset**:\n",
        "   - Place all your images in a single folder (e.g., `./places2`).\n",
        "   - Ensure the images are in common formats like `.jpg`, `.png`, or `.jpeg`.\n",
        "\n",
        "2. **Update the Code**:\n",
        "   - Use the provided code to load, preprocess, and generate masks for your dataset.\n",
        "   - Set the `DATA_DIR` variable to the path of your dataset folder.\n",
        "\n",
        "3. **Run the Dataset Preparation Code**:\n",
        "   - Execute the code to create the `train_dataset` and `val_dataset`.\n",
        "   - Verify that the dataset is created successfully by visualizing a sample batch.\n",
        "\n",
        "---\n",
        "\n",
        "### **Step 2: Design the Generator**\n",
        "The generator is responsible for filling in the missing regions of the image. Here's how to design it:\n",
        "\n",
        "1. **Input**:\n",
        "   - Concatenate the masked image and binary mask as input.\n",
        "\n",
        "2. **Encoder**:\n",
        "   - Use convolutional layers to extract features from the input.\n",
        "   - Apply instance normalization and ReLU activation.\n",
        "\n",
        "3. **Cascaded Modulation Blocks**:\n",
        "   - Implement global and spatially adaptive modulation to balance global coherence and local details.\n",
        "\n",
        "4. **Attention Mechanism**:\n",
        "   - Use self-attention or non-local blocks to propagate information from known to unknown regions.\n",
        "\n",
        "5. **Decoder**:\n",
        "   - Use transposed convolutions or upsampling layers to reconstruct the image.\n",
        "   - Apply skip connections from the encoder to preserve details.\n",
        "\n",
        "6. **Output Layer**:\n",
        "   - Use a convolutional layer with a `tanh` activation to generate the final image.\n",
        "\n",
        "---\n",
        "\n",
        "### **Step 3: Design the Discriminator**\n",
        "The discriminator evaluates the realism of the inpainted images.\n",
        "\n",
        "1. **Input**:\n",
        "   - Concatenate the inpainted image and ground truth image (for conditional GANs).\n",
        "\n",
        "2. **Layers**:\n",
        "   - Use convolutional layers with leaky ReLU activation.\n",
        "   - Apply spectral normalization for stable training.\n",
        "\n",
        "3. **Output**:\n",
        "   - Use a fully connected layer to output a realism score.\n",
        "\n",
        "---\n",
        "\n",
        "### **Step 4: Define Loss Functions**\n",
        "1. **Adversarial Loss**:\n",
        "   - Use Wasserstein loss or hinge loss for stable training.\n",
        "\n",
        "2. **Reconstruction Loss**:\n",
        "   - Use L1 or L2 loss to ensure pixel-level accuracy.\n",
        "\n",
        "3. **Perceptual Loss**:\n",
        "   - Use a pre-trained VGG network to compare high-level features.\n",
        "\n",
        "4. **Total Loss**:\n",
        "   - Combine all losses with appropriate weights.\n",
        "\n",
        "---\n",
        "\n",
        "### **Step 5: Train the Model**\n",
        "1. **Optimizers**:\n",
        "   - Use Adam optimizer for both the generator and discriminator.\n",
        "\n",
        "2. **Training Loop**:\n",
        "   - Alternate between training the generator and discriminator.\n",
        "   - Use gradient penalty for Wasserstein GANs.\n",
        "\n",
        "3. **Evaluation**:\n",
        "   - Visualize inpainted images during training.\n",
        "   - Use metrics like PSNR and SSIM for quantitative evaluation.\n",
        "\n",
        "---\n",
        "\n",
        "### **Step 6: Deploy the Model**\n",
        "1. **Save the Trained Model**:\n",
        "   - Save the generator for inference.\n",
        "\n",
        "2. **Inference**:\n",
        "   - Load the generator and use it to inpaint new images.\n"
      ],
      "metadata": {
        "id": "NowDdYzorwVX"
      }
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "BAHTnRHFBikD"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}