Georgek17's picture
Upload folder using huggingface_hub
3efa3dc verified
import numpy as np
import pandas as pd
import cv2
import pydicom
from flask import Flask, request, jsonify
from tensorflow.keras.models import load_model
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1" # Force CPU mode for TensorFlow
# Initialize Flask app
Pneumonia_predictor_api = Flask("Pneumonia predictor")
# Load trained model (must match TF/Keras version used during training)
model = load_model("final_model.keras")
@Pneumonia_predictor_api.get('/')
def home():
return "Welcome to the Pneumonia Prediction API!"
@Pneumonia_predictor_api.route('/v1/Pneumonia_prediction', methods=['POST'])
def predict_pneumonia():
# --------- STEP 1: CHECK FILE ---------
if 'file' not in request.files:
return jsonify({'error': 'No DICOM file uploaded'}), 400
dcm_file = request.files['file']
dcm_file.seek(0) # ensure pointer at beginning
# --------- STEP 2: READ DICOM SAFELY ---------
try:
dicom_data = pydicom.dcmread(dcm_file, force=True) # force=True fixes header issues
image = dicom_data.pixel_array.astype(np.float32)
except Exception as e:
return jsonify({"error": f"Invalid DICOM file: {str(e)}"}), 400
# --------- STEP 3: PREPROCESS IMAGE ---------
image = cv2.resize(image, (224, 224))
# Convert grayscale → RGB 3-channel
if len(image.shape) == 2:
image = np.stack([image] * 3, axis=-1)
# Normalize 0–1
max_val = np.max(image)
if max_val > 0:
image = image / max_val
image = np.expand_dims(image, axis=0) # (1,224,224,3)
# --------- STEP 4: PREDICT ---------
try:
prob = float(model.predict(image)[0][0]) # probability
pred_class = int(prob >= 0.5) # predicted class
except Exception as e:
return jsonify({"error": f"Model prediction failed: {str(e)}"}), 500
return jsonify({
"Predicted_Class": pred_class,
"Probability": prob,
"Message": "Prediction completed"
})
# Run Flask locally
if __name__ == '__main__':
Pneumonia_predictor_api.run(debug=True)