update the handler.py and test_handler file for bugs and error fixes

handler.py

@@ -1,3 +1,13 @@
+"""
+This module handles the endpoint for image upscaling using the Real-ESRGAN model.
+
+Required Environment Variables:
+- TILING_SIZE: The size of the tiles for processing images. Set to 0 to disable tiling.
+- AWS_ACCESS_KEY_ID: AWS access key for S3 access.
+- AWS_SECRET_ACCESS_KEY: AWS secret key for S3 access.
+- BUCKET_NAME: The name of the S3 bucket where images will be uploaded.
+
+"""
 import torch
 from PIL import Image
 from io import BytesIO
@@ -15,194 +25,259 @@ import torch
 import base64
 import requests
 import logging
+import time
 
 
 class EndpointHandler:
     def __init__(self, path=""):
+        """
+        Initializes the EndpointHandler class, setting up the Real-ESRGAN model and S3 client.
+
+        Args:
+            path (str): Optional path to the model weights. Defaults to an empty string.
         
-        self.tiling_size = int(os.environ["TILING_SIZE"])
+        This constructor performs the following actions:
+        - Configures logging based on environment variables.
+        - Retrieves the tiling size from environment variables.
+        - Initializes the Real-ESRGAN model with specified parameters, including scale, model path, and architecture.
+        - Sets up the S3 client using AWS credentials from environment variables.
+        - Logs the initialization process and any errors encountered during setup.
+        """
+        
+        # get the logging level from environment variables
+        logging.basicConfig(level=logging.INFO, format='%(levelname)s - %(message)s')
+        self.logger = logging.getLogger(__name__)
+
 
+        self.tiling_size = int(os.environ["TILING_SIZE"])
+        # self.model_path = f"/repository/weights/Real-ESRGAN-x4plus.pth"
+        self.max_image_size = 1400 * 1400
+        self.model_path = f"/workspace/real-esrgan/weights/Real-ESRGAN-x4plus.pth"
+        
+        
+        # log model path and tiling size
+        self.logger.info(f"model_path: {self.model_path}")
+        if self.tiling_size == 0: self.logger.info("TILING_SIZE is 0, not using tiling")
+        else: self.logger.info(f"TILING_SIZE is {self.tiling_size}, using tiling")
+        
+        
         # Initialize the Real-ESRGAN model with specified parameters
-        self.model = RealESRGANer(
-            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",
-            # Initialize the RRDBNet model architecture with specified parameters
-            model= RRDBNet(num_in_ch=3,
+        start_time = time.time()
+        self.logger.info(f"initializing model")
+        try:
+            self.model = RealESRGANer(
+                scale=4,  # Scale factor for the model
+                # Path to the pre-trained model weights
+                model_path=self.model_path,
+                # 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=self.tiling_size,
-            tile_pad=0,
-            half=True,
-        )
+                tile=self.tiling_size,
+                tile_pad=0,
+                half=True,
+            )
+            self.logger.info(f"model initialized in {time.time() - start_time} seconds")
+        except Exception as e:
+            self.logger.error(f"Error initializing model: {e}")
+            raise e
         
-        # Initialize the S3 client with AWS credentials from environment variables
-        self.s3 = boto3.client('s3', 
+        
+        try: 
+            # 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"]
+            # Get the S3 bucket name from environment variables
+            self.bucket_name = os.environ["S3_BUCKET_NAME"]
+        except Exception as e:
+            self.logger.error(f"Error initializing S3 client: {e}")
+            raise e
         
-        # 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]]:
-        
+        """
+        Processes the input data to upscale an image using the Real-ESRGAN model.
+
+        Args:
+            data (Any): A dictionary containing the input data. It should include:
+                - 'inputs': A dictionary with the following keys:
+                    - 'image_url' (str): The URL of the image to be upscaled.
+                    - 'outscale' (float): The scaling factor for the upscaling process.
+
+        Returns:
+            Dict[str, List[float]]: A dictionary containing the results of the upscaling process, which includes:
+                - 'image_url' (str | None): The URL of the upscaled image or None if an error occurred.
+                - 'image_key' (str | None): The key for the uploaded image in S3 or None if an error occurred.
+                - 'error' (str | None): An error message if an error occurred, otherwise None.
+        """
+
+        ############################################################
+        # get inputs and download image
+        ############################################################
+        self.logger.info(">>> 1/7: GETTING INPUTS....")
         try:
-            ############################################################
-            # 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}")
+            image_url = inputs["image_url"]
+        except Exception as e:
+            self.logger.error(f"Error getting inputs: {e}")
+            return {"image_url": None, "image_key": None, "error": f"Failed to get inputs: {e}"}
+        
+        # download image
+        try:
+            self.logger.info(f"downloading image from URL: {image_url}")
+            image = self.download_image_url(image_url)
+        except Exception as e:
+            self.logger.error(f"Error downloading image from URL: {image_url}. Exception: {e}")
+            return {"image_url": None, "image_key": None, "error": f"Failed to download image: {e}"}
+        
+        
+        ############################################################
+        # run assertions
+        ############################################################
+        self.logger.info(">>> 2/7: RUNNING ASSERTIONS ON IMAGE....")
+        
+        # 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
+        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] <  self.max_image_size, f"Image is too large: {in_size}: {in_size[0] * in_size[1]} is greater than {self.max_image_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 {"image_url": None, "image_key": None, "error": str(e)}
+        
+        
+        ############################################################
+        # Convert RGB to BGR (PIL uses RGB, OpenCV expects BGR)
+        ############################################################
+        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 {"image_url": None, "image_key": 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:
+            self.logger.error(f"Unsupported image mode: {in_mode}")
+            return {"image_url": None, "image_key": None, "error": f"Unsupported image mode: {in_mode}"}
             
-            # 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}"}
+        
+        ############################################################
+        # 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 {"image_url": None, "image_key": None, "error": "Image enhancement failed."}
+        # debug
+        self.logger.info(f"output.shape: {output.shape}")
 
-            
-            ############################################################
-            # run assertions
-            ############################################################
-            self.logger.info(">>> 2/7: RUNNING ASSERTIONS ON IMAGE....")
-            
-            # 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
-            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)}
-            
-            
-            ############################################################
-            # Convert RGB to BGR (PIL uses RGB, OpenCV expects BGR)
-            ############################################################
-            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:
-                self.logger.error(f"Unsupported image mode: {in_mode}")
-                return {"out_image": None, "error": f"Unsupported image mode: {in_mode}"}
-            
-            
-            ############################################################
-            # 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
-            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:
-                    output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
-                elif out_shape[2] == 4:
-                    output = cv2.cvtColor(output, cv2.COLOR_BGRA2RGBA)
-            else:
-                output = cv2.cvtColor(output, cv2.COLOR_GRAY2RGB)
-            
-            
-            ############################################################
-            # convert to PIL image
-            ############################################################
-            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}"}
-            
-            
-            ############################################################
-            # 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
-                    }
+
+        ############################################################
+        # 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:
+                output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
+            elif out_shape[2] == 4:
+                output = cv2.cvtColor(output, cv2.COLOR_BGRA2RGBA)
+        else:
+            output = cv2.cvtColor(output, cv2.COLOR_GRAY2RGB)
+        
+        
+        ############################################################
+        # convert to PIL image
+        ############################################################
+        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 {"image_url": None, "image_key": None, "error": f"Failed to convert upscaled image to PIL: {e}"}
         
