Update app.py

app.py

@@ -77,6 +77,12 @@ def find_or_create_folder(folder_name, parent=None):
     file = drive_service.files().create(body=file_metadata, fields='id').execute()
     return file.get('id')
 
+def get_image_bytes(image):
+    buf = BytesIO()
+    image.save(buf, format="PNG")
+    buf.seek(0)
+    return buf
+
 def process_image(uploaded_file):
     try:
         safe_name = uploaded_file.name.replace(" ", "_")
@@ -138,12 +144,6 @@ def process_image(uploaded_file):
     except:
         return None
 
-def get_image_bytes(image):
-    buf = BytesIO()
-    image.save(buf, format="PNG")
-    buf.seek(0)
-    return buf
-
 # 🗂️ Interface principal
 st.title("IA Model Segmentation")
 upload_option = st.radio("Choose upload type:", ["Single image", "Image folder"])
@@ -160,6 +160,20 @@ if upload_option == "Single image":
                 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.download_button(
+                    label="📥 Baixar Imagem Segmentada",
+                    data=result["Segmentada"],
+                    file_name="imagem_segmentada.png",
+                    mime="image/png"
+                )
+
+                st.download_button(
+                    label="📥 Baixar Polígono",
+                    data=result["Poligono"],
+                    file_name="poligono.png",
+                    mime="image/png"
+                )
             else:
                 st.warning("⚠️ No segmentation was detected in this image.")
 
@@ -214,22 +228,20 @@ if results:
         row = [imagem_escolhida, avaliacao, observacao]
         sheet.append_row(row)
 
-        if avaliacao in ["Aceitável", "Ruim", "Sem segmentação"]:
-            sufixo = "aceitavel" if avaliacao == "Aceitável" else "ruim" if avaliacao == "Ruim" else "sem_segmentacao"
+        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:
-                    # Sempre salva a original
                     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)
 
-                    # Só salva segmentada e polígono se houver segmentação
-                    if avaliacao != "Sem segmentação" and "Segmentada" in r and "Poligono" in r:
+                    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())))