|
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") |
|
|
|
|
|
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"]) |
|
|
|
if uploaded_file is not None: |
|
|
|
image = convert_to_opencv(uploaded_file) |
|
|
|
|
|
st.image(image, channels="BGR", caption="Uploaded Image", use_column_width=True) |
|
|
|
|
|
image_class = predict_single_image(image, model, hp) |
|
st.write(f"Image Class: {image_class}") |
|
|
|
def convert_to_opencv(uploaded_file): |
|
|
|
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): |
|
|
|
image = cv2.resize(image, [512, 512]) |
|
|
|
image = image / 255.0 |
|
|
|
image = image.astype(np.float32) |
|
return image |
|
|
|
def predict_single_image(image, model, hp): |
|
|
|
preprocessed_image = process_image_as_batch(image) |
|
|
|
preprocessed_image = tf.convert_to_tensor(preprocessed_image) |
|
|
|
preprocessed_image = tf.expand_dims(preprocessed_image, axis=0) |
|
|
|
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() |