app.py

import streamlit as st
import tensorflow as tf
import numpy as np
import cv2
from PIL import Image
from tensorflow.keras import layers, models
from tensorflow.keras.applications import EfficientNetB0
from tensorflow.keras.applications.efficientnet import preprocess_input
import joblib
import io
import os
from google.oauth2.credentials import Credentials
from googleapiclient.discovery import build
from googleapiclient.http import MediaFileUpload
from google.oauth2 import service_account

# Set page config
st.set_page_config(
    page_title="Stone Classification",
    page_icon="🪨",
    layout="wide"
)

# Custom CSS with improved styling
st.markdown("""
    <style>
    .main {
        padding: 2rem;
    }
    .stButton>button {
        width: 100%;
        margin-top: 1rem;
    }
    .prediction-card {
        padding: 2rem;
        border-radius: 0.5rem;
        background-color: #d7d7d9;
        margin: 1rem 0;
    }
    .top-predictions {
        margin-top: 2rem;
        padding: 1rem;
        background-color: white;
        border-radius: 0.5rem;
        box-shadow: 0 1px 3px rgba(0,0,0,0.12);
    }
    .survey-card {
        padding: 1rem;
        background-color: #f0f2f6;
        border-radius: 0.5rem;
        margin-top: 1rem;
    }
    </style>
    """, unsafe_allow_html=True)

from mega import Mega

# Đăng nhập vào tài khoản Mega
def upload_to_mega(file_path, folder_name):
    """
    Upload file to a specific folder on Mega.nz
    """
    try:
        # Đăng nhập vào tài khoản Mega
        mega = Mega()
        m = mega.login(os.getenv('EMAIL'), os.getenv('PASSWORD'))

        # Tìm thư mục đích
        folder = m.find(folder_name)

        if not folder:
            # Nếu thư mục không tồn tại, hiển thị thông báo lỗi
            return f"Thư mục '{folder_name}' không tồn tại!"

        # Tải tệp lên thư mục
        file = m.upload(file_path, folder[0])
        return f"Upload thành công! Link: {m.get_upload_link(file)}"

    except Exception as e:
        return f"Lỗi khi tải lên Mega: {str(e)}"


def main():
    st.title("🪨 Stone Classification")
    st.write("Upload an image of a stone to classify its type")

    # Load model and scaler
    model, scaler = load_model_and_scaler()
    if model is None or scaler is None:
        st.error("Failed to load model or scaler. Please ensure both files exist.")
        return

    # Initialize session state
    if 'predictions' not in st.session_state:
        st.session_state.predictions = None
    if 'uploaded_image' not in st.session_state:
        st.session_state.uploaded_image = None

    col1, col2 = st.columns(2)

    with col1:
        st.subheader("Upload Image")
        uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

        if uploaded_file is not None:
            try:
                image = Image.open(uploaded_file)
                st.image(image, caption="Uploaded Image", use_column_width=True)
                st.session_state.uploaded_image = image

                with st.spinner('Analyzing image...'):
                    processed_image = preprocess_image(image, scaler)
                    prediction = model.predict(processed_image, verbose=0)

                    class_names = ['10', '6.5', '7', '7.5', '8', '8.5', '9', '9.2', '9.5', '9.7']
                    st.session_state.predictions = get_top_predictions(prediction, class_names)

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

    with col2:
        st.subheader("Prediction Results")
        if st.session_state.predictions:
            # Display main prediction
            top_class, top_confidence = st.session_state.predictions[0]
            st.markdown(
                f"""
                <div class='prediction-card'>
                    <h3>Primary Prediction: Grade {top_class}</h3>
                    <h3>Confidence: {top_confidence:.2f}%</h3>
                </div>
                """,
                unsafe_allow_html=True
            )

            # Display confidence bar
            st.progress(top_confidence / 100)

            # Display top 5 predictions
            st.markdown("### Top 5 Predictions")
            st.markdown("<div class='top-predictions'>", unsafe_allow_html=True)

            for class_name, confidence in st.session_state.predictions:
                cols = st.columns([2, 6, 2])
                with cols[0]:
                    st.write(f"Grade {class_name}")
                with cols[1]:
                    st.progress(confidence / 100)
                with cols[2]:
                    st.write(f"{confidence:.2f}%")

            st.markdown("</div>", unsafe_allow_html=True)

            # User Survey
            st.markdown("<div class='survey-card'>", unsafe_allow_html=True)
            st.markdown("### Model Accuracy Survey")
            st.write("Mô hình có dự đoán chính xác màu sắc của đá trong ảnh này không?")

            # Accuracy Confirmation
            accuracy = st.radio(
                "Đánh giá độ chính xác",
                ["Chọn", "Chính xác", "Không chính xác"],
                index=0
            )

            if accuracy == "Không chính xác":
              # Color input for incorrect prediction
              correct_color = st.text_input(
                  "Vui lòng nhập màu sắc chính xác của đá:",
                  help="Ví dụ: 10, 9.7, 9.5, 9.2, v.v."
              )

              if st.button("Gửi phản hồi và tải ảnh"):
                  if correct_color and st.session_state.uploaded_image:
                      # Save the image temporarily
                      temp_image_path = f"temp_image_{hash(uploaded_file.name)}.png"
                      st.session_state.uploaded_image.save(temp_image_path)

                      # Upload to Mega.nz
                      upload_result = upload_to_mega(temp_image_path, correct_color)
                      if "Upload thành công" in upload_result:
                          st.success(upload_result)
                      else:
                          st.error(upload_result)

                      # Clean up temporary file
                      os.remove(temp_image_path)
                  else:
                      st.warning("Vui lòng nhập màu sắc chính xác")

            st.markdown("</div>", unsafe_allow_html=True)
        else:
            st.info("Upload an image to see the predictions")

    st.markdown("---")
    st.markdown("Made with ❤️ using Streamlit")

if __name__ == "__main__":
    main()