handler.py

import torch
from PIL import Image
from io import BytesIO
from realesrgan import RealESRGANer
from typing import Dict, List, Any
import os
from pathlib import Path
from basicsr.archs.rrdbnet_arch import RRDBNet
import numpy as np
import cv2
import PIL
import boto3
import uuid, io
import torch
import base64
import requests
import logging


class EndpointHandler:
    def __init__(self, path=""):
        
        self.tiling_size = int(os.environ["TILING_SIZE"])

        # 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,
                           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,
        )
        
        # 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 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 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....")
            
            # 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
                    }
        
        # handle unexpected errors
        except Exception as e:
            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."
        
        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
    
    def download_image_url(self, image_url):
        "Download the image from the url and return the image."
        
        response = requests.get(image_url)
        image = Image.open(BytesIO(response.content))
        return image