shdvfsdj

splitdata.py

@@ -0,0 +1,37 @@
+import os
+import shutil
+import random
+
+# ✅ Define paths
+data_dir = "PlantVillage"
+train_dir = "dataset/train"
+test_dir = "dataset/test"
+split_ratio = 0.8  # 80% train, 20% test
+
+# ✅ Ensure train/test folders exist
+os.makedirs(train_dir, exist_ok=True)
+os.makedirs(test_dir, exist_ok=True)
+
+# ✅ Split dataset
+for category in os.listdir(data_dir):  # Loop through plant disease categories
+    category_path = os.path.join(data_dir, category)
+    
+    if os.path.isdir(category_path):  # Ensure it's a folder
+        images = os.listdir(category_path)
+        random.shuffle(images)  # Shuffle to ensure randomness
+        
+        split_index = int(len(images) * split_ratio)
+        train_images = images[:split_index]
+        test_images = images[split_index:]
+
+        # ✅ Create category folders
+        os.makedirs(os.path.join(train_dir, category), exist_ok=True)
+        os.makedirs(os.path.join(test_dir, category), exist_ok=True)
+
+        # ✅ Move images
+        for img in train_images:
+            shutil.move(os.path.join(category_path, img), os.path.join(train_dir, category, img))
+        for img in test_images:
+            shutil.move(os.path.join(category_path, img), os.path.join(test_dir, category, img))
+
+print("✅ Dataset successfully split into train/test!")

train.py

@@ -0,0 +1,75 @@
+import os
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torchvision.transforms as transforms
+import torchvision.datasets as datasets
+from torch.utils.data import DataLoader
+from torchvision import models
+from tqdm import tqdm  # ✅ Progress bar
+
+# ✅ Define dataset paths
+train_dir = "dataset/train"
+test_dir = "dataset/test"
+
+# ✅ Optimized Transformations (Smaller image size)
+transform = transforms.Compose([
+    transforms.Resize((128, 128)),  # Reduce size for faster training
+    transforms.RandomHorizontalFlip(),
+    transforms.ToTensor(),
+    transforms.Normalize([0.5], [0.5])
+])
+
+# ✅ Load datasets
+train_dataset = datasets.ImageFolder(root=train_dir, transform=transform)
+test_dataset = datasets.ImageFolder(root=test_dir, transform=transform)
+
+# ✅ Get class names
+class_names = train_dataset.classes
+print(f"Class Names: {class_names}")
+
+# ✅ Save class names for later use in `app.py`
+with open("class_names.txt", "w") as f:
+    for name in class_names:
+        f.write(name + "\n")
+
+# ✅ Optimized DataLoaders (Smaller batch size)
+batch_size = 16  # Reduce batch size for speed
+train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
+test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
+
+# ✅ Use a Faster Model (MobileNetV2)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = models.mobilenet_v2(pretrained=True)
+model.classifier[1] = nn.Linear(model.classifier[1].in_features, len(class_names))
+model = model.to(device)
+
+# ✅ Define Loss Function & Optimizer
+criterion = nn.CrossEntropyLoss()
+optimizer = optim.Adam(model.parameters(), lr=0.001)
+
+# ✅ Train the Model with Progress Bar
+num_epochs = 3  # Reduce epochs for faster training
+for epoch in range(num_epochs):
+    model.train()
+    running_loss = 0.0
+    progress_bar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{num_epochs}", leave=False)
+
+    for images, labels in progress_bar:
+        images, labels = images.to(device), labels.to(device)
+
+        optimizer.zero_grad()
+        outputs = model(images)
+        loss = criterion(outputs, labels)
+        loss.backward()
+        optimizer.step()
+
+        running_loss += loss.item()
+        progress_bar.set_postfix(loss=f"{running_loss/len(train_loader):.4f}")
+
+    print(f"✅ Epoch [{epoch+1}/{num_epochs}] - Loss: {running_loss/len(train_loader):.4f}")
+
+# ✅ Save the Trained Model
+torch.save(model.state_dict(), "plant_disease_model.pth")
+print("✅ Model training complete and saved as plant_disease_model.pth")
+