Update app.py

app.py

@@ -8,42 +8,175 @@ from tensorflow.keras.applications import EfficientNetB0
 from tensorflow.keras.applications.efficientnet import preprocess_input
 import joblib
 import io
+import os
 
-# Set page config
-st.set_page_config(
-    page_title="Stone Classification",
-    page_icon="🪨",
-    layout="wide"
+# Add Cloudinary import
+import cloudinary
+import cloudinary.uploader
+from cloudinary.utils import cloudinary_url
+
+# Cloudinary Configuration
+cloudinary.config(
+    cloud_name = os.getenv("CLOUD"),
+    api_key = os.getenv("API"),
+    api_secret = os.getenv("SECRET"),
+    secure=True
 )
 
-# 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);
-    }
-    </style>
-    """, unsafe_allow_html=True)
-
-# Cache the model loading
-@st.cache_resource
+def upload_to_cloudinary(file_path, label):
+    """
+    Upload file to Cloudinary with specified label as folder
+    """
+    try:
+        # Upload to Cloudinary
+        upload_result = cloudinary.uploader.upload(
+            file_path,
+            folder=label,
+            public_id=f"{label}_{os.path.basename(file_path)}"
+        )
+
+        # Generate optimized URLs
+        optimize_url, _ = cloudinary_url(
+            upload_result['public_id'],
+            fetch_format="auto",
+            quality="auto"
+        )
+
+        auto_crop_url, _ = cloudinary_url(
+            upload_result['public_id'],
+            width=500,
+            height=500,
+            crop="auto",
+            gravity="auto"
+        )
+
+        return {
+            "upload_result": upload_result,
+            "optimize_url": optimize_url,
+            "auto_crop_url": auto_crop_url
+        }
+
+    except Exception as e:
+        return f"Error uploading to Cloudinary: {str(e)}"
+
+def main():
+    st.title("🨨 Phân loại đá")
+    st.write("Tải lên hình ảnh của một viên đá để phân loại loại của nó.")
+
+    # Load model and scaler
+    model, scaler = load_model_and_scaler()
+    if model is None or scaler is None:
+        st.error("Không thể tải mô hình hoặc bộ chuẩn hóa. Vui lòng đảm bảo rằng cả hai tệp đều tồn tại.")
+        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("Tải lên Hình ảnh")
+        uploaded_file = st.file_uploader("Chọn hình ảnh...", type=["jpg", "jpeg", "png"])
+
+        if uploaded_file is not None:
+            try:
+                image = Image.open(uploaded_file)
+                st.image(image, caption="Hình ảnh đã tải lên", use_column_width=True)
+                st.session_state.uploaded_image = image
+
+                with st.spinner('Đang phân tích hình ảnh...'):
+                    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"Lỗi khi xử lý hình ảnh: {str(e)}")
+
+    with col2:
+        st.subheader("Kết quả Dự đoán")
+        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>Dự đoán chính: Màu {top_class}</h3>
+                    <h3>Độ tin cậy: {top_confidence:.2f}%</h3>
+                </div>
+                """,
+                unsafe_allow_html=True
+            )
+
+            # Display confidence bar
+            st.progress(top_confidence / 100)
+
+            # Display top 5 predictions
+            st.markdown("### 5 Dự đoán hàng đầu")
+            st.markdown("<div class='top-predictions'>", unsafe_allow_html=True)
+
+            for class_name, confidence in st.session_state.predictions:
+                st.markdown(
+                    f"**Màu {class_name}: Độ tin cậy {confidence:.2f}%**"
+                )
+                st.progress(confidence / 100)
+
+            st.markdown("</div>", unsafe_allow_html=True)
+
+            # User Confirmation Section
+            st.markdown("### Xác nhận độ chính xác của mô hình")
+            st.write("Giúp chúng tôi cải thiện mô hình bằng cách xác nhận độ chính xác của dự đoán.")
+
+            # Accuracy Radio Button
+            accuracy_option = st.radio(
+                "Dự đoán có chính xác không?",
+                ["Chọn", "Chính xác", "Không chính xác"],
+                index=0
+            )
+
+            if accuracy_option == "Không chính xác":
+                # Input for correct grade
+                correct_grade = st.selectbox(
+                    "Chọn màu đá đúng:",
+                    ['10', '6.5', '7', '7.5', '8', '8.5', '9', '9.2', '9.5', '9.7'],
+                    index=None,
+                    placeholder="Chọn màu đúng"
+                )
+
+                # Upload button
+                if st.button("Tải lên Hình ảnh để sửa chữa"):
+                    if correct_grade 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)
+
+                        try:
+                            # Upload to Cloudinary
+                            cloudinary_result = upload_to_cloudinary(temp_image_path, correct_grade)
+
+                            if isinstance(cloudinary_result, dict):
+                                st.success(f"Hình ảnh đã được tải lên thành công cho màu {correct_grade}")
+                                st.write(f"URL công khai: {cloudinary_result['upload_result']['secure_url']}")
+                            else:
+                                st.error(cloudinary_result)
+
+                            # Clean up temporary file
+                            os.remove(temp_image_path)
+
+                        except Exception as e:
+                            st.error(f"Tải lên thất bại: {str(e)}")
+                    else:
+                        st.warning("Vui lòng chọn màu đúng trước khi tải lên.")
+        else:
+            st.info("Tải lên hình ảnh để xem các dự đoán.")
+
+    st.markdown("---")
+    st.markdown("Tạo bởi ❤️ với Streamlit")
+
 def load_model_and_scaler():
     """Load the trained model and scaler"""
     try:
