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)