CNN / cnnmodels.py
Gaykar's picture
Upload 3 files
7d3ce75 verified
from keras.layers import TFSMLayer
from PIL import Image
import numpy as np
import io
CLASS_NAMES = ['glioma', 'meningioma', 'no tumor', 'pituitary']
# Load models ONCE
classification_model = TFSMLayer(
"BrainTumorClassificationModel/model",
call_endpoint="serving_default"
)
def preprocess_image_bytes(image_bytes, target_size=(224, 224)):
image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
image = image.resize(target_size)
image = np.array(image).astype("float32") / 255.0
image = np.expand_dims(image, axis=0)
return image
def classify_tumor(image_bytes):
# 1. Preprocess the image
image = preprocess_image_bytes(image_bytes)
# 2. Run Classification directly
# Removed the detection_model block entirely
class_pred_tensor = classification_model(image)["output_0"]
class_pred = class_pred_tensor.numpy()
# 3. Get the index of the highest confidence score
idx = int(np.argmax(class_pred))
confidence = float(np.max(class_pred) * 100)
tumor_type = CLASS_NAMES[idx]
# 4. Determine has_tumor status based on the classification result
# We assume "no tumor" is a string in your CLASS_NAMES list
is_tumor_detected = tumor_type.lower() != "no tumor"
return {
"has_tumor": is_tumor_detected,
"tumor_type": tumor_type,
"confidence": confidence
}