viveksardey's picture
Upload folder using huggingface_hub
c90ae07 verified
Raw
History Blame Contribute Delete
2.8 kB
import os
# Force TensorFlow Keras compatibility
os.environ["TF_USE_LEGACY_KERAS"] = "1"
import numpy as np
import tensorflow as tf
from flask import Flask, request, jsonify
from model import build_model
pneumonia_detector_api = Flask(__name__)
IMG_SIZE = 224
MODEL_PATH = "pneumonia_resnet.weights.h5"
print("TensorFlow Version:", tf.__version__)
# -----------------------------
# SAFE MODEL LOADING
# -----------------------------
def load_trained_model():
try:
print("Building model architecture...")
model = build_model()
print("Loading weights...")
model.load_weights(
MODEL_PATH,
skip_mismatch=True
)
print("Model loaded successfully")
return model
except Exception as e:
print("Model loading failed:", e)
raise RuntimeError("Failed to load model")
model = load_trained_model()
# -----------------------------
# HOME ROUTE
# -----------------------------
@pneumonia_detector_api.route("/")
def home():
return "Pneumonia Detection API Running"
# -----------------------------
# IMAGE PREPROCESSING
# -----------------------------
def preprocess_image(img_array):
img = np.array(img_array)
# Ensure image has channel dimension
if len(img.shape) == 2:
img = np.expand_dims(img, axis=-1)
# Resize image
img = tf.image.resize(img, (IMG_SIZE, IMG_SIZE)).numpy()
# Normalize
img = img.astype("float32") / 255.0
# Convert grayscale → RGB
if img.shape[-1] == 1:
img = np.repeat(img, 3, axis=-1)
# Add batch dimension
img = np.expand_dims(img, axis=0)
return img
# -----------------------------
# PREDICTION API
# -----------------------------
@pneumonia_detector_api.route("/v1/predict", methods=["POST"])
def predict():
try:
if "file" not in request.files:
return jsonify({"error": "No file uploaded"}), 400
file = request.files["file"]
# Load .npy file
img_array = np.load(file)
processed_img = preprocess_image(img_array)
print("Input shape to model:", processed_img.shape)
prediction = model.predict(processed_img)
probability = float(prediction[0][0])
label = "Pneumonia" if probability >= 0.5 else "Normal"
return jsonify({
"Predicted_Class": label,
"Probability_Pneumonia": round(probability, 4)
})
except Exception as e:
print("Prediction error:", str(e))
return jsonify({
"error": str(e)
}), 500
# -----------------------------
# RUN SERVER
# -----------------------------
if __name__ == "__main__":
pneumonia_detector_api.run(
host="0.0.0.0",
port=7860,
debug=False
)