Update app.py

app.py

@@ -3,21 +3,39 @@ import torch
 from transformers import ViTImageProcessor, ViTForImageClassification
 from PIL import Image
 import numpy as np
+import logging
+
+# Configure logging
+logging.basicConfig(level=logging.DEBUG)
 
 # Load the pre-trained model and image processor
 model_name = "google/vit-base-patch16-224"
+logging.info("Loading image processor and model...")
 image_processor = ViTImageProcessor.from_pretrained(model_name)
 model = ViTForImageClassification.from_pretrained(model_name)
 
 # Define the prediction function
 def predict(image):
-    # Convert the dictionary to a PIL image
-    image = Image.fromarray(image.astype('uint8'), 'RGB')
-    inputs = image_processor(images=image, return_tensors="pt")
-    outputs = model(**inputs)
-    logits = outputs.logits
-    predicted_class_idx = logits.argmax(-1).item()
-    return model.config.id2label[predicted_class_idx]
+    try:
+        logging.info("Received image of type: %s", type(image))
+        # Convert the dictionary to a PIL image
+        if isinstance(image, dict):
+            logging.debug("Converting dictionary to NumPy array...")
+            image = np.array(image['image_data'])
+        logging.debug("Converting NumPy array to PIL image...")
+        image = Image.fromarray(image.astype('uint8'), 'RGB')
+        logging.debug("Image converted successfully.")
+
+        logging.info("Processing image...")
+        inputs = image_processor(images=image, return_tensors="pt")
+        outputs = model(**inputs)
+        logits = outputs.logits
+        predicted_class_idx = logits.argmax(-1).item()
+        logging.info("Prediction successful.")
+        return model.config.id2label[predicted_class_idx]
+    except Exception as e:
+        logging.error("Error during prediction: %s", e)
+        return str(e)
 
 # Create the Gradio interface
 iface = gr.Interface(