Edited streamlit_app.py

src/streamlit_app.py

@@ -1107,20 +1107,20 @@ def build_mask_batch_zip(selected_names, artifacts: dict) -> bytes:
     return buf.getvalue()
 
 
-def _zip_png_mask(zf: zipfile.ZipFile, member: str) -> np.ndarray:
-    with zf.open(member) as f:
-        return (np.array(Image.open(f).convert("L")) > 0).astype(np.uint8)
-
-
-def _pick_mask_member(files: dict, base_name: str):
-    for ext in (".png", ".tif", ".tiff"):
-        member = files.get(f"{base_name}{ext}")
-        if member:
-            return member
-    return None
-
-
-def import_correction_zip(zip_file, artifacts: dict) -> tuple[int, list[str]]:
+def _zip_png_mask(zf: zipfile.ZipFile, member: str) -> np.ndarray:
+    with zf.open(member) as f:
+        return (np.array(Image.open(f).convert("L")) > 0).astype(np.uint8)
+
+
+def _pick_mask_member(files: dict, base_name: str):
+    for ext in (".png", ".tif", ".tiff"):
+        member = files.get(f"{base_name}{ext}")
+        if member:
+            return member
+    return None
+
+
+def import_correction_zip(zip_file, artifacts: dict) -> tuple[int, list[str]]:
     submitted = 0
     errors = []
     artifact_by_safe = {safe_zip_stem(name): name for name in artifacts}
@@ -1137,14 +1137,14 @@ def import_correction_zip(zip_file, artifacts: dict) -> tuple[int, list[str]]:
 
     # Preferred format: folders produced by build_mask_batch_zip.
     folders = sorted({str(Path(m).parent).replace("\\", "/") for m in members if "/" in m})
-    for folder in folders:
-        if folder in ("", "."):
-            continue
-        files = {Path(m).name.lower(): m for m in members if str(Path(m).parent).replace("\\", "/") == folder}
-        nuc_member = _pick_mask_member(files, "nuclei_mask")
-        myo_member = _pick_mask_member(files, "myotube_mask")
-        if not (nuc_member and myo_member):
-            continue
+    for folder in folders:
+        if folder in ("", "."):
+            continue
+        files = {Path(m).name.lower(): m for m in members if str(Path(m).parent).replace("\\", "/") == folder}
+        nuc_member = _pick_mask_member(files, "nuclei_mask")
+        myo_member = _pick_mask_member(files, "myotube_mask")
+        if not (nuc_member and myo_member):
+            continue
         source = artifact_by_safe.get(Path(folder).name)
         meta_member = files.get("meta.json")
         if meta_member:
@@ -1154,12 +1154,12 @@ def import_correction_zip(zip_file, artifacts: dict) -> tuple[int, list[str]]:
                 source = meta.get("source_image", source)
             except Exception:
                 pass
