Update the codebase to log intermediate steps and to add try-except blocks

handler.py

@@ -14,6 +14,7 @@ import uuid, io
 import torch
 import base64
 import requests
+import logging
 
 
 class EndpointHandler:
@@ -48,48 +49,96 @@ class EndpointHandler:
         # Get the S3 bucket name from environment variables
         self.bucket_name = os.environ["S3_BUCKET_NAME"]
         
+        # get the logging level from environment variables
+        logging.basicConfig(level=logging.INFO, format='%(levelname)s - %(message)s')
+        self.logger = logging.getLogger(__name__)
+        
         
     def __call__(self, data: Any) -> Dict[str, List[float]]:
         
         try:
-            
-            # get inputs
+            ############################################################
+            # get inputs and download image
+            ############################################################
+            self.logger.info(">>> 1/7: GETTING INPUTS....")
             inputs = data.pop("inputs", data)
             
             # get outscale
             outscale = float(inputs.pop("outscale", 3))
+            self.logger.info(f"outscale: {outscale}")
+            
+            # download image
+            try:
+                self.logger.info(f"downloading image from URL: {inputs['image_url']}")
+                image = self.download_image_url(inputs['image_url'])
+            except Exception as e:
+                logging.error(f"Error downloading image from URL: {inputs['image_url']}. Exception: {e}")
+                return {"out_image": None, "error": f"Failed to download image: {e}"}
+
+            
+            ############################################################
+            # run assertions
+            ############################################################
+            self.logger.info(">>> 2/7: RUNNING ASSERTIONS ON IMAGE....")
             
-            # decode base64 image to PIL
-            image = self.download_image_url(inputs['image_url'])
+            # get image size and mode
             in_size, in_mode = image.size, image.mode
+            self.logger.info(f"image.size: {image.size}, image.mode: {image.mode}")
             
             # check image size and mode and return dict
-            assert in_mode in ["RGB", "RGBA", "L"], f"Unsupported image mode: {in_mode}"
-            if self.tiling_size == 0:
-                assert in_size[0] * in_size[1] <  1400*1400, f"Image is too large: {in_size}: {in_size[0] * in_size[1]} is greater than {self.tiling_size*self.tiling_size}"
-            assert outscale > 1 and outscale <=10, f"Outscale must be between 1 and 10: {outscale}"
+            try:
+                assert in_mode in ["RGB", "RGBA", "L"], f"Unsupported image mode: {in_mode}"
+                if self.tiling_size == 0:
+                    assert in_size[0] * in_size[1] <  1400*1400, f"Image is too large: {in_size}: {in_size[0] * in_size[1]} is greater than {self.tiling_size*self.tiling_size}"
+                assert outscale > 1 and outscale <= 10, f"Outscale must be between 1 and 10: {outscale}"
+            except AssertionError as e:
+                self.logger.error(f"Assertion error: {e}")
+                return {"out_image": None, "error": str(e)}
             
-            # debug
-            print(f"image.size: {in_size}, image.mode: {in_mode}, outscale: {outscale}")
             
+            ############################################################
             # Convert RGB to BGR (PIL uses RGB, OpenCV expects BGR)
-            opencv_image = np.array(image)
+            ############################################################
+            self.logger.info(f">>> 3/7: CONVERTING IMAGE TO OPENCV BGR/BGRA FORMAT....")
+            try:
+                opencv_image = np.array(image)
+            except Exception as e:
+                self.logger.error(f"Error converting image to opencv format: {e}")
+                return {"out_image": None, "error": f"Failed to convert image to opencv format: {e}"}
+            
+            # convert image to BGR
             if in_mode == "RGB":
+                self.logger.info(f"converting RGB image to BGR")
                 opencv_image = cv2.cvtColor(opencv_image, cv2.COLOR_RGB2BGR)
             elif in_mode == "RGBA":
+                self.logger.info(f"converting RGBA image to BGRA")
                 opencv_image = cv2.cvtColor(opencv_image, cv2.COLOR_RGBA2BGRA)
             elif in_mode == "L":
+                self.logger.info(f"converting grayscale image to BGR")
                 opencv_image = cv2.cvtColor(opencv_image, cv2.COLOR_GRAY2RGB)
             else:
-                raise ValueError(f"Unsupported image mode: {in_mode}")
+                self.logger.error(f"Unsupported image mode: {in_mode}")
+                return {"out_image": None, "error": f"Unsupported image mode: {in_mode}"}
+            
             
-            # enhance image
-            output, _ = self.model.enhance(opencv_image, outscale=outscale)
+            ############################################################
+            # upscale image
+            ############################################################
+            self.logger.info(f">>> 4/7: UPSCALING IMAGE....")
             
+            try:
+                output, _ = self.model.enhance(opencv_image, outscale=outscale)
+            except Exception as e:
+                self.logger.error(f"Error enhancing image: {e}")
+                return {"out_image": None, "error": "Image enhancement failed."}
             # debug
-            print(f"output.shape: {output.shape}")
+            self.logger.info(f"output.shape: {output.shape}")
+
 
+            ############################################################
             # convert to RGB/RGBA format
+            ############################################################
+            self.logger.info(f">>> 5/7: CONVERTING IMAGE TO RGB/RGBA FORMAT....")
             out_shape = output.shape
             if len(out_shape) == 3:
                 if out_shape[2] == 3:
@@ -99,37 +148,40 @@ class EndpointHandler:
             else:
                 output = cv2.cvtColor(output, cv2.COLOR_GRAY2RGB)
             
+            
+            ############################################################
             # convert to PIL image
