upload all files

.gitattributes

@@ -1,35 +1 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text

README copy.md

@@ -0,0 +1,28 @@
+---
+title: Pest and Disease Classification 🌿
+emoji: 🌱
+colorFrom: green
+colorTo: yellow
+sdk: gradio
+sdk_version: "4.44.1"
+app_file: app.py
+pinned: false
+---
+
+# 🌿 Pest and Disease Classification Demo
+
+This demo provides a simple web interface for classifying **pests and diseases in citrus leaves**.
+
+## 🧠 Model
+The model is based on a CNN backbone (ResNet50 by default) trained on a labeled dataset of citrus plant leaves.
+
+- **Framework:** PyTorch  
+- **Interface:** Gradio  
+- **Backbone:** ResNet50  
+- **Task:** Image classification  
+
+## 🚀 How to Use
+1. Click **“Upload Image”** and select a photo of a citrus leaf.  
+2. The app will output the **predicted pest or disease category** with confidence scores.  
+
+## 📂 Repository Structure

app.py

@@ -0,0 +1,94 @@
+"""
+Simple Demo for Pest and Disease Classification
+For Hugging Face Space Deployment
+"""
+
+import torch
+from PIL import Image
+import json
+import gradio as gr
+from torchvision import transforms
+
+from model import create_model
+
+
+class PestDiseasePredictor:
+    """Simple predictor class"""
+
+    def __init__(self, checkpoint_path, label_mapping_path, backbone='resnet50', device='cuda'):
+        self.device = torch.device(device if torch.cuda.is_available() else 'cpu')
+
+        # Load label mapping
+        with open(label_mapping_path, 'r', encoding='utf-8') as f:
+            mapping = json.load(f)
+            self.id_to_label = {int(k): v for k, v in mapping['id_to_label'].items()}
+            self.num_classes = mapping['num_classes']
+
+        # Load model
+        self.model = create_model(
+            num_classes=self.num_classes,
+            backbone=backbone,
+            pretrained=False
+        )
+
+        # Load checkpoint
+        checkpoint = torch.load(checkpoint_path, map_location=self.device)
+        self.model.load_state_dict(checkpoint['model_state_dict'])
+        self.model = self.model.to(self.device)
+        self.model.eval()
+
+        # Image transforms
+        self.transform = transforms.Compose([
+            transforms.Resize((224, 224)),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                                 std=[0.229, 0.224, 0.225])
+        ])
+
+        print(f"✅ Model loaded from {checkpoint_path}")
+        print(f"💻 Device: {self.device}")
+        print(f"📚 Classes: {self.num_classes}")
+
+    def predict(self, image):
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+
+        img_tensor = self.transform(image).unsqueeze(0).to(self.device)
+        with torch.no_grad():
+            outputs = self.model(img_tensor)
+            probs = torch.nn.functional.softmax(outputs, dim=1)[0].cpu().numpy()
+
+        results = {self.id_to_label[i]: float(p) for i, p in enumerate(probs)}
+        return dict(sorted(results.items(), key=lambda x: x[1], reverse=True))
+
+
+# ========== For Hugging Face Space ==========
+checkpoint_path = "checkpoints/best_efficientnet_b3.pth"
+label_mapping_path = "label_mapping.json"
+backbone = 'efficientnet_b3'
+device = "cuda"
+
+predictor = PestDiseasePredictor(
+    checkpoint_path=checkpoint_path,
+    label_mapping_path=label_mapping_path,
+    backbone=backbone,
+    device=device
+)
+
+def predict_image(image):
+    if image is None:
+        return None
+    return predictor.predict(image)
+
+demo = gr.Interface(
+    fn=predict_image,
+    inputs=gr.Image(type="pil", label="Upload Image"),
+    outputs=gr.Label(num_top_classes=10, label="Predictions"),
+    title="🌿 Pest and Disease Classification",
+    description="Upload an image of a citrus leaf to classify its pest or disease type.",
+    theme=gr.themes.Soft(),
+    allow_flagging="never"
+)
+
+if __name__ == "__main__":
+    demo.launch()

checkpoints/best_model_993.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e18b22b823871125c07933e128e7afed92110da462fee5781462fed1066b4e33
+size 138717293

dataset.py

@@ -0,0 +1,253 @@
+"""
+PyTorch Dataset and DataLoader for Pest and Disease Classification
+"""
+
+import torch
+from torch.utils.data import Dataset, DataLoader
+from torchvision import transforms
+from PIL import Image
+import pandas as pd
+import json
+from pathlib import Path
+
+
+class PestDiseaseDataset(Dataset):
+    """Custom Dataset for loading pest and disease images"""
+
+    def __init__(self, csv_file, label_mapping_file, split='train', transform=None):
+        """
+        Args:
+            csv_file (str): Path to CSV file with image paths and labels
+            label_mapping_file (str): Path to JSON file with label mappings
+            split (str): One of 'train', 'val', or 'test'
+            transform (callable, optional): Optional transform to be applied on images
+        """
+        self.df = pd.read_csv(csv_file)
+        self.df = self.df[self.df['split'] == split].reset_index(drop=True)
+
+        # Load label mapping
+        with open(label_mapping_file, 'r', encoding='utf-8') as f:
+            mapping = json.load(f)
+            self.label_to_id = mapping['label_to_id']
+            self.id_to_label = {int(k): v for k, v in mapping['id_to_label'].items()}
+            self.num_classes = mapping['num_classes']
+
+        self.transform = transform
+        self.split = split
+
+        print(f"Loaded {split} set: {len(self.df)} images")
+
+    def __len__(self):
+        return len(self.df)
+
+    def __getitem__(self, idx):
+        """
+        Returns:
+            image: Transformed image tensor
+            label: Label ID (integer)
+        """
+        row = self.df.iloc[idx]
+
+        # Load image
+        img_path = row['image_path']
+        image = Image.open(img_path).convert('RGB')
+
+        # Get label
+        label_name = row['label']
+        label = self.label_to_id[label_name]
+
+        # Apply transforms
+        if self.transform:
+            image = self.transform(image)
+
+        return image, label
+
+    def get_label_name(self, label_id):
+        """Convert label ID back to label name"""
+        return self.id_to_label[label_id]
+
+
+def get_transforms(split='train', img_size=224):
+    """
+    Get data augmentation transforms for different splits
+
+    Args:
+        split (str): 'train', 'val', or 'test'
+        img_size (int): Target image size (default: 224 for most pretrained models)
+
+    Returns:
+        transforms.Compose: Composed transforms
+    """
+    if split == 'train':
+        # Training: Apply data augmentation
+        return transforms.Compose([
+            transforms.Resize((img_size, img_size)),
+            transforms.RandomHorizontalFlip(p=0.5),
+            transforms.RandomVerticalFlip(p=0.3),
+            transforms.RandomRotation(degrees=30),
+            transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
+            transforms.RandomAffine(degrees=0, translate=(0.1, 0.1), scale=(0.9, 1.1)),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                               std=[0.229, 0.224, 0.225])
+        ])
+    else:
+        # Validation/Test: No augmentation, only resize and normalize
+        return transforms.Compose([
+            transforms.Resize((img_size, img_size)),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                               std=[0.229, 0.224, 0.225])
+        ])
+
+
+def get_dataloaders(csv_file='dataset.csv',
+                   label_mapping_file='label_mapping.json',
+                   batch_size=32,
+                   img_size=224,
+                   num_workers=4):
+    """
+    Create train, validation, and test dataloaders
+
+    Args:
+        csv_file (str): Path to dataset CSV
+        label_mapping_file (str): Path to label mapping JSON
+        batch_size (int): Batch size for training
+        img_size (int): Image size for models
+        num_workers (int): Number of workers for data loading
+
+    Returns:
+        dict: Dictionary containing 'train', 'val', 'test' dataloaders and 'num_classes'
+    """
+
+    # Create datasets
+    train_dataset = PestDiseaseDataset(
+        csv_file=csv_file,
+        label_mapping_file=label_mapping_file,
+        split='train',
+        transform=get_transforms('train', img_size)
+    )
+
+    val_dataset = PestDiseaseDataset(
+        csv_file=csv_file,
+        label_mapping_file=label_mapping_file,
+        split='val',
+        transform=get_transforms('val', img_size)
+    )
+
+    test_dataset = PestDiseaseDataset(
+        csv_file=csv_file,
+        label_mapping_file=label_mapping_file,
+        split='test',
+        transform=get_transforms('test', img_size)
+    )
+
+    # Create dataloaders
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=batch_size,
+        shuffle=True,
+        num_workers=num_workers,
+        pin_memory=True
+    )
+
+    val_loader = DataLoader(
+        val_dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=num_workers,
+        pin_memory=True
+    )
+
+    test_loader = DataLoader(
+        test_dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=num_workers,
+        pin_memory=True
+    )
+
+    return {
+        'train': train_loader,
+        'val': val_loader,
+        'test': test_loader,
+        'num_classes': train_dataset.num_classes,
+        'datasets': {
+            'train': train_dataset,
+            'val': val_dataset,
+            'test': test_dataset
+        }
+    }
+
+
+def calculate_class_weights(csv_file='dataset.csv', label_mapping_file='label_mapping.json'):
+    """
+    Calculate class weights for handling imbalanced dataset
+
+    Returns:
+        torch.Tensor: Class weights for loss function
+    """
+    df = pd.read_csv(csv_file)
+    train_df = df[df['split'] == 'train']
+
+    with open(label_mapping_file, 'r', encoding='utf-8') as f:
+        mapping = json.load(f)
+        label_to_id = mapping['label_to_id']
+        num_classes = mapping['num_classes']
+
+    # Count samples per class
+    class_counts = {}
+    for label in train_df['label']:
+        label_id = label_to_id[label]
+        class_counts[label_id] = class_counts.get(label_id, 0) + 1
+
+    # Calculate weights (inverse frequency)
+    total_samples = len(train_df)
+    weights = []
+    for i in range(num_classes):
+        count = class_counts.get(i, 1)
+        weight = total_samples / (num_classes * count)
+        weights.append(weight)
+
+    weights = torch.FloatTensor(weights)
+
+    print("\nClass weights:")
+    for i, w in enumerate(weights):
+        print(f"  Class {i}: {w:.4f}")
+
+    return weights
+
+
+if __name__ == "__main__":
+    """Test the dataloader"""
+    print("Testing Pest and Disease Dataloader")
+    print("=" * 60)
+
+    # Get dataloaders
+    loaders = get_dataloaders(batch_size=8, img_size=224, num_workers=0)
+
+    # Calculate class weights
+    class_weights = calculate_class_weights()
+
+    print("\n" + "=" * 60)
+    print("Testing batch loading...")
+    print("=" * 60)
+
+    # Test loading a batch from train set
+    train_loader = loaders['train']
+    train_dataset = loaders['datasets']['train']
+
+    for images, labels in train_loader:
+        print(f"\nBatch shape: {images.shape}")
+        print(f"Labels shape: {labels.shape}")
+        print(f"Image dtype: {images.dtype}")
+        print(f"Labels: {labels.tolist()}")
+        print(f"Label names: {[train_dataset.get_label_name(l.item()) for l in labels]}")
+
+        # Check value ranges
+        print(f"\nImage value range: [{images.min():.3f}, {images.max():.3f}]")
+        break
+
+    print("\n" + "=" * 60)
+    print("Dataloader test completed successfully!")
+    print("=" * 60)

demo.py

@@ -0,0 +1,163 @@
+"""
+Simple Demo for Pest and Disease Classification
+Upload an image and get prediction
+"""
+
+import torch
+from PIL import Image
+import json
+import argparse
+import gradio as gr
+from torchvision import transforms
+
+from model import create_model
+
+
+class PestDiseasePredictor:
+    """Simple predictor class"""
+
+    def __init__(self, checkpoint_path, label_mapping_path, backbone='resnet50', device='cuda'):
+        self.device = torch.device(device if torch.cuda.is_available() else 'cpu')
+
+        # Load label mapping
+        with open(label_mapping_path, 'r', encoding='utf-8') as f:
+            mapping = json.load(f)
+            self.id_to_label = {int(k): v for k, v in mapping['id_to_label'].items()}
+            self.num_classes = mapping['num_classes']
+
+        # Load model
+        self.model = create_model(
+            num_classes=self.num_classes,
+            backbone=backbone,
+            pretrained=False
+        )
+
+        # Load checkpoint
+        checkpoint = torch.load(checkpoint_path, map_location=self.device)
+        self.model.load_state_dict(checkpoint['model_state_dict'])
+        self.model = self.model.to(self.device)
+        self.model.eval()
+
+        # Image transforms
+        self.transform = transforms.Compose([
+            transforms.Resize((224, 224)),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                               std=[0.229, 0.224, 0.225])
+        ])
+
+        print(f"Model loaded from {checkpoint_path}")
+        print(f"Device: {self.device}")
+        print(f"Classes: {self.num_classes}")
+
+    def predict(self, image):
+        """
+        Predict class for input image
+
+        Args:
+            image: PIL Image
+
+        Returns:
+            dict: {class_name: probability}
+        """
+        # Preprocess
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+
+        img_tensor = self.transform(image).unsqueeze(0)
+        img_tensor = img_tensor.to(self.device)
+
+        # Predict
+        with torch.no_grad():
+            outputs = self.model(img_tensor)
+            probabilities = torch.nn.functional.softmax(outputs, dim=1)
+            probs = probabilities[0].cpu().numpy()
+
+        # Create results dictionary
+        results = {}
+        for idx, prob in enumerate(probs):
+            class_name = self.id_to_label[idx]
+            results[class_name] = float(prob)
+
+        # Sort by probability
+        results = dict(sorted(results.items(), key=lambda x: x[1], reverse=True))
+
+        return results
+
+
+def create_demo(predictor):
+    """Create Gradio interface"""
+
+    def predict_image(image):
+        """Prediction function for Gradio"""
+        if image is None:
+            return None
+
+        results = predictor.predict(image)
+        return results
+
+    # Create interface
+    demo = gr.Interface(
+        fn=predict_image,
+        inputs=gr.Image(type="pil", label="Upload Image"),
+        outputs=gr.Label(num_top_classes=10, label="Predictions"),
+        title="🌿 Pest and Disease Classification",
+        description="Upload an image of a citrus plant leaf to classify if it's healthy or has pests/diseases.",
+        examples=None,
+        theme=gr.themes.Soft(),
+        allow_flagging="never"
+    )
+
+    return demo
+
+
+def main(args):
+    """Main function"""
+    print("Starting Pest and Disease Classification Demo...")
+    print("=" * 60)
+
+    # Create predictor
+    predictor = PestDiseasePredictor(
+        checkpoint_path=args.checkpoint,
+        label_mapping_path=args.label_mapping,
+        backbone=args.backbone,
+        device=args.device
+    )
+
+    # Create and launch demo
+    demo = create_demo(predictor)
+
+    print("\n" + "=" * 60)
+    print("Launching demo...")
+    print("=" * 60)
+
+    demo.launch(
+        server_name=args.host,
+        server_port=args.port,
+        share=args.share
+    )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Demo for Pest and Disease Classification')
+
+    parser.add_argument('--checkpoint', type=str, default='checkpoints/best_model.pth',
+                       help='Path to model checkpoint')
+    parser.add_argument('--label_mapping', type=str, default='label_mapping.json',
+                       help='Path to label mapping JSON')
+    parser.add_argument('--backbone', type=str, default='resnet50',
+                       choices=['resnet50', 'resnet101', 'efficientnet_b0',
+                               'efficientnet_b3', 'mobilenet_v2'],
+                       help='Model backbone')
+    parser.add_argument('--device', type=str, default='cuda',
+                       choices=['cuda', 'cpu'],
+                       help='Device to use')
+    parser.add_argument('--host', type=str, default='127.0.0.1',
+                       help='Server host')
+    parser.add_argument('--port', type=int, default=7860,
+                       help='Server port')
+    parser.add_argument('--share', action='store_true',
+                       help='Create public link')
+
+    args = parser.parse_args()
+    main(args)

evaluate.py

@@ -0,0 +1,327 @@
+"""
+Evaluation script for Pest and Disease Classification
+Generate confusion matrix, classification report, and per-class metrics
+"""
+
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+import seaborn as sns
+from sklearn.metrics import confusion_matrix, classification_report, f1_score
+import argparse
+import json
+from pathlib import Path
+
+from dataset import get_dataloaders
+from model import create_model
+
+
+def evaluate_model(model, dataloader, device, dataset):
+    """
+    Evaluate model on a dataset
+
+    Returns:
+        predictions: List of predicted labels
+        true_labels: List of true labels
+        accuracy: Overall accuracy
+    """
+    model.eval()
+    all_preds = []
+    all_labels = []
+
+    with torch.no_grad():
+        for inputs, labels in dataloader:
+            inputs = inputs.to(device)
+            labels = labels.to(device)
+
+            outputs = model(inputs)
+            _, preds = torch.max(outputs, 1)
+
+            all_preds.extend(preds.cpu().numpy())
+            all_labels.extend(labels.cpu().numpy())
+
+    all_preds = np.array(all_preds)
+    all_labels = np.array(all_labels)
+    accuracy = np.mean(all_preds == all_labels)
+
+    return all_preds, all_labels, accuracy
+
+
+def plot_confusion_matrix(y_true, y_pred, class_names, save_path='confusion_matrix.png'):
+    """
+    Plot and save confusion matrix
+
+    Args:
+        y_true: True labels
+        y_pred: Predicted labels
+        class_names: List of class names
+        save_path: Path to save figure
+    """
+    cm = confusion_matrix(y_true, y_pred)
+
+    # Calculate percentages
+    cm_percent = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis] * 100
+
+    # Create figure
+    plt.figure(figsize=(12, 10))
+
+    # Plot with annotations
+    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
+                xticklabels=class_names,
+                yticklabels=class_names,
+                cbar_kws={'label': 'Count'})
+
+    plt.title('Confusion Matrix', fontsize=16, pad=20)
+    plt.ylabel('True Label', fontsize=12)
+    plt.xlabel('Predicted Label', fontsize=12)
+    plt.xticks(rotation=45, ha='right')
+    plt.yticks(rotation=0)
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=300, bbox_inches='tight')
+    print(f"Confusion matrix saved to {save_path}")
+
+    # Also save percentage version
+    plt.figure(figsize=(12, 10))
+    sns.heatmap(cm_percent, annot=True, fmt='.1f', cmap='Blues',
+                xticklabels=class_names,
+                yticklabels=class_names,
+                cbar_kws={'label': 'Percentage (%)'})
+
+    plt.title('Confusion Matrix (Percentage)', fontsize=16, pad=20)
+    plt.ylabel('True Label', fontsize=12)
+    plt.xlabel('Predicted Label', fontsize=12)
+    plt.xticks(rotation=45, ha='right')
+    plt.yticks(rotation=0)
+    plt.tight_layout()
+
+    save_path_percent = str(save_path).replace('.png', '_percent.png')
+    plt.savefig(save_path_percent, dpi=300, bbox_inches='tight')
+    print(f"Confusion matrix (percentage) saved to {save_path_percent}")
+
+    plt.close('all')
+
+    return cm
+
+
+def generate_classification_report(y_true, y_pred, class_names, save_path='classification_report.txt'):
+    """
+    Generate and save detailed classification report
+
+    Args:
+        y_true: True labels
+        y_pred: Predicted labels
+        class_names: List of class names
+        save_path: Path to save report
+    """
+    # Generate report
+    report = classification_report(
+        y_true, y_pred,
+        target_names=class_names,
+        digits=4
+    )
+
+    # Print to console
+    print("\n" + "=" * 80)
+    print("Classification Report")
+    print("=" * 80)
+    print(report)
+
+    # Save to file
+    with open(save_path, 'w', encoding='utf-8') as f:
+        f.write("Classification Report\n")
+        f.write("=" * 80 + "\n")
+        f.write(report)
+
+    print(f"\nClassification report saved to {save_path}")
+
+    # Calculate per-class metrics
+    from sklearn.metrics import precision_recall_fscore_support
+    precision, recall, f1, support = precision_recall_fscore_support(
+        y_true, y_pred, average=None
+    )
+
+    # Create detailed metrics dictionary
+    metrics = {}
+    for i, class_name in enumerate(class_names):
+        metrics[class_name] = {
+            'precision': float(precision[i]),
+            'recall': float(recall[i]),
+            'f1-score': float(f1[i]),
+            'support': int(support[i])
+        }
+
+    # Add overall metrics
+    metrics['overall'] = {
+        'accuracy': float(np.mean(y_true == y_pred)),
+        'macro_avg_f1': float(np.mean(f1)),
+        'weighted_avg_f1': float(f1_score(y_true, y_pred, average='weighted'))
+    }
+
+    # Save metrics as JSON
+    metrics_path = str(save_path).replace('.txt', '.json')
+    with open(metrics_path, 'w', encoding='utf-8') as f:
+        json.dump(metrics, f, indent=2, ensure_ascii=False)
+
+    print(f"Metrics JSON saved to {metrics_path}")
+
+    return metrics
+
+
+def plot_per_class_metrics(metrics, class_names, save_path='per_class_metrics.png'):
+    """
+    Plot per-class precision, recall, and F1-score
+
+    Args:
+        metrics: Dictionary of metrics
+        class_names: List of class names
+        save_path: Path to save figure
+    """
+    precision = [metrics[name]['precision'] for name in class_names]
+    recall = [metrics[name]['recall'] for name in class_names]
+    f1 = [metrics[name]['f1-score'] for name in class_names]
+
+    x = np.arange(len(class_names))
+    width = 0.25
+
+    fig, ax = plt.subplots(figsize=(14, 6))
+    ax.bar(x - width, precision, width, label='Precision', alpha=0.8)
+    ax.bar(x, recall, width, label='Recall', alpha=0.8)
+    ax.bar(x + width, f1, width, label='F1-Score', alpha=0.8)
+
+    ax.set_xlabel('Class', fontsize=12)
+    ax.set_ylabel('Score', fontsize=12)
+    ax.set_title('Per-Class Metrics', fontsize=14, pad=20)
+    ax.set_xticks(x)
+    ax.set_xticklabels(class_names, rotation=45, ha='right')
+    ax.legend()
+    ax.grid(axis='y', alpha=0.3)
+    ax.set_ylim([0, 1.1])
+
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=300, bbox_inches='tight')
+    print(f"Per-class metrics plot saved to {save_path}")
+    plt.close()
+
+
+def main(args):
+    """Main evaluation function"""
+    print("Pest and Disease Classification Evaluation")
+    print("=" * 80)
+    print(f"Configuration:")
+    print(f"  Checkpoint: {args.checkpoint}")
+    print(f"  Split: {args.split}")
+    print(f"  Batch size: {args.batch_size}")
+    print(f"  Device: {args.device}")
+    print("=" * 80)
+
+    # Set device
+    device = torch.device(args.device if torch.cuda.is_available() else 'cpu')
+    print(f"\nUsing device: {device}")
+
+    # Load data
+    print("\nLoading datasets...")
+    loaders = get_dataloaders(
+        csv_file=args.csv_file,
+        label_mapping_file=args.label_mapping,
+        batch_size=args.batch_size,
+        img_size=args.img_size,
+        num_workers=args.num_workers
+    )
+
+    # Get class names
+    dataset = loaders['datasets'][args.split]
+    class_names = [dataset.get_label_name(i) for i in range(dataset.num_classes)]
+    print(f"Classes: {class_names}")
+
+    # Create model
+    print(f"\nCreating model: {args.backbone}")
+    model = create_model(
+        num_classes=loaders['num_classes'],
+        backbone=args.backbone,
+        pretrained=False
+    )
+
+    # Load checkpoint
+    print(f"\nLoading checkpoint: {args.checkpoint}")
+    checkpoint = torch.load(args.checkpoint, map_location=device)
+    model.load_state_dict(checkpoint['model_state_dict'])
+    model = model.to(device)
+
+    if 'val_acc' in checkpoint:
+        print(f"Checkpoint validation accuracy: {checkpoint['val_acc']:.4f}")
+
+    # Evaluate
+    print(f"\nEvaluating on {args.split} set...")
+    dataloader = loaders[args.split]
+    predictions, true_labels, accuracy = evaluate_model(model, dataloader, device, dataset)
+
+    print(f"\n{args.split.capitalize()} Set Accuracy: {accuracy:.4f}")
+
+    # Create output directory
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(exist_ok=True)
+
+    # Generate confusion matrix
+    print("\nGenerating confusion matrix...")
+    cm = plot_confusion_matrix(
+        true_labels, predictions, class_names,
+        save_path=output_dir / f'confusion_matrix_{args.split}.png'
+    )
+
+    # Generate classification report
+    print("\nGenerating classification report...")
+    metrics = generate_classification_report(
+        true_labels, predictions, class_names,
+        save_path=output_dir / f'classification_report_{args.split}.txt'
+    )
+
+    # Plot per-class metrics
+    print("\nGenerating per-class metrics plot...")
+    plot_per_class_metrics(
+        metrics, class_names,
+        save_path=output_dir / f'per_class_metrics_{args.split}.png'
+    )
+
+    print("\n" + "=" * 80)
+    print("Evaluation complete!")
+    print(f"Results saved to {output_dir}/")
+    print("=" * 80)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Evaluate Pest and Disease Classifier')
+
+    # Data parameters
+    parser.add_argument('--csv_file', type=str, default='dataset.csv',
+                       help='Path to dataset CSV')
+    parser.add_argument('--label_mapping', type=str, default='label_mapping.json',
+                       help='Path to label mapping JSON')
+
+    # Model parameters
+    parser.add_argument('--checkpoint', type=str, default='checkpoints/best_model.pth',
+                       help='Path to model checkpoint')
+    parser.add_argument('--backbone', type=str, default='resnet50',
+                       choices=['resnet50', 'resnet101', 'efficientnet_b0',
+                               'efficientnet_b3', 'mobilenet_v2'],
+                       help='Model backbone')
+
+    # Evaluation parameters
+    parser.add_argument('--split', type=str, default='test',
+                       choices=['train', 'val', 'test'],
+                       help='Dataset split to evaluate')
+    parser.add_argument('--batch_size', type=int, default=16,
+                       help='Batch size')
+    parser.add_argument('--img_size', type=int, default=224,
+                       help='Image size')
+
+    # System parameters
+    parser.add_argument('--device', type=str, default='cuda',
+                       choices=['cuda', 'cpu'],
+                       help='Device to use')
+    parser.add_argument('--num_workers', type=int, default=4,
+                       help='Number of data loading workers')
+    parser.add_argument('--output_dir', type=str, default='evaluation_results',
+                       help='Directory to save results')
+
+    args = parser.parse_args()
+    main(args)

label_mapping.json

@@ -0,0 +1,27 @@
+{
+  "label_to_id": {
+    "介殼蟲": 0,
+    "健康植株-椪柑": 1,
+    "健康植株-茂谷柑": 2,
+    "油斑病": 3,
+    "潛葉蛾": 4,
+    "潰瘍病": 5,
+    "煤煙病": 6,
+    "薊馬": 7,
+    "蚜蟲": 8,
+    "黑點病": 9
+  },
+  "id_to_label": {
+    "0": "介殼蟲",
+    "1": "健康植株-椪柑",
+    "2": "健康植株-茂谷柑",
+    "3": "油斑病",
+    "4": "潛葉蛾",
+    "5": "潰瘍病",
+    "6": "煤煙病",
+    "7": "薊馬",
+    "8": "蚜蟲",
+    "9": "黑點病"
+  },
+  "num_classes": 10
+}

model.py

@@ -0,0 +1,159 @@
+"""
+Classification Models for Pest and Disease Detection
+Supports multiple pretrained backbones: ResNet, EfficientNet, MobileNet
+"""
+
+import torch
+import torch.nn as nn
+import torchvision.models as models
+
+
+class PestDiseaseClassifier(nn.Module):
+    """
+    General classifier with pretrained backbone for transfer learning
+    """
+
+    def __init__(self, num_classes=10, backbone='resnet50', pretrained=True, dropout=0.3):
+        """
+        Args:
+            num_classes (int): Number of output classes
+            backbone (str): Backbone architecture ('resnet50', 'resnet101', 'efficientnet_b0',
+                           'efficientnet_b3', 'mobilenet_v2')
+            pretrained (bool): Use pretrained weights
+            dropout (float): Dropout rate for regularization
+        """
+        super(PestDiseaseClassifier, self).__init__()
+
+        self.backbone_name = backbone
+        self.num_classes = num_classes
+
+        # Select backbone
+        if backbone == 'resnet50':
+            self.backbone = models.resnet50(pretrained=pretrained)
+            num_features = self.backbone.fc.in_features
+            self.backbone.fc = nn.Identity()
+
+        elif backbone == 'resnet101':
+            self.backbone = models.resnet101(pretrained=pretrained)
+            num_features = self.backbone.fc.in_features
+            self.backbone.fc = nn.Identity()
+
+        elif backbone == 'efficientnet_b0':
+            self.backbone = models.efficientnet_b0(pretrained=pretrained)
+            num_features = self.backbone.classifier[1].in_features
+            self.backbone.classifier = nn.Identity()
+
+        elif backbone == 'efficientnet_b3':
+            self.backbone = models.efficientnet_b3(pretrained=pretrained)
+            num_features = self.backbone.classifier[1].in_features
+            self.backbone.classifier = nn.Identity()
+
+        elif backbone == 'mobilenet_v2':
+            self.backbone = models.mobilenet_v2(pretrained=pretrained)
+            num_features = self.backbone.classifier[1].in_features
+            self.backbone.classifier = nn.Identity()
+
+        else:
+            raise ValueError(f"Unknown backbone: {backbone}")
+
+        # Custom classifier head
+        self.classifier = nn.Sequential(
+            nn.Dropout(dropout),
+            nn.Linear(num_features, 512),
+            nn.ReLU(inplace=True),
+            nn.Dropout(dropout),
+            nn.Linear(512, num_classes)
+        )
+
+        print(f"Model created: {backbone}")
+        print(f"  Features: {num_features}")
+        print(f"  Classes: {num_classes}")
+        print(f"  Pretrained: {pretrained}")
+
+    def forward(self, x):
+        """
+        Forward pass
+        Args:
+            x: Input tensor [batch_size, 3, H, W]
+        Returns:
+            logits: Output tensor [batch_size, num_classes]
+        """
+        features = self.backbone(x)
+        logits = self.classifier(features)
+        return logits
+
+    def freeze_backbone(self):
+        """Freeze backbone parameters for fine-tuning"""
+        for param in self.backbone.parameters():
+            param.requires_grad = False
+        print("Backbone frozen")
+
+    def unfreeze_backbone(self):
+        """Unfreeze backbone parameters"""
+        for param in self.backbone.parameters():
+            param.requires_grad = True
+        print("Backbone unfrozen")
+
+
+def create_model(num_classes=10, backbone='resnet50', pretrained=True, dropout=0.3):
+    """
+    Factory function to create model
+
+    Args:
+        num_classes (int): Number of classes
+        backbone (str): Model architecture
+        pretrained (bool): Use pretrained weights
+        dropout (float): Dropout rate
+
+    Returns:
+        model: PestDiseaseClassifier instance
+    """
+    model = PestDiseaseClassifier(
+        num_classes=num_classes,
+        backbone=backbone,
+        pretrained=pretrained,
+        dropout=dropout
+    )
+    return model
+
+
+def count_parameters(model):
+    """Count total and trainable parameters"""
+    total_params = sum(p.numel() for p in model.parameters())
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+    print(f"\nModel Parameters:")
+    print(f"  Total: {total_params:,}")
+    print(f"  Trainable: {trainable_params:,}")
+    print(f"  Non-trainable: {total_params - trainable_params:,}")
+
+    return total_params, trainable_params
+
+
+if __name__ == "__main__":
+    """Test model creation"""
+    print("Testing Pest and Disease Classification Models")
+    print("=" * 60)
+
+    # Test different backbones
+    backbones = ['resnet50', 'efficientnet_b0', 'mobilenet_v2']
+
+    for backbone in backbones:
+        print(f"\nTesting {backbone}...")
+        print("-" * 60)
+
+        model = create_model(num_classes=10, backbone=backbone, pretrained=True)
+        count_parameters(model)
+
+        # Test forward pass
+        dummy_input = torch.randn(2, 3, 224, 224)
+        with torch.no_grad():
+            output = model(dummy_input)
+
+        print(f"  Input shape: {dummy_input.shape}")
+        print(f"  Output shape: {output.shape}")
+        print(f"  Output range: [{output.min():.3f}, {output.max():.3f}]")
+
+    print("\n" + "=" * 60)
+    print("Model test completed successfully!")
+    print("=" * 60)

prepare_dataset.py

@@ -0,0 +1,247 @@
+"""
+Pest and Disease Classification Dataset Preparation Script
+- Scan data folders
+- Analyze image distribution
+- Generate train/val/test CSV files
+"""
+
+import pandas as pd
+import numpy as np
+from pathlib import Path
+from PIL import Image
+from sklearn.model_selection import train_test_split
+import json
+
+# Configuration parameters
+DATA_DIR = "Data"
+OUTPUT_CSV = "dataset.csv"
+TRAIN_RATIO = 0.7
+VAL_RATIO = 0.15
+TEST_RATIO = 0.15
+RANDOM_SEED = 42
+
+# Set random seed
+np.random.seed(RANDOM_SEED)
+
+def scan_dataset():
+    """Scan dataset and collect all image information"""
+    data_list = []
+    image_sizes = []
+
+    category_mapping = {
+        "A.健康植株": {
+            "椪柑": "健康植株-椪柑",
+            "茂谷柑": "健康植株-茂谷柑"
+        },
+        "B.病害": {
+            "1.病害-潰瘍病": "潰瘍病",
+            "2.病害-煤煙病": "煤煙病",
+            "3.病害-油斑病": "油斑病",
+            "4.病害-黑點病2": "黑點病"
+        },
+        "C.蟲害": {
+            "1.蟲害-薊馬": "薊馬",
+            "2.蟲害-潛葉蛾": "潛葉蛾",
+            "3.蟲害-蚜蟲": "蚜蟲",
+            "4.蟲害-介殼蟲": "介殼蟲"
+        }
+    }
+
+    print("Scanning dataset...")
+
+    for main_dir in ["A.健康植株", "B.病害", "C.蟲害"]:
+        main_path = Path(DATA_DIR) / main_dir
+
+        if not main_path.exists():
+            print(f"Warning: {main_path} does not exist")
+            continue
+
+        # Iterate through subdirectories
+        for sub_dir in main_path.iterdir():
+            if not sub_dir.is_dir():
+                continue
+
+            # Determine class label
+            try:
+                label = category_mapping[main_dir][sub_dir.name]
+                print(f"  Processing: {main_dir}/{sub_dir.name} -> {label}")
+            except KeyError:
+                print(f"  Warning: Unknown subdirectory {main_dir}/{sub_dir.name}, skipping...")
+                continue
+
+            # Store plant type info
+            if main_dir == "A.健康植株":
+                plant_type = sub_dir.name  # Ponkan or Murcott
+            else:
+                plant_type = "柑橘"
+
+            # Scan images (case-insensitive)
+            image_files = (list(sub_dir.glob("*.jpg")) + list(sub_dir.glob("*.JPG")) +
+                          list(sub_dir.glob("*.jpeg")) + list(sub_dir.glob("*.JPEG")) +
+                          list(sub_dir.glob("*.png")) + list(sub_dir.glob("*.PNG")))
+
+            for img_path in image_files:
+                try:
+                    # Get image dimensions
+                    with Image.open(img_path) as img:
+                        width, height = img.size
+                        image_sizes.append((width, height))
+
+                    data_list.append({
+                        'image_path': str(img_path),
+                        'label': label,
+                        'main_category': main_dir.split('.')[1],
+                        'plant_type': plant_type,
+                        'width': width,
+                        'height': height
+                    })
+                except Exception as e:
+                    print(f"Warning: Cannot read {img_path}: {e}")
+
+    return data_list, image_sizes
+
+def analyze_dataset(data_list, image_sizes):
+    """Analyze dataset statistics"""
+    df = pd.DataFrame(data_list)
+
+    print("\n" + "="*60)
+    print("Dataset Statistics")
+    print("="*60)
+
+    # Overall statistics
+    print(f"\nTotal images: {len(df)}")
+    print(f"\nClass distribution:")
+    label_counts = df['label'].value_counts()
+    for label, count in label_counts.items():
+        print(f"  {label}: {count} images ({count/len(df)*100:.1f}%)")
+
+    # Image size analysis
+    if image_sizes:
+        widths, heights = zip(*image_sizes)
+        print(f"\nImage size analysis:")
+        print(f"  Width: min={min(widths)}, max={max(widths)}, avg={np.mean(widths):.0f}")
+        print(f"  Height: min={min(heights)}, max={max(heights)}, avg={np.mean(heights):.0f}")
+
+        # Check size consistency
+        unique_sizes = set(image_sizes)
+        print(f"  Unique sizes: {len(unique_sizes)}")
+        if len(unique_sizes) <= 5:
+            print(f"  Main sizes: {list(unique_sizes)[:5]}")
+
+    # Check class imbalance
+    max_count = label_counts.max()
+    min_count = label_counts.min()
+    imbalance_ratio = max_count / min_count
+    print(f"\nClass imbalance ratio: {imbalance_ratio:.2f}x")
+    if imbalance_ratio > 3:
+        print("  Warning: Severe class imbalance detected. Consider using weighted loss or data augmentation")
+
+    return df
+
+def split_dataset(df, train_ratio=0.7, val_ratio=0.15, test_ratio=0.15):
+    """Split dataset into train/val/test sets with stratified sampling"""
+    assert abs(train_ratio + val_ratio + test_ratio - 1.0) < 1e-6, "Ratios must sum to 1"
+
+    print("\n" + "="*60)
+    print("Splitting Dataset (Stratified Sampling)")
+    print("="*60)
+
+    # First split out test set
+    train_val_df, test_df = train_test_split(
+        df,
+        test_size=test_ratio,
+        stratify=df['label'],
+        random_state=RANDOM_SEED
+    )
+
+    # Then split train and validation from remaining data
+    val_ratio_adjusted = val_ratio / (train_ratio + val_ratio)
+    train_df, val_df = train_test_split(
+        train_val_df,
+        test_size=val_ratio_adjusted,
+        stratify=train_val_df['label'],
+        random_state=RANDOM_SEED
+    )
+
+    # Add split column
+    train_df = train_df.copy()
+    val_df = val_df.copy()
+    test_df = test_df.copy()
+
+    train_df['split'] = 'train'
+    val_df['split'] = 'val'
+    test_df['split'] = 'test'
+
+    # Merge all splits
+    final_df = pd.concat([train_df, val_df, test_df], ignore_index=True)
+
+    # Display class distribution for each split
+    print(f"\nTrain set: {len(train_df)} images ({len(train_df)/len(df)*100:.1f}%)")
+    print(train_df['label'].value_counts().to_string())
+
+    print(f"\nValidation set: {len(val_df)} images ({len(val_df)/len(df)*100:.1f}%)")
+    print(val_df['label'].value_counts().to_string())
+
+    print(f"\nTest set: {len(test_df)} images ({len(test_df)/len(df)*100:.1f}%)")
+    print(test_df['label'].value_counts().to_string())
+
+    return final_df
+
+def save_dataset(df, output_path):
+    """Save dataset CSV and label mapping"""
+    # Save complete CSV
+    df.to_csv(output_path, index=False, encoding='utf-8-sig')
+    print(f"\nDataset saved to: {output_path}")
+
+    # Create label to ID mapping
+    unique_labels = sorted(df['label'].unique())
+    label_to_id = {label: idx for idx, label in enumerate(unique_labels)}
+    id_to_label = {idx: label for label, idx in label_to_id.items()}
+
+    # Save label mapping
+    mapping_file = "label_mapping.json"
+    with open(mapping_file, 'w', encoding='utf-8') as f:
+        json.dump({
+            'label_to_id': label_to_id,
+            'id_to_label': id_to_label,
+            'num_classes': len(unique_labels)
+        }, f, ensure_ascii=False, indent=2)
+
+    print(f"Label mapping saved to: {mapping_file}")
+    print(f"\nLabel mapping ({len(unique_labels)} classes):")
+    for label, idx in label_to_id.items():
+        print(f"  {idx}: {label}")
+
+    return label_to_id
+
+def main():
+    """Main function"""
+    print("Pest and Disease Dataset Preparation Tool")
+    print("="*60)
+
+    # 1. Scan dataset
+    data_list, image_sizes = scan_dataset()
+
+    if not data_list:
+        print("Error: No images found!")
+        return
+
+    # 2. Analyze dataset
+    df = analyze_dataset(data_list, image_sizes)
+
+    # 3. Split dataset
+    final_df = split_dataset(df, TRAIN_RATIO, VAL_RATIO, TEST_RATIO)
+
+    # 4. Save dataset
+    label_to_id = save_dataset(final_df, OUTPUT_CSV)
+
+    print("\n" + "="*60)
+    print("Dataset preparation completed!")
+    print("="*60)
+    print("\nNext steps:")
+    print("  1. Check dataset.csv and label_mapping.json")
+    print("  2. Run data loader test script")
+    print("  3. Start model training")
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,10 @@
+torch
+torchvision
+pillow
+gradio==4.39.0
+huggingface_hub==0.25.2
+rich
+seaborn
+pathlib
+pandas
+pydantic==2.10.6

