code: complete eval pipeline (7 metrics + per-class + Wilcoxon) + Swin-UNet/TransUNet networks; remove backups/obsolete
Browse filesThis view is limited to 50 files because it contains too many changes. See raw diff
- code/envs/seggen-controlnet.yml +18 -0
- code/envs/umamba.yml +21 -12
- code/framework/common/subset.py +0 -40
- code/framework/config.py +0 -7
- code/framework/config.py.bak +0 -109
- code/framework/data/loaders.py +0 -1
- code/framework/data/loaders.py.bak +0 -40
- code/framework/data/unified_dataset.py +1 -8
- code/framework/data/unified_dataset.py.bak +0 -180
- code/framework/eval_ckpt.py +0 -120
- code/framework/metrics/boundary.py +0 -68
- code/framework/{common → synth}/__init__.py +0 -0
- code/framework/synth/generative_baselines.py +112 -0
- code/scripts/a100_nnum_eval512.sh +47 -0
- code/scripts/a100_swin_transunet_3seed_eval512.py +58 -0
- code/scripts/h800_cache_data.sh +14 -0
- code/scripts/h800_fetch_data.py +31 -0
- code/scripts/h800_parallel_extract.sh +38 -0
- code/scripts/h800_run_unified512.py +80 -0
- code/scripts/h800_setup_seggen.sh +28 -0
- code/scripts/h800_swin_transunet_eval512.py +68 -0
- code/scripts/hf_update_unified512.py +26 -0
- code/scripts/hf_upload_gensegdataset.py +23 -0
- code/scripts/hf_upload_tars.py +31 -0
- code/scripts/p1/backbones.py +44 -0
- code/scripts/p1/fd_lever.py +99 -0
- code/scripts/p1/fd_results.json +42 -0
- code/scripts/p1/fid_and_viz.py +130 -0
- code/scripts/p1/fid_fixed.py +37 -0
- code/scripts/p1/fid_results.json +14 -0
- code/scripts/p1/gen_prdc.json +74 -0
- code/scripts/p1/gen_prdc.py +58 -0
- code/scripts/p1/make_jit_vs_fd.py +65 -0
- code/scripts/p1/p1_busi_master.py +166 -0
- code/scripts/p1/p1_busi_results.json +152 -0
- code/scripts/p1/p1_full_metrics.json +182 -0
- code/scripts/p1/p1_gen_queue.sh +32 -0
- code/scripts/p1/p1_master.py +167 -0
- code/scripts/p1/p1_results.json +302 -0
- code/scripts/p1/smoke_backbone.py +76 -0
- code/scripts/p1/smoke_pixelgen.py +70 -0
- code/scripts/p1/train_fd_patched.py +179 -0
- code/sota/Swin-Unet/README.md +64 -0
- code/sota/Swin-Unet/config.py +229 -0
- code/sota/Swin-Unet/configs/swin_tiny_patch4_window7_224_lite.yaml +12 -0
- code/sota/Swin-Unet/datasets/README.md +29 -0
- code/sota/Swin-Unet/datasets/__init__.py +0 -0
- code/sota/Swin-Unet/datasets/dataset_synapse.py +82 -0
- code/sota/Swin-Unet/lists/lists_Synapse/all.lst +30 -0
- 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
|
| 2 |
-
#
|
| 3 |
-
#
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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
|