Update app.py

app.py

@@ -27,6 +27,7 @@ model.eval()  # Set model to evaluation mode
 try:
     state_dict = torch.load(model_path, map_location=torch.device('cpu'))
     model.load_state_dict(state_dict)
+    print("Model loaded successfully.")
 except RuntimeError as e:
     print("Error loading state_dict:", e)
     print("Ensure that the saved model architecture matches ResNet50.")
@@ -43,32 +44,43 @@ preprocess = transforms.Compose([
 ])
 
 # Load labels
-with open("config.json") as f:
-    labels = json.load(f)
+try:
+    with open("config.json") as f:
+        labels = json.load(f)
+    print("Labels loaded successfully.")
+except Exception as e:
+    print("Error loading labels:", e)
 
 # Function to predict image class
 def predict(image):
-    # Convert the uploaded file to a PIL image
-    input_image = image.convert("RGB")
-    
-    # Preprocess the image
-    input_tensor = preprocess(input_image)
-    input_batch = input_tensor.unsqueeze(0)  # Add batch dimension
-    
-    # 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')
-    
-    # Perform inference
-    with torch.no_grad():
-        output = model(input_batch)
-    
-    # Get the predicted class with the highest score
-    _, predicted_idx = torch.max(output, 1)
-    predicted_class = labels[str(predicted_idx.item())]
+    try:
+        # Convert the uploaded file to a PIL image
+        input_image = image.convert("RGB")
+        
+        # Preprocess the image
+        input_tensor = preprocess(input_image)
+        input_batch = input_tensor.unsqueeze(0)  # Add batch dimension
+        
+        # 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')
+        else:
+            print("GPU not available, using CPU.")
+        
+        # Perform inference
+        with torch.no_grad():
+            output = model(input_batch)
+        
+        # Get the predicted class with the highest score
+        _, predicted_idx = torch.max(output, 1)
+        predicted_class = labels[str(predicted_idx.item())]
 
-    return 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."
 
 # Set up the Gradio interface
 iface = gr.Interface(