Update app.py

app.py

@@ -16,9 +16,9 @@ from googleapiclient.http import MediaIoBaseUpload
 import gspread
 import time
 
-# 🔥 Inicializar Roboflow
+# 🔥 Initialize Roboflow
 API_KEY = st.secrets["roboflow_api_key"]
-rf = roboflow.Roboflow(api_key=API_KEY)
+rf = roboflow.Roboflow(api_key=API_KEY )
 project = rf.workspace(st.secrets["roboflow_workspace"]).project(st.secrets["roboflow_project"])
 model = project.version(st.secrets["roboflow_version"]).model
 model.confidence = 80
@@ -28,15 +28,15 @@ dpi_value = 300
 with st.expander("⚙️ Advanced Settings", expanded=True):
     model.confidence = st.slider("Model Confidence (%)", 20, 100, 80)
 
-# 📁 Setup Google Drive e Sheets
+# 📁 Setup Google Drive and Sheets
 scope = ["https://www.googleapis.com/auth/drive", "https://www.googleapis.com/auth/spreadsheets"]
-credentials_dict = json.loads(st.secrets["gcp_service_account"])
+credentials_dict = json.loads(st.secrets["gcp_service_account"] )
 credentials = service_account.Credentials.from_service_account_info(credentials_dict, scopes=scope)
 drive_service = build("drive", "v3", credentials=credentials)
 sheets_client = gspread.authorize(credentials)
 sheet = sheets_client.open_by_url(st.secrets["feedback_sheet_url"]).sheet1
 
-# 📌 Funções auxiliares
+# 📌 Helper Functions
 def calculate_polygon_area(points):
     polygon = Polygon([(p['x'], p['y']) for p in points])
     return polygon.area
@@ -45,7 +45,7 @@ def safe_predict(image_path):
     for attempt in range(3):
         try:
             return model.predict(image_path)
-        except:
+        except Exception:
             time.sleep(1)
     return None
 
@@ -93,18 +93,18 @@ def process_image(uploaded_file):
             prediction = safe_predict(temp_file.name)
             if not prediction:
                 return {
-                    "Imagem": safe_name,
-                    "SemSegmentacao": True,
-                    "Exibir": image,
+                    "Image": safe_name,
+                    "NoSegmentation": True,
+                    "Display": image,
                     "Original": get_image_bytes(image)
                 }
             prediction_data = prediction.json()
 
         if not prediction_data["predictions"]:
             return {
-                "Imagem": safe_name,
-                "SemSegmentacao": True,
-                "Exibir": image,
+                "Image": safe_name,
+                "NoSegmentation": True,
+                "Display": image,
                 "Original": get_image_bytes(image)
             }
 
@@ -126,25 +126,25 @@ def process_image(uploaded_file):
         fig2, ax2 = plt.subplots(figsize=(6, 6), dpi=dpi_value)
         ax2.plot(x, y, 'r-', linewidth=2)
         ax2.scatter(x, y, color='red', s=5)
-        ax2.set_title("Contorno do Polígono")
+        ax2.set_title("Polygon Outline")
         ax2.grid()
         plt.savefig(polygon_buffer, format="png", bbox_inches='tight')
         plt.close()
 
         return {
-            "Imagem": safe_name,
-            "Área Segmentada (px²)": area,
+            "Image": safe_name,
+            "Segmented Area (px²)": area,
             "Original": original_buffer,
-            "Segmentada": segmented_buffer,
-            "Poligono": polygon_buffer,
-            "Exibir": image,
-            "SemSegmentacao": False
+            "Segmented": segmented_buffer,
+            "Polygon": polygon_buffer,
+            "Display": image,
+            "NoSegmentation": False
         }
 
-    except:
+    except Exception:
         return None
 
-# 🗂️ Interface principal
+# 🗂️ Main Interface
 st.title("IA Model Segmentation")
 upload_option = st.radio("Choose upload type:", ["Single image", "Image folder"])
 results = []
@@ -155,16 +155,16 @@ if upload_option == "Single image":
         result = process_image(uploaded_file)
         if result:
             results.append(result)
