app.py

import os
import base64
from io import BytesIO
from typing import Optional

import numpy as np
import tensorflow as tf
from flask import Flask, request, jsonify, render_template
from PIL import Image, UnidentifiedImageError

app = Flask(__name__)
model = tf.keras.models.load_model('best_model.keras')

class_names = [
    'Pepper__bell___Bacterial_spot',
    'Pepper__bell___healthy',
    'Potato___Early_blight',
    'Potato___healthy',
    'Potato___Late_blight',
    'Tomato_Bacterial_spot',
    'Tomato_Early_blight',
    'Tomato_Late_blight',
    'Tomato_Leaf_Mold',
    'Tomato_Septoria_leaf_spot',
    'Tomato_Spider_mites_Two_spotted_spider_mite',
    'Tomato__Target_Spot',
    'Tomato__Tomato_YellowLeaf__Curl_Virus',
    'Tomato__Tomato_mosaic_virus',
    'Tomato_healthy'
]

CONFIDENCE_THRESHOLD = 0.60
ALLOWED_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp"}
MAX_FILE_SIZE_BYTES = 10 * 1024 * 1024  # 10MB — generous for a phone photo, blocks abuse

# ---------------------------------------------------------------------------
# Grad-CAM setup — built ONCE at startup, not per-request (same principle as
# loading the model itself at import time: expensive graph construction
# shouldn't happen inside a route function).
# ---------------------------------------------------------------------------

GRADCAM_LAYER_NAME = "top_conv"  # per the training notebook — last conv layer in EfficientNetB0

# Grad-CAM roughly doubles per-request memory (a second forward pass through
# EfficientNetB0 plus gradient tape overhead) — confirmed via Render's Events
# tab to OOM-kill the worker on the 512MB free tier even after merging the
# prediction and Grad-CAM passes into one. Feature-flagged off by default so
# the hosted deploy stays stable; set ENABLE_GRADCAM=true locally to build
# and verify it, e.g. for a demo recording or local client walkthrough.
ENABLE_GRADCAM = os.environ.get("ENABLE_GRADCAM", "false").lower() == "true"


def _find_base_model_layer(outer_model: tf.keras.Model, conv_layer_name: str):
    """Finds the nested sub-model layer (EfficientNetB0) that itself contains
    `conv_layer_name`, and returns (base_model_layer, index_in_outer_model).
    A flat model.get_layer(name) lookup misses this because EfficientNetB0
    is nested as a single opaque layer inside the outer model — its internal
    layers aren't visible at the top level."""
    for index, layer in enumerate(outer_model.layers):
        if hasattr(layer, "layers"):  # this layer is itself a model
            try:
                layer.get_layer(conv_layer_name)
                return layer, index
            except ValueError:
                continue
    return None, None


def _accepts_training_kwarg(layer: tf.keras.layers.Layer) -> bool:
    """Dropout and BatchNorm behave differently in training vs. inference —
    force inference behavior so Grad-CAM isn't affected by random dropout."""
    import inspect
    try:
        return "training" in inspect.signature(layer.call).parameters
    except (TypeError, ValueError):
        return False


_base_model_layer, _base_model_index = _find_base_model_layer(model, GRADCAM_LAYER_NAME) if ENABLE_GRADCAM else (None, None)

last_conv_layer_model = None
classifier_model = None

if not ENABLE_GRADCAM:
    print("INFO: Grad-CAM disabled (ENABLE_GRADCAM is not set to 'true') — running in lean/production mode.")
elif _base_model_layer is not None:
    try:
        conv_layer_output = _base_model_layer.get_layer(GRADCAM_LAYER_NAME).output

        # Model #1: base model's own input -> its own conv output.
        # Valid because both tensors live inside EfficientNetB0's own graph.
        last_conv_layer_model = tf.keras.Model(
            inputs=_base_model_layer.input,
            outputs=conv_layer_output,
        )

        # Model #2: replay the remaining classifier layers (pooling, dropout,
        # dense) on a FRESH input shaped like the conv output. Calling a layer
        # object on a new tensor is always valid, regardless of which graph
        # that layer originally belonged to — this is what sidesteps the
        # nested-submodel connection problem entirely.
        classifier_input = tf.keras.Input(shape=conv_layer_output.shape[1:])
        x = classifier_input
        for classifier_layer in model.layers[_base_model_index + 1:]:
            x = classifier_layer(x, training=False) if _accepts_training_kwarg(classifier_layer) else classifier_layer(x)
        classifier_model = tf.keras.Model(classifier_input, x)
    except (ValueError, AttributeError) as e:
        print(f"WARNING: Grad-CAM setup failed ({e}) — Grad-CAM disabled.")
        last_conv_layer_model = None
        classifier_model = None
else:
    print(f"WARNING: layer '{GRADCAM_LAYER_NAME}' not found in any nested sub-model — Grad-CAM disabled.")

grad_model = (last_conv_layer_model, classifier_model) if last_conv_layer_model is not None else None


