Write the output image to S3 bucket and return image_url and filename

handler.py

@@ -9,8 +9,8 @@ from basicsr.archs.rrdbnet_arch import RRDBNet
 import numpy as np
 import cv2
 import PIL
-
-
+import boto3
+import uuid, io
 import torch
 import base64
 
@@ -18,24 +18,31 @@ import base64
 class EndpointHandler:
     def __init__(self, path=""):
 
+        # Initialize the Real-ESRGAN model with specified parameters
         self.model = RealESRGANer(
-            scale=4,
+            scale=4,  # Scale factor for the model
+            # Path to the pre-trained model weights
             model_path=f"/repository/weights/Real-ESRGAN-x4plus.pth",
-            # model_path=f"/workspace/real-esrgan/weights/Real-ESRGAN-x4plus.pth",
-            # dni_weight=dni_weight,
-            model= RRDBNet(num_in_ch=3, 
-                           num_out_ch=3, 
-                           num_feat=64, 
-                           num_block=23, 
-                           num_grow_ch=32, 
+            # Initialize the RRDBNet model architecture with specified parameters
+            model= RRDBNet(num_in_ch=3,
+                           num_out_ch=3,
+                           num_feat=64,
+                           num_block=23,
+                           num_grow_ch=32,
                            scale=4
                            ),
             tile=1000,
             tile_pad=20,
-            # pre_pad=args.pre_pad,
             half=True,
-            # gpu_id=args.gpu_id
         )
+        
+        # Initialize the S3 client with AWS credentials from environment variables
+        self.s3 = boto3.client('s3', 
+                               aws_access_key_id=os.environ['AWS_ACCESS_KEY_ID'],
+                               aws_secret_access_key=os.environ['AWS_SECRET_ACCESS_KEY'],
+                               )
+        # Get the S3 bucket name from environment variables
+        self.bucket_name = os.environ["S3_BUCKET_NAME"]
 
     def __call__(self, data: Any) -> Dict[str, List[float]]:
         
@@ -89,12 +96,14 @@ class EndpointHandler:
             img_byte_arr = BytesIO()
             output = Image.fromarray(output)
             
-            # save to BytesIO
-            output.save(img_byte_arr, format='PNG')
-            img_str = base64.b64encode(img_byte_arr.getvalue())
-            img_str = img_str.decode()
+            # # save to BytesIO
+            # output.save(img_byte_arr, format='PNG')
+            # img_str = base64.b64encode(img_byte_arr.getvalue())
+            # img_str = img_str.decode()
+            image_url, key = self.upload_to_s3(output)
             
-            return {"out_image": img_str, 
+            return {"image_url": image_url,
+                    "image_key": key,
                     "error": None
                     }
         
@@ -114,3 +123,20 @@ class EndpointHandler:
         except Exception as e:
             print(f"Exception: {e}")
             return {"out_image": None, "error": "An unexpected error occurred"}
+
+    def upload_to_s3(self, image):
+        "Upload the image to s3 and return the url."
+        
+        prefix = str(uuid.uuid4())
+        # Save the image to an in-memory file
+        in_mem_file = io.BytesIO()
+        image.save(in_mem_file, 'PNG')
+        in_mem_file.seek(0)
+
+        # Upload the image to s3
+        key = f"{prefix}.png"
+        self.s3.upload_fileobj(in_mem_file, Bucket=self.bucket_name, Key=key)
+        image_url = "https://{0}.s3.amazonaws.com/{1}".format(self.bucket_name, key)
+        
+        # return the url and the key    
+        return image_url, key

requirements.txt

@@ -7,4 +7,5 @@ Pillow
 torch>=1.7
 torchvision==0.16.2
 tqdm
-realesrgan
+realesrgan
+boto3

test_handler.ipynb

@@ -20,7 +20,7 @@
     "from io import BytesIO\n",
     "from PIL import Image\n",
     "import cv2\n",
-    "import random\n"
+    "import random"
    ]
   },
   {
@@ -48,7 +48,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [
     {
@@ -62,14 +62,11 @@
      "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"
+      "\tTile 1/2\n",
+      "\tTile 2/2\n",
+      "\tTile 1/2\n",
+      "\tTile 2/2\n",
+      "output.shape: (5170, 12000, 4)\n"
      ]
     }
    ],
@@ -91,9 +88,10 @@
     "\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"
+    "    # 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",
+    "    break"
    ]
   },
   {
@@ -106,7 +104,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },
@@ -124,5 +122,5 @@
   }
  },
  "nbformat": 4,
- "nbformat_minor": 2
+ "nbformat_minor": 4
 }