eval-part1.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "33c92487",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "import glob\n",
    "\n",
    "import sys\n",
    "sys.path.insert(1, './src')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c79abd7d",
   "metadata": {},
   "source": [
    "### Load Diffusion pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "19ed7564",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading Diffusion pipeline from:\n",
      "    - /media/share/Datasets/diffusers_out/2024-08-01_10-12-17_MLBriefs24_4_conditional_CA_downsampledmask\n",
      "\n",
      "VQ-VAE loaded\n",
      "U-Net model loaded\n",
      "Diffusion pipeline is ready\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ubuntu/anaconda3/envs/diffusion/lib/python3.9/site-packages/diffusers/configuration_utils.py:244: FutureWarning: It is deprecated to pass a pretrained model name or path to `from_config`.If you were trying to load a scheduler, please use <class 'diffusers.schedulers.scheduling_ddpm.DDPMScheduler'>.from_pretrained(...) instead. Otherwise, please make sure to pass a configuration dictionary instead. This functionality will be removed in v1.0.0.\n",
      "  deprecate(\"config-passed-as-path\", \"1.0.0\", deprecation_message, standard_warn=False)\n"
     ]
    }
   ],
   "source": [
    "from utils_maps import *\n",
    "\n",
    "exp_id = \"2024-07-25_14-27-03_MLBriefs24_2_conditional_vanilla_downsample_mask\"\n",
    "exp_id = \"2024-07-24_22-37-54_MLBriefs24_1_conditional_vanilla_encode_mask\"\n",
    "exp_id = \"2024-07-26_04-40-24_MLBriefs24_0_unconditional\"\n",
    "exp_id = \"2024-08-01_10-12-17_MLBriefs24_4_conditional_CA_downsampledmask\"\n",
    "#exp_id = \"2024-08-02_14-35-03_MLBriefs24_5_conditional_CA_encodedmask\"\n",
    "model_path = f\"/media/share/Datasets/diffusers_out/{exp_id}\"\n",
    "\n",
    "pipeline = load_pipeline(model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bddd0aa8",
   "metadata": {},
   "source": [
    "### Generate N images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "143a0bf7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ba470f8b88a14229afc2d4647468adca",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Resolving data files:   0%|          | 0/1098 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_2606223/336401626.py:28: FutureWarning: Accessing config attribute `sample_size` directly via 'VQModel' object attribute is deprecated. Please access 'sample_size' over 'VQModel's config object instead, e.g. 'unet.config.sample_size'.\n",
      "  im_size = pipeline.vae.sample_size\n",
      "  0%|                                                                                                                                                                  | 0/1 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[0.3804, 0.3804, 0.3961,  ..., 0.3725, 0.3647, 0.3647],\n",
      "         [0.3882, 0.3882, 0.3961,  ..., 0.3725, 0.3647, 0.3647],\n",
      "         [0.3882, 0.3882, 0.3882,  ..., 0.3647, 0.3725, 0.3647],\n",
      "         ...,\n",
      "         [0.3882, 0.3882, 0.3882,  ..., 0.7490, 0.7490, 0.7490],\n",
      "         [0.3882, 0.3882, 0.3882,  ..., 0.7490, 0.7490, 0.7490],\n",
      "         [0.3882, 0.3882, 0.3882,  ..., 0.7490, 0.7490, 0.7490]],\n",
      "\n",
      "        [[0.6235, 0.6235, 0.6392,  ..., 0.6157, 0.6078, 0.6078],\n",
      "         [0.6314, 0.6314, 0.6392,  ..., 0.6157, 0.6078, 0.6000],\n",
      "         [0.6314, 0.6314, 0.6314,  ..., 0.6157, 0.6000, 0.5922],\n",
      "         ...,\n",
      "         [0.6314, 0.6314, 0.6314,  ..., 0.7098, 0.7098, 0.7098],\n",
      "         [0.6314, 0.6314, 0.6314,  ..., 0.7098, 0.7098, 0.7098],\n",
      "         [0.6314, 0.6314, 0.6314,  ..., 0.7098, 0.7098, 0.7098]],\n",
      "\n",
      "        [[0.9765, 0.9765, 1.0000,  ..., 0.9765, 0.9765, 0.9765],\n",
      "         [0.9922, 0.9922, 1.0000,  ..., 0.9843, 0.9843, 0.9843],\n",
      "         [0.9922, 0.9922, 0.9922,  ..., 1.0000, 1.0000, 1.0000],\n",
      "         ...,\n",
      "         [0.9922, 0.9922, 0.9922,  ..., 0.6784, 0.6627, 0.6627],\n",
      "         [0.9922, 0.9922, 0.9922,  ..., 0.6784, 0.6627, 0.6627],\n",
      "         [0.9922, 0.9922, 0.9922,  ..., 0.6784, 0.6627, 0.6627]]])\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9c040d5d5178400e9cdf273e98dcd8ac",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/10 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:02<00:00,  2.24s/it]\n"
     ]
    },
    {
     "ename": "UnboundLocalError",
     "evalue": "local variable 'ue' referenced before assignment",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mUnboundLocalError\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[10], line 92\u001b[0m\n\u001b[1;32m     85\u001b[0m val_dataset \u001b[38;5;241m=\u001b[39m load_dataset(pipeline\u001b[38;5;241m.\u001b[39margs\u001b[38;5;241m.\u001b[39mval_dataset_name,\n\u001b[1;32m     86\u001b[0m     pipeline\u001b[38;5;241m.\u001b[39margs\u001b[38;5;241m.\u001b[39mdataset_config_name,\n\u001b[1;32m     87\u001b[0m     cache_dir\u001b[38;5;241m=\u001b[39mpipeline\u001b[38;5;241m.\u001b[39margs\u001b[38;5;241m.\u001b[39mcache_dir,\n\u001b[1;32m     88\u001b[0m     split\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtrain\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m     89\u001b[0m     )\n\u001b[1;32m     90\u001b[0m val_dataset\u001b[38;5;241m.\u001b[39mset_transform(parse_maps_val)\n\u001b[0;32m---> 92\u001b[0m \u001b[43mgenerate_N_images\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpipeline\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mout_dir\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mn_images\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbatch_size\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mbsz\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minference_steps\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m10\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mval_dataset\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mval_dataset\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     95\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDONE\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
      "Cell \u001b[0;32mIn[10], line 53\u001b[0m, in \u001b[0;36mgenerate_N_images\u001b[0;34m(pipeline, out_dir, n_images, batch_size, inference_steps, val_dataset)\u001b[0m\n\u001b[1;32m     50\u001b[0m     \u001b[38;5;66;03m# save output images\u001b[39;00m\n\u001b[1;32m     51\u001b[0m     save_pipeline_output_to_disk(output, batch_idx, batch_size, out_dir)\n\u001b[0;32m---> 53\u001b[0m ue \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1\u001b[39m\n\u001b[1;32m     55\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m n_images\u001b[38;5;241m%\u001b[39mbatch_size \u001b[38;5;241m>\u001b[39m \u001b[38;5;241m0\u001b[39m:\n\u001b[1;32m     56\u001b[0m     batch_idx \u001b[38;5;241m=\u001b[39m n_images\u001b[38;5;241m/\u001b[39m\u001b[38;5;241m/\u001b[39mbatch_size\n",
      "\u001b[0;31mUnboundLocalError\u001b[0m: local variable 'ue' referenced before assignment"
     ]
    }
   ],
   "source": [
    "from PIL import Image\n",
    "from tqdm import tqdm\n",
    "import json\n",
    "from datasets import load_dataset\n",
    "from utils_maps import *\n",
    "\n",
    "def save_pipeline_output_to_disk(output, batch_idx, batch_size, out_dir):\n",
    "    images, cond_images = output\n",
    "    os.makedirs(out_dir + \"/ims\", exist_ok=True)\n",
    "    os.makedirs(out_dir + \"/cond_ims\", exist_ok=True)\n",
    "    for idx in range(min(batch_size, images.shape[0])):\n",
    "        local_d = {}\n",
    "        global_img_idx = batch_idx * batch_size + idx\n",
    "        img = Image.fromarray((images[idx]*255).astype(np.uint8))\n",
    "        img.save(out_dir + f\"/ims/{global_img_idx:05d}.png\")\n",
    "        if cond_images is not None:\n",
    "            #cond_images[idx] = (cond_images[idx] * 0.5) + 0.5\n",
    "            img = Image.fromarray((cond_images[idx]*255).astype(np.uint8))\n",
    "            img.save(out_dir + f\"/cond_ims/{global_img_idx:05d}.png\")\n",
    "\n",
    "def generate_N_images(pipeline, out_dir, n_images, batch_size=16, inference_steps=None, val_dataset=None):\n",
    "    pipeline.nodule_attributes = True\n",
    "    generator = torch.Generator(device=pipeline.device).manual_seed(0)\n",
    "    if inference_steps is None:\n",
    "        inference_steps = pipeline.scheduler.config.num_train_timesteps\n",
    "\n",
    "    im_size = pipeline.vae.sample_size\n",
    "    for batch_idx in tqdm(range(n_images//batch_size)):\n",
    "        # prepare condition images\n",
    "        if val_dataset is not None:\n",
    "            sample_indices = np.arange(batch_idx*batch_size, (batch_idx+1)*batch_size)\n",
    "            len_val = len(val_dataset)\n",
    "            sample_indices[sample_indices >= len_val] = sample_indices[sample_indices >= len_val] % len_val\n",
    "            input_condition_imgs = [val_dataset[int(i)][\"condition\"] for i in sample_indices]\n",
    "            input_condition_imgs = torch.stack(input_condition_imgs)\n",
    "        else:\n",
    "            input_condition_imgs = None\n",
    "        # run pipeline in inference (sample random noise and denoise)\n",
    "        output = pipeline(\n",
    "            input_condition_imgs=None if pipeline.args.unconditional else input_condition_imgs,\n",
    "            generator=generator,\n",
    "            batch_size=batch_size,\n",
    "            num_inference_steps=inference_steps,\n",
    "            output_type=\"numpy\",\n",
    "            return_dict=False\n",
    "        )\n",
    "        # save output images\n",
    "        save_pipeline_output_to_disk(output, batch_idx, batch_size, out_dir)\n",
    "        \n",
    "    if n_images%batch_size > 0:\n",
    "        batch_idx = n_images//batch_size\n",
    "        batch_idx = batch_idx + 1 if batch_idx > 0 else 0\n",
    "        # prepare condition images\n",
    "        if val_dataset is not None:\n",
    "            sample_indices = np.arange(batch_idx*batch_size, (batch_idx+1)*batch_size)\n",
    "            len_val = len(val_dataset)\n",
    "            sample_indices[sample_indices >= len_val] = sample_indices[sample_indices >= len_val] % len_val\n",
    "            input_condition_imgs = [val_dataset[int(i)][\"condition\"] for i in sample_indices]\n",
    "            input_condition_imgs = torch.stack(input_condition_imgs)\n",
    "            input_condition_imgs = input_condition_imgs[:n_images%batch_size]\n",
    "        else:\n",
    "            input_condition_imgs = None\n",
    "        # run pipeline in inference (sample random noise and denoise)\n",
    "        output = pipeline(\n",
    "            input_condition_imgs=None if pipeline.args.unconditional else input_condition_imgs,\n",
    "            generator=generator,\n",
    "            batch_size=n_images%batch_size,\n",
    "            num_inference_steps=inference_steps,\n",
    "            output_type=\"numpy\",\n",
    "            return_dict=False\n",
    "        )\n",
    "        # save output images\n",
    "        save_pipeline_output_to_disk(output, batch_idx, batch_size, out_dir)\n",
    "\n",
    "\n",
    "out_dir = model_path + \"/fid_synthetic2\"\n",
    "n_images = 2048\n",
    "bsz = 16\n",
    "\n",
    "val_dataset = load_dataset(pipeline.args.val_dataset_name,\n",
    "    pipeline.args.dataset_config_name,\n",
    "    cache_dir=pipeline.args.cache_dir,\n",
    "    split=\"train\",\n",
    "    )\n",
    "val_dataset.set_transform(parse_maps_val)\n",
    "\n",
    "generate_N_images(pipeline, out_dir, n_images, batch_size=bsz, inference_steps=1000, val_dataset=val_dataset)\n",
    "\n",
    "print(\"DONE\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "bb056cec",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2048/2048 [04:24<00:00,  7.75it/s]\n",
      "  0%|▌                                                                                                                                                      | 8/2048 [00:00<02:39, 12.78it/s]\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[12], line 24\u001b[0m\n\u001b[1;32m     22\u001b[0m out_p \u001b[38;5;241m=\u001b[39m img_p\u001b[38;5;241m.\u001b[39mreplace(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfid_synthetic\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfid_synthetic_2\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m     23\u001b[0m os\u001b[38;5;241m.\u001b[39mmakedirs(os\u001b[38;5;241m.\u001b[39mpath\u001b[38;5;241m.\u001b[39mdirname(out_p), exist_ok\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[0;32m---> 24\u001b[0m \u001b[43mImage\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfromarray\u001b[49m\u001b[43m(\u001b[49m\u001b[43mimg\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msave\u001b[49m\u001b[43m(\u001b[49m\u001b[43mout_p\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/anaconda3/envs/diffusion/lib/python3.9/site-packages/PIL/Image.py:2436\u001b[0m, in \u001b[0;36mImage.save\u001b[0;34m(self, fp, format, **params)\u001b[0m\n\u001b[1;32m   2434\u001b[0m         fp \u001b[38;5;241m=\u001b[39m builtins\u001b[38;5;241m.\u001b[39mopen(filename, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mr+b\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m   2435\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 2436\u001b[0m         fp \u001b[38;5;241m=\u001b[39m \u001b[43mbuiltins\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mopen\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfilename\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mw+b\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m   2438\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m   2439\u001b[0m     save_handler(\u001b[38;5;28mself\u001b[39m, fp, filename)\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "# rewrite stuff\n",
    "\n",
    "exp_id = \"2024-07-24_22-37-54_MLBriefs24_1_conditional_vanilla_encode_mask\"\n",
    "exp_id = \"2024-07-25_14-27-03_MLBriefs24_2_conditional_vanilla_downsample_mask\"\n",
    "exp_id = \"2024-08-02_14-35-03_MLBriefs24_5_conditional_CA_encodedmask\"\n",
    "exp_id = \"2024-08-01_10-12-17_MLBriefs24_4_conditional_CA_downsampledmask\"\n",
    "\n",
    "out_dir = model_path + \"/fid_synthetic\"\n",
    "\n",
    "for img_p in tqdm(sorted(glob.glob(out_dir + \"/cond_ims/*.png\"))):\n",
    "    img = np.array(Image.open(img_p)) / 255.\n",
    "    img = (img * 0.5) + 0.5\n",
    "    img = (img * 255).astype(np.uint8)\n",
    "    out_p = img_p.replace(\"fid_synthetic\", \"fid_synthetic_2\")\n",
    "    os.makedirs(os.path.dirname(out_p), exist_ok=True)\n",
    "    Image.fromarray(img).save(out_p)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1fa3352b",
   "metadata": {},
   "source": [
    "### Compute FID"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "d70b30aa",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1098/1098 [00:03<00:00, 327.03it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "real images loaded. shape is torch.Size([1098, 3, 256, 256])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1098/1098 [00:03<00:00, 332.47it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fake images loaded. shape is torch.Size([1098, 3, 256, 256])\n",
      "[FID] Computing activations from 1098 samples\n",
      "-3.487417836822715e-05\n"
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "from PIL import Image\n",
    "import torch\n",
    "from torchvision import transforms\n",
    "\n",
    "\n",
    "exp_id = \"2024-07-26_04-40-24_MLBriefs24_0_unconditional\"\n",
    "exp_id = \"2024-07-24_22-37-54_MLBriefs24_1_conditional_vanilla_encode_mask\"\n",
    "exp_id = \"2024-07-25_14-27-03_MLBriefs24_2_conditional_vanilla_downsample_mask\"\n",
    "exp_id = \"2024-08-02_14-35-03_MLBriefs24_5_conditional_CA_encodedmask\"\n",
    "exp_id = \"2024-08-01_10-12-17_MLBriefs24_4_conditional_CA_downsampledmask\"\n",
    "model_path = f\"/media/share/Datasets/diffusers_out/{exp_id}\"\n",
    "\n",
    "fake_images_dir = f\"{model_path}/fid_synthetic/ims\"\n",
    "fake_images_dir = \"/home/ubuntu/projects/phase-iv-ai/maps/val_ims\"\n",
    "real_images_dir = \"/home/ubuntu/projects/phase-iv-ai/maps/val_ims\"\n",
    "n_images = 2048\n",
    "\n",
    "def read_image_dir_for_FID_computation(input_dir, resolution=256, n_images=None,\n",
    "                                       img_extension=\".png\", split=\"train\"):\n",
    "    augmentations = transforms.Compose(\n",
    "        [\n",
    "            transforms.Resize(resolution, interpolation=transforms.InterpolationMode.BILINEAR),\n",
    "            #transforms.ToTensor(),\n",
    "            #transforms.Normalize([0.5], [0.5]),\n",
    "        ]\n",
    "    )\n",
    "    \n",
    "    assert os.path.exists(input_dir)\n",
    "    img_paths = sorted(glob.glob(os.path.join(input_dir, f\"*{img_extension}\")))\n",
    "    images = []\n",
    "    for p in tqdm(img_paths):\n",
    "        img = np.array(Image.open(p))\n",
    "        if split == \"train\":\n",
    "            w = img.shape[1]\n",
    "            img = img[:, :w//2, :]\n",
    "        img = torch.from_numpy(img).permute(2, 0, 1).type(torch.float32) / 255.\n",
    "        if img.shape[1] != resolution:\n",
    "            img = augmentations(img)\n",
    "        images.append(img)\n",
    "    images = torch.stack(images)\n",
    "    # we need at least 2048 images to compute the FID\n",
    "    #https://github.com/mseitzer/pytorch-fid/issues/13\n",
    "    if (n_images is not None) and (images.shape[0] < n_images):\n",
    "        diff = n_images - images.shape[0]\n",
    "        images = torch.cat([images, images[:diff]], dim=0)\n",
    "    return images\n",
    "\n",
    "# read real images\n",
    "real_images = read_image_dir_for_FID_computation(real_images_dir, img_extension=\".png\", split=\"real\")\n",
    "print(f\"real images loaded. shape is {real_images.shape}\")\n",
    "\n",
    "fake_images = read_image_dir_for_FID_computation(fake_images_dir, img_extension=\".png\", split=\"fake\")\n",
    "print(f\"fake images loaded. shape is {fake_images.shape}\")\n",
    "\n",
    "\n",
    "from fid import FID\n",
    "\n",
    "# requires !pip install scipy==1.9.1\n",
    "fid = FID(real_images.cuda(), device=\"cuda\")\n",
    "fid_score = fid.calculate_FID(fake_images.cuda())\n",
    "print(fid_score)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b437172",
   "metadata": {},
   "source": [
    "### Save images to disk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "76ac990e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ok\n"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "for idx, image in enumerate(fake_images):\n",
    "    im = image.permute(1,2,0).cpu().numpy()\n",
    "    Image.fromarray((im*255).astype(np.uint8)).save(f\"/home/ubuntu/projects/phase-iv-ai/maps/val_ims/{idx+1}.png\")\n",
    "    \n",
    "for idx, image in enumerate(real_images):\n",
    "    im = image.permute(1,2,0).cpu().numpy()\n",
    "    Image.fromarray((im*255).astype(np.uint8)).save(f\"/home/ubuntu/projects/phase-iv-ai/maps/train_ims/{idx+1}.png\")\n",
    "    \n",
    "print(\"ok\")\n",
    "\"\"\""
   ]
  }
 ],
 "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}