{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "531bfd29",
   "metadata": {},
   "source": [
    "In this notebook, we transform raw datasets to parquet format to enable faster loading speed during training and evaluation.\n",
    "\n",
    "The raw format of released datasets is as follows:\n",
    "```python\n",
    "# train set\n",
    "/train/real/...\n",
    "/train/fake/...\n",
    "/train/masks/...\n",
    "# valid set\n",
    "/valid/real/...\n",
    "/valid/fake/...\n",
    "/valid/masks/...\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "8bd7e9d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from datasets import Dataset, DatasetDict\n",
    "from datasets import Features, Image, Value\n",
    "from typing import List, Optional\n",
    "\n",
    "\n",
    "def load_images_from_dir(directory: str) -> List[str]:\n",
    "    return [\n",
    "        os.path.join(directory, fname)\n",
    "        for fname in os.listdir(directory)\n",
    "        if fname.endswith((\"jpg\", \"jpeg\", \"png\", \"tif\"))\n",
    "    ]\n",
    "\n",
    "\n",
    "def create_split(root_dir: str, split: str) -> Optional[Dataset]:\n",
    "    fake_dir = os.path.join(root_dir, split, \"fake\")\n",
    "    masks_dir = os.path.join(root_dir, split, \"masks\")\n",
    "    real_dir = os.path.join(root_dir, split, \"real\")\n",
    "\n",
    "    if all(not os.path.isdir(p) for p in [fake_dir, masks_dir, real_dir]):\n",
    "        return None\n",
    "\n",
    "    print(f\"Split: {split},\", end=\" \")\n",
    "    fake_images, real_images, mask_images = [], [], []\n",
    "    if os.path.isdir(fake_dir):\n",
    "        fake_images = load_images_from_dir(fake_dir)\n",
    "        print(f\"Fake images: {len(fake_images)}\", end=\"\")\n",
    "    if os.path.isdir(masks_dir):\n",
    "        mask_images = load_images_from_dir(masks_dir)\n",
    "        print(f\", Masks: {len(mask_images)}\", end=\"\")\n",
    "        assert len(fake_images) == len(mask_images)\n",
    "    if os.path.isdir(real_dir):\n",
    "        real_images = load_images_from_dir(real_dir)\n",
    "        print(f\", Real images: {len(real_images)}\", end=\"\")\n",
    "    print()\n",
    "\n",
    "    return Dataset.from_dict(\n",
    "        {\n",
    "            \"path\": fake_images + real_images,\n",
    "            \"image\": fake_images + real_images,\n",
    "            \"mask\": mask_images + [None] * len(real_images),\n",
    "        },\n",
    "        features=Features(\n",
    "            {\"path\": Value(dtype=\"string\"), \"image\": Image(), \"mask\": Image()}\n",
    "        ),\n",
    "    )\n",
    "\n",
    "\n",
    "def create_dataset(root_dir: str) -> DatasetDict:\n",
    "    return DatasetDict(\n",
    "        {\n",
    "            split: d\n",
    "            for split in [\"train\", \"valid\", \"test\"]\n",
    "            if (d := create_split(root_dir, split)) is not None\n",
    "        }\n",
    "    )\n",
    "\n",
    "\n",
    "# replace with your own dataset path\n",
    "root_dir = \"/gemini/space/lye/track1\"\n",
    "save_dir = \"/gemini/space/jyc/track1\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1d6f1c7",
   "metadata": {},
   "source": [
    "We merge `real/` and `fake/` into `images` column for simplity. A image is real if there is no corresponding mask."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "07009f1e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Split: train, Fake images: 798831, Masks: 798831, Real images: 156100\n",
      "Split: valid, Fake images: 199708, Masks: 199708, Real images: 39025\n",
      "Split: test, Images: 222847\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['path', 'image', 'mask'],\n",
       "        num_rows: 954931\n",
       "    })\n",
       "    valid: Dataset({\n",
       "        features: ['path', 'image', 'mask'],\n",
       "        num_rows: 238733\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['path', 'image'],\n",
       "        num_rows: 222847\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset = create_dataset(root_dir)\n",
    "dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3aa7de84",
   "metadata": {},
   "source": [
    "Then save processed datasets to parquet."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd6b20bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.makedirs(save_dir, exist_ok=True)\n",
    "for split in dataset:\n",
    "    dataset[split].to_parquet(os.path.join(save_dir, f\"{split}.parquet\"))\n",
    "    print(f\"Saved {split} split to {save_dir}/{split}.parquet\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f63933c8",
   "metadata": {},
   "source": [
    "Load from processed datasets to do whatever you want."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "4af7f346",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['path', 'image', 'mask'],\n",
       "    num_rows: 954931\n",
       "})"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import os\n",
    "from datasets import load_dataset\n",
    "\n",
    "trainset = load_dataset(\"parquet\", data_dir=save_dir, split=\"train\")\n",
    "trainset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3c84f0a",
   "metadata": {},
   "source": [
    "Since the forged components are usually smaller in proportion compared to the real ones, this leads to class imbalance.\n",
    "For optimal training performance, hyper parameters such as `pixel_forge_weight` and `cls_forge_weight` in `src.loupe.configuration_loupe.LoupeConfig` must be appropriately configured. These parameters control the weights of forged pixels and forged images.\n",
    "\n",
    "Once suitable parameters are found using the following code snippet, you can set them in `configs/model/cls.yaml` or `configs/model/seg.yaml`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40a5ec91",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "19d416f59f20464692ee95bddefdaded",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Computing mask stats (num_proc=8):   0%|          | 0/5000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cls_forge_weight: 0.16920000000000002\n",
      "patch_forge_weight: 0.9294853830073696\n",
      "pixel_forge_weight: 0.9160308902282281\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "from PIL import Image\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "cls_forge_weight: float  # the ratio of forged images to total images.\n",
    "# the ratio of forged patches to total patches across all images.\n",
    "patch_forge_weight: float\n",
    "# the ratio of forged pixels to total pixels across fake images.\n",
    "pixel_forge_weight: float\n",
    "\n",
    "num_subset_samples = min(5000, len(trainset))\n",
    "subset = trainset.shuffle().select(range(num_subset_samples))\n",
    "image_size, patch_size = 336, 14\n",
    "\n",
    "\n",
    "def compute_mask_stats(example):\n",
    "\n",
    "    if example[\"mask\"] is None:\n",
    "        return {\n",
    "            \"is_forge\": 0,\n",
    "            \"forge_pixel_sum\": 0.0,\n",
    "            \"total_pixel_count\": 0,\n",
    "            \"forge_patch_sum\": 0.0,\n",
    "        }\n",
    "\n",
    "    mask = example[\"mask\"].convert(\"L\").resize((image_size, image_size), Image.NEAREST)\n",
    "    mask_np = np.array(mask, dtype=np.float32)\n",
    "\n",
    "    if mask_np.max() != mask_np.min():\n",
    "        mask_np = (mask_np - mask_np.min()) / (mask_np.max() - mask_np.min())\n",
    "    else:\n",
    "        mask_np[:] = 0.0\n",
    "\n",
    "    forged_pixel_sum = mask_np.sum()\n",
    "    total_pixels = mask_np.size\n",
    "\n",
    "    reshaped = mask_np.reshape(\n",
    "        image_size // patch_size, patch_size, image_size // patch_size, patch_size\n",
    "    )\n",
    "    patches = reshaped.transpose(0, 2, 1, 3)\n",
    "    forged_patch_sum = (patches != 0).sum(axis=(2, 3)) / (patch_size * patch_size)\n",
    "    forged_patch_sum = forged_patch_sum.sum()\n",
    "\n",
    "    return {\n",
    "        \"is_forge\": 1,\n",
    "        \"forge_pixel_sum\": forged_pixel_sum,\n",
    "        \"total_pixel_count\": total_pixels,\n",
    "        \"forge_patch_sum\": forged_patch_sum,\n",
    "    }\n",
    "\n",
    "\n",
    "processed = subset.map(compute_mask_stats, num_proc=8, desc=\"Computing mask stats\")\n",
    "\n",
    "num_forge_images = sum(processed[\"is_forge\"])\n",
    "num_forge_pixels = sum(processed[\"forge_pixel_sum\"])\n",
    "num_total_pixels = sum(processed[\"total_pixel_count\"])\n",
    "num_forge_patches = sum(processed[\"forge_patch_sum\"])\n",
    "num_total_patches = len(processed) * (image_size // patch_size) ** 2\n",
    "\n",
    "cls_forge_weight = 1 - num_forge_images / len(processed)\n",
    "patch_forge_weight = 1 - num_forge_patches / num_total_patches\n",
    "pixel_forge_weight = 1 - num_forge_pixels / num_total_pixels\n",
    "\n",
    "print(\"cls_forge_weight:\", cls_forge_weight)\n",
    "print(\"patch_forge_weight:\", patch_forge_weight)\n",
    "print(\"pixel_forge_weight:\", pixel_forge_weight)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "loupe2",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}