app.py

import os
import io
import time
import logging
import threading
from flask import Flask, request, jsonify, send_file
from flask_cors import CORS
from PIL import Image
import torch
from transformers import ViTImageProcessor, ViTForImageClassification

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s - %(message)s"
)
logger = logging.getLogger("PlantClassifierAPI")

# ---------------------------------------------------------------------------
# Hugging Face repository configuration
# ---------------------------------------------------------------------------
HF_MODEL_REPO  = os.environ.get("HF_MODEL_REPO", "")
SPACE_ID       = os.environ.get("SPACE_ID", "Karim31003/classifier")
HF_CACHE_DIR   = os.environ.get("HF_CACHE_DIR", "/tmp/hf")
HF_TOKEN       = os.environ.get("HF_TOKEN", None)

# Initialize Flask App
app = Flask(__name__)
CORS(app)

# Limit request payload to 10MB
app.config['MAX_CONTENT_LENGTH'] = 10 * 1024 * 1024

DEBUG_SAVE_IMAGES = os.environ.get("DEBUG_SAVE_IMAGES", "false").lower() == "true"
UPLOAD_FOLDER = os.path.join(os.path.dirname(os.path.abspath(__file__)), "uploads")
ALLOWED_EXTENSIONS = {'.jpg', '.jpeg', '.png'}
ALLOWED_MIME_TYPES = {'image/jpeg', 'image/png', 'application/octet-stream'}

# Global cache for loaded model
models_cache = {
    "processor": None,
    "model": None
}

# Detailed load reports for debugging
model_load_reports = []

# Force CPU device
device = torch.device("cpu")
logger.info(f"Using hardware device: {device}")
models_lock = threading.RLock()
models_ready = threading.Event()
models_load_error = None

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
# Note: Since the model is currently local, we look in the Leaf-Disease-Predictor folder
LOCAL_MODEL_DIR = os.path.join(BASE_DIR, "Leaf-Disease-Predictor")

# ---------------------------------------------------------------------------
# Startup model pre-loading
# ---------------------------------------------------------------------------

def preload_model():
    """Load and cache the ViT model at server startup."""
    global models_load_error, model_load_reports

    with models_lock:
        start_time = time.time()
        logger.info("=== Starting ViT model pre-load ===")
        model_load_reports = []

        try:
            # Prefer local directory first, otherwise fallback to Hugging Face Hub if configured
            model_path = LOCAL_MODEL_DIR if os.path.isdir(LOCAL_MODEL_DIR) else (HF_MODEL_REPO or SPACE_ID)
            
            logger.info(f"Loading processor and model from {model_path} ...")
            processor = ViTImageProcessor.from_pretrained(model_path)
            model = ViTForImageClassification.from_pretrained(model_path)
            
            model.to(device)
            model.eval()

            models_cache["processor"] = processor
            models_cache["model"] = model

            model_load_reports.append({
                "model": "ViT_Leaf_Disease_Predictor",
                "status": "loaded"
            })
            logger.info("ViT model loaded successfully.")

        except Exception as exc:
            models_load_error = str(exc)
            logger.critical(f"ViT model loading failed: {exc}", exc_info=True)
            model_load_reports.append({
                "model": "ViT_Leaf_Disease_Predictor",
                "status": "error",
                "error": str(exc)
            })

        elapsed = time.time() - start_time
        logger.info(f"=== Model pre-load done in {elapsed:.1f}s ===")

        if models_load_error is None:
            models_ready.set()


def ensure_models_loaded():
    """Trigger model loading on demand if not already done."""
    global models_load_error

    if models_cache["model"] is not None:
        models_ready.set()
        return True

    with models_lock:
        if models_cache["model"] is not None:
            models_ready.set()
            return True
        try:
            preload_model()
        except Exception as exc:
            models_load_error = str(exc)
            logger.error(f"Model loading failed: {exc}", exc_info=True)
            return False

    return models_cache["model"] is not None and models_load_error is None


# ---------------------------------------------------------------------------
# Flask helpers
# ---------------------------------------------------------------------------

def allowed_file(filename, mime_type):
    ext = os.path.splitext(filename)[1].lower()
    valid = ext in ALLOWED_EXTENSIONS and mime_type in ALLOWED_MIME_TYPES
    if not valid:
        logger.warning(f"File rejected — name: '{filename}', ext: '{ext}', mime: '{mime_type}'")
    return valid


