gerd-lightweight-vit / inference_example.py
ezAhmed's picture
Upload folder using huggingface_hub
5852205 verified
import tensorflow as tf
from PIL import Image
import numpy as np
import requests
from io import BytesIO
# Load model
model = tf.keras.models.load_model('model.keras')
# Class names
CLASS_NAMES = ['Esophagitis', 'GERD', 'Normal', 'Ulcer']
def predict_from_url(image_url):
response = requests.get(image_url)
image = Image.open(BytesIO(response.content))
return predict_image(image)
def predict_from_path(image_path):
image = Image.open(image_path)
return predict_image(image)
def predict_image(image):
# Preprocess
image = image.convert('RGB')
image = image.resize((224, 224))
image_array = np.array(image) / 255.0
image_array = np.expand_dims(image_array, axis=0)
# Predict
predictions = model.predict(image_array)[0]
# Format results
results = {
CLASS_NAMES[i]: float(predictions[i])
for i in range(len(CLASS_NAMES))
}
predicted_class = CLASS_NAMES[np.argmax(predictions)]
confidence = float(np.max(predictions))
return {
'predicted_class': predicted_class,
'confidence': confidence,
'all_predictions': results
}
# Example usage
if __name__ == "__main__":
result = predict_from_path('test_image.jpg')
print(f"Prediction: {result['predicted_class']}")
print(f"Confidence: {result['confidence']:.2%}")
print(f"All probabilities: {result['all_predictions']}")