Update app.py

app.py

@@ -36,10 +36,6 @@ tokenizer = AutoTokenizer.from_pretrained(
 )
 
 
-from PIL import Image, ImageEnhance
-import numpy as np
-import torch
-
 def preprocessing(image):
     """Apply three enhancement filters without resizing or cropping."""
     
@@ -119,11 +115,25 @@ def vision_ai_api(image, doc_type):
         
     return results
 
-
+def ensure_numpy(image):
+    """Ensure image is a valid NumPy array."""
+    if isinstance(image, torch.Tensor):
+        # Convert PyTorch tensor to NumPy array
+        image = image.permute(1, 2, 0).cpu().numpy()
+    elif isinstance(image, Image.Image):
+        # Convert PIL image to NumPy array
+        image = np.array(image)
+    
+    if len(image.shape) == 2:  
+        # Convert grayscale to 3-channel image
+        image = np.stack([image] * 3, axis=-1)
+    
+    return image
+    
 def predict(image):
     """Pipeline: Preprocess -> Detect -> Crop -> Vision AI API."""
     processed_image = preprocessing(image)
-    rotated_image = imageRotation(processed_image)  # Placeholder for rotation
+    rotated_image = ensure_numpy(rotated_image)
     detected_image, labels, bounding_boxes = detect_document(rotated_image)
 
     cropped_images = crop_image(rotated_image, bounding_boxes)