Update app.py

app.py

@@ -54,33 +54,41 @@ except Exception as e:
 # Function to predict image class
 def predict(image):
     try:
+        print("Starting prediction...")
+        
         # Convert the uploaded file to a PIL image
         input_image = image.convert("RGB")
+        print(f"Image converted to RGB: {input_image.size}")
         
         # Preprocess the image
         input_tensor = preprocess(input_image)
         input_batch = input_tensor.unsqueeze(0)  # Add batch dimension
+        print(f"Input tensor shape after unsqueeze: {input_batch.shape}")
         
         # Check if a GPU is available and move the input and model to GPU
         if torch.cuda.is_available():
             input_batch = input_batch.to('cuda')
             model.to('cuda')
+            print("Using GPU for inference.")
         else:
             print("GPU not available, using CPU.")
         
         # Perform inference
         with torch.no_grad():
             output = model(input_batch)
+        print(f"Inference output shape: {output.shape}")
         
         # Get the predicted class with the highest score
         _, predicted_idx = torch.max(output, 1)
         predicted_class = labels[str(predicted_idx.item())]
+        print(f"Predicted class index: {predicted_idx.item()}")
+        print(f"Predicted class: {predicted_class}")
 
         return f"Predicted class: {predicted_class}"
     
     except Exception as e:
         print(f"Error during prediction: {e}")
-        return "An error occurred during prediction. Please try again."
+        return f"An error occurred during prediction: {e}"
 
 # Set up the Gradio interface
 iface = gr.Interface(