MaybeRichard commited on
Commit
1a18f22
·
verified ·
1 Parent(s): 4e4d484

code: complete eval pipeline (7 metrics + per-class + Wilcoxon) + Swin-UNet/TransUNet networks; remove backups/obsolete

Browse files
This view is limited to 50 files because it contains too many changes.   See raw diff
Files changed (50) hide show
  1. code/envs/seggen-controlnet.yml +18 -0
  2. code/envs/umamba.yml +21 -12
  3. code/framework/common/subset.py +0 -40
  4. code/framework/config.py +0 -7
  5. code/framework/config.py.bak +0 -109
  6. code/framework/data/loaders.py +0 -1
  7. code/framework/data/loaders.py.bak +0 -40
  8. code/framework/data/unified_dataset.py +1 -8
  9. code/framework/data/unified_dataset.py.bak +0 -180
  10. code/framework/eval_ckpt.py +0 -120
  11. code/framework/metrics/boundary.py +0 -68
  12. code/framework/{common → synth}/__init__.py +0 -0
  13. code/framework/synth/generative_baselines.py +112 -0
  14. code/scripts/a100_nnum_eval512.sh +47 -0
  15. code/scripts/a100_swin_transunet_3seed_eval512.py +58 -0
  16. code/scripts/h800_cache_data.sh +14 -0
  17. code/scripts/h800_fetch_data.py +31 -0
  18. code/scripts/h800_parallel_extract.sh +38 -0
  19. code/scripts/h800_run_unified512.py +80 -0
  20. code/scripts/h800_setup_seggen.sh +28 -0
  21. code/scripts/h800_swin_transunet_eval512.py +68 -0
  22. code/scripts/hf_update_unified512.py +26 -0
  23. code/scripts/hf_upload_gensegdataset.py +23 -0
  24. code/scripts/hf_upload_tars.py +31 -0
  25. code/scripts/p1/backbones.py +44 -0
  26. code/scripts/p1/fd_lever.py +99 -0
  27. code/scripts/p1/fd_results.json +42 -0
  28. code/scripts/p1/fid_and_viz.py +130 -0
  29. code/scripts/p1/fid_fixed.py +37 -0
  30. code/scripts/p1/fid_results.json +14 -0
  31. code/scripts/p1/gen_prdc.json +74 -0
  32. code/scripts/p1/gen_prdc.py +58 -0
  33. code/scripts/p1/make_jit_vs_fd.py +65 -0
  34. code/scripts/p1/p1_busi_master.py +166 -0
  35. code/scripts/p1/p1_busi_results.json +152 -0
  36. code/scripts/p1/p1_full_metrics.json +182 -0
  37. code/scripts/p1/p1_gen_queue.sh +32 -0
  38. code/scripts/p1/p1_master.py +167 -0
  39. code/scripts/p1/p1_results.json +302 -0
  40. code/scripts/p1/smoke_backbone.py +76 -0
  41. code/scripts/p1/smoke_pixelgen.py +70 -0
  42. code/scripts/p1/train_fd_patched.py +179 -0
  43. code/sota/Swin-Unet/README.md +64 -0
  44. code/sota/Swin-Unet/config.py +229 -0
  45. code/sota/Swin-Unet/configs/swin_tiny_patch4_window7_224_lite.yaml +12 -0
  46. code/sota/Swin-Unet/datasets/README.md +29 -0
  47. code/sota/Swin-Unet/datasets/__init__.py +0 -0
  48. code/sota/Swin-Unet/datasets/dataset_synapse.py +82 -0
  49. code/sota/Swin-Unet/lists/lists_Synapse/all.lst +30 -0
  50. code/sota/Swin-Unet/lists/lists_Synapse/test_vol.txt +12 -0
code/envs/seggen-controlnet.yml ADDED
@@ -0,0 +1,18 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Isolated env for ControlNet (pytorch-lightning 1.5 + old torch). Needs SD v1.5 ckpt.
2
+ # Generative-augmentation baseline only; never mix with the main env.
3
+ name: seggen-controlnet
4
+ channels: [conda-forge]
5
+ dependencies:
6
+ - python=3.10
7
+ - pip
8
+ - pip:
9
+ # torch 1.12.1 (cu113 wheels bundle their own runtime):
10
+ # pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113
11
+ - pytorch-lightning==1.5.0
12
+ - omegaconf==2.1.1
13
+ - einops==0.3.0
14
+ - transformers==4.19.2
15
+ - open-clip-torch==2.0.2
16
+ - gradio==3.16.2
17
+ - opencv-python-headless
18
+ - basicsr==1.4.2
code/envs/umamba.yml CHANGED
@@ -1,16 +1,25 @@
1
- # Isolated env for U-Mamba reference baseline. mamba_ssm needs compilation; the repo
2
- # targets CUDA 11.8 + torch 2.0.1. Compilation on newer stacks may fail — pin to 11.8.
3
- # AMP can NaN in Mamba layers: use the NoAMP trainer variant.
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4
  name: umamba
5
  channels: [conda-forge]
6
  dependencies:
7
  - python=3.10
8
- - pip
9
- - pip:
10
- # pip install torch==2.0.1 torchvision==0.15.2 --index-url https://download.pytorch.org/whl/cu118
11
- # then (in order):
12
- # pip install causal-conv1d>=1.2.0
13
- # pip install mamba-ssm --no-cache-dir
14
- # cd sota/U-Mamba/umamba && pip install -e .
15
- - simpleitk
16
- - nibabel
 
1
+ # Isolated env for U-Mamba reference baseline (nnU-Net v2.1.1 + Mamba).
2
+ # VALIDATED on the A100 server (2026-06-05): mamba CUDA kernel + UMambaBot trainer
3
+ # run end-to-end. mamba_ssm/causal-conv1d are installed from PREBUILT WHEELS to
4
+ # avoid local CUDA compilation (system nvcc 12.8 != torch cu118).
5
+ #
6
+ # Reproduce with these exact commands (NOT `conda env create -f`; pip-driven):
7
+ #
8
+ # conda create -n umamba python=3.10 -y && conda activate umamba
9
+ # pip install torch==2.0.1 torchvision==0.15.2 --index-url https://download.pytorch.org/whl/cu118
10
+ # pip install https://github.com/Dao-AILab/causal-conv1d/releases/download/v1.2.0.post2/causal_conv1d-1.2.0.post2+cu118torch2.0cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
11
+ # pip install https://github.com/state-spaces/mamba/releases/download/v1.2.0.post1/mamba_ssm-1.2.0.post1+cu118torch2.0cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
12
+ # pip install -e sota/U-Mamba/umamba
13
+ # # critical pins (otherwise: transformers drops GreedySearchDecoderOnlyOutput; opencv/numpy clash):
14
+ # pip install "numpy<2" "transformers==4.38.2" "opencv-python-headless==4.9.0.80"
15
+ #
16
+ # Train (uses nnU-Net format from framework/nnunet_convert.py; A100 only — bf16/sm_80):
17
+ # export CUDA_VISIBLE_DEVICES=4
18
+ # nnUNetv2_plan_and_preprocess -d <ID> -c 2d # umamba env does its OWN 2.1.1 preprocess
19
+ # cp <raw>/Dataset<ID>_*/splits_final.json $nnUNet_preprocessed/Dataset<ID>_*/splits_final.json
20
+ # nnUNetv2_train <ID> 2d 0 -tr nnUNetTrainerUMambaBot
21
+ # # AMP can NaN in Mamba: if so use -tr nnUNetTrainerUMambaEncNoAMP
22
  name: umamba
23
  channels: [conda-forge]
24
  dependencies:
25
  - python=3.10
 
 
 
 
 
 
 
 
 
code/framework/common/subset.py DELETED
@@ -1,40 +0,0 @@
1
- """Deterministic low-data subset selection.
2
-
3
- The SAME function is used by (a) the segmentation trainer's real-data loader and
4
- (b) the pixel-diffusion generator's data loader, so that a low-data experiment at
5
- fraction f trains the generator on EXACTLY the real images the segmenter sees —
6
- no leakage, no mismatch, fully reproducible.
7
-
8
- Selection is driven only by (fraction, fraction_seed) and the SORTED list of
9
- items, so it is independent of the training seed: the 3 segmentation seeds all
10
- use the identical real subset, only their weight init differs.
11
- """
12
- from __future__ import annotations
13
-
14
- import math
15
- import random
16
- from typing import List, Sequence, TypeVar
17
-
18
- T = TypeVar("T")
19
-
20
-
21
- def select_fraction(items: Sequence[T], fraction: float, seed: int = 0) -> List[T]:
22
- """Return a deterministic subset of `items` of size ceil(fraction*N).
23
-
24
- items: any sequence (e.g. list of (image_path, mask_path) tuples). It is
25
- sorted by repr() first so order of discovery never affects the subset.
26
- fraction: in (0, 1]. >=1 returns all items (sorted copy).
27
- seed: subset seed (NOT the training seed). Fixed across training seeds.
28
- """
29
- ordered = sorted(items, key=lambda x: repr(x))
30
- n = len(ordered)
31
- if fraction >= 1.0 or n == 0:
32
- return ordered
33
- if fraction <= 0.0:
34
- raise ValueError(f"fraction must be in (0,1], got {fraction}")
35
- k = max(1, math.ceil(fraction * n))
36
- rng = random.Random(seed)
37
- idx = list(range(n))
38
- rng.shuffle(idx)
39
- keep = sorted(idx[:k])
40
- return [ordered[i] for i in keep]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
code/framework/config.py CHANGED
@@ -33,13 +33,6 @@ class Config:
33
  # Points at a dir laid out like a split: <synth_train_dir>/{images,masks}/.
34
  synth_train_dir: str = "" # "" = real data only (no generative augmentation)
35
 
36
- # low-data regime: train on a deterministic fraction of the REAL train split.
37
- # The SAME subset is used by the PixDiff generator (framework.common.subset), so
38
- # real and real+synth arms never disagree. fraction_seed is FIXED across the 3
39
- # training seeds (only weight init varies), so the data subset stays constant.
40
- train_fraction: float = 1.0 # 1.0 = full data; e.g. 0.1 = 10%
41
- fraction_seed: int = 0 # subset seed (independent of `seed`)
42
-
43
  # ---- augmentation (conventional baseline tier) ----
44
  aug: str = "standard" # none | standard | strong (albumentations online)
45
  aug_backend: str = "albumentations" # albumentations | monai
 
33
  # Points at a dir laid out like a split: <synth_train_dir>/{images,masks}/.
34
  synth_train_dir: str = "" # "" = real data only (no generative augmentation)
35
 
 
 
 
 
 
 
 
36
  # ---- augmentation (conventional baseline tier) ----
37
  aug: str = "standard" # none | standard | strong (albumentations online)
38
  aug_backend: str = "albumentations" # albumentations | monai
code/framework/config.py.bak DELETED
@@ -1,109 +0,0 @@
1
- """Unified experiment configuration.
2
-
3
- A single dataclass drives every run. Values can come from (in priority order):
4
- 1. command-line flags (argparse) 2. a YAML file (--config) 3. dataclass defaults.
5
-
6
- The same config object is used by train.py / test.py so that a training run and
7
- its evaluation are guaranteed to agree on dataset, model, image size, etc.
8
- """
9
- from __future__ import annotations
10
-
11
- import argparse
12
- import dataclasses
13
- from dataclasses import dataclass, field, asdict
14
- from typing import Optional, List
15
-
16
- import yaml
17
-
18
-
19
- @dataclass
20
- class Config:
21
- # ---- experiment identity ----
22
- exp_name: str = "default" # results/<exp_name>/<dataset>/<arch>/seed<seed>/
23
- seed: int = 0
24
-
25
- # ---- data ----
26
- data_root: str = "dataset/processed_unified"
27
- dataset: str = "cvc_clinicdb" # folder name under data_root
28
- protocol: str = "official" # e.g. official / fold01 ...
29
- in_channels: int = 0 # 0 = auto-detect from metadata/first image
30
- num_classes: int = 0 # 0 = auto-detect from metadata/masks (incl. background)
31
- img_size: int = 256 # square resize target (Swin/TransUNet need 224)
32
- # extra synthetic (image,mask) pairs to MERGE into the train split.
33
- # Points at a dir laid out like a split: <synth_train_dir>/{images,masks}/.
34
- synth_train_dir: str = "" # "" = real data only (no generative augmentation)
35
-
36
- # ---- augmentation (conventional baseline tier) ----
37
- aug: str = "standard" # none | standard | strong (albumentations online)
38
- aug_backend: str = "albumentations" # albumentations | monai
39
- normalize: str = "auto" # auto(imagenet for RGB, 0.5 for gray) | imagenet | none
40
-
41
- # ---- model ----
42
- arch: str = "unet" # see models/registry.py REGISTRY
43
- encoder: str = "resnet34" # SMP encoder name (ignored by non-SMP archs)
44
- encoder_weights: str = "imagenet" # imagenet | none
45
- pretrained_ckpt: str = "" # ViT/Swin pretrain for transunet/swinunet (optional)
46
-
47
- # ---- optimization ----
48
- epochs: int = 100
49
- batch_size: int = 16 # per-GPU batch size
50
- lr: float = 1e-4
51
- weight_decay: float = 1e-4
52
- optimizer: str = "adamw" # adamw | sgd
53
- scheduler: str = "poly" # poly | cosine | none
54
- warmup_epochs: int = 0
55
- loss: str = "ce_dice" # ce_dice | ce | dice
56
- num_workers: int = 8
57
- grad_clip: float = 0.0 # 0 = disabled
58
-
59
- # ---- precision / hardware ----
60
- amp: str = "bf16" # bf16(A100+) | fp16(V100) | fp32
61
- # DDP is driven by torchrun env vars (RANK/WORLD_SIZE/LOCAL_RANK); nothing to set here.
62
-
63
- # ---- evaluation / logging ----
64
- val_interval: int = 5 # epochs between validations
65
- min_epochs: int = 0 # never early-stop before this many epochs
66
- patience: int = 0 # early-stop after this many epochs w/o val improvement (0 = off)
67
- save_interval: int = 0 # 0 = only save best + last
68
- include_background: bool = False # include class 0 in reported Dice/IoU
69
- compute_hd95: bool = True
70
- out_root: str = "results"
71
- resume: str = "" # path to checkpoint to resume from
72
- visualize: bool = True # save overlays at test time
73
- vis_max: int = 32 # max number of overlay images to save
74
-
75
- def out_dir(self) -> str:
76
- return f"{self.out_root}/{self.exp_name}/{self.dataset}_{self.protocol}/{self.arch}/seed{self.seed}"
77
-
78
- def to_yaml(self, path: str) -> None:
79
- with open(path, "w") as f:
80
- yaml.safe_dump(asdict(self), f, sort_keys=False, allow_unicode=True)
81
-
82
- @classmethod
83
- def from_args(cls, argv: Optional[List[str]] = None) -> "Config":
84
- # First pass: only grab --config so YAML can set defaults that flags then override.
85
- pre = argparse.ArgumentParser(add_help=False)
86
- pre.add_argument("--config", type=str, default="")
87
- known, _ = pre.parse_known_args(argv)
88
-
89
- base = cls()
90
- if known.config:
91
- with open(known.config) as f:
92
- ydata = yaml.safe_load(f) or {}
93
- base = dataclasses.replace(base, **{k: v for k, v in ydata.items()
94
- if k in {f.name for f in dataclasses.fields(cls)}})
95
-
96
- p = argparse.ArgumentParser(parents=[pre],
97
- description="SegGen unified segmentation framework")
98
- for f in dataclasses.fields(cls):
99
- default = getattr(base, f.name)
100
- if f.type is bool or isinstance(default, bool):
101
- # support --flag / --no-flag
102
- p.add_argument(f"--{f.name}", dest=f.name, action="store_true", default=default)
103
- p.add_argument(f"--no-{f.name}", dest=f.name, action="store_false")
104
- else:
105
- p.add_argument(f"--{f.name}", type=type(default) if default is not None else str,
106
- default=default)
107
- ns = p.parse_args(argv)
108
- kwargs = {f.name: getattr(ns, f.name) for f in dataclasses.fields(cls)}
109
- return cls(**kwargs)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
code/framework/data/loaders.py CHANGED
@@ -17,7 +17,6 @@ def build_dataset(cfg, split: str) -> UnifiedSegDataset:
17
  data_root=cfg.data_root, dataset=cfg.dataset, protocol=cfg.protocol, split=split,
18
  transform=None, in_channels=cfg.in_channels, num_classes=cfg.num_classes,
19
  synth_dir=synth,
20
- train_fraction=cfg.train_fraction, fraction_seed=cfg.fraction_seed,
21
  )
22
  ds.transform = build_transform(cfg.img_size, ds.in_channels, train=train,
23
  aug=cfg.aug, normalize=cfg.normalize)
 
17
  data_root=cfg.data_root, dataset=cfg.dataset, protocol=cfg.protocol, split=split,
18
  transform=None, in_channels=cfg.in_channels, num_classes=cfg.num_classes,
19
  synth_dir=synth,
 
20
  )
21
  ds.transform = build_transform(cfg.img_size, ds.in_channels, train=train,
22
  aug=cfg.aug, normalize=cfg.normalize)
code/framework/data/loaders.py.bak DELETED
@@ -1,40 +0,0 @@
1
- """Build datasets / dataloaders from a Config, consistent across train & test."""
2
- from __future__ import annotations
3
-
4
- from torch.utils.data import DataLoader
5
- from torch.utils.data.distributed import DistributedSampler
6
-
7
- from .unified_dataset import UnifiedSegDataset
8
- from .transforms import build_transform
9
- from ..engine.distributed import is_dist
10
-
11
-
12
- def build_dataset(cfg, split: str) -> UnifiedSegDataset:
13
- train = (split == "train")
14
- synth = cfg.synth_train_dir if train else ""
15
- # construct without transform first so in_channels/num_classes auto-detect runs
16
- ds = UnifiedSegDataset(
17
- data_root=cfg.data_root, dataset=cfg.dataset, protocol=cfg.protocol, split=split,
18
- transform=None, in_channels=cfg.in_channels, num_classes=cfg.num_classes,
19
- synth_dir=synth,
20
- )
21
- ds.transform = build_transform(cfg.img_size, ds.in_channels, train=train,
22
- aug=cfg.aug, normalize=cfg.normalize)
23
- return ds
24
-
25
-
26
- def build_loader(cfg, split: str, ds: UnifiedSegDataset) -> DataLoader:
27
- train = (split == "train")
28
- sampler = None
29
- if is_dist():
30
- sampler = DistributedSampler(ds, shuffle=train, drop_last=train)
31
- return DataLoader(
32
- ds,
33
- batch_size=cfg.batch_size,
34
- shuffle=(train and sampler is None),
35
- sampler=sampler,
36
- num_workers=cfg.num_workers,
37
- pin_memory=True,
38
- drop_last=(train and sampler is None),
39
- persistent_workers=cfg.num_workers > 0,
40
- )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
code/framework/data/unified_dataset.py CHANGED
@@ -124,8 +124,7 @@ class UnifiedSegDataset(Dataset):
124
  def __init__(self, data_root: str, dataset: str, protocol: str, split: str,
125
  transform: Optional[Callable] = None,
126
  in_channels: int = 0, num_classes: int = 0,
127
- synth_dir: str = "",
128
- train_fraction: float = 1.0, fraction_seed: int = 0):
129
  self.data_root = data_root
130
  self.dataset = dataset
131
  self.split = split
@@ -142,12 +141,6 @@ class UnifiedSegDataset(Dataset):
142
  if not pairs:
143
  raise RuntimeError(f"no (image,mask) pairs found in {split_dir}")
144
 
145
- # low-data: deterministically subsample REAL train pairs (the generator uses
146
- # the SAME select_fraction, so no leakage) BEFORE merging any synthetic data.
147
- if split == "train" and train_fraction < 1.0:
148
- from ..common.subset import select_fraction
149
- pairs = select_fraction(pairs, train_fraction, fraction_seed)
150
-
151
  # optionally merge synthetic (image,mask) pairs into the (train) split
152
  if synth_dir and os.path.isdir(synth_dir):
