Update streamlit_app.py

streamlit_app.py

@@ -263,27 +263,32 @@ transform = transforms.Compose([
 
 def predict_butterfly(image, threshold=0.5):
     """Predict butterfly species from image"""
-    if image is None:
-        return None, None
-    
-    if isinstance(image, np.ndarray):
-        image = Image.fromarray(image)
-    if image.mode != 'RGB':
-        image = image.convert('RGB')
-    
-    input_tensor = transform(image).unsqueeze(0)
+    try:
+        if image is None:
+            return None, None
 
-    with torch.no_grad():
-        output = model(input_tensor)
-        probabilities = torch.nn.functional.softmax(output[0], dim=0)
+        if isinstance(image, np.ndarray):
+            image = Image.fromarray(image)
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+
+        input_tensor = transform(image).unsqueeze(0)
+
+        with torch.no_grad():
+            output = model(input_tensor)
+            probabilities = torch.nn.functional.softmax(output[0], dim=0)
+
+            confidence, pred = torch.max(probabilities, 0)
+            if confidence.item() < threshold:
+                return None, confidence.item()
+
+            predicted_class = class_names[pred.item()]
+            return predicted_class, confidence.item()
+    
+    except Exception as e:
+        st.error(f"Prediction error: {str(e)}")
+        return None, None
 
-        confidence, pred = torch.max(probabilities, 0)
-        if confidence.item() < threshold:
-            return None, confidence.item()  # Too uncertain
-        
-        predicted_class = class_names[pred.item()]
-        
-    return predicted_class, confidence.item()
 
         
 except Exception as e: