{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/workspace/miniconda3/envs/real-esrgan-hf/lib/python3.10/site-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": [
    "import os\n",
    "from handler import EndpointHandler\n",
    "import base64\n",
    "from io import BytesIO\n",
    "from PIL import Image\n",
    "import cv2\n",
    "import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"AWS_ACCESS_KEY_ID\"] = \"\"\n",
    "os.environ[\"AWS_SECRET_ACCESS_KEY\"] = \"\"\n",
    "os.environ[\"S3_BUCKET_NAME\"] = \"\"\n",
    "os.environ[\"TILING_SIZE\"] = \"0\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO - model_path: /workspace/real-esrgan/weights/Real-ESRGAN-x4plus.pth\n",
      "INFO - TILING_SIZE is 0, not using tiling\n",
      "INFO - initializing model\n",
      "INFO - model initialized in 90.91657543182373 seconds\n"
     ]
    }
   ],
   "source": [
    "# init handler\n",
    "my_handler = EndpointHandler(path=\".\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO - >>> 1/8: GETTING INPUTS....\n",
      "INFO - outscale: 2.3\n",
      "INFO - apply_thresholding: True\n",
      "INFO - downloading image from URL: https://jiffy-staging-transfers.imgix.net/2/ai-research/test-ghost-pixels/10816322.png\n",
      "INFO - >>> 2/8: RUNNING ASSERTIONS ON IMAGE....\n",
      "INFO - image.size: (1080, 1080), image.mode: RGBA\n",
      "INFO - >>> 3/8: CONVERTING IMAGE TO OPENCV BGR/BGRA FORMAT....\n",
      "INFO - converting RGBA image to BGRA\n",
      "INFO - >>> 4/8: UPSCALING IMAGE....\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "https://jiffy-staging-transfers.imgix.net/2/ai-research/test-ghost-pixels/10816322.png 2.3\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO - output.shape: (2484, 2484, 4)\n",
      "INFO - >>> 5/8: CONVERTING IMAGE TO RGB/RGBA FORMAT....\n",
      "INFO - >>> 6/8: CONVERTING IMAGE TO PIL....\n",
      "INFO - >>> 7/8: APPLYING THRESHOLDING....\n",
      "INFO - input image mode: RGBA, it has alpha channel, applying thresholding\n",
      "INFO - Applying edge-aware clamping\n",
      "INFO - Input image size: (1080, 1080)\n",
      "INFO - Edge-aware clamping applied, image size: (2484, 2484)\n",
      "INFO - >>> 8/8: UPLOADING IMAGE TO S3....\n",
      "INFO - image uploaded to s3: https://jiffy-staging-upscaled-images.s3.amazonaws.com/0cd7ea5b-a6fc-4cbb-a536-57d9dbcc12b4.png\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "https://jiffy-staging-upscaled-images.s3.amazonaws.com/0cd7ea5b-a6fc-4cbb-a536-57d9dbcc12b4.png 0cd7ea5b-a6fc-4cbb-a536-57d9dbcc12b4.png\n"
     ]
    }
   ],
   "source": [
    "img_dir = \"test_data/\"\n",
    "img_urls = [#\"https://jiffy-transfers.imgix.net/2/attachments/r267odvvfmkp6c5lccj1y6f9trb0\",\n",
    "#             \"https://jiffy-staging-transfers.imgix.net/2/development/attachments/zo31eau0ykhbwoddrjtlbyz6w9mp\", # larger than > 1.96M pixels\n",
    "#             \"https://jiffy-staging-transfers.imgix.net/2/development/attachments/b8ecchms9rr9wk3g71kfpfprqg1v\", # larger than > 1.96M pixels,\n",
    "#             \"https://jiffy-staging-transfers.imgix.net/2/ai-research/test-ghost-pixels/9B500B09.png\", # thresholding example 1\n",
    "            \"https://jiffy-staging-transfers.imgix.net/2/ai-research/test-ghost-pixels/10816322.png\"  # thresholding example 2\n",
    "\n",
    "           ]\n",
    "\n",
    "out_scales = [#4, \n",
    "              #3, \n",
    "              #2, \n",
    "              #2,\n",
    "              2.3\n",
    "]\n",
    "\n",
    "for img_url, outscale in zip(img_urls, out_scales):\n",
    "    print(img_url, outscale)\n",
    "    # create payload\n",
    "    payload = {\n",
    "        \"inputs\": {\"image_url\": img_url, \n",
    "                    \"outscale\": outscale,\n",
    "                   \"apply_thresholding\": True,\n",
    "                    }\n",
    "            }\n",
    "    \n",
    "    output_payload = my_handler(payload)\n",
    "    print(output_payload[\"image_url\"], output_payload[\"image_key\"])\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}