diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index a7b501f..572da8a 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -9,7 +9,7 @@ ci: repos: - repo: https://github.com/pre-commit/pre-commit-hooks - rev: v5.0.0 + rev: v6.0.0 hooks: - id: end-of-file-fixer - id: trailing-whitespace @@ -26,18 +26,18 @@ repos: args: ['--autofix', '--no-sort-keys', '--indent=4'] - id: end-of-file-fixer - id: mixed-line-ending - - repo: https://github.com/psf/black - rev: "24.10.0" + - repo: https://github.com/psf/black-pre-commit-mirror + rev: "26.10.0" hooks: - id: black - id: black-jupyter - repo: https://github.com/pycqa/isort - rev: 5.13.2 + rev: 9.0.2 hooks: - id: isort args: ["--profile", "black"] - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.8.6 + rev: v0.16.10 hooks: - id: ruff args: ['--fix'] diff --git a/vista2d/download_preprocessor/all_file_downloader.py b/vista2d/download_preprocessor/all_file_downloader.py index c8f4370..a423a80 100644 --- a/vista2d/download_preprocessor/all_file_downloader.py +++ b/vista2d/download_preprocessor/all_file_downloader.py @@ -49,7 +49,7 @@ def download_files(url_dict, directory): print(f"Saved to {file_path}") except Exception as e: - print(f"Failed to download from {key} ({url}). Reason: {str(e)}") + print(f"Failed to download from {key} ({url}). Reason: {e!s}") def main(): diff --git a/vista2d/download_preprocessor/generate_json.py b/vista2d/download_preprocessor/generate_json.py index 0453072..a84ce32 100644 --- a/vista2d/download_preprocessor/generate_json.py +++ b/vista2d/download_preprocessor/generate_json.py @@ -1164,7 +1164,6 @@ def main(): json_base_name = os.path.join(data_root_path, "json_files") tissuenet_json_files(dataset_dir=dataset_path, json_f_path=json_base_name) - return None if __name__ == "__main__": diff --git a/vista2d/download_preprocessor/process_data.py b/vista2d/download_preprocessor/process_data.py index 856dfc2..912a47d 100644 --- a/vista2d/download_preprocessor/process_data.py +++ b/vista2d/download_preprocessor/process_data.py @@ -469,7 +469,6 @@ def main(): dataset_path = os.path.join(data_root_path, key) tissuenet_process_files(dataset_dir=dataset_path) - return None if __name__ == "__main__": diff --git a/vista2d/scripts/cell_distributed_weighted_sampler.py b/vista2d/scripts/cell_distributed_weighted_sampler.py index dabc179..33b2ae6 100644 --- a/vista2d/scripts/cell_distributed_weighted_sampler.py +++ b/vista2d/scripts/cell_distributed_weighted_sampler.py @@ -12,7 +12,8 @@ # based on Pytorch DistributedSampler and WeightedRandomSampler combined import math -from typing import Iterator, Optional, Sequence, TypeVar +from collections.abc import Iterator, Sequence +from typing import TypeVar import torch import torch.distributed as dist @@ -29,8 +30,8 @@ def __init__( dataset: Dataset, weights: Sequence[float], num_samples: int, - num_replicas: Optional[int] = None, - rank: Optional[int] = None, + num_replicas: int | None = None, + rank: int | None = None, shuffle: bool = True, seed: int = 0, drop_last: bool = False, diff --git a/vista2d/scripts/workflow.py b/vista2d/scripts/workflow.py index 46076ca..699faff 100644 --- a/vista2d/scripts/workflow.py +++ b/vista2d/scripts/workflow.py @@ -854,7 +854,6 @@ def train(self): logger.warning( f"error copying {best_ckpt_path} {intermediate_ckpt_path} {err}" ) - pass if lr_scheduler is not None: lr_scheduler.step() @@ -1315,7 +1314,7 @@ def checkpoint_save(self, ckpt: str, model: torch.nn.Module, **kwargs): if self.config("compile", False): # remove key prefix of compiled models state_dict = OrderedDict( - (k[len("_orig_mod.") :] if k.startswith("_orig_mod.") else k, v) + (k.removeprefix("_orig_mod."), v) for k, v in state_dict.items() ) diff --git a/vista3d/README.md b/vista3d/README.md index 71deffe..68301e3 100644 --- a/vista3d/README.md +++ b/vista3d/README.md @@ -18,7 +18,7 @@ limitations under the License.
## News! -[10/27/2025] We release NV-Segment-CTMR, a joint CT-MR automatic segmentation model trained on over 30K CT and MRI scans, supporting over 300 classes. +[10/27/2025] We release NV-Segment-CTMR, a joint CT-MR automatic segmentation model trained on over 30K CT and MRI scans, supporting over 300 classes. [03/12/2025] We provide VISTA3D as a baseline for the challenge "CVPR 2025: Foundation Models for Interactive 3D Biomedical Image Segmentation"([link](https://www.codabench.org/competitions/5263/)). The simplified code based on MONAI 1.4 is provided in the [here](./cvpr_workshop/). @@ -40,8 +40,8 @@ limitations under the License. | **License** | [Commercial Friendly](https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-license/) | Same as VISTA3D | [Non-Commercial](https://developer.download.nvidia.com/licenses/NVIDIA-OneWay-Noncommercial-License-22Mar2022.pdf?t=eyJscyI6InJlZiIsImxzZCI6IlJFRi1naXRodWIuY29tL252aWRpYS1ob2xvc2NhbiJ9) | ``` -We recommend users to use NV-Segment-CTMR for large scale automatic segmentation for CT and MRI scans because it is trained with large and diverse datasets. For CT tumor or interactive refinement, user should try NV-Segment-CT. -``` +We recommend users to use NV-Segment-CTMR for large scale automatic segmentation for CT and MRI scans because it is trained with large and diverse datasets. For CT tumor or interactive refinement, user should try NV-Segment-CT. +``` **VISTA3D/NV-Segment-CT** ([`Paper`](https://arxiv.org/pdf/2406.05285)) is a foundation model trained systematically on 11,454 volumes encompassing 127 types of human anatomical structures and various lesions. The model provides State-of-the-art performances on: @@ -93,22 +93,22 @@ python -m monai.bundle run --config_file="['configs/inference.json', 'configs/ba # Automatic Batch segmentation for the whole folder with multi-gpu support. mgpu_inference.json is below. change nproc_per_node to your GPU number. torchrun --nproc_per_node=2 --nnodes=1 -m monai.bundle run --config_file="['configs/inference.json', 'configs/batch_inference.json', 'configs/mgpu_inference.json']" --input_dir="example/" --output_dir="example/" ``` -#### Interactive segmentation +#### Interactive segmentation ```bash # Points must be three dimensional (x,y,z) in the shape of [[x,y,z],...,[x,y,z]]. Point labels can only be -1(ignore), 0(negative), 1(positive) and 2(negative for special overlaped class like tumor), 3(positive for special class). Only supporting 1 class per inference. The output 255 represents NaN value which means not processed region. cd NV-Segment-CT python -m monai.bundle run --config_file configs/inference.json --input_dict "{'image':'example/spleen_03.nii.gz','points':[[128,128,16], [100,100,16]],'point_labels':[1, 0]}" ``` -**NOTE** MONAI bundle accepts multiple json config files and input arguments. The latter configs/arguments will overide the previous configs/arguments if they have overlapping keys. +**NOTE** MONAI bundle accepts multiple json config files and input arguments. The latter configs/arguments will overide the previous configs/arguments if they have overlapping keys. ## 1.2 **NV-Segment-CTMR**[[Github]](https://github.com/NVIDIA-Medtech/NV-Segment-CTMR/tree/main/NV-Segment-CTMR)[[Huggingface]](https://huggingface.co/nvidia/NV-Segment-CTMR/tree/main) -Please read the complete usage in the NV-Segment-CTMR [[Github]](https://github.com/NVIDIA-Medtech/NV-Segment-CTMR/tree/main/NV-Segment-CTMR) repo. +Please read the complete usage in the NV-Segment-CTMR [[Github]](https://github.com/NVIDIA-Medtech/NV-Segment-CTMR/tree/main/NV-Segment-CTMR) repo. -We defined 345 classes as in [metadata.json](https://github.com/NVIDIA-Medtech/NV-Segment-CTMR/blob/main/NV-Segment-CTMR/configs/metadata.json) and details in [label_dict.json](https://github.com/NVIDIA-Medtech/NV-Segment-CTMR/blob/main/NV-Segment-CTMR/configs/label_dict.json). It shows the label organ name, index, training dataset, modality and evaluation dice score. If a class only comes from CT training dataset, it may not perform well on MRI, but the actual performance will vary case by case. We support three type of segment everything "CT_BODY", "MRI_BODY", and "MRI_BRAIN". +We defined 345 classes as in [metadata.json](https://github.com/NVIDIA-Medtech/NV-Segment-CTMR/blob/main/NV-Segment-CTMR/configs/metadata.json) and details in [label_dict.json](https://github.com/NVIDIA-Medtech/NV-Segment-CTMR/blob/main/NV-Segment-CTMR/configs/label_dict.json). It shows the label organ name, index, training dataset, modality and evaluation dice score. If a class only comes from CT training dataset, it may not perform well on MRI, but the actual performance will vary case by case. We support three type of segment everything "CT_BODY", "MRI_BODY", and "MRI_BRAIN". -- "CT_BODY" is the previous VISTA3D bundle supported 132 CT classes. Same as NV-Segment-CT everything prompts. -- "MRI_BODY" shares the same 50 label class as TotalsegmentatorMR. +- "CT_BODY" is the previous VISTA3D bundle supported 132 CT classes. Same as NV-Segment-CT everything prompts. +- "MRI_BODY" shares the same 50 label class as TotalsegmentatorMR. - "MRI_BRAIN" is trained on skull stripped [LUMIR](https://github.com/JHU-MedImage-Reg/LUMIR_L2R) dataset and will segment 133 brain MRI substructures. We followed [MIR Preprocessing](https://github.com/junyuchen245/MIR/tree/main/tutorials/brain_MRI_preprocessing) tutorials and put the corresponding components into this repo. `All contrasts of brain MRI are supported` ### Quick Start @@ -144,7 +144,7 @@ For local `scripts.infer` usage in this research code, inference now prepares th ## 3. VISTA3D results postprocessing with [ShapeKit](https://arxiv.org/pdf/2506.24003) VISTA3D is trained with binary segmentation, and may produce false positives due to weak false positive supervision. ShapeKit solves this problem with sophisticated postprocessing. ShapeKit requires segmentation mask for each class. VISTA3D by default performs argmax and collaps overlapping classes. Change the `monai.apps.vista3d.transforms.VistaPostTransformd` in `inference.json` to save each class segmentation as a separate channel. Then follow [ShapeKit](https://github.com/BodyMaps/ShapeKit) codebase for processing. ```json -{ +{ "_target_": "Activationsd", "sigmoid": true, "keys": "pred" diff --git a/vista3d/README_research.md b/vista3d/README_research.md index df015ef..95bed29 100644 --- a/vista3d/README_research.md +++ b/vista3d/README_research.md @@ -80,7 +80,7 @@ To segment everything, run ```bash export CUDA_VISIBLE_DEVICES=0; python -m scripts.infer --config_file 'configs/infer.yaml' - infer_everything --image_file 'example-1.nii.gz' ``` -To segment based on point clicks, provide `point` and `point_label`. +To segment based on point clicks, provide `point` and `point_label`. ```bash export CUDA_VISIBLE_DEVICES=0; python -m scripts.infer --config_file 'configs/infer.yaml' - infer --image_file 'example-1.nii.gz' --point "[[128,128,16],[100,100,6]]" --point_label "[1,0]" --save_mask true ``` diff --git a/vista3d/cvpr_workshop/infer_cvpr.py b/vista3d/cvpr_workshop/infer_cvpr.py index 61d3536..0b3d7ef 100755 --- a/vista3d/cvpr_workshop/infer_cvpr.py +++ b/vista3d/cvpr_workshop/infer_cvpr.py @@ -16,6 +16,7 @@ from train_cvpr import ROI_SIZE + def convert_clicks(alldata): # indexes = list(alldata.keys()) # data = [alldata[i] for i in indexes] diff --git a/vista3d/cvpr_workshop/train_cvpr.py b/vista3d/cvpr_workshop/train_cvpr.py index ac099f6..1fe9eb1 100755 --- a/vista3d/cvpr_workshop/train_cvpr.py +++ b/vista3d/cvpr_workshop/train_cvpr.py @@ -7,10 +7,10 @@ import numpy as np import torch import torch.distributed as dist -import torch.optim as optim from monai.apps.vista3d.sampler import sample_prompt_pairs from monai.data import DataLoader from monai.networks.nets import vista3d132 +from torch import optim from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data import Dataset from torch.utils.tensorboard import SummaryWriter @@ -22,7 +22,8 @@ import matplotlib.pyplot as plt NUM_PATCHES_PER_IMAGE = 2 -ROI_SIZE= [128, 128, 128] +ROI_SIZE = [128, 128, 128] + def plot_to_tensorboard(writer, epoch, inputs, labels, points, outputs): """ @@ -109,7 +110,7 @@ def __getitem__(self, idx): keys=["image", "label"], label_key="label", num_classes=label.max() + 1, - ratios=tuple(float(i > 0) for i in range(label.max()+1)), + ratios=tuple(float(i > 0) for i in range(label.max() + 1)), num_samples=NUM_PATCHES_PER_IMAGE, ), monai.transforms.RandScaleIntensityd( @@ -137,28 +138,30 @@ def __getitem__(self, idx): mode=["constant", "constant"], keys=["image", "label"], spatial_size=ROI_SIZE, - ) + ), ] ) data = transforms(data) return data + import re + def get_latest_epoch(directory): # Pattern to match filenames like 'model_epoch.pth' - pattern = re.compile(r'model_epoch(\d+)\.pth') + pattern = re.compile(r"model_epoch(\d+)\.pth") max_epoch = -1 for filename in os.listdir(directory): match = pattern.match(filename) if match: epoch = int(match.group(1)) - if epoch > max_epoch: - max_epoch = epoch + max_epoch = max(max_epoch, epoch) return max_epoch if max_epoch != -1 else None + # Training function def train(): json_file = "allset.json" # Update with your JSON file @@ -169,7 +172,6 @@ def train(): start_epoch = get_latest_epoch(checkpoint_dir) start_checkpoint = "./CPRR25_vista3D_model_final_10percent_data.pth" - os.makedirs(checkpoint_dir, exist_ok=True) dist.init_process_group(backend="nccl") world_size = int(os.environ["WORLD_SIZE"]) @@ -189,11 +191,12 @@ def train(): model.load_state_dict(pretrained_ckpt, strict=True) else: print(f"Resuming from epoch {start_epoch}") - pretrained_ckpt = torch.load(os.path.join(checkpoint_dir, f"model_epoch{start_epoch}.pth")) - model.load_state_dict(pretrained_ckpt['model'], strict=True) + pretrained_ckpt = torch.load( + os.path.join(checkpoint_dir, f"model_epoch{start_epoch}.pth") + ) + model.load_state_dict(pretrained_ckpt["model"], strict=True) model = DDP(model, device_ids=[local_rank], find_unused_parameters=True) - optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=1.0e-05) lr_scheduler = monai.optimizers.WarmupCosineSchedule( optimizer=optimizer, @@ -265,10 +268,16 @@ def train(): if local_rank == 0: writer.add_scalar("loss", loss.item(), step) if local_rank == 0 and (epoch + 1) % save_interval == 0: - checkpoint_path = os.path.join(checkpoint_dir, f"model_epoch{epoch + 1}.pth") + checkpoint_path = os.path.join( + checkpoint_dir, f"model_epoch{epoch + 1}.pth" + ) if world_size > 1: torch.save( - {"model": model.module.state_dict(), "epoch": epoch + 1, "step": step}, + { + "model": model.module.state_dict(), + "epoch": epoch + 1, + "step": step, + }, checkpoint_path, ) print( diff --git a/vista3d/scripts/huggingface_download.py b/vista3d/scripts/huggingface_download.py index edf97da..5f3506a 100644 --- a/vista3d/scripts/huggingface_download.py +++ b/vista3d/scripts/huggingface_download.py @@ -1,7 +1,6 @@ import os import shutil from pathlib import Path -from typing import Optional def _is_rank_zero() -> bool: @@ -17,7 +16,7 @@ def touch_huggingface_download_counter( filename: str, revision: str = "main", rank_zero_only: bool = True, -) -> Optional[str]: +) -> str | None: """Force a tiny Hugging Face file request without re-downloading model weights.""" if rank_zero_only and not _is_rank_zero(): @@ -26,7 +25,9 @@ def touch_huggingface_download_counter( try: from huggingface_hub import hf_hub_download except ImportError: - print("[vista3d] warning: huggingface_hub is not installed; skipping Hugging Face download counter touch.") + print( + "[vista3d] warning: huggingface_hub is not installed; skipping Hugging Face download counter touch." + ) return None try: @@ -37,8 +38,10 @@ def touch_huggingface_download_counter( revision=revision, force_download=True, ) - except Exception as exc: # noqa: BLE001 - print(f"[vista3d] warning: could not touch Hugging Face download counter for {repo_id}/{filename}: {exc}") + except Exception as exc: + print( + f"[vista3d] warning: could not touch Hugging Face download counter for {repo_id}/{filename}: {exc}" + ) return None print(f"[vista3d] touched Hugging Face download counter for {repo_id}/{filename}") @@ -60,14 +63,18 @@ def prepare_huggingface_checkpoint( if rank_zero_only and not _is_rank_zero(): return str(local_path) - touch_huggingface_download_counter(repo_id, counter_filename, revision, rank_zero_only=False) + touch_huggingface_download_counter( + repo_id, counter_filename, revision, rank_zero_only=False + ) if local_path.exists(): return str(local_path) try: from huggingface_hub import hf_hub_download except ImportError as exc: - raise RuntimeError(f"{local_path} does not exist and huggingface_hub is not installed; cannot download {repo_id}.") from exc + raise RuntimeError( + f"{local_path} does not exist and huggingface_hub is not installed; cannot download {repo_id}." + ) from exc checkpoint_path = hf_hub_download( repo_id=repo_id, diff --git a/vista3d/scripts/infer.py b/vista3d/scripts/infer.py index 76bfb06..782a70f 100644 --- a/vista3d/scripts/infer.py +++ b/vista3d/scripts/infer.py @@ -145,7 +145,6 @@ def __init__(self, config_file="./configs/infer.yaml", **override): self.save_transforms = transforms.Compose(save_transforms) self.prev_mask = None self.batch_data = None - return def clear_cache(self): self.prev_mask = None diff --git a/vista3d/scripts/slic_process_sam.py b/vista3d/scripts/slic_process_sam.py index c73b69c..00ca7d1 100644 --- a/vista3d/scripts/slic_process_sam.py +++ b/vista3d/scripts/slic_process_sam.py @@ -65,7 +65,6 @@ def __init__(self): self.device ) self.model = SamPredictor(self.sam) - return @torch.no_grad() @torch.cuda.amp.autocast() diff --git a/vista3d/scripts/sliding_window.py b/vista3d/scripts/sliding_window.py index 67b2061..59e017e 100644 --- a/vista3d/scripts/sliding_window.py +++ b/vista3d/scripts/sliding_window.py @@ -13,8 +13,8 @@ import copy import itertools -from collections.abc import Callable, Mapping, Sequence -from typing import Any, Iterable +from collections.abc import Callable, Iterable, Mapping, Sequence +from typing import Any import monai import numpy as np @@ -42,7 +42,7 @@ tqdm, _ = optional_import("tqdm", name="tqdm") _nearest_mode = "nearest-exact" if pytorch_after(1, 11) else "nearest" -__all__ = ["sliding_window_inference", "point_based_window_inferer"] +__all__ = ["point_based_window_inferer", "sliding_window_inference"] def get_window_idx_c(p, roi, s): @@ -355,21 +355,21 @@ def sliding_window_inference( mode=look_up_option(padding_mode, PytorchPadMode), value=cval, ) - if "labels" in kwargs.keys() and kwargs["labels"] is not None: + if "labels" in kwargs and kwargs["labels"] is not None: kwargs["labels"] = F.pad( kwargs["labels"], pad=pad_size, mode=look_up_option(padding_mode, PytorchPadMode), value=cval, ) - if "prev_mask" in kwargs.keys() and kwargs["prev_mask"] is not None: + if "prev_mask" in kwargs and kwargs["prev_mask"] is not None: kwargs["prev_mask"] = F.pad( kwargs["prev_mask"], pad=pad_size, mode=look_up_option(padding_mode, PytorchPadMode), value=cval, ) - if "point_coords" in kwargs.keys() and kwargs["point_coords"] is not None: + if "point_coords" in kwargs and kwargs["point_coords"] is not None: kwargs["point_coords"] = kwargs["point_coords"] + torch.tensor( [pad_size[-2], pad_size[-4], pad_size[-6]] ).to(kwargs["point_coords"].device) diff --git a/vista3d/scripts/train.py b/vista3d/scripts/train.py index 0c89d16..7edb772 100644 --- a/vista3d/scripts/train.py +++ b/vista3d/scripts/train.py @@ -17,9 +17,9 @@ import sys import time import warnings +from collections.abc import Sequence from datetime import timedelta from functools import partial -from typing import Optional, Sequence, Union import monai import nibabel as nib @@ -115,7 +115,7 @@ def loss_wrapper(pred, label, loss_function): return loss_function(pred, label) -def run(config_file: Optional[Union[str, Sequence[str]]] = None, **override): +def run(config_file: str | Sequence[str] | None = None, **override): # Initialize distributed and scale parameters based on GPU memory if torch.cuda.device_count() > 1: dist.init_process_group( @@ -979,23 +979,14 @@ def run(config_file: Optional[Union[str, Sequence[str]]] = None, **override): yaml.dump([dict_file], stream=out_file) logger.debug( - "Current epoch: {} current mean dice: {:.4f} best mean dice: {:.4f} at epoch {}".format( - epoch, avg_metric, best_metric, best_metric_epoch - ) + f"Current epoch: {epoch} current mean dice: {avg_metric:.4f} best mean dice: {best_metric:.4f} at epoch {best_metric_epoch}" ) current_time = time.time() elapsed_time = (current_time - start_time) / 60.0 with open(os.path.join(ckpt_path, "accuracy_history.csv"), "a") as f: f.write( - "{:d}\t{:.5f}\t{:.5f}\t{:.5f}\t{:.1f}\t{:d}\n".format( - epoch, - avg_metric, - loss_torch_epoch, - lr, - elapsed_time, - idx_iter, - ) + f"{epoch:d}\t{avg_metric:.5f}\t{loss_torch_epoch:.5f}\t{lr:.5f}\t{elapsed_time:.1f}\t{idx_iter:d}\n" ) if world_size > 1: @@ -1016,7 +1007,6 @@ def run(config_file: Optional[Union[str, Sequence[str]]] = None, **override): if world_size > 1: dist.destroy_process_group() - return if __name__ == "__main__": diff --git a/vista3d/scripts/train_finetune.py b/vista3d/scripts/train_finetune.py index 02a5c15..bd848c1 100644 --- a/vista3d/scripts/train_finetune.py +++ b/vista3d/scripts/train_finetune.py @@ -18,9 +18,9 @@ import sys import time import warnings +from collections.abc import Sequence from datetime import timedelta from functools import partial -from typing import Optional, Sequence, Union import monai import nibabel as nib @@ -54,7 +54,7 @@ nib.imageglobals.logger.setLevel(40) -def run(config_file: Optional[Union[str, Sequence[str]]] = None, **override): +def run(config_file: str | Sequence[str] | None = None, **override): # Initialize distributed and scale parameters based on GPU memory if torch.cuda.device_count() > 1: dist.init_process_group( @@ -765,9 +765,7 @@ def run(config_file: Optional[Union[str, Sequence[str]]] = None, **override): yaml.dump([dict_file], stream=out_file) logger.debug( - "Current epoch: {} current mean dice: {:.4f} best mean dice: {:.4f} at epoch {}".format( - epoch, avg_metric, best_metric, best_metric_epoch - ) + f"Current epoch: {epoch} current mean dice: {avg_metric:.4f} best mean dice: {best_metric:.4f} at epoch {best_metric_epoch}" ) current_time = time.time() @@ -776,14 +774,7 @@ def run(config_file: Optional[Union[str, Sequence[str]]] = None, **override): os.path.join(ckpt_path, "accuracy_history.csv"), "a" ) as f: f.write( - "{:d}\t{:.5f}\t{:.5f}\t{:.5f}\t{:.1f}\t{:d}\n".format( - epoch, - avg_metric, - loss_torch_epoch, - lr, - elapsed_time, - idx_iter, - ) + f"{epoch:d}\t{avg_metric:.5f}\t{loss_torch_epoch:.5f}\t{lr:.5f}\t{elapsed_time:.1f}\t{idx_iter:d}\n" ) if world_size > 1: @@ -804,7 +795,6 @@ def run(config_file: Optional[Union[str, Sequence[str]]] = None, **override): if world_size > 1: dist.destroy_process_group() - return if __name__ == "__main__": diff --git a/vista3d/scripts/utils/trans_utils.py b/vista3d/scripts/utils/trans_utils.py index ec3ac4d..b943625 100644 --- a/vista3d/scripts/utils/trans_utils.py +++ b/vista3d/scripts/utils/trans_utils.py @@ -46,13 +46,13 @@ exposure, has_skimage = optional_import("skimage.exposure") __all__ = [ - "get_largest_connected_component_mask", "RandCropByLabelClassesShiftd", - "erode3d", - "erode2d", - "dilate3d", "VistaPostTransform", "convert_points_to_disc", + "dilate3d", + "erode2d", + "erode3d", + "get_largest_connected_component_mask", ] diff --git a/vista3d/scripts/utils/workflow_utils.py b/vista3d/scripts/utils/workflow_utils.py index 49692f9..e710601 100644 --- a/vista3d/scripts/utils/workflow_utils.py +++ b/vista3d/scripts/utils/workflow_utils.py @@ -389,7 +389,7 @@ def get_next_points_val( if mapped: # if in the global index, some supported classes need modification. neg_id, pos_id = get_point_label(prompt_class[id]) - if prompt_class[id].item() in MERGE_LIST.keys(): + if prompt_class[id].item() in MERGE_LIST: for m in MERGE_LIST[prompt_class[id].item()]: _gt = torch.logical_or(_gt, (gt == m)[0, 0]) fn_mask = torch.logical_and(_gt, pred[id][0] < pred_thresh) @@ -464,12 +464,12 @@ def get_next_points_auto_point( # if in the global index, some supported classes need modification. prompt_class is the local index if class_vector is not None: neg_id, pos_id = get_point_label(class_vector[id]) - if class_vector[id].item() in MERGE_LIST.keys(): + if class_vector[id].item() in MERGE_LIST: for m in MERGE_LIST[class_vector[id].item()]: _gt = torch.logical_or(_gt, (gt == m)[0, 0]) else: neg_id, pos_id = get_point_label(prompt_class[id]) - if prompt_class[id].item() in MERGE_LIST.keys(): + if prompt_class[id].item() in MERGE_LIST: for m in MERGE_LIST[prompt_class[id].item()]: _gt = torch.logical_or(_gt, (gt == m)[0, 0]) if (prompt_class[id] == 0 and not include_background) or _gt.sum() == 0: @@ -568,7 +568,7 @@ def get_next_points( if mapped: # if in the global index, some supported classes need modification. neg_id, pos_id = get_point_label(prompt_class[id]) - if prompt_class[id].item() in MERGE_LIST.keys(): + if prompt_class[id].item() in MERGE_LIST: for m in MERGE_LIST[prompt_class[id].item()]: _gt = torch.logical_or(_gt, (gt == m)[0, 0]) fn_mask = torch.logical_and(_gt, pred[id][0] < pred_thresh) diff --git a/vista3d/scripts/validation/val_multigpu_autopoint_patch.py b/vista3d/scripts/validation/val_multigpu_autopoint_patch.py index 61065e3..0a40734 100644 --- a/vista3d/scripts/validation/val_multigpu_autopoint_patch.py +++ b/vista3d/scripts/validation/val_multigpu_autopoint_patch.py @@ -17,9 +17,9 @@ import logging import os import sys +from collections.abc import Sequence from datetime import timedelta from functools import partial -from typing import Optional, Sequence, Union import monai import numpy as np @@ -42,7 +42,7 @@ from .build_vista3d_eval_only import vista_model_registry -def run(config_file: Optional[Union[str, Sequence[str]]] = None, **override): +def run(config_file: str | Sequence[str] | None = None, **override): # Initialize distributed and scale parameters based on GPU memory if torch.cuda.device_count() > 1: dist.init_process_group( @@ -432,7 +432,6 @@ def run(config_file: Optional[Union[str, Sequence[str]]] = None, **override): logger.debug(f"Final Evaluated Cases {_final_count}") dist.destroy_process_group() - return if __name__ == "__main__": diff --git a/vista3d/scripts/validation/val_multigpu_point_iterative.py b/vista3d/scripts/validation/val_multigpu_point_iterative.py index 2ac19f2..354442b 100644 --- a/vista3d/scripts/validation/val_multigpu_point_iterative.py +++ b/vista3d/scripts/validation/val_multigpu_point_iterative.py @@ -15,9 +15,9 @@ import logging import os import sys +from collections.abc import Sequence from datetime import timedelta from functools import partial -from typing import Optional, Sequence, Union import monai import torch @@ -40,7 +40,7 @@ from ..utils.workflow_utils import generate_prompt_pairs_val, get_next_points_val -def run(config_file: Optional[Union[str, Sequence[str]]] = None, **override): +def run(config_file: str | Sequence[str] | None = None, **override): # Initialize distributed and scale parameters based on GPU memory if torch.cuda.device_count() > 1: dist.init_process_group( @@ -382,7 +382,6 @@ def run(config_file: Optional[Union[str, Sequence[str]]] = None, **override): dist.barrier() dist.destroy_process_group() - return if __name__ == "__main__": diff --git a/vista3d/scripts/validation/val_multigpu_point_patch.py b/vista3d/scripts/validation/val_multigpu_point_patch.py index b02daf2..f1fa5d7 100644 --- a/vista3d/scripts/validation/val_multigpu_point_patch.py +++ b/vista3d/scripts/validation/val_multigpu_point_patch.py @@ -17,9 +17,9 @@ import logging import os import sys +from collections.abc import Sequence from datetime import timedelta from functools import partial -from typing import Optional, Sequence, Union import monai import numpy as np @@ -46,7 +46,7 @@ from ..utils.workflow_utils import sample_points_patch_val -def run(config_file: Optional[Union[str, Sequence[str]]] = None, **override): +def run(config_file: str | Sequence[str] | None = None, **override): # Initialize distributed and scale parameters based on GPU memory if torch.cuda.device_count() > 1: dist.init_process_group( @@ -460,7 +460,6 @@ def run(config_file: Optional[Union[str, Sequence[str]]] = None, **override): logger.debug(f"Final Evaluated Cases {_final_count}") dist.destroy_process_group() - return if __name__ == "__main__": diff --git a/vista3d/scripts/validation/val_multigpu_sam2_point_iterative.py b/vista3d/scripts/validation/val_multigpu_sam2_point_iterative.py index f47699a..5b467b0 100644 --- a/vista3d/scripts/validation/val_multigpu_sam2_point_iterative.py +++ b/vista3d/scripts/validation/val_multigpu_sam2_point_iterative.py @@ -18,8 +18,8 @@ import os import random import sys +from collections.abc import Sequence from datetime import timedelta -from typing import Optional, Sequence, Union import monai import numpy as np @@ -141,7 +141,7 @@ def get_points_from_false_pred(pred, gt, num_point=1): return _point, _label -def run(config_file: Optional[Union[str, Sequence[str]]] = None, **override): +def run(config_file: str | Sequence[str] | None = None, **override): # Initialize distributed and scale parameters based on GPU memory if torch.cuda.device_count() > 1: dist.init_process_group( @@ -530,7 +530,6 @@ def run(config_file: Optional[Union[str, Sequence[str]]] = None, **override): dist.barrier() dist.destroy_process_group() - return if __name__ == "__main__": diff --git a/vista3d/vista3d/modeling/class_head.py b/vista3d/vista3d/modeling/class_head.py index a55de51..9e1ea8d 100644 --- a/vista3d/vista3d/modeling/class_head.py +++ b/vista3d/vista3d/modeling/class_head.py @@ -1,6 +1,6 @@ import monai import torch -import torch.nn as nn +from torch import nn class Class_Mapping_Classify(nn.Module): diff --git a/vista3d/vista3d/modeling/point_head.py b/vista3d/vista3d/modeling/point_head.py index d8e6463..0108098 100644 --- a/vista3d/vista3d/modeling/point_head.py +++ b/vista3d/vista3d/modeling/point_head.py @@ -2,8 +2,8 @@ import numpy as np import torch -import torch.nn as nn from monai.utils import optional_import +from torch import nn from .sam_blocks import MLP, PositionEmbeddingRandom, TwoWayTransformer diff --git a/vista3d/vista3d/modeling/sam_blocks.py b/vista3d/vista3d/modeling/sam_blocks.py index d7bc413..6fef109 100644 --- a/vista3d/vista3d/modeling/sam_blocks.py +++ b/vista3d/vista3d/modeling/sam_blocks.py @@ -17,7 +17,7 @@ # LICENSE file in the root directory of this source tree. import math -from typing import Any, Optional, Tuple, Type +from typing import Any import numpy as np import torch @@ -32,7 +32,7 @@ def __init__( embedding_dim: int, num_heads: int, mlp_dim: int, - activation: Type[nn.Module] = nn.ReLU, + activation: type[nn.Module] = nn.ReLU, attention_downsample_rate: int = 2, ) -> None: """ @@ -76,7 +76,7 @@ def forward( image_embedding: Tensor, image_pe: Tensor, point_embedding: Tensor, - ) -> Tuple[Tensor, Tensor]: + ) -> tuple[Tensor, Tensor]: """ Args: image_embedding (torch.Tensor): image to attend to. Should be shape @@ -124,7 +124,7 @@ def __init__( embedding_dim: int, num_heads: int, mlp_dim: int = 2048, - activation: Type[nn.Module] = nn.ReLU, + activation: type[nn.Module] = nn.ReLU, attention_downsample_rate: int = 2, skip_first_layer_pe: bool = False, ) -> None: @@ -162,7 +162,7 @@ def __init__( def forward( self, queries: Tensor, keys: Tensor, query_pe: Tensor, key_pe: Tensor - ) -> Tuple[Tensor, Tensor]: + ) -> tuple[Tensor, Tensor]: # Self attention block if self.skip_first_layer_pe: queries = self.self_attn(q=queries, k=queries, v=queries) @@ -259,7 +259,7 @@ class PositionEmbeddingRandom(nn.Module): Positional encoding using random spatial frequencies. """ - def __init__(self, num_pos_feats: int = 64, scale: Optional[float] = None) -> None: + def __init__(self, num_pos_feats: int = 64, scale: float | None = None) -> None: super().__init__() if scale is None or scale <= 0.0: scale = 1.0 @@ -280,7 +280,7 @@ def _pe_encoding(self, coords: torch.Tensor) -> torch.Tensor: # [bs=1, N=2, 128+128=256] return torch.cat([torch.sin(coords), torch.cos(coords)], dim=-1) - def forward(self, size: Tuple[int, int, int]) -> torch.Tensor: + def forward(self, size: tuple[int, int, int]) -> torch.Tensor: """Generate positional encoding for a grid of the specified size.""" h, w, d = size device: Any = self.positional_encoding_gaussian_matrix.device @@ -295,7 +295,7 @@ def forward(self, size: Tuple[int, int, int]) -> torch.Tensor: return pe.permute(3, 0, 1, 2) # C x H x W def forward_with_coords( - self, coords_input: torch.Tensor, image_size: Tuple[int, int] + self, coords_input: torch.Tensor, image_size: tuple[int, int] ) -> torch.Tensor: """Positionally encode points that are not normalized to [0,1].""" coords = coords_input.clone() @@ -310,7 +310,7 @@ def __init__( self, embedding_dim: int, mlp_dim: int, - act: Type[nn.Module] = nn.GELU, + act: type[nn.Module] = nn.GELU, ) -> None: super().__init__() self.lin1 = nn.Linear(embedding_dim, mlp_dim) diff --git a/vista3d/vista3d/modeling/segresnetds.py b/vista3d/vista3d/modeling/segresnetds.py index e8c96d5..2c9bdc5 100644 --- a/vista3d/vista3d/modeling/segresnetds.py +++ b/vista3d/vista3d/modeling/segresnetds.py @@ -12,15 +12,14 @@ from __future__ import annotations from collections.abc import Callable -from typing import Union import numpy as np import torch -import torch.nn as nn from monai.networks.blocks.upsample import UpSample from monai.networks.layers.factories import Act, Conv, Norm, split_args from monai.networks.layers.utils import get_act_layer, get_norm_layer from monai.utils import UpsampleMode, has_option +from torch import nn __all__ = ["SegResNetDS2"] @@ -464,7 +463,7 @@ def is_valid_shape(self, x): def _forward( self, x: torch.Tensor, with_point, with_label - ) -> Union[None, torch.Tensor, list[torch.Tensor]]: + ) -> None | torch.Tensor | list[torch.Tensor]: if self.preprocess is not None: x = self.preprocess(x) @@ -522,7 +521,7 @@ def _forward( def forward( self, x: torch.Tensor, with_point=True, with_label=True, **kwargs - ) -> Union[None, torch.Tensor, list[torch.Tensor]]: + ) -> None | torch.Tensor | list[torch.Tensor]: return self._forward(x, with_point, with_label) def set_auto_grad(self, auto_freeze=False, point_freeze=False): diff --git a/vista3d/vista3d/modeling/vista3d.py b/vista3d/vista3d/modeling/vista3d.py index a67b660..cf4ed64 100644 --- a/vista3d/vista3d/modeling/vista3d.py +++ b/vista3d/vista3d/modeling/vista3d.py @@ -14,11 +14,11 @@ import monai import numpy as np import torch -import torch.nn as nn from monai.utils import optional_import from scripts.utils.trans_utils import convert_points_to_disc from scripts.utils.trans_utils import get_largest_connected_component_mask as lcc from scripts.utils.workflow_utils import sample_points_patch_val +from torch import nn rearrange, _ = optional_import("einops", name="rearrange") NINF_VALUE = -9999