quantumbit/mnist-classifier-api

1from flask import Flask, request, jsonify2import numpy as np3import tensorflow as tf4from PIL import Image5import io6import base647import re8import joblib9import os10 11app = Flask(__name__)12 13# Load all models - use absolute paths for Hugging Face14MODEL_DIR = os.path.join(os.getcwd(), "models")15models = {16    "cnn": tf.keras.models.load_model(os.path.join(MODEL_DIR, "mnist_cnn_model.h5")),17    "svm": joblib.load(os.path.join(MODEL_DIR, "mnist_svm.pkl")),18    "logistic": joblib.load(os.path.join(MODEL_DIR, "mnist_logistic_regression.pkl")),19    "random_forest": joblib.load(os.path.join(MODEL_DIR, "mnist_random_forest.pkl"))20}21 22# Preprocess image before prediction23def preprocess_image(image, model_type):24    image = image.resize((28, 28)).convert('L')25    img_array = np.array(image) / 255.026    27    if model_type == "cnn":28        return np.expand_dims(np.expand_dims(img_array, axis=0), axis=-1)29    else:30        return img_array.flatten().reshape(1, -1)31 32def create_simulated_scores(predicted_digit):33    scores = [0.01] * 1034    remaining = 1.0 - sum(scores)35    scores[predicted_digit] += remaining36    return scores37 38@app.route('/')39def home():40    return jsonify({41        "message": "MNIST Classifier API",42        "available_models": list(models.keys()),43        "endpoints": {44            "/predict": "POST - Send image and model_type",45            "/get_classification_report": "POST - Get model metrics"46        }47    })48 49@app.route('/predict', methods=['POST'])50def predict():51    try:52        data = request.json['image']53        model_type = request.json['model_type']54        55        # Process image directly without saving56        img_data = re.sub('^data:image/png;base64,', '', data)57        img = Image.open(io.BytesIO(base64.b64decode(img_data)))58        processed_image = preprocess_image(img, model_type)59        60        if model_type not in models:61            return jsonify({'error': 'Model not found'})62 63        model = models[model_type]64        65        if model_type == "cnn":66            prediction = model.predict(processed_image)67            predicted_digit = np.argmax(prediction)68            confidence_scores = prediction[0].tolist()69            score_type = "probability"70            71        elif model_type == "svm":72            predicted_digit = model.predict(processed_image)[0]73            if hasattr(model, "decision_function"):74                try:75                    decision_scores = model.decision_function(processed_image)76                    if len(decision_scores.shape) == 2:77                        confidence_scores = decision_scores[0].tolist()78                    else:79                        confidence_scores = [0] * 1080                        for i in range(10):81                            confidence_scores[i] = sum(1 for score in decision_scores[0] if score > 0)82                    min_score = min(confidence_scores)83                    if min_score < 0:84                        confidence_scores = [score - min_score for score in confidence_scores]85                    score_type = "decision_distance"86                except Exception:87                    confidence_scores = create_simulated_scores(int(predicted_digit))88                    score_type = "simulated"89            else:90                confidence_scores = create_simulated_scores(int(predicted_digit))91                score_type = "simulated"92        93        else:94            predicted_digit = model.predict(processed_image)[0]95            if hasattr(model, "predict_proba"):96                try:97                    confidence_scores = model.predict_proba(processed_image)[0].tolist()98                    score_type = "probability"99                except Exception:100                    confidence_scores = create_simulated_scores(int(predicted_digit))101                    score_type = "simulated"102            else:103                confidence_scores = create_simulated_scores(int(predicted_digit))104                score_type = "simulated"105 106        return jsonify({107            'digit': int(predicted_digit),108            'confidence_scores': confidence_scores,109            'score_type': score_type110        })111    112    except Exception as e:113        return jsonify({'error': str(e)})114 115if __name__ == '__main__':116    app.run(host='0.0.0.0', port=7860)