notebooks/SwinT-ChARM-Perceptual.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```\n",
    "Copyright 2024 Nikolai Körber. All Rights Reserved.\n",
    "\n",
    "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "you may not use this file except in compliance with the License.\n",
    "You may obtain a copy of the License at\n",
    "\n",
    "    https://www.apache.org/licenses/LICENSE-2.0\n",
    "\n",
    "Unless required by applicable law or agreed to in writing, software\n",
    "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "See the License for the specific language governing permissions and\n",
    "limitations under the License.\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# EGIC (TF 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"0\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install immutabledict\n",
    "!pip install tensorflow-datasets\n",
    "!pip install tensorflow-compression~=$(pip show tensorflow | perl -p -0777 -e 's/.*Version: (\\d+\\.\\d+).*/\\1.0/sg')\n",
    "!pip install protobuf==3.20.0\n",
    "!pip install \"numpy<1.24\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test that TFC works correctly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# If you encounter any issues, see https://github.com/tensorflow/compression (Installation/ Colab)\n",
    "!python -m tensorflow_compression.all_tests"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## TensorBoard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the TensorBoard notebook extension\n",
    "%load_ext tensorboard\n",
    "%tensorboard --logdir /tf/notebooks/EGIC/res/train_kkshms2024 --host 0.0.0.0 --port=6011"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Start Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Stage 1 (warm-up)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# caution: parsing arguments might be ignored for kkshms2024 (e.g. model definition)\n",
    "!python /tf/notebooks/EGIC/src/kkshms2024.py -V --model_path /tf/notebooks/EGIC/res/kkshms2024 train --max_support_slices 10 --warm_up 1 --lambda 2 --lr 1e-4 --epochs 200 --steps_per_epoch 10000 --train_path /tf/notebooks/EGIC/res/train_kkshms2024"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Stage 2 (full learning objective)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# caution: parsing arguments might be ignored for kkshms2024 (e.g. model definition)\n",
    "!python /tf/notebooks/EGIC/src/kkshms2024.py -V --model_path /tf/notebooks/EGIC/res/kkshms2024 train --max_support_slices 10 --warm_up 0 --lambda 2 --lr 1e-5 --epochs 300 --steps_per_epoch 10000 --init_disc 1 --train_path /tf/notebooks/EGIC/res/train_kkshms2024"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compression\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# evaluate single image\n",
    "!python /tf/notebooks/EGIC/src/kkshms2024.py -V --model_path /tf/notebooks/EGIC/res/kkshms2024 compress /tf/notebooks/EGIC/res/data/kodak/kodim19.png /tf/notebooks/EGIC/res/eval/kodim19.tfci"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# compress only\n",
    "!python /tf/notebooks/EGIC/src/kkshms2024.py --model_path /tf/notebooks/EGIC/res/kkshms2024 compress /tf/notebooks/EGIC/res/data/kodak/kodim19.png /tf/notebooks/EGIC/res/eval/kodim19.tfci"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# decompress only\n",
    "!python /tf/notebooks/EGIC/src/kkshms2024.py --model_path /tf/notebooks/EGIC/res/kkshms2024 decompress /tf/notebooks/EGIC/res/eval/kodim19.tfci /tf/notebooks/EGIC/res/eval/kodim19_hat.png"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.image as mpimg\n",
    "\n",
    "plt.rcParams[\"figure.figsize\"] = (16,8)\n",
    "\n",
    "x = mpimg.imread('/tf/notebooks/EGIC/res/data/kodak/kodim19.png')\n",
    "x_hat = mpimg.imread('/tf/notebooks/EGIC/res/eval/kodim19_hat.png')\n",
    "\n",
    "f = plt.figure()\n",
    "f.add_subplot(1,2, 1)\n",
    "plt.imshow(x)\n",
    "plt.title('x')\n",
    "f.add_subplot(1,2, 2)\n",
    "plt.imshow(x_hat)\n",
    "plt.title('x_hat')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate Kodak"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# only bpp, PSNR -> we compute FID and KID following https://github.com/facebookresearch/NeuralCompression/blob/main/projects/illm/eval_folder_example.py\n",
    "!python /tf/notebooks/EGIC/src/kkshms2024.py -V --model_path /tf/notebooks/EGIC/res/kkshms2024 evaluate_ds --images_glob=../res/data/kodak/*.png --out_dir=../res/eval/SwinT-ChARM-Perceptual/kodak/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate DIV2K"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# only bpp, PSNR -> we compute FID and KID following https://github.com/facebookresearch/NeuralCompression/blob/main/projects/illm/eval_folder_example.py\n",
    "!python /tf/notebooks/EGIC/src/kkshms2024.py -V --model_path /tf/notebooks/EGIC/res/kkshms2024 evaluate_ds --images_glob=../res/data/DIV2K_valid_HR/*.png --out_dir=../res/eval/SwinT-ChARM-Perceptual/DIV2K_valid_HR/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate CLIC 2020"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# only bpp, PSNR -> we compute FID and KID following https://github.com/facebookresearch/NeuralCompression/blob/main/projects/illm/eval_folder_example.py\n",
    "!python /tf/notebooks/EGIC/src/kkshms2024.py -V --model_path /tf/notebooks/EGIC/res/kkshms2024 evaluate_ds --images_glob=../res/data/clic2020/*/*.png --out_dir=../res/eval/SwinT-ChARM-Perceptual/clic2020/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.0rc1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}