diff --git a/.github/workflows/deploy.yml b/.github/workflows/deploy.yml new file mode 100644 index 0000000..54c52e6 --- /dev/null +++ b/.github/workflows/deploy.yml @@ -0,0 +1,41 @@ +name: Deploy to GitHub Pages + +on: + push: + branches: + - main + +permissions: + contents: write + +defaults: + run: + working-directory: docs + +jobs: + deploy: + name: Deploy to GitHub Pages + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + + - uses: actions/setup-node@v4 + with: + node-version: 20 + cache: npm + cache-dependency-path: docs/package-lock.json + + - name: Install dependencies + run: npm ci + + - name: Build website + run: npm run build + + - name: Deploy to GitHub Pages + uses: peaceiris/actions-gh-pages@v4 + with: + github_token: ${{ secrets.GITHUB_TOKEN }} + publish_dir: docs/build + commit_message: "Deploy website" diff --git a/.github/workflows/singularity.yml b/.github/workflows/singularity.yml index f309d25..9fe895f 100644 --- a/.github/workflows/singularity.yml +++ b/.github/workflows/singularity.yml @@ -8,36 +8,35 @@ jobs: build-test-containers: runs-on: ubuntu-latest strategy: - # Keep going on other deployments if anything bloops fail-fast: false matrix: + # Must match an existing tag on quay.io/singularity/singularity (full semver, e.g. 4.1.0 not 4.0). singularity_version: - - '3.8.1' + - '4.1.0' container: image: quay.io/singularity/singularity:v${{ matrix.singularity_version }} options: --privileged steps: - - name: Check out code for the container builds - uses: actions/checkout@v2 + uses: actions/checkout@v4 with: submodules: recursive - - name: Build Container + - name: Verify Singularity definition exists run: | - if [ -f Singularity ]; then - sudo -E singularity build container.sif Singularity - else - echo "Singularity is not found." - echo "Present working directory: $PWD" - ls - fi - + if [ ! -f Singularity ]; then + echo "Singularity definition file not found in $(pwd)" + ls -la + exit 1 + fi + + - name: Build Container + run: sudo -E singularity build container.sif Singularity + - name: Set Tag Based on Branch Name run: | - # Get the branch name and set it as the tag tag=${GITHUB_REF##*/}-singularity echo "Tag is $tag." echo "tag=$tag" >> $GITHUB_ENV @@ -45,6 +44,6 @@ jobs: - name: Login and Deploy Container if: (github.event_name != 'pull_request') run: | - echo ${{ secrets.GITHUB_TOKEN }} | singularity remote login -u ${{ secrets.GHCR_USERNAME }} --password-stdin oras://ghcr.io - container_name=$(echo ${GITHUB_REPOSITORY} | tr '[:upper:]' '[:lower:]') - singularity push container.sif oras://ghcr.io/${container_name}:${tag} + echo ${{ secrets.GITHUB_TOKEN }} | singularity remote login -u ${{ secrets.GHCR_USERNAME }} --password-stdin oras://ghcr.io + container_name=$(echo ${GITHUB_REPOSITORY} | tr '[:upper:]' '[:lower:]') + singularity push container.sif oras://ghcr.io/${container_name}:${tag} diff --git a/.github/workflows/test-deploy.yml b/.github/workflows/test-deploy.yml new file mode 100644 index 0000000..001515d --- /dev/null +++ b/.github/workflows/test-deploy.yml @@ -0,0 +1,31 @@ +name: Test deployment + +on: + pull_request: + branches: + - main + +defaults: + run: + working-directory: docs + +jobs: + test-deploy: + name: Test deployment + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + + - uses: actions/setup-node@v4 + with: + node-version: 20 + cache: npm + cache-dependency-path: docs/package-lock.json + + - name: Install dependencies + run: npm ci + + - name: Test build website + run: npm run build diff --git a/.gitignore b/.gitignore index 859c32a..0f0ef1c 100644 --- a/.gitignore +++ b/.gitignore @@ -164,3 +164,9 @@ cython_debug/ # DynamicDet run files runs +gui/.venv/ +partinet-env/ + +# Container images and scratch shell scripts +*.sif +*.sh diff --git a/Dockerfile b/Dockerfile index e084a47..c5e0be1 100644 --- a/Dockerfile +++ b/Dockerfile @@ -8,4 +8,4 @@ RUN python -m pip install --no-cache-dir --no-cache /opt/PartiNet LABEL AUTHORS Mihin Perera, Edward Yang, Julie Iskander LABEL MAINTAINERS Mihin Perera, Edward Yang, Julie Iskander -LABEL VERSION v1.0.1 +LABEL VERSION v1.1.0 diff --git a/README.md b/README.md index b5f83ee..e1275dc 100644 --- a/README.md +++ b/README.md @@ -17,6 +17,7 @@ Use our pretrained model at [Model Weights](https://huggingface.co/MihinP/PartiN - Seamless integration with cryoSPARC and RELION workflows - Confidence-based particle filtering - Visual detection validation +- Browser GUI (`partinet gui`) with optional Slurm job submission ## Prerequisites @@ -112,6 +113,7 @@ Available commands: - `denoise`: Clean input micrographs - `detect`: Identify particles - `star`: Generate STAR files +- `gui`: Launch the Gradio GUI (local or Slurm execution) - `train`: Train custom models diff --git a/Singularity b/Singularity index a6ffbb6..5039681 100644 --- a/Singularity +++ b/Singularity @@ -18,4 +18,4 @@ from: python:3.9.19-slim-bookworm %labels AUTHORS Mihin Perera, Edward Yang, Julie Iskander MAINTAINERS Mihin Perera, Edward Yang, Julie Iskander - VERSION v1.0.1 \ No newline at end of file + VERSION v1.1.0 \ No newline at end of file diff --git a/docs/docs/getting-started.md b/docs/docs/getting-started.md index 9ec1cbf..b8bc4ec 100644 --- a/docs/docs/getting-started.md +++ b/docs/docs/getting-started.md @@ -6,6 +6,14 @@ sidebar_position: 3 This guide walks you through your first PartiNet analysis using the three-stage pipeline. We'll process cryo-EM micrographs from start to finish. +:::tip New to PartiNet? +The **[GUI](gui.md)** is the easiest way to get started — no command-line experience needed. Launch it with `partinet gui` and run all three stages from your browser. + +![PartiNet GUI overview](/img/gui_overview.png) + +The CLI walkthrough below is for users who prefer scripting, are working in containers, or need to run PartiNet on an HPC cluster. +::: + ## Prerequisites Before starting, ensure you have: @@ -24,27 +32,29 @@ project_directory/ │ ├── micrograph1.mrc │ ├── micrograph2.mrc │ └── ... -├── denoised/ # Created by denoise stage -│ ├── micrograph1.mrc -│ ├── micrograph2.mrc +├── denoised/ # Created by denoise stage (default: .png) +│ ├── micrograph1.png +│ ├── micrograph2.png │ └── ... ├── exp/ # Created by detect stage │ ├── labels/ # Detection coordinates (YOLO format) │ │ ├── micrograph1.txt │ │ ├── micrograph2.txt │ │ └── ... -│ ├── micrograph1.png # Micrographs with detections drawn -│ ├── micrograph2. +│ ├── micrograph1.png # Micrographs with detections drawn +│ ├── micrograph2.png │ └── ... └── partinet_particles.star # CryoSPARC-style STAR file (created by star stage) ``` **Pipeline Flow:** 1. **Input** → `motion_corrected/` (your micrographs) -2. **Stage 1** → `denoised/` (cleaned micrographs) +2. **Stage 1** → `denoised/` (cleaned micrographs, PNG by default) 3. **Stage 2** → `exp*/` (detections + visualizations) 4. **Stage 3** → `*.star` (final particle coordinates) +See [Detect — skip denoise](stages/detect.md#skip-denoise-raw-mrc) if picking directly on raw MRC without denoising. + ## Stage 1: Denoise The first stage removes noise from your micrographs and improves signal-to-noise ratios: @@ -53,8 +63,8 @@ The first stage removes noise from your micrographs and improves signal-to-noise ```shell title="Local Installation" partinet denoise \ - --source /data/my_project/motion_corrected \ - --project /data/my_project + --source /data/partinet_picking/motion_corrected \ + --project /data/partinet_picking ``` @@ -100,9 +110,9 @@ The final stage converts detections to STAR format and applies confidence filter ```shell title="Local Installation" partinet star \ - --labels /data/my_project/exp/labels \ - --images /data/my_project/denoised \ - --output /data/my_project/partinet_particles.star \ + --labels /data/partinet_picking/exp/labels \ + --images /data/partinet_picking/denoised \ + --output /data/partinet_picking/partinet_particles.star \ --conf 0.1 ``` @@ -118,7 +128,7 @@ partinet star \ After running all three stages, you'll have: 1. **Denoised micrographs** (`denoised/`) - Cleaned input for particle detection -2. **Detection visualizations** (`exp/*.mrc`) - Micrographs with particle boxes drawn +2. **Detection visualizations** (`exp/*.png`) - Micrographs with particle boxes drawn 3. **Detection coordinates** (`exp/labels/*.txt`) - Raw detection data 4. **STAR file** (`*.star`) - Final particle coordinates ready for downstream processing diff --git a/docs/docs/gui.md b/docs/docs/gui.md new file mode 100644 index 0000000..f0615d0 --- /dev/null +++ b/docs/docs/gui.md @@ -0,0 +1,185 @@ +--- +sidebar_position: 4 +--- + +# GUI + +PartiNet includes a browser-based graphical interface built with Gradio. It covers all three pipeline stages — Denoise, Detect, and Star File — and is designed as a beginner-friendly alternative to the CLI. + +![PartiNet GUI overview](/img/gui_overview.png) + +:::tip Recommended installation +Run the GUI from a **local Python installation** ([Installation](installation.md) — pip, conda, or a virtualenv on your workstation or cluster login node). Apptainer/Singularity and Docker images are intended for the **CLI** pipeline on HPC; running `partinet gui` inside a container is possible but not recommended. + +For large-scale or automated processing, the CLI remains the recommended approach. + +A local install avoids extra bind-mount and cache setup (`--no-home`, Hugging Face theme cache), keeps Slurm tools (`sbatch`, `scancel`) on the host PATH, and makes SSH port forwarding straightforward. On HPC, use the GUI locally on the login node and submit compute-heavy work via **Slurm** +::: + +## Launching the GUI + +```shell title="Local Python installation" +partinet gui +``` + +Gradio will print the URL to the terminal. On a login node, use `127.0.0.1` or `localhost` — not `0.0.0.0`, which is only a bind address. + +``` +Running on local URL: http://127.0.0.1:7860 +``` + +Open that URL in your browser on the login node, or set up SSH port forwarding from your laptop: + +```shell +ssh -L 7860:127.0.0.1:7860 user@login-node +``` + +Then open `http://localhost:7860` in your local browser. + +### Options + +| Option | Default | Description | +|--------|---------|-------------| +| `--host` | `127.0.0.1` | Host address to bind the server to | +| `--port` | auto | Port to run the GUI on; auto-selects a free port if omitted | +| `--share` | off | Create a temporary public Gradio link (useful for remote sharing) | + +```shell +# Pin a specific port +partinet gui --port 7860 + +# Generate a public share link +partinet gui --share +``` + +## Project Directory + +The **Project directory** field at the top of the GUI is set once and auto-fills all stage-specific paths: + +| Field | Auto-filled value | +|-------|-----------------| +| Denoise → Project directory | `` | +| Detect → Denoised images directory | `/denoised` | +| Detect → Project directory | `` | +| Star File → Labels directory | `/exp*/labels` (latest detect run) | +| Star File → Denoised images directory | `/denoised` | +| Star File → Output STAR file | `/particles.star` | + +Individual fields remain editable if you need to override a path for a specific stage. + +## Running jobs + +Open the **Running jobs** accordion (below the project directory) to see active jobs submitted from the GUI for the current project. The table refreshes automatically every five seconds, or click **Refresh**. + +| Column | Description | +|--------|-------------| +| Stage | `denoise`, `detect`, or `star` | +| Mode | `local` or `slurm` | +| ID | Slurm job ID or local process ID | +| Status | `running`, `pending`, etc. | +| Started | UTC timestamp when the job was submitted | + +To cancel a job, select it from **Select job** and click **Cancel selected**. Local jobs are terminated; Slurm jobs receive `scancel`. Job metadata is stored in `/.partinet_jobs/manifest.json` so Slurm jobs remain visible (and cancellable) after a browser refresh. + +Selecting a job in the dropdown shows its live log in **Job log** (refreshed every five seconds while selected). Logs are read from the stage log file (`partinet_denoise.log`, etc.) and, for Slurm jobs, the batch script output under `.partinet_jobs/`. + +Only jobs submitted from this GUI session for the current project directory are listed. + +## Appearance + +The GUI uses the [lone17/kotaemon](https://huggingface.co/spaces/lone17/kotaemon) Gradio theme and opens in **dark mode** by default (via `?__theme=dark` on first load). You can switch to light or system mode using Gradio’s theme toggle in the footer. + +The theme is downloaded from the Hugging Face Hub on first launch (requires network once; cached afterward). Star File statistics plots use a dark matplotlib style to match the UI. + +Optional environment variables at launch: + +- `PARTINET_PROJECT` — pre-fill the project directory +- `PARTINET_WEIGHTS` — pre-fill the Detect weights path +- `PARTINET_SLURM_CONFIG` — path to a user-owned YAML file with default Slurm settings + +## Execution modes (Local / Slurm) + +Open **Execution settings** to choose how each stage runs: + +| Mode | Description | +|------|-------------| +| **Local** | Run `partinet` as a subprocess on the machine hosting the GUI | +| **Slurm** | Write a batch script under `/.partinet_jobs/` and submit with `sbatch` | + +All Slurm fields (partition, account, time, CPUs, GPUs, memory) are optional — leave them blank to use your cluster defaults. Use **Job setup script** for site-specific setup (`module load`, `conda activate`, container `apptainer exec`, etc.) without hardcoding paths in the PartiNet source. + +Logs stream in the browser from `{project}/partinet_denoise.log`, `partinet_detect.log`, or `partinet_star.log`. Slurm job IDs are shown when a batch job is submitted. + +On HPC, run the GUI on a login or visualization node (with SSH port forwarding) and submit Denoise, Detect, and Star File jobs to compute nodes via Slurm. + +## Stage 1 · Denoise + +Applies the PartiNet Wiener filter denoiser to raw micrographs and saves cleaned images to `project/denoised/`. + +**Key inputs:** + +| Field | Description | +|-------|-------------| +| Raw micrographs directory | Folder of `.mrc` files from motion correction | +| Project directory | All PartiNet outputs are written here | +| Output image format | `png` (recommended), `jpg`, or `mrc` | +| CPU workers | Number of parallel workers; leave blank to auto-detect | + +Progress is streamed live into the log window during processing. + +## Stage 2 · Detect + +Runs the DynamicDet particle detector on denoised micrographs. Results are saved to `project/exp/`. + +**Key inputs:** + +| Field | Description | +|-------|-------------| +| Model weights (.pt) | Pre-trained weights file — download from HuggingFace | +| Denoised images directory | Output of the Denoise stage (`project/denoised/`) | +| Project directory | Detection outputs saved here | +| GPU device(s) | Leave blank to auto-detect; use `cpu` if no GPU is available | +| Confidence threshold | Lower = more picks. Recommended: 0.0–0.3 | +| IOU threshold | Controls removal of overlapping boxes. Default 0.2 | + +Advanced options (image size, dynamic threshold) are available under the collapsible **Advanced options** panel. + +## Stage 3 · Star File + +Loads detections from a Detect run, provides interactive statistics and a micrograph preview, then generates a STAR file for CryoSPARC or RELION. + +![Star File metrics plots](/img/gui_star_plot.png) + +### Load Detections + +Enter the labels directory and denoised images directory, then click **Load Detections**. The GUI displays: + +- Total detection count and particles-per-micrograph summary +- Confidence score distribution histogram +- Box size distribution histogram +- Particles per micrograph bar chart (updates with threshold) +- Confidence vs box size bivariate histogram + +### Threshold and Preview + +Use the **Confidence threshold** slider to filter detections interactively. The plots and retained-particle count update in real time. + +Select a micrograph from the **Preview micrograph** dropdown to see detection boxes drawn on the image. Box colour indicates confidence: green = high, red = low. + +![Micrograph preview with confidence-coloured detection boxes](/img/gui_micrograph.png) + +### Generate STAR File + +Click **Generate STAR File** to write a CryoSPARC-compatible `.star` file. The confidence threshold set above is applied at export time. + +**RELION output** is available under the optional accordion — enable it to also produce a `pick.star` and per-micrograph coordinate files in the RELION project directory format. + + + + + +## What's Next + +- [Denoise (CLI reference)](stages/denoise.md) +- [Detect (CLI reference)](stages/detect.md) +- [Star File (CLI reference)](stages/star.md) diff --git a/docs/docs/stages/denoise.md b/docs/docs/stages/denoise.md index 9114117..9bf4412 100644 --- a/docs/docs/stages/denoise.md +++ b/docs/docs/stages/denoise.md @@ -12,10 +12,12 @@ Denoising can vastly improve particle picking by helping to increase signal to n ### Required Parameters -| Parameter | Description | Example | -|-----------|-------------|---------| -| `--source` | Directory containing motion-corrected micrographs in MRC format | `/data/partinet_picking/motion_corrected` | -| `--project` | Parent project directory where all outputs will be saved | `/data/partinet_picking` | +| Parameter | Role | Example | +|-----------|------|---------| +| `--source` | Input folder for this stage (motion-corrected `.mrc` files) | `/data/partinet_picking/motion_corrected` | +| `--project` | Dataset root where outputs and logs are written | `/data/partinet_picking` | + +Denoised micrographs are always written to `/denoised/`. The `--source` directory is never modified. ### Optional Parameters @@ -166,7 +168,7 @@ partinet denoise \ ### Different Output Formats -By default PartiNet outputs denoised images in `png` format. This is necessary for compatibility with the detection architecture. `png` is a lossless compression, however micrographs are normalised from 32 bit depth `mrc` files to 8 bit `png`. `jpg` is also available (as a legacy format) but is not recommended for use due to lossy compression. +By default PartiNet outputs denoised images in `png` format. This is the recommended format for the Detect stage. `png` is lossless; micrographs are normalised from 32-bit MRC to 8-bit PNG. `jpg` is available but not recommended due to lossy compression. ```shell # JPEG format (smaller file size, lossy compression) diff --git a/docs/docs/stages/detect.md b/docs/docs/stages/detect.md index 8dc15f2..6b8333e 100644 --- a/docs/docs/stages/detect.md +++ b/docs/docs/stages/detect.md @@ -105,8 +105,30 @@ partinet detect \ ### Supported Formats -- **PNG files** (`.png`) -- **JPEG files** (`.jpg`, `.jpeg`) +**Standard workflow (recommended):** + +- **PNG files** (`.png`) from the Denoise stage — use with `denoised_micrographs.pt` + +**Skip-denoise workflow:** + +- **MRC files** (`.mrc`) motion-corrected micrographs — use with `raw_micrographs.pt` +- Raw MRC inputs are per-micrograph min–max normalised to 8-bit before inference (not the same as EMAN2 or CryoSegNet display) + +### Skip denoise (raw MRC) + +If you want to pick directly on motion-corrected micrographs without denoising: + +```shell +partinet detect \ + --weight /path/to/raw_micrographs.pt \ + --source /data/partinet_picking/motion_corrected \ + --project /data/partinet_picking \ + --device 0 +``` + +Use `raw_micrographs.pt` only with raw MRC inputs. For denoised PNG, use `denoised_micrographs.pt`. + +When generating a STAR file after skip-denoise, set `--images` to the **same folder** used as detect `--source` (e.g. `motion_corrected/`, not `denoised/`). ### Directory Structure @@ -133,16 +155,16 @@ partinet_picking/ │ ├── micrograph2.mrc │ └── ... ├── denoised/ # 🧹 Created by denoise stage -│ ├── micrograph1.mrc -│ ├── micrograph2.mrc +│ ├── micrograph1.png +│ ├── micrograph2.png │ └── ... ├── exp/ # 🎯 Created by detect stage │ ├── labels/ # 📋 Detection coordinates │ │ ├── micrograph1.txt │ │ ├── micrograph2.txt │ │ └── ... -│ ├── micrograph1.mrc # 🖼️ Micrographs with detections drawn -│ ├── micrograph2.mrc +│ ├── micrograph1.png # 🖼️ Micrographs with detections drawn +│ ├── micrograph2.png │ └── ... └── partinet_detect.log ``` diff --git a/docs/docs/stages/star.md b/docs/docs/stages/star.md index 5f0aad6..5ef1017 100644 --- a/docs/docs/stages/star.md +++ b/docs/docs/stages/star.md @@ -40,12 +40,14 @@ docker run --gpus all -v /data:/data \ | Parameter | Type | Required | Description | |-----------|------|----------|-------------| | `--labels` | Path | Yes | Directory containing the particle coordinate files (`.txt` format) from the detection stage | -| `--images` | Path | Yes | Directory containing the denoised micrographs corresponding to the labels | +| `--images` | Path | Yes | Directory containing micrographs used during Detect — usually `/denoised/`; if you skipped denoise, use the same folder as detect `--source` | | `--output` | Path | Yes | Output path for the generated STAR file | | `--conf` | Float | Yes | Confidence threshold for filtering particle coordinates (0.0-1.0) | ## Input Requirements +**Important:** `--images` must match the folder passed to `partinet detect --source`. After the standard pipeline that is `/denoised/`; after skip-denoise on raw MRC it is `/motion_corrected/` (or wherever your `.mrc` files live). + At this stage of the pipeline, your directory structure should look like: ``` partinet_picking/ diff --git a/docs/docusaurus.config.ts b/docs/docusaurus.config.ts index ae3af78..c06a985 100644 --- a/docs/docusaurus.config.ts +++ b/docs/docusaurus.config.ts @@ -15,7 +15,7 @@ const config: Config = { }, url: 'https://wehi-researchcomputing.github.io', - baseUrl: '/', + baseUrl: '/PartiNet/', organizationName: 'WEHI-ResearchComputing', projectName: 'PartiNet', deploymentBranch: 'gh-pages', diff --git a/docs/static/img/gui_micrograph.png b/docs/static/img/gui_micrograph.png new file mode 100644 index 0000000..4bdc456 Binary files /dev/null and b/docs/static/img/gui_micrograph.png differ diff --git a/docs/static/img/gui_overview.png b/docs/static/img/gui_overview.png new file mode 100644 index 0000000..80b3cd3 Binary files /dev/null and b/docs/static/img/gui_overview.png differ diff --git a/docs/static/img/gui_star_plot.png b/docs/static/img/gui_star_plot.png new file mode 100644 index 0000000..b6cb2c8 Binary files /dev/null and b/docs/static/img/gui_star_plot.png differ diff --git a/partinet/DynamicDet/detect.py b/partinet/DynamicDet/detect.py index 46affa6..fd0e801 100644 --- a/partinet/DynamicDet/detect.py +++ b/partinet/DynamicDet/detect.py @@ -33,10 +33,6 @@ sh.setLevel(logging.INFO) logger.addHandler(sh) -fh = logging.FileHandler('partinet_detect.log') -fh.setFormatter(formatter) -fh.setLevel(logging.INFO) -logger.addHandler(fh) sys.stdout.reconfigure(line_buffering=True) sys.stderr.reconfigure(line_buffering=True) @@ -69,6 +65,16 @@ def detect(opt, save_img=False): save_dir = Path(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok)) (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) + log_path = Path(opt.project) / 'partinet_detect.log' + for h in logger.handlers[:]: + if isinstance(h, logging.FileHandler): + logger.removeHandler(h) + h.close() + fh = logging.FileHandler(log_path) + fh.setFormatter(formatter) + fh.setLevel(logging.INFO) + logger.addHandler(fh) + set_logging() devices = [f'cuda:{i}' for i in range(torch.cuda.device_count())] if torch.cuda.is_available() else ['cpu'] logger.info(f"Using devices: {devices}") diff --git a/partinet/DynamicDet/utils/datasets.py b/partinet/DynamicDet/utils/datasets.py index 9e45cb0..9cdbe5d 100644 --- a/partinet/DynamicDet/utils/datasets.py +++ b/partinet/DynamicDet/utils/datasets.py @@ -30,6 +30,7 @@ from partinet.DynamicDet.utils.general import check_requirements, xyxy2xywh, xywh2xyxy, xywhn2xyxy, xyn2xy, segment2box, segments2boxes, \ resample_segments, clean_str from partinet.DynamicDet.utils.torch_utils import torch_distributed_zero_first +from partinet.process_utils.image_io import load_micrograph_for_detect from partinet.process_utils.guided_denoiser import transform # Parameters @@ -191,11 +192,7 @@ def __next__(self): self.count += 1 # support cryo-EM micrographs saved as MRC if path.lower().endswith('.mrc'): - img_mrc = mrcfile.read(path) - img0 = transform(img_mrc).astype(np.uint8) - # `transform` returns a single-channel image; networks expect BGR - if img0.ndim == 2: - img0 = cv2.cvtColor(img0, cv2.COLOR_GRAY2BGR) + img0 = load_micrograph_for_detect(path) else: img0 = cv2.imread(path) # BGR assert img0 is not None, 'Image Not Found ' + path @@ -684,10 +681,7 @@ def load_image(self, index): path = self.img_files[index] # support uncompressed mrc micrographs if path.lower().endswith('.mrc'): - img_mrc = mrcfile.read(path) - img = transform(img_mrc).astype(np.uint8) - if img.ndim == 2: - img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) + img = load_micrograph_for_detect(path) else: img = cv2.imread(path) # BGR assert img is not None, 'Image Not Found ' + path diff --git a/partinet/__init__.py b/partinet/__init__.py index 7d78337..ba11800 100644 --- a/partinet/__init__.py +++ b/partinet/__init__.py @@ -1,7 +1,7 @@ import click -import sys, os +import sys -__version__ = "1.0.1" +__version__ = "1.1.0" DYNAMICDET_AVAILABLE_MODELS = ["yolov7", "yolov7x", "yolov7-w6", "yolov7-e6", "yolov7-d6", "yolov7-e6e"] @@ -105,6 +105,15 @@ def denoise(source, project, num_workers,img_format): import partinet.process_utils.pooled_denoise_proc partinet.process_utils.pooled_denoise_proc.main(source,project,num_workers,img_format) +@main.command() +@click.option('--host', default='127.0.0.1', show_default=True, help='Host to bind the GUI server to') +@click.option('--port', default=None, type=int, help='Port to run the GUI on (default: auto-select)') +@click.option('--share', is_flag=True, help='Create a public Gradio share link') +def gui(host, port, share): + """Launch the PartiNet Gradio GUI.""" + from partinet.gui import launch_gui + launch_gui(host=host, port=port, share=share) + @main.group() def train(): diff --git a/partinet/gui/__init__.py b/partinet/gui/__init__.py new file mode 100644 index 0000000..48d5cd4 --- /dev/null +++ b/partinet/gui/__init__.py @@ -0,0 +1,3 @@ +from partinet.gui.app import launch_gui + +__all__ = ["launch_gui"] diff --git a/partinet/gui/app.py b/partinet/gui/app.py new file mode 100644 index 0000000..730b7ea --- /dev/null +++ b/partinet/gui/app.py @@ -0,0 +1,779 @@ +import os + +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import gradio as gr + +from partinet.gui.job_registry import cancel_job, jobs_dropdown_update, jobs_markdown, read_job_log +from partinet.gui.job_runner import JobSpec, SlurmOptions, slurm_field_defaults, stream_job +from partinet.process_utils.image_io import is_micrograph_file, micrograph_dimensions, load_micrograph_for_detect + +# ── Branding ────────────────────────────────────────────────────────────────── + +_THEME_NAME = "lone17/kotaemon" + +_DEFAULT_DARK_HEAD = """ + +""" + + +def _resolve_theme(): + try: + return gr.themes.Base.from_hub(_THEME_NAME) + except Exception: + return gr.themes.Default() + + +_THEME = _resolve_theme() + +_HEADER_HTML = """ +
+
PartiNet
+
+ Automated cryo-EM particle picker  ·  + Run each stage in order: 1 · Denoise → 2 · Detect → 3 · Star File +
+
+""" + + +def _env_default(name: str) -> str: + return os.environ.get(name, "").strip() + + +def _slurm_from_ui(mode, partition, account, time_limit, cpus, gpus, mem, extra, preamble, partinet_cmd): + return SlurmOptions( + partition=(partition or "").strip(), + account=(account or "").strip(), + time_limit=(time_limit or "").strip(), + cpus=str(cpus).strip() if cpus not in (None, "") else "", + gpus=str(gpus).strip() if gpus not in (None, "") else "", + mem=(mem or "").strip(), + extra_sbatch=(extra or "").strip(), + preamble=(preamble or "").strip(), + partinet_cmd=(partinet_cmd or "partinet").strip() or "partinet", + ), _normalize_mode(mode) + + +def _normalize_mode(mode: str) -> str: + m = (mode or "Local").strip().lower() + return "slurm" if m == "slurm" else "local" + + +def _project_log(project: str, name: str) -> str: + return os.path.join(project.strip(), name) + + +def _run_job(spec: JobSpec): + yield from stream_job(spec) + + +# ── Stage runners ──────────────────────────────────────────────────────────── + +def run_denoise(source, project, img_format, num_workers, mode, partition, account, time_limit, cpus, gpus, mem, extra, preamble, partinet_cmd): + source, project = source.strip(), project.strip() + if not source: + yield "ERROR: Raw micrographs directory is required." + return + if not project: + yield "ERROR: Project directory is required." + return + + os.makedirs(project, exist_ok=True) + cmd = ["partinet", "denoise", "--source", source, "--project", project, "--img_format", img_format] + if num_workers not in (None, ""): + cmd.extend(["--num_workers", str(int(num_workers))]) + + slurm, job_mode = _slurm_from_ui(mode, partition, account, time_limit, cpus, gpus, mem, extra, preamble, partinet_cmd) + spec = JobSpec( + stage="denoise", + command=cmd, + project_dir=project, + log_path=_project_log(project, "partinet_denoise.log"), + mode=job_mode, + slurm=slurm, + ) + yield from _run_job(spec) + + +def run_detect(weight, source, project, conf_thres, iou_thres, device, img_size, dy_thres, mode, partition, account, time_limit, cpus, gpus, mem, extra, preamble, partinet_cmd): + weight, source, project = weight.strip(), source.strip(), project.strip() + if not weight: + yield "ERROR: Model weights path is required." + return + if not source: + yield "ERROR: Source images directory is required." + return + if not project: + yield "ERROR: Project directory is required." + return + + cmd = [ + "partinet", "detect", + "--weight", weight, + "--source", source, + "--project", project, + "--conf-thres", str(conf_thres), + "--iou-thres", str(iou_thres), + "--img-size", str(int(img_size)), + "--dy-thres", str(dy_thres), + "--exist-ok", + ] + if device.strip(): + cmd.extend(["--device", device.strip()]) + + slurm, job_mode = _slurm_from_ui(mode, partition, account, time_limit, cpus, gpus, mem, extra, preamble, partinet_cmd) + spec = JobSpec( + stage="detect", + command=cmd, + project_dir=project, + log_path=_project_log(project, "partinet_detect.log"), + mode=job_mode, + slurm=slurm, + ) + yield from _run_job(spec) + + +def run_star(labels, images, output, conf, relion, relion_project_dir, mrc_prefix, mode, partition, account, time_limit, cpus, gpus, mem, extra, preamble, partinet_cmd): + labels, images, output = labels.strip(), images.strip(), output.strip() + relion_project_dir = (relion_project_dir or "").strip() + mrc_prefix = (mrc_prefix or "").strip() + + if not labels: + yield "ERROR: Labels directory is required." + return + if not images: + yield "ERROR: Images directory is required." + return + if not output: + yield "ERROR: Output STAR file path is required." + return + if relion and not relion_project_dir: + yield "ERROR: RELION project directory is required when RELION output is enabled." + return + + project = os.path.dirname(os.path.abspath(output)) + cmd = [ + "partinet", "star", + "--labels", labels, + "--images", images, + "--output", output, + "--conf", str(conf), + ] + if relion: + cmd.extend(["--relion", "--relion-project-dir", relion_project_dir]) + if mrc_prefix: + cmd.extend(["--mrc-prefix", mrc_prefix]) + + slurm, job_mode = _slurm_from_ui(mode, partition, account, time_limit, cpus, gpus, mem, extra, preamble, partinet_cmd) + spec = JobSpec( + stage="star", + command=cmd, + project_dir=project, + log_path=_project_log(project, "partinet_star.log"), + mode=job_mode, + slurm=slurm, + ) + yield from _run_job(spec) + + +# ── Star File analysis ─────────────────────────────────────────────────────── + +_IMG_EXTS = (".png", ".jpg", ".jpeg", ".tif", ".tiff", ".mrc") + + +def _parse_label_file(path): + dets = [] + with open(path) as f: + for line in f: + parts = line.strip().split() + if len(parts) >= 5: + dets.append({ + "x": float(parts[1]), "y": float(parts[2]), + "w": float(parts[3]), "h": float(parts[4]), + "conf": float(parts[5]) if len(parts) >= 6 else 1.0, + }) + return dets + + +def _img_size(path): + try: + return micrograph_dimensions(path) + except (ValueError, OSError): + return 4096, 4096 + + +def load_detections(labels_dir, images_dir): + labels_dir, images_dir = labels_dir.strip(), images_dir.strip() + _err = lambda msg: (None, msg, None, None, None, None, gr.update(choices=[]), "") + + if not labels_dir or not images_dir: + return _err("Both directories are required.") + if not os.path.isdir(labels_dir): + return _err(f"Labels directory not found: `{labels_dir}`") + if not os.path.isdir(images_dir): + return _err(f"Images directory not found: `{images_dir}`") + + label_files = sorted(f for f in os.listdir(labels_dir) if f.endswith(".txt")) + if not label_files: + return _err("No .txt label files found in labels directory.") + + mics, all_confs, all_sizes = [], [], [] + for lf in label_files: + stem = os.path.splitext(lf)[0] + img_path = next( + (os.path.join(images_dir, stem + ext) for ext in _IMG_EXTS + if os.path.exists(os.path.join(images_dir, stem + ext))), + None, + ) + if img_path is None: + continue + w, h = _img_size(img_path) + dets = _parse_label_file(os.path.join(labels_dir, lf)) + for d in dets: + all_confs.append(d["conf"]) + all_sizes.append(max(d["w"] * w, d["h"] * h)) + mics.append({"name": stem, "img": img_path, "w": w, "h": h, "dets": dets}) + + if not mics: + return (None, "No label files had matching images in the images directory.", None, None, None, None, gr.update(choices=[]), "") + + state = {"mics": mics, "confs": all_confs, "sizes": all_sizes} + n, m = len(all_confs), len(mics) + counts = sorted(len(mic["dets"]) for mic in mics) + summary = ( + f"**{n:,} detections across {m:,} micrographs** — " + f"mean {n/m:.0f} · median {counts[m//2]} · " + f"range {counts[0]}–{counts[-1]} particles/micrograph" + ) + + DEFAULT = 0.1 + choices = [mic["name"] for mic in mics] + return ( + state, summary, + _conf_plot(all_confs, DEFAULT), + _size_plot(all_sizes), + _mic_count_plot(mics, DEFAULT), + _bivariate_plot(all_confs, all_sizes), + gr.update(choices=choices, value=choices[0]), + _retained_text(all_confs, DEFAULT), + ) + + +def _conf_plot(confs, threshold): + with plt.style.context("dark_background"): + fig, ax = plt.subplots(figsize=(5, 3)) + ax.hist(confs, bins=50, color="steelblue", edgecolor="white", linewidth=0.5) + ax.axvline(threshold, color="crimson", linestyle="--", linewidth=1.5, label=f"Threshold {threshold:.2f}") + ax.legend(fontsize=8) + ax.set_xlabel("Confidence score") + ax.set_ylabel("Detections") + ax.set_title("Confidence distribution") + fig.tight_layout() + return fig + + +def _size_plot(sizes_px): + with plt.style.context("dark_background"): + fig, ax = plt.subplots(figsize=(5, 3)) + ax.hist(sizes_px, bins=50, color="darkorange", edgecolor="white", linewidth=0.5) + ax.set_xlabel("Box size (px)") + ax.set_ylabel("Detections") + ax.set_title("Box size distribution") + fig.tight_layout() + return fig + + +def _mic_count_plot(mics, threshold): + with plt.style.context("dark_background"): + counts = [sum(1 for d in m["dets"] if d["conf"] >= threshold) for m in mics] + fig, ax = plt.subplots(figsize=(9, 3)) + ax.bar(range(len(counts)), counts, color="mediumseagreen", width=1.0, linewidth=0) + ax.set_xlabel("Micrograph index") + ax.set_ylabel("Particles") + ax.set_title(f"Particles per micrograph (threshold = {threshold:.2f})") + fig.tight_layout() + return fig + + +def _bivariate_plot(confs, sizes): + with plt.style.context("dark_background"): + fig, ax = plt.subplots(figsize=(5, 3)) + h = ax.hist2d(confs, sizes, bins=50, cmap="viridis") + fig.colorbar(h[3], ax=ax, label="Detections") + ax.set_xlabel("Confidence score") + ax.set_ylabel("Box size (px)") + ax.set_title("Confidence vs box size") + fig.tight_layout() + return fig + + +def _retained_text(confs, threshold): + total = len(confs) + kept = sum(1 for c in confs if c >= threshold) + pct = 100 * kept / total if total > 0 else 0 + return f"**{kept:,} / {total:,} particles retained** ({pct:.1f}%)" + + +def _draw_detections(mic, threshold, max_px=1000): + from PIL import Image as _PIL, ImageDraw + import cv2 + if mic["img"].lower().endswith(".mrc"): + bgr = load_micrograph_for_detect(mic["img"]) + img = _PIL.fromarray(cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)) + else: + img = _PIL.open(mic["img"]).convert("RGB") + w, h = img.size + scale = min(1.0, max_px / max(w, h)) + dw, dh = int(w * scale), int(h * scale) + if scale < 1.0: + img = img.resize((dw, dh), _PIL.LANCZOS) + draw = ImageDraw.Draw(img) + n_shown = 0 + for d in mic["dets"]: + if d["conf"] < threshold: + continue + cx, cy = d["x"] * dw, d["y"] * dh + bw, bh = d["w"] * dw, d["h"] * dh + x1, y1 = int(cx - bw / 2), int(cy - bh / 2) + x2, y2 = int(cx + bw / 2), int(cy + bh / 2) + c = d["conf"] + color = (int((1 - c) * 220), int(c * 220), 60) + draw.rectangle([x1, y1, x2, y2], outline=color, width=2) + n_shown += 1 + return img, n_shown + + +def update_threshold(state, threshold, mic_name): + if state is None: + return "", None, None, None, "" + retained = _retained_text(state["confs"], threshold) + conf_fig = _conf_plot(state["confs"], threshold) + mic_count_fig = _mic_count_plot(state["mics"], threshold) + img, mic_stats = None, "" + if mic_name: + mic = next((m for m in state["mics"] if m["name"] == mic_name), None) + if mic: + img, n = _draw_detections(mic, threshold) + mic_stats = f"**{n}** particles shown" + return retained, conf_fig, mic_count_fig, img, mic_stats + + +def update_micrograph(state, mic_name, threshold): + if state is None or not mic_name: + return None, "" + mic = next((m for m in state["mics"] if m["name"] == mic_name), None) + if mic is None: + return None, "" + img, n = _draw_detections(mic, threshold) + return img, f"**{n}** particles shown" + + +# ── Helpers for global project directory ───────────────────────────────────── + +def _find_labels_dirs(project_dir): + import glob as _glob + p = (project_dir or "").strip() + if not p or not os.path.isdir(p): + return [] + exp_dirs = sorted( + [d for d in _glob.glob(os.path.join(p, "exp*")) if os.path.isdir(d)], + key=os.path.getmtime, + reverse=True, + ) + return [ + os.path.join(d, "labels") + for d in exp_dirs + if os.path.isdir(os.path.join(d, "labels")) + ] + + +def refresh_labels(project_dir): + dirs = _find_labels_dirs(project_dir) + return dirs[0] if dirs else "" + + +def update_project_dir(project_dir): + p = (project_dir or "").strip() + if not p: + return ("", "", "", "", "", "", jobs_markdown(""), jobs_dropdown_update(""), "") + denoised = os.path.join(p, "denoised") + dirs = _find_labels_dirs(p) + labels_val = dirs[0] if dirs else os.path.join(p, "exp", "labels") + return ( + p, # d1_project + denoised, # d2_source + p, # d2_project + labels_val, # d3_labels + denoised, # d3_images + os.path.join(p, "particles.star"), # d3_output + jobs_markdown(p), + jobs_dropdown_update(p), + "", + ) + + +def refresh_jobs(project_dir, job_key=None): + log = read_job_log(project_dir, job_key or "") + return jobs_markdown(project_dir), jobs_dropdown_update(project_dir), log + + +def view_job_log(project_dir, job_key): + return read_job_log(project_dir, job_key or "") + + +def cancel_selected_job(project_dir, job_key): + message = cancel_job(job_key or "", project_dir) + markdown, dropdown, log = refresh_jobs(project_dir, job_key) + return markdown, dropdown, log, message + + +# ── Gradio UI ──────────────────────────────────────────────────────────────── + +def build_app(): + with gr.Blocks(title="PartiNet") as app: + gr.HTML(_HEADER_HTML) + + gr_project = gr.Textbox( + label="Project directory", + value=_env_default("PARTINET_PROJECT"), + placeholder="/path/to/my_project", + info="Set once — auto-fills project paths in all three stages below", + ) + + with gr.Accordion("Running jobs", open=False): + jobs_markdown_out = gr.Markdown("Set a project directory to view jobs.") + job_select = gr.Dropdown( + label="Select job", + choices=[], + interactive=True, + info="Select a running job to view its log below", + ) + job_log_view = gr.Textbox( + label="Job log", + lines=16, + max_lines=40, + interactive=False, + ) + with gr.Row(): + jobs_refresh_btn = gr.Button("Refresh", variant="secondary") + jobs_cancel_btn = gr.Button("Cancel selected", variant="stop") + jobs_cancel_status = gr.Markdown("") + jobs_timer = gr.Timer(5) + + with gr.Accordion("Execution settings (Local / Slurm)", open=False): + gr.Markdown( + "Run stages on this machine (**Local**) or submit batch jobs (**Slurm**). " + "Resource defaults update per stage when you switch tabs " + "(denoise/detect: 32 CPUs, 100G RAM; detect adds 4 GPUs; star: 16 CPUs, 64G RAM). " + "Leave fields blank to use those defaults. " + "Optional cluster-wide overrides: set `PARTINET_SLURM_CONFIG` to a YAML path." + ) + exec_mode = gr.Radio(["Local", "Slurm"], value="Local", label="Execution mode") + with gr.Row(): + slurm_partition = gr.Textbox(label="Partition", placeholder="") + slurm_account = gr.Textbox(label="Account", placeholder="") + slurm_time = gr.Textbox(label="Time limit", placeholder="HH:MM:SS") + with gr.Row(): + slurm_cpus = gr.Textbox(label="CPUs per task", value="32", placeholder="32") + slurm_gpus = gr.Textbox(label="GPUs (gres count)", value="", placeholder="") + slurm_mem = gr.Textbox(label="Memory", value="100G", placeholder="100G") + slurm_extra = gr.Textbox( + label="Extra #SBATCH lines", + placeholder="#SBATCH --constraint=...", + lines=2, + ) + slurm_preamble = gr.Textbox( + label="Job setup script", + placeholder="# module load ...\\n# source activate ...", + lines=3, + ) + slurm_partinet_cmd = gr.Textbox( + label="PartiNet executable", + value="partinet", + placeholder="partinet", + ) + + slurm_inputs = [ + exec_mode, slurm_partition, slurm_account, slurm_time, + slurm_cpus, slurm_gpus, slurm_mem, slurm_extra, slurm_preamble, slurm_partinet_cmd, + ] + + with gr.Tabs(): + + # ── 1. Denoise ─────────────────────────────────────────────────── + with gr.Tab("1 · Denoise") as tab_denoise: + gr.Markdown( + "Improve signal-to-noise in raw micrographs using a Wiener filter. " + "Output images are saved to `project/denoised/`." + ) + with gr.Row(): + d1_source = gr.Textbox( + label="Raw micrographs directory", + placeholder="/path/to/motion_corrected", + value="", + info="Folder of .mrc files from RELION or CryoSPARC motion correction", + ) + d1_project = gr.Textbox( + label="Project directory", + placeholder="/path/to/my_project", + value=_env_default("PARTINET_PROJECT"), + info="All PartiNet outputs for this dataset will be written here", + ) + with gr.Row(): + d1_fmt = gr.Dropdown( + choices=["png", "jpg", "mrc"], + value="png", + label="Output image format", + info="PNG recommended — lossless and directly compatible with Detect", + ) + d1_workers = gr.Number( + label="CPU workers (blank = auto)", + value=None, + precision=0, + minimum=1, + info="Parallel workers for denoising; auto uses half the available CPUs", + ) + d1_btn = gr.Button("▶ Run Denoise", variant="primary") + d1_log = gr.Textbox( + label="Log output", + lines=14, + max_lines=30, + interactive=False, + ) + d1_evt = d1_btn.click( + run_denoise, + inputs=[d1_source, d1_project, d1_fmt, d1_workers] + slurm_inputs, + outputs=d1_log, + ) + + # ── 2. Detect ──────────────────────────────────────────────────── + with gr.Tab("2 · Detect") as tab_detect: + gr.Markdown( + "Locate particles in denoised micrographs using the DynamicDet model. " + "Results are saved to `project/exp/`." + ) + with gr.Row(): + d2_weight = gr.Textbox( + label="Model weights (.pt)", + placeholder="/path/to/model.pt", + value=_env_default("PARTINET_WEIGHTS"), + info="Pre-trained weights — download from HuggingFace or use your own trained model", + ) + d2_source = gr.Textbox( + label="Denoised images directory", + placeholder="/path/to/my_project/denoised", + value="", + info="Output of the Denoise step (project/denoised/)", + ) + with gr.Row(): + d2_project = gr.Textbox( + label="Project directory", + placeholder="/path/to/my_project", + value=_env_default("PARTINET_PROJECT"), + info="Same project directory used in Denoise", + ) + d2_device = gr.Textbox( + label="GPU device(s)", + value="", + placeholder="0 or 0,1,2,3 or cpu", + info="Leave blank to auto-detect; use 'cpu' if no GPU is available", + ) + with gr.Row(): + d2_conf = gr.Slider( + 0.0, 1.0, value=0.1, step=0.01, + label="Confidence threshold", + info="Lower → more picks (including false positives). Recommended: 0.0–0.3", + ) + d2_iou = gr.Slider( + 0.0, 1.0, value=0.2, step=0.01, + label="IOU threshold", + info="Controls removal of overlapping detections. Default 0.2 works for most datasets", + ) + with gr.Accordion("Advanced options", open=False): + with gr.Row(): + d2_imgsize = gr.Number( + label="Inference image size (px)", + value=1280, + precision=0, + info="Must match the size used during training (default: 1280)", + ) + d2_dy = gr.Slider( + 0.0, 1.0, value=0.5, step=0.01, + label="Dynamic threshold", + info="Router threshold between easy/hard micrograph detectors", + ) + d2_btn = gr.Button("▶ Run Detect", variant="primary") + d2_log = gr.Textbox( + label="Log output", + lines=14, + max_lines=30, + interactive=False, + ) + d2_evt = d2_btn.click( + run_detect, + inputs=[ + d2_weight, d2_source, d2_project, d2_conf, d2_iou, d2_device, d2_imgsize, d2_dy, + ] + slurm_inputs, + outputs=d2_log, + ) + + # ── 3. Star File ───────────────────────────────────────────────── + with gr.Tab("3 · Star File") as tab_star: + gr.Markdown( + "Load detections from a Detect run, explore statistics, set a confidence " + "threshold interactively, then generate a STAR file for **CryoSPARC** or **RELION**." + ) + + # --- Load --- + with gr.Row(): + d3_labels = gr.Textbox( + label="Labels directory", + placeholder="/path/to/project/exp/labels", + value="", + info="Folder of .txt detection files from the Detect step", + ) + d3_images = gr.Textbox( + label="Denoised images directory", + placeholder="/path/to/project/denoised", + value="", + info="Used to resolve image dimensions for coordinate conversion", + ) + d3_state = gr.State(None) + d3_load_btn = gr.Button("Load Detections", variant="secondary") + d3_summary = gr.Markdown("") + + # --- Statistics --- + with gr.Row(): + d3_conf_plot = gr.Plot(format="png", label="Confidence distribution") + d3_size_plot = gr.Plot(format="png", label="Box size distribution") + with gr.Row(): + d3_mic_count_plot = gr.Plot(format="png", label="Particles per micrograph") + d3_bivariate_plot = gr.Plot(format="png", label="Confidence vs box size") + + # --- Threshold + preview --- + d3_thresh = gr.Slider( + 0.0, 1.0, value=0.1, step=0.01, + label="Confidence threshold", + info="Adjust to filter detections — updates plots and image preview in real time", + ) + d3_retained = gr.Markdown("") + d3_mic_select = gr.Dropdown(label="Preview micrograph", choices=[], interactive=True) + d3_mic_stats = gr.Markdown("") + d3_preview = gr.Image(label="Detection preview", type="pil") + + # --- STAR generation --- + gr.Markdown("---\n### Generate STAR File\nThe confidence threshold above is applied.") + d3_output = gr.Textbox( + label="Output STAR file", + placeholder="/path/to/project/particles.star", + value="", + info="CryoSPARC-compatible STAR file will be written here", + ) + with gr.Accordion("RELION output (optional)", open=False): + d3_relion = gr.Checkbox(label="Also generate RELION-format output", value=False) + d3_relion_dir = gr.Textbox( + label="RELION project directory", + placeholder="/path/to/relion_project", + info="Creates /partinet/pick.star and per-micrograph coordinate files", + ) + d3_mrc_prefix = gr.Textbox( + label="MRC path prefix", + placeholder="MotionCorr/job003/movies", + value="", + info="Prepended to micrograph names in the RELION STAR file", + ) + d3_star_btn = gr.Button("▶ Generate STAR File", variant="primary") + d3_log = gr.Textbox( + label="Log output", + lines=8, + max_lines=20, + interactive=False, + ) + + # --- Events --- + d3_load_btn.click( + load_detections, + inputs=[d3_labels, d3_images], + outputs=[d3_state, d3_summary, d3_conf_plot, d3_size_plot, d3_mic_count_plot, d3_bivariate_plot, d3_mic_select, d3_retained], + ) + d3_thresh.change( + update_threshold, + inputs=[d3_state, d3_thresh, d3_mic_select], + outputs=[d3_retained, d3_conf_plot, d3_mic_count_plot, d3_preview, d3_mic_stats], + ) + d3_mic_select.change( + update_micrograph, + inputs=[d3_state, d3_mic_select, d3_thresh], + outputs=[d3_preview, d3_mic_stats], + ) + d3_evt = d3_star_btn.click( + run_star, + inputs=[ + d3_labels, d3_images, d3_output, d3_thresh, d3_relion, d3_relion_dir, d3_mrc_prefix, + ] + slurm_inputs, + outputs=d3_log, + ) + + jobs_outputs = [jobs_markdown_out, job_select, job_log_view] + jobs_refresh_btn.click(refresh_jobs, inputs=[gr_project, job_select], outputs=jobs_outputs) + jobs_timer.tick(refresh_jobs, inputs=[gr_project, job_select], outputs=jobs_outputs) + job_select.change(view_job_log, inputs=[gr_project, job_select], outputs=job_log_view) + jobs_cancel_btn.click( + cancel_selected_job, + inputs=[gr_project, job_select], + outputs=[jobs_markdown_out, job_select, job_log_view, jobs_cancel_status], + cancels=[d1_evt, d2_evt, d3_evt], + ) + + gr_project.change( + update_project_dir, + inputs=[gr_project], + outputs=[ + d1_project, d2_source, d2_project, d3_labels, d3_images, d3_output, + jobs_markdown_out, job_select, job_log_view, + ], + ) + + slurm_resource_outputs = [slurm_cpus, slurm_gpus, slurm_mem] + tab_denoise.select(lambda: slurm_field_defaults("denoise"), outputs=slurm_resource_outputs) + tab_detect.select(lambda: slurm_field_defaults("detect"), outputs=slurm_resource_outputs) + tab_star.select(lambda: slurm_field_defaults("star"), outputs=slurm_resource_outputs) + + return app + + +def launch_gui(host="127.0.0.1", port=None, share=False): + import click + + app = build_app() + app.queue() + browse_host = "127.0.0.1" if host in ("0.0.0.0", "::") else host + if port is not None: + click.echo(f"PartiNet GUI — open http://{browse_host}:{port} on this machine.") + click.echo("Over SSH, forward the port from your laptop, then browse to localhost:") + click.echo(f" ssh -L {port}:127.0.0.1:{port} user@login-node") + click.echo(f" http://localhost:{port}") + else: + click.echo(f"PartiNet GUI — open http://{browse_host}: on this machine.") + click.echo( + "Over SSH, forward the port Gradio prints below, then open http://localhost: " + "in your local browser." + ) + app.launch( + server_name=host, + server_port=port, + share=share, + theme=_THEME, + head=_DEFAULT_DARK_HEAD, + ssr_mode=False, + ) diff --git a/partinet/gui/job_registry.py b/partinet/gui/job_registry.py new file mode 100644 index 0000000..916f241 --- /dev/null +++ b/partinet/gui/job_registry.py @@ -0,0 +1,341 @@ +"""Track and cancel GUI-submitted PartiNet jobs.""" + +from __future__ import annotations + +import json +import os +import signal +import subprocess +import threading +import uuid +from dataclasses import asdict, dataclass, field +from datetime import datetime, timezone +from typing import TYPE_CHECKING, Dict, List, Optional + +if TYPE_CHECKING: + from partinet.gui.job_runner import JobSpec + + +_ACTIVE = {"running", "pending"} + + +@dataclass +class JobRecord: + job_key: str + stage: str + mode: str + project_dir: str + log_path: str + script_path: str = "" + slurm_job_id: str = "" + pid: Optional[int] = None + status: str = "running" + submitted_at: str = "" + + @classmethod + def from_dict(cls, data: dict) -> "JobRecord": + pid = data.get("pid") + return cls( + job_key=str(data["job_key"]), + stage=str(data["stage"]), + mode=str(data["mode"]), + project_dir=str(data["project_dir"]), + log_path=str(data["log_path"]), + script_path=str(data.get("script_path") or ""), + slurm_job_id=str(data.get("slurm_job_id") or ""), + pid=int(pid) if pid not in (None, "") else None, + status=str(data.get("status") or "running"), + submitted_at=str(data.get("submitted_at") or ""), + ) + + def to_dict(self) -> dict: + return asdict(self) + + def display_id(self) -> str: + if self.mode == "slurm" and self.slurm_job_id: + return self.slurm_job_id + if self.pid is not None: + return str(self.pid) + return self.job_key[:8] + + +@dataclass +class _RegistryState: + lock: threading.Lock = field(default_factory=threading.Lock) + cancel_flags: Dict[str, bool] = field(default_factory=dict) + local_procs: Dict[str, subprocess.Popen] = field(default_factory=dict) + project_by_key: Dict[str, str] = field(default_factory=dict) + + +_STATE = _RegistryState() + + +def _jobs_dir(project_dir: str) -> str: + path = os.path.join(project_dir, ".partinet_jobs") + os.makedirs(path, exist_ok=True) + return path + + +def _manifest_path(project_dir: str) -> str: + return os.path.join(_jobs_dir(project_dir), "manifest.json") + + +def _utc_now() -> str: + return datetime.now(timezone.utc).replace(microsecond=0).isoformat() + + +def _pid_alive(pid: Optional[int]) -> bool: + if pid is None: + return False + try: + os.kill(pid, 0) + except OSError: + return False + return True + + +def _slurm_active(job_id: str) -> bool: + if not job_id: + return False + try: + proc = subprocess.run( + ["squeue", "-h", "-j", job_id], + capture_output=True, + text=True, + check=False, + ) + return bool(proc.stdout.strip()) + except FileNotFoundError: + return False + + +def _load_manifest(project_dir: str) -> Dict[str, JobRecord]: + path = _manifest_path(project_dir) + if not os.path.isfile(path): + return {} + try: + with open(path, "r", encoding="utf-8") as fh: + data = json.load(fh) + except (OSError, json.JSONDecodeError): + return {} + jobs = data.get("jobs") if isinstance(data, dict) else None + if not isinstance(jobs, dict): + return {} + return {key: JobRecord.from_dict(val) for key, val in jobs.items() if isinstance(val, dict)} + + +def _save_manifest(project_dir: str, records: Dict[str, JobRecord]) -> None: + path = _manifest_path(project_dir) + payload = {"jobs": {key: rec.to_dict() for key, rec in records.items()}} + tmp = path + ".tmp" + with open(tmp, "w", encoding="utf-8") as fh: + json.dump(payload, fh, indent=2) + fh.write("\n") + os.replace(tmp, path) + + +def _reconcile_record(record: JobRecord) -> JobRecord: + if record.status not in _ACTIVE: + return record + if record.mode == "slurm": + if record.slurm_job_id and _slurm_active(record.slurm_job_id): + return record + if record.slurm_job_id: + record.status = "completed" + return record + if record.pid is None: + return record + if _pid_alive(record.pid): + return record + record.status = "completed" + return record + + +def register_job(spec: "JobSpec", mode: str, script_path: str = "") -> JobRecord: + job_key = uuid.uuid4().hex + record = JobRecord( + job_key=job_key, + stage=spec.stage, + mode=mode, + project_dir=spec.project_dir, + log_path=spec.log_path, + script_path=script_path, + status="running", + submitted_at=_utc_now(), + ) + with _STATE.lock: + _STATE.cancel_flags[job_key] = False + _STATE.project_by_key[job_key] = spec.project_dir + records = _load_manifest(spec.project_dir) + records[job_key] = record + _save_manifest(spec.project_dir, records) + return record + + +def attach_local_pid(job_key: str, proc: subprocess.Popen) -> None: + with _STATE.lock: + _STATE.local_procs[job_key] = proc + project_dir = _STATE.project_by_key.get(job_key) + if not project_dir: + return + records = _load_manifest(project_dir) + if job_key in records: + records[job_key].pid = proc.pid + _save_manifest(project_dir, records) + + +def attach_slurm_id(job_key: str, slurm_job_id: str, script_path: str = "") -> None: + with _STATE.lock: + project_dir = _STATE.project_by_key.get(job_key) + if not project_dir: + return + records = _load_manifest(project_dir) + if job_key not in records: + return + records[job_key].slurm_job_id = slurm_job_id + if script_path: + records[job_key].script_path = script_path + _save_manifest(project_dir, records) + + +def is_cancelled(job_key: str) -> bool: + with _STATE.lock: + return _STATE.cancel_flags.get(job_key, False) + + +def complete_job(job_key: str, status: str) -> None: + with _STATE.lock: + _STATE.cancel_flags.pop(job_key, None) + _STATE.local_procs.pop(job_key, None) + project_dir = _STATE.project_by_key.pop(job_key, None) + if not project_dir: + return + records = _load_manifest(project_dir) + if job_key in records: + records[job_key].status = status + _save_manifest(project_dir, records) + + +def list_jobs(project_dir: str) -> List[JobRecord]: + p = (project_dir or "").strip() + if not p or not os.path.isdir(p): + return [] + with _STATE.lock: + records = _load_manifest(p) + updated: Dict[str, JobRecord] = {} + active: List[JobRecord] = [] + for key, rec in records.items(): + rec = _reconcile_record(rec) + updated[key] = rec + if rec.status in _ACTIVE: + active.append(rec) + _save_manifest(p, updated) + active.sort(key=lambda r: r.submitted_at, reverse=True) + return active + + +def cancel_job(job_key: str, project_dir: str = "") -> str: + if not job_key: + return "No job selected." + slurm_id = "" + proc = None + local_pid = None + with _STATE.lock: + _STATE.cancel_flags[job_key] = True + proc = _STATE.local_procs.get(job_key) + resolved_project = (project_dir or "").strip() or _STATE.project_by_key.get(job_key, "") + if resolved_project: + records = _load_manifest(resolved_project) + rec = records.get(job_key) + if rec: + slurm_id = rec.slurm_job_id + local_pid = rec.pid + rec.status = "cancelled" + _save_manifest(resolved_project, records) + + if proc is not None and proc.poll() is None: + proc.terminate() + try: + proc.wait(timeout=5) + except subprocess.TimeoutExpired: + proc.kill() + elif local_pid is not None: + try: + os.kill(local_pid, signal.SIGTERM) + except OSError: + pass + + if slurm_id: + try: + subprocess.run(["scancel", slurm_id], capture_output=True, text=True, check=False) + except FileNotFoundError: + return f"Cancelled locally; scancel not found for Slurm job {slurm_id}." + + return "Job cancellation requested." + + +def get_job_record(project_dir: str, job_key: str) -> Optional[JobRecord]: + p = (project_dir or "").strip() + if not p or not job_key: + return None + records = _load_manifest(p) + rec = records.get(job_key) + if rec is None: + return None + return _reconcile_record(rec) + + +def _read_tail(path: str, max_bytes: int = 512_000) -> str: + if not os.path.isfile(path): + return "" + with open(path, "rb") as fh: + fh.seek(0, os.SEEK_END) + size = fh.tell() + fh.seek(max(0, size - max_bytes)) + return fh.read().decode("utf-8", errors="replace") + + +def read_job_log(project_dir: str, job_key: str) -> str: + if not job_key: + return "Select a job to view its log." + rec = get_job_record(project_dir, job_key) + if rec is None: + return "Job not found." + parts = [f"**{rec.stage}** · {rec.mode} · {rec.display_id()} · {rec.status}"] + if rec.log_path: + chunk = _read_tail(rec.log_path) + if chunk: + parts.append(f"\n--- {rec.log_path} ---\n{chunk.rstrip()}") + if rec.script_path: + slurm_out = rec.script_path.replace(".sh", "_slurm.out") + chunk = _read_tail(slurm_out) + if chunk: + parts.append(f"\n--- {slurm_out} ---\n{chunk.rstrip()}") + if len(parts) == 1: + parts.append("\nNo log output yet.") + return "\n".join(parts) + + +def jobs_markdown(project_dir: str) -> str: + if not (project_dir or "").strip(): + return "Set a project directory to view jobs." + jobs = list_jobs(project_dir) + if not jobs: + return "No running jobs for this project." + lines = ["| Stage | Mode | ID | Status | Started |", "| --- | --- | --- | --- | --- |"] + for job in jobs: + lines.append( + f"| {job.stage} | {job.mode} | {job.display_id()} | {job.status} | {job.submitted_at} |" + ) + return "\n".join(lines) + + +def jobs_dropdown_update(project_dir: str): + import gradio as gr + + jobs = list_jobs(project_dir) + choice_list = [(f"{j.stage} · {j.mode} · {j.display_id()}", j.job_key) for j in jobs] + return gr.update( + choices=choice_list, + value=jobs[0].job_key if jobs else None, + ) diff --git a/partinet/gui/job_runner.py b/partinet/gui/job_runner.py new file mode 100644 index 0000000..8ba333e --- /dev/null +++ b/partinet/gui/job_runner.py @@ -0,0 +1,422 @@ +"""Local and Slurm job execution for the PartiNet GUI.""" + +from __future__ import annotations + +import logging +import os +import queue +import shlex +import subprocess +import sys +import threading +import time +import traceback +from dataclasses import dataclass, field +from datetime import datetime +from typing import Callable, Dict, Iterator, List, Optional + +import yaml + +from partinet.gui.job_registry import ( + attach_local_pid, + attach_slurm_id, + complete_job, + is_cancelled, + register_job, +) + +POLL_INTERVAL = 1.0 +_SLURM_TERMINAL = {"COMPLETED", "FAILED", "CANCELLED", "TIMEOUT", "NODE_FAIL", "OUT_OF_MEMORY", "PREEMPTED"} + + +@dataclass +class SlurmOptions: + partition: str = "" + account: str = "" + time_limit: str = "" + cpus: str = "" + gpus: str = "" + mem: str = "" + extra_sbatch: str = "" + preamble: str = "" + partinet_cmd: str = "partinet" + + +STAGE_SLURM_DEFAULTS: Dict[str, SlurmOptions] = { + "denoise": SlurmOptions(cpus="32", mem="100G"), + "detect": SlurmOptions(cpus="32", mem="100G", gpus="4"), + "star": SlurmOptions(cpus="16", mem="64G"), +} + + +@dataclass +class JobSpec: + stage: str + command: List[str] + project_dir: str + log_path: str + mode: str = "local" + slurm: Optional[SlurmOptions] = None + run_fn: Optional[Callable[[], None]] = None + loggers: List[str] = field(default_factory=list) + + +def load_slurm_defaults() -> SlurmOptions: + """Load optional user-owned YAML defaults from PARTINET_SLURM_CONFIG.""" + path = os.environ.get("PARTINET_SLURM_CONFIG", "") + if not path or not os.path.isfile(path): + return SlurmOptions() + with open(path, "r", encoding="utf-8") as fh: + data = yaml.safe_load(fh) or {} + if not isinstance(data, dict): + return SlurmOptions() + fields = {f.name for f in SlurmOptions.__dataclass_fields__.values()} # type: ignore[attr-defined] + return SlurmOptions(**{k: str(v) for k, v in data.items() if k in fields and v is not None}) + + +def merge_slurm_options(base: SlurmOptions, override: SlurmOptions) -> SlurmOptions: + """GUI fields override config-file defaults when non-empty.""" + merged = {} + for name in SlurmOptions.__dataclass_fields__: # type: ignore[attr-defined] + gui_val = getattr(override, name, "") + cfg_val = getattr(base, name, "") + merged[name] = gui_val if str(gui_val).strip() else cfg_val + return SlurmOptions(**merged) + + +def resolve_slurm_options(stage: str, ui: Optional[SlurmOptions]) -> SlurmOptions: + """Merge stage defaults, optional YAML config, then GUI overrides.""" + stage_base = STAGE_SLURM_DEFAULTS.get(stage, SlurmOptions()) + return merge_slurm_options( + merge_slurm_options(stage_base, load_slurm_defaults()), + ui or SlurmOptions(), + ) + + +def slurm_field_defaults(stage: str) -> tuple[str, str, str]: + opts = STAGE_SLURM_DEFAULTS.get(stage, SlurmOptions()) + return opts.cpus, opts.gpus, opts.mem + + +class _QueueHandler(logging.Handler): + def __init__(self, q: queue.Queue): + super().__init__() + self.q = q + + def emit(self, record): + self.q.put(self.format(record) + "\n") + + +def _tail_file(path: str, offset: int) -> tuple[str, int]: + if not os.path.exists(path): + return "", offset + with open(path, "r", encoding="utf-8", errors="replace") as fh: + fh.seek(offset) + chunk = fh.read() + return chunk, fh.tell() + + +def _stream_local_callable(spec: JobSpec, job_key: str) -> Iterator[str]: + fmt = logging.Formatter("%(asctime)s - %(message)s") + log_q: queue.Queue = queue.Queue() + handler = _QueueHandler(log_q) + handler.setFormatter(fmt) + + for name in spec.loggers: + lg = logging.getLogger(name) + lg.addHandler(handler) + if lg.level == logging.NOTSET or lg.level > logging.INFO: + lg.setLevel(logging.INFO) + + result = {"done": False, "error": None} + + def _target(): + try: + if spec.run_fn is None: + raise RuntimeError("Local job missing run_fn") + spec.run_fn() + except Exception: + result["error"] = traceback.format_exc() + finally: + result["done"] = True + + threading.Thread(target=_target, daemon=True).start() + output = "" + offset = 0 + while not result["done"]: + if is_cancelled(job_key): + output += "\n--- CANCELLED ---\n" + yield output + complete_job(job_key, "cancelled") + for name in spec.loggers: + logging.getLogger(name).removeHandler(handler) + return + time.sleep(POLL_INTERVAL) + while not log_q.empty(): + output += log_q.get_nowait() + chunk, offset = _tail_file(spec.log_path, offset) + if chunk: + output += chunk + yield output + + while not log_q.empty(): + output += log_q.get_nowait() + chunk, offset = _tail_file(spec.log_path, offset) + if chunk: + output += chunk + + for name in spec.loggers: + logging.getLogger(name).removeHandler(handler) + + if result["error"]: + output += f"\n--- FAILED ---\n{result['error']}" + complete_job(job_key, "failed") + else: + output += "\n--- Complete ---" + complete_job(job_key, "completed") + yield output + + +def _kill_local_proc(proc: subprocess.Popen) -> None: + if proc.poll() is not None: + return + proc.terminate() + try: + proc.wait(timeout=5) + except subprocess.TimeoutExpired: + proc.kill() + + +def _stream_local_subprocess(spec: JobSpec, job_key: str) -> Iterator[str]: + slurm = spec.slurm or SlurmOptions() + cmd = list(spec.command) + if slurm.partinet_cmd and slurm.partinet_cmd != "partinet": + cmd = [slurm.partinet_cmd] + cmd[1:] + + env = os.environ.copy() + proc = subprocess.Popen( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + env=env, + ) + attach_local_pid(job_key, proc) + output = f"Running: {shlex.join(cmd)}\n" + yield output + assert proc.stdout is not None + try: + while True: + if is_cancelled(job_key): + _kill_local_proc(proc) + output += "\n--- CANCELLED ---\n" + yield output + complete_job(job_key, "cancelled") + return + line = proc.stdout.readline() + if line: + output += line + yield output + elif proc.poll() is not None: + break + else: + time.sleep(0.05) + remaining = proc.stdout.read() + if remaining: + output += remaining + if is_cancelled(job_key): + output += "\n--- CANCELLED ---\n" + complete_job(job_key, "cancelled") + elif proc.returncode != 0: + output += f"\n--- FAILED --- (exit code {proc.returncode})\n" + complete_job(job_key, "failed") + else: + output += "\n--- Complete ---" + complete_job(job_key, "completed") + yield output + finally: + if proc.poll() is None: + _kill_local_proc(proc) + + +def _jobs_dir(project_dir: str) -> str: + path = os.path.join(project_dir, ".partinet_jobs") + os.makedirs(path, exist_ok=True) + return path + + +def _write_slurm_script(spec: JobSpec, slurm: SlurmOptions) -> str: + jobs_dir = _jobs_dir(spec.project_dir) + stamp = datetime.now().strftime("%Y%m%d_%H%M%S") + script_path = os.path.join(jobs_dir, f"{spec.stage}_{stamp}.sh") + slurm_out = os.path.join(jobs_dir, f"{spec.stage}_{stamp}_slurm.out") + + lines = ["#!/bin/bash", f"#SBATCH --job-name=partinet_{spec.stage}", f"#SBATCH --output={slurm_out}"] + if slurm.partition.strip(): + lines.append(f"#SBATCH --partition={slurm.partition.strip()}") + if slurm.account.strip(): + lines.append(f"#SBATCH --account={slurm.account.strip()}") + if slurm.time_limit.strip(): + lines.append(f"#SBATCH --time={slurm.time_limit.strip()}") + if slurm.cpus.strip(): + lines.append(f"#SBATCH --cpus-per-task={slurm.cpus.strip()}") + if slurm.gpus.strip(): + lines.append(f"#SBATCH --gres=gpu:{slurm.gpus.strip()}") + if slurm.mem.strip(): + lines.append(f"#SBATCH --mem={slurm.mem.strip()}") + for extra in slurm.extra_sbatch.splitlines(): + extra = extra.strip() + if extra: + lines.append(extra if extra.startswith("#SBATCH") else f"#SBATCH {extra}") + + lines.append("set -euo pipefail") + if slurm.preamble.strip(): + lines.append(slurm.preamble.strip()) + + cmd = list(spec.command) + if slurm.partinet_cmd and slurm.partinet_cmd != "partinet": + cmd[0] = slurm.partinet_cmd + lines.append(f"cd {shlex.quote(spec.project_dir)}") + lines.append(shlex.join(cmd)) + + with open(script_path, "w", encoding="utf-8") as fh: + fh.write("\n".join(lines) + "\n") + os.chmod(script_path, 0o750) + return script_path + + +def _slurm_state(job_id: str) -> str: + try: + proc = subprocess.run( + ["squeue", "-h", "-j", job_id, "-o", "%T"], + capture_output=True, + text=True, + check=False, + ) + state = proc.stdout.strip().splitlines() + if state and state[0]: + return state[0].strip() + except FileNotFoundError: + pass + try: + proc = subprocess.run( + ["sacct", "-j", job_id, "-n", "-X", "-o", "State"], + capture_output=True, + text=True, + check=False, + ) + for line in proc.stdout.splitlines(): + token = line.strip().split()[0] if line.strip() else "" + if token: + return token + except FileNotFoundError: + pass + return "UNKNOWN" + + +def _slurm_active(job_id: str) -> bool: + try: + proc = subprocess.run( + ["squeue", "-h", "-j", job_id], + capture_output=True, + text=True, + check=False, + ) + return bool(proc.stdout.strip()) + except FileNotFoundError: + return False + + +def _stream_slurm(spec: JobSpec, job_key: str) -> Iterator[str]: + slurm = resolve_slurm_options(spec.stage, spec.slurm) + script_path = _write_slurm_script(spec, slurm) + slurm_out = script_path.replace(".sh", "_slurm.out") + output = f"Submitting Slurm job\nScript: {script_path}\n" + + try: + proc = subprocess.run( + ["sbatch", "--parsable", script_path], + capture_output=True, + text=True, + check=True, + ) + except FileNotFoundError: + complete_job(job_key, "failed") + yield output + "\n--- FAILED ---\nsbatch not found on PATH\n" + return + except subprocess.CalledProcessError as exc: + complete_job(job_key, "failed") + yield output + f"\n--- FAILED ---\n{exc.stderr or exc.stdout}\n" + return + + job_id = proc.stdout.strip().split(";")[0].strip() + attach_slurm_id(job_key, job_id, script_path) + output += f"Job ID: {job_id}\n" + yield output + + log_offset = 0 + slurm_offset = 0 + while True: + if is_cancelled(job_key): + try: + subprocess.run(["scancel", job_id], capture_output=True, text=True, check=False) + except FileNotFoundError: + pass + chunk, log_offset = _tail_file(spec.log_path, log_offset) + if chunk: + output += chunk + chunk, slurm_offset = _tail_file(slurm_out, slurm_offset) + if chunk: + output += chunk + output += "\n--- CANCELLED ---\n" + yield output + complete_job(job_key, "cancelled") + return + + time.sleep(POLL_INTERVAL) + chunk, log_offset = _tail_file(spec.log_path, log_offset) + if chunk: + output += chunk + chunk, slurm_offset = _tail_file(slurm_out, slurm_offset) + if chunk: + output += chunk + if chunk: + yield output + + if _slurm_active(job_id): + continue + + state = _slurm_state(job_id) + if state in _SLURM_TERMINAL: + chunk, log_offset = _tail_file(spec.log_path, log_offset) + if chunk: + output += chunk + chunk, slurm_offset = _tail_file(slurm_out, slurm_offset) + if chunk: + output += chunk + if state == "COMPLETED": + output += "\n--- Complete ---" + complete_job(job_key, "completed") + elif state == "CANCELLED": + output += "\n--- CANCELLED ---\n" + complete_job(job_key, "cancelled") + else: + output += f"\n--- FAILED --- (Slurm state: {state})\n" + complete_job(job_key, "failed") + yield output + return + + +def stream_job(spec: JobSpec) -> Iterator[str]: + """Run a job locally or via Slurm and yield growing log text.""" + mode = (spec.mode or "local").strip().lower() + record = register_job(spec, mode) + job_key = record.job_key + if mode == "slurm": + yield from _stream_slurm(spec, job_key) + return + if spec.run_fn is not None: + yield from _stream_local_callable(spec, job_key) + return + yield from _stream_local_subprocess(spec, job_key) diff --git a/partinet/process_utils/image_io.py b/partinet/process_utils/image_io.py new file mode 100644 index 0000000..ae96b22 --- /dev/null +++ b/partinet/process_utils/image_io.py @@ -0,0 +1,57 @@ +"""Shared micrograph read helpers for detect, star export, and GUI preview.""" + +from __future__ import annotations + +import os +from typing import Tuple + +import cv2 +import mrcfile +import numpy as np + +from partinet.process_utils.guided_denoiser import transform + +_RASTER_EXTS = {".png", ".jpg", ".jpeg", ".tif", ".tiff", ".bmp", ".webp"} + + +def _mrc_array_shape(data: np.ndarray) -> Tuple[int, int]: + """Return (width, height) from a 2D or (1, H, W) MRC array.""" + if data.ndim == 2: + h, w = data.shape + elif data.ndim == 3 and data.shape[0] == 1: + _, h, w = data.shape + else: + raise ValueError(f"Unsupported MRC array shape: {data.shape}") + return int(w), int(h) + + +def micrograph_dimensions(path: str) -> Tuple[int, int]: + """Return pixel (width, height) for STAR coordinate scaling.""" + ext = os.path.splitext(path)[1].lower() + if ext == ".mrc": + with mrcfile.open(path, permissive=True) as mrc: + return _mrc_array_shape(mrc.data) + from PIL import Image as PILImage + with PILImage.open(path) as img: + return img.size + + +def load_micrograph_for_detect(path: str) -> np.ndarray: + """Load a micrograph as BGR uint8 using the same rules as detect LoadImages.""" + ext = os.path.splitext(path)[1].lower() + if ext == ".mrc": + img_mrc = mrcfile.read(path) + img_mrc = np.asarray(img_mrc, dtype=np.float32) + img0 = transform(img_mrc).astype(np.uint8) + if img0.ndim == 2: + img0 = cv2.cvtColor(img0, cv2.COLOR_GRAY2BGR) + return img0 + img0 = cv2.imread(path) + if img0 is None: + raise ValueError(f"Could not read image: {path}") + return img0 + + +def is_micrograph_file(filename: str) -> bool: + ext = os.path.splitext(filename)[1].lower() + return ext in _RASTER_EXTS or ext == ".mrc" diff --git a/partinet/process_utils/pooled_denoise_proc.py b/partinet/process_utils/pooled_denoise_proc.py index 6739e49..0f7d76d 100644 --- a/partinet/process_utils/pooled_denoise_proc.py +++ b/partinet/process_utils/pooled_denoise_proc.py @@ -7,8 +7,28 @@ import argparse import gc from concurrent.futures import ProcessPoolExecutor -from typing import List, Tuple +from typing import List, Tuple, Optional import mrcfile +import sys + +logger = logging.getLogger("partinet_denoise") + + +def _configure_denoise_logging(log_path: str) -> None: + """Attach file and stream handlers once per run; remove stale handlers first.""" + fmt = logging.Formatter("%(asctime)s - %(message)s") + logger.setLevel(logging.INFO) + logger.propagate = False + for h in logger.handlers[:]: + logger.removeHandler(h) + h.close() + fh = logging.FileHandler(log_path) + fh.setFormatter(fmt) + logger.addHandler(fh) + sh = logging.StreamHandler(sys.stdout) + sh.setFormatter(fmt) + logger.addHandler(sh) + # Function to perform CLAHE-based denoising @@ -33,15 +53,21 @@ def clahe_denoise(args: Tuple[str, str, str]) -> None: mrcfile.write(dest_path,data=denoised) else: cv2.imwrite(dest_path, denoised) - logging.info(f"Processed image {src_path} to dest. {dest_path}") + logger.info(f"Processed image {src_path} to dest. {dest_path}") del denoised gc.collect() except Exception as e: - logging.error(f"Failed to process {src_path}: {str(e)}") + logger.error(f"Failed to process {src_path}: {str(e)}") # Function to process all files in a directory -def process_directory(micrographs_dir: str, clahe_denoised_dir: str, max_workers: int, img_format: str) -> None: +def process_directory( + micrographs_dir: str, + clahe_denoised_dir: str, + max_workers: int, + img_format: str, + log_path: Optional[str] = None, +) -> None: """ Processes all `.mrc` files in the given directory using parallel workers for denoising. @@ -56,7 +82,7 @@ def process_directory(micrographs_dir: str, clahe_denoised_dir: str, max_workers - Existing denoised images are skipped. """ os.makedirs(clahe_denoised_dir, exist_ok=True) - logging.info(f"Directory ready: {clahe_denoised_dir}") + logger.info(f"Directory ready: {clahe_denoised_dir}") tasks: List[Tuple[str, str, str]] = [] # Iterate through files in the directory @@ -65,15 +91,21 @@ def process_directory(micrographs_dir: str, clahe_denoised_dir: str, max_workers src_path = os.path.join(micrographs_dir, file_name) dest_path = os.path.join(clahe_denoised_dir, file_name.replace(".mrc", "."+img_format)) if os.path.exists(dest_path): - logging.info(f"{dest_path} already exists!") + logger.info(f"{dest_path} already exists!") else: tasks.append((src_path, dest_path, img_format)) - # Parallel processing of tasks - with ProcessPoolExecutor(max_workers=max_workers) as executor: + if not tasks: + logger.info("No new micrographs to denoise.") + return + + pool_kwargs = {} + if log_path: + pool_kwargs = {"initializer": _configure_denoise_logging, "initargs": (log_path,)} + with ProcessPoolExecutor(max_workers=max_workers, **pool_kwargs) as executor: futures = [executor.submit(clahe_denoise, task) for task in tasks] - for future in futures: - future.result() + for future in futures: + future.result() gc.collect() # Main function @@ -92,14 +124,10 @@ def main(source_dir: str, project_dir: str, ncpu: int, img_format: str) -> None: - Logging is configured to write messages to a log file and the console. - Ensures at least one CPU is used for processing. """ - # Configure logging - logging.basicConfig(filename="partinet_denoise.log", level=logging.INFO, format='%(asctime)s - %(message)s') - logging.getLogger().addHandler(logging.StreamHandler()) - # Prepare output directory and log file denoise_dir = os.path.join(project_dir, "denoised") - logger_name = project_dir + "/partinet_denoise.log" - logging.basicConfig(filename=logger_name, level=logging.INFO, format='%(asctime)s - %(message)s') + log_path = os.path.join(project_dir, "partinet_denoise.log") + _configure_denoise_logging(log_path) # Determine number of available CPUs max_available_cpus = multiprocessing.cpu_count() @@ -111,12 +139,11 @@ def main(source_dir: str, project_dir: str, ncpu: int, img_format: str) -> None: else: ncpu = max_workers - logging.info(f"Using {ncpu} workers out of {max_available_cpus} available CPUs.") - logging.info(f"Processing raw micrographs in {source_dir}") - logging.info(f"Saving denoised micrographs in {denoise_dir}") + logger.info(f"Using {ncpu} workers out of {max_available_cpus} available CPUs.") + logger.info(f"Processing raw micrographs in {source_dir}") + logger.info(f"Saving denoised micrographs in {denoise_dir}") - # Process the directory - process_directory(source_dir, denoise_dir, ncpu, img_format) + process_directory(source_dir, denoise_dir, ncpu, img_format, log_path=log_path) # Command-line argument parsing diff --git a/partinet/process_utils/star_file.py b/partinet/process_utils/star_file.py index 008c58a..735cb5e 100644 --- a/partinet/process_utils/star_file.py +++ b/partinet/process_utils/star_file.py @@ -1,12 +1,35 @@ +import logging import math import os +import sys import pandas as pd import csv -import cv2 import argparse from typing import List, Dict, Tuple, Optional from multiprocessing import Pool, cpu_count +from partinet.process_utils.image_io import micrograph_dimensions + +logger = logging.getLogger("partinet.process_utils.star_file") + + +def _configure_star_logging(star_out_path: str) -> str: + project_dir = os.path.dirname(os.path.abspath(star_out_path)) or "." + log_path = os.path.join(project_dir, "partinet_star.log") + fmt = logging.Formatter("%(asctime)s - %(message)s") + logger.setLevel(logging.INFO) + logger.propagate = False + for h in logger.handlers[:]: + logger.removeHandler(h) + h.close() + fh = logging.FileHandler(log_path) + fh.setFormatter(fmt) + logger.addHandler(fh) + sh = logging.StreamHandler(sys.stdout) + sh.setFormatter(fmt) + logger.addHandler(sh) + return log_path + def yolo_to_starfile(yolo_coords: Dict[str, float], image_width: int, image_height: int, diameters: List[int]) -> Tuple[int, int, int]: x_center = math.ceil(yolo_coords["x_centre"] * image_width) y_center = math.ceil(yolo_coords["y_centre"] * image_height) @@ -34,10 +57,10 @@ def process_image(args_tuple) -> List[Tuple[str, int, int, int]]: # skip missing labels return [] - image = cv2.imread(os.path.join(images_path, image_file)) - if image is None: + try: + img_width, img_height = micrograph_dimensions(os.path.join(images_path, image_file)) + except (ValueError, OSError): return [] - img_width, img_height = image.shape[1], image.shape[0] custom_headers = ["class", "x_centre", "y_centre", "width", "height", "conf"] labels = pd.read_csv(label_file_path, header=None, names=custom_headers, sep=r"\s+") @@ -99,6 +122,8 @@ def relion_write(all_rows: List[Tuple[str, int, int, int]], pick_out: str, coord write_relion_pick_star(mapping, pick_out) def main(labels_path: str, images_path: str, star_out_path: str, conf_thresh: float, relion: bool = False, relion_project_dir: Optional[str] = None, relion_pick: Optional[str] = None, relion_coord_dir: Optional[str] = None, mrc_prefix: str = "") -> None: + _configure_star_logging(star_out_path) + logger.info(f"Generating STAR file from labels in {labels_path}") image_files = [f for f in os.listdir(images_path) if os.path.splitext(f)[1].lower() in [".mrc", ".tif", ".tiff", ".png", ".jpg", ".jpeg"]] args_list = [(img_file, labels_path, images_path, conf_thresh) for img_file in image_files] @@ -107,11 +132,11 @@ def main(labels_path: str, images_path: str, star_out_path: str, conf_thresh: fl all_rows = [row for result in results for row in result] if not all_rows: - print("No particle rows produced.") + logger.info("No particle rows produced.") return write_cryosparc_star(all_rows, star_out_path) - print(f"Wrote cryosparc-compatible star to: {star_out_path}") + logger.info(f"Wrote CryoSPARC-compatible STAR file to: {star_out_path}") if relion: if relion_project_dir is None: @@ -121,8 +146,8 @@ def main(labels_path: str, images_path: str, star_out_path: str, conf_thresh: fl relion_coorddir = os.path.join(relion_partinet, "movies") os.makedirs(relion_coorddir, exist_ok=True) relion_write(all_rows, relion_pickstar, relion_coorddir, mrc_prefix) - print(f"Wrote relion pick.star: {relion_pickstar}") - print(f"Wrote relion per-micrograph stars under: {relion_coorddir}") + logger.info(f"Wrote RELION pick.star: {relion_pickstar}") + logger.info(f"Wrote RELION per-micrograph stars under: {relion_coorddir}") def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Generate STAR files from YOLO labels (CryoSPARC and optional RELION)") diff --git a/pyproject.toml b/pyproject.toml index 574878d..4b75be5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,11 +40,17 @@ dependencies = [ "tensorboard", "scikit-learn", "mrcfile", + "gradio", +] + +[project.optional-dependencies] +dev = [ "pytest", ] [project.scripts] partinet = "partinet:main" +partinet-gui = "partinet.gui:launch_gui" [build-system] requires = ["setuptools", "Click"] @@ -55,7 +61,7 @@ where = ["."] include = ["partinet*"] [tool.setuptools.package-data] -"*" = ["dy-*.yaml", "hyp.*.yaml"] +"partinet" = ["dy-*.yaml", "hyp.*.yaml"] [tool.setuptools.dynamic] version = {attr = "partinet.__version__"} diff --git a/tests/test_datasets_mrc.py b/tests/test_datasets_mrc.py new file mode 100644 index 0000000..63d1868 --- /dev/null +++ b/tests/test_datasets_mrc.py @@ -0,0 +1,39 @@ +import sys, os +import numpy as np +import mrcfile +from pathlib import Path + +# ensure local `partinet` package is imported instead of any installed version +sys.path.insert(0, os.getcwd()) + +from partinet.DynamicDet.utils.datasets import LoadImages + + +def test_loadimages_reads_mrc(tmp_path): + """LoadImages should accept an uncompressed MRC and return a 3-channel image. + + The network expects BGR input; raw MRCs are single-channel so the loader + converts to BGR. This test exercises the path added for issue #45. + """ + arr = (np.random.rand(16, 16) * 255).astype(np.uint8) + fname = tmp_path / "micro.mrc" + with mrcfile.new(fname, overwrite=True) as mrc: + mrc.set_data(arr.astype(np.float32)) + + loader = LoadImages(str(tmp_path)) + assert len(loader) == 1 + path, img, img0, cap = next(iter(loader)) + # CHW, three channels + assert img.ndim == 3 and img.shape[0] == 3 + # original returned image should also be 3‑channel BGR + assert img0.ndim == 3 and img0.shape[2] == 3 + assert Path(path).suffix.lower() == ".mrc" + + +def test_save_filename_helper(): + """`_save_filename` should map .mrc inputs to .jpg output names.""" + from partinet.DynamicDet import detect + p1 = Path("foo.mrc") + p2 = Path("bar.png") + assert detect._save_filename(p1) == "foo.png" + assert detect._save_filename(p2) == "bar.png" diff --git a/tests/test_image_io.py b/tests/test_image_io.py new file mode 100644 index 0000000..a280013 --- /dev/null +++ b/tests/test_image_io.py @@ -0,0 +1,22 @@ +import numpy as np +import mrcfile +import pytest + +from partinet.process_utils.image_io import micrograph_dimensions, load_micrograph_for_detect + + +def test_micrograph_dimensions_mrc(tmp_path): + arr = np.random.rand(20, 30).astype(np.float32) + path = tmp_path / "mic.mrc" + with mrcfile.new(path, overwrite=True) as mrc: + mrc.set_data(arr) + assert micrograph_dimensions(str(path)) == (30, 20) + + +def test_load_micrograph_for_detect_mrc(tmp_path): + arr = (np.random.rand(16, 16) * 100).astype(np.float32) + path = tmp_path / "mic.mrc" + with mrcfile.new(path, overwrite=True) as mrc: + mrc.set_data(arr) + img = load_micrograph_for_detect(str(path)) + assert img.ndim == 3 and img.shape[2] == 3 diff --git a/tests/test_job_registry.py b/tests/test_job_registry.py new file mode 100644 index 0000000..005c638 --- /dev/null +++ b/tests/test_job_registry.py @@ -0,0 +1,85 @@ +import os +from unittest import mock + +import pytest + +from partinet.gui import job_registry +from partinet.gui.job_runner import JobSpec + +register_job = job_registry.register_job +list_jobs = job_registry.list_jobs +cancel_job = job_registry.cancel_job +attach_local_pid = job_registry.attach_local_pid +_manifest_path = job_registry._manifest_path + + +def _spec_for(project): + return JobSpec( + stage="denoise", + command=["partinet", "denoise", "--source", "/data", "--project", project], + project_dir=project, + log_path=os.path.join(project, "partinet_denoise.log"), + mode="local", + ) + + +def test_register_and_list_active_job(tmp_path): + project = str(tmp_path / "project") + os.makedirs(project) + rec = register_job(_spec_for(project), "local") + jobs = list_jobs(project) + assert len(jobs) == 1 + assert jobs[0].job_key == rec.job_key + assert jobs[0].stage == "denoise" + + +def test_manifest_round_trip(tmp_path): + project = str(tmp_path / "project") + os.makedirs(project) + rec = register_job(_spec_for(project), "slurm") + path = _manifest_path(project) + assert os.path.isfile(path) + jobs = list_jobs(project) + assert jobs[0].job_key == rec.job_key + + +def test_cancel_local_proc(tmp_path): + project = str(tmp_path / "project") + os.makedirs(project) + rec = register_job(_spec_for(project), "local") + proc = mock.Mock() + proc.poll.return_value = None + proc.pid = 4242 + attach_local_pid(rec.job_key, proc) + msg = cancel_job(rec.job_key, project) + assert "cancellation" in msg.lower() + proc.terminate.assert_called_once() + + +def test_cancel_slurm_calls_scancel(tmp_path): + project = str(tmp_path / "project") + os.makedirs(project) + rec = register_job(_spec_for(project), "slurm") + job_registry.attach_slurm_id(rec.job_key, "999001", "/tmp/job.sh") + with mock.patch.object(job_registry.subprocess, "run") as run: + cancel_job(rec.job_key, project) + assert any(call.args[0][:2] == ["scancel", "999001"] for call in run.call_args_list) + + +def test_list_jobs_empty_without_project(tmp_path): + assert list_jobs("") == [] + assert list_jobs(str(tmp_path / "missing")) == [] + + +def test_read_job_log(tmp_path): + project = str(tmp_path / "project") + os.makedirs(project) + log_path = os.path.join(project, "partinet_denoise.log") + with open(log_path, "w", encoding="utf-8") as fh: + fh.write("denoise started\n") + rec = register_job(_spec_for(project), "local") + records = job_registry._load_manifest(project) + records[rec.job_key].log_path = log_path + job_registry._save_manifest(project, records) + text = job_registry.read_job_log(project, rec.job_key) + assert "denoise started" in text diff --git a/tests/test_job_runner.py b/tests/test_job_runner.py new file mode 100644 index 0000000..eb3610b --- /dev/null +++ b/tests/test_job_runner.py @@ -0,0 +1,106 @@ +import os +from unittest import mock + +import pytest + +from partinet.gui import job_runner + +JobSpec = job_runner.JobSpec +SlurmOptions = job_runner.SlurmOptions +_write_slurm_script = job_runner._write_slurm_script +load_slurm_defaults = job_runner.load_slurm_defaults +resolve_slurm_options = job_runner.resolve_slurm_options +stream_job = job_runner.stream_job + + +def test_load_slurm_defaults_empty_without_config(monkeypatch): + monkeypatch.delenv("PARTINET_SLURM_CONFIG", raising=False) + opts = load_slurm_defaults() + assert opts.partition == "" + assert opts.partinet_cmd == "partinet" + + +def test_write_slurm_script_has_no_site_defaults(tmp_path): + project = tmp_path / "project" + project.mkdir() + spec = JobSpec( + stage="denoise", + command=["partinet", "denoise", "--source", "/data/mrc", "--project", str(project)], + project_dir=str(project), + log_path=str(project / "partinet_denoise.log"), + slurm=SlurmOptions(partinet_cmd="partinet"), + ) + script = _write_slurm_script(spec, resolve_slurm_options("denoise", spec.slurm)) + text = open(script, encoding="utf-8").read() + assert "#SBATCH --partition=" not in text + assert "#SBATCH --cpus-per-task=32" in text + assert "#SBATCH --mem=100G" in text + assert "partinet denoise" in text + assert str(project) in text + + +def test_stage_slurm_defaults(tmp_path): + project = tmp_path / "project" + project.mkdir() + for stage, cpus, mem, gpus in [ + ("denoise", "32", "100G", ""), + ("detect", "32", "100G", "4"), + ("star", "16", "64G", ""), + ]: + spec = JobSpec( + stage=stage, + command=["partinet", stage], + project_dir=str(project), + log_path=str(project / f"{stage}.log"), + slurm=SlurmOptions(partinet_cmd="partinet"), + ) + slurm = resolve_slurm_options(stage, spec.slurm) + assert slurm.cpus == cpus + assert slurm.mem == mem + assert slurm.gpus == gpus + + +def test_stream_job_local_subprocess(tmp_path): + log_path = tmp_path / "out.log" + with mock.patch.object(job_runner.subprocess, "Popen") as popen: + proc = mock.Mock() + stdout = mock.Mock() + stdout.readline.side_effect = ["line1\n", ""] + stdout.read.return_value = "" + proc.stdout = stdout + proc.poll.return_value = 0 + proc.returncode = 0 + proc.pid = 1234 + popen.return_value = proc + spec = JobSpec( + stage="star", + command=["partinet", "star", "--labels", "a", "--images", "b", "--output", str(log_path), "--conf", "0.1"], + project_dir=str(tmp_path), + log_path=str(log_path), + mode="local", + ) + chunks = list(stream_job(spec)) + assert any("Running:" in c for c in chunks) + assert chunks[-1].endswith("--- Complete ---") + + +def test_stream_job_local_cancel(tmp_path): + log_path = tmp_path / "out.log" + with mock.patch.object(job_runner.subprocess, "Popen") as popen, mock.patch.object( + job_runner, "is_cancelled", side_effect=[False, True] + ): + proc = mock.Mock() + proc.stdout = mock.Mock(readline=mock.Mock(return_value="")) + proc.poll.return_value = None + proc.pid = 5678 + popen.return_value = proc + spec = JobSpec( + stage="denoise", + command=["partinet", "denoise", "--source", "/data", "--project", str(tmp_path)], + project_dir=str(tmp_path), + log_path=str(log_path), + mode="local", + ) + chunks = list(stream_job(spec)) + assert chunks[-1].endswith("--- CANCELLED ---\n") + proc.terminate.assert_called()