Added validation and handling for image shapes and types

hy3dshape/hy3dshape/preprocessors.py

@@ -44,13 +44,34 @@ class ImageProcessorV2:
         Returns:
             ndarray: output image, float/uint8 [H, W, 3/4]
         """
-
-        if image.shape[-1] == 4:
+        
+        # Validate image shape before accessing dimensions
+        if not isinstance(image, np.ndarray):
+            raise ValueError(f"Expected numpy array for image, got {type(image)}")
+            
+        # Handle different image shapes
+        if len(image.shape) == 2:
+            # Convert grayscale to RGB
+            image = np.stack([image, image, image], axis=-1)
+            mask = np.ones_like(image[..., 0:1]) * 255
+            image = np.concatenate([image, mask], axis=-1)
+            mask = mask[..., 0]
+        elif len(image.shape) >= 3 and image.shape[-1] == 4:
+            # RGBA image
             mask = image[..., 3]
-        else:
+        elif len(image.shape) >= 3 and image.shape[-1] == 3:
+            # RGB image
             mask = np.ones_like(image[..., 0:1]) * 255
             image = np.concatenate([image, mask], axis=-1)
             mask = mask[..., 0]
+        elif len(image.shape) >= 3 and image.shape[-1] == 1:
+            # Single channel image
+            mask = np.ones_like(image) * 255
+            image = np.concatenate([image, image, image, mask], axis=-1)
+            mask = mask[..., 0]
+        else:
+            # Handle unexpected shape
+            raise ValueError(f"Unexpected image shape: {image.shape}. Expected 2D array or array with 1, 3, or 4 channels.")
 
         H, W, C = image.shape
 
@@ -105,7 +126,11 @@ class ImageProcessorV2:
             # Handle numpy array input
             if image.size == 0:
                 raise ValueError("Input image array is empty")
-            image, mask = self.recenter(image, border_ratio=border_ratio)
+                
+            try:
+                image, mask = self.recenter(image, border_ratio=border_ratio)
+            except Exception as e:
+                raise ValueError(f"Failed to process numpy array with shape {image.shape}: {str(e)}")
         else:
             # Handle any other type by trying to convert to numpy array
             try:
@@ -119,6 +144,10 @@ class ImageProcessorV2:
                 
                 if image.size == 0:
                     raise ValueError("Converted image array is empty")
+                    
+                # Print debug info before recenter
+                print(f"Debug: image shape before recenter: {image.shape}, type: {type(image)}")
+                
                 image, mask = self.recenter(image, border_ratio=border_ratio)
             except Exception as e:
                 raise ValueError(f"Could not process image input of type {type(image)}: {e}")