import streamlit as st
from PIL import Image
import requests
from transformers import AutoModelForImageClassification, AutoImageProcessor
import torch
from io import BytesIO

# Load the model and processor
repo_name = "vishalkatheriya18/convnextv2-tiny-1k-224-finetuned-bottomwear"
image_processor = AutoImageProcessor.from_pretrained(repo_name)
model = AutoModelForImageClassification.from_pretrained(repo_name)

# Streamlit UI
st.title("Bottomwear Image Classification with Streamlit")

# Option to select image input method
method = st.radio("Select input method", ("Upload Image", "Paste Image URL"))

image = None

if method == "Upload Image":
    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
    if uploaded_file is not None:
        image = Image.open(uploaded_file)
elif method == "Paste Image URL":
    url = st.text_input("Paste image URL here...")
    if url:
        response = requests.get(url)
        if response.status_code == 200:
            image = Image.open(BytesIO(response.content))

if image:
    resized_image = image.resize((500, 300))
    st.image(resized_image, caption="Uploaded Image", use_column_width=True)
    st.write("")
    st.write("Classifying...")

    # Prepare image for the model
    encoding = image_processor(image.convert("RGB"), return_tensors="pt")
    with torch.no_grad():
        logits = model(**encoding).logits

    # Get predicted class
    predicted_class_idx = logits.argmax(-1).item()
    st.write("Predicted class:", model.config.id2label[predicted_class_idx])