-        # handle unexpected errors
+        
+        ############################################################
+        # 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"An unexpected error occurred: {e}")
-            return {"out_image": None, "error": f"An unexpected error occurred: {e}"}
+            self.logger.error(f"Error uploading image to s3: {e}")
+            return {"image_url": None, "image_key": None, "error": f"Failed to upload image to s3: {e}"}
+        
+        
+        return {"image_url": image_url,
+                "image_key": key,
+                "error": None
+                }
+        
 
+    def upload_to_s3(self, image: Image.Image) -> tuple[str, str]:
+        """
+        Upload the image to S3 and return the URL and key.
+        
+        Args:
+            image (Image.Image): The image to upload.
 
-    def upload_to_s3(self, image):
-        "Upload the image to s3 and return the url."
+        Returns:
+            tuple[str, str]: A tuple containing the image URL and the S3 key.
+        """
         
         prefix = str(uuid.uuid4())
         # Save the image to an in-memory file
         in_mem_file = io.BytesIO()
-        image.save(in_mem_file, 'PNG')
+        image.save(in_mem_file, format='PNG')
         in_mem_file.seek(0)
 
-        # Upload the image to s3
+        # 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)
+        image_url = f"https://{self.bucket_name}.s3.amazonaws.com/{key}"
         
-        # return the url and the key    
+        # Return the URL and the key    
         return image_url, key
     
-    def download_image_url(self, image_url):
-        "Download the image from the url and return the image."
-        
+    def download_image_url(self, image_url: str) -> Image.Image:
+        """
+        Downloads an image from the specified URL and returns it as a PIL Image.
+
+        Args:
+            image_url (str): The URL of the image to download.
+
+        Returns:
+            Image.Image: The downloaded image as a PIL Image.
+        """
         response = requests.get(image_url)
         image = Image.open(BytesIO(response.content))
         return image

test_handler.ipynb

@@ -40,7 +40,18 @@
    "cell_type": "code",
    "execution_count": 3,
    "metadata": {},
-   "outputs": [],
+   "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 1.977891206741333 seconds\n"
+     ]
+    }
+   ],
    "source": [
     "# init handler\n",
     "my_handler = EndpointHandler(path=\".\")"
@@ -67,14 +78,14 @@
       "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"
+      "INFO - image uploaded to s3: https://jiffy-staging-upscaled-images.s3.amazonaws.com/25b91e15-b785-47ca-81a6-ad5fbdf8b92a.png\n"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "https://upscale-process-results.s3.amazonaws.com/fe21c683-ee2d-4e1d-9fbc-af56823c664c.png fe21c683-ee2d-4e1d-9fbc-af56823c664c.png\n"
+      "https://jiffy-staging-upscaled-images.s3.amazonaws.com/25b91e15-b785-47ca-81a6-ad5fbdf8b92a.png 25b91e15-b785-47ca-81a6-ad5fbdf8b92a.png\n"
      ]
     }
    ],