Fruit Deep Learning Classifier 🍎🍌🥕

A deep learning classifier that tells apart healthy vs rotten produce across 14 fruits and vegetables — 28 classes in total — built with PyTorch using transfer learning on ImageNet-pretrained backbones.

Two backbones are compared side by side in their own subfolders:

• EfficientNetV2/ — EfficientNet-V2-S
• MobileNetV2/ — MobileNet-V2 (lightweight, mobile-friendly)

Classes

Apple · Banana · Bellpepper · Carrot · Cucumber · Grape · Guava · Jujube · Mango · Orange · Pomegranate · Potato · Strawberry · Tomato — each in Healthy and Rotten variants.

Approach

1. Load an ImageNet-pretrained backbone from torchvision.models.
2. Freeze the feature extractor's weights.
3. Replace the final classifier head with a fresh nn.Linear(..., 28).
4. Train only the new head with Adam (lr = 0.001) and CrossEntropyLoss.
5. Save the trained weights to best_model.pth.

Standard ImageNet preprocessing: resize → center-crop to 224×224 → normalize with mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225].

Repository layout

FruitDeepLearning-Classifier/
├── train_and_evaluate.ipynb   # end-to-end notebook: Kaggle download → train → eval → predict
├── EfficientNetV2/
│   ├── main.py                # transfer-learning training loop
│   └── predict.py             # single-image inference with best_model.pth
└── MobileNetV2/
    ├── main.py                # transfer-learning training loop
    ├── predict.py             # single-image inference
    └── reduce_dataset.py      # subsamples each class folder to N images

Prerequisites

You'll need Kaggle API credentials so the notebook can download the dataset:

1. Sign in at kaggle.com → Account → Create New API Token → downloads kaggle.json.
2. Drop it into your OS's kaggle config folder:
   • Linux/macOS: ~/.kaggle/kaggle.json
   • Windows: %USERPROFILE%\.kaggle\kaggle.json
3. Restrict the file's permissions (Linux/macOS only): chmod 600 ~/.kaggle/kaggle.json.

If you're on Google Colab, upload kaggle.json when the notebook prompts you — no manual placement needed.

A CUDA-capable GPU is recommended but not required — the scripts fall back to CPU automatically. Training on CPU is significantly slower (Colab's free-tier GPU is more than enough).

Quickest path — run the notebook

Open train_and_evaluate.ipynb in Jupyter, Colab, or VS Code. It:

1. Installs the required packages.
2. Downloads the Kaggle dataset via kagglehub (using the credentials from the Prerequisites section above).
3. Builds train/val splits with augmentation on the training half only.
4. Loads a pretrained backbone (MobileNet-V2 or EfficientNet-V2-S — toggle via the BACKBONE variable), freezes it, and trains a fresh classifier head.
5. Plots loss/accuracy curves and runs a top-3 prediction on a sample image.

Note on committed outputs: the notebook is committed without cell outputs. The published version doesn't ship a trained best_model.pth — you generate one by running the notebook end-to-end on your own machine after the setup above.

Dataset

Expected layout — one subfolder per class (compatible with torchvision.datasets.ImageFolder):

data/
├── Apple__Healthy/
├── Apple__Rotten/
├── Banana__Healthy/
├── Banana__Rotten/
└── ...

The Fruit and Vegetable Disease (Healthy vs Rotten) Kaggle dataset matches this structure directly.

reduce_dataset.py (in MobileNetV2/) subsamples the full dataset down to a fixed number of images per class (default 500) so training stays quick on modest hardware.

Getting started

pip install torch torchvision pillow

Train:

cd EfficientNetV2         # or MobileNetV2
python main.py            # trains and saves best_model.pth

Predict on a single image:

python predict.py         # loads best_model.pth and classifies the test image

Point the image_path in predict.py at your own image before running.

The training scripts assume a CUDA-capable GPU when available — they fall back to CPU automatically, though training on CPU will be slow.

🔗 More

• Author: @pop123-ux
• Medium write-ups: medium.com/@Pop123