import streamlit as st
from PIL import Image
from transformers import pipeline

# Load the pre-trained model
classifier = pipeline("image-classification", model="https://teachablemachine.withgoogle.com/models/lcNO3nb0s/")

st.title("Korean Jelly Identifier")

uploaded_file = st.file_uploader("Choose an image...", type="jpg")

if uploaded_file is not None:
    image = Image.open(uploaded_file)
    st.image(image, caption='Uploaded Image.', use_column_width=True)
    st.write("")
    st.write("Classifying...")

    # Classify the image
    results = classifier(image)

    jelly_type = results[0]['label']
    sugar_level = get_sugar_level(jelly_type)
    hazard = get_hazard_level(sugar_level)

    st.write(f'Jelly Type: {jelly_type}')
    st.write(f'Sugar Level: {sugar_level}')
    st.write(f'Hazard: {hazard}')

def get_sugar_level(jelly_type):
    # Dummy data for demonstration purposes
    sugar_data = {
        'jellyA': 10,
        'jellyB': 20,
        'jellyC': 30
    }
    return sugar_data.get(jelly_type, 0)

def get_hazard_level(sugar_level):
    if sugar_level > 25:
        return 'Red (High Hazard)'
    elif sugar_level > 15:
        return 'Yellow (Moderate Hazard)'
    else:
        return 'Green (Low Hazard)'