Simplify tensor handling with more robust type checking

app.py

@@ -46,66 +46,55 @@ def detect_warning_lamp(image, history: list[tuple[str, str]], system_message):
         history.append((None, "Please upload an image first."))
         return history
     
-    # 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)
+        
+        # Get model prediction
+        pred_class, pred_idx, probs = learn_inf.predict(img)
+        
+        # Convert tensors to Python types safely
+        pred_class_str = str(pred_class)  # Convert class name to string
+        
+        # Format the prediction results
+        response = f"Detected Warning Lamp: {pred_class_str}"
+        
+        # Try to add confidence if possible
+        try:
+            # Get the index as an integer
+            if isinstance(pred_idx, torch.Tensor):
+                idx = pred_idx.item()
+            else:
+                idx = int(pred_idx)
+                
+            # Get the confidence value
+            if isinstance(probs, torch.Tensor) and idx < len(probs):
+                confidence = probs[idx].item()
+                response += f"\nConfidence: {confidence:.2%}"
+        except Exception as conf_error:
+            print(f"Could not calculate confidence: {conf_error}")
+        
+        # Add probabilities for all classes if possible
         try:
-            # Convert PIL image to FastAI compatible format
-            img = PILImage(image)
-            
-            # 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)
+            for i, cls in enumerate(learn_inf.dls.vocab):
+                if i < len(probs):
+                    if isinstance(probs, torch.Tensor):
+                        prob_value = probs[i].item()
+                    else:
+                        prob_value = float(probs[i])
                     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 prob_error:
+            print(f"Could not list all probabilities: {prob_error}")
             
-        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
+        # 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
 
 # Create a custom interface with image upload
 with gr.Blocks(title="Warning Lamp Detector", theme=gr.themes.Soft()) as demo: