Update model with improved regularization and data augmentation

README.md

@@ -7,25 +7,76 @@ tags:
 - pytorch
 - resnet
 - beans
+- agriculture
+- plant-disease
 datasets:
 - beans
 library_name: pytorch
 pipeline_tag: image-classification
+metrics:
+- accuracy
 ---
 
-# ResNet18 fine-tuned on Beans dataset
+# ResNet18 Fine-tuned on Beans Dataset
 
-This model was trained in Google Colab using a GPU and tracked with MLflow.
+This model was trained in Google Colab using a T4 GPU and tracked with MLflow.
 
-Dataset: [Beans](https://huggingface.co/datasets/AI-Lab-Makerere/beans)
+## Model Details
 
-## Classes:
+**Dataset:** [Beans](https://huggingface.co/datasets/AI-Lab-Makerere/beans)
+
+**Classes:**
 - Healthy
 - Bean Rust
 - Angular Leaf Spot
 
-Validation Accuracy: 0.9774
+**Validation Accuracy:** 0.9173
+
+## Training Configuration
+
+**Overfitting Prevention Techniques:**
+- Data augmentation (rotation, flip, crop, color jitter)
+- Dropout (30%)
+- L2 regularization (weight decay: 1e-4)
+- Learning rate scheduling (ReduceLROnPlateau)
+- Best model selection based on validation accuracy
+
+**Hyperparameters:**
+- Learning Rate: 5e-05
+- Epochs: 5
+- Batch Size: 32
+- Weight Decay: 0.0001
+- Dropout: 0.3
+- Optimizer: Adam
+
+## Artifacts
+
+- `resnet18_beans.pth` - PyTorch model weights
+- `per_class_metrics.csv` - Detailed per-class metrics
+- `confusion_matrix.png` - Confusion matrix visualization
+
+## Usage
+
+Download and load the model:
+
+    from huggingface_hub import hf_hub_download
+    import torch
+    from torchvision import models
+    from torch import nn
+    
+    model_path = hf_hub_download(
+        repo_id="vGiacomov/image-classifier-beans",
+        filename="resnet18_beans.pth"
+    )
+    
+    model = models.resnet18()
+    model.fc = nn.Sequential(
+        nn.Dropout(0.3),
+        nn.Linear(model.fc.in_features, 3)
+    )
+    model.load_state_dict(torch.load(model_path, map_location="cpu"))
+    model.eval()
+
+## License
 
-## Artifacts:
-- per_class_metrics.csv
-- confusion_matrix.png
+Apache 2.0

per_class_metrics.csv

@@ -1,7 +1,7 @@
 ,precision,recall,f1-score,support
-0,0.9772727272727273,0.9772727272727273,0.9772727272727273,44.0
-1,0.9782608695652174,1.0,0.989010989010989,45.0
-2,0.9767441860465116,0.9545454545454546,0.9655172413793104,44.0
-accuracy,0.9774436090225563,0.9774436090225563,0.9774436090225563,0.9774436090225563
-macro avg,0.9774259276281522,0.9772727272727272,0.9772669858876756,133.0
-weighted avg,0.9774322053870774,0.9774436090225563,0.9773552866630388,133.0
+Healthy,1.0,0.7954545454545454,0.8860759493670886,44.0
+Bean Rust,0.8035714285714286,1.0,0.8910891089108911,45.0
+Angular Leaf Spot,1.0,0.9545454545454546,0.9767441860465116,44.0
+accuracy,0.9172932330827067,0.9172932330827067,0.9172932330827067,0.9172932330827067
+macro avg,0.9345238095238096,0.9166666666666666,0.9179697481081638,133.0
+weighted avg,0.9335392051557464,0.9172932330827067,0.9177676380390113,133.0

resnet18_beans.pth

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:bc36a065da9c253480a10f50676af3e8e7e37b19a1c7b4ae39459542f290acc0
-size 44792331
+oid sha256:f2ba2b958dcf443e818257191c69f9cc804dc0df26934064e371c0f51fb4b5d2
+size 44792395