vit_fruit_ripeness_classifier

ViT Fruit Ripeness Classifier

A fruit ripeness classification model combining Vision Transformer (ViT) feature extraction with Logistic Regression for fast, accurate inference on CPU or GPU.

Model Description

This model classifies the ripeness condition of apples, bananas, and oranges into three categories:

• Fresh - Ready to eat
• Unripe - Needs more time to ripen
• Rotten - Past optimal consumption

Architecture

• Feature Extractor: ViT Base Patch-16 (google/vit-base-patch16-224)
• Classifier: Scikit-learn Logistic Regression
• Feature Dimension: 768-dim pooled output from ViT
• Total Classes: 9 (3 fruits × 3 ripeness states)

Supported Classes

Quick Start

Installation

pip install torch torchvision transformers scikit-learn pillow joblib numpy huggingface_hub

First Run Notes

• Automatic Downloads: The model files (~350-400MB) download automatically on first run
• No Manual Downloads: You don't need to manually download any model files
• Internet Required: Only for the first run; subsequent runs work offline using cached files
• Time: First run takes 2-5 minutes for downloads; later runs are instant

Single Image Inference

import json
import joblib
from pathlib import Path
from PIL import Image
import torch
import numpy as np
from huggingface_hub import hf_hub_download, HfApi
from transformers import AutoImageProcessor, ViTModel
import warnings

# ----------------- CONFIG -----------------
REPO_ID = "Meeteshn/vit_fruit_ripeness_classifier"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
NESTED_FOLDER = "vit_fruit_ripeness_updated"  # your repo uses this nested folder
TOP_K = 5
# ------------------------------------------

def hf_download_try(repo_id: str, filename: str, nested_folder: str = NESTED_FOLDER):
    """
    Try to download `filename` from repo root, then from nested_folder/filename.
    Returns local path to downloaded file or raises an informative error.
    """
    candidates = [filename, f"{nested_folder}/{filename}"]
    last_exc = None
    for f in candidates:
        try:
            print(f"Trying to download '{f}' from '{repo_id}'...")
            path = hf_hub_download(repo_id=repo_id, filename=f)
            print("Downloaded:", path)
            return path
        except Exception as e:
            print(f"Not found at '{f}': {e}")
            last_exc = e
    raise RuntimeError(f"Could not download '{filename}' from repo '{repo_id}'. Last error: {last_exc}")

def load_processor_and_backbone(repo_id: str, nested_folder: str = NESTED_FOLDER, device: str = DEVICE):
    """
    Try several likely subfolder locations for processor/backbone.
    Returns (processor, backbone).
    """
    # candidate subfolders for processor
    proc_candidates = [
        "processor",
        f"{nested_folder}/processor",
        "",  # no subfolder (root)
    ]
    last_exc = None
    for sub in proc_candidates:
        try:
            if sub == "":
                print(f"Trying AutoImageProcessor.from_pretrained('{repo_id}')")
                processor = AutoImageProcessor.from_pretrained(repo_id, use_fast=True)
            else:
                print(f"Trying AutoImageProcessor.from_pretrained('{repo_id}', subfolder='{sub}')")
                processor = AutoImageProcessor.from_pretrained(repo_id, subfolder=sub, use_fast=True)
            # now try backbone with matching guessed subfolder
            backbone_sub = sub.replace("processor", "vit_backbone") if sub and "processor" in sub else ("vit_backbone" if sub == "" else f"{nested_folder}/vit_backbone")
            try:
                print(f"Trying ViTModel.from_pretrained('{repo_id}', subfolder='{backbone_sub}')")
                backbone = ViTModel.from_pretrained(repo_id, subfolder=backbone_sub)
            except Exception as e_backbone:
                # final fallback: try root vit_backbone
                print(f"Backbone attempt failed for sub='{backbone_sub}': {e_backbone}. Trying root 'vit_backbone'.")
                backbone = ViTModel.from_pretrained(repo_id, subfolder="vit_backbone")
            backbone.to(device)
            backbone.eval()
            print(f"Loaded processor/backbone from subfolder='{sub or 'root'}'")
            return processor, backbone
        except Exception as e:
            print(f"Processor load failed for sub='{sub}': {e}")
            last_exc = e
    # ultimate fallback: official ViT from hub
    warnings.warn("Could not load processor/backbone from repo; falling back to official 'google/vit-base-patch16-224'.")
    processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224", use_fast=True)
    backbone = ViTModel.from_pretrained("google/vit-base-patch16-224")
    backbone.to(device)
    backbone.eval()
    return processor, backbone

# ----------------- Load assets (robust) -----------------
processor, backbone = load_processor_and_backbone(REPO_ID, nested_folder=NESTED_FOLDER, device=DEVICE)