-        if source not in artifacts:
-            errors.append(f"Skipped {folder}: source image is not in current analysis.")
-            continue
-        try:
-            nuc_arr = _zip_png_mask(zf, nuc_member)
-            myo_arr = _zip_png_mask(zf, myo_member)
+        if source not in artifacts:
+            errors.append(f"Skipped {folder}: source image is not in current analysis.")
+            continue
+        try:
+            nuc_arr = _zip_png_mask(zf, nuc_member)
+            myo_arr = _zip_png_mask(zf, myo_member)
             orig_rgb = artifacts[source]["rgb_u8"]
             add_to_queue(orig_rgb, nuc_mask=nuc_arr, myo_mask=myo_arr,
                          reason="user_correction_batch",
@@ -1212,6 +1212,79 @@ def import_correction_zip(zip_file, artifacts: dict) -> tuple[int, list[str]]:
 # Model  (architecture identical to training script)
 # ─────────────────────────────────────────────────────────────────────────────
 
+
+def correction_stem(upload_name: str) -> str:
+    stem = Path(upload_name).stem
+    lower = stem.lower()
+    for suffix in (
+        "_nuclei_mask", "-nuclei-mask", "_nuclei", "-nuclei", "_nuc", "-nuc",
+        "_myotube_mask", "-myotube-mask", "_myotubes_mask", "-myotubes-mask",
+        "_myotube", "-myotube", "_myotubes", "-myotubes", "_myo", "-myo",
+    ):
+        if lower.endswith(suffix):
+            return stem[: -len(suffix)].strip()
+    return stem.strip()
+
+
+def uploaded_mask_to_array(uploaded_file) -> np.ndarray:
+    return (np.array(Image.open(uploaded_file).convert("L")) > 0).astype(np.uint8)
+
+
+def submit_uploaded_correction_pairs(nuc_files, myo_files, artifacts: dict, selected_name: str):
+    nuc_files = nuc_files or []
+    myo_files = myo_files or []
+    if len(nuc_files) != len(myo_files):
+        return 0, [f"Upload the same number of nuclei and myotube masks ({len(nuc_files)} nuclei, {len(myo_files)} myotube)."]
+
+    by_stem = {Path(name).stem: name for name in artifacts}
+    by_safe = {safe_zip_stem(name): name for name in artifacts}
+    nuc_by_stem = {correction_stem(f.name): f for f in nuc_files}
+    myo_by_stem = {correction_stem(f.name): f for f in myo_files}
+    common_stems = sorted(set(nuc_by_stem) & set(myo_by_stem))
+
+    if len(nuc_files) == len(myo_files) == 1 and not common_stems:
+        common_stems = [selected_name]
+        nuc_by_stem[selected_name] = nuc_files[0]
+        myo_by_stem[selected_name] = myo_files[0]
+
+    submitted = 0
+    errors = []
+    missing_nuc = sorted(set(myo_by_stem) - set(nuc_by_stem))
+    missing_myo = sorted(set(nuc_by_stem) - set(myo_by_stem))
+    for stem in missing_nuc:
+        errors.append(f"Skipped {stem}: no matching nuclei mask.")
+    for stem in missing_myo:
+        errors.append(f"Skipped {stem}: no matching myotube mask.")
+
+    for stem in common_stems:
+        source = by_stem.get(stem) or by_safe.get(stem)
+        if source is None and len(common_stems) == 1:
+            source = selected_name
+        if source not in artifacts:
+            errors.append(f"Skipped {stem}: no matching analyzed image is loaded.")
+            continue
+
+        try:
+            nuc_arr = uploaded_mask_to_array(nuc_by_stem[stem])
+            myo_arr = uploaded_mask_to_array(myo_by_stem[stem])
+            add_to_queue(
+                artifacts[source]["rgb_u8"],
+                nuc_mask=nuc_arr,
+                myo_mask=myo_arr,
+                reason="user_correction_batch" if len(common_stems) > 1 else "user_correction",
+                metadata={
+                    "source_image": source,
+                    "nuclei_file": nuc_by_stem[stem].name,
+                    "myotube_file": myo_by_stem[stem].name,
+                    "timestamp": datetime.now().isoformat(),
+                },
+            )
+            submitted += 1
+        except Exception as exc:
+            errors.append(f"Skipped {stem}: {exc}")
+
+    return submitted, errors
+
 class DoubleConv(nn.Module):
     def __init__(self, in_ch, out_ch):
         super().__init__()
@@ -1927,38 +2000,11 @@ if not private_mode and names:
 if enable_al and not private_mode:
     st.subheader("🧠 Submit corrected labels (Active Learning)")
     st.caption(
-        "Download the predicted masks, correct them in ImageJ/FIJI or any "
+        "Download predicted masks, correct them in ImageJ/FIJI, napari, or any "
         "image editor (white = foreground, black = background), then upload "
-        "the corrected versions below."
-    )
-
-    st.download_button(
-        "Download batch correction ZIP",
-        data=build_mask_batch_zip(names, st.session_state.artifacts),
-        file_name=f"myoseg_correction_batch_{datetime.now().strftime('%Y%m%d_%H%M%S')}.zip",
-        mime="application/zip",
-        key="download_correction_batch",
-        help="Includes one folder per image with image.png, nuclei_mask.png, myotube_mask.png, and meta.json.",
+        "one corrected pair or many corrected pairs below."
     )
 
-    batch_zip = st.file_uploader(
-        "Upload corrected batch ZIP",
-        type=["zip"],
-        key="batch_corr_zip",
-        help="Re-upload the ZIP after editing each folder's nuclei_mask.png and myotube_mask.png.",
-    )
-    if st.button("Submit batch corrections", type="secondary"):
-        if batch_zip is None:
-            st.error("Please upload a corrected ZIP batch first.")
-        else:
-            n_submitted, batch_errors = import_correction_zip(batch_zip, st.session_state.artifacts)
-            if n_submitted:
-                st.success(f"Saved {n_submitted} corrected image pair(s) to corrections/.")
-            if batch_errors:
-                st.warning("Some files were skipped:\n" + "\n".join(f"- {e}" for e in batch_errors[:10]))
-            if not n_submitted and not batch_errors:
-                st.error("No correction pairs were found in the ZIP.")
-
     al_pick = st.selectbox("Correct masks for image", names, key="al_pick")
 
     # Download buttons for current masks — lets users edit and re-upload
@@ -2004,27 +2050,38 @@ if enable_al and not private_mode:
 
     acol1, acol2 = st.columns(2)
     with acol1:
-        corr_nuc = st.file_uploader("Upload corrected NUCLEI mask",
-                                    type=["png", "tif", "tiff"], key="nuc_corr")
+        corr_nuc = st.file_uploader(
+            "Upload corrected NUCLEI mask(s)",
+            type=["png", "tif", "tiff"],
+            accept_multiple_files=True,
+            key="nuc_corr",
+            help="Upload one nuclei mask for the selected image, or select many nuclei masks for batch correction.",
+        )
     with acol2:
-        corr_myo = st.file_uploader("Upload corrected MYOTUBE mask",
-                                    type=["png", "tif", "tiff"], key="myo_corr")
+        corr_myo = st.file_uploader(
+            "Upload corrected MYOTUBE mask(s)",
+            type=["png", "tif", "tiff"],
+            accept_multiple_files=True,
+            key="myo_corr",
+            help="Upload the matching myotube mask(s). Batch files are paired by filename stem.",
+        )
 
     if st.button("✅ Submit corrections", type="primary"):
-        if corr_nuc is None or corr_myo is None:
-            st.error("Please upload BOTH a nuclei mask and a myotube mask.")
+        if not corr_nuc or not corr_myo:
+            st.error("Please upload BOTH nuclei and myotube mask files.")
         else:
-            orig_rgb   = st.session_state.artifacts[al_pick]["rgb_u8"]
-            nuc_arr    = (np.array(Image.open(corr_nuc).convert("L")) > 0).astype(np.uint8)
-            myo_arr    = (np.array(Image.open(corr_myo).convert("L")) > 0).astype(np.uint8)
-            add_to_queue(orig_rgb, nuc_mask=nuc_arr, myo_mask=myo_arr,
-                         reason="user_correction",
-                         metadata={"source_image": al_pick,
-                                   "timestamp": datetime.now().isoformat()})
-            st.success(
-                f"✅ Corrections for **{al_pick}** saved to `corrections/`. "
-                "The model will retrain at the next scheduled cycle."
+            n_submitted, correction_errors = submit_uploaded_correction_pairs(
+                corr_nuc, corr_myo, st.session_state.artifacts, al_pick
             )
+            if n_submitted == 1:
+                st.success("Saved 1 corrected image pair to corrections/.")
+            elif n_submitted > 1:
+                st.success(f"Saved {n_submitted} corrected image pairs to corrections/.")
+            if correction_errors:
+                st.warning("Some files were skipped:\n" + "\n".join(f"- {e}" for e in correction_errors[:12]))
+            if not n_submitted and not correction_errors:
+                st.error("No matching correction pairs were found.")
+
 elif enable_al and private_mode:
     st.info(
         "🔒 Active learning and training submissions are disabled in Private mode. "