update app.py

app.py

@@ -1,42 +1,57 @@
 import streamlit as st
 import torch
 from torchvision import transforms
-from PIL import Image
+from PIL import Image, ImageDraw, ImageFont
 import numpy as np
 import time
 
-# Simplified model definition
+# Simplified YOLO-style model definition (Pillow-only version)
 class PlantDiseaseModel(torch.nn.Module):
     def __init__(self, num_classes=2):
         super().__init__()
+        # Example backbone (replace with your actual model architecture)
         self.model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
         self.model.classes = num_classes  # Set your number of classes
 
     def forward(self, x):
         return self.model(x)
 
-# Load model (simplified for Hugging Face)
+# Load model 
 @st.cache_resource
 def load_model():
-    model = PlantDiseaseModel(num_classes=2)  # Update with your actual class count
+    model = PlantDiseaseModel(num_classes=2)  # Update class count
     try:
         model.load_state_dict(torch.load('best_model.pth', map_location='cpu'))
     except:
-        st.warning("Couldn't load custom weights, using pretrained")
+        st.warning("Using pretrained weights (custom weights not found)")
     return model
 
-# Image preprocessing
-def preprocess_image(image):
-    transform = transforms.Compose([
-        transforms.Resize(640),
-        transforms.ToTensor(),
-    ])
-    return transform(image).unsqueeze(0)
+# Draw bounding boxes with Pillow
+def draw_boxes_pillow(image, predictions):
+    """Draw boxes/labels on image using Pillow only"""
+    draw = ImageDraw.Draw(image)
+    try:
+        font = ImageFont.load_default()
+        for _, row in predictions.iterrows():
+            xmin, ymin, xmax, ymax = row['xmin'], row['ymin'], row['xmax'], row['ymax']
+            label = f"{row['name']} {row['confidence']:.2f}"
+            
+            # Draw rectangle
+            draw.rectangle([xmin, ymin, xmax, ymax], outline="red", width=3)
+            
+            # Draw label background
+            text_width, text_height = font.getsize(label)
+            draw.rectangle([xmin, ymin-text_height, xmin+text_width, ymin], fill="red")
+            
+            # Draw text
+            draw.text((xmin, ymin-text_height), label, fill="white", font=font)
+    except Exception as e:
+        st.error(f"Error drawing boxes: {str(e)}")
+    return image
 
-# Main app
 def main():
     st.set_page_config(page_title="Plant Disease Detector", layout="wide")
-    st.title("🌱 Plant Disease Detection")
+    st.title("🌱 Plant Disease Detection (Tomato or Corn Maiza)")
     
     model = load_model()
     
@@ -44,32 +59,40 @@ def main():
     uploaded_file = st.file_uploader("Upload a plant image...", type=["jpg", "jpeg", "png"])
     
     if uploaded_file is not None:
-        # Display original image
+        # Load with Pillow
         image = Image.open(uploaded_file).convert("RGB")
-        st.image(image, caption="Uploaded Image", use_column_width=True)
+        col1, col2 = st.columns(2)
+        
+        with col1:
+            st.image(image, caption="Original Image", use_column_width=True)
         
         # Process and predict
         with st.spinner("Analyzing..."):
-            # Preprocess
-            input_tensor = preprocess_image(image)
+            # Convert to tensor (Pillow-compatible preprocessing)
+            transform = transforms.Compose([
+                transforms.Resize(640),
+                transforms.ToTensor(),
+            ])
+            input_tensor = transform(image).unsqueeze(0)
             
             # Predict
             with torch.no_grad():
                 results = model(input_tensor)
             
-            # Convert results to image
-            results_img = np.array(results.render()[0])  # Get first image with boxes
-            
-            # Display results
-            st.image(results_img, caption="Detection Results", use_column_width=True)
+            # Convert results to Pandas
+            try:
+                results_df = results.pandas().xyxy[0]
+                
+                # Draw boxes using Pillow
+                output_image = image.copy()
+                output_image = draw_boxes_pillow(output_image, results_df)
+                
+                with col2:
+                    st.image(output_image, caption="Detection Results", use_column_width=True)
+                    st.dataframe(results_df[['name', 'confidence', 'xmin', 'ymin', 'xmax', 'ymax']])
             
-            # Show predictions
-            st.subheader("📊 Detection Results")
-            if hasattr(results, 'pandas'):
-                df = results.pandas().xyxy[0]  # Convert to pandas DataFrame
-                st.dataframe(df[['name', 'confidence', 'xmin', 'ymin', 'xmax', 'ymax']])
-            else:
-                st.write("No detections found")
+            except Exception as e:
+                st.error(f"Prediction error: {str(e)}")
 
 if __name__ == "__main__":
     main()