ErnestBeckham's picture
updated app
a0c5de0
import streamlit as st
import tensorflow as tf
import cv2
import numpy as np
from huggingface_hub import from_pretrained_keras
model = from_pretrained_keras('ErnestBeckham/ViT-Lungs')
hp = {}
hp['class_names'] = ["lung_aca", "lung_n", "lung_scc"]
def main():
st.title("Lung Cancer Classification")
# Upload image through drag and drop
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
if uploaded_file is not None:
# Convert the uploaded file to OpenCV format
image = convert_to_opencv(uploaded_file)
# Display the uploaded image
st.image(image, channels="BGR", caption="Uploaded Image", use_column_width=True)
# Display the image shape
image_class = predict_single_image(image, model, hp)
st.write(f"Image Class: {image_class}")
def convert_to_opencv(uploaded_file):
# Read the uploaded file using OpenCV
image_bytes = uploaded_file.read()
np_arr = np.frombuffer(image_bytes, np.uint8)
image = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
return image
def process_image_as_batch(image):
#resize the image
image = cv2.resize(image, [512, 512])
#scale the image
image = image / 255.0
#change the data type of image
image = image.astype(np.float32)
return image
def predict_single_image(image, model, hp):
# Preprocess the image
preprocessed_image = process_image_as_batch(image)
# Convert the preprocessed image to a TensorFlow tensor if needed
preprocessed_image = tf.convert_to_tensor(preprocessed_image)
# Add an extra batch dimension (required for model.predict)
preprocessed_image = tf.expand_dims(preprocessed_image, axis=0)
# Make the prediction
predictions = model.predict(preprocessed_image)
np.around(predictions)
y_pred_classes = np.argmax(predictions, axis=1)
class_name = hp['class_names'][y_pred_classes[0]]
return class_name
if __name__ == "__main__":
main()