Add retry logic and robust tensor handling for intermittent failures

app.py

@@ -4,6 +4,7 @@ from fastai.learner import load_learner
 from pathlib import Path
 import pandas as pd
 import os
+import time
 
 """
 Warning Lamp Detector using FastAI
@@ -44,31 +45,67 @@ def detect_warning_lamp(image, history: list[tuple[str, str]], system_message):
     if image is None:
         history.append((None, "Please upload an image first."))
         return history
-        
-    try:
-        # Convert PIL image to FastAI compatible format
-        img = PILImage(image)
-        
-        # Get model prediction
-        pred_class, pred_idx, probs = learn_inf.predict(img)
-        
-        # Format the prediction results
-        confidence = float(probs[pred_idx])  # Convert to float for better formatting
-        response = f"Detected Warning Lamp: {pred_class}\nConfidence: {confidence:.2%}"
-        
-        # Add probabilities for all classes
-        response += "\n\nProbabilities for all classes:"
-        for i, (cls, prob) in enumerate(zip(learn_inf.dls.vocab, probs)):
-            response += f"\n- {cls}: {float(prob):.2%}"
+    
+    # Maximum number of retries
+    max_retries = 3
+    retry_count = 0
+    
+    while retry_count < max_retries:
+        try:
+            # Convert PIL image to FastAI compatible format
+            img = PILImage(image)
             
-        # Update chat history
-        history.append((None, response))
-        return history
-    except Exception as e:
-        error_msg = f"Error processing image: {str(e)}"
-        print(f"Exception in detect_warning_lamp: {e}")
-        history.append((None, error_msg))
-        return history
+            # Get model prediction
+            pred_class, pred_idx, probs = learn_inf.predict(img)
+            
+            # Try different approaches to handle tensor conversion
+            try:
+                # First approach - direct conversion
+                confidence = float(probs[pred_idx])
+            except Exception as e1:
+                print(f"First conversion approach failed: {e1}")
+                try:
+                    # Second approach - convert index first
+                    idx = int(pred_idx)
+                    confidence = float(probs[idx])
+                except Exception as e2:
+                    print(f"Second conversion approach failed: {e2}")
+                    # Third approach - use item() method if available
+                    if hasattr(probs[pred_idx], 'item'):
+                        confidence = probs[pred_idx].item()
+                    else:
+                        # Last resort - use the max probability
+                        confidence = float(max(probs))
+            
+            # Format the prediction results
+            response = f"Detected Warning Lamp: {pred_class}\nConfidence: {confidence:.2%}"
+            
+            # Add probabilities for all classes
+            response += "\n\nProbabilities for all classes:"
+            for i, (cls, prob) in enumerate(zip(learn_inf.dls.vocab, probs)):
+                try:
+                    prob_value = float(prob)
+                    response += f"\n- {cls}: {prob_value:.2%}"
+                except Exception as prob_error:
+                    print(f"Error converting probability for {cls}: {prob_error}")
+                    response += f"\n- {cls}: N/A"
+            
+            # Update chat history
+            history.append((None, response))
+            return history
+            
+        except Exception as e:
+            retry_count += 1
+            print(f"Attempt {retry_count} failed with error: {e}")
+            
+            if retry_count < max_retries:
+                print(f"Retrying in 1 second...")
+                time.sleep(1)  # Wait a bit before retrying
+            else:
+                error_msg = f"Error processing image after {max_retries} attempts: {str(e)}"
+                print(f"All retries failed: {error_msg}")
+                history.append((None, error_msg))
+                return history
 
 # Create a custom interface with image upload
 with gr.Blocks(title="Warning Lamp Detector", theme=gr.themes.Soft()) as demo: