Update app.py

app.py

@@ -8,6 +8,11 @@ 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(
@@ -39,183 +44,80 @@ st.markdown("""
         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)
 
-# Cache the model loading
-@st.cache_resource
-def load_model_and_scaler():
-    """Load the trained model and scaler"""
+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:
-        model = tf.keras.models.load_model('mlp_model.h5')
-        # Tải scaler đã lưu
-        scaler = joblib.load('standard_scaler.pkl')
-        return model, scaler
+        # Đă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:
-        st.error(f"Error loading model or scaler: {str(e)}")
-        return None, None
-
-def color_histogram(image, bins=16):
-    """Calculate color histogram features"""
-    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)
-        percentages = counts / len(labels)
-        return np.concatenate([centers.flatten(), percentages])
-    except Exception:
-        return np.zeros(k * 4)
-
-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])
-
-@st.cache_resource
-def create_vit_feature_extractor():
-    """Create and cache the 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)
-
-def extract_features(image):
-    """Extract all features from an image"""
-    # Traditional features
-    hist_features = color_histogram(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()])
-
-def preprocess_image(image, scaler):
-    """Preprocess the uploaded image"""
-    # 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 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:
@@ -230,14 +132,14 @@ def main():
                 """,
                 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]:
@@ -246,11 +148,50 @@ def main():
                     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")