Fix tensor index conversion in predict function

app.py

@@ -48,19 +48,14 @@ def detect_warning_lamp(image, history: list[tuple[str, str]], system_message):
         # Get model prediction
         pred_class, pred_idx, probs = learn_inf.predict(img)
         
-        # Convert tensor outputs to Python types
-        pred_class = str(pred_class)  # Convert class name to string
-        pred_idx = int(pred_idx)  # Convert index to integer
-        probs = [float(p) for p in probs]  # Convert probabilities to float list
-        
         # Format the prediction results
-        confidence = probs[pred_idx]  # Get confidence for predicted class
+        confidence = float(probs[int(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 cls, prob in zip(learn_inf.dls.vocab, probs):
-            response += f"\n- {cls}: {prob:.2%}"
+        for i, (cls, prob) in enumerate(zip(learn_inf.dls.vocab, probs)):
+            response += f"\n- {cls}: {float(prob):.2%}"
             
         # Update chat history
         history.append((None, response))