# ---------------------------------------------------------------------------
# Flask routes
# ---------------------------------------------------------------------------

@app.errorhandler(413)
def request_entity_too_large(error):
    return jsonify({"success": False, "message": "File size exceeds the 10 MB limit"}), 413


@app.route("/", methods=["GET"])
def index():
    return send_file("main.html")


@app.route("/health", methods=["GET"])
def health_check():
    return jsonify({
        "status": "healthy" if models_cache["model"] is not None and models_load_error is None else "starting",
        "device": str(device),
        "vit_loaded": models_cache["model"] is not None,
        "model_load_reports": model_load_reports,
        "debug_save_images": DEBUG_SAVE_IMAGES,
        "models_loading": models_cache["model"] is None,
        "load_error": models_load_error,
        "hf_model_repo": HF_MODEL_REPO or "(not configured)",
        "space_id": SPACE_ID or "(not detected)",
    }), 200


@app.route("/predict", methods=["POST"])
def predict():
    """Inference: Single ViT model for plant and disease classification."""
    if not ensure_models_loaded():
        return jsonify({
            "success": False,
            "message": models_load_error or "Model is still loading. Please retry shortly."
        }), 503

    if 'image' not in request.files:
        return jsonify({"success": False, "message": "No file in form-data key 'image'"}), 400

    file = request.files['image']
    if file.filename == '':
        return jsonify({"success": False, "message": "No file selected for upload"}), 400

    mime_type = file.content_type
    logger.info(f"Received file: '{file.filename}', MIME: '{mime_type}'")
    if not allowed_file(file.filename, mime_type):
        return jsonify({
            "success": False,
            "message": "Invalid file format. Only JPG, JPEG, and PNG are allowed."
        }), 400

    try:
        img_bytes = file.read()

        if DEBUG_SAVE_IMAGES:
            os.makedirs(UPLOAD_FOLDER, exist_ok=True)
            debug_filename = f"img_{int(time.time() * 1000)}{os.path.splitext(file.filename)[1]}"
            debug_path = os.path.join(UPLOAD_FOLDER, debug_filename)
            with open(debug_path, "wb") as fh:
                fh.write(img_bytes)
            logger.info(f"Debug: saved image to {debug_path}")

        img = Image.open(io.BytesIO(img_bytes)).convert("RGB")

        processor = models_cache["processor"]
        model = models_cache["model"]
        
        if model is None or processor is None:
            return jsonify({"success": False, "message": "ViT model is not loaded"}), 500

        # Process the image and predict
        inputs = processor(images=img, return_tensors="pt").to(device)
        with torch.no_grad():
            outputs = model(**inputs)
            logits = outputs.logits
            probs = torch.nn.functional.softmax(logits, dim=-1)
            predicted_class_idx = logits.argmax(-1).item()
            confidence = probs[0, predicted_class_idx].item()
            # id2label keys may be int or str depending on how config.json was saved
            id2label = model.config.id2label
            predicted_label = (
                id2label.get(predicted_class_idx)
                or id2label.get(str(predicted_class_idx))
                or f"class_{predicted_class_idx}"
            )

        # Extract plant type and disease from label (e.g., "Tomato_Bacterial_spot")
        if "_" in predicted_label:
            parts = predicted_label.split("_", 1)
            plant_type = parts[0]
            disease = parts[1].replace("_", " ")
        else:
            plant_type = "Unknown"
            disease = predicted_label

        # Format detection result matching previous frontend expectations
        # Without bounding box since ViT classifies the whole image
        detection = {
            "plant_type": plant_type,
            "detection_confidence": round(confidence, 4),
            "box": [0, 0, img.width, img.height], # Fallback full-image box
            "disease": disease,
            "disease_confidence": round(confidence, 4),
            "raw_label": predicted_label
        }

        return jsonify({"success": True, "detections": [detection]}), 200

    except Exception as exc:
        logger.error(f"Inference pipeline failed: {exc}", exc_info=True)
        return jsonify({
            "success": False,
            "message": f"Server error during inference: {exc}"
        }), 500


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    # Kick off model loading in a background thread so the HTTP server starts immediately.
    threading.Thread(target=preload_model, daemon=True).start()
    port = int(os.environ.get("PORT", 7860))
    app.run(host="0.0.0.0", port=port, debug=False)