train.py

@@ -0,0 +1,286 @@
+"""
+Simple Training Script for Pest and Disease Classification
+Using Rich for progress display
+"""
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from pathlib import Path
+import json
+import argparse
+from rich.console import Console
+from rich.progress import Progress, SpinnerColumn, BarColumn, TextColumn, TimeRemainingColumn
+from rich.table import Table
+from rich.panel import Panel
+
+from dataset import get_dataloaders, calculate_class_weights
+from model import create_model
+
+
+console = Console()
+
+
+def train_epoch(model, dataloader, criterion, optimizer, device, progress, task):
+    """Train for one epoch with progress bar"""
+    model.train()
+    running_loss = 0.0
+    running_corrects = 0
+    total_samples = 0
+
+    for inputs, labels in dataloader:
+        inputs = inputs.to(device)
+        labels = labels.to(device)
+
+        optimizer.zero_grad()
+        outputs = model(inputs)
+        loss = criterion(outputs, labels)
+        _, preds = torch.max(outputs, 1)
+
+        loss.backward()
+        optimizer.step()
+
+        running_loss += loss.item() * inputs.size(0)
+        running_corrects += torch.sum(preds == labels.data)
+        total_samples += inputs.size(0)
+
+        progress.update(task, advance=1)
+
+    epoch_loss = running_loss / total_samples
+    epoch_acc = running_corrects.double() / total_samples
+    return epoch_loss, epoch_acc.item()
+
+
+def validate_epoch(model, dataloader, criterion, device, progress, task):
+    """Validate for one epoch with progress bar"""
+    model.eval()
+    running_loss = 0.0
+    running_corrects = 0
+    total_samples = 0
+
+    with torch.no_grad():
+        for inputs, labels in dataloader:
+            inputs = inputs.to(device)
+            labels = labels.to(device)
+
+            outputs = model(inputs)
+            loss = criterion(outputs, labels)
+            _, preds = torch.max(outputs, 1)
+
+            running_loss += loss.item() * inputs.size(0)
+            running_corrects += torch.sum(preds == labels.data)
+            total_samples += inputs.size(0)
+
+            progress.update(task, advance=1)
+
+    epoch_loss = running_loss / total_samples
+    epoch_acc = running_corrects.double() / total_samples
+    return epoch_loss, epoch_acc.item()
+
+
+def train_model(model, train_loader, val_loader, criterion, optimizer,
+                num_epochs, device, save_dir):
+    """
+    Simple training loop with Rich progress display
+    """
+    save_dir = Path(save_dir)
+    save_dir.mkdir(exist_ok=True)
+
+    best_val_acc = 0.0
+    history = {
+        'train_loss': [],
+        'train_acc': [],
+        'val_loss': [],
+        'val_acc': []
+    }
+
+    console.print("\n[bold green]Starting Training[/bold green]")
+
+    for epoch in range(num_epochs):
+        console.print(f"\n[bold cyan]Epoch {epoch+1}/{num_epochs}[/bold cyan]")
+
+        with Progress(
+            SpinnerColumn(),
+            TextColumn("[progress.description]{task.description}"),
+            BarColumn(),
+            TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
+            TimeRemainingColumn(),
+            console=console
+        ) as progress:
+
+            # Training
+            train_task = progress.add_task(
+                "[red]Training...",
+                total=len(train_loader)
+            )
+            train_loss, train_acc = train_epoch(
+                model, train_loader, criterion, optimizer,
+                device, progress, train_task
+            )
+
+            # Validation
+            val_task = progress.add_task(
+                "[green]Validating...",
+                total=len(val_loader)
+            )
+            val_loss, val_acc = validate_epoch(
+                model, val_loader, criterion, device,
+                progress, val_task
+            )
+
+        # Create results table
+        table = Table(show_header=True, header_style="bold magenta")
+        table.add_column("Split", style="cyan")
+        table.add_column("Loss", justify="right", style="yellow")
+        table.add_column("Accuracy", justify="right", style="green")
+
+        table.add_row("Train", f"{train_loss:.4f}", f"{train_acc:.4f}")
+        table.add_row("Val", f"{val_loss:.4f}", f"{val_acc:.4f}")
+
+        console.print(table)
+
+        # Save history
+        history['train_loss'].append(train_loss)
+        history['train_acc'].append(train_acc)
+        history['val_loss'].append(val_loss)
+        history['val_acc'].append(val_acc)
+
+        # Save best model
+        if val_acc > best_val_acc:
+            best_val_acc = val_acc
+            torch.save({
+                'epoch': epoch,
+                'model_state_dict': model.state_dict(),
+                'optimizer_state_dict': optimizer.state_dict(),
+                'val_acc': val_acc,
+                'val_loss': val_loss,
+            }, save_dir / 'best_model.pth')
+            console.print(f"[bold green]✓ Saved best model (Val Acc: {val_acc:.4f})[/bold green]")
+
+        # Save checkpoint every 10 epochs
+        if (epoch + 1) % 10 == 0:
+            torch.save({
+                'epoch': epoch,
+                'model_state_dict': model.state_dict(),
+                'optimizer_state_dict': optimizer.state_dict(),
+                'val_acc': val_acc,
+                'val_loss': val_loss,
+            }, save_dir / f'checkpoint_epoch_{epoch+1}.pth')
+            console.print(f"[yellow]Checkpoint saved at epoch {epoch+1}[/yellow]")
+
+    # Save training history
+    with open(save_dir / 'training_history.json', 'w') as f:
+        json.dump(history, f, indent=2)
+
+    console.print(f"\n[bold green]Training Complete![/bold green]")
+    console.print(f"[bold]Best Val Acc: {best_val_acc:.4f}[/bold]")
+    console.print(f"[bold]Results saved to: {save_dir}/[/bold]")
+
+    return model, history
+
+
+def main(args):
+    """Main training function"""
+    # Print configuration
+    config_panel = Panel.fit(
+        f"""[bold]Configuration[/bold]
+Backbone: {args.backbone}
+Batch Size: {args.batch_size}
+Image Size: {args.img_size}
+Epochs: {args.epochs}
+Learning Rate: {args.lr}
+Optimizer: {args.optimizer}
+Device: {args.device}
+Class Weights: {args.use_class_weights}""",
+        title="Training Settings",
+        border_style="blue"
+    )
+    console.print(config_panel)
+
+    # Set device
+    device = torch.device(args.device if torch.cuda.is_available() else 'cpu')
+    console.print(f"\n[bold]Using device: {device}[/bold]")
+
+    # Load data
+    console.print("\n[bold]Loading datasets...[/bold]")
+    loaders = get_dataloaders(
+        csv_file=args.csv_file,
+        label_mapping_file=args.label_mapping,
+        batch_size=args.batch_size,
+        img_size=args.img_size,
+        num_workers=args.num_workers
+    )
+
+    # Create model
+    console.print(f"\n[bold]Creating model: {args.backbone}[/bold]")
+    model = create_model(
+        num_classes=loaders['num_classes'],
+        backbone=args.backbone,
+        pretrained=True,
+        dropout=args.dropout
+    )
+    model = model.to(device)
+
+    # Loss function
+    if args.use_class_weights:
+        class_weights = calculate_class_weights(args.csv_file, args.label_mapping)
+        class_weights = class_weights.to(device)
+        criterion = nn.CrossEntropyLoss(weight=class_weights)
+        console.print("[bold]Using weighted CrossEntropyLoss[/bold]")
+    else:
+        criterion = nn.CrossEntropyLoss()
+        console.print("[bold]Using CrossEntropyLoss[/bold]")
+
+    # Optimizer
+    if args.optimizer == 'adam':
+        optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
+    elif args.optimizer == 'adamw':
+        optimizer = optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
+    elif args.optimizer == 'sgd':
+        optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=0.9,
+                             weight_decay=args.weight_decay)
+
+    # Train model
+    model, history = train_model(
+        model=model,
+        train_loader=loaders['train'],
+        val_loader=loaders['val'],
+        criterion=criterion,
+        optimizer=optimizer,
+        num_epochs=args.epochs,
+        device=device,
+        save_dir=args.save_dir
+    )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Simple Training for Pest and Disease Classifier')
+
+    # Data parameters
+    parser.add_argument('--csv_file', type=str, default='dataset.csv')
+    parser.add_argument('--label_mapping', type=str, default='label_mapping.json')
+
+    # Model parameters
+    parser.add_argument('--backbone', type=str, default='resnet50',
+                       choices=['resnet50', 'resnet101', 'efficientnet_b0',
+                               'efficientnet_b3', 'mobilenet_v2'])
+    parser.add_argument('--dropout', type=float, default=0.3)
+
+    # Training parameters
+    parser.add_argument('--batch_size', type=int, default=64)
+    parser.add_argument('--img_size', type=int, default=224)
+    parser.add_argument('--epochs', type=int, default=50)
+    parser.add_argument('--lr', type=float, default=0.001)
+    parser.add_argument('--optimizer', type=str, default='adamw',
+                       choices=['adam', 'adamw', 'sgd'])
+    parser.add_argument('--weight_decay', type=float, default=0.01)
+    parser.add_argument('--use_class_weights', action='store_true')
+
+    # System parameters
+    parser.add_argument('--device', type=str, default='cuda',
+                       choices=['cuda', 'cpu'])
+    parser.add_argument('--num_workers', type=int, default=8)
+    parser.add_argument('--save_dir', type=str, default='checkpoints')
+
+    args = parser.parse_args()
+    main(args)