Spaces:
Sleeping
Sleeping
| from flask import Flask, request, jsonify | |
| import numpy as np | |
| import tensorflow as tf | |
| from PIL import Image | |
| import io | |
| import base64 | |
| import re | |
| import joblib | |
| import os | |
| app = Flask(__name__) | |
| # Load all models - use absolute paths for Hugging Face | |
| MODEL_DIR = os.path.join(os.getcwd(), "models") | |
| models = { | |
| "cnn": tf.keras.models.load_model(os.path.join(MODEL_DIR, "mnist_cnn_model.h5")), | |
| "svm": joblib.load(os.path.join(MODEL_DIR, "mnist_svm.pkl")), | |
| "logistic": joblib.load(os.path.join(MODEL_DIR, "mnist_logistic_regression.pkl")), | |
| "random_forest": joblib.load(os.path.join(MODEL_DIR, "mnist_random_forest.pkl")) | |
| } | |
| # Preprocess image before prediction | |
| def preprocess_image(image, model_type): | |
| image = image.resize((28, 28)).convert('L') | |
| img_array = np.array(image) / 255.0 | |
| if model_type == "cnn": | |
| return np.expand_dims(np.expand_dims(img_array, axis=0), axis=-1) | |
| else: | |
| return img_array.flatten().reshape(1, -1) | |
| def create_simulated_scores(predicted_digit): | |
| scores = [0.01] * 10 | |
| remaining = 1.0 - sum(scores) | |
| scores[predicted_digit] += remaining | |
| return scores | |
| def home(): | |
| return jsonify({ | |
| "message": "MNIST Classifier API", | |
| "available_models": list(models.keys()), | |
| "endpoints": { | |
| "/predict": "POST - Send image and model_type", | |
| "/get_classification_report": "POST - Get model metrics" | |
| } | |
| }) | |
| def predict(): | |
| try: | |
| data = request.json['image'] | |
| model_type = request.json['model_type'] | |
| # Process image directly without saving | |
| img_data = re.sub('^data:image/png;base64,', '', data) | |
| img = Image.open(io.BytesIO(base64.b64decode(img_data))) | |
| processed_image = preprocess_image(img, model_type) | |
| if model_type not in models: | |
| return jsonify({'error': 'Model not found'}) | |
| model = models[model_type] | |
| if model_type == "cnn": | |
| prediction = model.predict(processed_image) | |
| predicted_digit = np.argmax(prediction) | |
| confidence_scores = prediction[0].tolist() | |
| score_type = "probability" | |
| elif model_type == "svm": | |
| predicted_digit = model.predict(processed_image)[0] | |
| if hasattr(model, "decision_function"): | |
| try: | |
| decision_scores = model.decision_function(processed_image) | |
| if len(decision_scores.shape) == 2: | |
| confidence_scores = decision_scores[0].tolist() | |
| else: | |
| confidence_scores = [0] * 10 | |
| for i in range(10): | |
| confidence_scores[i] = sum(1 for score in decision_scores[0] if score > 0) | |
| min_score = min(confidence_scores) | |
| if min_score < 0: | |
| confidence_scores = [score - min_score for score in confidence_scores] | |
| score_type = "decision_distance" | |
| except Exception: | |
| confidence_scores = create_simulated_scores(int(predicted_digit)) | |
| score_type = "simulated" | |
| else: | |
| confidence_scores = create_simulated_scores(int(predicted_digit)) | |
| score_type = "simulated" | |
| else: | |
| predicted_digit = model.predict(processed_image)[0] | |
| if hasattr(model, "predict_proba"): | |
| try: | |
| confidence_scores = model.predict_proba(processed_image)[0].tolist() | |
| score_type = "probability" | |
| except Exception: | |
| confidence_scores = create_simulated_scores(int(predicted_digit)) | |
| score_type = "simulated" | |
| else: | |
| confidence_scores = create_simulated_scores(int(predicted_digit)) | |
| score_type = "simulated" | |
| return jsonify({ | |
| 'digit': int(predicted_digit), | |
| 'confidence_scores': confidence_scores, | |
| 'score_type': score_type | |
| }) | |
| except Exception as e: | |
| return jsonify({'error': str(e)}) | |
| if __name__ == '__main__': | |
| app.run(host='0.0.0.0', port=7860) |