Update app.py

app.py

@@ -93,6 +93,32 @@ def preprocess_image(image):
     
     return img_array
 
+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('[email protected]', '01283315889')
+
+        # 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 get_top_predictions(prediction, class_names, top_k=5):
     """Get top k predictions with their probabilities"""
     # Get indices of top k predictions
@@ -150,39 +176,78 @@ def main():
     
     with col2:
         st.subheader("Prediction Results")
-        if st.session_state.predictions is not None:
-            # Create a card-like container for results
-            results_container = st.container()
-            with results_container:
-                # Display main prediction
-                st.markdown("<div class='prediction-card'>", unsafe_allow_html=True)
-                top_class, top_confidence = st.session_state.predictions[0]
-                st.markdown(f"### Primary Prediction: Grade {top_class}")
-                st.markdown(f"### Confidence: {top_confidence:.2f}%")
-                st.markdown("</div>", unsafe_allow_html=True)
-                
-                # Display confidence bar for top prediction
-                st.progress(top_confidence / 100)
-                
-                # Display top 5 predictions
-                st.markdown("### Top 5 Predictions")
-                st.markdown("<div class='top-predictions'>", unsafe_allow_html=True)
-                
-                # Create a Streamlit container for the predictions
-                for class_name, confidence in st.session_state.predictions:
-                    col_label, col_bar, col_value = st.columns([2, 6, 2])
-                    with col_label:
-                        st.write(f"Grade {class_name}")
-                    with col_bar:
-                        st.progress(confidence / 100)
-                    with col_value:
-                        st.write(f"{confidence:.2f}%")
-                
-                st.markdown("</div>", unsafe_allow_html=True)
+        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 and click 'Predict' to see the results")
-    
-    # Footer
+            st.info("Upload an image to see the predictions")
+
     st.markdown("---")
     st.markdown("Made with ❤️ using Streamlit")