# Download sklearn artifacts (try root then nested)
scaler_path = hf_download_try(REPO_ID, "scaler.joblib", nested_folder=NESTED_FOLDER)
clf_path    = hf_download_try(REPO_ID, "logistic_model.joblib", nested_folder=NESTED_FOLDER)
metadata_path = hf_download_try(REPO_ID, "metadata.json", nested_folder=NESTED_FOLDER)

scaler = joblib.load(scaler_path)
clf = joblib.load(clf_path)
metadata = json.loads(Path(metadata_path).read_text(encoding="utf-8"))
classes = metadata["classes"]

# ----------------- Prediction function -----------------
def predict(image_path: str):
    """Predict ripeness condition for a single image."""
    img = Image.open(image_path).convert("RGB")
    inputs = processor(images=img, return_tensors="pt")
    pixel_values = inputs["pixel_values"].to(DEVICE)

    with torch.no_grad():
        out = backbone(pixel_values=pixel_values, return_dict=True)
        pooled = getattr(out, "pooler_output", None)
        if pooled is None:
            pooled = out.last_hidden_state[:, 0, :]
        feat = pooled.cpu().numpy()

    feat_scaled = scaler.transform(feat)
    # get probabilities (works for sklearn logistic / classifiers with predict_proba)
    if hasattr(clf, "predict_proba"):
        probs = clf.predict_proba(feat_scaled)[0]
    else:
        # fallback for classifiers without predict_proba
        dec = clf.decision_function(feat_scaled)[0]
        exp = np.exp(dec - np.max(dec))
        probs = exp / exp.sum()

    idx = int(np.argmax(probs))
    return classes[idx], float(probs[idx]), {classes[i]: float(probs[i]) for i in range(len(classes))}

# ----------------- Example usage -----------------
if __name__ == "__main__":
    sample_image = "my_apple.jpg"  # change as needed
    label, prob, all_probs = predict(sample_image)
    print(f"Prediction: {label} ({prob*100:.2f}%)")
    print("\nTop probabilities:")
    for cls, p in sorted(all_probs.items(), key=lambda x: -x[1])[:TOP_K]:
        print(f"  {cls}: {p*100:.2f}%")

Batch Prediction

from pathlib import Path
import csv

def batch_predict(folder_path: str, output_csv: str = "predictions.csv"):
    """Predict ripeness for all images in a folder."""
    folder = Path(folder_path)
    
    with open(output_csv, "w", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        writer.writerow(["filename", "predicted_label", "probability"])
        
        for img_path in sorted(folder.rglob("*")):
            if img_path.suffix.lower() not in [".jpg", ".jpeg", ".png", ".bmp"]:
                continue
            
            label, prob, _ = predict(str(img_path))
            writer.writerow([img_path.name, label, f"{prob*100:.2f}%"])
    
    print(f"Predictions saved to {output_csv}")

# Usage
batch_predict("path/to/images")

Example Output

Prediction: rottenapples (71.24%)

Top 5 probabilities:
  rottenapples: 71.24%
  rottenbanana: 12.35%
  freshapples: 6.12%
  unripe apple: 4.89%
  freshoranges: 2.31%

Repository Structure

vit_fruit_ripeness_updated/
├── processor/              # AutoImageProcessor configuration
├── vit_backbone/          # ViT feature extractor weights
├── logistic_model.joblib  # Trained classifier
├── scaler.joblib          # Feature scaler
├── metadata.json          # Class labels and metadata
└── features_extracted.npz # (Optional) Cached features

Performance Notes

• Best Results: Fruit is centered and clearly visible in the image
• Works Well: Smartphone photos, typical market/kitchen images
• Challenges: Heavy background clutter, extreme lighting conditions, unusual fruit varieties
• Uncertainty Handling: Use top-K probabilities to assess prediction confidence

Use Cases

• Quality control in fruit sorting facilities
• Smart grocery shopping apps
• Food waste reduction systems
• Educational tools for agriculture
• Retail inventory management

Technical Details

• Input Size: 224×224 pixels (automatically resized)
• Inference Speed: ~50-100ms per image (GPU), ~200-500ms (CPU)
• Memory Usage: ~500MB (model weights)
• Training: No training required for inference

License

MIT License - See LICENSE file for details

Author

Meetesh Nagrecha

Acknowledgments

• Base model: google/vit-base-patch16-224
• Framework: Hugging Face Transformers & Scikit-learn

---

Citation

If you use this model in your research, please cite:

@misc{vit-fruit-ripeness-classifier,
  author = {Nagrecha, Meetesh},
  title = {ViT Fruit Ripeness Classifier},
  year = {2024},
  publisher = {Hugging Face},
  howpublished = {\url{https://huggingface.co/Meeteshn/vit_fruit_ripeness_classifier}}
}