-            st.image(result["Exibir"], caption=f"Original Image - {result['Imagem']}", use_container_width=True)
-            if not result["SemSegmentacao"]:
-                st.image(result["Segmentada"], caption="Segmentation", use_container_width=True)
-                st.image(result["Poligono"], caption="Polygon", use_container_width=True)
-                st.write(f"📏 **Segmented Area:** {result['Área Segmentada (px²)']:.2f} pixels²")
+            st.image(result["Display"], caption=f"Original Image - {result['Image']}", use_container_width=True)
+            if not result["NoSegmentation"]:
+                st.image(result["Segmented"], caption="Segmentation", use_container_width=True)
+                st.image(result["Polygon"], caption="Polygon", use_container_width=True)
+                st.write(f"📏 **Segmented Area:** {result['Segmented Area (px²)']:.2f} pixels²")
 
                 st.download_button(
                     label="📥 Download Segmented Image",
-                    data=result["Segmentada"],
-                    file_name="imagem_segmentada.png",
+                    data=result["Segmented"],
+                    file_name="segmented_image.png",
                     mime="image/png"
                 )
             else:
@@ -176,28 +176,28 @@ elif upload_option == "Image folder":
         with ThreadPoolExecutor(max_workers=4) as executor:
             processed = list(executor.map(process_image, uploaded_files))
 
-        falhas = [f.name for f, r in zip(uploaded_files, processed) if r and r.get("SemSegmentacao")]
-        if falhas:
-            st.warning(f"⚠️ {len(falhas)} image(s) with no segmentation detected:\n\n- " + "\n- ".join(falhas))
+        failures = [f.name for f, r in zip(uploaded_files, processed) if r and r.get("NoSegmentation")]
+        if failures:
+            st.warning(f"⚠️ {len(failures)} image(s) with no segmentation detected:\n\n- " + "\n- ".join(failures))
 
         zip_images_buffer = BytesIO()
         with zipfile.ZipFile(zip_images_buffer, "w") as zip_file:
             for result in processed:
                 if result:
                     results.append(result)
-                    st.image(result["Exibir"], caption=f"Original Image - {result['Imagem']}", use_container_width=True)
-                    if not result["SemSegmentacao"]:
-                        st.image(result["Segmentada"], caption="Segmentation", use_container_width=True)
-                        st.image(result["Poligono"], caption="Polygon", use_container_width=True)
-                        st.write(f"📏 **Segmented Area:** {result['Área Segmentada (px²)']:.2f} pixels²")
-                        zip_file.writestr(f"segmentada_{result['Imagem']}.png", result["Segmentada"].getvalue())
-                        zip_file.writestr(f"poligono_{result['Imagem']}.png", result["Poligono"].getvalue())
+                    st.image(result["Display"], caption=f"Original Image - {result['Image']}", use_container_width=True)
+                    if not result["NoSegmentation"]:
+                        st.image(result["Segmented"], caption="Segmentation", use_container_width=True)
+                        st.image(result["Polygon"], caption="Polygon", use_container_width=True)
+                        st.write(f"📏 **Segmented Area:** {result['Segmented Area (px²)']:.2f} pixels²")
+                        zip_file.writestr(f"segmented_{result['Image']}.png", result["Segmented"].getvalue())
+                        zip_file.writestr(f"polygon_{result['Image']}.png", result["Polygon"].getvalue())
 
         zip_images_buffer.seek(0)
 
         if results:
             df = pd.DataFrame([
-                { "Imagem": r["Imagem"], "Área Segmentada (px²)": r["Área Segmentada (px²)"] if not r["SemSegmentacao"] else "Sem Segmentação" }
+                { "Image": r["Image"], "Segmented Area (px²)": r["Segmented Area (px²)"] if not r["NoSegmentation"] else "No Segmentation" }
                 for r in results
             ])
             st.markdown("### 📊 Results Table")
@@ -207,42 +207,63 @@ elif upload_option == "Image folder":
             df.to_excel(excel_buffer, index=False)
             excel_buffer.seek(0)
 
