update handler file: add checks for image size, mode. Add exceptions

handler.py

@@ -8,13 +8,11 @@ from pathlib import Path
 from basicsr.archs.rrdbnet_arch import RRDBNet
 import numpy as np
 import cv2
+import PIL
 
 
 import torch
 import base64
-# torch.cuda.empty_cache()
-# torch.cuda.set_per_process_memory_fraction(0.99)
-# os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:64,garbage_collection_threshold:0.7"
 
 
 class EndpointHandler:
@@ -22,9 +20,16 @@ class EndpointHandler:
 
         self.model = RealESRGANer(
             scale=4,
-            model_path=f"/repository/weights/Real-ESRGAN-x4plus.pth",
+            # model_path=f"/repository/weights/Real-ESRGAN-x4plus.pth",
+            model_path="/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, scale=4),
+            model= RRDBNet(num_in_ch=3, 
+                           num_out_ch=3, 
+                           num_feat=64, 
+                           num_block=23, 
+                           num_grow_ch=32, 
+                           scale=4
+                           ),
             tile=0,
             tile_pad=10,
             # pre_pad=args.pre_pad,
@@ -33,28 +38,80 @@ class EndpointHandler:
         )
 
     def __call__(self, data: Any) -> Dict[str, List[float]]:
-        inputs = data.pop("inputs", data)
-        outscale = 3
         
-        # decode base64 image to PIL
-        image = Image.open(BytesIO(base64.b64decode(inputs['image'])))
-        # Convert PIL image to NumPy array
-        opencv_image = np.array(image)
-        # Convert RGB to BGR (PIL uses RGB, OpenCV expects BGR)
-        opencv_image = cv2.cvtColor(opencv_image, cv2.COLOR_RGB2BGR)
-        output, _ = self.model.enhance(opencv_image, outscale=outscale)
-        
-        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)
+        try:
+            
+            # get inputs
+            inputs = data.pop("inputs", data)
+            
+            # get outscale
+            outscale = float(inputs.pop("outscale", 3))
+            
+            # decode base64 image to PIL
+            image = Image.open(BytesIO(base64.b64decode(inputs['image'])))
+            in_size, in_mode = image.size, image.mode
+            
+            # check image size and mode and return dict
+            assert in_mode in ["RGB", "RGBA", "L"], f"Unsupported image mode: {in_mode}"
+            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 {1400*1400}"
+            assert outscale > 1 and outscale <=10, f"Outscale must be between 1 and 10: {outscale}"
+            
+            # 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)
+            if in_mode == "RGB":
+                opencv_image = cv2.cvtColor(opencv_image, cv2.COLOR_RGB2BGR)
+            elif in_mode == "RGBA":
+                opencv_image = cv2.cvtColor(opencv_image, cv2.COLOR_RGBA2BGRA)
+            elif in_mode == "L":
+                opencv_image = cv2.cvtColor(opencv_image, cv2.COLOR_GRAY2RGB)
+            else:
+                raise ValueError(f"Unsupported image mode: {in_mode}")
+            
+            # enhance image
+            output, _ = self.model.enhance(opencv_image, outscale=outscale)
+            
+            # debug
+            print(f"output.shape: {output.shape}")
+
+            # convert 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
+            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()
+            
+            return {"out_image": img_str, 
+                    "error": None
+                    }
         
-        img_byte_arr = BytesIO()
-        output = Image.fromarray(output)
-        output.save(img_byte_arr, format='PNG')
-        img_str = base64.b64encode(img_byte_arr.getvalue())
-        return {"out_image": img_str.decode()}
+        # 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"}
+        except Exception as e:
+            print(f"Exception: {e}")
+            return {"out_image": None, "error": "An unexpected error occurred"}