Update app.py

app.py

@@ -49,7 +49,7 @@ def preprocess_image(image):
     logger.debug(f"Preprocessed image tensor shape: {image_tensor.shape}")
     return image_tensor
 
-# Define prediction function with detailed output
+# Define prediction function with detailed output and error handling
 def predict_xray(image):
     try:
         if image is None:
@@ -61,12 +61,20 @@ def predict_xray(image):
         probs = torch.nn.functional.softmax(outputs, dim=1)[0]  # Softmax over all conditions
         results = {conditions[i]: float(probs[i].cpu().numpy()) * 100 for i in range(len(conditions))}
 
+        # Handle case where predicted class might not be in our list of conditions
         most_likely_condition = max(results, key=results.get)
-        confidence = results[most_likely_condition]
-
+        
+        # Ensure the predicted condition is in the valid conditions list
+        if most_likely_condition not in conditions:
+            most_likely_condition = "Other"
+            confidence = 0.0
+        else:
+            confidence = results[most_likely_condition]
+
+        # Create a detailed summary of results
         summary = f"**Summary**: Based on the X-ray analysis, the most likely diagnosis is: <b>{most_likely_condition}</b> with a confidence of <b>{confidence:.2f}%</b>."
 
-        # Enhanced condition details
+        # Enhanced condition details for each disease/condition
         condition_details = {
             "Normal": {"description": "No abnormal signs detected.", "recommendation": "Routine check-ups recommended."},
             "Pneumonia": {"description": "Lung inflammation detected, possibly infectious.", "recommendation": "Seek medical attention for treatment."},
@@ -78,6 +86,7 @@ def predict_xray(image):
             # Add the rest of the conditions as needed...
         }
 
+        # Display results in a clear format
         detailed_results = "<ul class='result-list'>"
         for condition, prob in results.items():
             detailed_results += f"<li><b>{condition}:</b> {prob:.2f}%</li>"