Oldmajor's picture
initial commit
b847fb1
import streamlit as st
import keras.utils as image
import numpy as np
import io
import tensorflow as tf
import base64
# Load TFLite model and allocate tensors
interpreter = tf.lite.Interpreter(model_path="model/model_float16_quant.tflite")
interpreter.allocate_tensors()
st.title("Image Classifier")
uploaded_file = st.file_uploader("Upload an image...", type=["jpg", "png", "jpeg"])
img_data = None
if uploaded_file is not None:
# Use in-memory image
image_stream = io.BytesIO(uploaded_file.read())
img = image.load_img(image_stream, target_size=(224, 224))
# Convert PIL Image to data URL
img_buffer = io.BytesIO()
img.save(img_buffer, format="PNG")
img_data = base64.b64encode(img_buffer.getvalue()).decode()
# Preprocess the image
img_array = image.img_to_array(img)
x_mean = img_array.mean()
x_std = img_array.std()
img_array = (img_array - x_mean) / x_std
img_array = np.expand_dims(img_array, axis=0)
# Set input tensor for the TFLite interpreter
input_details = interpreter.get_input_details()
interpreter.set_tensor(input_details[0]['index'], img_array)
# Invoke the interpreter
interpreter.invoke()
# Get the prediction result
output_details = interpreter.get_output_details()
predictions = interpreter.get_tensor(output_details[0]['index'])
predicted_class = np.argmax(predictions, axis=1)
class_dict = {
0: 'Actinic keratoses',
1: 'Basal cell carcinoma',
2: 'Benign keratosis-like lesions',
3: 'Dermatofibroma_df',
4: 'Melanocytic nevi',
5: 'Vascular lesions',
6: 'Dermatofibroma_mel'
}
result = class_dict[predicted_class[0]]
st.write(f"Prediction: {result}")
# Display the uploaded image
st.image(img, caption='Uploaded Image', use_column_width=True)