import streamlit as st
import tensorflow as tf
from tensorflow.keras.applications.resnet50 import preprocess_input
from tensorflow.keras.preprocessing import image
import numpy as np
from PIL import Image
from huggingface_hub import hf_hub_download

# Download the model from Hugging Face
model_path = hf_hub_download(repo_id="AliAmr0/Kidney-Classification-Using-Resnet50", filename="resnet50_kidney_ct_augmented.h5")
model = tf.keras.models.load_model(model_path)

# Class labels (change based on your model's labels)
labels = ["Cyst", "Normal", "Stone", "Tumor"]

def predict(img):
    # Resize and preprocess image to fit ResNet50 input format
    img = img.resize((224, 224))  # ResNet50 expects 224x224 images
    img_array = image.img_to_array(img)
    img_array = np.expand_dims(img_array, axis=0)
    img_array = preprocess_input(img_array)

    # Model prediction
    prediction = model.predict(img_array)
    predicted_class = np.argmax(prediction, axis=1)
    
    return labels[predicted_class[0]]

# Streamlit interface
st.title("TensorFlow Image Classification with ResNet50")
st.write("Upload an image to classify")

uploaded_image = st.file_uploader("Choose an image...", type=["jpg", "png", "jpeg"])

if uploaded_image is not None:
    img = Image.open(uploaded_image)
    st.image(img, caption="Uploaded Image", use_column_width=True)
    
    # Make prediction
    prediction = predict(img)
    st.write(f"Prediction: {prediction}")