import streamlit as st
import torch
import torchvision.transforms as transforms
from torchvision.transforms import ToTensor, Normalize, Compose
from PIL import Image
import numpy as np
from galaxy_classification import GalaxyCNN  # Modelini import et

# Streamlit sayfa yapılandırması
st.set_page_config(
    page_title="Galaxy Morphology Classification",
    page_icon="✨",
    layout="wide",
    initial_sidebar_state="expanded",
)

# Modeli yükle
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = GalaxyCNN().to(device)
model.load_state_dict(torch.load("galaxy_cnn_model.pt", map_location=device))
model.eval()

# Görsel işleme fonksiyonu
def process_image(image):
    transform = Compose([ToTensor(), Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])  # 3 kanal için normalize
    image = image.convert("RGB").resize((64, 64))  # Modelin beklediği boyutta olmalı
    image = transform(image).unsqueeze(0).to(device)  # Model için uygun forma getir
    return image

# Başlık ve alt başlık
st.title("🔭 Galaxy Morphology Classifier")
st.subheader("Upload a galaxy image and predict its morphological class!")

# Açıklama
st.write("""
### How It Works
1. Upload a galaxy image.
2. Click the "Classify Image" button.
3. The model predicts the galaxy's morphological class.
""")

# Görsel yükleme alanı
uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])

if uploaded_file:
    image = Image.open(uploaded_file)
    st.image(image, caption="Uploaded Image", use_column_width=True)

    # Tahmin butonu
    if st.button("🔎 Classify Image"):
        # Görseli işle
        processed_image = process_image(image)

        # Model ile tahmin yap
        with torch.no_grad():
            output = model(processed_image)
            predicted_class = torch.argmax(output, dim=1).item()

        # Sınıf isimleri
        class_names = ["Smooth Round", "Smooth Cigar", "Edge-on Disk", "Spiral"]
        st.write(f"**Prediction:** {class_names[predicted_class]} 🎯")