Create app.py

app.py

@@ -0,0 +1,160 @@
+import os
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader
+from torchvision import datasets, transforms
+from tqdm import tqdm
+import matplotlib.pyplot as plt
+import timm
+
+# Data augmentation and normalization
+transform_train = transforms.Compose([
+    transforms.RandomResizedCrop(224, scale=(0.8, 1.0)),
+    transforms.RandomHorizontalFlip(),
+    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.2),
+    transforms.RandomRotation(15),
+    transforms.RandomAffine(degrees=15, translate=(0.1, 0.1)),
+    transforms.GaussianBlur(kernel_size=3),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                         std=[0.229, 0.224, 0.225]),
+])
+
+transform_val = transforms.Compose([
+    transforms.Resize(224),
+    transforms.CenterCrop(224),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                         std=[0.229, 0.224, 0.225]),
+])
+
+# Dataset loading
+train_dir = 'D:\\Dataset\\Potato Leaf Disease Dataset in Uncontrolled Environment'
+full_ds = datasets.ImageFolder(train_dir, transform=transform_train)
+train_size = int(0.8 * len(full_ds))
+val_size = len(full_ds) - train_size
+train_ds, val_ds = torch.utils.data.random_split(full_ds, [train_size, val_size])
+val_ds.dataset.transform = transform_val  # Apply validation transforms
+
+train_loader = DataLoader(train_ds, batch_size=32, shuffle=True, num_workers=4)
+val_loader = DataLoader(val_ds, batch_size=32, shuffle=False, num_workers=4)
+
+# Device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Model definition with custom classification head (optional improvement)
+model = timm.create_model('mobilenetv3_large_100', pretrained=True)
+in_features = model.classifier.in_features
+model.classifier = nn.Sequential(
+    nn.Linear(in_features, 512),
+    nn.ReLU(),
+    nn.Dropout(0.3),
+    nn.Linear(512, len(full_ds.classes))
+)
+model.to(device)
+
+# Loss and optimizer
+criterion = nn.CrossEntropyLoss()
+optimizer = optim.Adam(model.parameters(), lr=1e-4)
+scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', factor=0.5, patience=3)
+
+# MixUp augmentation
+def mixup_data(x, y, alpha=1.0):
+    if alpha > 0:
+        lam = np.random.beta(alpha, alpha)
+    else:
+        lam = 1
+    batch_size = x.size(0)
+    index = torch.randperm(batch_size).to(x.device)
+    mixed_x = lam * x + (1 - lam) * x[index, :]
+    y_a, y_b = y, y[index]
+    return mixed_x, y_a, y_b, lam
+
+def mixup_criterion(criterion, pred, y_a, y_b, lam):
+    return lam * criterion(pred, y_a) + (1 - lam) * criterion(pred, y_b)
+
+# Training function with MixUp
+def train_epoch(model, train_loader, criterion, optimizer):
+    model.train()
+    running_loss, correct_preds, total_preds = 0.0, 0, 0
+    for inputs, labels in tqdm(train_loader, desc="Training Epoch", leave=False):
+        inputs, labels = inputs.to(device), labels.to(device)
+        inputs, targets_a, targets_b, lam = mixup_data(inputs, labels, alpha=1.0)
+
+        optimizer.zero_grad()
+        outputs = model(inputs)
+        loss = mixup_criterion(criterion, outputs, targets_a, targets_b, lam)
+        loss.backward()
+        optimizer.step()
+
+        _, preds = torch.max(outputs, 1)
+        correct_preds += (lam * preds.eq(targets_a).sum().item()
+                          + (1 - lam) * preds.eq(targets_b).sum().item())
+        total_preds += labels.size(0)
+        running_loss += loss.item()
+
+    return running_loss / len(train_loader), correct_preds / total_preds
+
+# Validation function
+def validate_epoch(model, val_loader, criterion):
+    model.eval()
+    running_loss, correct_preds, total_preds = 0.0, 0, 0
+    with torch.no_grad():
+        for inputs, labels in tqdm(val_loader, desc="Validating Epoch", leave=False):
+            inputs, labels = inputs.to(device), labels.to(device)
+            outputs = model(inputs)
+            loss = criterion(outputs, labels)
+            _, preds = torch.max(outputs, 1)
+            correct_preds += (preds == labels).sum().item()
+            total_preds += labels.size(0)
+            running_loss += loss.item()
+    return running_loss / len(val_loader), correct_preds / total_preds
+
+# Plotting
+def plot_metrics(train_loss, val_loss, train_acc, val_acc):
+    epochs = range(1, len(train_loss) + 1)
+    plt.figure(figsize=(12, 5))
+    plt.subplot(1, 2, 1)
+    plt.plot(epochs, train_loss, label='Training Loss')
+    plt.plot(epochs, val_loss, label='Validation Loss')
+    plt.xlabel('Epochs')
+    plt.ylabel('Loss')
+    plt.legend()
+    plt.subplot(1, 2, 2)
+    plt.plot(epochs, train_acc, label='Training Accuracy')
+    plt.plot(epochs, val_acc, label='Validation Accuracy')
+    plt.xlabel('Epochs')
+    plt.ylabel('Accuracy')
+    plt.legend()
+    plt.show()
+
+# Training loop
+num_epochs = 20
+train_losses, val_losses = [], []
+train_accuracies, val_accuracies = [], []
+
+for epoch in range(num_epochs):
+    print(f"\nEpoch {epoch+1}/{num_epochs}")
+    train_loss, train_acc = train_epoch(model, train_loader, criterion, optimizer)
+    val_loss, val_acc = validate_epoch(model, val_loader, criterion)
+    scheduler.step(val_acc)
+
+    print(f"Train Loss: {train_loss:.4f}, Accuracy: {train_acc:.4f}")
+    print(f"Val   Loss: {val_loss:.4f}, Accuracy: {val_acc:.4f}")
+
+    train_losses.append(train_loss)
+    val_losses.append(val_loss)
+    train_accuracies.append(train_acc)
+    val_accuracies.append(val_acc)
+
+plot_metrics(train_losses, val_losses, train_accuracies, val_accuracies)
+best_val_acc = 0.0
+save_path = 'D:\\Dataset\\Potato Leaf Disease Dataset in Uncontrolled Environment\\best_model.pth'
+os.makedirs(os.path.dirname(save_path), exist_ok=True)
+
+if val_acc > best_val_acc:
+    best_val_acc = val_acc
+    torch.save(model.state_dict(), save_path)
+    print(f"✅ Best model saved with val_acc: {val_acc:.4f}")