{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8d2437e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from PIL import Image\n",
    "from RealESRGAN import RealESRGAN\n",
    "import gradio as gr\n",
    "import numpy as np\n",
    "import tempfile\n",
    "import time\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "871d9b94",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_model(scale):\n",
    "    model = RealESRGAN(device, scale=scale)\n",
    "    weights_path = f'weights/RealESRGAN_x{scale}.pth'\n",
    "    try:\n",
    "        model.load_weights(weights_path, download=True)\n",
    "        print(f\"Weights for scale {scale} loaded successfully.\")\n",
    "    except Exception as e:\n",
    "        print(f\"Error loading weights for scale {scale}: {e}\")\n",
    "        model.load_weights(weights_path, download=False)\n",
    "    return model\n",
    "\n",
    "model2 = load_model(2)\n",
    "model4 = load_model(4)\n",
    "model8 = load_model(8)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c891d4b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def enhance_image(image, scale):\n",
    "    try:\n",
    "        print(f\"Enhancing image with scale {scale}...\")\n",
    "        start_time = time.time()\n",
    "        image_np = np.array(image.convert('RGB'))\n",
    "        print(f\"Image converted to numpy array: shape {image_np.shape}, dtype {image_np.dtype}\")\n",
    "        \n",
    "        if scale == '2x':\n",
    "            result = model2.predict(image_np)\n",
    "        elif scale == '4x':\n",
    "            result = model4.predict(image_np)\n",
    "        else:\n",
    "            result = model8.predict(image_np)\n",
    "            \n",
    "        enhanced_image = Image.fromarray(np.uint8(result))\n",
    "        print(f\"Image enhanced in {time.time() - start_time:.2f} seconds\")\n",
    "        return enhanced_image\n",
    "    except Exception as e:\n",
    "        print(f\"Error enhancing image: {e}\")\n",
    "        return image\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9073bff6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def muda_dpi(input_image, dpi):\n",
    "    dpi_tuple = (dpi, dpi)\n",
    "    image = Image.fromarray(input_image.astype('uint8'), 'RGB')\n",
    "    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png')\n",
    "    image.save(temp_file, format='PNG', dpi=dpi_tuple)\n",
    "    temp_file.close()\n",
    "    return Image.open(temp_file.name)\n",
    "\n",
    "def resize_image(input_image, width, height):\n",
    "    image = Image.fromarray(input_image.astype('uint8'), 'RGB')\n",
    "    resized_image = image.resize((width, height))\n",
    "    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png')\n",
    "    resized_image.save(temp_file, format='PNG')\n",
    "    temp_file.close()\n",
    "    return Image.open(temp_file.name)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e470926d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_image(input_image, enhance, scale, adjust_dpi, dpi, resize, width, height):\n",
    "    original_image = Image.fromarray(input_image.astype('uint8'), 'RGB')\n",
    "    \n",
    "    if enhance:\n",
    "        original_image = enhance_image(original_image, scale)\n",
    "    \n",
    "    if adjust_dpi:\n",
    "        original_image = muda_dpi(np.array(original_image), dpi)\n",
    "        \n",
    "    if resize:\n",
    "        original_image = resize_image(np.array(original_image), width, height)\n",
    "        \n",
    "    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png')\n",
    "    original_image.save(temp_file.name)\n",
    "    return original_image, temp_file.name\n",
    "\n",
    "iface = gr.Interface(\n",
    "    fn=process_image,\n",
    "    inputs=[\n",
    "        gr.Image(label=\"Upload\"),\n",
    "        gr.Checkbox(label=\"Enhance Image\"),\n",
    "        gr.Radio(['2x', '4x', '8x'], type=\"value\", value='2x', label='Select Resolution model'),\n",
    "        gr.Checkbox(label=\"Apply DPI\"),\n",
    "        gr.Number(label=\"DPI\", value=300),\n",
    "        gr.Checkbox(label=\"Apply Resize\"),\n",
    "        gr.Number(label=\"Width\", value=512),\n",
    "        gr.Number(label=\"Height\", value=512)\n",
    "    ],\n",
    "    outputs=[\n",
    "        gr.Image(label=\"Final Image\"),\n",
    "        gr.File(label=\"Download Final Image\")\n",
    "    ],\n",
    "    title=\"Image Enhancer\",\n",
    "    description=\"Sorry for the inconvenience. The model is currently running on the CPU, which might affect performance. We appreciate your understanding.\",\n",
    "    theme=\"Yntec/HaleyCH_Theme_Orange\"\n",
    ")\n",
    "\n",
    "iface.launch(debug=True)\n"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 5
}