{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c1da833-fccc-47f7-b4b4-8f4939f1551e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import imageio\n",
    "from scipy.fft import fft2, ifft2, fftshift, fftfreq, ifftshift"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ee617c5-0765-4ae0-8dee-8d0777c47135",
   "metadata": {},
   "outputs": [],
   "source": [
    "def gaussian_filter(k=5, sigma=1.0):\n",
    "    radius = k // 2\n",
    "    samples = np.arange(-radius, radius+1)\n",
    "\n",
    "    x, y = np.meshgrid(samples, samples)\n",
    "    kernel = np.exp(-(x**2 + y**2) / (2 * (sigma**2)))\n",
    "    return kernel/kernel.sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8095226-562c-4a7d-817f-15813214c18d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def convolve_in_freq_domain(image, kernel):\n",
    "    radius = kernel.shape[0] // 2\n",
    "    \n",
    "    image_pad = np.pad(image, ((radius, radius), (radius, radius)), 'reflect')\n",
    "    P, Q = image_pad.shape\n",
    "    fourier_image = fftshift(fft2(image_pad, s=(P, Q)))\n",
    "    filter_f = fftshift(fft2(kernel, s=(P, Q)))\n",
    "    filtered_f = fourier_image * (filter_f)\n",
    "    filtered_image = ifft2(ifftshift(filtered_f))\n",
    "\n",
    "    return filtered_image.real"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24c4bd2a-42c4-4f6b-a2d0-308cc956a016",
   "metadata": {},
   "outputs": [],
   "source": [
    "img = imageio.imread('ball.jpeg', pilmode='L')\n",
    "plt.imshow(img, cmap='gray')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c205f52-2fac-4ec6-af84-3b113e5d3a69",
   "metadata": {},
   "outputs": [],
   "source": [
    "h = gaussian_filter(k=13, sigma=2.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7a46de4-398c-45cf-adfd-e896ef363a08",
   "metadata": {},
   "outputs": [],
   "source": [
    "img_f = convolve_in_freq_domain(img, h)\n",
    "plt.imshow(img_f, cmap='gray')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34b01d56-981a-4e93-9d8d-db92b63c7e7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def inverse_filter(image, kernel):\n",
    "    radius = kernel.shape[0] // 2\n",
    "    \n",
    "    image_pad = np.pad(image, ((radius, radius), (radius, radius)), 'reflect')\n",
    "    P, Q = image_pad.shape\n",
    "    fourier_image = fftshift(fft2(image_pad, s=(P, Q)))\n",
    "    filter_f = fftshift(fft2(kernel, s=(P, Q)))\n",
    "    filtered_f = fourier_image / (filter_f)\n",
    "    filtered_image = ifft2(ifftshift(filtered_f))\n",
    "\n",
    "    return filtered_image.real"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f783e88-c3dc-4a8b-8b04-30f5cc54a807",
   "metadata": {},
   "outputs": [],
   "source": [
    "img_f_fixed = inverse_filter(img_f, h)\n",
    "plt.imshow(img_f_fixed, cmap='gray')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a92d961f-6e7f-4065-b148-7a532b8b5fb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "img_noisy = img + np.random.randint(10, size=img.shape)\n",
    "plt.imshow(img_noisy, cmap='gray')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1cebef39-2932-4cd3-9427-32daa79665ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "img_f = convolve_in_freq_domain(img_noisy, h)\n",
    "plt.imshow(img_f, cmap='gray')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7be27978-02df-41f1-b685-dee7ffcd7623",
   "metadata": {},
   "outputs": [],
   "source": [
    "img_f_fixed = inverse_filter(img_f, h)\n",
    "plt.imshow(img_f_fixed, cmap='gray')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d9241e9-d513-4da7-bbe2-bab9fcaea05a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def pseudoinverse_filter(image, kernel, l):\n",
    "    radius = kernel.shape[0] // 2\n",
    "    \n",
    "    image_pad = np.pad(image, ((radius, radius), (radius, radius)), 'reflect')\n",
    "    P, Q = image_pad.shape\n",
    "    fourier_image = fftshift(fft2(image_pad, s=(P, Q)))\n",
    "    filter_f = fftshift(fft2(kernel, s=(P, Q)))\n",
    "    filter_f[np.abs(filter_f) <= l] = l\n",
    "    filtered_f = fourier_image / (filter_f)\n",
    "    filtered_image = ifft2(ifftshift(filtered_f))\n",
    "\n",
    "    return filtered_image.real"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc06498a-82da-4e6f-8aba-11198988807f",
   "metadata": {},
   "outputs": [],
   "source": [
    "img_f_fixed = pseudoinverse_filter(img_f, h, 0.0001)\n",
    "plt.imshow(img_f_fixed, cmap='gray')"
   ]
  }
 ],
 "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.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
