webapp.py

import matplotlib
matplotlib.use('Agg')

from flask import Flask, render_template, request
import base64
import io
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image

from sklearn.datasets import fetch_openml
from sklearn.decomposition import PCA
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

# Ensure these paths match your local project structure
from src.multiclass.models import KNNScratch, LinearSVM, OneVsRestClassifier
from src.multiclass.model_files.naive_bayes_pipeline import GaussianNB
from src.multiclass.model_files.decision_tree_pipeline import DecisionTreeScratch
from src.multiclass.model_files.logistic_regression import SoftmaxRegressionScratch
from src.multiclass.model_files.kernel_svm import MulticlassKernelSVM

app = Flask(__name__)

MODELS = {}
DATA = {}
PCA_TRANSFORMER = None
SCALER = None

def load_mnist():
    """Fetches MNIST and splits it into training and testing sets."""
    print("Loading MNIST from OpenML...")
    mnist = fetch_openml('mnist_784', version=1, as_frame=False, parser='auto')
    X = mnist.data.astype(np.float32) / 255.0
    y = mnist.target.astype(int)

    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42
    )
    return X_train, X_test, y_train, y_test

def initialize_app():
    """Initializes models and pre-calculates PCA features for the test set."""
    global MODELS, DATA, PCA_TRANSFORMER, SCALER

    if MODELS:
        return

    X_train_raw, X_test_raw, y_train, y_test = load_mnist()

    # 1. Scaling
    SCALER = StandardScaler()
    X_train_scaled = SCALER.fit_transform(X_train_raw)
    X_test_scaled = SCALER.transform(X_test_raw)

    # 2. PCA Dimensionality Reduction
    PCA_TRANSFORMER = PCA(n_components=100)
    X_train_pca = PCA_TRANSFORMER.fit_transform(X_train_scaled)
    X_test_pca = PCA_TRANSFORMER.transform(X_test_scaled)

    # 3. Store Test Data
    DATA["x_test_features"] = X_test_pca
    DATA["y_test"] = y_test
    DATA["sample_count"] = len(X_test_pca)
    # Store raw pixels (0-1 range) for visual display
    DATA["x_test_images"] = X_test_raw  

    # 4. Training subset (8000 samples for speed)
    X_train_small = X_train_pca[:8000]
    y_train_small = y_train[:8000]

    # 5. Model Instantiation
    MODELS = {
        "KNN": KNNScratch(k=5, weighted=True, batch_size=100),
        "SVM": OneVsRestClassifier(
            estimator_class=LinearSVM,
            learning_rate=0.01,
            lambda_param=0.001,
            n_iters=500
        ),
        "Naive Bayes": GaussianNB(eps=1e-3),
        "Decision Tree": DecisionTreeScratch(max_depth=12, criterion='entropy'),
        "Logistic Regression": SoftmaxRegressionScratch(
            learning_rate=0.05,
            iterations=1000
        ),
        "Kernel SVM": MulticlassKernelSVM(
            C=1.0,
            gamma=None,
            tol=1e-3,
            max_passes=5,
            random_state=42
        )
    }

    for name, model in MODELS.items():
        print(f"Training {name}...")
        if name == "Kernel SVM":
            # Use smaller subset for Kernel SVM
            model.fit(X_train_pca[:1500], y_train[:1500])
        else:
            model.fit(X_train_small, y_train_small)

def generate_base64_plot(img_array):
    """Converts a 28x28 numpy array to a base64 string for HTML display."""
    fig, ax = plt.subplots(figsize=(3, 3))
    ax.imshow(img_array.reshape(28, 28), cmap='gray')
    ax.axis('off')
    
    buf = io.BytesIO()
    # Save with tight bounding box to remove white margins
    fig.savefig(buf, format='png', bbox_inches='tight', pad_inches=0)
    buf.seek(0)
    img_b64 = base64.b64encode(buf.read()).decode('utf-8')
    plt.close(fig)
    return img_b64

def preprocess_uploaded_image(image_file):
    """Processes an external image file into the model's feature space."""
    try:
        img = Image.open(image_file).convert('L').resize((28, 28))
        img_array = np.array(img).astype(np.float32) / 255.0
        
        # Raw image for display
        img_display_b64 = generate_base64_plot(img_array)

        # Transformation for prediction
        feat = img_array.reshape(1, -1)
        feat = SCALER.transform(feat)
        feat = PCA_TRANSFORMER.transform(feat)

        return feat, img_display_b64
    except Exception as e:
        print("Upload Preprocess Error:", e)
        return None, None

@app.route("/", methods=["GET", "POST"])
def index():
    initialize_app()

    model_name = request.form.get("model", "KNN")
    sample_index = request.form.get("sample_index", "")
    uploaded_file = request.files.get("image")

    error = None
    prediction = None
    actual = None
    selected_image = None   # Matches HTML: {% if selected_image %}
    uploaded_image = None   # Matches HTML: {% if uploaded_image %}

    model = MODELS.get(model_name)

    # --- CASE A: USER UPLOADED AN IMAGE ---
    if uploaded_file and uploaded_file.filename:
        features, img_b64 = preprocess_uploaded_image(uploaded_file)
        if features is not None and model:
            prediction = int(model.predict(features)[0])
            uploaded_image = img_b64
        else:
            error = "Could not process the uploaded image."

    # --- CASE B: USER ENTERED AN INDEX ---
    elif sample_index:
        try:
            idx = int(sample_index)
            if 0 <= idx < DATA["sample_count"]:
                features = DATA["x_test_features"][idx:idx+1]
                prediction = int(model.predict(features)[0])
                actual = int(DATA["y_test"][idx])
                
                # Get the raw image data (not PCA) for visual display
                raw_pixels = DATA["x_test_images"][idx]
                selected_image = generate_base64_plot(raw_pixels)
            else:
                error = f"Index must be between 0 and {DATA['sample_count'] - 1}."
        except ValueError:
            error = "Please enter a valid integer for the sample index."

    return render_template(
        "index.html",
        model_name=model_name,
        sample_index=sample_index,
        models=list(MODELS.keys()),
        sample_count=DATA.get("sample_count", 0),
        prediction=prediction,
        actual=actual,
        error=error,
        selected_image=selected_image,
        uploaded_image=uploaded_image
    )

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5001, debug=True)