import streamlit as st
import torch
from transformers import AutoProcessor, PaliGemmaForConditionalGeneration
from PIL import Image

model_id = "brucewayne0459/paligemma_derm"
processor = AutoProcessor.from_pretrained(model_id)
model = PaliGemmaForConditionalGeneration.from_pretrained(model_id)
model.eval()

device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)

# Logo (Hugging Face)
st.markdown(
    """
    <style>
    .huggingface-logo {
        display: flex;
        justify-content: center;
        margin-bottom: 20px;
    }
    .huggingface-logo img {
        width: 150px;
    }
    </style>
    <div class="huggingface-logo">
        <img src="https://huggingface.co/front/assets/huggingface_logo-noborder.svg" alt="Hugging Face Logo">
    </div>
    """,
    unsafe_allow_html=True,
)

# App Title
st.title("VisionDerm")
st.write("Upload an image or use your camera to identify the skin condition.")

# Layout
col1, col2 = st.columns([3, 2])

with col1:
    # File uploader
    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
    # Camera input
    camera_photo = st.camera_input("Take a photo")
    prompt = 'Identify the skin condition?'

# Choose input image
input_image = None
if camera_photo:
    input_image = Image.open(camera_photo)
elif uploaded_file:
    try:
        # Open and convert uploaded file to RGB
        input_image = Image.open(uploaded_file).convert("RGB")
        input_image = input_image.copy()  # Detach from file pointer
    except Exception as e:
        st.error(f"Error loading image: {str(e)}")
        input_image = None

with col2:
    if input_image:
        try:
            resized_image = input_image.resize((300, 300))
            st.image(resized_image, caption="Selected Image (300x300)", use_container_width=True)

            # Resize the image for processing (512x512 pixels)
            max_size = (512, 512)
            processed_image = input_image.resize(max_size)

            with st.spinner("Processing..."):
                inputs = processor(
                    text=prompt,
                    images=processed_image,
                    return_tensors="pt",
                    padding="longest"
                ).to(device)

                default_max_tokens = 50
                with torch.no_grad():
                    outputs = model.generate(**inputs, max_new_tokens=default_max_tokens)

                decoded_output = processor.decode(outputs[0], skip_special_tokens=True)
                if prompt in decoded_output:
                    decoded_output = decoded_output.replace(prompt, "").strip()
                decoded_output = decoded_output.title()

                # Display the result
                st.success("Analysis Complete!")
                st.write("**Model Output:**", decoded_output)

        except Exception as e:
            st.error(f"Error: {str(e)}")

st.markdown("---")

# Team Information
st.info("""
### Team: Mahasigma Berprestasi
- **Muhammad Karov Ardava Barus** ; 103052300001
- **Akmal Yaasir Fauzaan** ; 103052300008
- **Farand Diy Dat Mahazalfaa** ; 103052300050
- **Hauzan Rafi Attallah**; 103052330011
""")