Update app.py

app.py

@@ -1,160 +1,51 @@
-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
+import gradio as gr
+from torchvision import transforms
+from PIL import Image
 
-# 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
+# Define class labels
+class_names = ['Bacteria', 'Fungi', 'Healthy', 'Nematode', 'Pest', 'Phytopthora', 'Virus']
 
-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
+# Load model
 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 = timm.create_model('mobilenetv3_large_100', pretrained=False)
 model.classifier = nn.Sequential(
-    nn.Linear(in_features, 512),
+    nn.Linear(model.classifier.in_features, 512),
     nn.ReLU(),
     nn.Dropout(0.3),
-    nn.Linear(512, len(full_ds.classes))
+    nn.Linear(512, len(class_names))
 )
+model.load_state_dict(torch.load('best_model.pth', map_location=device))
 model.to(device)
+model.eval()
 
-# 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
+# Transform for input image
+transform = transforms.Compose([
+    transforms.Resize(256),
+    transforms.CenterCrop(224),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406],
+                         [0.229, 0.224, 0.225])
+])
 
-# Validation function
-def validate_epoch(model, val_loader, criterion):
-    model.eval()
-    running_loss, correct_preds, total_preds = 0.0, 0, 0
+# Inference function
+def predict(image):
+    image = transform(image).unsqueeze(0).to(device)
     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)
+        outputs = model(image)
+        _, predicted = torch.max(outputs, 1)
+        confidence = torch.softmax(outputs, dim=1)[0][predicted.item()].item()
+    return {class_names[predicted.item()]: float(confidence)}
+
+# Gradio interface
+interface = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Label(num_top_classes=3),
+    title="Potato Leaf Disease Classification",
+    description="Upload an image of a potato leaf to detect the disease type."
+)
 
-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}")
+interface.launch()