Updated the streamlit app and removed the train.py

src/streamlit_app.py

@@ -42,8 +42,6 @@ import time
 import zipfile
 import hashlib
 import re
-import subprocess
-import sys
 from datetime import datetime
 from pathlib import Path
 
@@ -104,7 +102,7 @@ def _safe_binary_closing(mask, footprint):
 # CONFIG  ← edit these two lines to match your HF model repo
 # ─────────────────────────────────────────────────────────────────────────────
 MODEL_REPO_ID  = "skarugu/myotube-unet"
-MODEL_FILENAME = os.environ.get("MODEL_FILENAME", "model_final.pt")
+MODEL_FILENAME = "model_final.pt"
 
 CONF_FLAG_THR   = 0.60    # images below this confidence are queued for retraining
 QUEUE_DIR       = Path("retrain_queue")
@@ -1411,153 +1409,28 @@ class UNet(nn.Module):
         return self.out(x)
 
 
-def _group_count(channels: int) -> int:
-    for groups in (8, 4, 2, 1):
-        if channels % groups == 0:
-            return groups
-    return 1
-
-
-class ResidualBlock(nn.Module):
-    def __init__(self, in_ch: int, out_ch: int):
-        super().__init__()
-        self.conv1 = nn.Conv2d(in_ch, out_ch, 3, padding=1, bias=False)
-        self.norm1 = nn.GroupNorm(_group_count(out_ch), out_ch)
-        self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False)
-        self.norm2 = nn.GroupNorm(_group_count(out_ch), out_ch)
-        self.skip = nn.Identity() if in_ch == out_ch else nn.Conv2d(in_ch, out_ch, 1, bias=False)
-
-    def forward(self, x):
-        identity = self.skip(x)
-        x = torch.nn.functional.silu(self.norm1(self.conv1(x)))
-        x = self.norm2(self.conv2(x))
-        return torch.nn.functional.silu(x + identity)
-
-
-class AttentionGate(nn.Module):
-    def __init__(self, gate_ch: int, skip_ch: int, inter_ch: int):
-        super().__init__()
-        self.gate = nn.Sequential(
-            nn.Conv2d(gate_ch, inter_ch, 1, bias=True),
-            nn.GroupNorm(_group_count(inter_ch), inter_ch),
-        )
-        self.skip = nn.Sequential(
-            nn.Conv2d(skip_ch, inter_ch, 1, bias=True),
-            nn.GroupNorm(_group_count(inter_ch), inter_ch),
-        )
-        self.psi = nn.Sequential(nn.Conv2d(inter_ch, 1, 1, bias=True), nn.Sigmoid())
-
-    def forward(self, gate, skip):
-        alpha = self.psi(torch.nn.functional.silu(self.gate(gate) + self.skip(skip)))
-        return skip * alpha
-
-
-def _center_pad_like(x, ref):
-    diff_y = ref.size(2) - x.size(2)
-    diff_x = ref.size(3) - x.size(3)
-    if diff_x or diff_y:
-        x = torch.nn.functional.pad(
-            x,
-            [diff_x // 2, diff_x - diff_x // 2, diff_y // 2, diff_y - diff_y // 2],
-        )
-    return x
-
-
-class AttentionResUNetV2(nn.Module):
-    def __init__(self, in_ch: int = 2, out_ch: int = 2, base: int = 32):
-        super().__init__()
-        self.e1 = ResidualBlock(in_ch, base)
-        self.e2 = ResidualBlock(base, base * 2)
-        self.e3 = ResidualBlock(base * 2, base * 4)
-        self.e4 = ResidualBlock(base * 4, base * 8)
-        self.pool = nn.MaxPool2d(2)
-        self.bridge = ResidualBlock(base * 8, base * 16)
-
-        self.up4 = nn.ConvTranspose2d(base * 16, base * 8, 2, 2)
-        self.att4 = AttentionGate(base * 8, base * 8, base * 4)
-        self.dec4 = ResidualBlock(base * 16, base * 8)
-        self.up3 = nn.ConvTranspose2d(base * 8, base * 4, 2, 2)
-        self.att3 = AttentionGate(base * 4, base * 4, base * 2)
-        self.dec3 = ResidualBlock(base * 8, base * 4)
-        self.up2 = nn.ConvTranspose2d(base * 4, base * 2, 2, 2)
-        self.att2 = AttentionGate(base * 2, base * 2, base)
-        self.dec2 = ResidualBlock(base * 4, base * 2)
-        self.up1 = nn.ConvTranspose2d(base * 2, base, 2, 2)
-        self.att1 = AttentionGate(base, base, max(base // 2, 1))
-        self.dec1 = ResidualBlock(base * 2, base)
-        self.out = nn.Conv2d(base, out_ch, 1)
-
-    def forward(self, x):
-        e1 = self.e1(x)
-        e2 = self.e2(self.pool(e1))
-        e3 = self.e3(self.pool(e2))
-        e4 = self.e4(self.pool(e3))
-        b = self.bridge(self.pool(e4))
-
-        x = _center_pad_like(self.up4(b), e4)
-        x = self.dec4(torch.cat([x, self.att4(x, e4)], dim=1))
-        x = _center_pad_like(self.up3(x), e3)
-        x = self.dec3(torch.cat([x, self.att3(x, e3)], dim=1))
-        x = _center_pad_like(self.up2(x), e2)
-        x = self.dec2(torch.cat([x, self.att2(x, e2)], dim=1))
-        x = _center_pad_like(self.up1(x), e1)
-        x = self.dec1(torch.cat([x, self.att1(x, e1)], dim=1))
-        return self.out(x)
-
-
 @st.cache_resource
 def load_model(device: str):
-    active_filename = MODEL_FILENAME
-    try:
-        local = hf_hub_download(
-            repo_id=MODEL_REPO_ID,
-            filename=active_filename,
-            force_download=True,
-        )
-    except Exception as exc:
-        missing_requested_model = (
-            active_filename != "model_final.pt"
-            and (
-                "EntryNotFound" in exc.__class__.__name__
-                or "404" in str(exc)
-                or "Entry Not Found" in str(exc)
-            )
-        )
-        if not missing_requested_model:
-            raise
-        st.warning(
-            f"`{active_filename}` is not in `{MODEL_REPO_ID}` yet. "
-            "Using `model_final.pt` so the app can stay open while v2 is trained."
-        )
-        active_filename = "model_final.pt"
-        local = hf_hub_download(
-            repo_id=MODEL_REPO_ID,
-            filename=active_filename,
-            force_download=True,
-        )
+    local = hf_hub_download(
+        repo_id=MODEL_REPO_ID,
+        filename=MODEL_FILENAME,
+        force_download=True,
+    )
     file_sha = sha256_file(local)
     mtime    = time.ctime(os.path.getmtime(local))
     size_mb  = os.path.getsize(local) / 1e6
 
     st.sidebar.markdown("### 🔍 Model debug")
     st.sidebar.caption(f"Repo: `{MODEL_REPO_ID}`")
-    st.sidebar.caption(f"Requested file: `{MODEL_FILENAME}`")
-    st.sidebar.caption(f"Loaded file: `{active_filename}`")
+    st.sidebar.caption(f"File: `{MODEL_FILENAME}`")
     st.sidebar.caption(f"Size: {size_mb:.2f} MB")
     st.sidebar.caption(f"Modified: {mtime}")
     st.sidebar.caption(f"SHA256: `{file_sha[:20]}…`")
 
     ckpt  = torch.load(local, map_location=device)
     state = ckpt["model"] if isinstance(ckpt, dict) and "model" in ckpt else ckpt
-    args = ckpt.get("args", {}) if isinstance(ckpt, dict) else {}
-    arch = ckpt.get("arch", args.get("arch", "unet_v1")) if isinstance(ckpt, dict) else "unet_v1"
-    base = int(args.get("base", 32)) if isinstance(args, dict) else 32
-
-    if arch == "attention_resunet_v2":
-        model = AttentionResUNetV2(in_ch=2, out_ch=2, base=base)
-    else:
-        model = UNet(in_ch=2, out_ch=2, base=32)
-    st.sidebar.caption(f"Architecture: `{arch}`")
+    model = UNet(in_ch=2, out_ch=2, base=32)
+    st.sidebar.caption("Architecture: `unet_v1`")
     model.load_state_dict(state)
     model.to(device).eval()
     return model
@@ -1585,147 +1458,6 @@ h1,h2,h3,h4 { color:#90caf9; }
 st.title("🔬 MyoSeg — Myotube & Nuclei Analyser")
 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}")
-        st.download_button(
-            "Download model_v2.pt",
-            data=latest_model.read_bytes(),
-            file_name="model_v2.pt",
-            mime="application/octet-stream",
-            key=f"{key_prefix}_download_model",
-        )
-    if latest_metrics.exists():
-        try:
-            metrics_text = latest_metrics.read_text()
-            st.json(json.loads(metrics_text))
-            st.download_button(
-                "Download metrics_v2.json",
-                data=metrics_text,
-                file_name="metrics_v2.json",
-                mime="application/json",
-                key=f"{key_prefix}_download_metrics",
-            )
-        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
 # ─────────────────────────────────────────────────────────────────────────────
@@ -1919,10 +1651,6 @@ with st.sidebar:
 
 
 # ─────────────────────────────────────────────────────────────────────────────
-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 = (
@@ -2187,64 +1915,6 @@ with c3:
 
 st.divider()
 
-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
-
 # ─────────────────────────────────────────────────────────────────────────────
 # PER-IMAGE PREVIEW  +  ANIMATED METRICS
 # ─────────────────────────────────────────────────────────────────────────────
@@ -2637,82 +2307,3 @@ with st.expander("🔧 Self-training queue status"):
         import subprocess
         subprocess.Popen(["python", "self_train.py", "--manual"])
         st.info("Retraining started in the background. Check terminal / logs for progress.")
-
-
-with st.expander("Train v2 model"):
-    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)
-    with c2:
-        v2_batch = st.selectbox("v2 batch size", [2, 4], index=0)
-    with c3:
-        v2_base = st.selectbox("v2 base width", [24, 32], index=1)
-
-    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):
-        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}")
-        st.download_button(
-            "Download model_v2.pt",
-            data=latest_model.read_bytes(),
-            file_name="model_v2.pt",
-            mime="application/octet-stream",
-            key="v2_bottom_download_model",
-        )
-    if latest_metrics.exists():
-        try:
-            metrics_text = latest_metrics.read_text()
-            st.json(json.loads(metrics_text))
-            st.download_button(
-                "Download metrics_v2.json",
-                data=metrics_text,
-                file_name="metrics_v2.json",
-                mime="application/json",
-                key="v2_bottom_download_metrics",
-            )
-        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",
-        )
-        st.text_area("Latest v2 training log", log_text[-6000:], height=260)