@@ -60,34 +193,34 @@ def color_histogram(image, bins=16):
     hist_r = cv2.calcHist([image], [0], None, [bins], [0, 256]).flatten()
     hist_g = cv2.calcHist([image], [1], None, [bins], [0, 256]).flatten()
     hist_b = cv2.calcHist([image], [2], None, [bins], [0, 256]).flatten()
-    
+
     hist_r = hist_r / (np.sum(hist_r) + 1e-7)
     hist_g = hist_g / (np.sum(hist_g) + 1e-7)
     hist_b = hist_b / (np.sum(hist_b) + 1e-7)
-    
+
     return np.concatenate([hist_r, hist_g, hist_b])
 
 def color_moments(image):
     """Calculate color moments features"""
     img = image.astype(np.float32) / 255.0
     moments = []
-    
+
     for i in range(3):
         channel = img[:,:,i]
         mean = np.mean(channel)
         std = np.std(channel) + 1e-7
         skewness = np.mean(((channel - mean) / std) ** 3) if std != 0 else 0
         moments.extend([mean, std, skewness])
-    
+
     return np.array(moments)
 
 def dominant_color_descriptor(image, k=3):
     """Calculate dominant color descriptor"""
     pixels = image.reshape(-1, 3).astype(np.float32)
-    
+
     criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 100, 0.2)
     flags = cv2.KMEANS_RANDOM_CENTERS
-    
+
     try:
         _, labels, centers = cv2.kmeans(pixels, k, None, criteria, 10, flags)
         unique, counts = np.unique(labels, return_counts=True)
@@ -100,16 +233,16 @@ def color_coherence_vector(image, k=3):
     """Calculate color coherence vector"""
     gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
     gray = np.uint8(gray)
-    
+
     _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
     num_labels, labels = cv2.connectedComponents(binary)
-    
+
     ccv = []
     for i in range(1, min(k+1, num_labels)):
         region_mask = (labels == i)
         total_pixels = np.sum(region_mask)
         ccv.extend([total_pixels, total_pixels])
-    
+
     ccv.extend([0] * (2 * k - len(ccv)))
     return np.array(ccv[:2*k])
 
@@ -119,13 +252,13 @@ def create_vit_feature_extractor():
     input_shape = (256, 256, 3)
     inputs = layers.Input(shape=input_shape)
     x = layers.Lambda(preprocess_input)(inputs)
-    
+
     base_model = EfficientNetB0(
         include_top=False,
         weights='imagenet',
         input_tensor=x
     )
-    
+
     x = layers.GlobalAveragePooling2D()(base_model.output)
     return models.Model(inputs=inputs, outputs=x)
 
@@ -136,21 +269,21 @@ def extract_features(image):
     moment_features = color_moments(image)
     dominant_features = dominant_color_descriptor(image)
     ccv_features = color_coherence_vector(image)
-    
+
     traditional_features = np.concatenate([
         hist_features,
         moment_features,
         dominant_features,
         ccv_features
     ])
-    
+
     # Deep features using ViT
     feature_extractor = create_vit_feature_extractor()
     vit_features = feature_extractor.predict(
         np.expand_dims(image, axis=0),
         verbose=0
     )
-    
+
     # Combine all features
     return np.concatenate([traditional_features, vit_features.flatten()])
 
@@ -159,100 +292,25 @@ def preprocess_image(image, scaler):
     # Convert to RGB if needed
     if image.mode != 'RGB':
         image = image.convert('RGB')
-    
+
     # Convert to numpy array and resize
     img_array = np.array(image)
     img_array = cv2.resize(img_array, (256, 256))
     img_array = img_array.astype('float32') / 255.0
-    
+
     # Extract all features
     features = extract_features(img_array)
-    
+
     # Scale features using the provided scaler
     scaled_features = scaler.transform(features.reshape(1, -1))
-    
-    return scaled_features
 
-def get_top_predictions(prediction, class_names, top_k=5):
-    """Get top k predictions with their probabilities"""
-    top_indices = prediction.argsort()[0][-top_k:][::-1]
-    return [
-        (class_names[i], float(prediction[0][i]) * 100)
-        for i in top_indices
-    ]
+    return scaled_features
 
-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
-    
-    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)
-                
-                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)
-        else:
-            st.info("Upload an image to see the predictions")
-    
-    st.markdown("---")
-    st.markdown("Made with ❤️ using Streamlit")
+def get_top_predictions(prediction, class_names):
+    # Extract the top 5 predictions with confidence values
+    probabilities = tf.nn.softmax(prediction[0]).numpy()
+    top_indices = np.argsort(probabilities)[-5:][::-1]
+    return [(class_names[i], probabilities[i] * 100) for i in top_indices]
 
 if __name__ == "__main__":
-    main()
+    main()