-            img_byte_arr = BytesIO()
-            output = Image.fromarray(output)
+            ############################################################
+            self.logger.info(f">>> 6/7: CONVERTING IMAGE TO PIL....")
+            try:
+                img_byte_arr = BytesIO()
+                output = Image.fromarray(output)
+            except Exception as e:
+                self.logger.error(f"Error converting upscaled image to PIL: {e}")
+                return {"out_image": None, "error": f"Failed to convert upscaled image to PIL: {e}"}
+            
             
-            # # 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)
+            ############################################################
+            # upload to s3
+            ############################################################
+            self.logger.info(f">>> 7/7: UPLOADING IMAGE TO S3....")
+            try:
+                image_url, key = self.upload_to_s3(output)
+                self.logger.info(f"image uploaded to s3: {image_url}")
+            except Exception as e:
+                self.logger.error(f"Error uploading image to s3: {e}")
+                return {"out_image": None, "error": f"Failed to upload image to s3: {e}"}
             
             return {"image_url": image_url,
                     "image_key": key,
                     "error": None
                     }
         
-        # handle errors
-        except AssertionError as e:
-            print(f"AssertionError: {e}")
-            return {"out_image": None, "error": str(e)}
-        except KeyError as e:
-            print(f"KeyError: {e}")
-            return {"out_image": None, "error": f"Missing key: {e}"}
-        except ValueError as e:
-            print(f"ValueError: {e}")
-            return {"out_image": None, "error": str(e)}
-        except PIL.UnidentifiedImageError as e:
-            print(f"PIL.UnidentifiedImageError: {e}")
-            return {"out_image": None, "error": "Invalid image format"}
+        # handle unexpected errors
         except Exception as e:
-            print(f"Exception: {e}")
-            return {"out_image": None, "error": "An unexpected error occurred"}
+            self.logger.error(f"An unexpected error occurred: {e}")
+            return {"out_image": None, "error": f"An unexpected error occurred: {e}"}
+
 
     def upload_to_s3(self, image):
         "Upload the image to s3 and return the url."

test_handler.ipynb

@@ -33,7 +33,7 @@
     "os.environ[\"AWS_ACCESS_KEY_ID\"] = \"\"\n",
     "os.environ[\"AWS_SECRET_ACCESS_KEY\"] = \"\"\n",
     "os.environ[\"S3_BUCKET_NAME\"] = \"\"\n",
-    "os.environ[\"TILING_SIZE\"] = \"1000\""
+    "os.environ[\"TILING_SIZE\"] = \"0\""
    ]
   },
   {
@@ -52,22 +52,29 @@
    "metadata": {},
    "outputs": [
     {
-     "name": "stdout",
+     "name": "stderr",
      "output_type": "stream",
      "text": [
-      "image.size: (1024, 1024), image.mode: RGB, outscale: 4.0\n"
+      "INFO - >>> 1/7: GETTING INPUTS....\n",
+      "INFO - outscale: 4.0\n",
+      "INFO - downloading image from URL: https://jiffy-transfers.imgix.net/2/attachments/r267odvvfmkp6c5lccj1y6f9trb0?ixlib=rb-0.3.5\n",
+      "INFO - >>> 2/7: RUNNING ASSERTIONS ON IMAGE....\n",
+      "INFO - image.size: (1024, 1024), image.mode: RGB\n",
+      "INFO - >>> 3/7: CONVERTING IMAGE TO OPENCV BGR/BGRA FORMAT....\n",
+      "INFO - converting RGB image to BGR\n",
+      "INFO - >>> 4/7: UPSCALING IMAGE....\n",
+      "INFO - output.shape: (4096, 4096, 3)\n",
+      "INFO - >>> 5/7: CONVERTING IMAGE TO RGB/RGBA FORMAT....\n",
+      "INFO - >>> 6/7: CONVERTING IMAGE TO PIL....\n",
+      "INFO - >>> 7/7: UPLOADING IMAGE TO S3....\n",
+      "INFO - image uploaded to s3: https://upscale-process-results.s3.amazonaws.com/fe21c683-ee2d-4e1d-9fbc-af56823c664c.png\n"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "\tTile 1/4\n",
-      "\tTile 2/4\n",
-      "\tTile 3/4\n",
-      "\tTile 4/4\n",
-      "output.shape: (4096, 4096, 3)\n",
-      "https://jiffy-staging-upscaled-images.s3.amazonaws.com/d91323cb-0801-45b7-8109-9739212037ed.png d91323cb-0801-45b7-8109-9739212037ed.png\n"
+      "https://upscale-process-results.s3.amazonaws.com/fe21c683-ee2d-4e1d-9fbc-af56823c664c.png fe21c683-ee2d-4e1d-9fbc-af56823c664c.png\n"
      ]
     }
    ],
@@ -78,7 +85,7 @@
     "            # \"https://jiffy-staging-transfers.imgix.net/2/development/attachments/b8ecchms9rr9wk3g71kfpfprqg1v?ixlib=rb-0.3.5\" # larger than > 1.96M pixels\n",
     "           ]\n",
     "\n",
-    "out_scales = [4, 3, 2]\n",
+    "out_scales = [4]#, 3, 2]\n",
     "for img_url, outscale in zip(img_urls, out_scales):\n",
     "    # create payload\n",
     "    payload = {\n",
@@ -88,7 +95,8 @@
     "            }\n",
     "    \n",
     "    output_payload = my_handler(payload)\n",
-    "    print(output_payload[\"image_url\"], output_payload[\"image_key\"])\n"
+    "    print(output_payload[\"image_url\"], output_payload[\"image_key\"])\n",
+    "    \n"
    ]
   },
   {