Fixed the base64 to PIL image function

gradio_app.py

@@ -141,8 +141,28 @@ def base64_to_pil_image(base64_string: str) -> Image.Image:
         if base64_string.startswith('data:image'):
             base64_string = base64_string.split(',')[1]
         
-        image_data = base64.b64decode(base64_string)
+        # Ensure we have valid base64 data
+        # Add padding if necessary
+        missing_padding = len(base64_string) % 4
+        if missing_padding:
+            base64_string += '=' * (4 - missing_padding)
+            
+        # Decode base64 data
+        try:
+            image_data = base64.b64decode(base64_string)
+        except Exception as e:
+            raise ValueError(f"Failed to decode base64 string: {str(e)}")
+            
+        # Ensure we have valid image data
+        if not image_data or len(image_data) == 0:
+            raise ValueError("Empty image data after base64 decoding")
+            
+        # Open as PIL Image
         image = Image.open(io.BytesIO(image_data))
+        
+        # Ensure consistent format - convert to RGBA
+        image = image.convert("RGBA")
+        
         return image
     except Exception as e:
         raise HTTPException(status_code=400, detail=f"Invalid image data: {str(e)}")

hy3dshape/hy3dshape/preprocessors.py

@@ -86,18 +86,54 @@ class ImageProcessorV2:
         return result, mask
 
     def load_image(self, image, border_ratio=0.15, to_tensor=True):
+        # Check for None or empty input
+        if image is None:
+            raise ValueError('Input image is None')
+            
+        # Handle different input types
         if isinstance(image, str):
+            # Handle file path
             image = cv2.imread(image, cv2.IMREAD_UNCHANGED)
+            if image is None:
+                raise ValueError(f'Failed to load image from path')
             image, mask = self.recenter(image, border_ratio=border_ratio)
             image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
         elif isinstance(image, Image.Image):
+            # Handle PIL Image
             image = image.convert("RGBA")
             image = np.asarray(image)
             image, mask = self.recenter(image, border_ratio=border_ratio)
-
-        image = cv2.resize(image, (self.size, self.size), interpolation=cv2.INTER_CUBIC)
-        mask = cv2.resize(mask, (self.size, self.size), interpolation=cv2.INTER_NEAREST)
-        mask = mask[..., np.newaxis]
+        elif isinstance(image, np.ndarray):
+            # Handle numpy array
+            if image.size == 0 or image.ndim < 2:
+                raise ValueError(f'Invalid numpy array shape: {image.shape}')
+            
+            # Ensure image has proper channels (3 for RGB or 4 for RGBA)
+            if image.ndim == 2:  # Grayscale
+                image = np.stack([image, image, image], axis=2)
+            elif image.shape[-1] not in [3, 4]:
+                raise ValueError(f'Invalid image channels: {image.shape[-1]}')
+                
+            if image.shape[-1] == 3:  # RGB, add alpha channel
+                alpha = np.ones((*image.shape[:2], 1), dtype=image.dtype) * 255
+                image = np.concatenate([image, alpha], axis=-1)
+                
+            image, mask = self.recenter(image, border_ratio=border_ratio)
+        else:
+            # Handle other types or unknown formats
+            raise TypeError(f'Unsupported image type: {type(image)}')
+        
+        # Validate image before resize
+        if image is None or image.size == 0 or image.ndim < 2:
+            raise ValueError(f'Invalid image after preprocessing, shape: {getattr(image, "shape", None)}')
+            
+        # Resize image and mask
+        try:
+            image = cv2.resize(image, (self.size, self.size), interpolation=cv2.INTER_CUBIC)
+            mask = cv2.resize(mask, (self.size, self.size), interpolation=cv2.INTER_NEAREST)
+            mask = mask[..., np.newaxis]
+        except Exception as e:
+            raise ValueError(f'Failed to resize image: {e}, image shape: {getattr(image, "shape", None)}')
 
         if to_tensor:
             image = array_to_tensor(image)