153
  sp = _pair_by_glob(synth_dir if os.path.isdir(os.path.join(synth_dir, "images"))
 
124
  def __init__(self, data_root: str, dataset: str, protocol: str, split: str,
125
  transform: Optional[Callable] = None,
126
  in_channels: int = 0, num_classes: int = 0,
127
+ synth_dir: str = ""):
 
128
  self.data_root = data_root
129
  self.dataset = dataset
130
  self.split = split
 
141
  if not pairs:
142
  raise RuntimeError(f"no (image,mask) pairs found in {split_dir}")
143
 
 
 
 
 
 
 
144
  # optionally merge synthetic (image,mask) pairs into the (train) split
145
  if synth_dir and os.path.isdir(synth_dir):
146
  sp = _pair_by_glob(synth_dir if os.path.isdir(os.path.join(synth_dir, "images"))
code/framework/data/unified_dataset.py.bak DELETED
@@ -1,180 +0,0 @@
1
- """Dataset reader for the standardized `processed_unified` layout.
2
-
3
- Expected layout (see dataset/SEGMENTATION_WORKSPACE_README.md):
4
- <data_root>/<dataset>/<protocol>/<split>/images/*.png
5
- <data_root>/<dataset>/<protocol>/<split>/masks/*.png
6
- <data_root>/<dataset>/metadata.json (optional, preferred)
7
- <data_root>/<dataset>/manifest.jsonl (optional)
8
-
9
- Returns per item: {"image": FloatTensor[C,H,W], "mask": LongTensor[H,W], "name": str}.
10
-
11
- Binary and multi-class masks are both supported: masks keep their integer class
12
- ids (0..C-1). Auto-detection of in_channels / num_classes falls back to scanning
13
- files when metadata is absent, so the loader is robust to missing metadata.
14
- """
15
- from __future__ import annotations
16
-
17
- import json
18
- import os
19
- from glob import glob
20
- from typing import Optional, Callable, List, Tuple
21
-
22
- import numpy as np
23
- import cv2
24
- from torch.utils.data import Dataset
25
-
26
-
27
- _MODALITY_CHANNELS = { # hint table; only used when metadata lacks in_channels
28
- "rgb": 3, "fundus": 3, "colonoscopy": 3, "endoscopy": 3, "histopathology": 3,
29
- "ultrasound": 1, "mri": 1, "ct": 1, "grayscale": 1,
30
- }
31
-
32
- # Documented class counts (incl. background). metadata.json on the server has no
33
- # num_classes field, so this table is the fast, reliable primary source; unknown
34
- # datasets fall back to a FULL scan of the mask set (accurate but slower).
35
- _KNOWN_NUM_CLASSES = {
36
- "cvc_clinicdb": 2, "kvasir_seg": 2, "fives": 2, "busi": 2,
37
- "refuge2": 3, "acdc_png": 4,
38
- "idridd_segmentation": 6, "pannuke_semantic": 6,
39
- }
40
-
41
-
42
- def _read_metadata(data_root: str, dataset: str) -> dict:
43
- path = os.path.join(data_root, dataset, "metadata.json")
44
- if os.path.isfile(path):
45
- try:
46
- with open(path) as f:
47
- return json.load(f)
48
- except Exception:
49
- return {}
50
- return {}
51
-
52
-
53
- def _pair_from_manifest(split_dir: str, manifest: str) -> Optional[List[Tuple[str, str]]]:
54
- if not os.path.isfile(manifest):
55
- return None
56
- pairs = []
57
- base = os.path.dirname(manifest)
58
- with open(manifest) as f:
59
- for line in f:
60
- line = line.strip()
61
- if not line:
62
- continue
63
- rec = json.loads(line)
64
- img = rec.get("image") or rec.get("image_path") or rec.get("img")
65
- msk = rec.get("mask") or rec.get("mask_path") or rec.get("label")
66
- if img is None or msk is None:
67
- return None
68
- # manifest paths may be relative to dataset root or absolute
69
- ip = img if os.path.isabs(img) else os.path.join(base, img)
70
- mp = msk if os.path.isabs(msk) else os.path.join(base, msk)
71
- # only keep entries that fall under this split dir
72
- if os.path.normpath(split_dir) in os.path.normpath(ip):
73
- pairs.append((ip, mp))
74
- return pairs or None
75
-
76
-
77
- def _pair_by_glob(split_dir: str) -> List[Tuple[str, str]]:
78
- img_dir = os.path.join(split_dir, "images")
79
- msk_dir = os.path.join(split_dir, "masks")
80
- imgs = sorted(glob(os.path.join(img_dir, "*")))
81
- pairs = []
82
- for ip in imgs:
83
- stem = os.path.splitext(os.path.basename(ip))[0]
84
- # mask may share extension or be .png
85
- cands = glob(os.path.join(msk_dir, stem + ".*"))
86
- if not cands:
87
- continue
88
- pairs.append((ip, cands[0]))
89
- return pairs
90
-
91
-
92
- def detect_in_channels(meta: dict, sample_img: Optional[str]) -> int:
93
- if meta.get("in_channels"):
94
- return int(meta["in_channels"])
95
- mod = str(meta.get("modality", "")).lower()
96
- for k, v in _MODALITY_CHANNELS.items():
97
- if k in mod:
98
- return v
99
- if sample_img and os.path.isfile(sample_img):
100
- im = cv2.imread(sample_img, cv2.IMREAD_UNCHANGED)
101
- if im is not None and im.ndim == 3 and im.shape[2] >= 3:
102
- return 3
103
- return 1
104
-
105
-
106
- def detect_num_classes(meta: dict, mask_paths: List[str], dataset: str = "") -> int:
107
- if dataset in _KNOWN_NUM_CLASSES:
108
- return _KNOWN_NUM_CLASSES[dataset]
109
- if meta.get("num_classes"):
110
- return int(meta["num_classes"])
111
- # unknown dataset: scan ALL masks so a rare class is never missed
112
- vals = set()
113
- for mp in mask_paths:
114
- m = cv2.imread(mp, cv2.IMREAD_GRAYSCALE)
115
- if m is not None:
116
- vals.update(np.unique(m).tolist())
117
- if not vals:
118
- return 2
119
- maxv = max(vals)
120
- return int(maxv) + 1 if maxv >= 1 else 2
121
-
122
-
123
- class UnifiedSegDataset(Dataset):
124
- def __init__(self, data_root: str, dataset: str, protocol: str, split: str,
125
- transform: Optional[Callable] = None,
126
- in_channels: int = 0, num_classes: int = 0,
127
- synth_dir: str = ""):
128
- self.data_root = data_root
129
- self.dataset = dataset
130
- self.split = split
131
- self.transform = transform
132
-
133
- split_dir = os.path.join(data_root, dataset, protocol, split)
134
- if not os.path.isdir(split_dir):
135
- raise FileNotFoundError(
136
- f"split dir not found: {split_dir}\n"
137
- f"(data is prepared separately; see dataset/ scripts)")
138
-
139
- manifest = os.path.join(data_root, dataset, "manifest.jsonl")
140
- pairs = _pair_from_manifest(split_dir, manifest) or _pair_by_glob(split_dir)
141
- if not pairs:
142
- raise RuntimeError(f"no (image,mask) pairs found in {split_dir}")
143
-
144
- # optionally merge synthetic (image,mask) pairs into the (train) split
145
- if synth_dir and os.path.isdir(synth_dir):
146
- sp = _pair_by_glob(synth_dir if os.path.isdir(os.path.join(synth_dir, "images"))
147
- else os.path.dirname(synth_dir))
148
- pairs = pairs + sp
149
-
150
- self.pairs = pairs
151
- meta = _read_metadata(data_root, dataset)
152
- self.in_channels = in_channels or detect_in_channels(meta, pairs[0][0])
153
- self.num_classes = num_classes or detect_num_classes(meta, [p[1] for p in pairs], dataset)
154
-
155
- def __len__(self) -> int:
156
- return len(self.pairs)
157
-
158
- def _load_image(self, path: str) -> np.ndarray:
159
- if self.in_channels == 1:
160
- im = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
161
- if im is None:
162
- raise IOError(f"cannot read image {path}")
163
- return im[:, :, None] # H,W,1
164
- im = cv2.imread(path, cv2.IMREAD_COLOR) # BGR
165
- if im is None:
166
- raise IOError(f"cannot read image {path}")
167
- return cv2.cvtColor(im, cv2.COLOR_BGR2RGB) # H,W,3
168
-
169
- def __getitem__(self, idx: int):
170
- ip, mp = self.pairs[idx]
171
- image = self._load_image(ip)
172
- mask = cv2.imread(mp, cv2.IMREAD_GRAYSCALE)
173
- if mask is None:
174
- raise IOError(f"cannot read mask {mp}")
175
- mask = mask.astype(np.int64)
176
-
177
- if self.transform is not None:
178
- image, mask = self.transform(image, mask)
179
- return {"image": image, "mask": mask,
180
- "name": os.path.splitext(os.path.basename(ip))[0]}
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
code/framework/eval_ckpt.py DELETED
@@ -1,120 +0,0 @@
1
- """Evaluate a trained framework checkpoint on a chosen dataset/split, with the full
2
- metric set INCLUDING boundary metrics. Two uses, one script (no retraining):
3
-
4
- * C7 (boundary fidelity, in-domain): point --eval_dataset at the same dataset to
5
- recompute boundary-Dice / NSD from a saved best.pth.
6
- * C4 (cross-center): set --eval_dataset/--eval_protocol to a DIFFERENT but
7
- label-compatible dataset (e.g. train busi? no — train cvc_clinicdb, eval
8
- kvasir_seg; both binary polyp). num_classes/in_channels must match.
9
-
10
- Run from project root (…/NPJ), env seggen:
11
- CUDA_VISIBLE_DEVICES=5 python -m framework.eval_ckpt \
12
- --ckpt results/baselines/cvc_clinicdb_official/unet/seed0/best.pth \
13
- --arch unet --encoder resnet50 \
14
- --data_root $DR --dataset cvc_clinicdb --protocol official \
15
- --eval_dataset kvasir_seg --eval_protocol official \
16
- --out_json results/crosscenter/cvc2kvasir_unet_seed0.json
17
- """
18
- from __future__ import annotations
19
-
20
- import argparse
21
- import json
22
- import os
23
- import sys
24
-
25
- sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
26
-
27
- import numpy as np
28
- import torch
29
- from torch.utils.data import DataLoader
30
-
31
- from framework.config import Config # noqa: F401 (kept for parity / future YAML use)
32
- from framework.models.registry import build_model, required_img_size
33
- from framework.data.unified_dataset import UnifiedSegDataset
34
- from framework.data.transforms import build_transform
35
- from framework.metrics.metrics import per_image_metrics
36
- from framework.metrics.boundary import boundary_metrics
37
-
38
-
39
- def get_args():
40
- p = argparse.ArgumentParser("Evaluate a checkpoint (in-domain or cross-center) + boundary metrics")
41
- p.add_argument("--ckpt", required=True)
42
- p.add_argument("--arch", default="unet")
43
- p.add_argument("--encoder", default="resnet50")
44
- p.add_argument("--data_root", required=True)
45
- p.add_argument("--dataset", required=True, help="dataset the ckpt was TRAINED on (sets in_ch/num_classes)")
46
- p.add_argument("--protocol", required=True)
47
- p.add_argument("--eval_dataset", default="", help="dataset to evaluate ON (default = --dataset)")
48
- p.add_argument("--eval_protocol", default="", help="default = --protocol")
49
- p.add_argument("--split", default="test")
50
- p.add_argument("--img_size", type=int, default=256)
51
- p.add_argument("--normalize", default="auto")
52
- p.add_argument("--tol", type=float, default=2.0)
53
- p.add_argument("--batch_size", type=int, default=16)
54
- p.add_argument("--num_workers", type=int, default=6)
55
- p.add_argument("--out_json", default="")
56
- return p.parse_args()
57
-
58
-
59
- def main():
60
- a = get_args()
61
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
62
- img_size = required_img_size(a.arch) or a.img_size
63
- eval_ds_name = a.eval_dataset or a.dataset
64
- eval_proto = a.eval_protocol or a.protocol
65
-
66
- # in_ch / num_classes are fixed by the TRAIN dataset (how the model was built)
67
- train_probe = UnifiedSegDataset(a.data_root, a.dataset, a.protocol, "test", transform=None)
68
- in_ch, n_cls = train_probe.in_channels, train_probe.num_classes
69
-
70
- model = build_model(a.arch, in_channels=in_ch, num_classes=n_cls, img_size=img_size,
71
- encoder=a.encoder, encoder_weights="none", pretrained_ckpt="")
72
- ckpt = torch.load(a.ckpt, map_location="cpu", weights_only=False)
73
- model.load_state_dict(ckpt["model"])
74
- model.to(device).eval()
75
-
76
- tf = build_transform(img_size, in_ch, train=False, aug="none", normalize=a.normalize)
77
- ds = UnifiedSegDataset(a.data_root, eval_ds_name, eval_proto, a.split,
78
- transform=tf, in_channels=in_ch, num_classes=n_cls)
79
- if ds.num_classes != n_cls:
80
- raise ValueError(f"label mismatch: train num_classes={n_cls} vs eval={ds.num_classes} "
81
- f"({a.dataset}->{eval_ds_name}) — not label-compatible for cross-center.")
82
- cross = (eval_ds_name != a.dataset) or (eval_proto != a.protocol)
83
- print(f"[eval] ckpt={os.path.basename(a.ckpt)} train={a.dataset}/{a.protocol} "
84
- f"eval={eval_ds_name}/{eval_proto}/{a.split} cross_center={cross} "
85
- f"in_ch={in_ch} num_classes={n_cls} n={len(ds)}", flush=True)
86
-
87
- loader = DataLoader(ds, batch_size=a.batch_size, shuffle=False, num_workers=a.num_workers)
88
- recs = []
89
- for batch in loader:
90
- img = batch["image"].to(device, non_blocking=True)
91
- msk = batch["mask"].numpy()
92
- with torch.no_grad():
93
- pred = model(img).argmax(1).cpu().numpy()
94
- for i in range(pred.shape[0]):
95
- m = per_image_metrics(pred[i], msk[i], n_cls,
96
- include_background=False, compute_hd95=True)
97
- b = boundary_metrics(pred[i], msk[i], n_cls, tol=a.tol)
98
- m.update(b)
99
- recs.append(m)
100
-
101
- keys = [k for k, val in recs[0].items() if isinstance(val, (int, float))] # skip per_class dict
102
- summary = {}
103
- for k in keys:
104
- v = np.array([r[k] for r in recs], dtype=np.float64)
105
- v = v[~np.isnan(v)]
106
- summary[k] = {"mean": round(float(v.mean()), 4) if v.size else None,
107
- "std": round(float(v.std()), 4) if v.size else None, "n": int(v.size)}
108
- out = {"ckpt": a.ckpt, "train": f"{a.dataset}/{a.protocol}",
109
- "eval": f"{eval_ds_name}/{eval_proto}/{a.split}", "cross_center": cross,
110
- "num_images": len(ds), "metrics": summary}
111
- print(json.dumps(out, indent=2), flush=True)
112
- if a.out_json:
113
- os.makedirs(os.path.dirname(os.path.abspath(a.out_json)) or ".", exist_ok=True)
114
- with open(a.out_json, "w") as f:
115
- json.dump(out, f, indent=2)
116
- print(f"[eval] wrote {a.out_json}", flush=True)
117
-
118
-
119
- if __name__ == "__main__":
120
- main()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
code/framework/metrics/boundary.py DELETED
@@ -1,68 +0,0 @@
1
- """Boundary-localized segmentation metrics (for the C7 boundary-fidelity analysis).
2
-
3
- Two standard measures, computed per foreground class then averaged:
4
- * Normalized Surface Dice (NSD) @ tol: fraction of pred/gt surface points within
5
- `tol` pixels of the other surface (Nikolov et al.); the medical-standard
6
- boundary metric.
7
- * Boundary-Dice @ tol: Dice between the tol-dilated pred and gt boundaries.
8
-
9
- Recomputable post-hoc from saved predictions/checkpoints, so adding it never
10
- requires retraining (see framework/eval_ckpt.py).
11
- """
12
- from __future__ import annotations
13
-
14
- import numpy as np
15
-
16
- try:
17
- from scipy.ndimage import binary_erosion, binary_dilation, distance_transform_edt
18
- _HAVE_SCIPY = True
19
- except Exception: # pragma: no cover
20
- _HAVE_SCIPY = False
21
-
22
-
23
- def _surface(bin_mask: np.ndarray) -> np.ndarray:
24
- return bin_mask ^ binary_erosion(bin_mask)
25
-
26
-
27
- def _nsd_binary(pred: np.ndarray, gt: np.ndarray, tol: float) -> float:
28
- sp, sg = _surface(pred), _surface(gt)
29
- ssum = sp.sum() + sg.sum()
30
- if sp.sum() == 0 and sg.sum() == 0:
31
- return 1.0
32
- if ssum == 0:
33
- return 0.0
34
- dt_to_gt = distance_transform_edt(~sg)
35
- dt_to_pred = distance_transform_edt(~sp)
36
- pred_close = (dt_to_gt[sp] <= tol).sum()
37
- gt_close = (dt_to_pred[sg] <= tol).sum()
38
- return float((pred_close + gt_close) / ssum)
39
-
40
-
41
- def _bdice_binary(pred: np.ndarray, gt: np.ndarray, tol: int) -> float:
42
- sp, sg = _surface(pred), _surface(gt)
43
- if sp.sum() == 0 and sg.sum() == 0:
44
- return 1.0
45
- spd = binary_dilation(sp, iterations=int(tol))
46
- sgd = binary_dilation(sg, iterations=int(tol))
47
- denom = spd.sum() + sgd.sum()
48
- if denom == 0:
49
- return 0.0
50
- return float(2.0 * (spd & sgd).sum() / denom)
51
-
52
-
53
- def boundary_metrics(pred: np.ndarray, gt: np.ndarray, num_classes: int,
54
- tol: float = 2.0) -> dict:
55
- """Mean over foreground classes (1..num_classes-1). NaN if scipy missing."""
56
- if not _HAVE_SCIPY:
57
- return {"nsd": float("nan"), "boundary_dice": float("nan")}
58
- nsds, bdices = [], []
59
- for c in range(1, num_classes):
60
- p, g = (pred == c), (gt == c)
61
- if g.sum() == 0 and p.sum() == 0:
62
- continue
63
- nsds.append(_nsd_binary(p, g, tol))
64
- bdices.append(_bdice_binary(p, g, int(round(tol))))
65
- return {
66
- "nsd": float(np.mean(nsds)) if nsds else float("nan"),
67
- "boundary_dice": float(np.mean(bdices)) if bdices else float("nan"),
68
- }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
code/framework/{common → synth}/__init__.py RENAMED
File without changes
code/framework/synth/generative_baselines.py ADDED
@@ -0,0 +1,112 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Orchestration for the generative-augmentation SOTA baselines (category B).
2
+
3
+ These methods are compared against our SegGen method. Each runs in its OWN conda
4
+ env (see envs/) because their dependency stacks conflict with the main framework.
5
+ The shared contract: every generator must emit paired (image, mask) into
6
+
7
+ <data_root>/<dataset>/<protocol>/synth_<method>/{images,masks}/
8
+
9
+ which the unified trainer then merges into the train split via --synth_train_dir.
10
+
11
+ Kept baselines:
12
+ * SegGuidedDiff (diffusion, mask->image, medical, modern stack) -- best fit, USE-AS-IS
13
+ * SPADE (GAN, mask->image) -- ADAPT (needs sync_bn)
14
+ * ControlNet (diffusion, SD-finetune, mask->image) -- ADAPT (needs SD ckpt)
15
+
16
+ Dropped (per scoping): StyleGAN2-ADA (no masks), LDM (dep hell + AE training).
17
+
18
+ This module only BUILDS the commands + assembles the standard synth dir; it does
19
+ not import the repos (they live in separate envs). Run the printed commands in the
20
+ matching env, then call assemble_synth_dir() (env-agnostic) to lay out pairs.
21
+ """
22
+ from __future__ import annotations
23
+
24
+ import os
25
+ import shutil
26
+ from glob import glob
27
+
28
+ SOTA = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "sota"))
29
+
30
+
31
+ def assemble_synth_dir(generated_images_dir: str, masks_source_dir: str,
32
+ out_dir: str, strip_prefix: str = "condon_",
33
+ link: bool = True) -> int:
34
+ """Pair each generated image with the real mask it was conditioned on.
35
+
36
+ Mask-conditioned generators name outputs after the conditioning mask
37
+ (SegGuidedDiff: 'condon_<maskname>.png'). We recover the mask name, copy/link
38
+ the matching real mask, and place both under out_dir/{images,masks}/.
39
+ Returns the number of pairs assembled.
40
+ """
41
+ img_out = os.path.join(out_dir, "images")
42
+ msk_out = os.path.join(out_dir, "masks")
43
+ os.makedirs(img_out, exist_ok=True)
44
+ os.makedirs(msk_out, exist_ok=True)
45
+
46
+ n = 0
47
+ for gp in sorted(glob(os.path.join(generated_images_dir, "*"))):
48
+ base = os.path.basename(gp)
49
+ stem = os.path.splitext(base)[0]
50
+ if strip_prefix and stem.startswith(strip_prefix):
51
+ mask_stem = stem[len(strip_prefix):]
52
+ else:
53
+ mask_stem = stem
54
+ cands = glob(os.path.join(masks_source_dir, mask_stem + ".*"))
55
+ if not cands:
56
+ continue
57
+ out_name = f"synth_{n:06d}"
58
+ dst_img = os.path.join(img_out, out_name + os.path.splitext(base)[1])
59
+ dst_msk = os.path.join(msk_out, out_name + os.path.splitext(cands[0])[1])
60
+ _place(gp, dst_img, link)
61
+ _place(cands[0], dst_msk, link)
62
+ n += 1
63
+ return n
64
+
65
+
66
+ def _place(src, dst, link):
67
+ if os.path.exists(dst):
68
+ os.remove(dst)
69
+ if link:
70
+ os.symlink(os.path.abspath(src), dst)
71
+ else:
72
+ shutil.copy2(src, dst)
73
+
74
+
75
+ # ---- command builders (printed into run.sh; run in the matching conda env) ----
76
+
77
+ def segguideddiff_cmds(data_root, dataset, protocol, num_classes, in_channels,
78
+ img_size=256, epochs=400, sample_size=1000):
79
+ repo = os.path.join(SOTA, "segmentation-guided-diffusion")
80
+ img_dir = f"{data_root}/{dataset}/{protocol}/train/images"
81
+ seg_dir = f"{data_root}/{dataset}/{protocol}/train/masks"
82
+ train = (f"cd {repo} && python main.py --mode train --model_type DDIM "
83
+ f"--img_size {img_size} --num_img_channels {in_channels} --dataset {dataset} "
84
+ f"--img_dir {img_dir} --seg_dir {seg_dir} --segmentation_guided "
85
+ f"--num_segmentation_classes {num_classes} --num_epochs {epochs}")
86
+ synth = (f"cd {repo} && python main.py --mode eval_many --model_type DDIM "
87
+ f"--img_size {img_size} --num_img_channels {in_channels} --dataset {dataset} "
88
+ f"--seg_dir {seg_dir} --segmentation_guided "
89
+ f"--num_segmentation_classes {num_classes} --eval_sample_size {sample_size}")
90
+ return train, synth
91
+
92
+
93
+ def spade_cmds(data_root, dataset, protocol, num_classes, img_size=256, niter=100):
94
+ repo = os.path.join(SOTA, "SPADE")
95
+ img_dir = f"{data_root}/{dataset}/{protocol}/train/images"
96
+ lab_dir = f"{data_root}/{dataset}/{protocol}/train/masks"
97
+ setup = (f"cd {repo}/models/networks && "
98
+ f"git clone https://github.com/vacancy/Synchronized-BatchNorm-PyTorch && "
99
+ f"cp -r Synchronized-BatchNorm-PyTorch/sync_batchnorm .")
100
+ train = (f"cd {repo} && python train.py --name {dataset}_spade --dataset_mode custom "
101
+ f"--label_dir {lab_dir} --image_dir {img_dir} --label_nc {num_classes} "
102
+ f"--no_instance --crop_size {img_size} --load_size {int(img_size*1.12)} --niter {niter}")
103
+ synth = (f"cd {repo} && python test.py --name {dataset}_spade --dataset_mode custom "
104
+ f"--label_dir {lab_dir} --image_dir {img_dir} --label_nc {num_classes} "
105
+ f"--no_instance --results_dir ./synth_{dataset}")
106
+ return setup, train, synth
107
+
108
+
109
+ def controlnet_notes():
110
+ return ("ControlNet: download SD v1.5 (~4GB), run tool_add_control.py, write a "
111
+ "MyDataset that colorizes integer masks to RGB hints + triples grayscale "
112
+ "images to 3ch, then tutorial_train.py. Run in env seggen-controlnet.")
code/scripts/a100_nnum_eval512.sh ADDED
@@ -0,0 +1,47 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env bash
2
+ # nnU-Net + U-Mamba @512 re-eval on a100 (CPU-only, does NOT touch GPUs).
3
+ # Re-scores cached predTs/predTs_umamba at eval_size 512 into results/unified512/
4
+ # so they join the unified-512 table. ~64 cells, parallel (concurrency MAXJ).
5
+ set -u
6
+ cd /home/wzhang/LSC/Code/NPJ
7
+ source /opt/anaconda3/etc/profile.d/conda.sh
8
+ conda activate seggen
9
+ export OMP_NUM_THREADS=4 MKL_NUM_THREADS=4 OPENBLAS_NUM_THREADS=4 NUMEXPR_NUM_THREADS=4
10
+ DR=/home/wzhang/LSC/Dataset/Segmentation/processed_unified
11
+ RAW=/home/wzhang/LSC/Code/NPJ/nnunet_workspace/raw
12
+ PRED_NN=/home/wzhang/LSC/Code/NPJ/nnunet_workspace/predTs
13
+ PRED_UM=/home/wzhang/LSC/Code/NPJ/nnunet_workspace/predTs_umamba
14
+ MAXJ=10
15
+
16
+ declare -A DS=(
17
+ [1]=cvc_clinicdb:official [2]=kvasir_seg:official [3]=fives:official
18
+ [4]=refuge2:official [5]=busi:fold01 [6]=idridd_segmentation:fold01
19
+ [7]=acdc_png:official [8]=pannuke_semantic:fold01 [9]=medsegdb_isic2018:holdout
20
+ [10]=medsegdb_kits19:fold01 [11]=pannuke_semantic:fold02 [12]=pannuke_semantic:fold03
21
+ )
22
+
23
+ run=0
24
+ for id in 1 2 3 4 5 6 7 8 9 10 11 12; do
25
+ IFS=: read -r ds proto <<< "${DS[$id]}"
26
+ for f in 0 1 2; do
27
+ for ap in "nnunet:$PRED_NN" "umamba:$PRED_UM"; do
28
+ arch=${ap%%:*}; pred=${ap#*:}
29
+ outdir=$pred/d${id}_f${f}
30
+ [ -d "$outdir" ] && ls -A "$outdir"/*.png >/dev/null 2>&1 || continue
31
+ (
32
+ python framework/nnunet_eval.py --data_root "$DR" --dataset "$ds" --protocol "$proto" \
33
+ --raw "$RAW" --dataset_id "$id" --fold "$f" --pred_dir "$outdir" --arch "$arch" \
34
+ --exp_name unified512 --eval_size 512 \
35
+ > /tmp/nnum512_${arch}_d${id}_f${f}.log 2>&1 \
36
+ && echo "[ok] $arch d${id}_f${f} ($ds $proto)" \
37
+ || echo "[FAIL] $arch d${id}_f${f} ($ds $proto)"
38
+ ) &
39
+ run=$((run+1))
40
+ if [ "$run" -ge "$MAXJ" ]; then wait -n; run=$((run-1)); fi
41
+ done
42
+ done
43
+ done
44
+ wait
45
+ echo "NNUM_EVAL512_DONE"
46
+ n=$(find results/unified512 -path '*/nnunet/*/metrics.json' -o -path '*/umamba/*/metrics.json' 2>/dev/null | wc -l)
47
+ echo "unified512 nnunet+umamba metrics.json count: $n"
code/scripts/a100_swin_transunet_3seed_eval512.py ADDED
@@ -0,0 +1,58 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Fill SwinUNet/TransUNet to FULL 3-seed @512 on a100 (their per-seed best.pth live here).
2
+ For every results/baselines/<cell>/{swinunet,transunet}/seed<s>/best.pth, copy it into the
3
+ unified512 tree and run eval_at_res.py --eval_size 512 --exp_name unified512. GPU 4/5 only
4
+ (A100 80G; PCI_BUS_ID). 64 evals. Then metrics.json get transferred to h800 + re-aggregated.
5
+ """
6
+ import os, glob, shutil, subprocess, time
7
+
8
+ CODE = "/home/wzhang/LSC/Code/NPJ"
9
+ DATA = "/home/wzhang/LSC/Dataset/Segmentation/processed_unified"
10
+ PY = "/opt/anaconda3/envs/seggen/bin/python"
11
+ BASE = CODE + "/results/baselines" # source per-seed weights
12
+ UNI = CODE + "/results/unified512" # eval_at_res writes here (out_root=results rel to CODE)
13
+ LOGD = "/tmp/sw_tr_3seed_logs"; os.makedirs(LOGD, exist_ok=True)
14
+ SLOTS = [4, 4, 4, 5, 5, 5] # GPU 4/5, 3 co-located evals each
15
+
16
+ jobs = []
17
+ for arch in ("swinunet", "transunet"):
18
+ for w in sorted(glob.glob(f"{BASE}/*/{arch}/seed*/best.pth")):
19
+ parts = w.split("/")
20
+ cell, seed = parts[-4], parts[-2] # <cell>, seedN
21
+ sd = int(seed.replace("seed", ""))
22
+ # parse cell -> dataset, protocol
23
+ ds, proto = None, None
24
+ for p in ("official", "holdout", "fold01", "fold02", "fold03"):
25
+ if cell.endswith("_" + p):
26
+ ds, proto = cell[:-(len(p) + 1)], p; break
27
+ out = f"{UNI}/{cell}/{arch}/{seed}"
28
+ jobs.append({"ds": ds, "proto": proto, "arch": arch, "seed": sd, "w": w,
29
+ "out": out, "mj": out + "/metrics.json", "tag": f"{cell}_{arch}_s{sd}"})
30
+
31
+ pending = [j for j in jobs if not os.path.isfile(j["mj"])]
32
+ print(f"[3seed] total={len(jobs)} done={len(jobs)-len(pending)} pending={len(pending)}", flush=True)
33
+
34
+ def make_cmd(j, gpu):
35
+ enc = "R50-ViT-B_16" if j["arch"] == "transunet" else "resnet50"
36
+ os.makedirs(j["out"], exist_ok=True)
37
+ shutil.copy(j["w"], j["out"] + "/best.pth")
38
+ return (f"export CUDA_DEVICE_ORDER=PCI_BUS_ID CUDA_VISIBLE_DEVICES={gpu} "
39
+ f"OMP_NUM_THREADS=8 MKL_NUM_THREADS=8 OPENBLAS_NUM_THREADS=8 && cd {CODE} && "
40
+ f"{PY} framework/eval_at_res.py --data_root {DATA} --dataset {j['ds']} "
41
+ f"--protocol {j['proto']} --arch {j['arch']} --seed {j['seed']} --eval_size 512 "
42
+ f"--exp_name unified512 --encoder {enc}")
43
+
44
+ running = {}; free = list(SLOTS); i = 0; ok = fail = 0
45
+ while i < len(pending) or running:
46
+ while free and i < len(pending):
47
+ gpu = free.pop(0); j = pending[i]; i += 1
48
+ lf = open(f"{LOGD}/{j['tag']}.log", "w")
49
+ p = subprocess.Popen(["bash", "-lc", make_cmd(j, gpu)], stdout=lf, stderr=subprocess.STDOUT)
50
+ running[id(p)] = (p, j, lf, gpu); print(f"[launch] gpu{gpu} {j['tag']}", flush=True)
51
+ time.sleep(6)
52
+ for k, (p, j, lf, gpu) in list(running.items()):
53
+ if p.poll() is not None:
54
+ lf.close(); okj = os.path.isfile(j["mj"]); ok += okj; fail += (not okj)
55
+ print(f"[finish] gpu{gpu} {j['tag']} rc={p.returncode} ok={okj}", flush=True)
56
+ del running[k]; free.append(gpu)
57
+ print(f"[3seed] ALL DONE ok={ok} fail={fail}", flush=True)
58
+ print("SWTR_3SEED_DONE", flush=True)
code/scripts/h800_cache_data.sh ADDED
@@ -0,0 +1,14 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env bash
2
+ # Cache the dataset from the slow JuiceFS share to local RAID /data/temp for fast
3
+ # training reads. Parallel per-dataset cp -a (preserves pannuke hard links), overlaps
4
+ # JuiceFS small-file read latency. Run detached; log to /tmp/cache_data.log.
5
+ SRC=/mnt/tidal-alsh-share2/dataset/qinshengqian/research/c3/NPJ-ACM/Data
6
+ DST=/data/temp/NPJ-ACM/Data
7
+ mkdir -p "$DST"
8
+ echo "[start] $(date +%T)"
9
+ for d in acdc_png busi cvc_clinicdb fives idridd_segmentation kvasir_seg \
10
+ medsegdb_isic2018 medsegdb_kits19 pannuke_semantic refuge2; do
11
+ ( cp -a "$SRC/$d" "$DST/" && echo "done $d $(date +%T)" ) &
12
+ done
13
+ wait
14
+ echo "CACHE_DONE $(date +%T)"
code/scripts/h800_fetch_data.py ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """On h800: download GenSegDataset tars from HF (via proxy+token), extract into Data/,
2
+ then remove the tars. Produces the processed_unified layout under Data/."""
3
+ import os, glob, tarfile
4
+ from huggingface_hub import snapshot_download
5
+
6
+ BASE = "/mnt/tidal-alsh-share2/dataset/qinshengqian/research/c3/NPJ-ACM/Data"
7
+ TARS = os.path.join(BASE, "_tars")
8
+
9
+ print("[1] downloading tars ...", flush=True)
10
+ snapshot_download("MaybeRichard/GenSegDataset", repo_type="dataset",
11
+ allow_patterns=["*.tar", "README.md"], local_dir=TARS)
12
+
13
+ print("[2] extracting ...", flush=True)
14
+ for t in sorted(glob.glob(os.path.join(TARS, "*.tar"))):
15
+ print(" extract", os.path.basename(t), flush=True)
16
+ with tarfile.open(t) as tf:
17
+ tf.extractall(BASE)
18
+
19
+ rd = os.path.join(TARS, "README.md")
20
+ if os.path.isfile(rd):
21
+ os.replace(rd, os.path.join(BASE, "README.md"))
22
+
23
+ print("[3] cleanup tars ...", flush=True)
24
+ for t in glob.glob(os.path.join(TARS, "*.tar")):
25
+ os.remove(t)
26
+ try:
27
+ os.rmdir(TARS)
28
+ except OSError:
29
+ pass
30
+
31
+ print("DONE_DATA", flush=True)
code/scripts/h800_parallel_extract.sh ADDED
@@ -0,0 +1,38 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env bash
2
+ # Parallel extraction of the remaining GenSegDataset tars on h800's slow network share.
3
+ # Big archives are split by member-list into N chunks, each extracted by a separate
4
+ # `tar -x -T <chunk>` process, to saturate the share's parallel small-file throughput.
5
+ set -u
6
+ BASE=/mnt/tidal-alsh-share2/dataset/qinshengqian/research/c3/NPJ-ACM/Data
7
+ TARS=$BASE/_tars
8
+ WORK=/tmp/pextract
9
+ mkdir -p "$WORK"
10
+
11
+ # dataset -> parallel chunk count (kits19 = most files)
12
+ launch_ds() {
13
+ local ds=$1 n=$2 tar="$TARS/$1.tar"
14
+ [ -f "$tar" ] || { echo "MISSING $tar"; return; }
15
+ if [ "$n" -le 1 ]; then
16
+ tar -xf "$tar" -C "$BASE" &
17
+ else
18
+ tar -tf "$tar" | grep -v '/$' > "$WORK/$ds.list"
19
+ split -n "l/$n" -d "$WORK/$ds.list" "$WORK/$ds.chunk."
20
+ for c in "$WORK/$ds.chunk."*; do
21
+ tar -xf "$tar" -C "$BASE" -T "$c" &
22
+ done
23
+ fi
24
+ }
25
+
26
+ echo "[start] $(date +%T) launching parallel extraction"
27
+ launch_ds medsegdb_kits19 8
28
+ launch_ds pannuke_semantic 4
29
+ launch_ds refuge2 1
30
+ echo "launched $(jobs -p | wc -l) parallel tar streams"
31
+ wait
32
+ echo "PEXTRACT_DONE $(date +%T)"
33
+
34
+ # cleanup tars + work
35
+ rm -f "$TARS"/*.tar
36
+ rmdir "$TARS" 2>/dev/null || true
37
+ rm -rf "$WORK"
38
+ echo "CLEANUP_DONE $(date +%T)"
code/scripts/h800_run_unified512.py ADDED
@@ -0,0 +1,80 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """8-GPU pool runner for the UNIFIED-512 conv retrain on h800 (L20Y x8).
2
+ Mirrors the baseline grid: 4 conv archs x the existing (dataset,protocol,seed) cells,
3
+ all at img_size 512. 1 job per GPU, resumable (skips cells whose metrics.json exists),
4
+ per-job log. occupy.py auto-yields. Run detached; tail /tmp/unified512_runner.log.
5
+ """
6
+ import os, sys, subprocess, time
7
+
8
+ CODE = "/mnt/tidal-alsh-share2/dataset/qinshengqian/research/c3/NPJ-ACM/Code"
9
+ DATA = "/data/temp/NPJ-ACM/Data"
10
+ WORK = "/data/temp/NPJ-ACM/work" # CWD; results -> WORK/results/unified512/...
11
+ PY = "/data/temp/miniconda3/envs/seggen/bin/python"
12
+ LOGD = WORK + "/logs_unified512"
13
+ PROXY = "http://10.140.15.68:3128"
14
+ NGPU, IMG, BATCH, EPOCHS = 8, 512, 8, 300
15
+ ARCHS = ["unet", "unetpp", "deeplabv3plus", "attention_unet"]
16
+ CELLS = [ # (dataset, protocol, [seeds]) -- matches existing baseline structure
17
+ ("acdc_png", "official", [0, 1, 2]),
18
+ ("busi", "fold01", [0, 1, 2]),
19
+ ("cvc_clinicdb", "official", [0, 1, 2]),
20
+ ("fives", "official", [0, 1, 2]),
21
+ ("idridd_segmentation", "fold01", [0, 1, 2]),
22
+ ("kvasir_seg", "official", [0, 1, 2]),
23
+ ("medsegdb_isic2018", "holdout", [0, 1, 2]),
24
+ ("medsegdb_kits19", "fold01", [0, 1, 2]),
25
+ ("refuge2", "official", [0, 1, 2]),
26
+ ("pannuke_semantic", "fold01", [0, 1, 2]),
27
+ ("pannuke_semantic", "fold02", [0]),
28
+ ("pannuke_semantic", "fold03", [0]),
29
+ ]
30
+ os.makedirs(LOGD, exist_ok=True)
31
+ os.makedirs(WORK, exist_ok=True)
32
+
33
+ jobs = []
34
+ for ds, proto, seeds in CELLS:
35
+ for arch in ARCHS:
36
+ for s in seeds:
37
+ out = f"{WORK}/results/unified512/{ds}_{proto}/{arch}/seed{s}/metrics.json"
38
+ jobs.append({"ds": ds, "proto": proto, "arch": arch, "seed": s, "out": out,
39
+ "tag": f"{ds}_{proto}_{arch}_s{s}"})
40
+
41
+ def make_cmd(j, gpu):
42
+ return (
43
+ f"export CUDA_DEVICE_ORDER=PCI_BUS_ID CUDA_VISIBLE_DEVICES={gpu} "
44
+ f"OMP_NUM_THREADS=8 MKL_NUM_THREADS=8 OPENBLAS_NUM_THREADS=8 "
45
+ f"https_proxy={PROXY} http_proxy={PROXY} && cd {WORK} && "
46
+ f"{PY} {CODE}/framework/train.py --data_root {DATA} --dataset {j['ds']} --protocol {j['proto']} "
47
+ f"--arch {j['arch']} --img_size {IMG} --batch_size {BATCH} --num_workers 8 --amp bf16 "
48
+ f"--exp_name unified512 --seed {j['seed']} --no-visualize --encoder resnet50 --encoder_weights imagenet "
49
+ f"--epochs {EPOCHS} && "
50
+ f"{PY} {CODE}/framework/test.py --data_root {DATA} --dataset {j['ds']} --protocol {j['proto']} "
51
+ f"--arch {j['arch']} --img_size {IMG} --exp_name unified512 --seed {j['seed']} --encoder resnet50"
52
+ )
53
+
54
+ pending = [j for j in jobs if not os.path.isfile(j["out"])]
55
+ print(f"[runner] total={len(jobs)} done={len(jobs)-len(pending)} pending={len(pending)}", flush=True)
56
+
57
+ running = {} # gpu -> (Popen, job, start)
58
+ free = list(range(NGPU))
59
+ done = ok = fail = 0
60
+ i = 0
61
+ while i < len(pending) or running:
62
+ while free and i < len(pending):
63
+ gpu = free.pop(0)
64
+ j = pending[i]; i += 1
65
+ lf = open(f"{LOGD}/{j['tag']}.log", "w")
66
+ p = subprocess.Popen(["bash", "-lc", make_cmd(j, gpu)], stdout=lf, stderr=subprocess.STDOUT)
67
+ running[gpu] = (p, j, time.time(), lf)
68
+ print(f"[launch] gpu{gpu} {j['tag']}", flush=True)
69
+ time.sleep(20)
70
+ for gpu, (p, j, st, lf) in list(running.items()):
71
+ if p.poll() is not None:
72
+ lf.close(); done += 1
73
+ okj = os.path.isfile(j["out"])
74
+ ok += okj; fail += (not okj)
75
+ mins = (time.time() - st) / 60
76
+ print(f"[finish] gpu{gpu} {j['tag']} rc={p.returncode} ok={okj} "
77
+ f"{mins:.0f}min ({done}/{len(pending)} done, {fail} failed)", flush=True)
78
+ del running[gpu]; free.append(gpu); free.sort()
79
+ print(f"[runner] ALL DONE. ok={ok} fail={fail} of {len(pending)}", flush=True)
80
+ print("UNIFIED512_RUNNER_DONE", flush=True)
code/scripts/h800_setup_seggen.sh ADDED
@@ -0,0 +1,28 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env bash
2
+ # Build the seggen conda env on h800 (L20Y / CUDA 12.8). torch installed FIRST (cu128)
3
+ # so SMP/MONAI don't pull a mismatched torch. Run detached; log to /tmp/seggen_env.log.
4
+ set -e
5
+ export https_proxy=http://10.140.15.68:3128 http_proxy=http://10.140.15.68:3128
6
+ CONDA=/data/temp/miniconda3
7
+ PROXY=http://10.140.15.68:3128
8
+
9
+ echo "[1] create env (python 3.11) -- conda-forge only, avoids defaults-channel ToS block"
10
+ $CONDA/bin/conda create -y -n seggen -c conda-forge --override-channels python=3.11 pip
11
+
12
+ PIP="$CONDA/envs/seggen/bin/pip"
13
+ echo "[2] torch cu128 (host CUDA 12.8, >2.6)"
14
+ $PIP install --proxy "$PROXY" torch torchvision --index-url https://download.pytorch.org/whl/cu128
15
+
16
+ echo "[3] seg stack (torch already present -> no reinstall)"
17
+ $PIP install --proxy "$PROXY" \
18
+ segmentation-models-pytorch albumentations==2.0.8 monai medpy \
19
+ opencv-python-headless numpy pyyaml timm einops ml-collections tqdm \
20
+ diffusers==0.21.4 datasets==2.14.5
21
+
22
+ echo "[4] verify"
23
+ $CONDA/envs/seggen/bin/python - <<'PY'
24
+ import torch, segmentation_models_pytorch, monai, albumentations, timm, cv2
25
+ print("torch", torch.__version__, "| cuda", torch.version.cuda, "| avail", torch.cuda.is_available(), "| ndev", torch.cuda.device_count())
26
+ print("smp ok, monai ok, albumentations ok, timm ok, cv2", cv2.__version__)
27
+ PY
28
+ echo "SEGGEN_ENV_DONE"
code/scripts/h800_swin_transunet_eval512.py ADDED
@@ -0,0 +1,68 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Phase-1 of unified-512 re-eval on h800: re-score SwinUNet/TransUNet at eval_size 512.
2
+ These are res-locked (224/256) so we DON'T retrain — we load their HF-curated best-seed
3
+ weights, run at native input, resize preds+GT to 512, write metrics.json into the
4
+ unified512 results tree. 12 cells x {swinunet, transunet} = 24 evals, 8-GPU pool.
5
+ """
6
+ import os, glob, shutil, subprocess, time
7
+
8
+ CODE = "/mnt/tidal-alsh-share2/dataset/qinshengqian/research/c3/NPJ-ACM/Code"
9
+ DATA = "/data/temp/NPJ-ACM/Data"
10
+ WORK = "/data/temp/NPJ-ACM/work"
11
+ PY = "/data/temp/miniconda3/envs/seggen/bin/python"
12
+ HFW = WORK + "/hf_weights/weights/framework" # <cell>/{swinunet,transunet}.pth
13
+ RES = WORK + "/results/unified512" # eval_at_res writes here (out_root=results rel to WORK)
14
+ LOGD = WORK + "/logs_eval512"; os.makedirs(LOGD, exist_ok=True)
15
+ PROXY = "http://10.140.15.68:3128"
16
+ PROTOS = ["official", "holdout", "fold01", "fold02", "fold03"]
17
+ NGPU = 8
18
+
19
+ def split_cell(cell):
20
+ for p in PROTOS:
21
+ if cell.endswith("_" + p):
22
+ return cell[:-(len(p) + 1)], p
23
+ raise ValueError("bad cell " + cell)
24
+
25
+ jobs = []
26
+ for cell_dir in sorted(glob.glob(HFW + "/*")):
27
+ cell = os.path.basename(cell_dir)
28
+ ds, proto = split_cell(cell)
29
+ for arch in ("swinunet", "transunet"):
30
+ w = f"{cell_dir}/{arch}.pth"
31
+ if not os.path.isfile(w):
32
+ continue
33
+ out = f"{RES}/{cell}/{arch}/seed0"
34
+ jobs.append({"ds": ds, "proto": proto, "arch": arch, "w": w, "out": out,
35
+ "tag": f"{cell}_{arch}", "mj": f"{out}/metrics.json"})
36
+
37
+ pending = [j for j in jobs if not os.path.isfile(j["mj"])]
38
+ print(f"[eval512] total={len(jobs)} done={len(jobs)-len(pending)} pending={len(pending)}", flush=True)
39
+
40
+ def make_cmd(j, gpu):
41
+ enc = "R50-ViT-B_16" if j["arch"] == "transunet" else "resnet50"
42
+ # place the HF weight as best.pth where eval_at_res.py expects it
43
+ os.makedirs(j["out"], exist_ok=True)
44
+ shutil.copy(j["w"], j["out"] + "/best.pth")
45
+ return (
46
+ f"export CUDA_DEVICE_ORDER=PCI_BUS_ID CUDA_VISIBLE_DEVICES={gpu} "
47
+ f"OMP_NUM_THREADS=8 MKL_NUM_THREADS=8 OPENBLAS_NUM_THREADS=8 "
48
+ f"https_proxy={PROXY} http_proxy={PROXY} && cd {WORK} && "
49
+ f"{PY} {CODE}/framework/eval_at_res.py --data_root {DATA} --dataset {j['ds']} "
50
+ f"--protocol {j['proto']} --arch {j['arch']} --seed 0 --eval_size 512 "
51
+ f"--exp_name unified512 --encoder {enc}"
52
+ )
53
+
54
+ running = {}; free = list(range(NGPU)); i = 0; ok = fail = 0
55
+ while i < len(pending) or running:
56
+ while free and i < len(pending):
57
+ gpu = free.pop(0); j = pending[i]; i += 1
58
+ lf = open(f"{LOGD}/{j['tag']}.log", "w")
59
+ p = subprocess.Popen(["bash", "-lc", make_cmd(j, gpu)], stdout=lf, stderr=subprocess.STDOUT)
60
+ running[gpu] = (p, j, lf); print(f"[launch] gpu{gpu} {j['tag']}", flush=True)
61
+ time.sleep(8)
62
+ for gpu, (p, j, lf) in list(running.items()):
63
+ if p.poll() is not None:
64
+ lf.close(); okj = os.path.isfile(j["mj"]); ok += okj; fail += (not okj)
65
+ print(f"[finish] gpu{gpu} {j['tag']} rc={p.returncode} ok={okj}", flush=True)
66
+ del running[gpu]; free.append(gpu); free.sort()
67
+ print(f"[eval512] ALL DONE ok={ok} fail={fail}", flush=True)
68
+ print("SWIN_TRANSUNET_EVAL512_DONE", flush=True)
code/scripts/hf_update_unified512.py ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Push the unified-512 deliverables to HF MaybeRichard/GenSeg-Baselines (replacing the
2
+ old 256 report): enhanced aggregate.py (code) + unified-512 summary.* + all metrics.json
3
+ (results/). Run with the write token + the REAL HF endpoint (local env defaults to the
4
+ hf-mirror download mirror, which does NOT accept authenticated writes).
5
+ HF_ENDPOINT=https://huggingface.co T=<write-token> python3 scripts/hf_update_unified512.py
6
+ """
7
+ import os
8
+ from huggingface_hub import HfApi
9
+
10
+ REPO = "MaybeRichard/GenSeg-Baselines"
11
+ LOCAL = "/home/richard/Documents/Code/ZJU/SegGen"
12
+ api = HfApi(token=os.environ["T"], endpoint="https://huggingface.co")
13
+
14
+ print("[1] code: enhanced aggregate.py")
15
+ api.upload_file(path_or_fileobj=f"{LOCAL}/framework/report/aggregate.py", repo_id=REPO,
16
+ repo_type="model", path_in_repo="code/framework/report/aggregate.py",
17
+ commit_message="aggregate.py: unified-512 intro + per-class + Wilcoxon")
18
+
19
+ print("[2] results: unified-512 summary.* + all metrics.json -> results/")
20
+ api.upload_folder(folder_path=f"{LOCAL}/results/unified512", repo_id=REPO, repo_type="model",
21
+ path_in_repo="results",
22
+ allow_patterns=["**/metrics.json", "summary.html", "summary.csv",
23
+ "summary.md", "summary.tex"],
24
+ commit_message="results: unified-512 (8 methods @512, per-class + significance)")
25
+
26
+ print("DONE")
code/scripts/hf_upload_gensegdataset.py ADDED
@@ -0,0 +1,23 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Upload GenSegDataset (processed_unified mirror + dataset card) to the Hugging
2
+ Face Hub as a PRIVATE dataset repo. Storage is content-addressed, so fold-duplicated
3
+ images collapse to unique blobs. Resumable via upload_large_folder."""
4
+ import sys
5
+ from huggingface_hub import HfApi
6
+
7
+ REPO = "MaybeRichard/GenSegDataset"
8
+ DATA = "/home/wzhang/LSC/Dataset/Segmentation/processed_unified"
9
+ CARD = "/home/wzhang/LSC/Dataset/Segmentation/GenSegDataset_README.md"
10
+
11
+ api = HfApi()
12
+ api.create_repo(REPO, repo_type="dataset", private=True, exist_ok=True)
13
+ print("repo ready:", REPO, flush=True)
14
+
15
+ # big data upload (resumable, dedups blobs, parallel)
16
+ api.upload_large_folder(repo_id=REPO, repo_type="dataset", folder_path=DATA)
17
+ print("folder uploaded", flush=True)
18
+
19
+ # dataset card last so it is the final README at the repo root
20
+ api.upload_file(path_or_fileobj=CARD, path_in_repo="README.md",
21
+ repo_id=REPO, repo_type="dataset",
22
+ commit_message="add dataset card")
23
+ print("UPLOAD_DONE", flush=True)
code/scripts/hf_upload_tars.py ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Recovery upload: ship GenSegDataset as ONE tar per subset (10 files) instead of
2
+ ~110k loose PNGs, to stay under HF's 128-commit/hour limit. Resets the partially
3
+ populated repo, uploads <subset>.tar + the dataset card."""
4
+ import os, subprocess
5
+ from huggingface_hub import HfApi
6
+
7
+ REPO = "MaybeRichard/GenSegDataset"
8
+ DATA = "/home/wzhang/LSC/Dataset/Segmentation/processed_unified"
9
+ TARS = "/home/wzhang/LSC/Dataset/Segmentation/hf_tars"
10
+ CARD = "/home/wzhang/LSC/Dataset/Segmentation/GenSegDataset_README.md"
11
+
12
+ os.makedirs(TARS, exist_ok=True)
13
+ for ds in sorted(os.listdir(DATA)):
14
+ if not os.path.isdir(os.path.join(DATA, ds)):
15
+ continue
16
+ out = os.path.join(TARS, ds + ".tar")
17
+ if os.path.exists(out) and os.path.getsize(out) > 0:
18
+ print("skip (exists):", ds, flush=True); continue
19
+ # -h dereferences symlinks so fold-shared images are materialized into the tar
20
+ subprocess.run(["tar", "-chf", out, "-C", DATA, ds], check=True)
21
+ print("tarred %s -> %.1f MB" % (ds, os.path.getsize(out) / 1e6), flush=True)
22
+
23
+ api = HfApi()
24
+ api.delete_repo(REPO, repo_type="dataset", missing_ok=True)
25
+ api.create_repo(REPO, repo_type="dataset", private=True, exist_ok=True)
26
+ print("repo reset:", REPO, flush=True)
27
+
28
+ api.upload_file(path_or_fileobj=CARD, path_in_repo="README.md",
29
+ repo_id=REPO, repo_type="dataset", commit_message="dataset card")
30
+ api.upload_large_folder(repo_id=REPO, repo_type="dataset", folder_path=TARS)
31
+ print("UPLOAD_DONE", flush=True)
code/scripts/p1/backbones.py ADDED
@@ -0,0 +1,44 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """build_backbone: instantiate one of {jit, pixelgen, deco, pixeldit} pixel-space
2
+ denoisers for mask-concat conditioning. Each returns a net callable as net(x, t, y)
3
+ -> (N, C>=img_ch, H, W); the caller slices [:, :img_ch] (backbone-agnostic decouple).
4
+ in_channels = img_channels + cond_channels; a single dummy class (num_classes=1).
5
+ Unified 'Base' tier (~130-150M) for the P1 backbone bake-off (P1 = native arch under a
6
+ common flow-matching objective; perceptual/DCT/FD losses are P2 levers, not used here)."""
7
+ import os
8
+ import sys
9
+
10
+ _SOTA = "/home/wzhang/LSC/Code/NPJ/sota"
11
+
12
+
13
+ def _add(path):
14
+ if path not in sys.path:
15
+ sys.path.insert(0, path)
16
+
17
+
18
+ def build_backbone(backbone: str, model_name: str, img_size: int,
19
+ in_channels: int, num_classes: int = 1):
20
+ bk = backbone.lower()
21
+ if bk == "jit":
22
+ _add(os.path.join(_SOTA, "JiT"))
23
+ from model_jit import JiT_models
24
+ return JiT_models[model_name](input_size=img_size, in_channels=in_channels,
25
+ num_classes=num_classes)
26
+ if bk == "pixelgen":
27
+ _add(os.path.join(_SOTA, "PixelGen", "src", "models", "transformer"))
28
+ import importlib
29
+ jit = importlib.import_module("JiT") # PixelGen's self-contained JiT.py
30
+ return jit.JiT_models[model_name](input_size=img_size, in_channels=in_channels,
31
+ num_classes=num_classes)
32
+ if bk == "deco":
33
+ _add(os.path.join(_SOTA, "DeCo", "src", "models", "transformer"))
34
+ from dit_c2i_DeCo import PixNerDiT
35
+ return PixNerDiT(in_channels=in_channels, patch_size=16, num_groups=12,
36
+ hidden_size=768, hidden_size_x=32, num_blocks=13,
37
+ num_cond_blocks=12, num_classes=num_classes)
38
+ if bk == "pixeldit":
39
+ _add(os.path.join(_SOTA, "PixelDiT"))
40
+ from pixdit_core.pixeldit_c2i import PixDiT
41
+ return PixDiT(in_channels=in_channels, num_groups=10, hidden_size=640,
42
+ pixel_hidden_size=16, patch_depth=9, pixel_depth=4,
43
+ patch_size=16, num_classes=num_classes)
44
+ raise ValueError(f"unknown backbone: {backbone}")
code/scripts/p1/fd_lever.py ADDED
@@ -0,0 +1,99 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """FD-lever ablation on the recommended P2 base (JiT): refine p1_jit_{ds} with FD loss
2
+ -> sample -> downstream. Compares +JiT-FD vs +JiT(native) vs real. 18 jobs on GPU0-5."""
3
+ import os, time, json, subprocess
4
+
5
+ ROOT = "/home/wzhang/LSC/Code/NPJ"
6
+ DR = "/home/wzhang/LSC/Dataset/Segmentation/processed_unified"
7
+ PY = "/opt/anaconda3/envs/seggen/bin/python"
8
+ GPUS = [0, 1, 2, 3, 4, 5]
9
+ os.chdir(ROOT)
10
+ LOGD = os.path.join(ROOT, "logs", "fdlever")
11
+ os.makedirs(LOGD, exist_ok=True)
12
+ def log(m):
13
+ line = f"[{time.strftime('%F %T')}] {m}"
14
+ open(os.path.join(LOGD, "status.md"), "a").write(line + "\n"); print(line, flush=True)
15
+
16
+ DSETS = {"isic": ("medsegdb_isic2018", "holdout", 2582), "kvasir": ("kvasir_seg", "official", 800)}
17
+ NS = [50, 100]; SEEDS = [0, 1, 2]
18
+ jobs = {}
19
+ def add(jid, cmd, deps=(), done_path=None, done_min=1):
20
+ jobs[jid] = {"cmd": cmd, "deps": list(deps), "done_path": done_path, "done_min": done_min,
21
+ "state": "pending", "tries": 0, "gpu": None}
22
+
23
+ for dk, (ds, proto, tot) in DSETS.items():
24
+ base = f"pretrained/pixdiff/p1_jit_{dk}.pt"
25
+ out = f"pretrained/pixdiff/p1_jitfd_{dk}.pt"
26
+ cmd = (f"{PY} -m framework.synth.pixdiff.train_fd --base_ckpt {base} --data_root {DR} "
27
+ f"--dataset {ds} --protocol {proto} --train_fraction 1.0 --epochs 150 --batch_size 16 "
28
+ f"--amp bf16 --fd_weight 0.5 --out_ckpt {out} --log_interval 100")
29
+ add(f"genfd_{dk}", cmd, done_path=os.path.join(ROOT, out))
30
+ for N in NS:
31
+ f = N / tot
32
+ sd = f"{DR}/{ds}/{proto}/synth_p1_jitfd_{dk}_f{N}"
33
+ cmd = (f"{PY} -m framework.synth.pixdiff.sample --ckpt {out} --data_root {DR} --dataset {ds} "
34
+ f"--protocol {proto} --train_fraction {f} --fraction_seed 0 --n_per_mask 4 --mask_aug "
35
+ f"--num_steps 50 --out_dir {sd}")
36
+ add(f"samp_jitfd_{dk}_N{N}", cmd, deps=[f"genfd_{dk}"], done_path=os.path.join(sd, "images"), done_min=N * 4)
37
+ for S in SEEDS:
38
+ exp = f"p1_jitfd_{dk}_N{N}"
39
+ mp = os.path.join(ROOT, f"results/{exp}/{ds}_{proto}/unet/seed{S}/metrics.json")
40
+ cmd = (f"{PY} framework/train.py --data_root {DR} --dataset {ds} --protocol {proto} --arch unet "
41
+ f"--encoder resnet50 --aug standard --epochs 400 --train_fraction {f} --fraction_seed 0 "
42
+ f"--synth_train_dir {sd} --exp_name {exp} --amp bf16 --seed {S} "
43
+ f"&& {PY} framework/test.py --data_root {DR} --dataset {ds} --protocol {proto} --arch unet "
44
+ f"--encoder resnet50 --aug standard --exp_name {exp} --seed {S}")
45
+ add(f"seg_jitfd_{dk}_N{N}_s{S}", cmd, deps=[f"samp_jitfd_{dk}_N{N}"], done_path=mp)
46
+
47
+ def is_done(j):
48
+ p = j["done_path"]
49
+ if not p or not os.path.exists(p): return False
50
+ if os.path.isdir(p):
51
+ try: return len(os.listdir(p)) >= j["done_min"]
52
+ except OSError: return False
53
+ return True
54
+ def aggregate():
55
+ res = {}
56
+ for dk, (ds, proto, tot) in DSETS.items():
57
+ for N in NS:
58
+ exp = f"p1_jitfd_{dk}_N{N}"; ious = []; dices = []
59
+ for S in SEEDS:
60
+ mp = f"results/{exp}/{ds}_{proto}/unet/seed{S}/metrics.json"
61
+ if os.path.exists(mp):
62
+ try:
63
+ m = json.load(open(mp))["metrics"]; ious.append(m["iou_mean"]); dices.append(m["dice_mean"])
64
+ except Exception: pass
65
+ if ious:
66
+ res[f"{dk}_N{N}_jitfd"] = {"iou_mean": sum(ious) / len(ious), "dice_mean": sum(dices) / len(dices),
67
+ "n_seeds": len(ious), "iou_seeds": ious}
68
+ json.dump(res, open(os.path.join(LOGD, "fd_results.json"), "w"), indent=2)
69
+
70
+ for jid, j in jobs.items():
71
+ if is_done(j): j["state"] = "done"
72
+ def deps_done(j): return all(jobs[d]["state"] == "done" for d in j["deps"])
73
+ running = {}; free = set(GPUS); last = 0
74
+ log(f"START {len(jobs)} jobs on {GPUS} ({sum(1 for j in jobs.values() if j['state']=='done')} pre-done)")
75
+ while True:
76
+ if all(j["state"] in ("done", "failed") for j in jobs.values()): break
77
+ for jid, j in jobs.items():
78
+ if not free: break
79
+ if j["state"] == "pending" and deps_done(j):
80
+ if is_done(j): j["state"] = "done"; continue
81
+ g = free.pop()
82
+ env = dict(os.environ, CUDA_DEVICE_ORDER="PCI_BUS_ID", CUDA_VISIBLE_DEVICES=str(g),
83
+ TORCHDYNAMO_DISABLE="1", PYTHONPATH=".", OMP_NUM_THREADS="4")
84
+ lf = open(os.path.join(LOGD, jid + ".log"), "a")
85
+ p = subprocess.Popen(j["cmd"], shell=True, env=env, stdout=lf, stderr=subprocess.STDOUT, cwd=ROOT)
86
+ running[g] = (jid, p, lf); j["state"] = "running"; j["gpu"] = g; j["tries"] += 1
87
+ log(f"LAUNCH {jid} GPU{g} try{j['tries']}")
88
+ for g, (jid, p, lf) in list(running.items()):
89
+ rc = p.poll()
90
+ if rc is None: continue
91
+ lf.close(); del running[g]; free.add(g); j = jobs[jid]
92
+ if is_done(j): j["state"] = "done"; log(f"DONE {jid}")
93
+ elif j["tries"] < 2: j["state"] = "pending"; log(f"RETRY {jid} rc={rc}")
94
+ else: j["state"] = "failed"; log(f"FAILED {jid} rc={rc}")
95
+ if time.time() - last > 180:
96
+ cnt = {s: sum(1 for j in jobs.values() if j["state"] == s) for s in ("done", "running", "pending", "failed")}
97
+ log(f"SUMMARY {cnt}"); aggregate(); last = time.time()
98
+ time.sleep(10)
99
+ aggregate(); log("ALL DONE"); print("FD_LEVER_DONE", flush=True)
code/scripts/p1/fd_results.json ADDED
@@ -0,0 +1,42 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "isic_N50_jitfd": {
3
+ "iou_mean": 0.807111643913793,
4
+ "dice_mean": 0.8785879691414072,
5
+ "n_seeds": 3,
6
+ "iou_seeds": [
7
+ 0.8075343832113462,
8
+ 0.8052902472456404,
9
+ 0.8085103012843926
10
+ ]
11
+ },
12
+ "isic_N100_jitfd": {
13
+ "iou_mean": 0.8194000537587307,
14
+ "dice_mean": 0.8883052260740664,
15
+ "n_seeds": 3,
16
+ "iou_seeds": [
17
+ 0.815320694079291,
18
+ 0.8190540649488038,
19
+ 0.8238254022480974
20
+ ]
21
+ },
22
+ "kvasir_N50_jitfd": {
23
+ "iou_mean": 0.7743735031913296,
24
+ "dice_mean": 0.8502822525848202,
25
+ "n_seeds": 3,
26
+ "iou_seeds": [
27
+ 0.7818383912178852,
28
+ 0.7792812954940015,
29
+ 0.7620008228621021
30
+ ]
31
+ },
32
+ "kvasir_N100_jitfd": {
33
+ "iou_mean": 0.8237618615778238,
34
+ "dice_mean": 0.8922327796820345,
35
+ "n_seeds": 3,
36
+ "iou_seeds": [
37
+ 0.8173440030838561,
38
+ 0.8195464262764907,
39
+ 0.8343951553731246
40
+ ]
41
+ }
42
+ }
code/scripts/p1/fid_and_viz.py ADDED
@@ -0,0 +1,130 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """FID (1-2k samples) per backbone x dataset + clear same-mask aligned viz.
2
+ A) fid-sample: train_fraction=1.0, mask_aug, n_per_mask -> ~1.6-2.6k synth; FID vs real train.
3
+ B) align-sample: f50 masks, NO aug, 1/mask -> all backbones share identical real masks -> aligned grid.
4
+ Then pytorch_fid per pair + build [mask|real|4 backbones] grids. GPU0-5 pool."""
5
+ import os, time, json, re, subprocess
6
+ import numpy as np
7
+ import matplotlib; matplotlib.use("Agg")
8
+ import matplotlib.pyplot as plt
9
+ from PIL import Image
10
+
11
+ ROOT = "/home/wzhang/LSC/Code/NPJ"; DR = "/home/wzhang/LSC/Dataset/Segmentation/processed_unified"
12
+ PY = "/opt/anaconda3/envs/seggen/bin/python"; GPUS = [0, 1, 2, 3, 4, 5]
13
+ os.chdir(ROOT); LOGD = os.path.join(ROOT, "logs", "fidviz"); os.makedirs(LOGD, exist_ok=True)
14
+ def log(m):
15
+ line = f"[{time.strftime('%F %T')}] {m}"; open(os.path.join(LOGD, "status.md"), "a").write(line + "\n"); print(line, flush=True)
16
+ # (ds, proto, total, npm_for_fid)
17
+ DSETS = {"isic": ("medsegdb_isic2018", "holdout", 2582, 1),
18
+ "kvasir": ("kvasir_seg", "official", 800, 2),
19
+ "busi": ("busi", "fold01", 545, 3)}
20
+ BKS = ["jit", "pixelgen", "deco", "pixeldit"]; LAB = {"jit": "JiT", "pixelgen": "PixelGen", "deco": "DeCo", "pixeldit": "PixelDiT"}
21
+ jobs = {}
22
+ def add(jid, cmd, deps=(), done_path=None, done_min=1):
23
+ jobs[jid] = {"cmd": cmd, "deps": list(deps), "done_path": done_path, "done_min": done_min, "state": "pending", "tries": 0, "gpu": None}
24
+
25
+ for bk in BKS:
26
+ for dk, (ds, proto, tot, npm) in DSETS.items():
27
+ ck = f"pretrained/pixdiff/p1_{bk}_{dk}.pt"
28
+ fsd = f"{DR}/{ds}/{proto}/synth_fid_{bk}_{dk}"
29
+ add(f"fidsamp_{bk}_{dk}",
30
+ f"{PY} -m framework.synth.pixdiff.sample --ckpt {ck} --data_root {DR} --dataset {ds} --protocol {proto} "
31
+ f"--train_fraction 1.0 --fraction_seed 0 --n_per_mask {npm} --mask_aug --num_steps 50 --out_dir {fsd}",
32
+ done_path=os.path.join(fsd, "images"), done_min=int(0.8 * tot * npm))
33
+ real = f"{DR}/{ds}/{proto}/train/images"
34
+ flog = os.path.join(LOGD, f"fid_{bk}_{dk}.log"); fok = os.path.join(LOGD, f"fid_{bk}_{dk}.ok")
35
+ add(f"fid_{bk}_{dk}",
36
+ f"{PY} -m pytorch_fid {real} {fsd}/images --device cuda --batch-size 64 > {flog} 2>&1 && grep -q FID {flog} && touch {fok}",
37
+ deps=[f"fidsamp_{bk}_{dk}"], done_path=fok)
38
+ f50 = 50 / tot; asd = f"{DR}/{ds}/{proto}/synth_align_{bk}_{dk}"
39
+ add(f"alignsamp_{bk}_{dk}",
40
+ f"{PY} -m framework.synth.pixdiff.sample --ckpt {ck} --data_root {DR} --dataset {ds} --protocol {proto} "
41
+ f"--train_fraction {f50} --fraction_seed 0 --n_per_mask 1 --num_steps 50 --out_dir {asd}",
42
+ done_path=os.path.join(asd, "images"), done_min=40)
43
+
44
+ def is_done(j):
45
+ p = j["done_path"]
46
+ if not p or not os.path.exists(p): return False
47
+ if os.path.isdir(p):
48
+ try: return len(os.listdir(p)) >= j["done_min"]
49
+ except OSError: return False
50
+ return True
51
+ for jid, j in jobs.items():
52
+ if is_done(j): j["state"] = "done"
53
+ def deps_done(j): return all(jobs[d]["state"] == "done" for d in j["deps"])
54
+ running = {}; free = set(GPUS); last = 0
55
+ log(f"START {len(jobs)} jobs on {GPUS}")
56
+ while True:
57
+ if all(j["state"] in ("done", "failed") for j in jobs.values()): break
58
+ for jid, j in jobs.items():
59
+ if not free: break
60
+ if j["state"] == "pending" and deps_done(j):
61
+ if is_done(j): j["state"] = "done"; continue
62
+ g = free.pop()
63
+ env = dict(os.environ, CUDA_DEVICE_ORDER="PCI_BUS_ID", CUDA_VISIBLE_DEVICES=str(g), TORCHDYNAMO_DISABLE="1", PYTHONPATH=".", OMP_NUM_THREADS="4")
64
+ lf = open(os.path.join(LOGD, jid + ".log"), "a")
65
+ p = subprocess.Popen(j["cmd"], shell=True, env=env, stdout=lf, stderr=subprocess.STDOUT, cwd=ROOT)
66
+ running[g] = (jid, p, lf); j["state"] = "running"; j["gpu"] = g; j["tries"] += 1
67
+ log(f"LAUNCH {jid} GPU{g} try{j['tries']}")
68
+ for g, (jid, p, lf) in list(running.items()):
69
+ rc = p.poll()
70
+ if rc is None: continue
71
+ lf.close(); del running[g]; free.add(g); j = jobs[jid]
72
+ if is_done(j): j["state"] = "done"; log(f"DONE {jid}")
73
+ elif j["tries"] < 2: j["state"] = "pending"; log(f"RETRY {jid} rc={rc}")
74
+ else: j["state"] = "failed"; log(f"FAILED {jid} rc={rc}")
75
+ if time.time() - last > 180:
76
+ cnt = {s: sum(1 for j in jobs.values() if j["state"] == s) for s in ("done", "running", "pending", "failed")}; log(f"SUMMARY {cnt}"); last = time.time()
77
+ time.sleep(8)
78
+
79
+ # ---- parse FID ----
80
+ fid = {}
81
+ for bk in BKS:
82
+ for dk in DSETS:
83
+ lg = os.path.join(LOGD, f"fid_{bk}_{dk}.log")
84
+ if os.path.exists(lg):
85
+ m = re.findall(r"FID:\s*([0-9.]+)", open(lg).read())
86
+ if m: fid[f"{dk}_{bk}"] = float(m[-1])
87
+ json.dump(fid, open(os.path.join(LOGD, "fid_results.json"), "w"), indent=2)
88
+ log(f"FID: {fid}")
89
+
90
+ # ---- aligned grids ([mask | real | 4 backbones], same real mask per column) ----
91
+ def rgb(p): return np.asarray(Image.open(p).convert("RGB").resize((256, 256)))
92
+ def gray(p): return np.asarray(Image.open(p).convert("L").resize((256, 256)))
93
+ def fmap(d):
94
+ p = os.path.join(d, "images"); m = {}
95
+ if os.path.isdir(p):
96
+ for f in sorted(os.listdir(p)):
97
+ if f.endswith(".png"): m.setdefault(f[:-4].split("__")[0], os.path.join(p, f))
98
+ return m
99
+ for dk, (ds, proto, tot, npm) in DSETS.items():
100
+ base = f"{DR}/{ds}/{proto}"; ri, rm = f"{base}/train/images", f"{base}/train/masks"
101
+ maps = {bk: fmap(f"{base}/synth_align_{bk}_{dk}") for bk in BKS}
102
+ common = set(os.path.splitext(f)[0] for f in os.listdir(ri) if f.endswith(".png"))
103
+ for bk in BKS: common &= set(maps[bk].keys())
104
+ common = sorted(common); ncol = min(6, len(common))
105
+ if ncol == 0: continue
106
+ idx = [round(i * (len(common) - 1) / (ncol - 1)) for i in range(ncol)] if ncol > 1 else [0]
107
+ cases = [common[i] for i in idx]
108
+ rows = [("Conditioning mask", "mask"), ("Real image", "real")] + [(LAB[bk], bk) for bk in BKS]
109
+ fig, ax = plt.subplots(len(rows), ncol, figsize=(ncol * 1.9, len(rows) * 1.95))
110
+ for r, (labr, kind) in enumerate(rows):
111
+ for c, bs in enumerate(cases):
112
+ a = ax[r][c]
113
+ try:
114
+ mk = gray(f"{rm}/{bs}.png")
115
+ if kind == "mask":
116
+ a.imshow(mk, cmap="gray")
117
+ elif kind == "real":
118
+ a.imshow(rgb(f"{ri}/{bs}.png")); a.contour((mk > 127).astype(float), levels=[0.5], colors=["#19f04b"], linewidths=1.0)
119
+ else:
120
+ a.imshow(rgb(maps[kind][bs])); a.contour((mk > 127).astype(float), levels=[0.5], colors=["#19f04b"], linewidths=1.0)
121
+ except Exception:
122
+ a.imshow(np.ones((256, 256, 3))); a.text(0.5, 0.5, "n/a", ha="center", va="center", transform=a.transAxes, fontsize=8)
123
+ a.set_xticks([]); a.set_yticks([])
124
+ for s in a.spines.values(): s.set_visible(False)
125
+ if c == 0: a.set_ylabel(labr, fontsize=10, rotation=90, va="center", labelpad=8, color=("#111" if r < 2 else "#1a3b8b"))
126
+ fig.suptitle(f"{dk.upper()} — same-mask aligned: every backbone generates the SAME real mask (row 1)\n"
127
+ f"Row2=real image; rows 3-6=each backbone's mask-conditioned synthesis (green=that mask). Proves mask guidance.", fontsize=10)
128
+ plt.tight_layout(rect=[0.02, 0, 1, 0.94]); plt.savefig(f"/tmp/p1_aligned_{dk}.png", dpi=145, bbox_inches="tight", facecolor="white")
129
+ log(f"aligned grid saved /tmp/p1_aligned_{dk}.png")
130
+ log("ALL DONE"); print("FIDVIZ_DONE", flush=True)
code/scripts/p1/fid_fixed.py ADDED
@@ -0,0 +1,37 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Corrected FID: resize REAL train images to 256 (synth already 256) so pytorch_fid's
2
+ collate doesn't choke on variable native sizes (Kvasir/BUSI). FID(real256, synth) per pair."""
3
+ import os, re, json, subprocess
4
+ from PIL import Image
5
+
6
+ ROOT = "/home/wzhang/LSC/Code/NPJ"; DR = "/home/wzhang/LSC/Dataset/Segmentation/processed_unified"
7
+ PY = "/opt/anaconda3/envs/seggen/bin/python"
8
+ DSETS = {"isic": ("medsegdb_isic2018", "holdout"), "kvasir": ("kvasir_seg", "official"), "busi": ("busi", "fold01")}
9
+ BKS = ["jit", "pixelgen", "deco", "pixeldit"]
10
+
11
+ real256 = {}
12
+ for dk, (ds, proto) in DSETS.items():
13
+ src = f"{DR}/{ds}/{proto}/train/images"; dst = f"/tmp/real256_{dk}"; os.makedirs(dst, exist_ok=True)
14
+ for f in os.listdir(src):
15
+ if f.lower().endswith((".png", ".jpg", ".jpeg")):
16
+ o = f"{dst}/{os.path.splitext(f)[0]}.png"
17
+ if not os.path.exists(o):
18
+ Image.open(f"{src}/{f}").convert("RGB").resize((256, 256)).save(o)
19
+ real256[dk] = dst
20
+ print(f"[resize] real {dk}: {len(os.listdir(dst))} imgs", flush=True)
21
+
22
+ fid = {}
23
+ for bk in BKS:
24
+ for dk, (ds, proto) in DSETS.items():
25
+ synth = f"{DR}/{ds}/{proto}/synth_fid_{bk}_{dk}/images"
26
+ if not os.path.isdir(synth) or len(os.listdir(synth)) < 100:
27
+ print(f"[skip] {dk} {bk}: synth missing/small", flush=True); continue
28
+ env = dict(os.environ, CUDA_DEVICE_ORDER="PCI_BUS_ID", CUDA_VISIBLE_DEVICES="0")
29
+ r = subprocess.run([PY, "-m", "pytorch_fid", real256[dk], synth, "--device", "cuda", "--batch-size", "50"],
30
+ capture_output=True, text=True, env=env)
31
+ m = re.findall(r"FID:\s*([0-9.]+)", r.stdout + r.stderr)
32
+ if m:
33
+ fid[f"{dk}_{bk}"] = round(float(m[-1]), 2); print(f"[FID] {dk} {bk} = {m[-1]}", flush=True)
34
+ else:
35
+ print(f"[FAIL] {dk} {bk}: {(r.stderr or r.stdout)[-300:]}", flush=True)
36
+ json.dump(fid, open(f"{ROOT}/logs/fidviz/fid_results.json", "w"), indent=2)
37
+ print("FID_FIXED_DONE", json.dumps(fid), flush=True)
code/scripts/p1/fid_results.json ADDED
@@ -0,0 +1,14 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "isic_jit": 107.96,
3
+ "kvasir_jit": 109.09,
4
+ "busi_jit": 139.28,
5
+ "isic_pixelgen": 107.63,
6
+ "kvasir_pixelgen": 118.44,
7
+ "busi_pixelgen": 142.91,
8
+ "isic_deco": 141.84,
9
+ "kvasir_deco": 114.43,
10
+ "busi_deco": 174.36,
11
+ "isic_pixeldit": 109.87,
12
+ "kvasir_pixeldit": 112.83,
13
+ "busi_pixeldit": 151.82
14
+ }
code/scripts/p1/gen_prdc.json ADDED
@@ -0,0 +1,74 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "isic_jit": {
3
+ "precision": 0.301,
4
+ "recall": 0.107,
5
+ "density": 0.122,
6
+ "coverage": 0.146
7
+ },
8
+ "kvasir_jit": {
9
+ "precision": 0.402,
10
+ "recall": 0.021,
11
+ "density": 0.169,
12
+ "coverage": 0.294
13
+ },
14
+ "busi_jit": {
15
+ "precision": 0.271,
16
+ "recall": 0.174,
17
+ "density": 0.103,
18
+ "coverage": 0.332
19
+ },
20
+ "isic_pixelgen": {
21
+ "precision": 0.319,
22
+ "recall": 0.101,
23
+ "density": 0.122,
24
+ "coverage": 0.14
25
+ },
26
+ "kvasir_pixelgen": {
27
+ "precision": 0.388,
28
+ "recall": 0.009,
29
+ "density": 0.159,
30
+ "coverage": 0.242
31
+ },
32
+ "busi_pixelgen": {
33
+ "precision": 0.275,
34
+ "recall": 0.261,
35
+ "density": 0.093,
36
+ "coverage": 0.29
37
+ },
38
+ "isic_deco": {
39
+ "precision": 0.14,
40
+ "recall": 0.04,
41
+ "density": 0.04,
42
+ "coverage": 0.048
43
+ },
44
+ "kvasir_deco": {
45
+ "precision": 0.287,
46
+ "recall": 0.01,
47
+ "density": 0.11,
48
+ "coverage": 0.195
49
+ },
50
+ "busi_deco": {
51
+ "precision": 0.081,
52
+ "recall": 0.042,
53
+ "density": 0.023,
54
+ "coverage": 0.088
55
+ },
56
+ "isic_pixeldit": {
57
+ "precision": 0.213,
58
+ "recall": 0.065,
59
+ "density": 0.077,
60
+ "coverage": 0.106
61
+ },
62
+ "kvasir_pixeldit": {
63
+ "precision": 0.324,
64
+ "recall": 0.027,
65
+ "density": 0.123,
66
+ "coverage": 0.201
67
+ },
68
+ "busi_pixeldit": {
69
+ "precision": 0.245,
70
+ "recall": 0.051,
71
+ "density": 0.077,
72
+ "coverage": 0.182
73
+ }
74
+ }
code/scripts/p1/gen_prdc.py ADDED
@@ -0,0 +1,58 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Generation-side Precision/Recall (Kynkaanniemi) + Density/Coverage (Naeem) on
2
+ InceptionV3 (pytorch-fid) features. Precision=fidelity (fake in real manifold),
3
+ Recall=diversity/coverage (real covered by fake). No sklearn: kNN via torch.cdist."""
4
+ import os, sys, json, random
5
+ import numpy as np, torch
6
+ from PIL import Image
7
+ sys.path.insert(0, "/home/wzhang/LSC/Code/NPJ")
8
+ from framework.synth.pixdiff.fd_loss import InceptionFeatures
9
+
10
+ DR = "/home/wzhang/LSC/Dataset/Segmentation/processed_unified"
11
+ DSETS = {"isic": ("medsegdb_isic2018", "holdout"), "kvasir": ("kvasir_seg", "official"), "busi": ("busi", "fold01")}
12
+ BKS = ["jit", "pixelgen", "deco", "pixeldit"]
13
+ dev = "cuda"; CAP = 2000; K = 5
14
+ inc = InceptionFeatures().to(dev).eval()
15
+
16
+ def feats(d, cap=CAP):
17
+ fs = sorted(f for f in os.listdir(d) if f.lower().endswith((".png", ".jpg", ".jpeg")))
18
+ if len(fs) > cap:
19
+ random.seed(0); fs = random.sample(fs, cap)
20
+ out = []
21
+ for i in range(0, len(fs), 64):
22
+ b = []
23
+ for f in fs[i:i + 64]:
24
+ im = Image.open(os.path.join(d, f)).convert("RGB").resize((256, 256))
25
+ b.append(torch.from_numpy(np.asarray(im)).permute(2, 0, 1).float() / 255.)
26
+ with torch.no_grad():
27
+ out.append(inc(torch.stack(b).to(dev)).cpu())
28
+ return torch.cat(out)
29
+
30
+ def knn_radius(X, k):
31
+ d = torch.cdist(X, X); d.fill_diagonal_(float("inf")); return d.kthvalue(k, dim=1).values
32
+
33
+ def prdc(R, F, k=K):
34
+ R, F = R.to(dev), F.to(dev)
35
+ rr = knn_radius(R, k); ff = knn_radius(F, k); drf = torch.cdist(R, F)
36
+ prec = (drf <= rr[:, None]).any(0).float().mean().item()
37
+ rec = (drf <= ff[None, :]).any(1).float().mean().item()
38
+ dens = ((drf <= rr[:, None]).sum(0).float().mean() / k).item()
39
+ cov = (drf <= rr[:, None]).any(1).float().mean().item()
40
+ return prec, rec, dens, cov
41
+
42
+ realf = {}
43
+ for dk, (ds, proto) in DSETS.items():
44
+ realf[dk] = feats(f"{DR}/{ds}/{proto}/train/images")
45
+ print(f"[real] {dk}: {realf[dk].shape}", flush=True)
46
+
47
+ res = {}
48
+ for bk in BKS:
49
+ for dk, (ds, proto) in DSETS.items():
50
+ sd = f"{DR}/{ds}/{proto}/synth_fid_{bk}_{dk}/images"
51
+ if not os.path.isdir(sd):
52
+ print(f"[skip] {dk} {bk}"); continue
53
+ F = feats(sd)
54
+ p, r, de, c = prdc(realf[dk], F)
55
+ res[f"{dk}_{bk}"] = {"precision": round(p, 3), "recall": round(r, 3), "density": round(de, 3), "coverage": round(c, 3)}
56
+ print(f"[PRDC] {dk} {bk}: {res[f'{dk}_{bk}']}", flush=True)
57
+ json.dump(res, open("/home/wzhang/LSC/Code/NPJ/logs/fidviz/gen_prdc.json", "w"), indent=2)
58
+ print("PRDC_DONE", flush=True)
code/scripts/p1/make_jit_vs_fd.py ADDED
@@ -0,0 +1,65 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Same-mask JiT(native) vs JiT-FD comparison: does FD-perceptual sharpen the synth?
2
+ Samples JiT-FD on the SAME no-aug f50 masks the native JiT align-set used, builds
3
+ [mask | real | JiT native | JiT-FD] grids for ISIC + Kvasir."""
4
+ import os, subprocess
5
+ import numpy as np
6
+ import matplotlib; matplotlib.use("Agg")
7
+ import matplotlib.pyplot as plt
8
+ from PIL import Image
9
+
10
+ ROOT = "/home/wzhang/LSC/Code/NPJ"; DR = "/home/wzhang/LSC/Dataset/Segmentation/processed_unified"
11
+ PY = "/opt/anaconda3/envs/seggen/bin/python"
12
+ DSETS = {"isic": ("medsegdb_isic2018", "holdout", 2582), "kvasir": ("kvasir_seg", "official", 800)}
13
+
14
+ def sample(ckpt, ds, proto, frac, out):
15
+ if os.path.isdir(out + "/images") and len(os.listdir(out + "/images")) >= 40:
16
+ print(f"[skip-sample] {out} exists", flush=True); return
17
+ env = dict(os.environ, CUDA_DEVICE_ORDER="PCI_BUS_ID", CUDA_VISIBLE_DEVICES="0",
18
+ TORCHDYNAMO_DISABLE="1", PYTHONPATH=".", OMP_NUM_THREADS="4")
19
+ subprocess.run([PY, "-m", "framework.synth.pixdiff.sample", "--ckpt", ckpt, "--data_root", DR,
20
+ "--dataset", ds, "--protocol", proto, "--train_fraction", str(frac),
21
+ "--fraction_seed", "0", "--n_per_mask", "1", "--num_steps", "50", "--out_dir", out],
22
+ env=env, cwd=ROOT, check=True)
23
+ print(f"[sampled] {out}", flush=True)
24
+
25
+ def fmap(d):
26
+ p = os.path.join(d, "images"); m = {}
27
+ if os.path.isdir(p):
28
+ for f in sorted(os.listdir(p)):
29
+ if f.endswith(".png"): m.setdefault(f[:-4].split("__")[0], os.path.join(p, f))
30
+ return m
31
+ def rgb(p): return np.asarray(Image.open(p).convert("RGB").resize((256, 256)))
32
+ def gray(p): return np.asarray(Image.open(p).convert("L").resize((256, 256)))
33
+
34
+ for dk, (ds, proto, tot) in DSETS.items():
35
+ f50 = 50 / tot
36
+ fd_out = f"{DR}/{ds}/{proto}/synth_alignfd_jitfd_{dk}"
37
+ sample(f"pretrained/pixdiff/p1_jitfd_{dk}.pt", ds, proto, f50, fd_out)
38
+ base = f"{DR}/{ds}/{proto}"; ri, rm = f"{base}/train/images", f"{base}/train/masks"
39
+ nat = fmap(f"{base}/synth_align_jit_{dk}"); fd = fmap(fd_out)
40
+ common = set(os.path.splitext(f)[0] for f in os.listdir(ri) if f.endswith(".png")) & set(nat) & set(fd)
41
+ common = sorted(common); ncol = min(6, len(common))
42
+ idx = [round(i * (len(common) - 1) / (ncol - 1)) for i in range(ncol)] if ncol > 1 else [0]
43
+ cases = [common[i] for i in idx]
44
+ rows = [("Conditioning mask", "mask"), ("Real", "real"), ("JiT (native, P1)", nat), ("JiT-FD (FD-感知)", fd)]
45
+ fig, ax = plt.subplots(len(rows), ncol, figsize=(ncol * 2.1, len(rows) * 2.15))
46
+ for r, (lab, src) in enumerate(rows):
47
+ for c, bs in enumerate(cases):
48
+ a = ax[r][c]
49
+ try:
50
+ mk = gray(f"{rm}/{bs}.png")
51
+ if src == "mask": a.imshow(mk, cmap="gray")
52
+ elif src == "real": a.imshow(rgb(f"{ri}/{bs}.png"))
53
+ else: a.imshow(rgb(src[bs]))
54
+ if src not in ("mask",): a.contour((mk > 127).astype(float), levels=[0.5], colors=["#19f04b"], linewidths=0.9)
55
+ except Exception:
56
+ a.imshow(np.ones((256, 256, 3))); a.text(0.5, 0.5, "n/a", ha="center", va="center", transform=a.transAxes)
57
+ a.set_xticks([]); a.set_yticks([])
58
+ for s in a.spines.values(): s.set_visible(False)
59
+ if c == 0: a.set_ylabel(lab, fontsize=11, rotation=90, va="center", labelpad=8,
60
+ color=("#111" if r < 2 else "#1a3b8b"), fontweight=("bold" if r == 3 else "normal"))
61
+ fig.suptitle(f"{dk.upper()} — 原生 JiT vs JiT-FD(同掩码):FD-感知精修是否更锐?", fontsize=12)
62
+ plt.tight_layout(rect=[0.03, 0, 1, 0.95])
63
+ out = f"/tmp/jit_vs_fd_{dk}.png"; plt.savefig(out, dpi=150, bbox_inches="tight", facecolor="white")
64
+ print(f"[grid] {out}", flush=True)
65
+ print("JIT_VS_FD_DONE", flush=True)
code/scripts/p1/p1_busi_master.py ADDED
@@ -0,0 +1,166 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """P1 master orchestrator: DAG scheduler over GPU0-5 for the backbone bake-off.
2
+ Phases: A) 8 generators (4 backbones x 2 datasets, amortized, 50k steps)
3
+ B) 16 sampling jobs (per gen x N in {50,100}, mask_aug n_per_mask=4)
4
+ C) 60 downstream seg runs (real + 4 backbones) x 2 ds x 2 N x 3 seeds
5
+ Single GPU per job (no DDP needed: 84 independent jobs). Retry-once on failure.
6
+ Resumable (skips done outputs). Rolling aggregate -> logs/p1master/p1_results.json."""
7
+ import os, sys, time, json, subprocess, statistics as st
8
+
9
+ ROOT = "/home/wzhang/LSC/Code/NPJ"
10
+ DR = "/home/wzhang/LSC/Dataset/Segmentation/processed_unified"
11
+ PY = "/opt/anaconda3/envs/seggen/bin/python"
12
+ GPUS = [0, 1, 2, 3, 4, 5]
13
+ os.chdir(ROOT)
14
+ LOGD = os.path.join(ROOT, "logs", "p1busi")
15
+ os.makedirs(LOGD, exist_ok=True)
16
+
17
+ def log(m):
18
+ line = f"[{time.strftime('%F %T')}] {m}"
19
+ with open(os.path.join(LOGD, "status.md"), "a") as f:
20
+ f.write(line + "\n")
21
+ print(line, flush=True)
22
+
23
+ DSETS = {"busi": ("busi", "fold01", 545)}
24
+ BKS = ["jit", "pixelgen", "deco", "pixeldit"]
25
+ NS = [50, 100]
26
+ SEEDS = [0, 1, 2]
27
+
28
+ jobs = {}
29
+ def add(jid, cmd, deps=(), done_path=None, done_min=1):
30
+ jobs[jid] = {"cmd": cmd, "deps": list(deps), "done_path": done_path,
31
+ "done_min": done_min, "state": "pending", "tries": 0, "gpu": None}
32
+
33
+ # Phase A: generators
34
+ for bk in BKS:
35
+ for dk, (ds, proto, tot) in DSETS.items():
36
+ out = f"pretrained/pixdiff/p1_{bk}_{dk}.pt"
37
+ cmd = (f"{PY} -m framework.synth.pixdiff.train --data_root {DR} --dataset {ds} "
38
+ f"--protocol {proto} --backbone {bk} --img_size 256 --batch_size 16 "
39
+ f"--epochs 100000 --max_steps 50000 --lr 1e-4 --amp bf16 "
40
+ f"--train_fraction 1.0 --fraction_seed 0 --out_ckpt {out} --log_interval 500")
41
+ add(f"gen_{bk}_{dk}", cmd, done_path=os.path.join(ROOT, out))
42
+
43
+ # Phase B: sampling
44
+ for bk in BKS:
45
+ for dk, (ds, proto, tot) in DSETS.items():
46
+ ck = f"pretrained/pixdiff/p1_{bk}_{dk}.pt"
47
+ for N in NS:
48
+ f = N / tot
49
+ sd = f"{DR}/{ds}/{proto}/synth_p1_{bk}_{dk}_f{N}"
50
+ cmd = (f"{PY} -m framework.synth.pixdiff.sample --ckpt {ck} --data_root {DR} "
51
+ f"--dataset {ds} --protocol {proto} --train_fraction {f} --fraction_seed 0 "
52
+ f"--n_per_mask 4 --mask_aug --num_steps 50 --out_dir {sd}")
53
+ add(f"samp_{bk}_{dk}_N{N}", cmd, deps=[f"gen_{bk}_{dk}"],
54
+ done_path=os.path.join(sd, "images"), done_min=N * 4)
55
+
56
+ # Phase C: downstream
57
+ def mpath(exp, ds, proto, S):
58
+ return os.path.join(ROOT, f"results/{exp}/{ds}_{proto}/unet/seed{S}/metrics.json")
59
+
60
+ def seg_cmd(ds, proto, f, exp, S, synth=None):
61
+ base = (f"{PY} framework/train.py --data_root {DR} --dataset {ds} --protocol {proto} "
62
+ f"--arch unet --encoder resnet50 --aug standard --epochs 400 "
63
+ f"--train_fraction {f} --fraction_seed 0 --exp_name {exp} --amp bf16 --seed {S}")
64
+ if synth:
65
+ base += f" --synth_train_dir {synth}"
66
+ test = (f"{PY} framework/test.py --data_root {DR} --dataset {ds} --protocol {proto} "
67
+ f"--arch unet --encoder resnet50 --aug standard --exp_name {exp} --seed {S}")
68
+ return base + " && " + test
69
+
70
+ for dk, (ds, proto, tot) in DSETS.items():
71
+ for N in NS:
72
+ f = N / tot
73
+ for S in SEEDS:
74
+ exp = f"p1_real_{dk}_N{N}"
75
+ add(f"seg_real_{dk}_N{N}_s{S}", seg_cmd(ds, proto, f, exp, S),
76
+ done_path=mpath(exp, ds, proto, S))
77
+ for bk in BKS:
78
+ sd = f"{DR}/{ds}/{proto}/synth_p1_{bk}_{dk}_f{N}"
79
+ for S in SEEDS:
80
+ exp = f"p1_{bk}_{dk}_N{N}"
81
+ add(f"seg_{bk}_{dk}_N{N}_s{S}", seg_cmd(ds, proto, f, exp, S, synth=sd),
82
+ deps=[f"samp_{bk}_{dk}_N{N}"], done_path=mpath(exp, ds, proto, S))
83
+
84
+ def is_done(j):
85
+ p = j["done_path"]
86
+ if not p or not os.path.exists(p):
87
+ return False
88
+ if os.path.isdir(p):
89
+ try:
90
+ return len(os.listdir(p)) >= j["done_min"]
91
+ except OSError:
92
+ return False
93
+ return True
94
+
95
+ def aggregate():
96
+ res = {}
97
+ for dk, (ds, proto, tot) in DSETS.items():
98
+ for N in NS:
99
+ for arm in ["real"] + BKS:
100
+ exp = f"p1_{arm}_{dk}_N{N}"
101
+ ious, dices = [], []
102
+ for S in SEEDS:
103
+ mp = mpath(exp, ds, proto, S)
104
+ if os.path.exists(mp):
105
+ try:
106
+ m = json.load(open(mp))["metrics"]
107
+ ious.append(m["iou_mean"]); dices.append(m["dice_mean"])
108
+ except Exception:
109
+ pass
110
+ if ious:
111
+ res[f"{dk}_N{N}_{arm}"] = {
112
+ "iou_mean": sum(ious) / len(ious), "dice_mean": sum(dices) / len(dices),
113
+ "n_seeds": len(ious), "iou_seeds": ious, "dice_seeds": dices}
114
+ json.dump(res, open(os.path.join(LOGD, "p1_results.json"), "w"), indent=2)
115
+
116
+ for jid, j in jobs.items():
117
+ if is_done(j):
118
+ j["state"] = "done"
119
+ def deps_done(j):
120
+ return all(jobs[d]["state"] == "done" for d in j["deps"])
121
+
122
+ running = {}
123
+ free = set(GPUS)
124
+ MAXTRIES = 2
125
+ log(f"START {len(jobs)} jobs on GPUs {GPUS} ({sum(1 for j in jobs.values() if j['state']=='done')} pre-done)")
126
+ last_summary = 0
127
+ while True:
128
+ if all(j["state"] in ("done", "failed") for j in jobs.values()):
129
+ break
130
+ for jid, j in jobs.items():
131
+ if not free:
132
+ break
133
+ if j["state"] == "pending" and deps_done(j):
134
+ if is_done(j):
135
+ j["state"] = "done"; continue
136
+ g = free.pop()
137
+ env = dict(os.environ, CUDA_DEVICE_ORDER="PCI_BUS_ID",
138
+ CUDA_VISIBLE_DEVICES=str(g), TORCHDYNAMO_DISABLE="1",
139
+ PYTHONPATH=".", OMP_NUM_THREADS="4")
140
+ lf = open(os.path.join(LOGD, jid + ".log"), "a")
141
+ p = subprocess.Popen(j["cmd"], shell=True, env=env, stdout=lf,
142
+ stderr=subprocess.STDOUT, cwd=ROOT)
143
+ running[g] = (jid, p, lf); j["state"] = "running"; j["gpu"] = g; j["tries"] += 1
144
+ log(f"LAUNCH {jid} GPU{g} try{j['tries']}")
145
+ for g, (jid, p, lf) in list(running.items()):
146
+ rc = p.poll()
147
+ if rc is None:
148
+ continue
149
+ lf.close(); del running[g]; free.add(g)
150
+ j = jobs[jid]
151
+ if is_done(j):
152
+ j["state"] = "done"; log(f"DONE {jid} rc={rc}")
153
+ elif j["tries"] < MAXTRIES:
154
+ j["state"] = "pending"; log(f"RETRY {jid} rc={rc}")
155
+ else:
156
+ j["state"] = "failed"; log(f"FAILED {jid} rc={rc}")
157
+ if time.time() - last_summary > 300:
158
+ cnt = {s: sum(1 for j in jobs.values() if j["state"] == s)
159
+ for s in ("done", "running", "pending", "failed")}
160
+ log(f"SUMMARY {cnt} | running={sorted(j['gpu'] for j in jobs.values() if j['state']=='running')}")
161
+ aggregate(); last_summary = time.time()
162
+ time.sleep(10)
163
+
164
+ aggregate()
165
+ log("ALL DONE")
166
+ print("P1_MASTER_DONE", flush=True)
code/scripts/p1/p1_busi_results.json ADDED
@@ -0,0 +1,152 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "busi_N50_real": {
3
+ "iou_mean": 0.5264299887746967,
4
+ "dice_mean": 0.6150173362667782,
5
+ "n_seeds": 3,
6
+ "iou_seeds": [
7
+ 0.5264031863116967,
8
+ 0.5078566908181252,
9
+ 0.5450300891942683
10
+ ],
11
+ "dice_seeds": [
12
+ 0.6158220746914147,
13
+ 0.5943401974078463,
14
+ 0.6348897367010737
15
+ ]
16
+ },
17
+ "busi_N50_jit": {
18
+ "iou_mean": 0.6034184804314947,
19
+ "dice_mean": 0.6819841561663282,
20
+ "n_seeds": 3,
21
+ "iou_seeds": [
22
+ 0.6151347744667092,
23
+ 0.5876539837190357,
24
+ 0.6074666831087394
25
+ ],
26
+ "dice_seeds": [
27
+ 0.689949667483754,
28
+ 0.6696459718689409,
29
+ 0.6863568291462898
30
+ ]
31
+ },
32
+ "busi_N50_pixelgen": {
33
+ "iou_mean": 0.5931307319807805,
34
+ "dice_mean": 0.6689307916999434,
35
+ "n_seeds": 3,
36
+ "iou_seeds": [
37
+ 0.5920417792314931,
38
+ 0.5799457559831734,
39
+ 0.6074046607276752
40
+ ],
41
+ "dice_seeds": [
42
+ 0.666457437158763,
43
+ 0.6579920482874959,
44
+ 0.6823428896535711
45
+ ]
46
+ },
47
+ "busi_N50_deco": {
48
+ "iou_mean": 0.5979877144880544,
49
+ "dice_mean": 0.6759540945542658,
50
+ "n_seeds": 3,
51
+ "iou_seeds": [
52
+ 0.5926858046978442,
53
+ 0.5927788031039587,
54
+ 0.6084985356623608
55
+ ],
56
+ "dice_seeds": [
57
+ 0.6700259199200342,
58
+ 0.6746954264901631,
59
+ 0.6831409372526002
60
+ ]
61
+ },
62
+ "busi_N50_pixeldit": {
63
+ "iou_mean": 0.607193388275135,
64
+ "dice_mean": 0.6873956052246571,
65
+ "n_seeds": 3,
66
+ "iou_seeds": [
67
+ 0.6123619477064693,
68
+ 0.611641987503641,
69
+ 0.5975762296152947
70
+ ],
71
+ "dice_seeds": [
72
+ 0.6897889057784018,
73
+ 0.6909085054783096,
74
+ 0.6814894044172604
75
+ ]
76
+ },
77
+ "busi_N100_real": {
78
+ "iou_mean": 0.5901633542778459,
79
+ "dice_mean": 0.672577797062739,
80
+ "n_seeds": 3,
81
+ "iou_seeds": [
82
+ 0.596288318471246,
83
+ 0.5925297498972647,
84
+ 0.581671994465027
85
+ ],
86
+ "dice_seeds": [
87
+ 0.6805305551414621,
88
+ 0.6720643475113856,
89
+ 0.6651384885353694
90
+ ]
91
+ },
92
+ "busi_N100_jit": {
93
+ "iou_mean": 0.638692683336804,
94
+ "dice_mean": 0.7151948555339698,
95
+ "n_seeds": 3,
96
+ "iou_seeds": [
97
+ 0.638767115763446,
98
+ 0.6464336918630936,
99
+ 0.6308772423838722
100
+ ],
101
+ "dice_seeds": [
102
+ 0.7157337836019346,
103
+ 0.7221307950206755,
104
+ 0.7077199879792992
105
+ ]
106
+ },
107
+ "busi_N100_pixelgen": {
108
+ "iou_mean": 0.6466251026809454,
109
+ "dice_mean": 0.7228950392456911,
110
+ "n_seeds": 3,
111
+ "iou_seeds": [
112
+ 0.6475516801145157,
113
+ 0.6421441721654421,
114
+ 0.6501794557628784
115
+ ],
116
+ "dice_seeds": [
117
+ 0.7219398382769914,
118
+ 0.7182700866504312,
119
+ 0.7284751928096508
120
+ ]
121
+ },
122
+ "busi_N100_deco": {
123
+ "iou_mean": 0.6285762758560305,
124
+ "dice_mean": 0.7061591091011407,
125
+ "n_seeds": 3,
126
+ "iou_seeds": [
127
+ 0.6474242343217572,
128
+ 0.6371748339087213,
129
+ 0.6011297593376128
130
+ ],
131
+ "dice_seeds": [
132
+ 0.7268623914793524,
133
+ 0.7158535119640104,
134
+ 0.6757614238600594
135
+ ]
136
+ },
137
+ "busi_N100_pixeldit": {
138
+ "iou_mean": 0.6540538953088814,
139
+ "dice_mean": 0.7316955478068173,
140
+ "n_seeds": 3,
141
+ "iou_seeds": [
142
+ 0.6616519138889801,
143
+ 0.6470095283658505,
144
+ 0.6535002436718134
145
+ ],
146
+ "dice_seeds": [
147
+ 0.7383933408398845,
148
+ 0.7258913137588962,
149
+ 0.730801988821671
150
+ ]
151
+ }
152
+ }
code/scripts/p1/p1_full_metrics.json ADDED
@@ -0,0 +1,182 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "isic_N50_real": {
3
+ "dice": 85.16,
4
+ "iou": 77.49,
5
+ "sensitivity": 85.34,
6
+ "precision": 90.16
7
+ },
8
+ "isic_N50_jit": {
9
+ "dice": 87.74,
10
+ "iou": 80.52,
11
+ "sensitivity": 88.36,
12
+ "precision": 90.69
13
+ },
14
+ "isic_N50_pixelgen": {
15
+ "dice": 87.9,
16
+ "iou": 80.71,
17
+ "sensitivity": 88.68,
18
+ "precision": 90.82
19
+ },
20
+ "isic_N50_deco": {
21
+ "dice": 87.27,
22
+ "iou": 79.96,
23
+ "sensitivity": 87.87,
24
+ "precision": 90.63
25
+ },
26
+ "isic_N50_pixeldit": {
27
+ "dice": 87.68,
28
+ "iou": 80.52,
29
+ "sensitivity": 88.33,
30
+ "precision": 90.58
31
+ },
32
+ "isic_N100_real": {
33
+ "dice": 87.31,
34
+ "iou": 80.18,
35
+ "sensitivity": 87.78,
36
+ "precision": 90.47
37
+ },
38
+ "isic_N100_jit": {
39
+ "dice": 88.52,
40
+ "iou": 81.66,
41
+ "sensitivity": 89.46,
42
+ "precision": 90.72
43
+ },
44
+ "isic_N100_pixelgen": {
45
+ "dice": 88.95,
46
+ "iou": 82.13,
47
+ "sensitivity": 90.15,
48
+ "precision": 90.83
49
+ },
50
+ "isic_N100_deco": {
51
+ "dice": 88.15,
52
+ "iou": 81.19,
53
+ "sensitivity": 89.85,
54
+ "precision": 89.97
55
+ },
56
+ "isic_N100_pixeldit": {
57
+ "dice": 88.6,
58
+ "iou": 81.7,
59
+ "sensitivity": 89.26,
60
+ "precision": 91.13
61
+ },
62
+ "kvasir_N50_real": {
63
+ "dice": 83.75,
64
+ "iou": 75.35,
65
+ "sensitivity": 83.39,
66
+ "precision": 88.81
67
+ },
68
+ "kvasir_N50_jit": {
69
+ "dice": 86.75,
70
+ "iou": 79.2,
71
+ "sensitivity": 86.93,
72
+ "precision": 89.85
73
+ },
74
+ "kvasir_N50_pixelgen": {
75
+ "dice": 86.16,
76
+ "iou": 78.42,
77
+ "sensitivity": 86.95,
78
+ "precision": 89.06
79
+ },
80
+ "kvasir_N50_deco": {
81
+ "dice": 88.03,
82
+ "iou": 80.84,
83
+ "sensitivity": 87.64,
84
+ "precision": 90.75
85
+ },
86
+ "kvasir_N50_pixeldit": {
87
+ "dice": 85.87,
88
+ "iou": 78.97,
89
+ "sensitivity": 85.23,
90
+ "precision": 89.44
91
+ },
92
+ "kvasir_N100_real": {
93
+ "dice": 86.95,
94
+ "iou": 79.5,
95
+ "sensitivity": 87.53,
96
+ "precision": 90.66
97
+ },
98
+ "kvasir_N100_jit": {
99
+ "dice": 88.5,
100
+ "iou": 81.62,
101
+ "sensitivity": 87.71,
102
+ "precision": 92.71
103
+ },
104
+ "kvasir_N100_pixelgen": {
105
+ "dice": 86.94,
106
+ "iou": 80.32,
107
+ "sensitivity": 86.37,
108
+ "precision": 91.14
109
+ },
110
+ "kvasir_N100_deco": {
111
+ "dice": 88.14,
112
+ "iou": 81.22,
113
+ "sensitivity": 88.95,
114
+ "precision": 90.6
115
+ },
116
+ "kvasir_N100_pixeldit": {
117
+ "dice": 89.52,
118
+ "iou": 82.83,
119
+ "sensitivity": 89.42,
120
+ "precision": 91.78
121
+ },
122
+ "busi_N50_real": {
123
+ "dice": 61.5,
124
+ "iou": 52.64,
125
+ "sensitivity": 59.37,
126
+ "precision": 72.68
127
+ },
128
+ "busi_N50_jit": {
129
+ "dice": 68.2,
130
+ "iou": 60.34,
131
+ "sensitivity": 68.58,
132
+ "precision": 73.26
133
+ },
134
+ "busi_N50_pixelgen": {
135
+ "dice": 66.89,
136
+ "iou": 59.31,
137
+ "sensitivity": 65.33,
138
+ "precision": 77.14
139
+ },
140
+ "busi_N50_deco": {
141
+ "dice": 67.6,
142
+ "iou": 59.8,
143
+ "sensitivity": 66.72,
144
+ "precision": 75.06
145
+ },
146
+ "busi_N50_pixeldit": {
147
+ "dice": 68.74,
148
+ "iou": 60.72,
149
+ "sensitivity": 67.82,
150
+ "precision": 75.91
151
+ },
152
+ "busi_N100_real": {
153
+ "dice": 67.26,
154
+ "iou": 59.02,
155
+ "sensitivity": 66.01,
156
+ "precision": 73.61
157
+ },
158
+ "busi_N100_jit": {
159
+ "dice": 71.52,
160
+ "iou": 63.87,
161
+ "sensitivity": 70.99,
162
+ "precision": 76.41
163
+ },
164
+ "busi_N100_pixelgen": {
165
+ "dice": 72.29,
166
+ "iou": 64.66,
167
+ "sensitivity": 72.49,
168
+ "precision": 77.16
169
+ },
170
+ "busi_N100_deco": {
171
+ "dice": 70.62,
172
+ "iou": 62.86,
173
+ "sensitivity": 71.52,
174
+ "precision": 74.61
175
+ },
176
+ "busi_N100_pixeldit": {
177
+ "dice": 73.17,
178
+ "iou": 65.41,
179
+ "sensitivity": 72.65,
180
+ "precision": 79.23
181
+ }
182
+ }
code/scripts/p1/p1_gen_queue.sh ADDED
@@ -0,0 +1,32 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env bash
2
+ # P1 generator-training queue: backbones x datasets, amortized (full-data), ~50k steps.
3
+ # Args: GPU [datasets] [backbones] e.g. p1_gen_queue.sh 5 "isic kvasir" "jit pixelgen deco pixeldit"
4
+ set -u
5
+ DR=/home/wzhang/LSC/Dataset/Segmentation/processed_unified
6
+ ROOT=/home/wzhang/LSC/Code/NPJ
7
+ PY=/opt/anaconda3/envs/seggen/bin/python
8
+ GPU=${1:-5}
9
+ ORDER=${2:-"isic kvasir"}
10
+ BKS=${3:-"jit pixelgen deco pixeldit"}
11
+ E="CUDA_DEVICE_ORDER=PCI_BUS_ID CUDA_VISIBLE_DEVICES=$GPU TORCHDYNAMO_DISABLE=1 PYTHONPATH=. OMP_NUM_THREADS=4"
12
+ cd "$ROOT" || exit 1
13
+ ST=logs/p1gen_status.md
14
+ log(){ echo "[$(date '+%F %T')] $*" >> "$ST"; }
15
+ declare -A DSMAP=( [isic]="medsegdb_isic2018 holdout" [kvasir]="kvasir_seg official" )
16
+ log "=== P1 GEN QUEUE START GPU-$GPU order=[$ORDER] bks=[$BKS] pid=$$ ==="
17
+ for dk in $ORDER; do
18
+ set -- ${DSMAP[$dk]}; ds=$1; proto=$2
19
+ for bk in $BKS; do
20
+ out=pretrained/pixdiff/p1_${bk}_${dk}.pt
21
+ if [ -f "$out" ]; then log "SKIP $bk/$dk (ckpt exists)"; continue; fi
22
+ log "TRAIN $bk/$dk -> $out"
23
+ env $E $PY -m framework.synth.pixdiff.train --data_root "$DR" --dataset "$ds" --protocol "$proto" \
24
+ --backbone "$bk" --img_size 256 --batch_size 16 --epochs 100000 --max_steps 50000 \
25
+ --lr 1e-4 --amp bf16 --train_fraction 1.0 --fraction_seed 0 \
26
+ --out_ckpt "$out" --log_interval 200 > logs/p1gen_${bk}_${dk}.log 2>&1
27
+ rc=$?
28
+ log " done $bk/$dk rc=$rc ckpt=$([ -f "$out" ] && echo ok || echo MISSING)"
29
+ done
30
+ done
31
+ log "=== P1 GEN QUEUE DONE GPU-$GPU ==="
32
+ echo "P1GEN_DONE_$GPU" >> "$ST"
code/scripts/p1/p1_master.py ADDED
@@ -0,0 +1,167 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """P1 master orchestrator: DAG scheduler over GPU0-5 for the backbone bake-off.
2
+ Phases: A) 8 generators (4 backbones x 2 datasets, amortized, 50k steps)
3
+ B) 16 sampling jobs (per gen x N in {50,100}, mask_aug n_per_mask=4)
4
+ C) 60 downstream seg runs (real + 4 backbones) x 2 ds x 2 N x 3 seeds
5
+ Single GPU per job (no DDP needed: 84 independent jobs). Retry-once on failure.
6
+ Resumable (skips done outputs). Rolling aggregate -> logs/p1master/p1_results.json."""
7
+ import os, sys, time, json, subprocess, statistics as st
8
+
9
+ ROOT = "/home/wzhang/LSC/Code/NPJ"
10
+ DR = "/home/wzhang/LSC/Dataset/Segmentation/processed_unified"
11
+ PY = "/opt/anaconda3/envs/seggen/bin/python"
12
+ GPUS = [0, 1, 2, 3, 4, 5]
13
+ os.chdir(ROOT)
14
+ LOGD = os.path.join(ROOT, "logs", "p1master")
15
+ os.makedirs(LOGD, exist_ok=True)
16
+
17
+ def log(m):
18
+ line = f"[{time.strftime('%F %T')}] {m}"
19
+ with open(os.path.join(LOGD, "status.md"), "a") as f:
20
+ f.write(line + "\n")
21
+ print(line, flush=True)
22
+
23
+ DSETS = {"isic": ("medsegdb_isic2018", "holdout", 2582),
24
+ "kvasir": ("kvasir_seg", "official", 800)}
25
+ BKS = ["jit", "pixelgen", "deco", "pixeldit"]
26
+ NS = [50, 100]
27
+ SEEDS = [0, 1, 2]
28
+
29
+ jobs = {}
30
+ def add(jid, cmd, deps=(), done_path=None, done_min=1):
31
+ jobs[jid] = {"cmd": cmd, "deps": list(deps), "done_path": done_path,
32
+ "done_min": done_min, "state": "pending", "tries": 0, "gpu": None}
33
+
34
+ # Phase A: generators
35
+ for bk in BKS:
36
+ for dk, (ds, proto, tot) in DSETS.items():
37
+ out = f"pretrained/pixdiff/p1_{bk}_{dk}.pt"
38
+ cmd = (f"{PY} -m framework.synth.pixdiff.train --data_root {DR} --dataset {ds} "
39
+ f"--protocol {proto} --backbone {bk} --img_size 256 --batch_size 16 "
40
+ f"--epochs 100000 --max_steps 50000 --lr 1e-4 --amp bf16 "
41
+ f"--train_fraction 1.0 --fraction_seed 0 --out_ckpt {out} --log_interval 500")
42
+ add(f"gen_{bk}_{dk}", cmd, done_path=os.path.join(ROOT, out))
43
+
44
+ # Phase B: sampling
45
+ for bk in BKS:
46
+ for dk, (ds, proto, tot) in DSETS.items():
47
+ ck = f"pretrained/pixdiff/p1_{bk}_{dk}.pt"
48
+ for N in NS:
49
+ f = N / tot
50
+ sd = f"{DR}/{ds}/{proto}/synth_p1_{bk}_{dk}_f{N}"
51
+ cmd = (f"{PY} -m framework.synth.pixdiff.sample --ckpt {ck} --data_root {DR} "
52
+ f"--dataset {ds} --protocol {proto} --train_fraction {f} --fraction_seed 0 "
53
+ f"--n_per_mask 4 --mask_aug --num_steps 50 --out_dir {sd}")
54
+ add(f"samp_{bk}_{dk}_N{N}", cmd, deps=[f"gen_{bk}_{dk}"],
55
+ done_path=os.path.join(sd, "images"), done_min=N * 4)
56
+
57
+ # Phase C: downstream
58
+ def mpath(exp, ds, proto, S):
59
+ return os.path.join(ROOT, f"results/{exp}/{ds}_{proto}/unet/seed{S}/metrics.json")
60
+
61
+ def seg_cmd(ds, proto, f, exp, S, synth=None):
62
+ base = (f"{PY} framework/train.py --data_root {DR} --dataset {ds} --protocol {proto} "
63
+ f"--arch unet --encoder resnet50 --aug standard --epochs 400 "
64
+ f"--train_fraction {f} --fraction_seed 0 --exp_name {exp} --amp bf16 --seed {S}")
65
+ if synth:
66
+ base += f" --synth_train_dir {synth}"
67
+ test = (f"{PY} framework/test.py --data_root {DR} --dataset {ds} --protocol {proto} "
68
+ f"--arch unet --encoder resnet50 --aug standard --exp_name {exp} --seed {S}")
69
+ return base + " && " + test
70
+
71
+ for dk, (ds, proto, tot) in DSETS.items():
72
+ for N in NS:
73
+ f = N / tot
74
+ for S in SEEDS:
75
+ exp = f"p1_real_{dk}_N{N}"
76
+ add(f"seg_real_{dk}_N{N}_s{S}", seg_cmd(ds, proto, f, exp, S),
77
+ done_path=mpath(exp, ds, proto, S))
78
+ for bk in BKS:
79
+ sd = f"{DR}/{ds}/{proto}/synth_p1_{bk}_{dk}_f{N}"
80
+ for S in SEEDS:
81
+ exp = f"p1_{bk}_{dk}_N{N}"
82
+ add(f"seg_{bk}_{dk}_N{N}_s{S}", seg_cmd(ds, proto, f, exp, S, synth=sd),
83
+ deps=[f"samp_{bk}_{dk}_N{N}"], done_path=mpath(exp, ds, proto, S))
84
+
85
+ def is_done(j):
86
+ p = j["done_path"]
87
+ if not p or not os.path.exists(p):
88
+ return False
89
+ if os.path.isdir(p):
90
+ try:
91
+ return len(os.listdir(p)) >= j["done_min"]
92
+ except OSError:
93
+ return False
94
+ return True
95
+
96
+ def aggregate():
97
+ res = {}
98
+ for dk, (ds, proto, tot) in DSETS.items():
99
+ for N in NS:
100
+ for arm in ["real"] + BKS:
101
+ exp = f"p1_{arm}_{dk}_N{N}"
102
+ ious, dices = [], []
103
+ for S in SEEDS:
104
+ mp = mpath(exp, ds, proto, S)
105
+ if os.path.exists(mp):
106
+ try:
107
+ m = json.load(open(mp))["metrics"]
108
+ ious.append(m["iou_mean"]); dices.append(m["dice_mean"])
109
+ except Exception:
110
+ pass
111
+ if ious:
112
+ res[f"{dk}_N{N}_{arm}"] = {
113
+ "iou_mean": sum(ious) / len(ious), "dice_mean": sum(dices) / len(dices),
114
+ "n_seeds": len(ious), "iou_seeds": ious, "dice_seeds": dices}
115
+ json.dump(res, open(os.path.join(LOGD, "p1_results.json"), "w"), indent=2)
116
+
117
+ for jid, j in jobs.items():
118
+ if is_done(j):
119
+ j["state"] = "done"
120
+ def deps_done(j):
121
+ return all(jobs[d]["state"] == "done" for d in j["deps"])
122
+
123
+ running = {}
124
+ free = set(GPUS)
125
+ MAXTRIES = 2
126
+ log(f"START {len(jobs)} jobs on GPUs {GPUS} ({sum(1 for j in jobs.values() if j['state']=='done')} pre-done)")
127
+ last_summary = 0
128
+ while True:
129
+ if all(j["state"] in ("done", "failed") for j in jobs.values()):
130
+ break
131
+ for jid, j in jobs.items():
132
+ if not free:
133
+ break
134
+ if j["state"] == "pending" and deps_done(j):
135
+ if is_done(j):
136
+ j["state"] = "done"; continue
137
+ g = free.pop()
138
+ env = dict(os.environ, CUDA_DEVICE_ORDER="PCI_BUS_ID",
139
+ CUDA_VISIBLE_DEVICES=str(g), TORCHDYNAMO_DISABLE="1",
140
+ PYTHONPATH=".", OMP_NUM_THREADS="4")
141
+ lf = open(os.path.join(LOGD, jid + ".log"), "a")
142
+ p = subprocess.Popen(j["cmd"], shell=True, env=env, stdout=lf,
143
+ stderr=subprocess.STDOUT, cwd=ROOT)
144
+ running[g] = (jid, p, lf); j["state"] = "running"; j["gpu"] = g; j["tries"] += 1
145
+ log(f"LAUNCH {jid} GPU{g} try{j['tries']}")
146
+ for g, (jid, p, lf) in list(running.items()):
147
+ rc = p.poll()
148
+ if rc is None:
149
+ continue
150
+ lf.close(); del running[g]; free.add(g)
151
+ j = jobs[jid]
152
+ if is_done(j):
153
+ j["state"] = "done"; log(f"DONE {jid} rc={rc}")
154
+ elif j["tries"] < MAXTRIES:
155
+ j["state"] = "pending"; log(f"RETRY {jid} rc={rc}")
156
+ else:
157
+ j["state"] = "failed"; log(f"FAILED {jid} rc={rc}")
158
+ if time.time() - last_summary > 300:
159
+ cnt = {s: sum(1 for j in jobs.values() if j["state"] == s)
160
+ for s in ("done", "running", "pending", "failed")}
161
+ log(f"SUMMARY {cnt} | running={sorted(j['gpu'] for j in jobs.values() if j['state']=='running')}")
162
+ aggregate(); last_summary = time.time()
163
+ time.sleep(10)
164
+
165
+ aggregate()
166
+ log("ALL DONE")
167
+ print("P1_MASTER_DONE", flush=True)
code/scripts/p1/p1_results.json ADDED
@@ -0,0 +1,302 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "isic_N50_real": {
3
+ "iou_mean": 0.7749372989665764,
4
+ "dice_mean": 0.8516048790137111,
5
+ "n_seeds": 3,
6
+ "iou_seeds": [
7
+ 0.7697344814191404,
8
+ 0.772716144763783,
9
+ 0.7823612707168056
10
+ ],
11
+ "dice_seeds": [
12
+ 0.8480779677176173,
13
+ 0.8503343900165478,
14
+ 0.8564022793069677
15
+ ]
16
+ },
17
+ "isic_N50_jit": {
18
+ "iou_mean": 0.8052051026596693,
19
+ "dice_mean": 0.8773763597758372,
20
+ "n_seeds": 3,
21
+ "iou_seeds": [
22
+ 0.8038682316971465,
23
+ 0.8033287650931844,
24
+ 0.8084183111886771
25
+ ],
26
+ "dice_seeds": [
27
+ 0.8763739862706927,
28
+ 0.8755590664589031,
29
+ 0.8801960265979158
30
+ ]
31
+ },
32
+ "isic_N50_pixelgen": {
33
+ "iou_mean": 0.8071234360528656,
34
+ "dice_mean": 0.8790093239202045,
35
+ "n_seeds": 3,
36
+ "iou_seeds": [
37
+ 0.8046804415401705,
38
+ 0.8071308385675863,
39
+ 0.80955902805084
40
+ ],
41
+ "dice_seeds": [
42
+ 0.8771928252649649,
43
+ 0.8781965482168992,
44
+ 0.8816385982787495
45
+ ]
46
+ },
47
+ "isic_N50_deco": {
48
+ "iou_mean": 0.7996043535547449,
49
+ "dice_mean": 0.8726861944567714,
50
+ "n_seeds": 3,
51
+ "iou_seeds": [
52
+ 0.7988670932254234,
53
+ 0.8020180961470182,
54
+ 0.797927871291793
55
+ ],
56
+ "dice_seeds": [
57
+ 0.8724373392477919,
58
+ 0.875006987196055,
59
+ 0.8706142569264674
60
+ ]
61
+ },
62
+ "isic_N50_pixeldit": {
63
+ "iou_mean": 0.8051857432382189,
64
+ "dice_mean": 0.8768404429697196,
65
+ "n_seeds": 3,
66
+ "iou_seeds": [
67
+ 0.8076451208631789,
68
+ 0.80299518868651,
69
+ 0.8049169201649677
70
+ ],
71
+ "dice_seeds": [
72
+ 0.8789570258524421,
73
+ 0.8757209786234587,
74
+ 0.8758433244332583
75
+ ]
76
+ },
77
+ "isic_N100_real": {
78
+ "iou_mean": 0.8018487332144956,
79
+ "dice_mean": 0.8731124297212501,
80
+ "n_seeds": 3,
81
+ "iou_seeds": [
82
+ 0.8018532993620694,
83
+ 0.8030565805042732,
84
+ 0.8006363197771443
85
+ ],
86
+ "dice_seeds": [
87
+ 0.8742042532437185,
88
+ 0.8744577536420762,
89
+ 0.8706752822779558
90
+ ]
91
+ },
92
+ "isic_N100_jit": {
93
+ "iou_mean": 0.8166433570078636,
94
+ "dice_mean": 0.8852389536608354,
95
+ "n_seeds": 3,
96
+ "iou_seeds": [
97
+ 0.8184147975006147,
98
+ 0.8152316048562138,
99
+ 0.8162836686667623
100
+ ],
101
+ "dice_seeds": [
102
+ 0.8866063380696619,
103
+ 0.8842458679141912,
104
+ 0.8848646549986529
105
+ ]
106
+ },
107
+ "isic_N100_pixelgen": {
108
+ "iou_mean": 0.8212823342374741,
109
+ "dice_mean": 0.8895036576449075,
110
+ "n_seeds": 3,
111
+ "iou_seeds": [
112
+ 0.8203996376523117,
113
+ 0.8215948549134239,
114
+ 0.8218525101466865
115
+ ],
116
+ "dice_seeds": [
117
+ 0.8893210218244473,
118
+ 0.8894422995119722,
119
+ 0.8897476515983034
120
+ ]
121
+ },
122
+ "isic_N100_deco": {
123
+ "iou_mean": 0.8119110110180822,
124
+ "dice_mean": 0.8814897596168004,
125
+ "n_seeds": 3,
126
+ "iou_seeds": [
127
+ 0.812474683629472,
128
+ 0.81098766389058,
129
+ 0.8122706855341947
130
+ ],
131
+ "dice_seeds": [
132
+ 0.8834257222572127,
133
+ 0.8798560384070768,
134
+ 0.881187518186112
135
+ ]
136
+ },
137
+ "isic_N100_pixeldit": {
138
+ "iou_mean": 0.8169910899725092,
139
+ "dice_mean": 0.8860181724664612,
140
+ "n_seeds": 3,
141
+ "iou_seeds": [
142
+ 0.8168252607788484,
143
+ 0.8168659787418466,
144
+ 0.8172820303968324
145
+ ],
146
+ "dice_seeds": [
147
+ 0.8855665474000166,
148
+ 0.8859441754776876,
149
+ 0.8865437945216795
150
+ ]
151
+ },
152
+ "kvasir_N50_real": {
153
+ "iou_mean": 0.7534671958018979,
154
+ "dice_mean": 0.8375380075145843,
155
+ "n_seeds": 3,
156
+ "iou_seeds": [
157
+ 0.7422559578123432,
158
+ 0.7636408134000111,
159
+ 0.7545048161933395
160
+ ],
161
+ "dice_seeds": [
162
+ 0.8244831334988446,
163
+ 0.849740539589318,
164
+ 0.8383903494555902
165
+ ]
166
+ },
167
+ "kvasir_N50_jit": {
168
+ "iou_mean": 0.7919630629212887,
169
+ "dice_mean": 0.867516315771424,
170
+ "n_seeds": 3,
171
+ "iou_seeds": [
172
+ 0.7929972233023507,
173
+ 0.7943493738869872,
174
+ 0.7885425915745284
175
+ ],
176
+ "dice_seeds": [
177
+ 0.8676604855354516,
178
+ 0.8655865846767077,
179
+ 0.8693018771021128
180
+ ]
181
+ },
182
+ "kvasir_N50_pixelgen": {
183
+ "iou_mean": 0.7842259846149005,
184
+ "dice_mean": 0.8615631488519918,
185
+ "n_seeds": 3,
186
+ "iou_seeds": [
187
+ 0.7841354377174787,
188
+ 0.7782063762034493,
189
+ 0.7903361399237733
190
+ ],
191
+ "dice_seeds": [
192
+ 0.858331250209457,
193
+ 0.8583852594567815,
194
+ 0.8679729368897371
195
+ ]
196
+ },
197
+ "kvasir_N50_deco": {
198
+ "iou_mean": 0.8083752894364583,
199
+ "dice_mean": 0.8803164279623182,
200
+ "n_seeds": 3,
201
+ "iou_seeds": [
202
+ 0.8104997825534876,
203
+ 0.8045413607010391,
204
+ 0.8100847250548484
205
+ ],
206
+ "dice_seeds": [
207
+ 0.8826893496480072,
208
+ 0.8778416284559154,
209
+ 0.8804183057830322
210
+ ]
211
+ },
212
+ "kvasir_N50_pixeldit": {
213
+ "iou_mean": 0.7897261916004509,
214
+ "dice_mean": 0.8586578118039875,
215
+ "n_seeds": 3,
216
+ "iou_seeds": [
217
+ 0.7733292625059578,
218
+ 0.800297355081953,
219
+ 0.7955519572134416
220
+ ],
221
+ "dice_seeds": [
222
+ 0.8420264952707284,
223
+ 0.8674089356274011,
224
+ 0.8665380045138328
225
+ ]
226
+ },
227
+ "kvasir_N100_real": {
228
+ "iou_mean": 0.7949746887997339,
229
+ "dice_mean": 0.8695256998060409,
230
+ "n_seeds": 3,
231
+ "iou_seeds": [
232
+ 0.7966858390583808,
233
+ 0.7968998232753884,
234
+ 0.7913384040654323
235
+ ],
236
+ "dice_seeds": [
237
+ 0.87442242361909,
238
+ 0.8720526205988572,
239
+ 0.8621020552001751
240
+ ]
241
+ },
242
+ "kvasir_N100_jit": {
243
+ "iou_mean": 0.8162035292350188,
244
+ "dice_mean": 0.884963785884059,
245
+ "n_seeds": 3,
246
+ "iou_seeds": [
247
+ 0.8044583298250717,
248
+ 0.8215198247885459,
249
+ 0.8226324330914389
250
+ ],
251
+ "dice_seeds": [
252
+ 0.8733364457591601,
253
+ 0.8902805661843565,
254
+ 0.8912743457086602
255
+ ]
256
+ },
257
+ "kvasir_N100_pixelgen": {
258
+ "iou_mean": 0.8032477917950044,
259
+ "dice_mean": 0.8693680901970753,
260
+ "n_seeds": 3,
261
+ "iou_seeds": [
262
+ 0.7990136881750312,
263
+ 0.8131907666134136,
264
+ 0.7975389205965682
265
+ ],
266
+ "dice_seeds": [
267
+ 0.8653448154322447,
268
+ 0.8767741359428084,
269
+ 0.8659853192161729
270
+ ]
271
+ },
272
+ "kvasir_N100_deco": {
273
+ "iou_mean": 0.8121859633543412,
274
+ "dice_mean": 0.8814163311797074,
275
+ "n_seeds": 3,
276
+ "iou_seeds": [
277
+ 0.8014463766191772,
278
+ 0.8207859220669117,
279
+ 0.8143255913769347
280
+ ],
281
+ "dice_seeds": [
282
+ 0.8737877150565937,
283
+ 0.885505301166682,
284
+ 0.8849559773158466
285
+ ]
286
+ },
287
+ "kvasir_N100_pixeldit": {
288
+ "iou_mean": 0.8282534914466804,
289
+ "dice_mean": 0.8952017769853248,
290
+ "n_seeds": 3,
291
+ "iou_seeds": [
292
+ 0.8263046284256831,
293
+ 0.8313217018221024,
294
+ 0.8271341440922554
295
+ ],
296
+ "dice_seeds": [
297
+ 0.8969579025391758,
298
+ 0.8957992526085868,
299
+ 0.8928481758082121
300
+ ]
301
+ }
302
+ }
code/scripts/p1/smoke_backbone.py ADDED
@@ -0,0 +1,76 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Smoke: port DeCo (PixNerDiT) or PixelDiT (PixDiT) into the PixDiff mask-concat scaffolding.
2
+ Backbone-agnostic decouple: build with in=img+cond, out defaults to in, take x_pred[:, :img_ch].
3
+ Usage: python smoke_backbone.py {deco|pixeldit}"""
4
+ import os, sys
5
+ sys.path.insert(0, "/home/wzhang/LSC/Code/NPJ")
6
+ import torch, torch.nn as nn
7
+ from torch.utils.data import DataLoader
8
+ from framework.synth.pixdiff.conditioning import build_conditioner
9
+ from framework.synth.pixdiff.data import MaskCondGenDataset
10
+
11
+ BK = sys.argv[1]
12
+ DECO = "/home/wzhang/LSC/Code/NPJ/sota/DeCo"
13
+ PIXELDIT = "/home/wzhang/LSC/Code/NPJ/sota/PixelDiT"
14
+ dev = "cuda"
15
+ DR = "/home/wzhang/LSC/Dataset/Segmentation/processed_unified"
16
+
17
+ ds = MaskCondGenDataset(DR, "medsegdb_isic2018", "holdout", img_size=256,
18
+ train_fraction=0.02, fraction_seed=0)
19
+ cond = build_conditioner("onehot", ds.num_classes).to(dev)
20
+ img_ch, K = ds.in_channels, cond.cond_channels
21
+ in_tot = img_ch + K
22
+ print(f"[{BK}] ds n={len(ds)} img_ch={img_ch} K={K} in_tot={in_tot}", flush=True)
23
+
24
+ if BK == "deco":
25
+ sys.path.insert(0, os.path.join(DECO, "src", "models", "transformer"))
26
+ from dit_c2i_DeCo import PixNerDiT
27
+ net = PixNerDiT(in_channels=in_tot, patch_size=16, num_groups=12, hidden_size=768,
28
+ hidden_size_x=32, num_blocks=13, num_cond_blocks=12, num_classes=1).to(dev)
29
+ elif BK == "pixeldit":
30
+ sys.path.insert(0, PIXELDIT)
31
+ from pixdit_core.pixeldit_c2i import PixDiT
32
+ net = PixDiT(in_channels=in_tot, num_groups=12, hidden_size=768, pixel_hidden_size=16,
33
+ patch_depth=12, pixel_depth=4, patch_size=16, num_classes=1).to(dev)
34
+ else:
35
+ raise SystemExit("backbone must be deco|pixeldit")
36
+
37
+ print(f"[{BK}] params={sum(p.numel() for p in net.parameters())/1e6:.1f}M", flush=True)
38
+ opt = torch.optim.AdamW(net.parameters(), lr=1e-4)
39
+ dl = DataLoader(ds, batch_size=4, shuffle=True, drop_last=True, num_workers=2)
40
+ it = iter(dl)
41
+ def get_batch():
42
+ global it
43
+ try: b = next(it)
44
+ except StopIteration: it = iter(dl); b = next(it)
45
+ return (b["image"], b["mask"]) if isinstance(b, dict) else (b[0], b[1])
46
+
47
+ net.train()
48
+ for step in range(20):
49
+ img, msk = get_batch(); img, msk = img.to(dev), msk.to(dev)
50
+ t = torch.sigmoid(torch.randn(img.size(0), device=dev) * 0.8 - 0.8).view(-1, 1, 1, 1)
51
+ e = torch.randn_like(img)
52
+ z = t * img + (1 - t) * e
53
+ v = (img - z) / (1 - t).clamp_min(5e-2)
54
+ c = cond(msk)
55
+ y = torch.zeros(img.size(0), dtype=torch.long, device=dev)
56
+ out = net(torch.cat([z, c], dim=1), t.flatten(), y)
57
+ assert out.dim() == 4 and out.shape[1] >= img_ch, f"bad out shape {tuple(out.shape)}"
58
+ x_pred = out[:, :img_ch]
59
+ v_pred = (x_pred - z) / (1 - t).clamp_min(5e-2)
60
+ loss = ((v - v_pred) ** 2).mean()
61
+ loss.backward(); opt.step(); opt.zero_grad()
62
+ if step % 5 == 0 or step == 19:
63
+ print(f"[{BK}] step {step:2d} loss {loss.item():.4f}", flush=True)
64
+
65
+ net.eval()
66
+ with torch.no_grad():
67
+ msk0 = msk[:2]; c0 = cond(msk0)
68
+ z = torch.randn(2, img_ch, 256, 256, device=dev)
69
+ ts = torch.linspace(0, 1, 11).tolist()
70
+ for i in range(10):
71
+ tc, dt = ts[i], ts[i + 1] - ts[i]
72
+ out = net(torch.cat([z, c0], dim=1), torch.full((2,), tc, device=dev),
73
+ torch.zeros(2, dtype=torch.long, device=dev))[:, :img_ch]
74
+ z = z + (out - z) / max(1 - tc, 5e-2) * dt
75
+ print(f"[{BK}] sample ok shape={tuple(z.shape)} range=({z.min():.2f},{z.max():.2f})", flush=True)
76
+ print(f"SMOKE_{BK.upper()}_PASS", flush=True)
code/scripts/p1/smoke_pixelgen.py ADDED
@@ -0,0 +1,70 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Smoke test: port PixelGen's JiT denoiser into the PixDiff mask-concat scaffolding.
2
+ Validates: import on server, in=img+cond/out=img build, flow-matching train step, sampling.
3
+ Tiny: 20 train steps + 10-step sample on ~50 ISIC images. No checkpoint written."""
4
+ import os, sys
5
+ sys.path.insert(0, "/home/wzhang/LSC/Code/NPJ") # framework.*
6
+ PG = "/home/wzhang/LSC/Code/NPJ/sota/PixelGen"
7
+ sys.path.insert(0, os.path.join(PG, "src", "models", "transformer")) # JiT.py
8
+ import torch, torch.nn as nn
9
+ from torch.utils.data import DataLoader
10
+ from JiT import JiT_models, FinalLayer # PixelGen denoiser
11
+ from framework.synth.pixdiff.conditioning import build_conditioner
12
+ from framework.synth.pixdiff.data import MaskCondGenDataset
13
+
14
+ dev = "cuda"
15
+ DR = "/home/wzhang/LSC/Dataset/Segmentation/processed_unified"
16
+ ds = MaskCondGenDataset(DR, "medsegdb_isic2018", "holdout", img_size=256,
17
+ train_fraction=0.02, fraction_seed=0)
18
+ print(f"[ds] n={len(ds)} in_ch={ds.in_channels} num_classes={ds.num_classes}", flush=True)
19
+ cond = build_conditioner("onehot", ds.num_classes).to(dev)
20
+ img_ch, K = ds.in_channels, cond.cond_channels
21
+
22
+ net = JiT_models["JiT-B/16"](input_size=256, in_channels=img_ch + K, num_classes=1).to(dev)
23
+ if net.out_channels != img_ch:
24
+ net.out_channels = img_ch
25
+ net.final_layer = FinalLayer(net.hidden_size, net.patch_size, img_ch).to(dev)
26
+ nn.init.constant_(net.final_layer.linear.weight, 0); nn.init.constant_(net.final_layer.linear.bias, 0)
27
+ nn.init.constant_(net.final_layer.adaLN_modulation[-1].weight, 0)
28
+ nn.init.constant_(net.final_layer.adaLN_modulation[-1].bias, 0)
29
+ print(f"[net] PixelGen JiT-B/16 in={img_ch+K} out={net.out_channels} params={sum(p.numel() for p in net.parameters())/1e6:.1f}M", flush=True)
30
+
31
+ opt = torch.optim.AdamW(net.parameters(), lr=1e-4)
32
+ dl = DataLoader(ds, batch_size=4, shuffle=True, drop_last=True, num_workers=2)
33
+ it = iter(dl)
34
+
35
+ def get_batch():
36
+ global it
37
+ try: b = next(it)
38
+ except StopIteration:
39
+ it = iter(dl); b = next(it)
40
+ if isinstance(b, dict): return b["image"], b["mask"]
41
+ return b[0], b[1]
42
+
43
+ net.train()
44
+ for step in range(20):
45
+ img, msk = get_batch(); img, msk = img.to(dev), msk.to(dev)
46
+ t = torch.sigmoid(torch.randn(img.size(0), device=dev) * 0.8 - 0.8).view(-1, 1, 1, 1)
47
+ e = torch.randn_like(img)
48
+ z = t * img + (1 - t) * e
49
+ v = (img - z) / (1 - t).clamp_min(5e-2)
50
+ c = cond(msk)
51
+ y = torch.zeros(img.size(0), dtype=torch.long, device=dev)
52
+ x_pred = net(torch.cat([z, c], dim=1), t.flatten(), y)
53
+ v_pred = (x_pred - z) / (1 - t).clamp_min(5e-2)
54
+ loss = ((v - v_pred) ** 2).mean()
55
+ loss.backward(); opt.step(); opt.zero_grad()
56
+ if step % 5 == 0 or step == 19:
57
+ print(f"[train] step {step:2d} loss {loss.item():.4f}", flush=True)
58
+
59
+ net.eval()
60
+ with torch.no_grad():
61
+ msk0 = msk[:2]; c0 = cond(msk0)
62
+ z = torch.randn(2, img_ch, 256, 256, device=dev)
63
+ ts = torch.linspace(0, 1, 11).tolist()
64
+ for i in range(10):
65
+ tc, dt = ts[i], ts[i + 1] - ts[i]
66
+ tt = torch.full((2,), tc, device=dev)
67
+ xp = net(torch.cat([z, c0], dim=1), tt, torch.zeros(2, dtype=torch.long, device=dev))
68
+ z = z + (xp - z) / max(1 - tc, 5e-2) * dt
69
+ print(f"[sample] ok shape={tuple(z.shape)} range=({z.min():.2f},{z.max():.2f})", flush=True)
70
+ print("SMOKE_PIXELGEN_PASS", flush=True)
code/scripts/p1/train_fd_patched.py ADDED
@@ -0,0 +1,179 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """FD-Loss post-training for a PixDiff generator (env: seggen, GPU 5).
2
+
3
+ Starts from a base PixDiff checkpoint and continues training with:
4
+ total = flow_matching_loss + fd_weight * normalized_FD_loss
5
+ where the FD term matches the generated x0 feature distribution to the real-image
6
+ reference distribution (Inception space), gated to low-noise timesteps (where x0 is
7
+ a meaningful image). This targets the blur/distribution gap the MSE objective leaves.
8
+
9
+ Run from project root (…/NPJ):
10
+ CUDA_VISIBLE_DEVICES=5 python -m framework.synth.pixdiff.train_fd \
11
+ --base_ckpt pretrained/pixdiff/kvasir_seg_official_f1.0.pt \
12
+ --data_root /home/wzhang/LSC/Dataset/Segmentation/processed_unified \
13
+ --dataset kvasir_seg --protocol official \
14
+ --epochs 200 --lr 2e-5 --fd_weight 0.5 \
15
+ --out_ckpt pretrained/pixdiff/kvasir_seg_official_f1.0_fd.pt
16
+ """
17
+ from __future__ import annotations
18
+
19
+ import argparse
20
+ import os
21
+ import sys
22
+ import time
23
+
24
+ sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..")))
25
+
26
+ import numpy as np
27
+ import torch
28
+ from torch.utils.data import DataLoader
29
+
30
+ from framework.synth.pixdiff.data import MaskCondGenDataset
31
+ from framework.synth.pixdiff.conditioning import build_conditioner
32
+ from framework.synth.pixdiff.mask_jit import MaskDenoiser
33
+ from framework.synth.pixdiff.fd_loss import (
34
+ InceptionFeatures, FeatureQueue, compute_frechet_distance_loss,
35
+ precompute_sigma_ref_sqrt, compute_ref_stats,
36
+ )
37
+
38
+
39
+ def get_args():
40
+ p = argparse.ArgumentParser("PixDiff FD-Loss post-training")
41
+ p.add_argument("--base_ckpt", required=True)
42
+ p.add_argument("--data_root", required=True)
43
+ p.add_argument("--dataset", required=True)
44
+ p.add_argument("--protocol", required=True)
45
+ p.add_argument("--train_fraction", type=float, default=1.0)
46
+ p.add_argument("--fraction_seed", type=int, default=0)
47
+ p.add_argument("--epochs", type=int, default=200)
48
+ p.add_argument("--batch_size", type=int, default=32)
49
+ p.add_argument("--lr", type=float, default=2e-5)
50
+ p.add_argument("--num_workers", type=int, default=6)
51
+ p.add_argument("--amp", default="bf16", choices=["bf16", "fp16", "fp32"])
52
+ # FD-Loss knobs
53
+ p.add_argument("--fd_weight", type=float, default=0.5)
54
+ p.add_argument("--fd_gate_t", type=float, default=0.5, help="apply FD only when t>=this (low noise)")
55
+ p.add_argument("--queue_size", type=int, default=512)
56
+ p.add_argument("--fd_norm_eps", type=float, default=1e-2)
57
+ p.add_argument("--lpips_weight", type=float, default=0.0)
58
+ p.add_argument("--dino_weight", type=float, default=0.0)
59
+ p.add_argument("--percep_gate_t", type=float, default=0.5, help="apply perceptual only when t>=this")
60
+ p.add_argument("--ref_stats", default="", help="npz of (mu,sigma); auto path + compute if empty")
61
+ p.add_argument("--ema_decay", type=float, default=0.9999)
62
+ p.add_argument("--seed", type=int, default=0)
63
+ p.add_argument("--out_ckpt", required=True)
64
+ p.add_argument("--log_interval", type=int, default=20)
65
+ return p.parse_args()
66
+
67
+
68
+ def main():
69
+ a = get_args()
70
+ torch.manual_seed(a.seed)
71
+ device = "cuda"
72
+ amp_dtype = {"bf16": torch.bfloat16, "fp16": torch.float16}.get(a.amp)
73
+
74
+ # ---- data ----
75
+ ds = MaskCondGenDataset(a.data_root, a.dataset, a.protocol, img_size=256,
76
+ train_fraction=a.train_fraction, fraction_seed=a.fraction_seed)
77
+ img_ch, n_cls = ds.in_channels, ds.num_classes
78
+ loader = DataLoader(ds, batch_size=a.batch_size, shuffle=True, drop_last=True,
79
+ num_workers=a.num_workers, pin_memory=True, persistent_workers=a.num_workers > 0)
80
+ print(f"[fd] {a.dataset}/{a.protocol} n={len(ds)} in_ch={img_ch} num_classes={n_cls}", flush=True)
81
+ if img_ch != 3:
82
+ print("[fd][warn] Inception expects 3ch; non-RGB dataset — FD features may be weak.", flush=True)
83
+
84
+ # ---- model from base ckpt ----
85
+ ckpt = torch.load(a.base_ckpt, map_location="cpu", weights_only=False)
86
+ cond = build_conditioner(ckpt.get("conditioner", "onehot"), n_cls)
87
+ model = MaskDenoiser(ckpt["model_name"], ckpt["img_size"], ckpt["img_channels"], cond,
88
+ noise_scale=ckpt.get("noise_scale", 1.0), ema_decay=a.ema_decay, backbone=ckpt.get("backbone", "jit")).to(device)
89
+ model.load_state_dict(ckpt["state_dict"])
90
+ model._ema = [e.to(device) for e in ckpt["ema"]] if ckpt.get("ema") is not None else None
91
+ if model._ema is None:
92
+ model.ema_init()
93
+ print(f"[fd] loaded base {a.base_ckpt}", flush=True)
94
+
95
+ # ---- FD machinery ----
96
+ inception = InceptionFeatures().to(device).eval()
97
+ queue = FeatureQueue(size=a.queue_size, feat_dim=inception.feat_dim).to(device)
98
+ percep = None
99
+ if a.lpips_weight > 0 or a.dino_weight > 0:
100
+ from framework.synth.pixdiff.perceptual import PerceptualLoss
101
+ percep = PerceptualLoss(use_lpips=a.lpips_weight > 0, use_dino=a.dino_weight > 0, device=device)
102
+ print(f"[fd] perceptual ON lpips_w={a.lpips_weight} dino_w={a.dino_weight} gate_t={a.percep_gate_t}", flush=True)
103
+
104
+ ref_path = a.ref_stats or a.out_ckpt.replace(".pt", "_refstats.npz")
105
+ if os.path.isfile(ref_path):
106
+ rs = np.load(ref_path); mu_ref_np, sigma_ref_np = rs["mu"], rs["sigma"]
107
+ print(f"[fd] loaded ref stats {ref_path}", flush=True)
108
+ else:
109
+ print("[fd] computing reference stats from real train images...", flush=True)
110
+ ref_loader = DataLoader(MaskCondGenDataset(a.data_root, a.dataset, a.protocol, img_size=256,
111
+ train_fraction=a.train_fraction, fraction_seed=a.fraction_seed,
112
+ hflip=False, vflip=False),
113
+ batch_size=a.batch_size, shuffle=False, num_workers=a.num_workers)
114
+ mu_ref_np, sigma_ref_np, nref = compute_ref_stats(ref_loader, inception, device)
115
+ os.makedirs(os.path.dirname(os.path.abspath(ref_path)) or ".", exist_ok=True)
116
+ np.savez(ref_path, mu=mu_ref_np, sigma=sigma_ref_np)
117
+ print(f"[fd] ref stats from {nref} imgs -> {ref_path}", flush=True)
118
+ mu_ref = torch.tensor(mu_ref_np, device=device, dtype=torch.float64)
119
+ sigma_ref = torch.tensor(sigma_ref_np, device=device, dtype=torch.float64)
120
+ sigma_ref_sqrt = precompute_sigma_ref_sqrt(sigma_ref)
121
+
122
+ opt = torch.optim.AdamW(model._trainable(), lr=a.lr, weight_decay=0.0)
123
+ os.makedirs(os.path.dirname(os.path.abspath(a.out_ckpt)) or ".", exist_ok=True)
124
+
125
+ def save():
126
+ torch.save({"model_name": ckpt["model_name"], "img_size": ckpt["img_size"],
127
+ "img_channels": img_ch, "num_classes": n_cls,
128
+ "conditioner": ckpt.get("conditioner", "onehot"),
129
+ "noise_scale": ckpt.get("noise_scale", 1.0),
130
+ "state_dict": model.state_dict(), "ema": model._ema, "args": vars(a)}, a.out_ckpt)
131
+ print(f"[fd] saved {a.out_ckpt}", flush=True)
132
+
133
+ step = 0
134
+ for epoch in range(a.epochs):
135
+ model.train(); t0 = time.time(); run_fm = run_fd = run_fdraw = 0.0
136
+ for batch in loader:
137
+ img = batch["image"].to(device, non_blocking=True)
138
+ mask = batch["mask"].to(device, non_blocking=True)
139
+ opt.zero_grad(set_to_none=True)
140
+ with torch.autocast("cuda", dtype=amp_dtype) if amp_dtype else _null():
141
+ fm_loss, x_pred, t = model(img, mask, return_pred=True)
142
+ # FD term on predicted clean image, gated to low noise
143
+ gate = t >= a.fd_gate_t
144
+ fd_loss = torch.zeros((), device=device); fd_raw = 0.0
145
+ if int(gate.sum()) >= 2:
146
+ xg = x_pred[gate].float()
147
+ feats = inception((xg.clamp(-1, 1) + 1) / 2) # (Ng,2048), grad flows
148
+ if queue.is_ready():
149
+ mu, sigma = queue.build_feats_stats(feats)
150
+ fd = compute_frechet_distance_loss(mu_ref, sigma_ref, mu, sigma, sigma_ref_sqrt)
151
+ fd_raw = float(fd); fd_loss = fd / (fd.detach() + a.fd_norm_eps)
152
+ queue.enqueue(feats)
153
+ total = fm_loss + a.fd_weight * fd_loss
154
+ pl_lpips = pl_dino = 0.0
155
+ if percep is not None:
156
+ pgate = t >= a.percep_gate_t
157
+ if int(pgate.sum()) >= 1:
158
+ pld = percep(x_pred[pgate].float(), img[pgate].float())
159
+ if "lpips" in pld:
160
+ total = total + a.lpips_weight * pld["lpips"]; pl_lpips = float(pld["lpips"])
161
+ if "dino" in pld:
162
+ total = total + a.dino_weight * pld["dino"]; pl_dino = float(pld["dino"])
163
+ total.backward(); opt.step(); model.ema_update()
164
+ run_fm += float(fm_loss); run_fd += float(fd_loss); run_fdraw += fd_raw; step += 1
165
+ if step % a.log_interval == 0:
166
+ print(f"[fd] ep{epoch} step{step} fm={float(fm_loss):.4f} fd_raw={fd_raw:.1f} "
167
+ f"lpips={pl_lpips:.3f} dino={pl_dino:.3f} qready={queue.is_ready()}", flush=True)
168
+ print(f"[fd] epoch {epoch} fm={run_fm/max(1,len(loader)):.4f} "
169
+ f"fd_raw={run_fdraw/max(1,len(loader)):.1f} ({time.time()-t0:.1f}s)", flush=True)
170
+ save()
171
+
172
+
173
+ class _null:
174
+ def __enter__(self): return self
175
+ def __exit__(self, *a): return False
176
+
177
+
178
+ if __name__ == "__main__":
179
+ main()
code/sota/Swin-Unet/README.md ADDED
@@ -0,0 +1,64 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Swin-Unet
2
+ [ECCVW2022] The codes for the work "Swin-Unet: Unet-like Pure Transformer for Medical Image Segmentation"(https://arxiv.org/abs/2105.05537). Our paper has been accepted by ECCV 2022 MEDICAL COMPUTER VISION WORKSHOP (https://mcv-workshop.github.io/). We updated the Reproducibility. I hope this will help you to reproduce the results.
3
+
4
+ ## 1. Download pre-trained swin transformer model (Swin-T)
5
+ * [Get pre-trained model in this link] (https://drive.google.com/drive/folders/1UC3XOoezeum0uck4KBVGa8osahs6rKUY?usp=sharing): Put pretrained Swin-T into folder "pretrained_ckpt/"
6
+
7
+ ## 2. Prepare data
8
+
9
+ - The datasets we used are provided by TransUnet's authors. [Get processed data in this link] (Synapse/BTCV: https://drive.google.com/drive/folders/1ACJEoTp-uqfFJ73qS3eUObQh52nGuzCd and ACDC: https://drive.google.com/drive/folders/1KQcrci7aKsYZi1hQoZ3T3QUtcy7b--n4).
10
+
11
+ ## 3. Environment
12
+
13
+ - Please prepare an environment with python=3.7, and then use the command "pip install -r requirements.txt" for the dependencies.
14
+
15
+ ## 4. Train/Test
16
+
17
+ - Run the train script on synapse dataset. The batch size we used is 24. If you do not have enough GPU memory, the bacth size can be reduced to 12 or 6 to save memory.
18
+
19
+ - Train
20
+
21
+ ```bash
22
+ sh train.sh
23
+ # or
24
+ python train.py --dataset Synapse --cfg configs/swin_tiny_patch4_window7_224_lite.yaml --root_path your DATA_DIR --max_epochs 150 --output_dir your OUT_DIR --img_size 224 --base_lr 0.05 --batch_size 24
25
+ ```
26
+
27
+ - Test
28
+
29
+ ```bash
30
+ sh test.sh
31
+ # or
32
+ python test.py --dataset Synapse --cfg configs/swin_tiny_patch4_window7_224_lite.yaml --is_saveni --volume_path your DATA_DIR --output_dir your OUT_DIR --max_epoch 150 --base_lr 0.05 --img_size 224 --batch_size 24
33
+ ```
34
+
35
+ ## Reproducibility
36
+
37
+
38
+ - Codes
39
+
40
+ Our trained model is stored on the Huawei cloud. The interns do not have the right to send any files out from the internal system, so I can't share our trained model weights. Regarding how to reproduce the segmentation results presented in the paper, we discovered that different GPU types would generate different results. In our code, we carefully set the random seed, so the results should be consistent when trained multiple times on the same type of GPU. If the training does not give the same segmentation results as in the paper, it is recommended to adjust the learning rate. And, the type of GPU we used in this work is Tesla v100. Finaly, pre-training is very important for pure transformer models. In our experiments, both the encoder and decoder are initialized with pretrained weights rather than initializing the encoder with pretrained weights only.
41
+
42
+ ## References
43
+ * [TransUnet](https://github.com/Beckschen/TransUNet)
44
+ * [SwinTransformer](https://github.com/microsoft/Swin-Transformer)
45
+
46
+ ## Citation
47
+
48
+ ```bibtex
49
+ @InProceedings{swinunet,
50
+ author = {Hu Cao and Yueyue Wang and Joy Chen and Dongsheng Jiang and Xiaopeng Zhang and Qi Tian and Manning Wang},
51
+ title = {Swin-Unet: Unet-like Pure Transformer for Medical Image Segmentation},
52
+ booktitle = {Proceedings of the European Conference on Computer Vision Workshops(ECCVW)},
53
+ year = {2022}
54
+ }
55
+
56
+ @misc{cao2021swinunet,
57
+ title={Swin-Unet: Unet-like Pure Transformer for Medical Image Segmentation},
58
+ author={Hu Cao and Yueyue Wang and Joy Chen and Dongsheng Jiang and Xiaopeng Zhang and Qi Tian and Manning Wang},
59
+ year={2021},
60
+ eprint={2105.05537},
61
+ archivePrefix={arXiv},
62
+ primaryClass={eess.IV}
63
+ }
64
+ ```
code/sota/Swin-Unet/config.py ADDED
@@ -0,0 +1,229 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # --------------------------------------------------------
2
+ # Swin Transformer
3
+ # Copyright (c) 2021 Microsoft
4
+ # Licensed under The MIT License [see LICENSE for details]
5
+ # Written by Ze Liu
6
+ # --------------------------------------------------------'
7
+
8
+ import os
9
+ import yaml
10
+ from yacs.config import CfgNode as CN
11
+
12
+ _C = CN()
13
+
14
+ # Base config files
15
+ _C.BASE = ['']
16
+
17
+ # -----------------------------------------------------------------------------
18
+ # Data settings
19
+ # -----------------------------------------------------------------------------
20
+ _C.DATA = CN()
21
+ # Batch size for a single GPU, could be overwritten by command line argument
22
+ _C.DATA.BATCH_SIZE = 128
23
+ # Path to dataset, could be overwritten by command line argument
24
+ _C.DATA.DATA_PATH = ''
25
+ # Dataset name
26
+ _C.DATA.DATASET = 'imagenet'
27
+ # Input image size
28
+ _C.DATA.IMG_SIZE = 224
29
+ # Interpolation to resize image (random, bilinear, bicubic)
30
+ _C.DATA.INTERPOLATION = 'bicubic'
31
+ # Use zipped dataset instead of folder dataset
32
+ # could be overwritten by command line argument
33
+ _C.DATA.ZIP_MODE = False
34
+ # Cache Data in Memory, could be overwritten by command line argument
35
+ _C.DATA.CACHE_MODE = 'part'
36
+ # Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.
37
+ _C.DATA.PIN_MEMORY = True
38
+ # Number of data loading threads
39
+ _C.DATA.NUM_WORKERS = 8
40
+
41
+ # -----------------------------------------------------------------------------
42
+ # Model settings
43
+ # -----------------------------------------------------------------------------
44
+ _C.MODEL = CN()
45
+ # Model type
46
+ _C.MODEL.TYPE = 'swin'
47
+ # Model name
48
+ _C.MODEL.NAME = 'swin_tiny_patch4_window7_224'
49
+ # Checkpoint to resume, could be overwritten by command line argument
50
+ _C.MODEL.PRETRAIN_CKPT = './pretrained_ckpt/swin_tiny_patch4_window7_224.pth'
51
+ _C.MODEL.RESUME = ''
52
+ # Number of classes, overwritten in data preparation
53
+ _C.MODEL.NUM_CLASSES = 1000
54
+ # Dropout rate
55
+ _C.MODEL.DROP_RATE = 0.0
56
+ # Drop path rate
57
+ _C.MODEL.DROP_PATH_RATE = 0.1
58
+ # Label Smoothing
59
+ _C.MODEL.LABEL_SMOOTHING = 0.1
60
+
61
+ # Swin Transformer parameters
62
+ _C.MODEL.SWIN = CN()
63
+ _C.MODEL.SWIN.PATCH_SIZE = 4
64
+ _C.MODEL.SWIN.IN_CHANS = 3
65
+ _C.MODEL.SWIN.EMBED_DIM = 96
66
+ _C.MODEL.SWIN.DEPTHS = [2, 2, 6, 2]
67
+ _C.MODEL.SWIN.DECODER_DEPTHS = [2, 2, 6, 2]
68
+ _C.MODEL.SWIN.NUM_HEADS = [3, 6, 12, 24]
69
+ _C.MODEL.SWIN.WINDOW_SIZE = 7
70
+ _C.MODEL.SWIN.MLP_RATIO = 4.
71
+ _C.MODEL.SWIN.QKV_BIAS = True
72
+ _C.MODEL.SWIN.QK_SCALE = None
73
+ _C.MODEL.SWIN.APE = False
74
+ _C.MODEL.SWIN.PATCH_NORM = True
75
+ _C.MODEL.SWIN.FINAL_UPSAMPLE= "expand_first"
76
+
77
+ # -----------------------------------------------------------------------------
78
+ # Training settings
79
+ # -----------------------------------------------------------------------------
80
+ _C.TRAIN = CN()
81
+ _C.TRAIN.START_EPOCH = 0
82
+ _C.TRAIN.EPOCHS = 300
83
+ _C.TRAIN.WARMUP_EPOCHS = 20
84
+ _C.TRAIN.WEIGHT_DECAY = 0.05
85
+ _C.TRAIN.BASE_LR = 5e-4
86
+ _C.TRAIN.WARMUP_LR = 5e-7
87
+ _C.TRAIN.MIN_LR = 5e-6
88
+ # Clip gradient norm
89
+ _C.TRAIN.CLIP_GRAD = 5.0
90
+ # Auto resume from latest checkpoint
91
+ _C.TRAIN.AUTO_RESUME = True
92
+ # Gradient accumulation steps
93
+ # could be overwritten by command line argument
94
+ _C.TRAIN.ACCUMULATION_STEPS = 0
95
+ # Whether to use gradient checkpointing to save memory
96
+ # could be overwritten by command line argument
97
+ _C.TRAIN.USE_CHECKPOINT = False
98
+
99
+ # LR scheduler
100
+ _C.TRAIN.LR_SCHEDULER = CN()
101
+ _C.TRAIN.LR_SCHEDULER.NAME = 'cosine'
102
+ # Epoch interval to decay LR, used in StepLRScheduler
103
+ _C.TRAIN.LR_SCHEDULER.DECAY_EPOCHS = 30
104
+ # LR decay rate, used in StepLRScheduler
105
+ _C.TRAIN.LR_SCHEDULER.DECAY_RATE = 0.1
106
+
107
+ # Optimizer
108
+ _C.TRAIN.OPTIMIZER = CN()
109
+ _C.TRAIN.OPTIMIZER.NAME = 'adamw'
110
+ # Optimizer Epsilon
111
+ _C.TRAIN.OPTIMIZER.EPS = 1e-8
112
+ # Optimizer Betas
113
+ _C.TRAIN.OPTIMIZER.BETAS = (0.9, 0.999)
114
+ # SGD momentum
115
+ _C.TRAIN.OPTIMIZER.MOMENTUM = 0.9
116
+
117
+ # -----------------------------------------------------------------------------
118
+ # Augmentation settings
119
+ # -----------------------------------------------------------------------------
120
+ _C.AUG = CN()
121
+ # Color jitter factor
122
+ _C.AUG.COLOR_JITTER = 0.4
123
+ # Use AutoAugment policy. "v0" or "original"
124
+ _C.AUG.AUTO_AUGMENT = 'rand-m9-mstd0.5-inc1'
125
+ # Random erase prob
126
+ _C.AUG.REPROB = 0.25
127
+ # Random erase mode
128
+ _C.AUG.REMODE = 'pixel'
129
+ # Random erase count
130
+ _C.AUG.RECOUNT = 1
131
+ # Mixup alpha, mixup enabled if > 0
132
+ _C.AUG.MIXUP = 0.8
133
+ # Cutmix alpha, cutmix enabled if > 0
134
+ _C.AUG.CUTMIX = 1.0
135
+ # Cutmix min/max ratio, overrides alpha and enables cutmix if set
136
+ _C.AUG.CUTMIX_MINMAX = None
137
+ # Probability of performing mixup or cutmix when either/both is enabled
138
+ _C.AUG.MIXUP_PROB = 1.0
139
+ # Probability of switching to cutmix when both mixup and cutmix enabled
140
+ _C.AUG.MIXUP_SWITCH_PROB = 0.5
141
+ # How to apply mixup/cutmix params. Per "batch", "pair", or "elem"
142
+ _C.AUG.MIXUP_MODE = 'batch'
143
+
144
+ # -----------------------------------------------------------------------------
145
+ # Testing settings
146
+ # -----------------------------------------------------------------------------
147
+ _C.TEST = CN()
148
+ # Whether to use center crop when testing
149
+ _C.TEST.CROP = True
150
+
151
+ # -----------------------------------------------------------------------------
152
+ # Misc
153
+ # -----------------------------------------------------------------------------
154
+ # Mixed precision opt level, if O0, no amp is used ('O0', 'O1', 'O2')
155
+ # overwritten by command line argument
156
+ _C.AMP_OPT_LEVEL = ''
157
+ # Path to output folder, overwritten by command line argument
158
+ _C.OUTPUT = ''
159
+ # Tag of experiment, overwritten by command line argument
160
+ _C.TAG = 'default'
161
+ # Frequency to save checkpoint
162
+ _C.SAVE_FREQ = 1
163
+ # Frequency to logging info
164
+ _C.PRINT_FREQ = 10
165
+ # Fixed random seed
166
+ _C.SEED = 0
167
+ # Perform evaluation only, overwritten by command line argument
168
+ _C.EVAL_MODE = False
169
+ # Test throughput only, overwritten by command line argument
170
+ _C.THROUGHPUT_MODE = False
171
+ # local rank for DistributedDataParallel, given by command line argument
172
+ _C.LOCAL_RANK = 0
173
+
174
+
175
+ def _update_config_from_file(config, cfg_file):
176
+ config.defrost()
177
+ with open(cfg_file, 'r') as f:
178
+ yaml_cfg = yaml.load(f, Loader=yaml.FullLoader)
179
+
180
+ for cfg in yaml_cfg.setdefault('BASE', ['']):
181
+ if cfg:
182
+ _update_config_from_file(
183
+ config, os.path.join(os.path.dirname(cfg_file), cfg)
184
+ )
185
+ print('=> merge config from {}'.format(cfg_file))
186
+ config.merge_from_file(cfg_file)
187
+ config.freeze()
188
+
189
+
190
+ def update_config(config, args):
191
+ _update_config_from_file(config, args.cfg)
192
+
193
+ config.defrost()
194
+ if args.opts:
195
+ config.merge_from_list(args.opts)
196
+
197
+ # merge from specific arguments
198
+ if args.batch_size:
199
+ config.DATA.BATCH_SIZE = args.batch_size
200
+ if args.zip:
201
+ config.DATA.ZIP_MODE = True
202
+ if args.cache_mode:
203
+ config.DATA.CACHE_MODE = args.cache_mode
204
+ if args.resume:
205
+ config.MODEL.RESUME = args.resume
206
+ if args.accumulation_steps:
207
+ config.TRAIN.ACCUMULATION_STEPS = args.accumulation_steps
208
+ if args.use_checkpoint:
209
+ config.TRAIN.USE_CHECKPOINT = True
210
+ if args.amp_opt_level:
211
+ config.AMP_OPT_LEVEL = args.amp_opt_level
212
+ if args.tag:
213
+ config.TAG = args.tag
214
+ if args.eval:
215
+ config.EVAL_MODE = True
216
+ if args.throughput:
217
+ config.THROUGHPUT_MODE = True
218
+
219
+ config.freeze()
220
+
221
+
222
+ def get_config(args):
223
+ """Get a yacs CfgNode object with default values."""
224
+ # Return a clone so that the defaults will not be altered
225
+ # This is for the "local variable" use pattern
226
+ config = _C.clone()
227
+ update_config(config, args)
228
+
229
+ return config
code/sota/Swin-Unet/configs/swin_tiny_patch4_window7_224_lite.yaml ADDED
@@ -0,0 +1,12 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ MODEL:
2
+ TYPE: swin
3
+ NAME: swin_tiny_patch4_window7_224
4
+ DROP_PATH_RATE: 0.2
5
+ PRETRAIN_CKPT: "./pretrained_ckpt/swin_tiny_patch4_window7_224.pth"
6
+ SWIN:
7
+ FINAL_UPSAMPLE: "expand_first"
8
+ EMBED_DIM: 96
9
+ DEPTHS: [ 2, 2, 2, 2 ]
10
+ DECODER_DEPTHS: [ 2, 2, 2, 1]
11
+ NUM_HEADS: [ 3, 6, 12, 24 ]
12
+ WINDOW_SIZE: 7
code/sota/Swin-Unet/datasets/README.md ADDED
@@ -0,0 +1,29 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Data Preparing
2
+
3
+ 1. Access to the synapse multi-organ dataset:
4
+ 1. Sign up in the [official Synapse website](https://www.synapse.org/#!Synapse:syn3193805/wiki/) and download the dataset. Convert them to numpy format, clip the images within [-125, 275], normalize each 3D image to [0, 1], and extract 2D slices from 3D volume for training cases while keeping the 3D volume in h5 format for testing cases.
5
+ 2. You can also send an Email directly to jienengchen01 AT gmail.com to request the preprocessed data for reproduction.
6
+ 2. The directory structure of the whole project is as follows:
7
+
8
+ ```bash
9
+ .
10
+ ├── TransUNet
11
+ │   ├──datasets
12
+ │   │    └── dataset_*.py
13
+ │   ├──train.py
14
+ │   ├──test.py
15
+ │   └──...
16
+ ├── model
17
+ │   └── vit_checkpoint
18
+ │   └── imagenet21k
19
+ │      ├── R50+ViT-B_16.npz
20
+ │      └── *.npz
21
+ └── data
22
+ └──Synapse
23
+ ├── test_vol_h5
24
+ │   ├── case0001.npy.h5
25
+ │   └── *.npy.h5
26
+ └── train_npz
27
+ ├── case0005_slice000.npz
28
+ └── *.npz
29
+ ```
code/sota/Swin-Unet/datasets/__init__.py ADDED
File without changes
code/sota/Swin-Unet/datasets/dataset_synapse.py ADDED
@@ -0,0 +1,82 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import random
3
+
4
+ import h5py
5
+ import numpy as np
6
+ import torch
7
+ from scipy import ndimage
8
+ from scipy.ndimage.interpolation import zoom
9
+ from torch.utils.data import Dataset
10
+
11
+
12
+ def random_rot_flip(image, label):
13
+ k = np.random.randint(0, 4)
14
+ image = np.rot90(image, k)
15
+ label = np.rot90(label, k)
16
+ axis = np.random.randint(0, 2)
17
+ image = np.flip(image, axis=axis).copy()
18
+ label = np.flip(label, axis=axis).copy()
19
+ return image, label
20
+
21
+
22
+ def random_rotate(image, label):
23
+ angle = np.random.randint(-20, 20)
24
+ image = ndimage.rotate(image, angle, order=0, reshape=False)
25
+ label = ndimage.rotate(label, angle, order=0, reshape=False)
26
+ return image, label
27
+
28
+
29
+ class RandomGenerator(object):
30
+ def __init__(self, output_size):
31
+ self.output_size = output_size
32
+
33
+ def __call__(self, sample):
34
+ image, label = sample['image'], sample['label']
35
+
36
+ if random.random() > 0.5:
37
+ image, label = random_rot_flip(image, label)
38
+ elif random.random() > 0.5:
39
+ image, label = random_rotate(image, label)
40
+ x, y = image.shape
41
+ if x != self.output_size[0] or y != self.output_size[1]:
42
+ image = zoom(image, (self.output_size[0] / x, self.output_size[1] / y), order=3) # why not 3?
43
+ label = zoom(label, (self.output_size[0] / x, self.output_size[1] / y), order=0)
44
+ image = torch.from_numpy(image.astype(np.float32)).unsqueeze(0)
45
+ label = torch.from_numpy(label.astype(np.float32))
46
+ sample = {'image': image, 'label': label.long()}
47
+ return sample
48
+
49
+
50
+ class Synapse_dataset(Dataset):
51
+ def __init__(self, base_dir, list_dir, split, transform=None):
52
+ self.transform = transform # using transform in torch!
53
+ self.split = split
54
+ self.sample_list = open(os.path.join(list_dir, self.split + '.txt')).readlines()
55
+ self.data_dir = base_dir
56
+
57
+ def __len__(self):
58
+ return len(self.sample_list)
59
+
60
+ def __getitem__(self, idx):
61
+ if self.split in ["train", "val"] or self.sample_list[idx].strip('\n').split(",")[0].endswith(".npz"):
62
+ slice_name = self.sample_list[idx].strip('\n').split(",")[0]
63
+ if slice_name.endswith(".npz"):
64
+ data_path = os.path.join(self.data_dir, slice_name)
65
+ else:
66
+ data_path = os.path.join(self.data_dir, slice_name + '.npz')
67
+ data = np.load(data_path)
68
+ try:
69
+ image, label = data['image'], data['label']
70
+ except:
71
+ image, label = data['data'], data['seg']
72
+ else:
73
+ vol_name = self.sample_list[idx].strip('\n')
74
+ filepath = self.data_dir + "/{}.npy.h5".format(vol_name)
75
+ data = h5py.File(filepath)
76
+ image, label = data['image'][:], data['label'][:]
77
+
78
+ sample = {'image': image, 'label': label}
79
+ if self.transform:
80
+ sample = self.transform(sample)
81
+ sample['case_name'] = self.sample_list[idx].strip('\n')
82
+ return sample
code/sota/Swin-Unet/lists/lists_Synapse/all.lst ADDED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ case0031.npy.h5
2
+ case0007.npy.h5
3
+ case0009.npy.h5
4
+ case0005.npy.h5
5
+ case0026.npy.h5
6
+ case0039.npy.h5
7
+ case0024.npy.h5
8
+ case0034.npy.h5
9
+ case0033.npy.h5
10
+ case0030.npy.h5
11
+ case0023.npy.h5
12
+ case0040.npy.h5
13
+ case0010.npy.h5
14
+ case0021.npy.h5
15
+ case0006.npy.h5
16
+ case0027.npy.h5
17
+ case0028.npy.h5
18
+ case0037.npy.h5
19
+ case0008.npy.h5
20
+ case0022.npy.h5
21
+ case0038.npy.h5
22
+ case0036.npy.h5
23
+ case0032.npy.h5
24
+ case0002.npy.h5
25
+ case0029.npy.h5
26
+ case0003.npy.h5
27
+ case0001.npy.h5
28
+ case0004.npy.h5
29
+ case0025.npy.h5
30
+ case0035.npy.h5
code/sota/Swin-Unet/lists/lists_Synapse/test_vol.txt ADDED
@@ -0,0 +1,12 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ case0008
2
+ case0022
3
+ case0038
4
+ case0036
5
+ case0032
6
+ case0002
7
+ case0029
8
+ case0003
9
+ case0001
10
+ case0004
11
+ case0025
12
+ case0035