Update app.py

app.py

@@ -1,44 +1,45 @@
-import streamlit as st
-from PIL import Image
-import torch
-from ultralytics import YOLO
-import numpy as np
-
-# Load the YOLOv8 model
-@st.cache_resource
-def load_model():
-    model = YOLO('yolov8n.pt')  # Path to your YOLOv8 trained weights
-    return model
-
-model = load_model()
 
-# Function to process image and detect weeds
-def detect_weeds(image):
-    # Run inference
+from PIL import Image
+import torchvision.transforms as transforms
+
+# لود مدل YOLOv5
+model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
+
+# پیش‌پردازش تصویر
+def preprocess_image(image_path):
+    image = Image.open(image_path)
+    transform = transforms.Compose([
+        transforms.Resize((640, 640)),
+        transforms.ToTensor()
+    ])
+    return transform(image).unsqueeze(0)
+
+# آموزش مدل
+def train_model(data_dir, epochs=10):
+    # آماده‌سازی داده‌ها
+    dataset = ...  # خواندن داده‌ها از data_dir
+    dataloader = ...  # ایجاد DataLoader
+
+    # تنظیم پارامترهای آموزش
+    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
+    criterion = torch.nn.CrossEntropyLoss()
+
+    for epoch in range(epochs):
+        for images, labels in dataloader:
+            optimizer.zero_grad()
+            outputs = model(images)
+            loss = criterion(outputs, labels)
+            loss.backward()
+            optimizer.step()
+        print(f'Epoch {epoch+1}/{epochs}, Loss: {loss.item()}')
+
+# تشخیص مناطق دارای گپ
+def detect_gaps(image_path):
+    image = preprocess_image(image_path)
     results = model(image)
-    
-    # Process the results and draw bounding boxes
-    result_img = image.copy()  # Start with the original image
-    for result in results:
-        annotated_img = result.plot()  # Plot boxes and labels
-        result_img = np.array(annotated_img)
-
-    return result_img
-
-# Streamlit app UI
-st.title("Weed Detection in Sugarcane Fields")
-st.write("Upload an image of the sugarcane field, and the model will detect weeds.")
-
-# Upload image
-uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
-
-if uploaded_file is not None:
-    image = Image.open(uploaded_file)
-    st.image(image, caption="Uploaded Image", use_column_width=True)
-
-    st.write("Detecting weeds...")
-    result_image = detect_weeds(np.array(image))  # Convert to array for processing
-    
-    # Display result
-    st.image(result_image, caption="Detected Weeds", use_column_width=True)
+    return results
 
+# مثال استفاده
+image_path = '/content/Sugarcane-Cultivation-in-Tamil-Nadu-1.jpg'
+results = detect_gaps(image_path)
+print(results)