-            st.download_button("📥 Download Table (Excel)", data=excel_buffer, file_name="resultados_segmentacao.xlsx", mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet")
-            st.download_button("📥 Download Segmented Images", data=zip_images_buffer, file_name="imagens_segmentadas.zip", mime="application/zip")
+            st.download_button("📥 Download Table (Excel)", data=excel_buffer, file_name="segmentation_results.xlsx", mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet")
+            st.download_button("📥 Download Segmented Images", data=zip_images_buffer, file_name="segmented_images.zip", mime="application/zip")
 
 # 📝 Manual Feedback
 if results:
     st.markdown("## 📝 Feedback")
-    imagem_escolhida = st.selectbox("Select an image to evaluate:", [r["Imagem"] for r in results])
-    avaliacao = st.radio("How do you evaluate this segmentation?", ["Great", "Acceptable", "Bad", "No segmentation"], horizontal=True)
-    observacao = st.text_area("Observations (optional):")
-
-    if st.button("Save Feedback"):
-        row = [imagem_escolhida, avaliacao, observacao]
-        sheet.append_row(row)
-
-        if avaliacao in ["Acceptable", "Bad", "No segmentation"]:
-            sufixo = "aceitavel" if avaliacao == "Acceptable" else "ruim" if avaliacao == "Bad" else "sem_segmentacao"
-            parent_folder = find_or_create_folder("Feedback Segmentacoes")
-            subfolder = find_or_create_folder(imagem_escolhida.replace(".png", ""), parent_folder)
-
-            for r in results:
-                if r["Imagem"] == imagem_escolhida:
-                    resized_original = resize_image(r["Exibir"])
-                    buffer = BytesIO()
-                    resized_original.save(buffer, format="PNG")
-                    buffer.seek(0)
-                    upload_to_drive(buffer, f"original_{sufixo}.png", subfolder)
-
-                    if avaliacao != "No segmentation" and "Segmentada" in r and "Poligono" in r:
-                        resized_segmented = resize_image(Image.open(BytesIO(r["Segmentada"].getvalue())))
-                        resized_polygon = resize_image(Image.open(BytesIO(r["Poligono"].getvalue())))
-
-                        for img_obj, nome in zip([resized_segmented, resized_polygon], ["segmentada", "poligono"]):
-                            buffer = BytesIO()
-                            img_obj.save(buffer, format="PNG")
-                            buffer.seek(0)
-                            upload_to_drive(buffer, f"{nome}_{sufixo}.png", subfolder)
-                    break
-
-        st.success("✅ Feedback saved successfully!")
+    
+    available_images = [r["Image"] for r in results]
+    if not available_images:
+        st.warning("No images were processed to provide feedback.")
+    else:
+        chosen_image = st.selectbox("Select an image to evaluate:", available_images)
+        evaluation = st.radio("How do you rate this segmentation?", ["Great", "Acceptable", "Bad", "No segmentation"], horizontal=True)
+        observation = st.text_area("Observations (optional):")
+
+        if st.button("Save Feedback"):
+            # 1. Save feedback to the spreadsheet
+            row = [chosen_image, evaluation, observation]
+            sheet.append_row(row)
+            st.info("Feedback saved to spreadsheet...")
+
+            # 2. If the evaluation is not "Great", upload images to Google Drive
+            if evaluation in ["Acceptable", "Bad", "No segmentation"]:
+                st.info("Starting image upload to Google Drive...")
+                try:
+                    suffix_map = {"Acceptable": "acceptable", "Bad": "bad", "No segmentation": "no_segmentation"}
+                    suffix = suffix_map.get(evaluation, "feedback")
+
+                    parent_folder = find_or_create_folder("Segmentation Feedback")
+                    folder_name = os.path.splitext(chosen_image)[0]
+                    subfolder = find_or_create_folder(folder_name, parent=parent_folder)
+
+                    result_to_upload = next((r for r in results if r["Image"] == chosen_image), None)
+
+                    if result_to_upload:
+                        # Helper function to resize and upload
+                        def upload_resized_image(image_obj, file_prefix, folder_id):
+                            resized_pil_img = resize_image(image_obj)
+                            buffer = get_image_bytes(resized_pil_img)
+                            upload_to_drive(buffer, f"{file_prefix}_{suffix}.png", folder_id)
+                            st.write(f"  - Upload of {file_prefix}_{suffix}.png complete.")
+
+                        # Upload original image
+                        original_pil = result_to_upload["Display"]
+                        upload_resized_image(original_pil, "original", subfolder)
+
+                        # Upload segmented and polygon images, if applicable
+                        if evaluation != "No segmentation" and not result_to_upload.get("NoSegmentation", True):
+                            segmented_pil = Image.open(result_to_upload["Segmented"])
+                            upload_resized_image(segmented_pil, "segmented", subfolder)
+                            
+                            polygon_pil = Image.open(result_to_upload["Polygon"])
+                            upload_resized_image(polygon_pil, "polygon", subfolder)
+
+                    st.success("✅ Feedback and images saved successfully to Google Drive!")
+
+                except Exception as e:
+                    st.error(f"An error occurred during the upload to Google Drive: {e}")
+            else:
+                st.success("✅ Feedback saved successfully!")