Update app.py

app.py

@@ -1,143 +1,47 @@
 import streamlit as st
-import torch
-import torch.nn as nn
-import torchvision.transforms as transforms
+from ultralytics import YOLO
 from PIL import Image
-import os
-import requests
+import tempfile
+import numpy as np
 
-# ======= Groq API Key (Embed Yours Here) =======
-GROQ_API_KEY = "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"
+st.set_page_config(page_title="YOLOv8 Inference App", layout="centered")
+st.title("🔍 YOLOv8 Image Detector")
 
-# ======= Device & Model Setup =======
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+MODEL_PATH = "best.pt"
 
-class PlantDiseaseCNN(nn.Module):
-    def __init__(self, num_classes):
-        super(PlantDiseaseCNN, self).__init__()
-        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
-        self.pool = nn.MaxPool2d(2, 2)
-        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
-        self.conv3 = nn.Conv2d(64, 64, kernel_size=3, padding=1)
-        self.fc1 = nn.Linear(64 * 32 * 32, 128)
-        self.fc2 = nn.Linear(128, num_classes)
-        self.relu = nn.ReLU()
+@st.cache_resource
+def load_model():
+    return YOLO(MODEL_PATH)
 
-    def forward(self, x):
-        x = self.pool(self.relu(self.conv1(x)))
-        x = self.pool(self.relu(self.conv2(x)))
-        x = self.pool(self.relu(self.conv3(x)))
-        x = x.view(x.size(0), -1)
-        x = self.relu(self.fc1(x))
-        x = self.fc2(x)
-        return x
+model = load_model()
 
-# ======= Class Names for Each Crop =======
-class_names = {
-    "Tomato Disease Detection": [
-        "Bacterial spot", "Early blight", "Late blight", "Leaf Mold",
-        "Septoria leaf spot", "Spider mites", "Target Spot", "Yellow Leaf Curl Virus",
-        "Mosaic virus", "Healthy"
-    ],
-    "Pepper Disease Detection": [
-        "Bacterial spot", "Healthy"
-    ],
-    "Potato Disease Detection": [
-        "Early blight", "Late blight", "Healthy"
-    ]
-}
+# Option to choose input method
+input_method = st.radio("Choose Image Input Method:", ("📤 Upload", "📷 Camera"))
 
-# ======= Load Models =======
-models_info = {
-    "Tomato Disease Detection": "tomato-detection-model.pth",
-    "Pepper Disease Detection": "pepper-detection-model.pth",
-    "Potato Disease Detection": "potato-detection-model.pth"
-}
+if input_method == "📤 Upload":
+    image_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
+    if image_file:
+        image = Image.open(image_file)
 
-models = {}
+elif input_method == "📷 Camera":
+    camera_image = st.camera_input("Take a picture")
+    if camera_image:
+        image = Image.open(camera_image)
 
-for crop, model_path in models_info.items():
-    num_classes = len(class_names[crop])
-    model = PlantDiseaseCNN(num_classes)
-    state_dict = torch.load(model_path, map_location=device)
-    model.load_state_dict(state_dict)
-    model.to(device)
-    model.eval()
-    models[crop] = model
+# Proceed if image is available
+if ('image' in locals()):
+    st.image(image, caption="🖼️ Input Image", use_column_width=True)
 
-# ======= Preprocessing =======
-transform = transforms.Compose([
-    transforms.Resize((256, 256)),
-    transforms.ToTensor(),
-    transforms.Normalize([0.5], [0.5])
-])
+    with st.spinner("🔍 Detecting objects..."):
+        results = model(image)
+        result_img = results[0].plot()
 
-def predict_image(image, model, crop_name):
-    image = image.convert("RGB")
-    image = transform(image).unsqueeze(0).to(device)
-    with torch.no_grad():
-        output = model(image)
-        _, predicted = torch.max(output, 1)
-    return class_names[crop_name][predicted.item()]
+    st.image(result_img, caption="🧠 Detection Result", use_column_width=True)
 
-# ======= Groq LLM Remedies =======
-def get_remedies_from_groq(disease_name):
-    headers = {
-        "Authorization": f"Bearer {GROQ_API_KEY}",
-        "Content-Type": "application/json",
-    }
-    payload = {
-        "model": "llama-3.1-8b-instant",
-        "messages": [
-            {
-                "role": "user",
-                "content": f"""The plant disease is **{disease_name}**.
-Give me professional agricultural advice in this format:
-1. Insecticides to use:
-2. Pesticides to use:
-3. Herbicides to use:
-4. Natural remedies:
-5. Preventive measures:
-Make it clear and line-by-line."""
-            }
-        ],
-        "temperature": 0.7
-    }
+    # Save detection result
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as tmp:
+        Image.fromarray(result_img).save(tmp.name)
+        st.download_button("📥 Download Output", data=open(tmp.name, "rb").read(), file_name="detected.jpg")
 
-    response = requests.post("https://api.groq.com/openai/v1/chat/completions", headers=headers, json=payload)
-    if response.status_code == 200:
-        return response.json()["choices"][0]["message"]["content"]
-    else:
-        return "❌ Failed to fetch remedies. Check your API key or network connection."
-
-# ======= Streamlit UI =======
-st.set_page_config(page_title="Plant Disease Detection", layout="centered")
-st.title("🌿 Plant Disease Detection & Remedies ")
-
-tab1, tab2, tab3 = st.tabs(["🍅 Tomato", "🌶️ Pepper", "🥔 Potato"])
-
-for tab, crop in zip([tab1, tab2, tab3], models_info.keys()):
-    with tab:
-        st.subheader(crop)
-        method = st.radio("Select Input Method:", ["Upload Image", "Use Camera"], key=f"{crop}_method")
-
-        image = None
-        if method == "Upload Image":
-            uploaded_file = st.file_uploader("Upload a leaf image", type=["jpg", "jpeg", "png"], key=f"{crop}_upload")
-            if uploaded_file:
-                image = Image.open(uploaded_file)
-
-        elif method == "Use Camera":
-            captured_image = st.camera_input("Capture a leaf image", key=f"{crop}_camera")
-            if captured_image:
-                image = Image.open(captured_image)
-
-        if image:
-            st.image(image, caption="Leaf Image", use_column_width=True)
-            prediction = predict_image(image, models[crop], crop)
-            st.success(f"🔍 Prediction: **{prediction}**")
-
-            with st.spinner("🤖 Fetching remedies ..."):
-                advice = get_remedies_from_groq(prediction)
-                st.markdown("### 🌱 Remedies & Prevention")
-                st.markdown(advice)
+else:
+    st.info("Please provide an image using Upload or Camera.")