{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.10/dist-packages/torchvision/transforms/functional_tensor.py:5: UserWarning: The torchvision.transforms.functional_tensor module is deprecated in 0.15 and will be **removed in 0.17**. Please don't rely on it. You probably just need to use APIs in torchvision.transforms.functional or in torchvision.transforms.v2.functional.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "from handler import EndpointHandler\n",
    "import base64\n",
    "from io import BytesIO\n",
    "from PIL import Image\n",
    "import cv2\n",
    "import random\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# helper decoder\n",
    "def decode_base64_image(image_string):\n",
    "  base64_image = base64.b64decode(image_string)\n",
    "  buffer = BytesIO(base64_image)\n",
    "  return  Image.open(buffer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# init handler\n",
    "my_handler = EndpointHandler(path=\".\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "image.size: (1200, 517), image.mode: RGBA, outscale: 10.0\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "output.shape: (5170, 12000, 4)\n",
      "out_image.size: (12000, 5170)\n",
      "image.size: (1056, 1068), image.mode: RGB, outscale: 3.0\n",
      "output.shape: (3204, 3168, 3)\n",
      "out_image.size: (3168, 3204)\n",
      "image.size: (1056, 1068), image.mode: L, outscale: 5.49\n",
      "output.shape: (5863, 5797, 3)\n",
      "out_image.size: (5797, 5863)\n"
     ]
    }
   ],
   "source": [
    "img_dir = \"test_data/\"\n",
    "img_names = [\"4121783.png\", \"FB_IMG_1725931665635.jpg\", \"FB_IMG_1725931665635_gray.jpg\"]\n",
    "out_scales = [10, 3, 5.49]\n",
    "for img_name, outscale in zip(img_names, out_scales):\n",
    "    image_path = img_dir + img_name\n",
    "    # create payload\n",
    "    with open(image_path, \"rb\") as i:\n",
    "        b64 = base64.b64encode(i.read())\n",
    "        b64 = b64.decode(\"utf-8\")\n",
    "        payload = {\n",
    "            \"inputs\": {\"image\": b64, \n",
    "                    \"outscale\": outscale\n",
    "                    }\n",
    "            }\n",
    "\n",
    "\n",
    "    output_payload = my_handler(payload)\n",
    "    out_image = decode_base64_image(output_payload[\"out_image\"])\n",
    "    print(f\"out_image.size: {out_image.size}\")\n",
    "    out_image.save(f\"test_data/outputs/{img_name.split('.')[0]}_outscale_{outscale}.png\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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": 2
}