File size: 2,425 Bytes
815a66f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 |
import gradio as gr
import numpy as np
from PIL import Image
import tensorflow as tf
from safetensors import safe_open
# ✅ Constants
IMG_SIZE = 224
CLASS_NAMES = ["Fractured", "Non-Fractured"]
SAFETENSOR_PATH = "osteologic.safetensors"
# ✅ Step 1: Rebuild architecture
def build_model():
inputs = tf.keras.Input(shape=(IMG_SIZE, IMG_SIZE, 3))
base_model = tf.keras.applications.MobileNetV2(weights=None, include_top=False, input_tensor=inputs)
x = base_model.output
x = tf.keras.layers.GlobalAveragePooling2D()(x)
x = tf.keras.layers.Dense(128, activation="relu", kernel_regularizer=tf.keras.regularizers.l2(0.001))(x)
x = tf.keras.layers.Dropout(0.5)(x)
outputs = tf.keras.layers.Dense(len(CLASS_NAMES), activation="softmax")(x)
model = tf.keras.Model(inputs, outputs)
return model
# ✅ Step 2: Load weights from .safetensors
def load_weights(model, path=SAFETENSOR_PATH):
with safe_open(path, framework="pt", device="cpu") as f:
for layer in model.layers:
if isinstance(layer, (tf.keras.layers.Conv2D, tf.keras.layers.Dense)):
w_key = f"{layer.name}.weight"
b_key = f"{layer.name}.bias"
if w_key in f.keys() and b_key in f.keys():
weights = f.get_tensor(w_key)
bias = f.get_tensor(b_key)
# Adjust shape if needed (PyTorch → TF)
if isinstance(layer, tf.keras.layers.Conv2D):
weights = weights.transpose(2, 3, 1, 0) # [out, in, h, w] → [h, w, in, out]
layer.set_weights([weights, bias])
return model
# ✅ Step 3: Build and load model
model = build_model()
model = load_weights(model)
# ✅ Step 4: Prediction function
def predict(image: Image.Image):
image = image.resize((IMG_SIZE, IMG_SIZE)).convert("RGB")
arr = np.array(image) / 255.0
arr = arr.reshape(1, IMG_SIZE, IMG_SIZE, 3)
preds = model.predict(arr)[0]
label = CLASS_NAMES[np.argmax(preds)]
confidence = round(float(np.max(preds)), 3)
return f"{label} ({confidence})"
# ✅ Step 5: Gradio interface
gr.Interface(
fn=predict,
inputs=gr.Image(type="pil", label="Upload Radiograph"),
outputs=gr.Text(label="Prediction"),
title="🦴 OsteoLogic Fracture Detector",
description="Upload a radiograph to detect fractures using safetensors-powered MobileNetV2."
).launch() |