UniFMIR.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hxuwBSFDhghn"
      },
      "source": [
        "# UniFMIR: Pre-training a Foundation Model for Universal Fluorescence Microscopy Image Restoration\n",
        "This demo allows you to run the models on your own images or the examples  from the paper. Please refer to the paper for more details."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YeRahcBXhrgO"
      },
      "source": [
        "## 1. Clone the repository and prepare environment"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "z0nciz5yhOKl",
        "outputId": "5f5b96b8-144f-482d-b2ab-a2d9f991d858"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Cloning into 'UNiFMIR'...\n",
            "remote: Enumerating objects: 211, done.\u001b[K\n",
            "remote: Counting objects: 100% (207/207), done.\u001b[K\n",
            "remote: Compressing objects: 100% (163/163), done.\u001b[K\n",
            "remote: Total 211 (delta 63), reused 172 (delta 33), pack-reused 4\u001b[K\n",
            "Receiving objects: 100% (211/211), 371.09 KiB | 10.60 MiB/s, done.\n",
            "Resolving deltas: 100% (63/63), done.\n"
          ]
        }
      ],
      "source": [
        "# Clone the repository\n",
        "!git clone https://github.com/cxm12/UNiFMIR"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lHVPuNvTp7VY",
        "outputId": "07a72690-0fc7-4c69-b9cf-d40feda6cfcf"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: gradio in /usr/local/lib/python3.10/dist-packages (3.48.0)\n",
            "Requirement already satisfied: einops==0.6.0 in /usr/local/lib/python3.10/dist-packages (0.6.0)\n",
            "Collecting timm==0.6.11\n",
            "  Downloading timm-0.6.11-py3-none-any.whl (548 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m548.7/548.7 kB\u001b[0m \u001b[31m7.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: torch>=1.7 in /usr/local/lib/python3.10/dist-packages (from timm==0.6.11) (2.0.1+cu118)\n",
            "Requirement already satisfied: torchvision in /usr/local/lib/python3.10/dist-packages (from timm==0.6.11) (0.15.2+cu118)\n",
            "Requirement already satisfied: pyyaml in /usr/local/lib/python3.10/dist-packages (from timm==0.6.11) (6.0.1)\n",
            "Requirement already satisfied: huggingface-hub in /usr/local/lib/python3.10/dist-packages (from timm==0.6.11) (0.18.0)\n",
            "Requirement already satisfied: aiofiles<24.0,>=22.0 in /usr/local/lib/python3.10/dist-packages (from gradio) (23.2.1)\n",
            "Requirement already satisfied: altair<6.0,>=4.2.0 in /usr/local/lib/python3.10/dist-packages (from gradio) (4.2.2)\n",
            "Requirement already satisfied: fastapi in /usr/local/lib/python3.10/dist-packages (from gradio) (0.103.2)\n",
            "Requirement already satisfied: ffmpy in /usr/local/lib/python3.10/dist-packages (from gradio) (0.3.1)\n",
            "Requirement already satisfied: gradio-client==0.6.1 in /usr/local/lib/python3.10/dist-packages (from gradio) (0.6.1)\n",
            "Requirement already satisfied: httpx in /usr/local/lib/python3.10/dist-packages (from gradio) (0.25.0)\n",
            "Requirement already satisfied: importlib-resources<7.0,>=1.3 in /usr/local/lib/python3.10/dist-packages (from gradio) (6.1.0)\n",
            "Requirement already satisfied: jinja2<4.0 in /usr/local/lib/python3.10/dist-packages (from gradio) (3.1.2)\n",
            "Requirement already satisfied: markupsafe~=2.0 in /usr/local/lib/python3.10/dist-packages (from gradio) (2.1.3)\n",
            "Requirement already satisfied: matplotlib~=3.0 in /usr/local/lib/python3.10/dist-packages (from gradio) (3.7.1)\n",
            "Requirement already satisfied: numpy~=1.0 in /usr/local/lib/python3.10/dist-packages (from gradio) (1.23.5)\n",
            "Requirement already satisfied: orjson~=3.0 in /usr/local/lib/python3.10/dist-packages (from gradio) (3.9.9)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.10/dist-packages (from gradio) (23.2)\n",
            "Requirement already satisfied: pandas<3.0,>=1.0 in /usr/local/lib/python3.10/dist-packages (from gradio) (1.5.3)\n",
            "Requirement already satisfied: pillow<11.0,>=8.0 in /usr/local/lib/python3.10/dist-packages (from gradio) (9.4.0)\n",
            "Requirement already satisfied: pydantic!=1.8,!=1.8.1,!=2.0.0,!=2.0.1,<3.0.0,>=1.7.4 in /usr/local/lib/python3.10/dist-packages (from gradio) (1.10.13)\n",
            "Requirement already satisfied: pydub in /usr/local/lib/python3.10/dist-packages (from gradio) (0.25.1)\n",
            "Requirement already satisfied: python-multipart in /usr/local/lib/python3.10/dist-packages (from gradio) (0.0.6)\n",
            "Requirement already satisfied: requests~=2.0 in /usr/local/lib/python3.10/dist-packages (from gradio) (2.31.0)\n",
            "Requirement already satisfied: semantic-version~=2.0 in /usr/local/lib/python3.10/dist-packages (from gradio) (2.10.0)\n",
            "Requirement already satisfied: typing-extensions~=4.0 in /usr/local/lib/python3.10/dist-packages (from gradio) (4.5.0)\n",
            "Requirement already satisfied: uvicorn>=0.14.0 in /usr/local/lib/python3.10/dist-packages (from gradio) (0.23.2)\n",
            "Requirement already satisfied: websockets<12.0,>=10.0 in /usr/local/lib/python3.10/dist-packages (from gradio) (11.0.3)\n",
            "Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from gradio-client==0.6.1->gradio) (2023.6.0)\n",
            "Requirement already satisfied: entrypoints in /usr/local/lib/python3.10/dist-packages (from altair<6.0,>=4.2.0->gradio) (0.4)\n",
            "Requirement already satisfied: jsonschema>=3.0 in /usr/local/lib/python3.10/dist-packages (from altair<6.0,>=4.2.0->gradio) (4.19.1)\n",
            "Requirement already satisfied: toolz in /usr/local/lib/python3.10/dist-packages (from altair<6.0,>=4.2.0->gradio) (0.12.0)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from huggingface-hub->timm==0.6.11) (3.12.4)\n",
            "Requirement already satisfied: tqdm>=4.42.1 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub->timm==0.6.11) (4.66.1)\n",
            "Requirement already satisfied: contourpy>=1.0.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib~=3.0->gradio) (1.1.1)\n",
            "Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.10/dist-packages (from matplotlib~=3.0->gradio) (0.12.1)\n",
            "Requirement already satisfied: fonttools>=4.22.0 in /usr/local/lib/python3.10/dist-packages (from matplotlib~=3.0->gradio) (4.43.1)\n",
            "Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib~=3.0->gradio) (1.4.5)\n",
            "Requirement already satisfied: pyparsing>=2.3.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib~=3.0->gradio) (3.1.1)\n",
            "Requirement already satisfied: python-dateutil>=2.7 in /usr/local/lib/python3.10/dist-packages (from matplotlib~=3.0->gradio) (2.8.2)\n",
            "Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.10/dist-packages (from pandas<3.0,>=1.0->gradio) (2023.3.post1)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests~=2.0->gradio) (3.3.0)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests~=2.0->gradio) (3.4)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests~=2.0->gradio) (2.0.6)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests~=2.0->gradio) (2023.7.22)\n",
            "Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch>=1.7->timm==0.6.11) (1.12)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch>=1.7->timm==0.6.11) (3.1)\n",
            "Requirement already satisfied: triton==2.0.0 in /usr/local/lib/python3.10/dist-packages (from torch>=1.7->timm==0.6.11) (2.0.0)\n",
            "Requirement already satisfied: cmake in /usr/local/lib/python3.10/dist-packages (from triton==2.0.0->torch>=1.7->timm==0.6.11) (3.27.6)\n",
            "Requirement already satisfied: lit in /usr/local/lib/python3.10/dist-packages (from triton==2.0.0->torch>=1.7->timm==0.6.11) (17.0.2)\n",
            "Requirement already satisfied: click>=7.0 in /usr/local/lib/python3.10/dist-packages (from uvicorn>=0.14.0->gradio) (8.1.7)\n",
            "Requirement already satisfied: h11>=0.8 in /usr/local/lib/python3.10/dist-packages (from uvicorn>=0.14.0->gradio) (0.14.0)\n",
            "Requirement already satisfied: anyio<4.0.0,>=3.7.1 in /usr/local/lib/python3.10/dist-packages (from fastapi->gradio) (3.7.1)\n",
            "Requirement already satisfied: starlette<0.28.0,>=0.27.0 in /usr/local/lib/python3.10/dist-packages (from fastapi->gradio) (0.27.0)\n",
            "Requirement already satisfied: httpcore<0.19.0,>=0.18.0 in /usr/local/lib/python3.10/dist-packages (from httpx->gradio) (0.18.0)\n",
            "Requirement already satisfied: sniffio in /usr/local/lib/python3.10/dist-packages (from httpx->gradio) (1.3.0)\n",
            "Requirement already satisfied: exceptiongroup in /usr/local/lib/python3.10/dist-packages (from anyio<4.0.0,>=3.7.1->fastapi->gradio) (1.1.3)\n",
            "Requirement already satisfied: attrs>=22.2.0 in /usr/local/lib/python3.10/dist-packages (from jsonschema>=3.0->altair<6.0,>=4.2.0->gradio) (23.1.0)\n",
            "Requirement already satisfied: jsonschema-specifications>=2023.03.6 in /usr/local/lib/python3.10/dist-packages (from jsonschema>=3.0->altair<6.0,>=4.2.0->gradio) (2023.7.1)\n",
            "Requirement already satisfied: referencing>=0.28.4 in /usr/local/lib/python3.10/dist-packages (from jsonschema>=3.0->altair<6.0,>=4.2.0->gradio) (0.30.2)\n",
            "Requirement already satisfied: rpds-py>=0.7.1 in /usr/local/lib/python3.10/dist-packages (from jsonschema>=3.0->altair<6.0,>=4.2.0->gradio) (0.10.4)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.10/dist-packages (from python-dateutil>=2.7->matplotlib~=3.0->gradio) (1.16.0)\n",
            "Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.10/dist-packages (from sympy->torch>=1.7->timm==0.6.11) (1.3.0)\n",
            "Installing collected packages: timm\n",
            "Successfully installed timm-0.6.11\n"
          ]
        }
      ],
      "source": [
        "# Install packages\n",
        "!pip install gradio einops==0.6.0 timm==0.6.11"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tqe-qP_CmDr2"
      },
      "source": [
        "## 2. Get the models and examples"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "z39i17CGmTtr",
        "outputId": "44413e37-71a6-4193-fbc9-c2b2565aab74"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "/content/UNiFMIR\n"
          ]
        }
      ],
      "source": [
        "# Get into the directory\n",
        "%cd UNiFMIR"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "je7mHsH0mGd-",
        "outputId": "dd2713c8-aac4-483a-d793-45bad92a6ae6"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "--2023-10-18 02:42:17--  https://github.com/cxm12/UNiFMIR/releases/download/2023.10.05/example.tgz\n",
            "Resolving github.com (github.com)... 140.82.121.4\n",
            "Connecting to github.com (github.com)|140.82.121.4|:443... connected.\n",
            "HTTP request sent, awaiting response... 302 Found\n",
            "Location: https://objects.githubusercontent.com/github-production-release-asset-2e65be/657524589/649ba765-a28d-4c0c-943d-c086ee6dbaec?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIWNJYAX4CSVEH53A%2F20231018%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20231018T024217Z&X-Amz-Expires=300&X-Amz-Signature=c25e8f519a13c6ea0668ff058fd30c5a39cab43ea07cc5a8d2f3db8f2ed41334&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=657524589&response-content-disposition=attachment%3B%20filename%3Dexample.tgz&response-content-type=application%2Foctet-stream [following]\n",
            "--2023-10-18 02:42:17--  https://objects.githubusercontent.com/github-production-release-asset-2e65be/657524589/649ba765-a28d-4c0c-943d-c086ee6dbaec?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIWNJYAX4CSVEH53A%2F20231018%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20231018T024217Z&X-Amz-Expires=300&X-Amz-Signature=c25e8f519a13c6ea0668ff058fd30c5a39cab43ea07cc5a8d2f3db8f2ed41334&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=657524589&response-content-disposition=attachment%3B%20filename%3Dexample.tgz&response-content-type=application%2Foctet-stream\n",
            "Resolving objects.githubusercontent.com (objects.githubusercontent.com)... 185.199.108.133, 185.199.109.133, 185.199.110.133, ...\n",
            "Connecting to objects.githubusercontent.com (objects.githubusercontent.com)|185.199.108.133|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 682938742 (651M) [application/octet-stream]\n",
            "Saving to: ‘example.tgz’\n",
            "\n",
            "example.tgz         100%[===================>] 651.30M  51.7MB/s    in 13s     \n",
            "\n",
            "2023-10-18 02:42:31 (50.6 MB/s) - ‘example.tgz’ saved [682938742/682938742]\n",
            "\n",
            "--2023-10-18 02:42:31--  https://github.com/cxm12/UNiFMIR/releases/download/2023.10.05/model.tgz\n",
            "Resolving github.com (github.com)... 140.82.121.4\n",
            "Connecting to github.com (github.com)|140.82.121.4|:443... connected.\n",
            "HTTP request sent, awaiting response... 302 Found\n",
            "Location: https://objects.githubusercontent.com/github-production-release-asset-2e65be/657524589/2cce36c6-eb83-4bc8-8c29-c6533eb9a18a?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIWNJYAX4CSVEH53A%2F20231018%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20231018T024231Z&X-Amz-Expires=300&X-Amz-Signature=a4d9eefadd15c57213e4966b2a2506ea16cc595e96a7600a8c9e5f5c994fccf6&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=657524589&response-content-disposition=attachment%3B%20filename%3Dmodel.tgz&response-content-type=application%2Foctet-stream [following]\n",
            "--2023-10-18 02:42:31--  https://objects.githubusercontent.com/github-production-release-asset-2e65be/657524589/2cce36c6-eb83-4bc8-8c29-c6533eb9a18a?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIWNJYAX4CSVEH53A%2F20231018%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20231018T024231Z&X-Amz-Expires=300&X-Amz-Signature=a4d9eefadd15c57213e4966b2a2506ea16cc595e96a7600a8c9e5f5c994fccf6&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=657524589&response-content-disposition=attachment%3B%20filename%3Dmodel.tgz&response-content-type=application%2Foctet-stream\n",
            "Resolving objects.githubusercontent.com (objects.githubusercontent.com)... 185.199.108.133, 185.199.109.133, 185.199.110.133, ...\n",
            "Connecting to objects.githubusercontent.com (objects.githubusercontent.com)|185.199.108.133|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 416285157 (397M) [application/octet-stream]\n",
            "Saving to: ‘model.tgz’\n",
            "\n",
            "model.tgz           100%[===================>] 397.00M  50.5MB/s    in 8.1s    \n",
            "\n",
            "2023-10-18 02:42:40 (49.1 MB/s) - ‘model.tgz’ saved [416285157/416285157]\n",
            "\n"
          ]
        }
      ],
      "source": [
        "# Get examples and models\n",
        "!wget https://github.com/cxm12/UNiFMIR/releases/download/2023.10.05/example.tgz\n",
        "!wget https://github.com/cxm12/UNiFMIR/releases/download/2023.10.05/model.tgz"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8nxef9BbmIyH",
        "outputId": "618b141c-e67e-4e22-8cdc-a18958d42dd4"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "experiment/SwinIR2t3_stage2VCD/model_best17.pt\n",
            "experiment/SwinIRCCPs/model_best.pt\n",
            "experiment/SwinIRDenoising_Planaria/model_best15.pt\n",
            "experiment/SwinIRER/model_best147.pt\n",
            "experiment/SwinIRF-actin/model_best181.pt\n",
            "experiment/SwinIRF-actin/model_best77.pt\n",
            "experiment/SwinIRIsotropic_Liver/model_best465.pt\n",
            "experiment/SwinIRMicrotubules/model_best.pt\n",
            "experiment/SwinIRmto1Denoising_Tribolium/model_best.pt\n",
            "experiment/SwinIRproj2stg_enlcn_2npzProjection_Flywing/model_best6.pt\n",
            "exampledata/\n",
            "exampledata/Proj/\n",
            "exampledata/Proj/Flywing.tif\n",
            "exampledata/Denoise/\n",
            "exampledata/Denoise/T/\n",
            "exampledata/Denoise/T/C3/\n",
            "exampledata/Denoise/T/C3/Tribolium_C3.tif\n",
            "exampledata/Denoise/T/C2/\n",
            "exampledata/Denoise/T/C2/Tribolium_C2.tif\n",
            "exampledata/Denoise/T/C1/\n",
            "exampledata/Denoise/T/C1/Tribolium_C1.tif\n",
            "exampledata/Denoise/P/\n",
            "exampledata/Denoise/P/C3/\n",
            "exampledata/Denoise/P/C3/Planaria_C3.tif\n",
            "exampledata/Denoise/P/C2/\n",
            "exampledata/Denoise/P/C2/Planaria_C2.tif\n",
            "exampledata/Denoise/P/C1/\n",
            "exampledata/Denoise/P/C1/Planaria_C1.tif\n",
            "exampledata/volumetricRec/\n",
            "exampledata/volumetricRec/VCD.tif\n",
            "exampledata/testdata.txt\n",
            "exampledata/Isotropic/\n",
            "exampledata/Isotropic/Liver.tif\n",
            "exampledata/BioSR/\n",
            "exampledata/BioSR/CCP.tif\n",
            "exampledata/BioSR/F-actin.tif\n",
            "exampledata/BioSR/Microtubes.tif\n"
          ]
        }
      ],
      "source": [
        "# Unzip the files\n",
        "!tar xzvf model.tgz\n",
        "!tar xzvf example.tgz"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WbgjwFjmkIoL"
      },
      "source": [
        "## 2. Run UniFMIR demo\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 646
        },
        "id": "uGNwQ3V0kNV0",
        "outputId": "ba2353c8-0da4-4168-d1c8-ebf5d02721c8"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Setting queue=True in a Colab notebook requires sharing enabled. Setting `share=True` (you can turn this off by setting `share=False` in `launch()` explicitly).\n",
            "\n",
            "Colab notebook detected. To show errors in colab notebook, set debug=True in launch()\n",
            "Running on public URL: https://24c8e3f37f4d85a35c.gradio.live\n",
            "\n",
            "This share link expires in 72 hours. For free permanent hosting and GPU upgrades, run `gradio deploy` from Terminal to deploy to Spaces (https://huggingface.co/spaces)\n"
          ]
        },
        {
          "data": {
            "text/html": [
              "<div><iframe src=\"https://24c8e3f37f4d85a35c.gradio.live\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "from app import *"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}