def predict_with_optional_gradcam(img_array: tf.Tensor):
    """Runs ONE forward pass (shared by prediction and Grad-CAM) instead of
    two. The original version called model.predict() for the prediction,
    then separately ran last_conv_layer_model + classifier_model again for
    Grad-CAM — a full second pass through EfficientNetB0. On a 512MB
    free-tier box, that duplicate pass plus TensorFlow's own baseline
    footprint was enough to OOM the worker. Reusing the classifier replay's
    own output as the prediction removes that duplicate pass entirely."""
    if grad_model is None:
        predictions = model.predict(img_array, verbose=0)
        return predictions, None

    last_conv_model, clf_model = grad_model
    with tf.GradientTape() as tape:
        conv_outputs = last_conv_model(img_array, training=False)
        tape.watch(conv_outputs)
        predictions = clf_model(conv_outputs, training=False)
        top_class = tf.argmax(predictions[0])
        class_score = predictions[:, top_class]

    grads = tape.gradient(class_score, conv_outputs)
    pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))
    conv_outputs_single = conv_outputs[0]
    heatmap = tf.reduce_sum(conv_outputs_single * pooled_grads, axis=-1)
    heatmap = tf.maximum(heatmap, 0)  # ReLU — only positive influence matters
    max_val = tf.reduce_max(heatmap)
    if max_val > 0:
        heatmap = heatmap / max_val

    return predictions.numpy(), heatmap.numpy()


def _apply_colormap(gray: np.ndarray) -> np.ndarray:
    """
    Hand-rolled 'hot'-style colormap (black -> red -> yellow -> white).
    Deliberately not using matplotlib or opencv here — both add real weight
    to a free-tier deploy for something this simple to compute by hand.
    gray: 2D array, values 0-1. Returns HxWx3 uint8 RGB.
    """
    r = np.clip(gray * 3.0, 0, 1)
    g = np.clip(gray * 3.0 - 1.0, 0, 1)
    b = np.clip(gray * 3.0 - 2.0, 0, 1)
    rgb = np.stack([r, g, b], axis=-1)
    return (rgb * 255).astype(np.uint8)


def make_gradcam_overlay(original_img: Image.Image, heatmap: np.ndarray, alpha: float = 0.4) -> str:
    """Resizes the heatmap up to the original image's size, colorizes it,
    blends it over the original photo, and returns a base64 PNG string
    ready to drop straight into an <img> src on the frontend."""
    heatmap_img = Image.fromarray((heatmap * 255).astype(np.uint8)).resize(
        original_img.size, resample=Image.BICUBIC
    )
    heatmap_rgb = _apply_colormap(np.array(heatmap_img) / 255.0)
    heatmap_rgb_img = Image.fromarray(heatmap_rgb).convert("RGB")

    base = original_img.convert("RGB")
    overlay = Image.blend(base, heatmap_rgb_img, alpha=alpha)

    buffer = BytesIO()
    overlay.save(buffer, format="PNG")
    return base64.b64encode(buffer.getvalue()).decode("utf-8")


# ---------------------------------------------------------------------------
# Routes
# ---------------------------------------------------------------------------

@app.route('/')
def home():
    return render_template('index.html')


@app.route('/predict', methods=['POST'])
def predict():
    if 'image' not in request.files:
        return jsonify({'error': 'No image sent'}), 400

    file = request.files['image']

    if file.filename == "":
        return jsonify({'error': 'No file selected'}), 400

    _, ext = os.path.splitext(file.filename.lower())
    if ext not in ALLOWED_EXTENSIONS:
        return jsonify({'error': f'Unsupported file type: {ext or "unknown"}'}), 400

    # Real byte count from the stream itself, not the (spoofable) Content-Length header
    file.stream.seek(0, os.SEEK_END)
    file_size = file.stream.tell()
    file.stream.seek(0)

    if file_size == 0:
        return jsonify({'error': 'Uploaded file is empty'}), 400
    if file_size > MAX_FILE_SIZE_BYTES:
        return jsonify({'error': f'File exceeds {MAX_FILE_SIZE_BYTES // (1024 * 1024)}MB limit'}), 413

    raw_bytes = file.read()

    try:
        # .verify() checks the file is a genuine, non-corrupt image, but it
        # exhausts the file object — reopen afterward to actually use it.
        Image.open(BytesIO(raw_bytes)).verify()
        original_img = Image.open(BytesIO(raw_bytes))
    except (UnidentifiedImageError, OSError):
        return jsonify({'error': 'File is corrupt or not a valid image'}), 400

    try:
        resized_img = original_img.convert("RGB").resize((224, 224))
        img_array = tf.keras.utils.img_to_array(resized_img)
        img_array = tf.expand_dims(img_array, 0)

        preds, heatmap = predict_with_optional_gradcam(img_array)
        confidence = float(np.max(preds))
        class_index = int(np.argmax(preds))
        class_name = class_names[class_index]
    except (RuntimeError, ValueError) as e:
        return jsonify({'error': f'Model inference failed: {str(e)}'}), 500

    grad_cam_image = None
    if heatmap is not None:
        try:
            grad_cam_image = make_gradcam_overlay(original_img, heatmap)
        except (RuntimeError, ValueError):
            # Explainability is a bonus, not the core feature — if it fails,
            # still return the prediction rather than failing the whole request.
            grad_cam_image = None

    if confidence < CONFIDENCE_THRESHOLD:
        return jsonify({
            'predicted_class': None,
            'confidence': round(confidence, 4),
            'warning': 'Low confidence — retake image in better lighting',
            'grad_cam_image': grad_cam_image
        })

    return jsonify({
        'predicted_class': class_name,
        'confidence': round(confidence, 4),
        'warning': None,
        'grad_cam_image': grad_cam_image
    })


@app.route('/health', methods=['GET'])
def health():
    return jsonify({'status': 'running'})


if __name__ == '__main__':
    app.run(debug=False)  # production — debug=True leaks stack traces to any visitor on error