Update app.py

app.py

@@ -4,27 +4,44 @@ from PIL import Image
 import torch
 import numpy as np
 import spaces
+import logging
+
+# Set up verbose logging
+logging.basicConfig(level=logging.DEBUG)
+logger = logging.getLogger(__name__)
+
 # Load model and processor from the Hugging Face Hub
 MODEL_REPO = "Rausda6/autotrain-yh172-uui7d"  # Replace with your actual model repo name
+logger.debug(f"Loading model from: {MODEL_REPO}")
 model = AutoModelForImageClassification.from_pretrained(MODEL_REPO)
 processor = AutoImageProcessor.from_pretrained(MODEL_REPO)
 
 labels = model.config.id2label
+
 @spaces.GPU
 def classify_image(img: Image.Image):
-    inputs = processor(images=img, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model(**inputs)
-        logits = outputs.logits
-        probs = torch.nn.functional.softmax(logits, dim=-1)[0]
-
-    # Build result dictionary with confidence values
-    probs_dict = {labels[i]: float(probs[i]) for i in range(len(probs))}
-    # Sort and format nicely
-    sorted_probs = sorted(probs_dict.items(), key=lambda x: x[1], reverse=True)
-    top_label, top_score = sorted_probs[0]
-
-    return {"Prediction": top_label, "Confidence": f"{top_score:.2%}"}, dict(sorted_probs)
+    logger.debug("Received image for classification.")
+    try:
+        inputs = processor(images=img, return_tensors="pt")
+        logger.debug(f"Processed inputs: {inputs}")
+        with torch.no_grad():
+            outputs = model(**inputs)
+            logger.debug(f"Model outputs: {outputs}")
+            logits = outputs.logits
+            probs = torch.nn.functional.softmax(logits, dim=-1)[0]
+        logger.debug(f"Probabilities: {probs}")
+
+        # Build result dictionary with confidence values
+        probs_dict = {labels[i]: float(probs[i]) for i in range(len(probs))}
+        # Sort and format nicely
+        sorted_probs = sorted(probs_dict.items(), key=lambda x: x[1], reverse=True)
+        top_label, top_score = sorted_probs[0]
+
+        logger.debug(f"Top prediction: {top_label} with confidence {top_score:.2%}")
+        return {"Prediction": top_label, "Confidence": f"{top_score:.2%}"}, dict(sorted_probs)
+    except Exception as e:
+        logger.exception("Error during classification")
+        raise e
 
 # Gradio interface
 demo = gr.Interface(