Updated Streamlit

src/streamlit_app.py

@@ -62,10 +62,42 @@ import matplotlib.patches as mpatches
 from huggingface_hub import hf_hub_download
 
 import scipy.ndimage as ndi
-from skimage.morphology import remove_small_objects, disk, closing, opening, binary_dilation, binary_erosion
+from skimage.morphology import remove_small_objects, disk
 from skimage import measure
 from skimage.segmentation import watershed, find_boundaries
-from skimage.feature import peak_local_max
+from skimage.feature import peak_local_max
+
+
+def _as_bool_c_array(mask) -> np.ndarray:
+    return np.ascontiguousarray(np.asarray(mask, dtype=bool))
+
+
+def _safe_binary_dilation(mask, footprint):
+    return ndi.binary_dilation(
+        _as_bool_c_array(mask),
+        structure=np.asarray(footprint, dtype=bool),
+    )
+
+
+def _safe_binary_erosion(mask, footprint):
+    return ndi.binary_erosion(
+        _as_bool_c_array(mask),
+        structure=np.asarray(footprint, dtype=bool),
+    )
+
+
+def _safe_binary_opening(mask, footprint):
+    return ndi.binary_opening(
+        _as_bool_c_array(mask),
+        structure=np.asarray(footprint, dtype=bool),
+    )
+
+
+def _safe_binary_closing(mask, footprint):
+    return ndi.binary_closing(
+        _as_bool_c_array(mask),
+        structure=np.asarray(footprint, dtype=bool),
+    )
 
 
 # ─────────────────────────────────────────────────────────────────────────────
