Update app.py

app.py

@@ -1,9 +1,7 @@
 import streamlit as st
-import roboflow
 import pandas as pd
 import matplotlib.pyplot as plt
 import zipfile
-import tempfile
 from shapely.geometry import Polygon
 from PIL import Image
 from io import BytesIO
@@ -11,21 +9,36 @@ from concurrent.futures import ThreadPoolExecutor
 from google.oauth2.credentials import Credentials
 from googleapiclient.discovery import build
 from googleapiclient.http import MediaIoBaseUpload
+from huggingface_hub import hf_hub_download
+from ultralytics import YOLO
 import gspread
+import numpy as np
 import time
 
+st.set_page_config(page_title="Scratch Assay Segmentation", layout="wide")
+
 APP_VERSION = "2.4"
+DEFAULT_IMGSZ = 640
+
+MODEL_OPTIONS = {
+    "2.4": "24.pt",
+    "3.7": "37.pt",
+}
+
 
 # =========================
-# Roboflow init
+# Local model init (Hugging Face private repo)
 # =========================
-API_KEY = st.secrets["roboflow_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
-model.overlap = 25
-dpi_value = 300
+@st.cache_resource
+def load_model(model_filename):
+    local_model_path = hf_hub_download(
+        repo_id=st.secrets["HF_MODEL_REPO"],
+        filename=model_filename,
+        repo_type="model",
+        token=st.secrets["HF_TOKEN"],
+    )
+    return YOLO(local_model_path)
+
 
 # =========================
 # Google Drive + Sheets (OAuth2)
@@ -43,6 +56,7 @@ drive_service = build("drive", "v3", credentials=credentials)
 sheets_client = gspread.authorize(credentials)
 sheet = sheets_client.open_by_url(st.secrets["feedback_sheet_url"]).sheet1
 
+
 # =========================
 # Helpers
 # =========================
@@ -51,10 +65,16 @@ def calculate_polygon_area(points):
     return polygon.area
 
 
-def safe_predict(image_path):
+def safe_predict(model, image_array, conf_threshold):
     for _ in range(3):
         try:
-            return model.predict(image_path)
+            results = model.predict(
+                source=image_array,
+                imgsz=DEFAULT_IMGSZ,
+                conf=conf_threshold,
+                verbose=False,
+            )
+            return results
         except Exception:
             time.sleep(1)
     return None
@@ -96,10 +116,11 @@ def get_image_bytes(image):
     return buf
 
 
-def process_image(uploaded_file, fov_um=None, pixel_size_um=None):
+def process_image(uploaded_file, model, model_confidence, fov_um=None, pixel_size_um=None):
     try:
         safe_name = uploaded_file.name.replace(" ", "_")
         image = Image.open(uploaded_file).convert("RGB")
+        image_np = np.array(image)
 
         width_px, _ = image.size
 
@@ -109,21 +130,22 @@ def process_image(uploaded_file, fov_um=None, pixel_size_um=None):
         elif fov_um is not None and fov_um > 0:
             effective_pixel_size_um = fov_um / float(width_px)
 
-        with tempfile.NamedTemporaryFile(suffix=".png", delete=True) as temp_file:
-            image.save(temp_file.name)
-            prediction = safe_predict(temp_file.name)
-            if not prediction:
-                return {
-                    "Imagem": safe_name,
-                    "Área Segmentada (px²)": None,
-                    "Área Segmentada (µm²)": None,
-                    "SemSegmentacao": True,
-                    "Exibir": image,
-                    "Original": get_image_bytes(image),
-                }
-            prediction_data = prediction.json()
-
-        if not prediction_data["predictions"]:
+        conf_threshold = model_confidence / 100.0
+        results = safe_predict(model, image_np, conf_threshold)
+
+        if not results or len(results) == 0:
+            return {
+                "Imagem": safe_name,
+                "Área Segmentada (px²)": None,
+                "Área Segmentada (µm²)": None,
+                "SemSegmentacao": True,
+                "Exibir": image,
+                "Original": get_image_bytes(image),
+            }
+
+        result = results[0]
+
+        if result.masks is None or len(result.masks.xyn) == 0:
             return {
                 "Imagem": safe_name,
                 "Área Segmentada (px²)": None,
@@ -133,12 +155,31 @@ def process_image(uploaded_file, fov_um=None, pixel_size_um=None):
                 "Original": get_image_bytes(image),
             }
 
-        points = prediction_data["predictions"][0]["points"]
+        best_idx = 0
+        if result.boxes is not None and result.boxes.conf is not None and len(result.boxes.conf) > 0:
+            best_idx = int(result.boxes.conf.argmax().item())
+
+        contour_norm = result.masks.xyn[best_idx]
+        if contour_norm is None or len(contour_norm) < 3:
+            return {
+                "Imagem": safe_name,
+                "Área Segmentada (px²)": None,
+                "Área Segmentada (µm²)": None,
+                "SemSegmentacao": True,
+                "Exibir": image,
+                "Original": get_image_bytes(image),
+            }
+
+        points = [
+            {"x": float(x * width_px), "y": float(y * image.height)}
+            for x, y in contour_norm
+        ]
+
         area_px2 = calculate_polygon_area(points)
 
         area_um2 = None
         if effective_pixel_size_um is not None:
-            area_um2 = area_px2 * (effective_pixel_size_um**2)
+            area_um2 = area_px2 * (effective_pixel_size_um ** 2)
 
         x = [p["x"] for p in points] + [points[0]["x"]]
         y = [p["y"] for p in points] + [points[0]["y"]]
@@ -146,7 +187,7 @@ def process_image(uploaded_file, fov_um=None, pixel_size_um=None):
         original_buffer = get_image_bytes(image)
 
         segmented_buffer = BytesIO()
-        fig, ax = plt.subplots(figsize=(6, 6), dpi=dpi_value)
+        fig, ax = plt.subplots(figsize=(6, 6), dpi=300)
         ax.imshow(image)
         ax.plot(x, y, color="red", linewidth=2)
         ax.axis("off")
@@ -154,7 +195,7 @@ def process_image(uploaded_file, fov_um=None, pixel_size_um=None):
         plt.close()
 
         polygon_buffer = BytesIO()
-        fig2, ax2 = plt.subplots(figsize=(6, 6), dpi=dpi_value)
+        fig2, ax2 = plt.subplots(figsize=(6, 6), dpi=300)
         ax2.plot(x, y, "r-", linewidth=2)
         ax2.scatter(x, y, color="red", s=5)
         ax2.set_title("Polygon contour")
@@ -239,8 +280,6 @@ def render_feedback_block(result, prefix_key=""):
 # =========================
 # Layout / UI
 # =========================
-st.set_page_config(page_title="Scratch Assay Segmentation", layout="wide")
-
 st.title("Scratch Assay Segmentation Tool")
 st.caption(f"Version {APP_VERSION} · Deep learning–based wound closure segmentation")
 
@@ -248,11 +287,19 @@ st.markdown("---")
 
 # Upload block
 st.markdown("### Input")
-upload_option = st.radio("Choose upload type:", ["Single image", "Image folder"], horizontal=True)
+col_input_1, col_input_2 = st.columns([2, 1])
+
+with col_input_1:
+    upload_option = st.radio("Choose upload type:", ["Single image", "Image folder"], horizontal=True)
+
+with col_input_2:
+    selected_model_label = st.selectbox("Model checkpoint", list(MODEL_OPTIONS.keys()), index=0)
+
+model = load_model(MODEL_OPTIONS[selected_model_label])
 
 # Advanced settings (collapsed by default)
 with st.expander("⚙️ Advanced Settings", expanded=False):
-    model.confidence = st.slider("Model confidence (%)", 20, 100, 80)
+    model_confidence = st.slider("Model confidence (%)", 20, 100, 80)
     st.markdown(
         "### Physical calibration (optional)\n"
         "Provide the physical scale for conversion from pixel area to physical units (µm²). "
@@ -298,7 +345,13 @@ if upload_option == "Single image":
         st.markdown("---")
         st.markdown("### Result")
 
-        result = process_image(uploaded_file, fov_um=fov_um, pixel_size_um=pixel_size_um)
+        result = process_image(
+            uploaded_file,
+            model=model,
+            model_confidence=model_confidence,
+            fov_um=fov_um,
+            pixel_size_um=pixel_size_um,
+        )
         if result:
             results.append(result)
 
@@ -331,6 +384,7 @@ if upload_option == "Single image":
             st.markdown("---")
             render_feedback_block(result, prefix_key="single_")
 
+
 # =========================
 # Folder
 # =========================
@@ -346,9 +400,16 @@ elif upload_option == "Image folder":
         st.markdown("### Processing")
 
         def process_wrapper(f):
-            return process_image(f, fov_um=fov_um, pixel_size_um=pixel_size_um)
+            return process_image(
+                f,
+                model=model,
+                model_confidence=model_confidence,
+                fov_um=fov_um,
+                pixel_size_um=pixel_size_um,
+            )
 
-        with ThreadPoolExecutor(max_workers=4) as executor:
+        # Local inference is more stable with a single worker.
+        with ThreadPoolExecutor(max_workers=1) as executor:
             processed = list(executor.map(process_wrapper, uploaded_files))
 
         falhas = [f.name for f, r in zip(uploaded_files, processed) if r and r.get("SemSegmentacao")]