@@ -149,17 +181,17 @@ def postprocess_masks(nuc_mask, myo_mask,
     # Nuclei — closing fills small gaps, then size filter
     nuc_bin = nuc_mask.astype(bool)
     if int(nuc_close_radius) > 0:
-        nuc_bin = closing(nuc_bin, disk(int(nuc_close_radius)))
+        nuc_bin = _safe_binary_closing(nuc_bin, disk(int(nuc_close_radius)))
     nuc_clean = remove_small_objects(nuc_bin, min_size=int(min_nuc_area)).astype(np.uint8)
 
     # Myotubes — opening + erode/dilate + size filter + shape filter
     myo_bin = myo_mask.astype(bool)
     if int(myo_open_radius) > 0:
-        myo_bin = opening(myo_bin, disk(int(myo_open_radius)))
+        myo_bin = _safe_binary_opening(myo_bin, disk(int(myo_open_radius)))
     if int(myo_erode_radius) > 0:
         se = disk(int(myo_erode_radius))
-        myo_bin = binary_erosion(myo_bin, se)
-        myo_bin = binary_dilation(myo_bin, se)   # re-dilate to restore size
+        myo_bin = _safe_binary_erosion(myo_bin, se)
+        myo_bin = _safe_binary_dilation(myo_bin, se)   # re-dilate to restore size
     myo_bin = remove_small_objects(myo_bin, min_size=int(min_myo_area))
     if float(min_myo_aspect_ratio) > 0:
         myo_bin = _filter_by_aspect_ratio(myo_bin, float(min_myo_aspect_ratio))
@@ -366,7 +398,7 @@ def classify_nucleus_in_myotube(nuc_coords: np.ndarray,
     nuc_footprint = np.zeros((H, W), dtype=bool)
     nuc_footprint[rows, cols] = True
 
-    nuc_dilated = binary_dilation(nuc_footprint, footprint=disk(ring_width))
+    nuc_dilated = _safe_binary_dilation(nuc_footprint, disk(ring_width))
     ring_mask   = nuc_dilated & ~nuc_footprint & myo_mask_full.astype(bool)
 
     if ring_mask.sum() < 4:
@@ -596,14 +628,14 @@ def make_outline_overlay(rgb_u8: np.ndarray,
     if myo_disp.max() > 0:
         myo_bounds = find_boundaries(myo_disp, mode='outer')
         if line_width > 1:
-            myo_bounds = binary_dilation(myo_bounds, footprint=disk(line_width - 1))
+            myo_bounds = _safe_binary_dilation(myo_bounds, disk(line_width - 1))
         out[myo_bounds] = myo_color
 
     # Nuclei outlines
     if nuc_disp.max() > 0:
         nuc_bounds = find_boundaries(nuc_disp, mode='outer')
         if line_width > 1:
-            nuc_bounds = binary_dilation(nuc_bounds, footprint=disk(max(line_width - 2, 0)))
+            nuc_bounds = _safe_binary_dilation(nuc_bounds, disk(max(line_width - 2, 0)))
         out[nuc_bounds] = nuc_color
 
     return out
@@ -1551,7 +1583,133 @@ h1,h2,h3,h4 { color:#90caf9; }
 """, unsafe_allow_html=True)
 
 st.title("🔬 MyoSeg — Myotube & Nuclei Analyser")
-device = "cuda" if torch.cuda.is_available() else "cpu"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+V2_RUN_DIR = Path("runs_v2")
+V2_PID_PATH = V2_RUN_DIR / "train_v2.pid"
+V2_LOG_PATH = V2_RUN_DIR / "train_v2_latest.log"
+
+
+def _pid_is_running(pid: int) -> bool:
+    if pid <= 0:
+        return False
+    try:
+        os.kill(pid, 0)
+        return True
+    except OSError:
+        return False
+
+
+def _v2_training_status():
+    if not V2_PID_PATH.exists():
+        return False, None
+    try:
+        pid = int(V2_PID_PATH.read_text().strip())
+    except Exception:
+        return False, None
+    return _pid_is_running(pid), pid
+
+
+def _start_v2_training(epochs: int, batch_size: int, base: int, image_size: int):
+    V2_RUN_DIR.mkdir(parents=True, exist_ok=True)
+    cmd = [
+        sys.executable,
+        "train_model_v2.py",
+        "--data_root", "data",
+        "--corrections_dir", "corrections",
+        "--out_dir", str(V2_RUN_DIR),
+        "--epochs", str(int(epochs)),
+        "--batch_size", str(int(batch_size)),
+        "--image_size", str(int(image_size)),
+        "--base", str(int(base)),
+        "--push_to_hub",
+        "--hf_repo", MODEL_REPO_ID,
+        "--hf_filename", "model_v2.pt",
+    ]
+    env = os.environ.copy()
+    env["PYTHONUNBUFFERED"] = "1"
+    with V2_LOG_PATH.open("ab", buffering=0) as log_f:
+        log_f.write(
+            f"\n\n=== MyoSeg v2 training started {datetime.now().isoformat()} ===\n".encode()
+        )
+        log_f.write(("Command: " + " ".join(cmd) + "\n\n").encode())
+        proc = subprocess.Popen(
+            cmd,
+            stdout=log_f,
+            stderr=subprocess.STDOUT,
+            cwd=".",
+            env=env,
+        )
+    V2_PID_PATH.write_text(str(proc.pid))
+    return proc.pid
+
+
+def render_v2_training_panel(image_size: int, key_prefix: str = "v2_top"):
+    train_script = Path("train_model_v2.py")
+    has_script = train_script.exists()
+    has_token = bool(os.environ.get("HF_TOKEN"))
+    running, pid = _v2_training_status()
+    latest_model = V2_RUN_DIR / "model_v2.pt"
+    latest_metrics = V2_RUN_DIR / "metrics_v2.json"
+
+    s1, s2, s3 = st.columns(3)
+    s1.metric("train_model_v2.py", "Found" if has_script else "Missing")
+    s2.metric("HF_TOKEN", "Ready" if has_token else "Missing")
+    s3.metric("v2 status", f"Running pid {pid}" if running else "Idle")
+
+    if not has_script:
+        st.error("Upload train_model_v2.py to the Space before starting v2 training.")
+    if not has_token:
+        st.warning("Add HF_TOKEN as a Space secret with write access before training.")
+
+    c1, c2, c3 = st.columns(3)
+    with c1:
+        v2_epochs = st.number_input(
+            "v2 epochs", min_value=5, max_value=200, value=80, step=5,
+            key=f"{key_prefix}_epochs",
+        )
+    with c2:
+        v2_batch = st.selectbox("v2 batch size", [2, 4], index=0, key=f"{key_prefix}_batch")
+    with c3:
+        v2_base = st.selectbox("v2 base width", [24, 32], index=1, key=f"{key_prefix}_base")
+
+    st.caption(
+        "This trains attention_resunet_v2 on data/images plus data/masks/Nuclei_m "
+        "and data/masks/Myotubes_m, then uploads model_v2.pt to the model repo."
+    )
+
+    start_disabled = running or not has_script or not has_token
+    if st.button("Start v2 training now", type="primary", disabled=start_disabled, key=f"{key_prefix}_start"):
+        new_pid = _start_v2_training(
+            epochs=int(v2_epochs),
+            batch_size=int(v2_batch),
+            base=int(v2_base),
+            image_size=int(image_size),
+        )
+        st.success(f"v2 training started in the background. Process id: {new_pid}")
+        st.info("Refresh this page to update status. Watch the log preview below for progress.")
+
+    if running:
+        st.info("v2 training is currently running. Avoid restarting the Space until it finishes.")
+
+    if latest_model.exists():
+        st.success(f"Local v2 model exists: {latest_model}")
+    if latest_metrics.exists():
+        try:
+            st.json(json.loads(latest_metrics.read_text()))
+        except Exception:
+            st.caption(f"Metrics file exists: {latest_metrics}")
+
+    if V2_LOG_PATH.exists():
+        log_text = V2_LOG_PATH.read_text(errors="replace")
+        st.download_button(
+            "Download v2 training log",
+            data=log_text,
+            file_name="train_v2_latest.log",
+            mime="text/plain",
+            key=f"{key_prefix}_download_log",
+        )
+        st.text_area("Latest v2 training log", log_text[-6000:], height=260, key=f"{key_prefix}_log")
 
 # ─────────────────────────────────────────────────────────────────────────────
 # SIDEBAR
@@ -1746,7 +1904,11 @@ with st.sidebar:
 
 
 # ─────────────────────────────────────────────────────────────────────────────
-# FILE UPLOADER
+with st.expander("Train v2 model", expanded=True):
+    render_v2_training_panel(image_size=int(image_size), key_prefix="v2_top")
+
+
+# FILE UPLOADER
 # ─────────────────────────────────────────────────────────────────────────────
 _uploader_label = (
     "Upload 1+ images (png / jpg / tif). Any filename is accepted. 🔒 Private mode is ON — images will not be stored."