├── IMDLBenCo ├── utils │ ├── __init__.py │ ├── misc.py │ └── paths.py ├── model_zoo │ ├── span │ │ ├── __init__.py │ │ ├── README.md │ │ ├── SPAN.py │ │ └── PixelAttention.py │ ├── trufor │ │ ├── __init__.py │ │ ├── cmx │ │ │ ├── __init__.py │ │ │ ├── utils │ │ │ │ ├── __init__.py │ │ │ │ └── init_func.py │ │ │ ├── LICENSE_CMX.txt │ │ │ ├── layer_utils.py │ │ │ └── decoders │ │ │ │ └── MLPDecoder.py │ │ ├── config.py │ │ └── README.md │ ├── cat_net │ │ ├── __init__.py │ │ ├── README.md │ │ ├── cat_net.py │ │ └── cat_net_post_function.py │ ├── iml_vit │ │ ├── __init__.py │ │ └── README.md │ ├── mantra_net │ │ ├── __init__.py │ │ ├── README.md │ │ └── modify_weights.py │ ├── mesorch │ │ ├── __init__.py │ │ ├── extractor │ │ │ ├── __init__.py │ │ │ ├── high_frequency_feature_extraction.py │ │ │ └── low_frequency_feature_extraction.py │ │ └── README.md │ ├── mvss_net │ │ ├── __init__.py │ │ └── README.md │ ├── pscc_net │ │ ├── __init__.py │ │ ├── README.md │ │ └── seg_hrnet_config.py │ ├── sparse_vit │ │ ├── __init__.py │ │ └── README.md │ ├── object_former │ │ ├── __init__.py │ │ └── README.md │ ├── README.md │ └── __init__.py ├── training_scripts │ ├── __init__.py │ ├── utils │ │ └── __init__.py │ ├── trainer │ │ └── __init__.py │ ├── schedular │ │ ├── __init.py │ │ └── cos_lr_schedular.py │ └── tester │ │ └── __init__.py ├── modules │ ├── extractors │ │ ├── __init__.py │ │ ├── bayar_conv.py │ │ ├── srm_filter.py │ │ ├── sobel.py │ │ └── high_frequency_feature_extraction.py │ ├── __init__.py │ └── backbones │ │ ├── __init__.py │ │ ├── unet.py │ │ ├── resnet.py │ │ ├── swin.py │ │ └── vit16.py ├── cli_funcs │ ├── __init__.py │ ├── cli_env.py │ ├── copy_funcs.py │ └── cli_init.py ├── statics │ ├── base │ │ ├── test_complexity_mymodel.sh │ │ ├── test_mymodel.sh │ │ ├── test_robust_mymodel.sh │ │ ├── test_save_image_mymodel.sh │ │ ├── train_mymodel.sh │ │ ├── README-IMDLBenCo.md │ │ └── mymodel.py │ ├── model_zoo │ │ ├── runs │ │ │ ├── test_complexity │ │ │ │ ├── test_complexity_imlvit.sh │ │ │ │ └── test_complexity_mvss.sh │ │ │ ├── demo_test_mesorch.sh │ │ │ ├── demo_test_iml_vit.sh │ │ │ ├── demo_test_mantra_net.sh │ │ │ ├── demo_test_object_former.sh │ │ │ ├── demo_test_cat_net.sh │ │ │ ├── demo_test_mvss.sh │ │ │ ├── demo_test_robustness_mantra_net.sh │ │ │ ├── demo_test_robustness_mesorch.sh │ │ │ ├── demo_test_robustness_iml_vit.sh │ │ │ ├── test_save_images │ │ │ │ ├── save_image_imlvit.sh │ │ │ │ ├── save_image_mvss.sh │ │ │ │ ├── save_image_catnet.sh │ │ │ │ └── save_image_pscc.sh │ │ │ ├── demo_test_robustness_cat_net.sh │ │ │ ├── demo_test_span.sh │ │ │ ├── demo_test_robustness_mvss.sh │ │ │ ├── demo_test_robustness_object_former.sh │ │ │ ├── demo_test_robustness_span.sh │ │ │ ├── demo_test_pscc.sh │ │ │ ├── demo_train_backbone.sh │ │ │ ├── demo_train_mantra_net.sh │ │ │ ├── demo_test_robustness_pscc.sh │ │ │ ├── demo_train_cat_net.sh │ │ │ ├── demo_catnet_protocol_mvss.sh │ │ │ ├── demo_train_object_former.sh │ │ │ ├── demo_train_mvss.sh │ │ │ ├── demo_train_span.sh │ │ │ ├── demo_train_sparse_vit.sh │ │ │ ├── demo_train_backbone_segformer.sh │ │ │ ├── demo_catnet_protocol_iml_vit.sh │ │ │ ├── demo_test_trufor.sh │ │ │ ├── demo_train_mesorch.sh │ │ │ ├── demo_test_robustness_trufor.sh │ │ │ ├── demo_train_iml_vit.sh │ │ │ ├── demo_train_pscc.sh │ │ │ └── demo_train_trufor.sh │ │ └── configs │ │ │ ├── trufor.yaml │ │ │ └── CAT_full.yaml │ └── dataset_json │ │ ├── test_datasets.json │ │ └── balanced_dataset.json ├── transforms │ ├── __init__.py │ ├── robustness_wrapper.py │ └── edge_mask_generator.py ├── datasets │ ├── __init__.py │ ├── dummy_dataset.py │ ├── utils.py │ ├── iml_datasets.py │ └── balanced_dataset.py ├── build_functions.py ├── __init__.py ├── evaluation │ ├── __init__.py │ ├── gradcam │ │ ├── grad_cam_hack.py │ │ ├── activations_and_gradients_hack.py │ │ └── grad_camera_visualize.py │ ├── temp.py │ └── abstract_class.py ├── version.py └── cli.py ├── codecov.yml ├── tests ├── .gitignore ├── test_registry_bug.py ├── test_metrics │ ├── test_datasets.json │ ├── README.md │ ├── test_mymodel.sh │ ├── test_mymodel_4_5.sh │ ├── test_mymodel_4_5_revised_imgAcc.sh │ ├── dataset.json │ ├── generate_dataset.py │ ├── mymodel.py │ └── reference.py ├── test_concat.py ├── test_registers_dataset.py ├── test_registers_module.py ├── test_version.py ├── test_evaluators.py ├── test_subdataset.py ├── test_model_auto_parser.py ├── test_json_datasets.py ├── gradcam.py ├── test_openfile.py ├── test_metrics_cpu.py ├── pytests │ ├── test_evaluators_empty.py │ └── test_F1.py ├── test_datasets_and_evaluation.py ├── test_appdirs.py └── test_tail_dataset.py ├── requirements-dev.txt ├── images ├── IMDL_BenCo.png └── IMDLBenCo_overview.png ├── release_to_pypi.sh ├── .gitignore ├── release_to_test_pypi.sh ├── requirements.txt ├── MANIFEST.in ├── .github ├── ISSUE_TEMPLATE │ ├── 3-feature-request.yaml │ ├── 2-ask-for-help.yaml │ └── 1-bug-reports.yaml └── workflows │ ├── python-publish.yml │ └── tests.yml ├── pyproject.toml └── setup_backup.py /IMDLBenCo/utils/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /codecov.yml: -------------------------------------------------------------------------------- 1 | comment: false 2 | -------------------------------------------------------------------------------- /IMDLBenCo/model_zoo/span/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /IMDLBenCo/model_zoo/trufor/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /IMDLBenCo/training_scripts/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /tests/.gitignore: -------------------------------------------------------------------------------- 1 | output_masks_shape -------------------------------------------------------------------------------- /IMDLBenCo/model_zoo/cat_net/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /IMDLBenCo/model_zoo/iml_vit/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /IMDLBenCo/model_zoo/mantra_net/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /IMDLBenCo/model_zoo/mesorch/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /IMDLBenCo/model_zoo/mvss_net/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /IMDLBenCo/model_zoo/pscc_net/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /IMDLBenCo/model_zoo/sparse_vit/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /IMDLBenCo/model_zoo/trufor/cmx/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /IMDLBenCo/modules/extractors/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /IMDLBenCo/model_zoo/object_former/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /IMDLBenCo/model_zoo/trufor/cmx/utils/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /IMDLBenCo/training_scripts/utils/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /IMDLBenCo/model_zoo/mesorch/extractor/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /requirements-dev.txt: -------------------------------------------------------------------------------- 1 | flake8 2 | pytest 3 | pytest-cov -------------------------------------------------------------------------------- /IMDLBenCo/modules/__init__.py: -------------------------------------------------------------------------------- 1 | from .backbones import * 2 | from .extractors import * -------------------------------------------------------------------------------- /tests/test_registry_bug.py: -------------------------------------------------------------------------------- 1 | from IMDLBenCo import MODELS 2 | MODELS.get("IML_VIT") -------------------------------------------------------------------------------- /images/IMDL_BenCo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/scu-zjz/IMDLBenCo/HEAD/images/IMDL_BenCo.png -------------------------------------------------------------------------------- /images/IMDLBenCo_overview.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/scu-zjz/IMDLBenCo/HEAD/images/IMDLBenCo_overview.png -------------------------------------------------------------------------------- /IMDLBenCo/model_zoo/mvss_net/README.md: -------------------------------------------------------------------------------- 1 | MVSS-Net can load pre-trained weights by it self. No need to download mannually. -------------------------------------------------------------------------------- /IMDLBenCo/training_scripts/trainer/__init__.py: -------------------------------------------------------------------------------- 1 | from .trainer import train_one_epoch 2 | 3 | 4 | __all__ = ["train_one_epoch"] -------------------------------------------------------------------------------- /tests/test_metrics/test_datasets.json: -------------------------------------------------------------------------------- 1 | { 2 | "Dummy": "/mnt/data0/xiaochen/workspace/IMDLBenCo_pure/test_metrics/dataset.json" 3 | } -------------------------------------------------------------------------------- /IMDLBenCo/cli_funcs/__init__.py: -------------------------------------------------------------------------------- 1 | from .cli_init import cli_init 2 | from .cli_env import cli_env 3 | 4 | 5 | __all__ = ["cli_env", "cli_init"] -------------------------------------------------------------------------------- /IMDLBenCo/model_zoo/pscc_net/README.md: -------------------------------------------------------------------------------- 1 | The HRNET pretrain weights are from the official repository: https://github.com/proteus1991/PSCC-Net. -------------------------------------------------------------------------------- /IMDLBenCo/training_scripts/schedular/__init.py: -------------------------------------------------------------------------------- 1 | from cos_lr_schedular import adjust_learning_rate 2 | 3 | __all__ = ['adjust_learning_rate'] -------------------------------------------------------------------------------- /tests/test_concat.py: -------------------------------------------------------------------------------- 1 | import torch 2 | a = None 3 | b = torch.tensor([1, 2, 3]) 4 | print(torch.cat([b], dim=0)) 5 | print(torch.cat([a, b], dim=0)) -------------------------------------------------------------------------------- /IMDLBenCo/model_zoo/span/README.md: -------------------------------------------------------------------------------- 1 | Please download pretrained weights IMTFEv4.pt from [ManTraNet-pytorch](https://github.com/RonyAbecidan/ManTraNet-pytorch) -------------------------------------------------------------------------------- /IMDLBenCo/statics/base/test_complexity_mymodel.sh: -------------------------------------------------------------------------------- 1 | python ./test_complexity.py \ 2 | --model MyModel \ 3 | --test_batch_size 1 \ 4 | --image_size 512 \ 5 | --if_resizing -------------------------------------------------------------------------------- /IMDLBenCo/training_scripts/tester/__init__.py: -------------------------------------------------------------------------------- 1 | from .tester import test_one_epoch, inference_and_save_one_epoch 2 | 3 | __all__ = ["test_one_epoch", "inference_and_save_one_epoch"] -------------------------------------------------------------------------------- /IMDLBenCo/modules/backbones/__init__.py: -------------------------------------------------------------------------------- 1 | from .resnet import * 2 | from .vit import * 3 | from .vit16 import * 4 | from .swin import * 5 | from .unet import * 6 | from .segformer import * 7 | -------------------------------------------------------------------------------- /release_to_pypi.sh: -------------------------------------------------------------------------------- 1 | # Only for Developers!!!!!!!!! 2 | rm -rf build dist *.egg-info 3 | python setup.py sdist bdist_wheel 4 | twine upload --repository-url https://upload.pypi.org/legacy/ dist/* -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | test 2 | __pycache__ 3 | .* 4 | output_dir* 5 | eval_dir* 6 | robust_dir* 7 | *.log 8 | *.png 9 | *.jpg 10 | dist_test.py 11 | ckpt/ 12 | 13 | *.egg-info 14 | build 15 | dist -------------------------------------------------------------------------------- /release_to_test_pypi.sh: -------------------------------------------------------------------------------- 1 | # Only for Developers!!!!!!!!! 2 | rm -rf build dist *.egg-info 3 | python setup.py sdist bdist_wheel 4 | twine upload --repository-url https://test.pypi.org/legacy/ dist/* -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | rich 2 | fvcore 3 | albumentations==1.3.0 4 | matplotlib 5 | numpy<2 6 | scikit_learn 7 | timm 8 | fvcore 9 | tensorboard 10 | grad_cam 11 | ttach 12 | appdirs 13 | colorama -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/runs/test_complexity/test_complexity_imlvit.sh: -------------------------------------------------------------------------------- 1 | python ./test_complexity.py \ 2 | --model IML_ViT \ 3 | --test_batch_size 1 \ 4 | --image_size 1024 \ 5 | --if_padding -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/runs/test_complexity/test_complexity_mvss.sh: -------------------------------------------------------------------------------- 1 | python ./test_complexity.py \ 2 | --model MVSSNet \ 3 | --test_batch_size 1 \ 4 | --edge_mask_width 7 \ 5 | --image_size 512 \ 6 | --if_resizing -------------------------------------------------------------------------------- /IMDLBenCo/model_zoo/object_former/README.md: -------------------------------------------------------------------------------- 1 | ObjectFormer is reproduced by us, and it requires processing the `ViT` backbone to obtain the pretrained weights. The detailed processing code can be found in [this issue](https://github.com/scu-zjz/IMDLBenCo/issues/41). -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | # for help : https://setuptools.pypa.io/en/latest/userguide/datafiles.html 2 | include IMDLBenCo/*.py 3 | recursive-include IMDLBenCo * 4 | 5 | include README.md 6 | include requirements.txt 7 | include LICENSE 8 | 9 | recursive-include configs * -------------------------------------------------------------------------------- /IMDLBenCo/transforms/__init__.py: -------------------------------------------------------------------------------- 1 | from .iml_transforms import get_albu_transforms, RandomCopyMove, RandomInpainting 2 | from .edge_mask_generator import EdgeMaskGenerator 3 | 4 | __all__ = ['get_albu_transforms', 'RandomCopyMove', 'RandomInpainting', 'EdgeMaskGenerator'] -------------------------------------------------------------------------------- /IMDLBenCo/model_zoo/mantra_net/README.md: -------------------------------------------------------------------------------- 1 | Please download pretrained weights MantraNetv4.pt from [ManTraNet-pytorch](https://github.com/RonyAbecidan/ManTraNet-pytorch) and put it under this directory. 2 | 3 | Then run modify_weights.py to modify the weight into appropriate form. -------------------------------------------------------------------------------- /IMDLBenCo/model_zoo/README.md: -------------------------------------------------------------------------------- 1 | # MODEL_ZOO 2 | Here is the implementation of all models in `model_zoo` of IMDLBenCo. 3 | 4 | You can find the guide to download the pretrained weights of each model under their **sub-folders** here, which is must when you want train each model by yourself. -------------------------------------------------------------------------------- /IMDLBenCo/datasets/__init__.py: -------------------------------------------------------------------------------- 1 | from .iml_datasets import ManiDataset, JsonDataset 2 | from .balanced_dataset import BalancedDataset 3 | from .utils import denormalize 4 | from .dummy_dataset import DummyDataset 5 | 6 | __all__ = ['ManiDataset', "JsonDataset", "BalancedDataset", "denormalize", "DummyDataset"] -------------------------------------------------------------------------------- /tests/test_registers_dataset.py: -------------------------------------------------------------------------------- 1 | import sys 2 | sys.path.append("./") 3 | # from IMDLBench.datasets import DATASETS 4 | from IMDLBenCo.registry import DATASETS 5 | 6 | print(DATASETS) 7 | obj = DATASETS.build("ManiDataset", path="/mnt/data0/public_datasets/IML/basic_eval_dataset", is_padding=True) 8 | print(obj) -------------------------------------------------------------------------------- /tests/test_registers_module.py: -------------------------------------------------------------------------------- 1 | import sys 2 | sys.path.append("./") 3 | # from IMDLBench.datasets import DATASETS 4 | from IMDLBenCo.registry import MODELS 5 | 6 | print(MODELS) 7 | # obj = MODELS.build("ManiDataset", path="/mnt/data0/public_datasets/IML/basic_eval_dataset", is_padding=True) 8 | # print(obj) -------------------------------------------------------------------------------- /IMDLBenCo/model_zoo/sparse_vit/README.md: -------------------------------------------------------------------------------- 1 | ## Checkpoint for SparseViT 2 | - You can download the recommend checkpoint from this link: 3 | - [https://huggingface.co/Sense-X/uniformer_image/tree/main](https://huggingface.co/Sense-X/uniformer_image/tree/main) 4 | - You can load this file directly to train SparseViT: `uniformer_base_in1k.pth` -------------------------------------------------------------------------------- /IMDLBenCo/model_zoo/mesorch/README.md: -------------------------------------------------------------------------------- 1 | Mesorch requires two sets of pretrained weights. The ConvNeXt weights will be automatically downloaded when `conv_pretrain` is set to `True`. 2 | As for the SegFormer weights, please visit [https://github.com/NVlabs/SegFormer](https://github.com/NVlabs/SegFormer) and manually download the **MiT-B3** pretrained model. 3 | -------------------------------------------------------------------------------- /IMDLBenCo/build_functions.py: -------------------------------------------------------------------------------- 1 | 2 | def build_from_cfg(self): 3 | # TODO 4 | pass 5 | 6 | 7 | # https://github.com/open-mmlab/mmdetection/blob/cfd5d3a985b0249de009b67d04f37263e11cdf3d/mmdet/datasets/transforms/transforms.py#L1646 8 | 9 | 10 | # https://github.com/open-mmlab/mmengine/blob/66fb81f7b392b2cd304fc1979d8af3cc71a011f5/mmengine/registry/build_functions.py#L17 -------------------------------------------------------------------------------- /IMDLBenCo/__init__.py: -------------------------------------------------------------------------------- 1 | from .version import __version__, version_info 2 | from .registry import Registry, MODELS, DATASETS, POSTFUNCS 3 | 4 | from .datasets import * # Registry all Datasets to the DATASETS register 5 | from .model_zoo import * 6 | from .modules import * 7 | 8 | __all__ = ['__version__', 'version_info', 'MODELS', "DATASETS", "POSTFUNCS"] 9 | 10 | 11 | 12 | -------------------------------------------------------------------------------- /tests/test_metrics/README.md: -------------------------------------------------------------------------------- 1 | ## 测试指标是否正确，交叉验证的脚本（仅供官方开发人员留档，IMDLBenCo的具体使用请参考https://scu-zjz.github.io/IMDLBenCo-doc/zh/guide/quickstart/2_load_ckpt.html） 2 | 用法： 3 | 1. 首先执行`generate_dataset.py`，然后获得一个测试用的数据集。 4 | 2. 修改`test_mymodel.sh`中的路径指向该dataset。 5 | 3. `Mymodel.py`中已经写好了一个20个sample的样例，且强制按照文件名有固定的label输出。 6 | 4. 运行`reference.py`获得CPU的sklearn在该数据集上的评估指标。 7 | 5. 运行`sh test_mymodel.sh`获得多卡输出。 8 | 6. 比较二者，以验证是否一致。 9 | -------------------------------------------------------------------------------- /tests/test_version.py: -------------------------------------------------------------------------------- 1 | 2 | from IMDLBenCo.version import __version__, parse_version_info 3 | 4 | current_version = '0.2.30' 5 | future_version = '0.2.30' 6 | 7 | print(current_version < future_version) 8 | 9 | 10 | parsed_current = parse_version_info(current_version) 11 | parsed_future = parse_version_info(future_version) 12 | print(parsed_current) 13 | print(parsed_future) 14 | print(parsed_current < parsed_future) -------------------------------------------------------------------------------- /tests/test_evaluators.py: -------------------------------------------------------------------------------- 1 | import sys 2 | from pprint import pprint 3 | sys.path.append(".") 4 | 5 | import IMDLBenCo 6 | from IMDLBenCo.evaluation import PixelF1, ImageF1 7 | 8 | list_e = [ 9 | PixelF1(), 10 | ImageF1() 11 | ] 12 | 13 | 14 | print(list_e[0], list_e[1]) 15 | 16 | 17 | 18 | def print_fun(obj): 19 | for i in obj: 20 | print(i) 21 | 22 | 23 | 24 | print_fun(list_e) -------------------------------------------------------------------------------- /IMDLBenCo/model_zoo/mantra_net/modify_weights.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | weight_path = 'MantraNetv4.pt' 4 | state_dict = torch.load(weight_path) 5 | 6 | new_state_dict = {'model': {}} 7 | for key, value in state_dict.items(): 8 | new_state_dict['model'][key] = value 9 | 10 | new_weight_path = '../../../eval_dir/checkpoint-0.pth' 11 | torch.save(new_state_dict, new_weight_path) 12 | 13 | print(f"Modified weights saved to {new_weight_path}") 14 | -------------------------------------------------------------------------------- /IMDLBenCo/statics/dataset_json/test_datasets.json: -------------------------------------------------------------------------------- 1 | { 2 | "Columbia": "/mnt/data0/public_datasets/IML/Columbia.json", 3 | "NIST16_1024": "/mnt/data0/public_datasets/IML/NIST16_1024", 4 | "NIST16_cleaned": "/mnt/data0/public_datasets/IML/NIST16_1024_cleaning", 5 | "coverage": "/mnt/data0/public_datasets/IML/coverage.json", 6 | "CASIAv1": "/mnt/data0/public_datasets/IML/CASIA1.0", 7 | "IMD20_1024": "/mnt/data0/public_datasets/IML/IMD_20_1024" 8 | } -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/configs/trufor.yaml: -------------------------------------------------------------------------------- 1 | CUDNN: 2 | BENCHMARK: false 3 | DETERMINISTIC: false 4 | ENABLED: false 5 | WORKERS: 16 6 | 7 | DATASET: 8 | NUM_CLASSES: 2 9 | MODEL: 10 | NAME: detconfcmx 11 | MODS: ('RGB','NP++') 12 | EXTRA: 13 | BACKBONE: mit_b2 14 | DECODER: MLPDecoder 15 | DECODER_EMBED_DIM: 512 16 | PREPRC: 'imagenet' 17 | BN_EPS: 0.001 18 | BN_MOMENTUM: 0.1 19 | DETECTION: 'confpool' 20 | CONF: true 21 | TEST: 22 | MODEL_FILE: '../weights/trufor.pth.tar' -------------------------------------------------------------------------------- /IMDLBenCo/model_zoo/cat_net/README.md: -------------------------------------------------------------------------------- 1 | Please refer to [CAT-Net official repository](https://github.com/mjkwon2021/CAT-Net) to download pre-trained weights. 2 | 3 | It should include two files: 4 | - `hrnetv2_w48_imagenet_pretrained.pth` 5 | - `DCT_djpeg.pth.tar` 6 | 7 | And you need to revise the path of following key-value pairs in `CAT_full.yaml` under the `config` dir after `benco init model_zoo`: 8 | ```yaml 9 | PRETRAINED_RGB: '/mnt/data0/bingkui/Cat-Net/hrnetv2_w48_imagenet_pretrained.pth' 10 | PRETRAINED_DCT: '/mnt/data0/bingkui/Cat-Net/DCT_djpeg.pth.tar' 11 | ``` -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/runs/demo_test_mesorch.sh: -------------------------------------------------------------------------------- 1 | base_dir="./eval_dir_mesorch_f1" 2 | mkdir -p ${base_dir} 3 | 4 | CUDA_VISIBLE_DEVICES=0,1,2,3 \ 5 | torchrun \ 6 | --standalone \ 7 | --nnodes=1 \ 8 | --nproc_per_node=4 \ 9 | ./test.py \ 10 | --model Mesorch \ 11 | --world_size 4 \ 12 | --test_data_json "./test_datasets.json" \ 13 | --checkpoint_path "./ckpt_mesorch/" \ 14 | --test_batch_size 2 \ 15 | --image_size 512 \ 16 | --if_resizing \ 17 | --output_dir ${base_dir}/ \ 18 | --log_dir ${base_dir}/ \ 19 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /tests/test_metrics/test_mymodel.sh: -------------------------------------------------------------------------------- 1 | base_dir="./eval_dir" 2 | mkdir -p ${base_dir} 3 | 4 | CUDA_VISIBLE_DEVICES=0,1,2,3 \ 5 | torchrun \ 6 | --standalone \ 7 | --nnodes=1 \ 8 | --nproc_per_node=4 \ 9 | ./test.py \ 10 | --model MyModel \ 11 | --world_size 1 \ 12 | --test_data_json /mnt/data0/xiaochen/workspace/IMDLBenCo_pure/test_metrics/test_datasets.json \ 13 | --checkpoint_path "./output_dir/" \ 14 | --test_batch_size 3 \ 15 | --image_size 512 \ 16 | --if_resizing \ 17 | --output_dir ${base_dir}/ \ 18 | --log_dir ${base_dir}/ \ 19 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /tests/test_metrics/test_mymodel_4_5.sh: -------------------------------------------------------------------------------- 1 | base_dir="./eval_dir_4_5" 2 | mkdir -p ${base_dir} 3 | 4 | CUDA_VISIBLE_DEVICES=0,1,2,3 \ 5 | torchrun \ 6 | --standalone \ 7 | --nnodes=1 \ 8 | --nproc_per_node=4 \ 9 | ./test.py \ 10 | --model MyModel \ 11 | --world_size 1 \ 12 | --test_data_json /mnt/data0/xiaochen/workspace/IMDLBenCo_pure/test_metrics/test_datasets.json \ 13 | --checkpoint_path "./output_dir/" \ 14 | --test_batch_size 5 \ 15 | --image_size 512 \ 16 | --if_resizing \ 17 | --output_dir ${base_dir}/ \ 18 | --log_dir ${base_dir}/ \ 19 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /IMDLBenCo/model_zoo/iml_vit/README.md: -------------------------------------------------------------------------------- 1 | - You should download MAE pre-trained weights from the official repo of the paper [Masked Autoencoders are scalable vision learners](https://arxiv.org/abs/2111.06377) before start training(fine-tuning) with our IML-ViT. 2 | - [official repo of MAE](https://github.com/facebookresearch/mae) 3 | - [download link of MAE pre-trained ViT-B](https://dl.fbaipublicfiles.com/mae/pretrain/mae_pretrain_vit_base.pth) 4 | 5 | - By default, our commend for training takes this dir as the place where pre-trained weights are. Therefore, you should place your downloaded weights in this directory, or re-direct to your custom dir of MAE weights. -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/runs/demo_test_iml_vit.sh: -------------------------------------------------------------------------------- 1 | base_dir="./eval_dir" 2 | mkdir -p ${base_dir} 3 | 4 | CUDA_VISIBLE_DEVICES=0 \ 5 | torchrun \ 6 | --standalone \ 7 | --nnodes=1 \ 8 | --nproc_per_node=1 \ 9 | ./test.py \ 10 | --model IML_ViT \ 11 | --edge_mask_width 7 \ 12 | --world_size 1 \ 13 | --test_data_json "./test_datasets.json" \ 14 | --checkpoint_path "/mnt/data0/xiaochen/workspace/IMDLBench_dev/output_dir" \ 15 | --test_batch_size 1 \ 16 | --image_size 1024 \ 17 | --if_padding \ 18 | --output_dir ${base_dir}/ \ 19 | --log_dir ${base_dir}/ \ 20 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /tests/test_metrics/test_mymodel_4_5_revised_imgAcc.sh: -------------------------------------------------------------------------------- 1 | base_dir="./eval_dir_4_3_imgacc_revised_new" 2 | mkdir -p ${base_dir} 3 | 4 | CUDA_VISIBLE_DEVICES=0,1,2,3 \ 5 | torchrun \ 6 | --standalone \ 7 | --nnodes=1 \ 8 | --nproc_per_node=4 \ 9 | ./test.py \ 10 | --model MyModel \ 11 | --world_size 1 \ 12 | --test_data_json /mnt/data0/xiaochen/workspace/IMDLBenCo_pure/test_metrics/test_datasets.json \ 13 | --checkpoint_path "./output_dir/" \ 14 | --test_batch_size 3 \ 15 | --image_size 512 \ 16 | --if_resizing \ 17 | --output_dir ${base_dir}/ \ 18 | --log_dir ${base_dir}/ \ 19 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/runs/demo_test_mantra_net.sh: -------------------------------------------------------------------------------- 1 | base_dir="./eval_dir" 2 | mkdir -p ${base_dir} 3 | 4 | CUDA_VISIBLE_DEVICES=1,2,3 \ 5 | torchrun \ 6 | --standalone \ 7 | --nnodes=1 \ 8 | --nproc_per_node=3 \ 9 | ./test.py \ 10 | --model MantraNet \ 11 | --edge_mask_width 7 \ 12 | --world_size 1 \ 13 | --test_data_json "./test_datasets.json" \ 14 | --checkpoint_path "/mnt/data0/xiaochen/workspace/IMDLBenCo_pure/check_span/output_dir" \ 15 | --test_batch_size 3 \ 16 | --image_size 512 \ 17 | --if_padding \ 18 | --output_dir ${base_dir}/ \ 19 | --log_dir ${base_dir}/ \ 20 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/runs/demo_test_object_former.sh: -------------------------------------------------------------------------------- 1 | base_dir="./eval_dir" 2 | mkdir -p ${base_dir} 3 | 4 | CUDA_VISIBLE_DEVICES=1 \ 5 | torchrun \ 6 | --standalone \ 7 | --nnodes=1 \ 8 | --nproc_per_node=1 \ 9 | ./test.py \ 10 | --model ObjectFormer \ 11 | --edge_mask_width 7 \ 12 | --world_size 1 \ 13 | --test_data_json "./test_datasets.json" \ 14 | --checkpoint_path "/mnt/data0/username/workspace/IMDLBench/output_test/standard" \ 15 | --test_batch_size 3 \ 16 | --image_size 224 \ 17 | --if_resizing \ 18 | --output_dir ${base_dir}/ \ 19 | --log_dir ${base_dir}/ \ 20 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /IMDLBenCo/statics/base/test_mymodel.sh: -------------------------------------------------------------------------------- 1 | base_dir="./eval_dir" 2 | mkdir -p ${base_dir} 3 | 4 | CUDA_VISIBLE_DEVICES=0,1,2,3 \ 5 | torchrun \ 6 | --standalone \ 7 | --nnodes=1 \ 8 | --nproc_per_node=4 \ 9 | ./test.py \ 10 | --model MyModel \ 11 | --MyModel_Customized_param 12345678 \ 12 | --pre_trained_weights '' \ 13 | --world_size 1 \ 14 | --test_data_json "./test_datasets.json" \ 15 | --checkpoint_path "./output_dir/" \ 16 | --test_batch_size 2 \ 17 | --image_size 512 \ 18 | --if_resizing \ 19 | --output_dir ${base_dir}/ \ 20 | --log_dir ${base_dir}/ \ 21 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/runs/demo_test_cat_net.sh: -------------------------------------------------------------------------------- 1 | base_dir="./eval_dir" 2 | mkdir -p ${base_dir} 3 | 4 | CUDA_VISIBLE_DEVICES=1,2,3 \ 5 | torchrun \ 6 | --standalone \ 7 | --nnodes=1 \ 8 | --nproc_per_node=3 \ 9 | ./test.py \ 10 | --model Cat_Net \ 11 | --world_size 1 \ 12 | --test_data_json "./test_datasets.json" \ 13 | --checkpoint_path "/home/bingkui/IMDLBenCo/output_dir_balance" \ 14 | --test_batch_size 3 \ 15 | --if_resizing \ 16 | --image_size 512 \ 17 | --cfg_file "/home/bingkui/workspace/IMDLBenCo/configs/CAT_full.yaml" \ 18 | --output_dir ${base_dir}/ \ 19 | --log_dir ${base_dir}/ \ 20 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/runs/demo_test_mvss.sh: -------------------------------------------------------------------------------- 1 | base_dir="./eval_dir_mvss" 2 | mkdir -p ${base_dir} 3 | 4 | CUDA_VISIBLE_DEVICES=1,2,3 \ 5 | torchrun \ 6 | --standalone \ 7 | --nnodes=1 \ 8 | --nproc_per_node=3 \ 9 | ./test.py \ 10 | --model MVSSNet \ 11 | --edge_mask_width 7 \ 12 | --world_size 1 \ 13 | --test_data_json "./test_datasets.json" \ 14 | --checkpoint_path "/mnt/data0/sulei/workspace/IMDLBench/output_dir_MVSS_bachsize32_best" \ 15 | --test_batch_size 3 \ 16 | --image_size 512 \ 17 | --no_model_eval \ 18 | --if_resizing \ 19 | --output_dir ${base_dir}/ \ 20 | --log_dir ${base_dir}/ \ 21 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/runs/demo_test_robustness_mantra_net.sh: -------------------------------------------------------------------------------- 1 | base_dir="./robust_dir" 2 | mkdir -p ${base_dir} 3 | 4 | CUDA_VISIBLE_DEVICES=0,1,2,3 \ 5 | torchrun \ 6 | --standalone \ 7 | --nnodes=1 \ 8 | --nproc_per_node=4 \ 9 | ./test_robust.py \ 10 | --model MantraNet \ 11 | --edge_mask_width 7 \ 12 | --world_size 1 \ 13 | --test_data_path "/mnt/data0/public_datasets/IML/CASIA1.0" \ 14 | --checkpoint_path "/home/zeyu/workspace/IMDLBenCo/eval_dir/checkpoint-0.pth" \ 15 | --test_batch_size 2 \ 16 | --image_size 1024 \ 17 | --if_resizing \ 18 | --output_dir ${base_dir}/ \ 19 | --log_dir ${base_dir}/ \ 20 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /IMDLBenCo/training_scripts/schedular/cos_lr_schedular.py: -------------------------------------------------------------------------------- 1 | import math 2 | 3 | def adjust_learning_rate(optimizer, epoch, args): 4 | """Decay the learning rate with half-cycle cosine after warmup""" 5 | if epoch < args.warmup_epochs: 6 | lr = args.lr * epoch / args.warmup_epochs 7 | else: 8 | lr = args.min_lr + (args.lr - args.min_lr) * 0.5 * \ 9 | (1. + math.cos(math.pi * (epoch - args.warmup_epochs) / (args.epochs - args.warmup_epochs))) 10 | for param_group in optimizer.param_groups: 11 | if "lr_scale" in param_group: 12 | param_group["lr"] = lr * param_group["lr_scale"] 13 | else: 14 | param_group["lr"] = lr 15 | return lr 16 | -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/runs/demo_test_robustness_mesorch.sh: -------------------------------------------------------------------------------- 1 | base_dir="./eval_robust_dir_mesorch" 2 | mkdir -p ${base_dir} 3 | 4 | CUDA_VISIBLE_DEVICES=0,1,2,3 \ 5 | torchrun \ 6 | --standalone \ 7 | --nnodes=1 \ 8 | --nproc_per_node=4 \ 9 | ./test_robust.py \ 10 | --model Mesorch \ 11 | --world_size 1 \ 12 | --test_data_path "/mnt/data0/public_datasets/IML/CASIA1.0" \ 13 | --checkpoint_path "/mnt/data0/xuekang/workspace/Mesorch/ckpt_mesorch/mesorch-98.pth" \ 14 | --test_batch_size 2 \ 15 | --image_size 512 \ 16 | --if_resizing \ 17 | --output_dir ${base_dir}/ \ 18 | --log_dir ${base_dir}/ \ 19 | --seed 42 \ 20 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/runs/demo_test_robustness_iml_vit.sh: -------------------------------------------------------------------------------- 1 | base_dir="./robust_dir" 2 | mkdir -p ${base_dir} 3 | 4 | CUDA_VISIBLE_DEVICES=0,1,2,3 \ 5 | torchrun \ 6 | --standalone \ 7 | --nnodes=1 \ 8 | --nproc_per_node=4 \ 9 | ./test_robust.py \ 10 | --model IML_ViT \ 11 | --edge_mask_width 7 \ 12 | --world_size 1 \ 13 | --test_data_path "/mnt/data0/public_datasets/IML/CASIA1.0" \ 14 | --checkpoint_path "/mnt/data0/xiaochen/workspace/IML-VIT-shuffle/rebuttal_trufor/checkpoint-188.pth" \ 15 | --test_batch_size 2 \ 16 | --image_size 1024 \ 17 | --if_padding \ 18 | --output_dir ${base_dir}/ \ 19 | --log_dir ${base_dir}/ \ 20 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/runs/test_save_images/save_image_imlvit.sh: -------------------------------------------------------------------------------- 1 | base_dir="./save_img_dir_imlvit" 2 | mkdir -p ${base_dir} 3 | 4 | CUDA_VISIBLE_DEVICES=0,1,2,3 \ 5 | torchrun \ 6 | --standalone \ 7 | --nnodes=1 \ 8 | --nproc_per_node=4 \ 9 | ./test_save_images.py \ 10 | --model IML_ViT \ 11 | --edge_mask_width 7 \ 12 | --world_size 1 \ 13 | --test_data_path "/mnt/data0/public_datasets/IML/CASIA1.0" \ 14 | --checkpoint_path "/mnt/data0/public_datasets/IML/IMDLBenCo_ckpt/iml_vit_casiav2.pth" \ 15 | --test_batch_size 2 \ 16 | --image_size 1024 \ 17 | --if_padding \ 18 | --output_dir ${base_dir}/ \ 19 | --log_dir ${base_dir}/ \ 20 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/runs/demo_test_robustness_cat_net.sh: -------------------------------------------------------------------------------- 1 | base_dir="./robust_dir" 2 | mkdir -p ${base_dir} 3 | 4 | CUDA_VISIBLE_DEVICES=0,1,2,3 \ 5 | torchrun \ 6 | --standalone \ 7 | --nnodes=1 \ 8 | --nproc_per_node=4 \ 9 | ./test_robust.py \ 10 | --model Cat_Net \ 11 | --world_size 1 \ 12 | --test_data_path "/mnt/data0/public_datasets/IML/CASIA1.0" \ 13 | --checkpoint_path "/home/bingkui/IMDLBenCo/output_dir_balance/checkpoint-44.pth" \ 14 | --test_batch_size 2 \ 15 | --image_size 512 \ 16 | --if_resizing \ 17 | --cfg_file "/home/bingkui/workspace/IMDLBenCo/configs/CAT_full.yaml" \ 18 | --output_dir ${base_dir}/ \ 19 | --log_dir ${base_dir}/ \ 20 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/runs/demo_test_span.sh: -------------------------------------------------------------------------------- 1 | base_dir="./eval_dir" 2 | mkdir -p ${base_dir} 3 | 4 | CUDA_VISIBLE_DEVICES=1,2,3 \ 5 | torchrun \ 6 | --standalone \ 7 | --nnodes=1 \ 8 | --nproc_per_node=3 \ 9 | ./test.py \ 10 | --model SPAN \ 11 | --edge_mask_width 7 \ 12 | --world_size 1 \ 13 | --test_data_json "./test_datasets.json" \ 14 | --checkpoint_path "/mnt/data0/xiaochen/workspace/IMDLBenCo_pure/check_span/output_dir" \ 15 | --test_batch_size 3 \ 16 | --image_size 224 \ 17 | --if_resizing \ 18 | --output_dir ${base_dir}/ \ 19 | --log_dir ${base_dir}/ \ 20 | --weight_path '/mnt/data0/xiaochen/workspace/IMDLBenCo_pure/check_span/IMTFEv4.pt' \ 21 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /IMDLBenCo/statics/base/test_robust_mymodel.sh: -------------------------------------------------------------------------------- 1 | base_dir="./eval_robust_dir" 2 | mkdir -p ${base_dir} 3 | 4 | CUDA_VISIBLE_DEVICES=0,1,2,3 \ 5 | torchrun \ 6 | --standalone \ 7 | --nnodes=1 \ 8 | --nproc_per_node=4 \ 9 | ./test_robust.py \ 10 | --model MyModel \ 11 | --MyModel_Customized_param 12345678 \ 12 | --pre_trained_weights '' \ 13 | --world_size 1 \ 14 | --test_data_path "/mnt/data0/public_datasets/IML/CASIA1.0" \ 15 | --checkpoint_path "/home/xiaochen/workspace/IMDLBenCo/eval_dir/checkpoint-0.pth" \ 16 | --test_batch_size 2 \ 17 | --image_size 512 \ 18 | --if_resizing \ 19 | --output_dir ${base_dir}/ \ 20 | --log_dir ${base_dir}/ \ 21 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/runs/demo_test_robustness_mvss.sh: -------------------------------------------------------------------------------- 1 | base_dir="./robust_dir_mvss" 2 | mkdir -p ${base_dir} 3 | 4 | CUDA_VISIBLE_DEVICES=0,1,2,3 \ 5 | torchrun \ 6 | --standalone \ 7 | --nnodes=1 \ 8 | --nproc_per_node=4 \ 9 | ./test_robust.py \ 10 | --model MVSSNet \ 11 | --edge_mask_width 7 \ 12 | --world_size 1 \ 13 | --test_data_path "/mnt/data0/public_datasets/IML/CASIA1.0" \ 14 | --checkpoint_path "/mnt/data0/sulei/workspace/IMDLBench/output_dir_MVSS_bachsize32_best/checkpoint-104.pth" \ 15 | --test_batch_size 2 \ 16 | --image_size 512 \ 17 | --no_model_eval \ 18 | --if_resizing \ 19 | --output_dir ${base_dir}/ \ 20 | --log_dir ${base_dir}/ \ 21 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/runs/test_save_images/save_image_mvss.sh: -------------------------------------------------------------------------------- 1 | base_dir="./save_img_dir_mvss" 2 | mkdir -p ${base_dir} 3 | 4 | CUDA_VISIBLE_DEVICES=0,1,2,3 \ 5 | torchrun \ 6 | --standalone \ 7 | --nnodes=1 \ 8 | --nproc_per_node=4 \ 9 | ./test_save_images.py \ 10 | --model MVSSNet \ 11 | --edge_mask_width 7 \ 12 | --world_size 1 \ 13 | --test_data_path "/mnt/data0/public_datasets/IML/CASIA1.0" \ 14 | --checkpoint_path "/mnt/data0/public_datasets/IML/IMDLBenCo_ckpt/checkpoint-mvss-casiav2.pth" \ 15 | --test_batch_size 2 \ 16 | --image_size 512 \ 17 | --no_model_eval \ 18 | --if_resizing \ 19 | --output_dir ${base_dir}/ \ 20 | --log_dir ${base_dir}/ \ 21 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /IMDLBenCo/evaluation/__init__.py: -------------------------------------------------------------------------------- 1 | from .temp import cal_confusion_matrix, generate_region_mask, cal_F1 2 | from .abstract_class import AbstractEvaluator 3 | from .F1 import ImageF1, PixelF1 4 | from .AUC import ImageAUC, PixelAUC 5 | from .IOU import PixelIOU 6 | from .Accuracy import ImageAccuracy, PixelAccuracy 7 | from .gradcam.grad_camera_visualize import grad_camera_visualize 8 | 9 | __all__ = [ 10 | # Below for develop 11 | 'cal_confusion_matrix', 12 | 'generate_region_mask', 13 | 'cal_F1', 14 | # Below for real-world senario 15 | 'AbstractEvaluator', 16 | 'ImageF1', 17 | 'PixelF1', 18 | 'ImageAUC', 19 | 'PixelAUC', 20 | 'PixelIOU', 21 | 'ImageAccuracy', 22 | 'PixelAccuracy', 23 | 'grad_camera_visualize' 24 | ] -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/runs/demo_test_robustness_object_former.sh: -------------------------------------------------------------------------------- 1 | base_dir="./robust_dir" 2 | mkdir -p ${base_dir} 3 | 4 | CUDA_VISIBLE_DEVICES=0 \ 5 | torchrun \ 6 | --standalone \ 7 | --nnodes=1 \ 8 | --nproc_per_node=1 \ 9 | ./test_robust.py \ 10 | --model ObjectFormer \ 11 | --edge_mask_width 7 \ 12 | --world_size 1 \ 13 | --test_data_path "/mnt/data0/public_datasets/IML/CASIA1.0" \ 14 | --checkpoint_path "/mnt/data0/username/workspace/IMDLBench/output_test/standard/output_dir_pretrain_224_1100_12_1e4_checkpoint-1052.pth" \ 15 | --test_batch_size 2 \ 16 | --image_size 224 \ 17 | --if_resizing \ 18 | --output_dir ${base_dir}/ \ 19 | --log_dir ${base_dir}/ \ 20 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/runs/demo_test_robustness_span.sh: -------------------------------------------------------------------------------- 1 | base_dir="./robust_dir" 2 | mkdir -p ${base_dir} 3 | 4 | CUDA_VISIBLE_DEVICES=0,1,2,3 \ 5 | torchrun \ 6 | --standalone \ 7 | --nnodes=1 \ 8 | --nproc_per_node=4 \ 9 | ./test_robust.py \ 10 | --model SPAN \ 11 | --edge_mask_width 7 \ 12 | --world_size 1 \ 13 | --test_data_path "/mnt/data0/public_datasets/IML/CASIA1.0" \ 14 | --checkpoint_path "/home/zeyu/workspace/IMDLBenCo/output_dir/checkpoint-0.pth" \ 15 | --test_batch_size 2 \ 16 | --image_size 224 \ 17 | --if_resizing \ 18 | --output_dir ${base_dir}/ \ 19 | --log_dir ${base_dir}/ \ 20 | --weight_path '/home/zeyu/workspace/IMDLBenCo/IMDLBenCo/model_zoo/span/IMTFEv4.pt' \ 21 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /IMDLBenCo/statics/base/test_save_image_mymodel.sh: -------------------------------------------------------------------------------- 1 | base_dir="./save_img_dir" 2 | mkdir -p ${base_dir} 3 | 4 | CUDA_VISIBLE_DEVICES=0,1,2,3 \ 5 | torchrun \ 6 | --standalone \ 7 | --nnodes=1 \ 8 | --nproc_per_node=4 \ 9 | ./test_save_images.py \ 10 | --model MyModel \ 11 | --MyModel_Customized_param 12345678 \ 12 | --pre_trained_weights '' \ 13 | --world_size 1 \ 14 | --test_data_path "/mnt/data0/public_datasets/IML/CASIA1.0" \ 15 | --checkpoint_path "/home/xiaochen/workspace/IMDLBenCo/eval_dir/checkpoint-0.pth" \ 16 | --test_batch_size 2 \ 17 | --image_size 512 \ 18 | --no_model_eval \ 19 | --if_resizing \ 20 | --output_dir ${base_dir}/ \ 21 | --log_dir ${base_dir}/ \ 22 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/runs/test_save_images/save_image_catnet.sh: -------------------------------------------------------------------------------- 1 | base_dir="./save_img_dir_catnet" 2 | mkdir -p ${base_dir} 3 | 4 | CUDA_VISIBLE_DEVICES=0,1,2,3 \ 5 | torchrun \ 6 | --standalone \ 7 | --nnodes=1 \ 8 | --nproc_per_node=4 \ 9 | ./test_save_images.py \ 10 | --model Cat_Net \ 11 | --cfg_file /mnt/data0/xiaochen/workspace/IMDLBenCo_pure/test_cat/configs/CAT_full.yaml \ 12 | --world_size 1 \ 13 | --test_data_path "/mnt/data0/public_datasets/IML/CASIA1.0" \ 14 | --checkpoint_path "/mnt/data0/public_datasets/IML/IMDLBenCo_ckpt/cat_net_cat_net.pth" \ 15 | --test_batch_size 2 \ 16 | --image_size 512 \ 17 | --if_resizing \ 18 | --output_dir ${base_dir}/ \ 19 | --log_dir ${base_dir}/ \ 20 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/runs/test_save_images/save_image_pscc.sh: -------------------------------------------------------------------------------- 1 | base_dir="./save_img/pscc" 2 | mkdir -p ${base_dir} 3 | 4 | CUDA_VISIBLE_DEVICES=0 \ 5 | torchrun \ 6 | --standalone \ 7 | --nnodes=1 \ 8 | --nproc_per_node=1 \ 9 | ./save_images.py \ 10 | --model PSCC_Net \ 11 | --edge_mask_width 7 \ 12 | --world_size 1 \ 13 | --test_data_path "/mnt/data0/public_datasets/IML/CASIA1.0" \ 14 | --pretrain_path "/mnt/data0/IMDLBenCo-main/new_workspace/weights/PSCC/hrnet_w18_small_v2.pth" \ 15 | --checkpoint_path "/mnt/data0/public_datasets/IML/IMDLBenCo_ckpt/pscc_catnet.pth" \ 16 | --test_batch_size 2 \ 17 | --image_size 256 \ 18 | --if_resizing \ 19 | --output_dir ${base_dir}/ \ 20 | --log_dir ${base_dir}/ \ 21 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/runs/demo_test_pscc.sh: -------------------------------------------------------------------------------- 1 | base_dir="./log/test_pscc" 2 | mkdir -p ${base_dir} 3 | 4 | CUDA_VISIBLE_DEVICES=1,2,3,4,5 \ 5 | torchrun \ 6 | --standalone \ 7 | --nnodes=1 \ 8 | --nproc_per_node=5 \ 9 | ./IMDLBenCo/training_scripts/test.py \ 10 | --model PSCC_Net \ 11 | --edge_mask_width 7 \ 12 | --pretrain_path "/mnt/data0/dubo/workspace/IMDLBenCo/IMDLBenCo/model_zoo/pscc/hrnet_w18_small_v2.pth" \ 13 | --world_size 1 \ 14 | --test_data_json "./runs/test_datasets.json" \ 15 | --checkpoint_path "/mnt/data0/dubo/workspace/IMDLBenCo/log/train_casiav2full_pscc" \ 16 | --test_batch_size 8 \ 17 | --image_size 256 \ 18 | --if_resizing \ 19 | --output_dir ${base_dir}/ \ 20 | --log_dir ${base_dir}/ \ 21 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/3-feature-request.yaml: -------------------------------------------------------------------------------- 1 | name: "✨ Feature Request" 2 | description: Suggest an idea for new features. 3 | labels: enhancement 4 | 5 | body: 6 | - type: textarea 7 | id: system-info 8 | attributes: 9 | label: System Info 10 | description: | 11 | Please share your system info with us. You can run the command **benco env** and copy-paste its output below. 12 | 请提供您的系统信息。您可以在命令行运行 **benco env** 并将其输出复制到该文本框中。 13 | placeholder: benco version, platform, python version, ... 14 | validations: 15 | required: true 16 | 17 | - type: textarea 18 | id: others 19 | attributes: 20 | label: Others 21 | description: | 22 | Describe feature you want detailly here. 23 | 详细描述你期待的新功能。 24 | validations: 25 | required: false 26 | -------------------------------------------------------------------------------- /IMDLBenCo/model_zoo/__init__.py: -------------------------------------------------------------------------------- 1 | from .iml_vit.iml_vit import IML_ViT 2 | from .cat_net.cat_net import Cat_Net 3 | from .cat_net.cat_net_post_function import cat_net_post_func 4 | from .mantra_net.mantranet import MantraNet 5 | from .mvss_net.mvssnet import MVSSNet 6 | from .object_former.object_former import ObjectFormer 7 | from .pscc_net.pscc_net import PSCC_Net 8 | from .span.SPAN import SPAN 9 | from .trufor.trufor import Trufor 10 | from .mesorch.mesorch import Mesorch 11 | from .sparse_vit.sparse_vit import SparseViT, SparseViTBackbone 12 | 13 | __all__ = [ 14 | 'IML_ViT', 15 | "Cat_Net", 16 | 'cat_net_post_func', 17 | "MantraNet", 18 | "MVSSNet", 19 | "ObjectFormer", 20 | "PSCC_Net", 21 | "SPAN", 22 | "Trufor", 23 | 'Mesorch', 24 | "SparseViT", 25 | "SparseViTBackbone" 26 | ] 27 | -------------------------------------------------------------------------------- /tests/test_subdataset.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.utils.data import DataLoader, Dataset 3 | 4 | # 创建一个简单的自定义Dataset 5 | class MyDataset(Dataset): 6 | def __init__(self): 7 | self.data = [1, 2, 3, 4, 5] 8 | 9 | def __len__(self): 10 | return len(self.data) 11 | 12 | def __getitem__(self, idx): 13 | return self.data[idx] 14 | 15 | # 实例化Dataset 16 | dataset = MyDataset() 17 | 18 | # 用DataLoader加载Dataset 19 | dataloader = DataLoader(dataset, batch_size=2, shuffle=True) 20 | 21 | print(len(dataloader.dataset)) 22 | # 获取DataLoader的Dataset引用 23 | dataset_from_dataloader = dataloader.dataset 24 | 25 | # 验证引用是否相同 26 | print(f"Original Dataset ID: {id(dataset)}") 27 | print(f"Dataset from DataLoader ID: {id(dataset_from_dataloader)}") 28 | print(f"Is same object: {dataset is dataset_from_dataloader}") -------------------------------------------------------------------------------- /IMDLBenCo/version.py: -------------------------------------------------------------------------------- 1 | __version__ = '0.1.42' 2 | short_version = __version__ 3 | 4 | def parse_version_info(version_str): 5 | """Parse a version string into a tuple. 6 | 7 | Args: 8 | version_str (str): The version string. 9 | Returns: 10 | tuple[int | str]: The version info, e.g., "1.3.0" is parsed into 11 | (1, 3, 0), and "2.0.0rc1" is parsed into (2, 0, 0, 'rc1'). 12 | """ 13 | version_info = [] 14 | for x in version_str.split('.'): 15 | if x.isdigit(): 16 | version_info.append(int(x)) 17 | elif x.find('rc') != -1: 18 | patch_version = x.split('rc') 19 | version_info.append(int(patch_version[0])) 20 | version_info.append(f'rc{patch_version[1]}') 21 | return tuple(version_info) 22 | 23 | version_info = parse_version_info(__version__) 24 | -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/runs/demo_train_backbone.sh: -------------------------------------------------------------------------------- 1 | base_dir="./output_dir" 2 | mkdir -p ${base_dir} 3 | 4 | CUDA_VISIBLE_DEVICES=3 \ 5 | torchrun \ 6 | --standalone \ 7 | --nnodes=1 \ 8 | --nproc_per_node=1 \ 9 | ./train.py \ 10 | --model bayar_resnet \ 11 | --world_size 1 \ 12 | --batch_size 1 \ 13 | --data_path /mnt/data0/public_datasets/IML/CASIA2.0 \ 14 | --epochs 200 \ 15 | --lr 1e-4 \ 16 | --image_size 1024 \ 17 | --if_resizing \ 18 | --min_lr 5e-7 \ 19 | --weight_decay 0.05 \ 20 | --edge_mask_width 7 \ 21 | --test_data_path "/mnt/data0/public_datasets/IML/CASIA1.0" \ 22 | --warmup_epochs 2 \ 23 | --output_dir ${base_dir}/ \ 24 | --log_dir ${base_dir}/ \ 25 | --accum_iter 8 \ 26 | --seed 42 \ 27 | --test_period 4 \ 28 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/runs/demo_train_mantra_net.sh: -------------------------------------------------------------------------------- 1 | base_dir="./output_dir" 2 | mkdir -p ${base_dir} 3 | 4 | CUDA_VISIBLE_DEVICES=4,5 \ 5 | torchrun \ 6 | --standalone \ 7 | --nnodes=1 \ 8 | --nproc_per_node=2 \ 9 | ./train.py \ 10 | --model MantraNet \ 11 | --world_size 1 \ 12 | --batch_size 1 \ 13 | --data_path /mnt/data0/public_datasets/IML/CASIA2.0 \ 14 | --epochs 200 \ 15 | --lr 1e-4 \ 16 | --image_size 512 \ 17 | --if_resizing \ 18 | --min_lr 5e-7 \ 19 | --weight_decay 0.05 \ 20 | --edge_mask_width 7 \ 21 | --test_data_path "/mnt/data0/public_datasets/IML/CASIA1.0" \ 22 | --warmup_epochs 2 \ 23 | --output_dir ${base_dir}/ \ 24 | --log_dir ${base_dir}/ \ 25 | --accum_iter 8 \ 26 | --seed 42 \ 27 | --test_period 4 \ 28 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/runs/demo_test_robustness_pscc.sh: -------------------------------------------------------------------------------- 1 | base_dir="./log/robust_pscc" 2 | mkdir -p ${base_dir} 3 | 4 | CUDA_VISIBLE_DEVICES=1,2,3,4,5 \ 5 | torchrun \ 6 | --standalone \ 7 | --nnodes=1 \ 8 | --nproc_per_node=5 \ 9 | ./IMDLBenCo/training_scripts/test_robust.py \ 10 | --model PSCC_Net \ 11 | --edge_mask_width 7 \ 12 | --pretrain_path "/mnt/data0/dubo/workspace/IMDLBenCo/IMDLBenCo/model_zoo/pscc/hrnet_w18_small_v2.pth" \ 13 | --world_size 1 \ 14 | --test_data_path "/mnt/data0/public_datasets/IML/CASIA1.0" \ 15 | --checkpoint_path "/mnt/data0/dubo/workspace/IMDLBenCo/log/train_casiav2full_pscc/checkpoint.pth" \ 16 | --test_batch_size 8 \ 17 | --image_size 256 \ 18 | --if_resizing \ 19 | --output_dir ${base_dir}/ \ 20 | --log_dir ${base_dir}/ \ 21 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /IMDLBenCo/statics/dataset_json/balanced_dataset.json: -------------------------------------------------------------------------------- 1 | [ 2 | [ 3 | "ManiDataset", 4 | "/mnt/data0/public_datasets/IML/CASIA2.0" 5 | ], 6 | [ 7 | "JsonDataset", 8 | "/mnt/data0/public_datasets/IML/FantasticReality_v1/FantasticReality.json" 9 | ], 10 | [ 11 | "ManiDataset", 12 | "/mnt/data0/public_datasets/IML/IMD_20_1024" 13 | ], 14 | [ 15 | "JsonDataset", 16 | "/mnt/data0/public_datasets/IML/tampCOCO/sp_COCO_list.json" 17 | ], 18 | [ 19 | "JsonDataset", 20 | "/mnt/data0/public_datasets/IML/tampCOCO/cm_COCO_list.json" 21 | ], 22 | [ 23 | "JsonDataset", 24 | "/mnt/data0/public_datasets/IML/tampCOCO/bcm_COCO_list.json" 25 | ], 26 | [ 27 | "JsonDataset", 28 | "/mnt/data0/public_datasets/IML/tampCOCO/bcmc_COCO_list.json" 29 | ] 30 | ] 31 | -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/runs/demo_train_cat_net.sh: -------------------------------------------------------------------------------- 1 | base_dir="./output_dir" 2 | mkdir -p ${base_dir} 3 | 4 | CUDA_VISIBLE_DEVICES=0,1,2,3 \ 5 | torchrun \ 6 | --standalone \ 7 | --nnodes=1 \ 8 | --nproc_per_node=4 \ 9 | ./train.py \ 10 | --model Cat_Net \ 11 | --world_size 1 \ 12 | --batch_size 32 \ 13 | --data_path "/mnt/data0/public_datasets/IML/CASIA2.0" \ 14 | --epochs 200 \ 15 | --lr 1e-4 \ 16 | --image_size 512 \ 17 | --if_resizing \ 18 | --min_lr 5e-7 \ 19 | --weight_decay 0.05 \ 20 | --edge_mask_width 7 \ 21 | --test_data_path "/mnt/data0/public_datasets/IML/CASIA1.0" \ 22 | --warmup_epochs 2 \ 23 | --output_dir ${base_dir}/ \ 24 | --log_dir ${base_dir}/ \ 25 | --accum_iter 8 \ 26 | --seed 42 \ 27 | --test_period 4 \ 28 | --cfg_file "./configs/CAT_full.yaml" \ 29 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/runs/demo_catnet_protocol_mvss.sh: -------------------------------------------------------------------------------- 1 | base_dir="./output_dir_mvss" 2 | mkdir -p ${base_dir} 3 | 4 | CUDA_VISIBLE_DEVICES=0 \ 5 | torchrun \ 6 | --standalone \ 7 | --nnodes=1 \ 8 | --nproc_per_node=1 \ 9 | ./train.py \ 10 | --model MVSSNet \ 11 | --world_size 1 \ 12 | --batch_size 1 \ 13 | --data_path ./balanced_dataset.json \ 14 | --epochs 200 \ 15 | --lr 2e-5 \ 16 | --image_size 512 \ 17 | --if_resizing \ 18 | --if_not_amp \ 19 | --find_unused_parameters \ 20 | --no_model_eval \ 21 | --min_lr 5e-7 \ 22 | --weight_decay 0.05 \ 23 | --edge_mask_width 7 \ 24 | --test_data_path "/mnt/data0/public_datasets/IML/CASIA1.0" \ 25 | --warmup_epochs 2 \ 26 | --output_dir ${base_dir}/ \ 27 | --log_dir ${base_dir}/ \ 28 | --accum_iter 8 \ 29 | --seed 42 \ 30 | --test_period 4 \ 31 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/runs/demo_train_object_former.sh: -------------------------------------------------------------------------------- 1 | base_dir="./output_dir" 2 | mkdir -p ${base_dir} 3 | 4 | CUDA_VISIBLE_DEVICES=3 \ 5 | torchrun \ 6 | --standalone \ 7 | --nnodes=1 \ 8 | --nproc_per_node=1 \ 9 | ./train.py \ 10 | --init_weight_path object_former/processed_model_weights.pth \ 11 | --model ObjectFormer \ 12 | --world_size 1 \ 13 | --batch_size 12 \ 14 | --data_path /mnt/data0/public_datasets/IML/CASIA2.0 \ 15 | --epochs 1100 \ 16 | --lr 1e-4 \ 17 | --image_size 224 \ 18 | --if_resizing \ 19 | --min_lr 5e-7 \ 20 | --weight_decay 0.05 \ 21 | --edge_mask_width 7 \ 22 | --test_data_path "/mnt/data0/public_datasets/IML/CASIA1.0" \ 23 | --warmup_epochs 2 \ 24 | --output_dir ${base_dir}/ \ 25 | --log_dir ${base_dir}/ \ 26 | --accum_iter 8 \ 27 | --seed 42 \ 28 | --test_period 4 \ 29 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /IMDLBenCo/statics/base/train_mymodel.sh: -------------------------------------------------------------------------------- 1 | base_dir="./output_dir" 2 | mkdir -p ${base_dir} 3 | 4 | CUDA_VISIBLE_DEVICES=3 \ 5 | torchrun \ 6 | --standalone \ 7 | --nnodes=1 \ 8 | --nproc_per_node=1 \ 9 | ./train.py \ 10 | --model MyModel \ 11 | --MyModel_Customized_param 12345678 \ 12 | --pre_trained_weights '' \ 13 | --world_size 1 \ 14 | --batch_size 1 \ 15 | --data_path /mnt/data0/public_datasets/IML/CASIA2.0 \ 16 | --epochs 200 \ 17 | --lr 1e-4 \ 18 | --image_size 512 \ 19 | --if_resizing \ 20 | --min_lr 5e-7 \ 21 | --weight_decay 0.05 \ 22 | --edge_mask_width 7 \ 23 | --test_data_path "/mnt/data0/public_datasets/IML/CASIA1.0" \ 24 | --warmup_epochs 2 \ 25 | --output_dir ${base_dir}/ \ 26 | --log_dir ${base_dir}/ \ 27 | --accum_iter 8 \ 28 | --seed 42 \ 29 | --test_period 4 \ 30 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/runs/demo_train_mvss.sh: -------------------------------------------------------------------------------- 1 | base_dir="./output_dir_mvss" 2 | mkdir -p ${base_dir} 3 | 4 | CUDA_VISIBLE_DEVICES=4 \ 5 | torchrun \ 6 | --standalone \ 7 | --nnodes=1 \ 8 | --nproc_per_node=1 \ 9 | ./train.py \ 10 | --model MVSSNet \ 11 | --world_size 1 \ 12 | --batch_size 1 \ 13 | --data_path /mnt/data0/public_datasets/IML/CASIA2.0 \ 14 | --epochs 200 \ 15 | --lr 2e-5 \ 16 | --image_size 512 \ 17 | --if_not_amp \ 18 | --find_unused_parameters \ 19 | --no_model_eval \ 20 | --if_resizing \ 21 | --min_lr 5e-7 \ 22 | --weight_decay 0.05 \ 23 | --edge_mask_width 7 \ 24 | --test_data_path "/mnt/data0/public_datasets/IML/CASIA1.0" \ 25 | --warmup_epochs 2 \ 26 | --output_dir ${base_dir}/ \ 27 | --log_dir ${base_dir}/ \ 28 | --accum_iter 8 \ 29 | --seed 42 \ 30 | --test_period 4 \ 31 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/runs/demo_train_span.sh: -------------------------------------------------------------------------------- 1 | base_dir="./output_dir" 2 | mkdir -p ${base_dir} 3 | 4 | CUDA_VISIBLE_DEVICES=0,1,2,3 \ 5 | torchrun \ 6 | --standalone \ 7 | --nnodes=1 \ 8 | --nproc_per_node=4 \ 9 | ./train.py \ 10 | --model SPAN \ 11 | --world_size 1 \ 12 | --batch_size 8 \ 13 | --data_path /mnt/data0/public_datasets/IML/CASIA2.0 \ 14 | --epochs 200 \ 15 | --lr 1e-4 \ 16 | --image_size 224 \ 17 | --if_resizing \ 18 | --min_lr 5e-7 \ 19 | --weight_decay 0.05 \ 20 | --edge_mask_width 7 \ 21 | --test_data_path "/mnt/data0/public_datasets/IML/CASIA1.0" \ 22 | --warmup_epochs 2 \ 23 | --output_dir ${base_dir}/ \ 24 | --log_dir ${base_dir}/ \ 25 | --accum_iter 8 \ 26 | --seed 42 \ 27 | --test_period 4 \ 28 | --weight_path '/mnt/data0/xiaochen/workspace/IMDLBenCo_pure/check_span/IMTFEv4.pt' \ 29 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/runs/demo_train_sparse_vit.sh: -------------------------------------------------------------------------------- 1 | base_dir="./output_dir" 2 | mkdir -p ${base_dir} 3 | 4 | torchrun \ 5 | --standalone \ 6 | --nnodes=1 \ 7 | --nproc_per_node=4 \ 8 | train.py \ 9 | --model SparseViT \ 10 | --world_size 1 \ 11 | --batch_size 16 \ 12 | --data_path /mnt/data0/public_datasets/IML/CASIA2.0 \ 13 | --epochs 200 \ 14 | --lr 2e-4 \ 15 | --min_lr 0 \ 16 | --weight_decay 0.05 \ 17 | --pretrained_path "/mnt/data0/xiaochen/workspace/IMDLBenCo_pure/test_sparse_vit/uniformer_image/uniformer_base_in1k.pth" \ 18 | --test_data_path "/mnt/data0/public_datasets/IML/CASIA1.0" \ 19 | --if_resizing \ 20 | --find_unused_parameters \ 21 | --warmup_epochs 4 \ 22 | --output_dir ${base_dir}/ \ 23 | --log_dir ${base_dir}/ \ 24 | --accum_iter 1 \ 25 | --seed 42 \ 26 | --test_period 4 \ 27 | --num_workers 12 \ 28 | 2> ${base_dir}/train_error.log 1>${base_dir}/train_log.log 29 | -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/runs/demo_train_backbone_segformer.sh: -------------------------------------------------------------------------------- 1 | base_dir="./output_dir" 2 | mkdir -p ${base_dir} 3 | 4 | CUDA_VISIBLE_DEVICES=3 \ 5 | torchrun \ 6 | --standalone \ 7 | --nnodes=1 \ 8 | --nproc_per_node=1 \ 9 | ./train.py \ 10 | --model segformer \ 11 | --pretrain_pth_path /mnt/data0/user/workspace/IMDLBench/IMDLBench/modules/baselines/mit_b2.pth \ 12 | --world_size 1 \ 13 | --batch_size 1 \ 14 | --data_path /mnt/data0/public_datasets/IML/CASIA2.0 \ 15 | --epochs 200 \ 16 | --lr 1e-4 \ 17 | --image_size 1024 \ 18 | --if_resizing \ 19 | --min_lr 5e-7 \ 20 | --weight_decay 0.05 \ 21 | --edge_mask_width 7 \ 22 | --test_data_path "/mnt/data0/public_datasets/IML/CASIA1.0" \ 23 | --warmup_epochs 2 \ 24 | --output_dir ${base_dir}/ \ 25 | --log_dir ${base_dir}/ \ 26 | --accum_iter 8 \ 27 | --seed 42 \ 28 | --test_period 4 \ 29 | # 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/runs/demo_catnet_protocol_iml_vit.sh: -------------------------------------------------------------------------------- 1 | base_dir="./output_dir" 2 | mkdir -p ${base_dir} 3 | 4 | CUDA_VISIBLE_DEVICES=4 \ 5 | torchrun \ 6 | --standalone \ 7 | --nnodes=1 \ 8 | --nproc_per_node=1 \ 9 | ./train.py \ 10 | --model IML-ViT \ 11 | --edge_lambda 20 \ 12 | --vit_pretrain_path /mnt/data0/xiaochen/workspace/IML-ViT/pretrained-weights/mae_pretrain_vit_base.pth \ 13 | --world_size 1 \ 14 | --batch_size 1 \ 15 | --data_path ./balanced_dataset.json \ 16 | --epochs 200 \ 17 | --lr 1e-4 \ 18 | --image_size 1024 \ 19 | --if_resizing \ 20 | --min_lr 5e-7 \ 21 | --weight_decay 0.05 \ 22 | --edge_mask_width 7 \ 23 | --test_data_path "/mnt/data0/public_datasets/IML/CASIA1.0" \ 24 | --warmup_epochs 2 \ 25 | --output_dir ${base_dir}/ \ 26 | --log_dir ${base_dir}/ \ 27 | --accum_iter 8 \ 28 | --seed 42 \ 29 | --test_period 4 \ 30 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/runs/demo_test_trufor.sh: -------------------------------------------------------------------------------- 1 | base_dir="./log/test_trufor" 2 | mkdir -p ${base_dir} 3 | 4 | CUDA_VISIBLE_DEVICES=1,2,3,4,5 \ 5 | torchrun \ 6 | --standalone \ 7 | --nnodes=1 \ 8 | --nproc_per_node=5 \ 9 | ./IMDLBenCo/training_scripts/test.py \ 10 | --model Trufor \ 11 | --edge_mask_width 7 \ 12 | --np_pretrain_weights "/mnt/data0/dubo/workspace/IMDLBenCo/IMDLBenCo/model_zoo/trufor/noiseprint.pth" \ 13 | --mit_b2_pretrain_weights "/mnt/data0/dubo/workspace/IMDLBenCo/IMDLBenCo/model_zoo/trufor/mit_b2.pth" \ 14 | --config_path "/mnt/data0/dubo/workspace/IMDLBenCo/configs/trufor.yaml" \ 15 | --world_size 1 \ 16 | --test_data_json "./runs/test_datasets.json" \ 17 | --checkpoint_path "/mnt/data0/dubo/workspace/IMDLBenCo/log/train_casiav2full_trufor" \ 18 | --test_batch_size 8 \ 19 | --image_size 512 \ 20 | --if_resizing \ 21 | --output_dir ${base_dir}/ \ 22 | --log_dir ${base_dir}/ \ 23 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /IMDLBenCo/model_zoo/trufor/config.py: -------------------------------------------------------------------------------- 1 | from yacs.config import CfgNode as CN 2 | 3 | _C = CN() 4 | 5 | _C.OUTPUT_DIR = '' 6 | _C.LOG_DIR = '' 7 | _C.GPUS = (0,) 8 | _C.WORKERS = 4 9 | 10 | # Cudnn parameters 11 | _C.CUDNN = CN() 12 | _C.CUDNN.BENCHMARK = True 13 | _C.CUDNN.DETERMINISTIC = False 14 | _C.CUDNN.ENABLED = True 15 | 16 | # Model parameters 17 | _C.MODEL = CN() 18 | _C.MODEL.NAME = 'detconfcmx' 19 | _C.MODEL.PRETRAINED = '' 20 | _C.MODEL.MODS = ('RGB','NP++') 21 | _C.MODEL.EXTRA = CN(new_allowed=True) 22 | _C.MODEL.EXTRA.DETECTION = None 23 | _C.MODEL.EXTRA.CONF = False 24 | 25 | # Dataset parameters 26 | _C.DATASET = CN() 27 | _C.DATASET.ROOT = '' 28 | _C.DATASET.TRAIN = [] 29 | _C.DATASET.VALID = [] 30 | _C.DATASET.NUM_CLASSES = 2 31 | 32 | # Testing parameters 33 | _C.TEST = CN() 34 | _C.TEST.MODEL_FILE = '' 35 | 36 | 37 | def update_config(cfg, args, config_path): 38 | cfg.defrost() 39 | cfg.merge_from_file(config_path) 40 | # if args.opts: 41 | # cfg.merge_from_list(args.opts) 42 | cfg.freeze() -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/runs/demo_train_mesorch.sh: -------------------------------------------------------------------------------- 1 | base_dir="./output_dir_mesorch" 2 | mkdir -p ${base_dir} 3 | 4 | CUDA_VISIBLE_DEVICES=0,1,2,3 \ 5 | torchrun \ 6 | --standalone \ 7 | --nnodes=1 \ 8 | --nproc_per_node=4 \ 9 | ./train.py \ 10 | --model Mesorch \ 11 | --conv_pretrain True \ 12 | --seg_pretrain_path "/mnt/data0/xuekang/workspace/segformer/mit_b3.pth" \ 13 | --world_size 4 \ 14 | --find_unused_parameters \ 15 | --batch_size 12 \ 16 | --data_path /mnt/data0/xuekang/workspace/Mesorch/balanced_dataset.json \ 17 | --epochs 150 \ 18 | --lr 1e-4 \ 19 | --image_size 512 \ 20 | --if_resizing \ 21 | --min_lr 5e-7 \ 22 | --weight_decay 0.05 \ 23 | --test_data_path "/mnt/data0/public_datasets/IML/CASIA1.0" \ 24 | --warmup_epochs 2 \ 25 | --output_dir ${base_dir}/ \ 26 | --log_dir ${base_dir}/ \ 27 | --accum_iter 2 \ 28 | --seed 42 \ 29 | --test_period 2 \ 30 | --num_workers 12 \ 31 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/runs/demo_test_robustness_trufor.sh: -------------------------------------------------------------------------------- 1 | base_dir="./log/robust_trufor" 2 | mkdir -p ${base_dir} 3 | 4 | CUDA_VISIBLE_DEVICES=1,2,3,4,5 \ 5 | torchrun \ 6 | --standalone \ 7 | --nnodes=1 \ 8 | --nproc_per_node=5 \ 9 | ./IMDLBenCo/training_scripts/test_robust.py \ 10 | --model Trufor \ 11 | --edge_mask_width 7 \ 12 | --np_pretrain_weights "/mnt/data0/dubo/workspace/IMDLBenCo/IMDLBenCo/model_zoo/trufor/noiseprint.pth" \ 13 | --mit_b2_pretrain_weights "/mnt/data0/dubo/workspace/IMDLBenCo/IMDLBenCo/model_zoo/trufor/mit_b2.pth" \ 14 | --config_path "/mnt/data0/dubo/workspace/IMDLBenCo/configs/trufor.yaml" \ 15 | --world_size 1 \ 16 | --test_data_path "/mnt/data0/public_datasets/IML/CASIA1.0" \ 17 | --checkpoint_path "/mnt/data0/dubo/workspace/IMDLBenCo/log/train_casiav2full_trufor/checkpoint.pth" \ 18 | --test_batch_size 8 \ 19 | --image_size 512 \ 20 | --if_resizing \ 21 | --output_dir ${base_dir}/ \ 22 | --log_dir ${base_dir}/ \ 23 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /IMDLBenCo/evaluation/gradcam/grad_cam_hack.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from .base_cam_hack import BaseCAM 4 | 5 | 6 | class GradCAM(BaseCAM): 7 | def __init__(self, model, target_layers, 8 | reshape_transform=None): 9 | super( 10 | GradCAM, 11 | self).__init__( 12 | model, 13 | target_layers, 14 | reshape_transform) 15 | 16 | def get_cam_weights(self, 17 | input_tensor, 18 | target_layer, 19 | target_category, 20 | activations, 21 | grads): 22 | # 2D image 23 | if len(grads.shape) == 4: 24 | return np.mean(grads, axis=(2, 3)) 25 | 26 | # 3D image 27 | elif len(grads.shape) == 5: 28 | return np.mean(grads, axis=(2, 3, 4)) 29 | 30 | else: 31 | raise ValueError("Invalid grads shape." 32 | "Shape of grads should be 4 (2D image) or 5 (3D image).") -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/runs/demo_train_iml_vit.sh: -------------------------------------------------------------------------------- 1 | base_dir="./output_dir" 2 | mkdir -p ${base_dir} 3 | # For MAE Pretrained weights please reference here: 4 | # https://dl.fbaipublicfiles.com/mae/pretrain/mae_pretrain_vit_base.pth 5 | CUDA_VISIBLE_DEVICES=0 \ 6 | torchrun \ 7 | --standalone \ 8 | --nnodes=1 \ 9 | --nproc_per_node=1 \ 10 | ./train.py \ 11 | --model IML_ViT \ 12 | --edge_lambda 20 \ 13 | --vit_pretrain_path /mnt/data0/xiaochen/workspace/IML-ViT/pretrained-weights/mae_pretrain_vit_base.pth \ 14 | --world_size 1 \ 15 | --batch_size 2 \ 16 | --data_path /mnt/data0/public_datasets/IML/CASIA2.0 \ 17 | --epochs 200 \ 18 | --lr 1e-4 \ 19 | --image_size 1024 \ 20 | --if_padding \ 21 | --min_lr 5e-7 \ 22 | --weight_decay 0.05 \ 23 | --edge_mask_width 7 \ 24 | --test_data_path "/mnt/data0/public_datasets/IML/CASIA1.0" \ 25 | --warmup_epochs 2 \ 26 | --output_dir ${base_dir}/ \ 27 | --log_dir ${base_dir}/ \ 28 | --accum_iter 8 \ 29 | --seed 42 \ 30 | --test_period 4 \ 31 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/runs/demo_train_pscc.sh: -------------------------------------------------------------------------------- 1 | base_dir="./output_dir_pscc" 2 | mkdir -p ${base_dir} 3 | # Check point to hrnet_w18_small_v2.pth please reference here: 4 | # https://github.com/proteus1991/PSCC-Net/blob/main/models/hrnet_w18_small_v2.pth 5 | 6 | CUDA_VISIBLE_DEVICES=4 \ 7 | torchrun \ 8 | --standalone \ 9 | --nnodes=1 \ 10 | --nproc_per_node=1 \ 11 | ./train.py \ 12 | --model PSCC_Net \ 13 | --world_size 1 \ 14 | --batch_size 16 \ 15 | --data_path /mnt/data0/public_datasets/IML/CASIA2.0 \ 16 | --pretrain_path "/mnt/data0/xiaochen/workspace/IMDLBenCo_pure/test_cat_net/hrnet_w18_small_v2.pth" \ 17 | --epochs 150 \ 18 | --lr 1e-4 \ 19 | --image_size 256 \ 20 | --if_resizing \ 21 | --min_lr 0 \ 22 | --weight_decay 0.05 \ 23 | --edge_mask_width 7 \ 24 | --if_predict_label \ 25 | --if_not_amp \ 26 | --test_data_path "/mnt/data0/public_datasets/IML/CASIA1.0" \ 27 | --warmup_epochs 2 \ 28 | --output_dir ${base_dir}/ \ 29 | --log_dir ${base_dir}/ \ 30 | --accum_iter 1 \ 31 | --seed 42 \ 32 | --test_period 4 \ 33 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /tests/test_model_auto_parser.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import inspect 3 | import sys 4 | sys.path.append("./") 5 | from IMDLBenCo import MODELS 6 | 7 | def create_argparser(model_class): 8 | parser = argparse.ArgumentParser(description=f"Arguments for {model_class.__name__}") 9 | 10 | # 获取模型的__init__方法的签名 11 | sig = inspect.signature(model_class.__init__) 12 | 13 | # 解析每个参数并添加到argparse 14 | for name, param in sig.parameters.items(): 15 | if name == 'self': 16 | continue 17 | arg_type = param.annotation if param.annotation != inspect.Parameter.empty else str 18 | default_value = param.default if param.default != inspect.Parameter.empty else None 19 | print(name, arg_type, default_value) 20 | 21 | if default_value is not None: 22 | parser.add_argument(f'--{name}', type=arg_type, default=default_value, help=f'{name} (default: {default_value})') 23 | else: 24 | parser.add_argument(f'--{name}', type=arg_type, required=True, help=f'{name} (required)') 25 | 26 | return parser 27 | 28 | 29 | iml_vit = MODELS.get("IML_ViT") 30 | create_argparser(iml_vit) -------------------------------------------------------------------------------- /IMDLBenCo/model_zoo/trufor/cmx/LICENSE_CMX.txt: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2022 Huayao Liu 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. -------------------------------------------------------------------------------- /IMDLBenCo/statics/base/README-IMDLBenCo.md: -------------------------------------------------------------------------------- 1 | # Brief Introduction for "benco init base" 2 | 3 | 4 | `benco init base` and `benco init` behave similarly. They both generate necessary training, testing, and robustness testing scripts in the current working directory. 5 | 6 | They also provide a `MyModel.py` file to demonstrate a typical development pattern. 7 | 8 | Our design pattern encourages you to appropriately **modify** the training script to meet your **personalized needs**, as long as you understand the functionality of the code. Feel free to modify the Python and shell scripts generated by `benco init`! 9 | 10 | ## Quick Start 11 | You can simply start training a single-layer convolutional model with the following command. 12 | 13 | >Currently, you may need to properly configure the dataset according to the README on the [homepage](https://github.com/scu-zjz/IMDLBenCo) before it can run correctly. 14 | 15 | ```shell 16 | sh train_mymodel.sh 17 | ``` 18 | 19 | 20 | Once the model is saved, you can test its performance and robustness with the following command. 21 | 22 | ```shell 23 | sh test_mymodel.sh 24 | ``` 25 | 26 | ```shell 27 | sh test_robust_mymodel.sh 28 | ``` -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/2-ask-for-help.yaml: -------------------------------------------------------------------------------- 1 | name: "🤔 ask for help" 2 | description: Ask for help to others here, like model performance, something won't work. 3 | labels: question 4 | 5 | body: 6 | - type: textarea 7 | id: system-info 8 | attributes: 9 | label: System Info 10 | description: | 11 | Please share your system info with us. You can run the command **benco env** and copy-paste its output below. 12 | 请提供您的系统信息。您可以在命令行运行 **benco env** 并将其输出复制到该文本框中。 13 | placeholder: benco version, platform, python version, ... 14 | validations: 15 | required: true 16 | 17 | - type: textarea 18 | id: reproduction 19 | attributes: 20 | label: Reproduction 21 | description: | 22 | Please provide entry shell script. 23 | 请提供入口的shell脚本。 24 | value: | 25 | ```text 26 | Put your XXXXX.sh here 27 | ``` 28 | validations: 29 | required: true 30 | 31 | - type: textarea 32 | id: others 33 | attributes: 34 | label: Others 35 | description: | 36 | Describe your issue detailly here. 37 | 详细描述你的问题。 38 | validations: 39 | required: false 40 | -------------------------------------------------------------------------------- /IMDLBenCo/datasets/dummy_dataset.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | import numpy as np 4 | from ..registry import DATASETS 5 | from .abstract_dataset import AbstractDataset 6 | 7 | 8 | @DATASETS.register_module() 9 | class DummyDataset(AbstractDataset): 10 | def _init_dataset_path(self, path): 11 | """ 12 | Dummy dataset for testing purposes. 13 | Useless, return None. 14 | """ 15 | self.entry_path = path 16 | tp_list = ["dummy_tp_path"] 17 | gt_list = ["dummy_gt_path"] 18 | return tp_list, gt_list 19 | def _get_image(self, index): 20 | """ 21 | Dummy dataset for testing purposes. 22 | return random genrated images and masks. 23 | shape is (H, W, C), value range from 0 to 255. 24 | """ 25 | tp_image = np.random.randint(0, 255, (512, 512, 3), dtype=np.uint8) 26 | gt_image = np.random.randint(0, 255, (512, 512, 3), dtype=np.uint8) 27 | label = np.random.randint(0, 2) 28 | tp_shape = (512, 512) 29 | gt_shape = (512, 512) 30 | tp_path = "dummy_tp_path" 31 | gt_path = "dummy_gt_path" 32 | return tp_image, gt_image, label, tp_shape, gt_shape, tp_path, gt_path 33 | -------------------------------------------------------------------------------- /IMDLBenCo/model_zoo/trufor/README.md: -------------------------------------------------------------------------------- 1 | ### Pretrain weights 2 | 3 | The segformer pretrain weights are from the official repository: https://github.com/NVlabs/SegFormer. 4 | 5 | We also provide SegFormer weights in the following link. As the author does not provide the pretrain weights of the NoisePrint++, we separate the weights from the checkpoint provided: 6 | 7 | 百度网盘:链接: https://pan.baidu.com/s/1SAXJMiWbUsssk7RtwiOdMQ?pwd=3hmr 提取码: 3hmr 8 | 9 | Google Drive:https://drive.google.com/drive/folders/1Q9RxEHsIcRWeZjJRBtAwW4au5QybIoP2?usp=sharing 10 | 11 | ### Training Phase 12 | 13 | 14 | Trufor's training process is divided into three stages: noiseprint++, localization, and detection. Due to the lack of pretraining data for noiseprint++, we carefully extracted the weights of noiseprint++ from the checkpoint provided in the official repository to train the latter two stages. Therefore, Trufor accepts a **phase** parameter (2 or 3) to distinguish between localization training and detection training. In phase 3, the model accepts the **det_resume_ckpt** parameter to load the weights obtained from localization training and continue training the detection head, while in phase 2, this parameter is set to empty. 15 | -------------------------------------------------------------------------------- /IMDLBenCo/utils/misc.py: -------------------------------------------------------------------------------- 1 | from collections import abc 2 | from typing import Any, Callable, Optional, Type, Union 3 | 4 | def is_seq_of(seq: Any, 5 | expected_type: Union[Type, tuple], 6 | seq_type: Type = None) -> bool: 7 | """Check whether it is a sequence of some type. 8 | 9 | Args: 10 | seq (Sequence): The sequence to be checked. 11 | expected_type (type or tuple): Expected type of sequence items. 12 | seq_type (type, optional): Expected sequence type. Defaults to None. 13 | 14 | Returns: 15 | bool: Return True if ``seq`` is valid else False. 16 | 17 | Examples: 18 | >>> from mmengine.utils import is_seq_of 19 | >>> seq = ['a', 'b', 'c'] 20 | >>> is_seq_of(seq, str) 21 | True 22 | >>> is_seq_of(seq, int) 23 | False 24 | """ 25 | if seq_type is None: 26 | exp_seq_type = abc.Sequence 27 | else: 28 | assert isinstance(seq_type, type) 29 | exp_seq_type = seq_type 30 | if not isinstance(seq, exp_seq_type): 31 | return False 32 | for item in seq: 33 | if not isinstance(item, expected_type): 34 | return False 35 | return True -------------------------------------------------------------------------------- /IMDLBenCo/utils/paths.py: -------------------------------------------------------------------------------- 1 | 2 | import os 3 | from pathlib import Path 4 | import appdirs 5 | 6 | class BencoPath: 7 | app_name = "IMDLBenCo" 8 | app_author = "IMDLBenCo-authors" 9 | 10 | @classmethod 11 | def get_data_storage_path(cls): 12 | storage_path = appdirs.user_data_dir(cls.app_name, cls.app_author) 13 | if not os.path.exists(storage_path): 14 | os.makedirs(storage_path) 15 | return storage_path 16 | 17 | @classmethod 18 | def get_package_dir(cls): 19 | return Path(__file__).parent.parent 20 | 21 | @classmethod 22 | def get_templates_dir(cls): 23 | return cls.get_package_dir() / 'training_scripts' 24 | 25 | @classmethod 26 | def get_dataset_json_dir(cls): 27 | return cls.get_package_dir() / 'statics' / 'dataset_json' 28 | 29 | @classmethod 30 | def get_init_base_dir(cls): 31 | return cls.get_package_dir() / 'statics' / 'base' 32 | 33 | @classmethod 34 | def get_model_zoo_runs_dir(cls): 35 | return cls.get_package_dir() / 'statics' / 'model_zoo' / 'runs' 36 | 37 | @classmethod 38 | def get_model_zoo_configs_dir(cls): 39 | return cls.get_package_dir() / 'statics' / 'model_zoo' / 'configs' -------------------------------------------------------------------------------- /tests/test_json_datasets.py: -------------------------------------------------------------------------------- 1 | import sys 2 | from pprint import pprint 3 | sys.path.append(".") 4 | from torch.utils.data import Dataset, DataLoader 5 | import IMDLBenCo 6 | import IMDLBenCo.datasets 7 | from IMDLBenCo.datasets import ManiDataset 8 | from IMDLBenCo.datasets.jpeg_dataset_deprecated import MetaCatnetDataset 9 | from IMDLBenCo.registry import DATASETS 10 | import torch 11 | 12 | 13 | 14 | 15 | # data = IMDLBench.datasets.mani_dataset("/mnt/data0/public_datasets/IML/CASIA2.0", is_padding=True, edge_width= 7) 16 | # ['/mnt/data0/public_datasets/IML/IMD_20_1024', MetaCatnetDataset], 17 | # ['/mnt/data0/public_datasets/IML/tampCOCO/sp_COCO_list.json', MetaCatnetDataset], 18 | # ['/mnt/data0/public_datasets/IML/tampCOCO/cm_COCO_list.json', MetaCatnetDataset], 19 | # ['/mnt/data0/public_datasets/IML/tampCOCO/bcm_COCO_list.json', MetaCatnetDataset], 20 | # ['/mnt/data0/public_datasets/IML/tampCOCO/bcmc_COCO_list.json', MetaCatnetDataset] 21 | data = MetaCatnetDataset("/mnt/data0/public_datasets/IML/IMD_20_1024", is_padding=True, edge_width= 7) 22 | 23 | 24 | batch_size = 1 25 | dataloader = DataLoader(data, batch_size=batch_size, shuffle=True) 26 | 27 | for i in data: 28 | for key in i.keys(): 29 | if isinstance(i[key], torch.Tensor): 30 | print(key, i[key].shape) 31 | exit() -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/runs/demo_train_trufor.sh: -------------------------------------------------------------------------------- 1 | base_dir="./output_dir" 2 | mkdir -p ${base_dir} 3 | # For Pretrained weights please reference here: 4 | # https://github.com/scu-zjz/IMDLBenCo/tree/main/IMDLBenCo/model_zoo/trufor 5 | 6 | CUDA_VISIBLE_DEVICES=4 \ 7 | torchrun \ 8 | --standalone \ 9 | --nnodes=1 \ 10 | --nproc_per_node=1 \ 11 | ./train.py \ 12 | --model Trufor \ 13 | --world_size 1 \ 14 | --batch_size 8 \ 15 | --data_path /mnt/data0/public_datasets/IML/CASIA2.0 \ 16 | --np_pretrain_weights "/mnt/data0/dubo/workspace/IMDLBenCo/IMDLBenCo/model_zoo/trufor/noiseprint.pth" \ 17 | --mit_b2_pretrain_weights "/mnt/data0/dubo/workspace/IMDLBenCo/IMDLBenCo/model_zoo/trufor/mit_b2.pth" \ 18 | --config_path "./configs/trufor.yaml" \ 19 | --phase 2 \ 20 | --det_resume_ckpt "" \ 21 | --epochs 150 \ 22 | --lr 4e-6 \ 23 | --if_predict_label \ 24 | --if_not_amp \ 25 | --find_unused_parameters \ 26 | --image_size 512 \ 27 | --if_resizing \ 28 | --min_lr 0 \ 29 | --weight_decay 0.05 \ 30 | --edge_mask_width 7 \ 31 | --test_data_path "/mnt/data0/public_datasets/IML/CASIA1.0" \ 32 | --warmup_epochs 2 \ 33 | --output_dir ${base_dir}/ \ 34 | --log_dir ${base_dir}/ \ 35 | --accum_iter 1 \ 36 | --seed 42 \ 37 | --test_period 4 \ 38 | 2> ${base_dir}/error.log 1>${base_dir}/logs.log -------------------------------------------------------------------------------- /tests/test_metrics/dataset.json: -------------------------------------------------------------------------------- 1 | [ 2 | [ 3 | "images/01.jpg", 4 | "masks/01_mask.jpg" 5 | ], 6 | [ 7 | "images/02.jpg", 8 | "masks/02_mask.jpg" 9 | ], 10 | [ 11 | "images/03.jpg", 12 | "masks/03_mask.jpg" 13 | ], 14 | [ 15 | "images/04.jpg", 16 | "masks/04_mask.jpg" 17 | ], 18 | [ 19 | "images/05.jpg", 20 | "masks/05_mask.jpg" 21 | ], 22 | [ 23 | "images/06.jpg", 24 | "masks/06_mask.jpg" 25 | ], 26 | [ 27 | "images/07.jpg", 28 | "masks/07_mask.jpg" 29 | ], 30 | [ 31 | "images/08.jpg", 32 | "masks/08_mask.jpg" 33 | ], 34 | [ 35 | "images/09.jpg", 36 | "masks/09_mask.jpg" 37 | ], 38 | [ 39 | "images/10.jpg", 40 | "masks/10_mask.jpg" 41 | ], 42 | [ 43 | "images/11.jpg", 44 | "Negative" 45 | ], 46 | [ 47 | "images/12.jpg", 48 | "Negative" 49 | ], 50 | [ 51 | "images/13.jpg", 52 | "Negative" 53 | ], 54 | [ 55 | "images/14.jpg", 56 | "Negative" 57 | ], 58 | [ 59 | "images/15.jpg", 60 | "Negative" 61 | ], 62 | [ 63 | "images/16.jpg", 64 | "Negative" 65 | ], 66 | [ 67 | "images/17.jpg", 68 | "Negative" 69 | ], 70 | [ 71 | "images/18.jpg", 72 | "Negative" 73 | ], 74 | [ 75 | "images/19.jpg", 76 | "Negative" 77 | ], 78 | [ 79 | "images/20.jpg", 80 | "Negative" 81 | ] 82 | ] -------------------------------------------------------------------------------- /IMDLBenCo/modules/extractors/bayar_conv.py: -------------------------------------------------------------------------------- 1 | from cv2 import detail_ImageFeatures 2 | import numpy as np 3 | import matplotlib.pyplot as plt 4 | from PIL import Image 5 | from collections import OrderedDict 6 | import torch 7 | from torch import nn 8 | import torch.nn.functional as F 9 | 10 | 11 | class BayerConv(nn.Module): 12 | def __init__(self, in_channel=3): 13 | super(BayerConv, self).__init__() 14 | self.BayarConv2D = nn.Conv2d(in_channel, 3, 5, 1, padding=2, bias=False) 15 | 16 | # Convert tensors to Parameters 17 | bayar_mask = np.ones(shape=(5, 5)) 18 | bayar_mask[2, 2] = 0 19 | self.bayar_mask = nn.Parameter(torch.tensor(bayar_mask, dtype=torch.float32), requires_grad=False) 20 | 21 | bayar_final = np.zeros((5, 5)) 22 | bayar_final[2, 2] = -1 23 | self.bayar_final = nn.Parameter(torch.tensor(bayar_final, dtype=torch.float32), requires_grad=False) 24 | 25 | def forward(self, x): 26 | self.BayarConv2D.weight.data *= self.bayar_mask 27 | self.BayarConv2D.weight.data *= torch.pow(self.BayarConv2D.weight.data.sum(axis=(2, 3)).view(3, 3, 1, 1), -1) 28 | self.BayarConv2D.weight.data += self.bayar_final 29 | 30 | return self.BayarConv2D(x) 31 | 32 | if __name__ == '__main__': 33 | torch.manual_seed(42) 34 | img = torch.randn(2,3,256,256).to(0) 35 | model = BayerConv().to(0) 36 | print(model(img)) -------------------------------------------------------------------------------- /tests/gradcam.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import sys 3 | sys.path.append(".") 4 | from IMDLBenCo.datasets import ManiDataset 5 | from IMDLBenCo.transforms import get_albu_transforms 6 | from IMDLBenCo.model_zoo import cat_net 7 | from IMDLBench.model_zoo.cat_net.cat_net_post_function import cat_net_post_func 8 | import torch 9 | from IMDLBenCo.evaluation import grad_camera_visualize 10 | 11 | if __name__ == '__main__': 12 | model = cat_net('/home/bingkui/IMDLBenCo/IMDLBench/training/CAT_full.yaml') # TODO 这里加载模型 13 | ckpt = '/home/bingkui/IMDLBenCo/output_dir_balance/checkpoint-44.pth' # TODO 这里填已经训练好的模型 14 | ckpt = torch.load(ckpt, map_location='cuda') 15 | model.load_state_dict(ckpt['model']) 16 | model.cuda() 17 | 18 | dataset = ManiDataset(path='/mnt/data0/public_datasets/IML/CASIA1.0', 19 | is_padding=False, 20 | is_resizing=True, 21 | output_size=(512, 512), 22 | common_transforms=get_albu_transforms('test'), 23 | edge_width=7)[0:1] 24 | #post_funcs=cat_net_post_func -> this argument is only for Cat-Net 25 | 26 | target_layers = [model.model.last_layer[-1]] 27 | grad_camera_visualize(model=model, 28 | image=dataset, 29 | target_layers=target_layers, # TODO 这里放你的模型结构中最后一个计算单元，用list装起来 30 | output_path='/home/bingkui/IMDLBenCo/images/CASIA1.0') # TODO 这里放图片输出的文件夹地址 31 | -------------------------------------------------------------------------------- /IMDLBenCo/model_zoo/trufor/cmx/layer_utils.py: -------------------------------------------------------------------------------- 1 | # %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% 2 | # Copyright (c) 2023 Image Processing Research Group of University Federico II of Naples ('GRIP-UNINA'). 3 | # 4 | # All rights reserved. 5 | # This work should only be used for nonprofit purposes. 6 | # 7 | # By downloading and/or using any of these files, you implicitly agree to all the 8 | # terms of the license, as specified in the document LICENSE.txt 9 | # (included in this package) and online at 10 | # http://www.grip.unina.it/download/LICENSE_OPEN.txt 11 | 12 | """ 13 | Created in September 2022 14 | @author: davide.cozzolino 15 | """ 16 | 17 | import torch 18 | import torch.nn.functional as F 19 | 20 | 21 | def weighted_statistics_pooling(x, log_w=None): 22 | b = x.shape[0] 23 | c = x.shape[1] 24 | x = x.view(b,c,-1) 25 | 26 | if log_w is None: 27 | log_w = torch.zeros((b,1,x.shape[-1]), device=x.device) 28 | else: 29 | assert log_w.shape[0]==b 30 | assert log_w.shape[1]==1 31 | log_w = log_w.view(b,1,-1) 32 | 33 | assert log_w.shape[-1]==x.shape[-1] 34 | 35 | log_w = F.log_softmax(log_w, dim=-1) 36 | x_min = -torch.logsumexp(log_w-x, dim=-1) 37 | x_max = torch.logsumexp(log_w+x, dim=-1) 38 | 39 | w = torch.exp(log_w) 40 | x_avg = torch.sum(w*x , dim=-1) 41 | x_msq = torch.sum(w*x*x, dim=-1) 42 | 43 | x = torch.cat((x_min, x_max, x_avg, x_msq), dim=1) 44 | 45 | return x -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = ["setuptools>=64.0", "wheel"] 3 | build-backend = "setuptools.build_meta" 4 | 5 | [project] 6 | name = "imdlbenco" 7 | authors = [ 8 | {name = "Xiaochen Ma", email = "xiaochen.ma.cs@gmail.com"}, 9 | ] 10 | description = "A comprehensive benchmark and code base for Image manipulation and localization." 11 | readme = {file = "README.md", content-type = "text/markdown"} 12 | requires-python = ">=3.7, <4" 13 | license = {text = "CC-BY-4.0"} 14 | classifiers = [ 15 | "Development Status :: 3 - Alpha", 16 | "Intended Audience :: Developers", 17 | "Intended Audience :: Science/Research", 18 | "Topic :: Security", 19 | "Topic :: Scientific/Engineering :: Artificial Intelligence", 20 | "License :: Free For Educational Use", 21 | "Programming Language :: Python :: 3", 22 | "Programming Language :: Python :: 3.7", 23 | "Programming Language :: Python :: 3.8", 24 | "Programming Language :: Python :: 3.9", 25 | "Programming Language :: Python :: 3.10", 26 | "Programming Language :: Python :: 3 :: Only", 27 | ] 28 | keywords = [ 29 | "AI", 30 | "artificial intelligence", 31 | "image forensics", 32 | "image manipulation localization", 33 | "image manipulation detection", 34 | ] 35 | dynamic = ["version", "dependencies"] 36 | 37 | [project.urls] 38 | Github = "https://github.com/scu-zjz/IMDLBenCo/" 39 | Documentation = "https://github.com/scu-zjz/IMDLBenCo-doc" 40 | "Bug Reports" = "https://github.com/scu-zjz/IMDLBenCo/issues" 41 | 42 | [project.entry-points."console_scripts"] 43 | benco = "IMDLBenCo.cli:main" 44 | 45 | [tool.setuptools] 46 | include-package-data = true 47 | packages = ["IMDLBenCo"] # 显式指定主包 48 | 49 | 50 | [tool.setuptools.dynamic] 51 | version = {attr = "IMDLBenCo.version.__version__"} 52 | dependencies = {file = "requirements.txt"} -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/1-bug-reports.yaml: -------------------------------------------------------------------------------- 1 | name: "🐞 Bug reports" 2 | description: Report Bug to help us improve IMDLBenCo 3 | labels: bug 4 | 5 | body: 6 | - type: markdown 7 | attributes: 8 | value: | 9 | Please briefly describe your issue, and try to use the search function with relevant keywords to look for solutions in existing issues first. 10 | 请简要描述你的困难，并优先活用搜索功能配合关键字在现有issue中寻找解决方案。 11 | 12 | - type: checkboxes 13 | id: reminder 14 | attributes: 15 | label: Reminder 16 | description: | 17 | Please ensure you have carefully searched the existing issues. 18 | 请确保您已经认真搜索过现有的 issues。 19 | options: 20 | - label: I have read the above rules and searched the existing issues. 21 | required: true 22 | 23 | - type: textarea 24 | id: system-info 25 | attributes: 26 | label: System Info 27 | description: | 28 | Please share your system info with us. You can run the command **benco env** and copy-paste its output below. 29 | 请提供您的系统信息。您可以在命令行运行 **benco env** 并将其输出复制到该文本框中。 30 | placeholder: benco version, platform, python version, ... 31 | validations: 32 | required: true 33 | 34 | - type: textarea 35 | id: reproduction 36 | attributes: 37 | label: Reproduction 38 | description: | 39 | Please provide entry arguments, error messages and stack traces that reproduces the problem. 40 | 请提供入口参数，错误日志以及异常堆栈以便于我们复现问题。 41 | value: | 42 | ```text 43 | Put your message here. 44 | ``` 45 | validations: 46 | required: true 47 | 48 | - type: textarea 49 | id: others 50 | attributes: 51 | label: Others 52 | description: | 53 | Describe your issue or bug detailly here. 54 | 详细描述你的问题或发现的bug。 55 | validations: 56 | required: false 57 | -------------------------------------------------------------------------------- /IMDLBenCo/modules/extractors/srm_filter.py: -------------------------------------------------------------------------------- 1 | from cv2 import detail_ImageFeatures 2 | import numpy as np 3 | import matplotlib.pyplot as plt 4 | from PIL import Image 5 | from collections import OrderedDict 6 | #Pytorch 7 | import torch 8 | from torch import nn 9 | import torch.nn.functional as F 10 | 11 | 12 | class SRMConv2D(nn.Module): 13 | def __init__(self): 14 | super(SRMConv2D,self).__init__() 15 | q = [4, 12, 2] # coefficient of the kernels 16 | kernel1 = np.array([ 17 | [0, 0, 0, 0, 0], 18 | [0,-1, 2,-1, 0], 19 | [0, 2,-4, 2, 0], 20 | [0,-1, 2,-1, 0], 21 | [0, 0, 0, 0, 0] 22 | ],dtype=np.float32) 23 | kernel2 = np.array([ 24 | [-1, 2,-2, 2,-1], 25 | [2, -6, 8,-6, 2], 26 | [-2, 8,-12,8,-2], 27 | [2, -6, 8,-6, 2], 28 | [-1, 2,-2, 2,-1] 29 | ],dtype=np.float32) 30 | kernel3 = np.array([ 31 | [0, 0, 0, 0, 0], 32 | [0, 0, 0, 0, 0], 33 | [0, 1,-2, 1, 0], 34 | [0, 0, 0, 0, 0], 35 | [0, 0, 0, 0, 0], 36 | ],dtype=np.float32) 37 | zero_kernel = np.zeros_like(kernel3) 38 | # shape (3,9,5,5) 39 | weight = torch.tensor(np.array([ 40 | [[kernel1 / q[0] if j == i else zero_kernel for j in range(3)] for i in range(3)], 41 | [[kernel2 / q[1] if j == i else zero_kernel for j in range(3)] for i in range(3)], 42 | [[kernel3 / q[2] if j == i else zero_kernel for j in range(3)] for i in range(3)], 43 | ]),dtype=torch.float32) 44 | weight = weight.reshape(-1,3,5,5) 45 | self.weight = torch.nn.Parameter(weight, requires_grad=False) 46 | 47 | def forward(self, x): 48 | with torch.no_grad(): 49 | return torch.nn.functional.conv2d(x, weight=self.weight, padding = 2) 50 | 51 | 52 | -------------------------------------------------------------------------------- /tests/test_openfile.py: -------------------------------------------------------------------------------- 1 | import os 2 | import subprocess 3 | import tempfile 4 | import platform 5 | 6 | def get_default_editor(): 7 | """Get the system's default editor.""" 8 | editor = os.getenv('EDITOR') 9 | if not editor: 10 | if platform.system() == 'Windows': 11 | editor = 'notepad' 12 | elif platform.system() == 'Darwin': # macOS 13 | editor = 'nano' # macOS does not have a default GUI editor 14 | else: # Linux and other Unix-like systems 15 | editor = 'code --wait' # 'nano' is a common text editor in Unix-like systems 16 | return editor 17 | 18 | def open_file_with_editor(file_path): 19 | """Open the file with the default editor.""" 20 | editor = get_default_editor() 21 | try: 22 | subprocess.run([editor, file_path]) 23 | except Exception as e: 24 | print(f"Error opening file with editor {editor}: {e}") 25 | 26 | def commit_message(): 27 | """Simulate the Git commit message process.""" 28 | with tempfile.NamedTemporaryFile(delete=False, mode='w+', suffix=".tmp") as temp_file: 29 | temp_file_path = temp_file.name 30 | temp_file.write("# Please enter the commit message for your changes. Lines starting\n") 31 | temp_file.write("# with '#' will be ignored, and an empty message aborts the commit.\n") 32 | temp_file.write("\n") 33 | temp_file.write("# This is a comment line.\n") 34 | temp_file.flush() 35 | open_file_with_editor(temp_file_path) 36 | 37 | with open(temp_file_path, 'r') as f: 38 | commit_msg = ''.join([line for line in f if not line.startswith('#')]).strip() 39 | 40 | os.remove(temp_file_path) 41 | 42 | if not commit_msg: 43 | print("Aborting commit due to empty commit message.") 44 | else: 45 | print(f"Commit message is:\n{commit_msg}") 46 | 47 | if __name__ == "__main__": 48 | commit_message() 49 | -------------------------------------------------------------------------------- /IMDLBenCo/model_zoo/pscc_net/seg_hrnet_config.py: -------------------------------------------------------------------------------- 1 | # ------------------------------------------------------------------------------ 2 | # Copyright (c) Microsoft 3 | # Licensed under the MIT License. 4 | # Written by Ke Sun (sunk@mail.ustc.edu.cn) 5 | # ------------------------------------------------------------------------------ 6 | 7 | from __future__ import absolute_import 8 | from __future__ import division 9 | from __future__ import print_function 10 | 11 | from yacs.config import CfgNode as CN 12 | 13 | # high_resoluton_net related params for segmentation 14 | HRNET = CN() 15 | HRNET.PRETRAINED_LAYERS = ['*'] 16 | HRNET.STEM_INPLANES = 64 17 | HRNET.FINAL_CONV_KERNEL = 1 18 | 19 | HRNET.STAGE1 = CN() 20 | HRNET.STAGE1.NUM_MODULES = 1 21 | HRNET.STAGE1.NUM_BRANCHES = 1 22 | HRNET.STAGE1.NUM_BLOCKS = [2] 23 | HRNET.STAGE1.NUM_CHANNELS = [64] 24 | HRNET.STAGE1.BLOCK = 'BOTTLENECK' 25 | HRNET.STAGE1.FUSE_METHOD = 'SUM' 26 | 27 | HRNET.STAGE2 = CN() 28 | HRNET.STAGE2.NUM_MODULES = 1 29 | HRNET.STAGE2.NUM_BRANCHES = 2 30 | HRNET.STAGE2.NUM_BLOCKS = [2, 2] 31 | HRNET.STAGE2.NUM_CHANNELS = [18, 36] 32 | HRNET.STAGE2.BLOCK = 'BASIC' 33 | HRNET.STAGE2.FUSE_METHOD = 'SUM' 34 | 35 | HRNET.STAGE3 = CN() 36 | HRNET.STAGE3.NUM_MODULES = 1 37 | HRNET.STAGE3.NUM_BRANCHES = 3 38 | HRNET.STAGE3.NUM_BLOCKS = [2, 2, 2] 39 | HRNET.STAGE3.NUM_CHANNELS = [18, 36, 72] 40 | HRNET.STAGE3.BLOCK = 'BASIC' 41 | HRNET.STAGE3.FUSE_METHOD = 'SUM' 42 | 43 | HRNET.STAGE4 = CN() 44 | HRNET.STAGE4.NUM_MODULES = 1 45 | HRNET.STAGE4.NUM_BRANCHES = 4 46 | HRNET.STAGE4.NUM_BLOCKS = [2, 2, 2, 2] 47 | HRNET.STAGE4.NUM_CHANNELS = [18, 36, 72, 144] 48 | HRNET.STAGE4.BLOCK = 'BASIC' 49 | HRNET.STAGE4.FUSE_METHOD = 'SUM' 50 | 51 | 52 | def get_hrnet_cfg(pretrain_path=None): 53 | """Get a yacs CfgNode object with default values for my_project.""" 54 | # Return a clone so that the defaults will not be altered 55 | # This is for the "local variable" use pattern 56 | HRNET.PRETRAINED = pretrain_path 57 | return HRNET.clone() -------------------------------------------------------------------------------- /IMDLBenCo/evaluation/gradcam/activations_and_gradients_hack.py: -------------------------------------------------------------------------------- 1 | class ActivationsAndGradients: 2 | """ Class for extracting activations and 3 | registering gradients from targetted intermediate layers """ 4 | 5 | def __init__(self, model, target_layers, reshape_transform): 6 | self.model = model 7 | self.gradients = [] 8 | self.activations = [] 9 | self.reshape_transform = reshape_transform 10 | self.handles = [] 11 | for target_layer in target_layers: 12 | self.handles.append( 13 | target_layer.register_forward_hook(self.save_activation)) 14 | # Because of https://github.com/pytorch/pytorch/issues/61519, 15 | # we don't use backward hook to record gradients. 16 | self.handles.append( 17 | target_layer.register_forward_hook(self.save_gradient)) 18 | 19 | def save_activation(self, module, input, output): 20 | activation = output 21 | 22 | if self.reshape_transform is not None: 23 | activation = self.reshape_transform(activation) 24 | self.activations.append(activation.cpu().detach()) 25 | 26 | def save_gradient(self, module, input, output): 27 | if not hasattr(output, "requires_grad") or not output.requires_grad: 28 | # You can only register hooks on tensor requires grad. 29 | return 30 | 31 | # Gradients are computed in reverse order 32 | def _store_grad(grad): 33 | if self.reshape_transform is not None: 34 | grad = self.reshape_transform(grad) 35 | self.gradients = [grad.cpu().detach()] + self.gradients 36 | 37 | output.register_hook(_store_grad) 38 | 39 | def __call__(self, data): 40 | self.gradients = [] 41 | self.activations = [] 42 | # import pdb;pdb.set_trace() 43 | return self.model(data) 44 | 45 | def release(self): 46 | for handle in self.handles: 47 | handle.remove() 48 | -------------------------------------------------------------------------------- /tests/test_metrics_cpu.py: -------------------------------------------------------------------------------- 1 | import sys 2 | from pprint import pprint 3 | sys.path.append(".") 4 | from torch.utils.data import Dataset, DataLoader 5 | import IMDLBenCo 6 | import IMDLBenCo.datasets 7 | from IMDLBenCo.datasets import ManiDataset 8 | from IMDLBenCo.datasets.jpeg_dataset_deprecated import MetaCatnetDataset 9 | from IMDLBenCo.registry import DATASETS 10 | import torch 11 | 12 | from IMDLBenCo.evaluation import generate_region_mask, cal_confusion_matrix, cal_F1 13 | from IMDLBenCo.datasets import denormalize 14 | 15 | 16 | IS_PADDING = True 17 | IS_RESIZING = False 18 | batch_size = 3 19 | 20 | # 主要测试这个样例数据集的指标： 21 | data = IMDLBenCo.datasets.ManiDataset("/mnt/data0/public_datasets/IML/basic_eval_dataset", 22 | is_padding=IS_PADDING, 23 | is_resizing=IS_RESIZING, 24 | edge_width=None) 25 | 26 | 27 | batch_size = 3 28 | dataloader = DataLoader(data, batch_size=batch_size, shuffle=False) 29 | 30 | F1_list = [] 31 | 32 | for i in dataloader: 33 | i['image'] = torch.mean(i['image'], dim=1, keepdim=True) # [3, 3, 1024, 1024] -> [3, 1, 1024, 1024] 34 | mask = i['mask'] 35 | pred = i['image'] 36 | shape_mask = i['shape_mask'] 37 | shape = i['shape'] 38 | print(shape) 39 | mask_crop = mask[:, :, 0:shape[0], 0:shape[1]] 40 | 41 | 42 | import matplotlib.pyplot as plt 43 | plt.subplot(3, 3, 1) 44 | plt.imshow(i['mask'][0][0]) 45 | plt.subplot(3, 3, 2) 46 | plt.imshow(i['mask'][1][0]) 47 | 48 | # test for evaluation split 49 | # regin_mask = genertate_region_mask(i['masks'], i['shapes']) 50 | plt.subplot(3, 2, 3) 51 | plt.title('shape_mask') 52 | plt.imshow(i['shape_mask'][0][0]) 53 | plt.subplot(3, 2, 4) 54 | plt.imshow(mask_crop[0]) 55 | 56 | plt.subplot(3, 2, 5) 57 | plt.imshow(i['image'][0][0]) 58 | plt.subplot(3, 2, 6) 59 | plt.imshow(i['image'][1][0]) 60 | 61 | 62 | plt.savefig("test.png") 63 | exit() 64 | 65 | 66 | -------------------------------------------------------------------------------- /IMDLBenCo/transforms/robustness_wrapper.py: -------------------------------------------------------------------------------- 1 | import albumentations as albu 2 | 3 | class AbstractTransformWrapper: 4 | def __init__(self, param_list): 5 | self.param_list = param_list 6 | self.index = 0 7 | 8 | def __iter__(self): 9 | return self 10 | 11 | def __str__(self): 12 | return self.__class__.__name__[:-7] # 返回类名 13 | 14 | def __next__(self): 15 | if self.index < len(self.param_list): 16 | param = self.param_list[self.index] 17 | self.index += 1 18 | 19 | if param == 0: 20 | return param, None 21 | else: 22 | return self._get_transform(param) 23 | else: 24 | self.index = 0 25 | raise StopIteration 26 | 27 | def _get_transform(self, param): 28 | raise NotImplementedError 29 | 30 | 31 | class GaussianBlurWrapper(AbstractTransformWrapper): 32 | def _get_transform(self, param): 33 | return param, albu.GaussianBlur( 34 | blur_limit=(param, param), 35 | always_apply=True, 36 | p=1.0 37 | ) 38 | 39 | class GaussianNoiseWrapper(AbstractTransformWrapper): 40 | def _get_transform(self, param): 41 | return param, albu.GaussNoise( 42 | var_limit=(param, param), 43 | always_apply=True, 44 | p=1.0 45 | ) 46 | 47 | class JpegCompressionWrapper(AbstractTransformWrapper): 48 | def _get_transform(self, param): 49 | return param, albu.JpegCompression( 50 | quality_lower = param-1, 51 | quality_upper = param, 52 | p=1.0 53 | ) 54 | 55 | if __name__ == "__main__": 56 | # 示例用法 57 | param_list = [90, 80, 0] # 传入一个装有int的列表，代表需要遍历的参数数值 58 | wrapper = JpegCompressionWrapper(param_list) 59 | 60 | for transform in wrapper: 61 | print(transform) 62 | 63 | 64 | for transform in wrapper: 65 | print(transform) 66 | 67 | print(str(wrapper)) 68 | 69 | -------------------------------------------------------------------------------- /IMDLBenCo/datasets/utils.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import torch 3 | from PIL import Image 4 | import numpy as np 5 | from PIL import Image 6 | import tempfile, os, io 7 | 8 | 9 | def import_jpegio() -> None: 10 | try: 11 | import jpegio as jio 12 | return jio 13 | except ImportError: 14 | raise ImportError( 15 | 'Please run "pip install jpegio" to install jpegio. This only support Linux system') 16 | 17 | def pil_loader(path: str) -> Image.Image: 18 | """PIL image loader 19 | 20 | Args: 21 | path (str): image path 22 | 23 | Returns: 24 | Image.Image: PIL image (after np.array(x) becomes [0,255] int8) 25 | """ 26 | # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835) 27 | with open(path, 'rb') as f: 28 | img = Image.open(f) 29 | return img.convert('RGB') 30 | 31 | 32 | def jpeg_loader(path : str): 33 | pass # TODO TBK 34 | 35 | def denormalize(image, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]): 36 | """denormalize image with mean and std 37 | """ 38 | image = image.clone().detach().cpu() 39 | image = image * torch.tensor(std).view(3, 1, 1) 40 | image = image + torch.tensor(mean).view(3, 1, 1) 41 | return image 42 | 43 | def convert_to_temp_jpeg(tensor): 44 | tensor = tensor.permute(1, 2, 0) 45 | img = Image.fromarray(tensor.numpy().astype('uint8')) 46 | 47 | with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as temp_file: 48 | img.save(temp_file, format='JPEG') 49 | temp_file_path = temp_file.name 50 | 51 | return temp_file_path 52 | 53 | def read_jpeg_from_memory(tensor): 54 | jio = import_jpegio() 55 | temp_file_path = convert_to_temp_jpeg(tensor) 56 | jpeg = jio.read(temp_file_path) 57 | os.remove(temp_file_path) 58 | 59 | return jpeg 60 | 61 | 62 | if __name__ == "__main__": 63 | # pass 64 | path = '/mnt/data0/public_datasets/IML/CASIA2.0/Tp/Tp_D_NRN_S_N_arc00090_art00010_00013.tif' 65 | print(read_jpeg_from_memory(path)) -------------------------------------------------------------------------------- /IMDLBenCo/evaluation/temp.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | import matplotlib.pyplot as plt 5 | from sklearn.metrics import roc_auc_score,roc_curve 6 | import numpy as np 7 | from torch.utils.data import DataLoader 8 | def generate_region_mask(masks ,batch_shape:torch.Tensor): 9 | """generate B 1 H W meaningful-region-mask for a batch of masks 10 | 11 | Args: 12 | batch_shape (_type_): list of tensor, e.g. [ tensor(384, 256), tensor(256, 384), ....] 13 | """ 14 | meaningful_mask = torch.zeros_like(masks) 15 | for idx, shape in enumerate(batch_shape): 16 | # shape = shape.tolist() 17 | meaningful_mask[idx, :, :shape[0], :shape[1]] = 1 # TODO 可能会出傻逼问题（长宽翻转）的地方 一定要做测试 18 | return meaningful_mask 19 | 20 | def cal_confusion_matrix(predict, target, region_mask, threshold=0.5): 21 | """compute local confusion matrix for a batch of predict and target masks 22 | Args: 23 | predict (_type_): _description_ 24 | target (_type_): _description_ 25 | region (_type_): _description_ 26 | 27 | Returns: 28 | TP, TN, FP, FN 29 | """ 30 | predict = (predict > threshold).float() 31 | TP = torch.sum(predict * target * region_mask, dim=(1, 2, 3)) 32 | TN = torch.sum((1-predict) * (1-target) * region_mask, dim=(1, 2, 3)) 33 | FP = torch.sum(predict * (1-target) * region_mask, dim=(1, 2, 3)) 34 | FN = torch.sum((1-predict) * target * region_mask, dim=(1, 2, 3)) 35 | return TP, TN, FP, FN 36 | 37 | def cal_F1(TP, TN, FP, FN): 38 | """_summary_ 39 | 40 | Args: 41 | TP (_type_): _description_ 42 | TN (_type_): _description_ 43 | FP (_type_): _description_ 44 | FN (_type_): _description_ 45 | 46 | Returns: 47 | _type_: _description_ 48 | """ 49 | precision = TP / (TP + FP + 1e-8) 50 | recall = TP / (TP + FN + 1e-8) 51 | F1 = 2 * precision * recall / (precision + recall + 1e-8) 52 | # F1 = torch.mean(F1) # fuse the Batch dimension 53 | return F1 54 | 55 | if __name__ == "__main__": 56 | pass -------------------------------------------------------------------------------- /tests/pytests/test_evaluators_empty.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import torch 3 | from unittest.mock import patch 4 | from IMDLBenCo.evaluation import ImageF1, ImageAUC, ImageAccuracy # 替换为你的实际路径 5 | 6 | @pytest.fixture(autouse=True) 7 | def patch_distributed_env(): 8 | with patch("IMDLBenCo.training_scripts.utils.misc.get_world_size", return_value=1), \ 9 | patch("IMDLBenCo.training_scripts.utils.misc.get_rank", return_value=0), \ 10 | patch("torch.distributed.all_gather", side_effect=lambda out_list, tensor: out_list.__setitem__(0, tensor.clone())): 11 | yield 12 | 13 | @pytest.mark.parametrize("EvaluatorClass, expected_value", [ 14 | (ImageF1, torch.tensor(1.0)), 15 | (ImageAUC, torch.tensor(1.0)), 16 | (ImageAccuracy, torch.tensor(1.0)), 17 | ]) 18 | def test_epoch_update_with_empty_predict(EvaluatorClass, expected_value): 19 | evaluator = EvaluatorClass() 20 | evaluator.remain_update(torch.tensor([[0.8], [0.2]]), torch.tensor([[1.0], [0.0]])) 21 | result = evaluator.epoch_update() 22 | # 转换成 float 以兼容 float32 / float64 / Tensor 类型 23 | result = result.item() if isinstance(result, torch.Tensor) else float(result) 24 | assert result == pytest.approx(expected_value, abs=1e-5), f"{EvaluatorClass.__name__} failed on remain-only data" 25 | 26 | @pytest.mark.parametrize("EvaluatorClass, expected_value", [ 27 | (ImageF1, torch.tensor(1.0)), 28 | (ImageAUC, torch.tensor(1.0)), 29 | (ImageAccuracy, torch.tensor(1.0)), 30 | ]) 31 | def test_epoch_update_with_empty_remain(EvaluatorClass, expected_value): 32 | evaluator = EvaluatorClass() 33 | evaluator.batch_update(torch.tensor([[0.9], [0.1]]), torch.tensor([[1.0], [0.0]])) 34 | result = evaluator.epoch_update() 35 | result = torch.tensor(result, dtype=torch.float32) if not isinstance(result, torch.Tensor) else result 36 | assert torch.isclose(result, expected_value, atol=1e-5), f"{EvaluatorClass.__name__} mismatch with only batch data" 37 | 38 | @pytest.mark.parametrize("EvaluatorClass", [ImageF1, ImageAUC, ImageAccuracy]) 39 | def test_epoch_update_with_both_empty(EvaluatorClass): 40 | evaluator = EvaluatorClass() 41 | with pytest.raises((RuntimeError)): 42 | _ = evaluator.epoch_update() 43 | -------------------------------------------------------------------------------- /IMDLBenCo/evaluation/abstract_class.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | """ 3 | 下面这个类很重要，可以自动管理在多卡之间的一个具体数值变量的reduce（显卡之间归并数据） 4 | 可以用于实现算法的参考。 5 | """ 6 | # from training.utils.misc import MetricLogger 7 | 8 | """ 9 | 改变这个接口的主要目的是image-level的指标和pixel-level的指标的计算方式不同 10 | """ 11 | class AbstractEvaluator(object): # 想了想没必要用nn.module 反而可能会引起一些其他的问题，比如梯度追踪，或者Parameter的追踪等等问题？？？（我有点困不是很确定） 12 | def __init__(self) -> None: 13 | self.name = None 14 | self.desc = None 15 | self.threshold = None 16 | def _check_pixel_level_params(self, predict, mask): 17 | if predict is None: 18 | raise ValueError(f"Detect none mask predict from the model, cannot calculate {self.name}. Please remove Pixel-level metrics from the script, or check the model output.") 19 | if mask is None: 20 | raise ValueError(f"Detect none mask label from the dataset, cannot calculate {self.name}. Please remove Pixel-level metrics from the script, or check the dataset output.") 21 | def _chekc_image_level_params(self, predict_label, label): 22 | if predict_label is None: 23 | raise ValueError(f"Detect none image-level predict label from the model, cannot calculate {self.name}. Please remove Image-level metrics from the script, or check the model output.") 24 | if label is None: 25 | raise ValueError(f"Detect none image-level binary label from the dataset, cannot calculate {self.name}. Please remove Image-level metrics from the script, or check the dataset output.") 26 | def batch_update(self, predict, pred_label, mask, shape_mask=None, *args, **kwargs): 27 | """ 28 | 本函数在每个batch结尾update。 29 | """ 30 | raise NotImplementedError 31 | def remain_update(self, predict, pred_label, mask, shape_mask=None, *args, **kwargs): 32 | """ 33 | 本函数在每个batch结尾update。 34 | 主要用于处理在最后一个batch之后的剩余数据。 35 | """ 36 | raise NotImplementedError 37 | def epoch_update(self): 38 | """ 39 | 理论上这个时候没有新的数据了，所以没有输入参数。 40 | 41 | 功能：在显卡之间收集所有在整个epoch内统计的指标，然后返回最终期望的信息。 42 | """ 43 | raise NotImplementedError 44 | def recovery(self): 45 | raise NotImplementedError 46 | 47 | 48 | 49 | 50 | 51 | -------------------------------------------------------------------------------- /tests/test_datasets_and_evaluation.py: -------------------------------------------------------------------------------- 1 | import sys 2 | from pprint import pprint 3 | sys.path.append(".") 4 | from torch.utils.data import Dataset, DataLoader 5 | import IMDLBenCo 6 | import IMDLBenCo.datasets 7 | from IMDLBenCo.datasets import ManiDataset 8 | from IMDLBenCo.datasets.jpeg_dataset_deprecated import MetaCatnetDataset 9 | from IMDLBenCo.registry import DATASETS 10 | import torch 11 | 12 | from IMDLBenCo.evaluation import generate_region_mask, cal_confusion_matrix, cal_F1 13 | from IMDLBenCo.datasets import denormalize 14 | 15 | 16 | data = IMDLBenCo.datasets.ManiDataset("/mnt/data0/public_datasets/IML/basic_eval_dataset", is_padding=True, edge_width= 7) 17 | # ['/mnt/data0/public_datasets/IML/IMD_20_1024', MetaCatnetDataset], 18 | # ['/mnt/data0/public_datasets/IML/tampCOCO/sp_COCO_list.json', MetaCatnetDataset], 19 | # ['/mnt/data0/public_datasets/IML/tampCOCO/cm_COCO_list.json', MetaCatnetDataset], 20 | # ['/mnt/data0/public_datasets/IML/tampCOCO/bcm_COCO_list.json', MetaCatnetDataset], 21 | # ['/mnt/data0/public_datasets/IML/tampCOCO/bcmc_COCO_list.json', MetaCatnetDataset] 22 | # data = MetaCatnetDataset("/mnt/data0/public_datasets/IML/IMD_20_1024", is_padding=True, edge_width= 7) 23 | 24 | 25 | batch_size = 3 26 | dataloader = DataLoader(data, batch_size=batch_size, shuffle=True) 27 | 28 | for i in dataloader: 29 | for key in i.keys(): 30 | if isinstance(i[key], torch.Tensor) and key != 'shapes': 31 | print(key, i[key].shape) 32 | else: 33 | print(key, i[key]) 34 | 35 | import matplotlib.pyplot as plt 36 | plt.subplot(3, 2, 1) 37 | plt.imshow(i['mask'][0][0]) 38 | plt.subplot(3, 2, 2) 39 | plt.imshow(i['mask'][1][0]) 40 | 41 | # test for evaluation split 42 | # regin_mask = genertate_region_mask(i['masks'], i['shapes']) 43 | plt.subplot(3, 2, 3) 44 | plt.title('shape_mask') 45 | plt.imshow(i['shape_mask'][0][0]) 46 | plt.subplot(3, 2, 4) 47 | plt.imshow(i['shape_mask'][1][0]) 48 | 49 | plt.subplot(3, 2, 5) 50 | plt.imshow(i['image'][0][0]) 51 | plt.subplot(3, 2, 6) 52 | plt.imshow(i['image'][1][0]) 53 | 54 | 55 | 56 | 57 | plt.savefig("test.png") 58 | exit() 59 | 60 | 61 | 62 | -------------------------------------------------------------------------------- /tests/pytests/test_F1.py: -------------------------------------------------------------------------------- 1 | import os 2 | import pytest 3 | import torch 4 | import torch.distributed as dist 5 | import torch.multiprocessing as mp 6 | from torch.utils.data import Dataset, DataLoader 7 | from torch.utils.data.distributed import DistributedSampler 8 | 9 | 10 | class DummyImageMaskDataset(Dataset): 11 | def __init__(self): 12 | # 假设每张图像是 3x32x32，mask 是 1x32x32 13 | self.images = [torch.rand(3, 32, 32) for _ in range(12)] 14 | self.masks = [torch.randint(0, 2, (1, 32, 32)) for _ in range(12)] 15 | 16 | def __len__(self): 17 | return len(self.images) 18 | 19 | def __getitem__(self, idx): 20 | return self.images[idx], self.masks[idx] 21 | 22 | @pytest.fixture(scope="session") 23 | def dummy_dataset(): 24 | return DummyImageMaskDataset() 25 | 26 | class DummyDataset(Dataset): 27 | def __init__(self): 28 | self.data = list(range(12)) # 100 个样本 29 | 30 | def __len__(self): 31 | return len(self.data) 32 | 33 | def __getitem__(self, idx): 34 | return torch.tensor([self.data[idx]], dtype=torch.float32) 35 | 36 | def setup(rank, world_size): 37 | os.environ['MASTER_ADDR'] = 'localhost' 38 | os.environ['MASTER_PORT'] = '12355' 39 | dist.init_process_group("gloo", rank=rank, world_size=world_size) 40 | 41 | def cleanup(): 42 | dist.destroy_process_group() 43 | 44 | def demo(rank, world_size): 45 | print(f"Rank {rank} starting.") 46 | setup(rank, world_size) 47 | 48 | # 构造 dataset 和 distributed sampler 49 | dataset = DummyDataset() 50 | sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=False) 51 | dataloader = DataLoader(dataset, batch_size=10, sampler=sampler) 52 | 53 | total = torch.tensor([0.0]) 54 | 55 | # 迭代数据并求和 56 | for batch in dataloader: 57 | print(f"Rank {rank} batch: {batch}") 58 | total += batch.sum() 59 | 60 | print(f"Rank {rank} local sum before reduce: {total.item()}") 61 | 62 | # 分布式 reduce：将所有 rank 的 total 相加，放到 rank=0 63 | dist.reduce(total, dst=0, op=dist.ReduceOp.SUM) 64 | 65 | if rank == 0: 66 | print(f"Rank {rank} total sum after reduce: {total.item()}") 67 | 68 | cleanup() 69 | 70 | if __name__ == "__main__": 71 | world_size = 4 72 | mp.spawn(demo, args=(world_size,), nprocs=world_size, join=True) 73 | -------------------------------------------------------------------------------- /IMDLBenCo/cli_funcs/cli_env.py: -------------------------------------------------------------------------------- 1 | import os 2 | import platform 3 | from colorama import init, Fore, Style 4 | from IMDLBenCo import __version__ 5 | 6 | import torch 7 | import timm 8 | import albumentations 9 | 10 | def is_torch_cuda_available(): 11 | """ 12 | Check if CUDA is available for PyTorch. 13 | """ 14 | return torch.cuda.is_available() 15 | 16 | def get_env_info(): 17 | info = { 18 | "`IMDLBenCo` version": __version__, 19 | "Platform": platform.platform(), 20 | "Python version": platform.python_version(), 21 | "PyTorch version": torch.__version__, 22 | "Torchvision version": torch.__version__, 23 | "TIMM version": timm.__version__, 24 | "Albumentations version": albumentations.__version__, 25 | } 26 | if is_torch_cuda_available(): 27 | info["PyTorch version"] += " (GPU)" 28 | info["GPU type"] = torch.cuda.get_device_name() 29 | info["GPU number"] = torch.cuda.device_count() 30 | info["GPU memory"] = f"{torch.cuda.mem_get_info()[1] / (1024**3):.2f}GB" 31 | 32 | 33 | try: 34 | import deepspeed # type: ignore 35 | 36 | info["DeepSpeed version"] = deepspeed.__version__ 37 | except Exception: 38 | pass 39 | 40 | try: 41 | import bitsandbytes # type: ignore 42 | 43 | info["Bitsandbytes version"] = bitsandbytes.__version__ 44 | except Exception: 45 | pass 46 | 47 | try: 48 | import vllm 49 | 50 | info["vLLM version"] = vllm.__version__ 51 | except Exception: 52 | pass 53 | try: 54 | import subprocess 55 | # get the dir of imdlbenco package 56 | imdlbenco_dir = os.path.dirname(os.path.abspath(__file__)) 57 | # move to this dir and get the git commit hash in a subprocess 58 | # but don't change the current working directory 59 | os.chdir(imdlbenco_dir) 60 | commit_info = subprocess.run(["git", "rev-parse", "HEAD"], capture_output=True, text=True, check=True) 61 | commit_hash = commit_info.stdout.strip() 62 | info["Git commit"] = commit_hash 63 | except Exception: 64 | pass 65 | 66 | print("\n" + "\n".join([f"- {key}: {value}" for key, value in info.items()]) + "\n") 67 | 68 | 69 | def cli_env(config): 70 | get_env_info() 71 | 72 | if __name__ == "__main__": 73 | print(get_env_info()) -------------------------------------------------------------------------------- /IMDLBenCo/cli_funcs/copy_funcs.py: -------------------------------------------------------------------------------- 1 | import shutil 2 | from colorama import init, Fore, Style 3 | from pathlib import Path 4 | 5 | def copy_file(source: Path, destination: Path): 6 | """ 7 | Copy a single file with path, with checking and asking about the existence. 8 | """ 9 | if not source.is_file(): 10 | print(f"Error: {source} is not a file.") 11 | return 12 | 13 | if destination.exists(): 14 | print(f"Warning: {destination.name} already exists in {destination.parent}.") 15 | user_input = input(f"Do you want to overwrite {destination.name}? (y/n): ").strip().lower() 16 | if user_input != 'y': 17 | print(f"Skipping {destination.name}.") 18 | return 19 | 20 | shutil.copy(source, destination) 21 | print(f"Copied {source.name} to {destination}.") 22 | 23 | def copy_files(source: Path, destination): 24 | """ 25 | Copy files under a path without recursion, with checking and asking about the existence. 26 | """ 27 | yes_to_all, none_to_all=False, False 28 | 29 | for template_file in source.iterdir(): 30 | if template_file.is_file(): 31 | if template_file.name == "__init__.py": 32 | continue # skip __init__.py files 33 | destination_file = Path(destination) / template_file.name 34 | if destination_file.exists(): 35 | if none_to_all: 36 | print(f' Skipping {template_file.name}.\n') 37 | continue 38 | if not yes_to_all: 39 | # Alert , whether overwrite? 40 | print(f' {Fore.YELLOW}Warning: {template_file.name} already exists in {destination}.{Style.RESET_ALL}') 41 | user_input = input(f' Do you want to overwrite {template_file.name}? (y/n/all/none): ').strip().lower() 42 | if user_input == 'all': 43 | yes_to_all = True 44 | elif user_input == 'none': 45 | none_to_all = True 46 | print(f' Skipping {template_file.name}.\n') 47 | continue 48 | elif user_input != 'y': 49 | print(f' Skipping {template_file.name}.\n') 50 | continue 51 | shutil.copy(template_file, destination_file) 52 | print(f' Copied {template_file.name} to {destination}.\n') -------------------------------------------------------------------------------- /IMDLBenCo/model_zoo/trufor/cmx/utils/init_func.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # encoding: utf-8 3 | # @Time : 2018/9/28 下午12:13 4 | # @Author : yuchangqian 5 | # @Contact : changqian_yu@163.com 6 | # @File : init_func.py.py 7 | import torch 8 | import torch.nn as nn 9 | 10 | def __init_weight(feature, conv_init, norm_layer, bn_eps, bn_momentum, 11 | **kwargs): 12 | for name, m in feature.named_modules(): 13 | if isinstance(m, (nn.Conv1d, nn.Conv2d, nn.Conv3d)): 14 | conv_init(m.weight, **kwargs) 15 | elif isinstance(m, norm_layer): 16 | m.eps = bn_eps 17 | m.momentum = bn_momentum 18 | nn.init.constant_(m.weight, 1) 19 | nn.init.constant_(m.bias, 0) 20 | 21 | 22 | def init_weight(module_list, conv_init, norm_layer, bn_eps, bn_momentum, 23 | **kwargs): 24 | if isinstance(module_list, list): 25 | for feature in module_list: 26 | __init_weight(feature, conv_init, norm_layer, bn_eps, bn_momentum, 27 | **kwargs) 28 | else: 29 | __init_weight(module_list, conv_init, norm_layer, bn_eps, bn_momentum, 30 | **kwargs) 31 | 32 | 33 | def group_weight(weight_group, module, norm_layer, lr): 34 | group_decay = [] 35 | group_no_decay = [] 36 | count = 0 37 | for m in module.modules(): 38 | if isinstance(m, nn.Linear): 39 | group_decay.append(m.weight) 40 | if m.bias is not None: 41 | group_no_decay.append(m.bias) 42 | elif isinstance(m, (nn.Conv1d, nn.Conv2d, nn.Conv3d, nn.ConvTranspose2d, nn.ConvTranspose3d)): 43 | group_decay.append(m.weight) 44 | if m.bias is not None: 45 | group_no_decay.append(m.bias) 46 | elif isinstance(m, norm_layer) or isinstance(m, nn.BatchNorm1d) or isinstance(m, nn.BatchNorm2d) \ 47 | or isinstance(m, nn.BatchNorm3d) or isinstance(m, nn.GroupNorm) or isinstance(m, nn.LayerNorm): 48 | if m.weight is not None: 49 | group_no_decay.append(m.weight) 50 | if m.bias is not None: 51 | group_no_decay.append(m.bias) 52 | elif isinstance(m, nn.Parameter): 53 | group_decay.append(m) 54 | 55 | assert len(list(module.parameters())) >= len(group_decay) + len(group_no_decay) 56 | weight_group.append(dict(params=group_decay, lr=lr)) 57 | weight_group.append(dict(params=group_no_decay, weight_decay=.0, lr=lr)) 58 | return weight_group -------------------------------------------------------------------------------- /IMDLBenCo/model_zoo/mesorch/extractor/high_frequency_feature_extraction.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.fft 4 | import math 5 | class HighDctFrequencyExtractor(nn.Module): 6 | def __init__(self, alpha=0.05): 7 | super(HighDctFrequencyExtractor, self).__init__() 8 | if alpha <= 0 or alpha >= 1: 9 | raise ValueError("alpha must be between 0 and 1 (exclusive)") 10 | self.alpha = alpha 11 | self.dct_matrix_h = None 12 | self.dct_matrix_w = None 13 | 14 | def create_dct_matrix(self, N): 15 | n = torch.arange(N, dtype=torch.float32).reshape((1, N)) 16 | k = torch.arange(N, dtype=torch.float32).reshape((N, 1)) 17 | dct_matrix = torch.sqrt(torch.tensor(2.0 / N)) * torch.cos(math.pi * k * (2 * n + 1) / (2 * N)) 18 | dct_matrix[0, :] = 1 / math.sqrt(N) 19 | return dct_matrix 20 | 21 | def dct_2d(self, x): 22 | H, W = x.size(-2), x.size(-1) 23 | if self.dct_matrix_h is None or self.dct_matrix_h.size(0) != H: 24 | self.dct_matrix_h = self.create_dct_matrix(H).to(x.device) 25 | if self.dct_matrix_w is None or self.dct_matrix_w.size(0) != W: 26 | self.dct_matrix_w = self.create_dct_matrix(W).to(x.device) 27 | 28 | return torch.matmul(self.dct_matrix_h, torch.matmul(x, self.dct_matrix_w.t())) 29 | 30 | def idct_2d(self, x): 31 | H, W = x.size(-2), x.size(-1) 32 | if self.dct_matrix_h is None or self.dct_matrix_h.size(0) != H: 33 | self.dct_matrix_h = self.create_dct_matrix(H).to(x.device) 34 | if self.dct_matrix_w is None or self.dct_matrix_w.size(0) != W: 35 | self.dct_matrix_w = self.create_dct_matrix(W).to(x.device) 36 | 37 | return torch.matmul(self.dct_matrix_h.t(), torch.matmul(x, self.dct_matrix_w)) 38 | 39 | def high_pass_filter(self, x, alpha): 40 | h, w = x.shape[-2:] 41 | mask = torch.ones(h, w, device=x.device) 42 | alpha_h, alpha_w = int(alpha * h), int(alpha * w) 43 | mask[:alpha_h, :alpha_w] = 0 44 | 45 | return x * mask 46 | 47 | def forward(self, x): 48 | xq = self.dct_2d(x) 49 | xq_high = self.high_pass_filter(xq, self.alpha) 50 | xh = self.idct_2d(xq_high) 51 | B = xh.shape[0] 52 | min_vals = xh.reshape(B, -1).min(dim=1, keepdim=True).values.view(B, 1, 1, 1) 53 | max_vals = xh.reshape(B, -1).max(dim=1, keepdim=True).values.view(B, 1, 1, 1) 54 | xh = (xh - min_vals) / (max_vals - min_vals) 55 | return xh -------------------------------------------------------------------------------- /IMDLBenCo/model_zoo/mesorch/extractor/low_frequency_feature_extraction.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.fft 4 | import math 5 | class LowDctFrequencyExtractor(nn.Module): 6 | def __init__(self, alpha=0.95): 7 | super(LowDctFrequencyExtractor, self).__init__() 8 | if alpha <= 0 or alpha >= 1: 9 | raise ValueError("alpha must be between 0 and 1 (exclusive)") 10 | self.alpha = alpha 11 | self.dct_matrix_h = None 12 | self.dct_matrix_w = None 13 | 14 | def create_dct_matrix(self, N): 15 | n = torch.arange(N, dtype=torch.float32).reshape((1, N)) 16 | k = torch.arange(N, dtype=torch.float32).reshape((N, 1)) 17 | dct_matrix = torch.sqrt(torch.tensor(2.0 / N)) * torch.cos(math.pi * k * (2 * n + 1) / (2 * N)) 18 | dct_matrix[0, :] = 1 / math.sqrt(N) 19 | return dct_matrix 20 | 21 | def dct_2d(self, x): 22 | H, W = x.size(-2), x.size(-1) 23 | if self.dct_matrix_h is None or self.dct_matrix_h.size(0) != H: 24 | self.dct_matrix_h = self.create_dct_matrix(H).to(x.device) 25 | if self.dct_matrix_w is None or self.dct_matrix_w.size(0) != W: 26 | self.dct_matrix_w = self.create_dct_matrix(W).to(x.device) 27 | 28 | return torch.matmul(self.dct_matrix_h, torch.matmul(x, self.dct_matrix_w.t())) 29 | 30 | def idct_2d(self, x): 31 | H, W = x.size(-2), x.size(-1) 32 | if self.dct_matrix_h is None or self.dct_matrix_h.size(0) != H: 33 | self.dct_matrix_h = self.create_dct_matrix(H).to(x.device) 34 | if self.dct_matrix_w is None or self.dct_matrix_w.size(0) != W: 35 | self.dct_matrix_w = self.create_dct_matrix(W).to(x.device) 36 | 37 | return torch.matmul(self.dct_matrix_h.t(), torch.matmul(x, self.dct_matrix_w)) 38 | 39 | def high_pass_filter(self, x, alpha): 40 | h, w = x.shape[-2:] 41 | mask = torch.ones(h, w, device=x.device) 42 | alpha_h, alpha_w = int(alpha * h), int(alpha * w) 43 | mask[-alpha_h:, -alpha_w:] = 0 44 | 45 | return x * mask 46 | 47 | def forward(self, x): 48 | xq = self.dct_2d(x) 49 | xq_high = self.high_pass_filter(xq, self.alpha) 50 | xh = self.idct_2d(xq_high) 51 | B = xh.shape[0] 52 | min_vals = xh.reshape(B, -1).min(dim=1, keepdim=True).values.view(B, 1, 1, 1) 53 | max_vals = xh.reshape(B, -1).max(dim=1, keepdim=True).values.view(B, 1, 1, 1) 54 | xh = (xh - min_vals) / (max_vals - min_vals) 55 | return xh -------------------------------------------------------------------------------- /tests/test_appdirs.py: -------------------------------------------------------------------------------- 1 | import appdirs 2 | import os 3 | """ expected outputs: 4 | 5 | User data directory: /home/xiaochen/.local/share/MyApp 6 | User config directory: /home/xiaochen/.config/MyApp 7 | User cache directory: /home/xiaochen/.cache/MyApp 8 | User log directory: /home/xiaochen/.cache/MyApp/log 9 | Site data directory: /usr/local/share/MyApp 10 | Traceback (most recent call last): 11 | File "/mnt/data0/xiaochen/workspace/IMDLBenCo_pure/IMDLBenCo/tests/test_appdirs.py", line 49, in 12 | test_appdirs() 13 | File "/mnt/data0/xiaochen/workspace/IMDLBenCo_pure/IMDLBenCo/tests/test_appdirs.py", line 37, in test_appdirs 14 | os.makedirs(site_data, exist_ok=True) 15 | File "", line 225, in makedirs 16 | PermissionError: [Errno 13] Permission denied: '/usr/local/share/MyApp' 17 | 18 | """ 19 | def test_appdirs(): 20 | app_name = "MyApp" 21 | app_author = "MyCompany" 22 | 23 | print("Testing appdirs functionalities...\n") 24 | 25 | # User data directory 26 | user_data = appdirs.user_data_dir(app_name, app_author) 27 | print(f"User data directory: {user_data}") 28 | os.makedirs(user_data, exist_ok=True) 29 | assert os.path.exists(user_data) 30 | 31 | # User config directory 32 | user_config = appdirs.user_config_dir(app_name, app_author) 33 | print(f"User config directory: {user_config}") 34 | os.makedirs(user_config, exist_ok=True) 35 | assert os.path.exists(user_config) 36 | 37 | # User cache directory 38 | user_cache = appdirs.user_cache_dir(app_name, app_author) 39 | print(f"User cache directory: {user_cache}") 40 | os.makedirs(user_cache, exist_ok=True) 41 | assert os.path.exists(user_cache) 42 | 43 | # User log directory 44 | user_log = appdirs.user_log_dir(app_name, app_author) 45 | print(f"User log directory: {user_log}") 46 | os.makedirs(user_log, exist_ok=True) 47 | assert os.path.exists(user_log) 48 | 49 | # # Site data directory 50 | # site_data = appdirs.site_data_dir(app_name, app_author) 51 | # print(f"Site data directory: {site_data}") 52 | # os.makedirs(site_data, exist_ok=True) 53 | # assert os.path.exists(site_data) 54 | 55 | # # Site config directory 56 | # site_config = appdirs.site_config_dir(app_name, app_author) 57 | # print(f"Site config directory: {site_config}") 58 | # os.makedirs(site_config, exist_ok=True) 59 | # assert os.path.exists(site_config) 60 | 61 | print("\nAll appdirs functionalities have been tested successfully.") 62 | 63 | if __name__ == "__main__": 64 | test_appdirs() 65 | -------------------------------------------------------------------------------- /tests/test_metrics/generate_dataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | import json 3 | import random 4 | import numpy as np 5 | from PIL import Image 6 | 7 | # 配置参数 8 | IMAGE_SIZE = (512, 512) # (height, width) 9 | NUM_IMAGES = 20 10 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 11 | 12 | # 创建目录 13 | os.makedirs(os.path.join(BASE_DIR, "images"), exist_ok=True) 14 | os.makedirs(os.path.join(BASE_DIR, "masks"), exist_ok=True) 15 | 16 | def generate_random_rectangle(img_size): 17 | """生成随机位置和大小的矩形区域（在图像中心附近）""" 18 | h, w = img_size 19 | # 随机中心点偏移范围 20 | max_offset = min(h, w) // 6 21 | cx = w//2 + random.randint(-max_offset, max_offset) 22 | cy = h//2 + random.randint(-max_offset, max_offset) 23 | 24 | # 随机尺寸 25 | rect_w = random.randint(min(h, w)//4, min(h, w)//2) 26 | rect_h = random.randint(min(h, w)//4, min(h, w)//2) 27 | 28 | # 计算边界 29 | x1 = max(0, cx - rect_w//2) 30 | y1 = max(0, cy - rect_h//2) 31 | x2 = min(w, cx + rect_w//2) 32 | y2 = min(h, cy + rect_h//2) 33 | return (x1, y1, x2, y2) 34 | 35 | def generate_mask_with_blocks(img_size): 36 | """生成带随机黑块的噪声mask""" 37 | # 生成基础噪声 38 | mask = np.random.choice([0, 255], img_size, p=[0.5, 0.5]).astype(np.uint8) 39 | 40 | # 添加随机黑块 41 | for _ in range(random.randint(1, 3)): # 1-3个黑块 42 | block_w = random.randint(20, 80) 43 | block_h = random.randint(20, 80) 44 | x = random.randint(0, img_size[1] - block_w) 45 | y = random.randint(0, img_size[0] - block_h) 46 | mask[y:y+block_h, x:x+block_w] = 0 47 | return mask 48 | 49 | # 生成数据集 50 | dataset = [] 51 | 52 | for idx in range(NUM_IMAGES): 53 | # 生成基础噪声图像 54 | img = np.random.randint(0, 256, (*IMAGE_SIZE, 3), dtype=np.uint8) 55 | 56 | # 添加随机矩形区域 57 | rect = generate_random_rectangle(IMAGE_SIZE) 58 | img[rect[1]:rect[3], rect[0]:rect[2], :] = 0 59 | 60 | # 保存图像 61 | img_name = f"{idx+1:02d}.jpg" 62 | img_path = os.path.join("images", img_name) 63 | Image.fromarray(img).save(os.path.join(BASE_DIR, img_path)) 64 | 65 | # 处理前10张图的mask 66 | if idx < 10: 67 | mask = generate_mask_with_blocks(IMAGE_SIZE) 68 | mask_name = f"{idx+1:02d}_mask.jpg" 69 | mask_path = os.path.join("masks", mask_name) 70 | Image.fromarray(mask).save(os.path.join(BASE_DIR, mask_path)) 71 | dataset.append([img_path, mask_path]) 72 | else: 73 | dataset.append([img_path, "Negative"]) 74 | 75 | # 保存JSON文件 76 | with open(os.path.join(BASE_DIR, "dataset.json"), "w") as f: 77 | json.dump(dataset, f, indent=2) 78 | 79 | print("数据集生成完成！") -------------------------------------------------------------------------------- /.github/workflows/python-publish.yml: -------------------------------------------------------------------------------- 1 | # This workflow will upload a Python Package to PyPI when a release is created 2 | # For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python#publishing-to-package-registries 3 | 4 | # This workflow uses actions that are not certified by GitHub. 5 | # They are provided by a third-party and are governed by 6 | # separate terms of service, privacy policy, and support 7 | # documentation. 8 | 9 | name: Upload Python Package 10 | 11 | on: 12 | release: 13 | types: [published] 14 | 15 | permissions: 16 | contents: read 17 | 18 | jobs: 19 | release-build: 20 | runs-on: ubuntu-latest 21 | 22 | steps: 23 | - uses: actions/checkout@v4 24 | 25 | - uses: actions/setup-python@v5 26 | with: 27 | python-version: "3.12" 28 | 29 | - name: Build release distributions 30 | run: | 31 | python -m pip install --upgrade pip 32 | pip install -r requirements.txt 33 | pip install build twine 34 | python -m build --sdist --wheel 35 | 36 | - name: Verify distribution 37 | run: | 38 | twine check dist/* 39 | 40 | - name: Upload distributions 41 | uses: actions/upload-artifact@v4 42 | with: 43 | name: release-dists 44 | path: dist/ 45 | 46 | pypi-publish: 47 | runs-on: ubuntu-latest 48 | needs: 49 | - release-build 50 | permissions: 51 | # IMPORTANT: this permission is mandatory for trusted publishing 52 | id-token: write 53 | 54 | # Dedicated environments with protections for publishing are strongly recommended. 55 | # For more information, see: https://docs.github.com/en/actions/deployment/targeting-different-environments/using-environments-for-deployment#deployment-protection-rules 56 | environment: 57 | name: pypi 58 | # OPTIONAL: uncomment and update to include your PyPI project URL in the deployment status: 59 | # url: https://pypi.org/p/YOURPROJECT 60 | # 61 | # ALTERNATIVE: if your GitHub Release name is the PyPI project version string 62 | # ALTERNATIVE: exactly, uncomment the following line instead: 63 | # url: https://pypi.org/project/YOURPROJECT/${{ github.event.release.name }} 64 | 65 | steps: 66 | - name: Retrieve release distributions 67 | uses: actions/download-artifact@v4 68 | with: 69 | name: release-dists 70 | path: dist/ 71 | 72 | - name: Publish release distributions to PyPI 73 | uses: pypa/gh-action-pypi-publish@release/v1 74 | with: 75 | packages-dir: dist/ 76 | -------------------------------------------------------------------------------- /tests/test_metrics/mymodel.py: -------------------------------------------------------------------------------- 1 | from IMDLBenCo.registry import MODELS 2 | import torch.nn as nn 3 | import torch 4 | import numpy as np 5 | 6 | """ 7 | === Image-level Metrics === 8 | F1 Score: 0.4211 9 | AUC Score: 0.4300 10 | Accuracy: 0.4500 11 | 12 | === Pixel-level Metrics === 13 | F1 Score: 0.1983 14 | Accuracy: 0.5892 15 | """ 16 | 17 | @MODELS.register_module() 18 | class MyModel(nn.Module): 19 | def __init__(self): 20 | super(MyModel, self).__init__() 21 | # 放一个虚假的参数用于欺骗DDP 22 | self.param = nn.Parameter(torch.randn(1)) 23 | def forward(self, image : torch.Tensor, mask : torch.Tensor, name, **kwargs): 24 | """ 25 | 对于固定的name,测试数据集返回固定的label确保可以复现： 26 | name=["01.jpg", "02.jpg", "03.jpg", "04.jpg"... "20.jpg"], 27 | label是固定的0~1的数值 28 | """ 29 | # 将整个[b,c,h,w]的tensor拍扁成1维 30 | pred_mask = torch.mean(image, dim=1, keepdim=True) 31 | # 利用图片的左上角的像素值来保存该图片的label预测结果。小trick 32 | 33 | # 转回numpy并保存这个mask 34 | pred_mask = pred_mask.clip(0, 1) 35 | # 尽可能保证向后计算指标的图，与保存的图像一致（不会被int卡掉太多精度） 36 | pred_mask = (pred_mask * 255).to(torch.uint8).to(torch.float32) / 255 37 | pred_tensor_mask = pred_mask.clone() 38 | pred_mask = pred_mask.cpu().numpy() 39 | # print(pred_mask.shape) 40 | pred_mask = np.uint8(pred_mask * 255) 41 | 42 | # 计算一个mse损失，没啥用 43 | loss = (image - mask) ** 2 44 | 45 | pred_label_dict = { 46 | "01.jpg": 0.1, 47 | "02.jpg": 0.2, 48 | "03.jpg": 0.3, 49 | "04.jpg": 0.4, 50 | "05.jpg": 0.5, 51 | "06.jpg": 0.6, 52 | "07.jpg": 0.7, 53 | "08.jpg": 0.1, 54 | "09.jpg": 0.9, 55 | "10.jpg": 1.0, 56 | "11.jpg": 0.1, 57 | "12.jpg": 0.2, 58 | "13.jpg": 0.3, 59 | "14.jpg": 0.4, 60 | "15.jpg": 0.5, 61 | "16.jpg": 0.6, 62 | "17.jpg": 0.7, 63 | "18.jpg": 0.8, 64 | "19.jpg": 0.9, 65 | "20.jpg": 1.0, 66 | } 67 | # 用name做映射，然后返回0~1的 tensor作为label 68 | pred_label = [] 69 | for i in range(len(name)): 70 | pred_label.append(pred_label_dict[name[i]]) 71 | pred_label = torch.tensor(pred_label).to(image.device) 72 | output_dict = { 73 | "backward_loss": loss, 74 | # predicted mask, will calculate for metrics automatically 75 | "pred_mask": pred_tensor_mask, 76 | # predicted binaray label, will calculate for metrics automatically 77 | "pred_label": pred_label, 78 | "visual_loss": { 79 | }, 80 | "visual_image": { 81 | } 82 | } 83 | return output_dict 84 | 85 | if __name__ == "__main__": 86 | print(MODELS) -------------------------------------------------------------------------------- /IMDLBenCo/modules/extractors/sobel.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torchvision.transforms as transforms 4 | import matplotlib.pyplot as plt 5 | import numpy as np 6 | import inspect 7 | class SobelFilter(nn.Module): 8 | def __init__(self,in_chan=3, out_chan=1,norm = None ): 9 | super(SobelFilter, self).__init__() 10 | if norm is not None and (not inspect.isclass(norm) or not issubclass(norm, nn.Module)): 11 | raise ValueError("norm must be a class derived from nn.Module or None") 12 | 13 | filter_x = np.array([ 14 | [1, 0, -1], 15 | [2, 0, -2], 16 | [1, 0, -1], 17 | ]).astype(np.float32) 18 | filter_y = np.array([ 19 | [1, 2, 1], 20 | [0, 0, 0], 21 | [-1, -2, -1], 22 | ]).astype(np.float32) 23 | 24 | filter_x = filter_x.reshape((1, 1, 3, 3)) 25 | filter_x = np.repeat(filter_x, in_chan, axis=1) 26 | filter_x = np.repeat(filter_x, out_chan, axis=0) 27 | 28 | filter_y = filter_y.reshape((1, 1, 3, 3)) 29 | filter_y = np.repeat(filter_y, in_chan, axis=1) 30 | filter_y = np.repeat(filter_y, out_chan, axis=0) 31 | 32 | filter_x = torch.from_numpy(filter_x) 33 | filter_y = torch.from_numpy(filter_y) 34 | filter_x = nn.Parameter(filter_x, requires_grad=False) 35 | filter_y = nn.Parameter(filter_y, requires_grad=False) 36 | 37 | sobel_kernel_x = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32).unsqueeze(0).unsqueeze(0) 38 | sobel_kernel_y = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=torch.float32).unsqueeze(0).unsqueeze(0) 39 | conv_x = nn.Conv2d(in_chan, out_chan, kernel_size=3, stride=1, padding=1, bias=False) 40 | conv_x.weight = filter_x 41 | conv_y = nn.Conv2d(in_chan, out_chan, kernel_size=3, stride=1, padding=1, bias=False) 42 | conv_y.weight = filter_y 43 | norm1 = nn.Identity() if norm is None else norm(out_chan) 44 | norm2 = nn.Identity() if norm is None else norm(out_chan) 45 | 46 | self.sobel_x = nn.Sequential(conv_x, norm1) 47 | self.sobel_y = nn.Sequential(conv_y, norm2) 48 | 49 | 50 | self.sobel_kernel_x = sobel_kernel_x.expand(1, 3, 3, 3) 51 | self.sobel_kernel_y = sobel_kernel_y.expand(1, 3, 3, 3) 52 | 53 | self.conv_x = nn.Conv2d(in_channels=3, out_channels=1, kernel_size=3, padding=1, bias=False) 54 | self.conv_y = nn.Conv2d(in_channels=3, out_channels=1, kernel_size=3, padding=1, bias=False) 55 | 56 | 57 | self.conv_x.weight = nn.Parameter(self.sobel_kernel_x, requires_grad=False) 58 | self.conv_y.weight = nn.Parameter(self.sobel_kernel_y, requires_grad=False) 59 | 60 | def forward(self, x): 61 | grad_x = self.conv_x(x) 62 | grad_y = self.conv_y(x) 63 | grad = torch.sqrt(grad_x**2 + grad_y**2) 64 | return grad 65 | 66 | -------------------------------------------------------------------------------- /IMDLBenCo/model_zoo/cat_net/cat_net.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from .network_CAT import get_seg_model, CrossEntropy 4 | 5 | from yacs.config import CfgNode 6 | from torch.nn import functional as F 7 | import yaml 8 | import random 9 | 10 | from IMDLBenCo.registry import MODELS 11 | 12 | @MODELS.register_module() 13 | class Cat_Net(nn.Module): 14 | """ 15 | Distribute the loss on multi-gpu to reduce 16 | the memory cost in the main gpu. 17 | You can check the following discussion. 18 | https://discuss.pytorch.org/t/dataparallel-imbalanced-memory-usage/22551/21 19 | """ 20 | 21 | def __init__(self, cfg_file): 22 | super(Cat_Net, self).__init__() 23 | cfg = None 24 | with open(cfg_file, "r") as f: 25 | yaml_cfg = yaml.safe_load(f) 26 | cfg = CfgNode(yaml_cfg) 27 | self.model = get_seg_model(cfg) 28 | self.loss = CrossEntropy(ignore_label=cfg.TRAIN.IGNORE_LABEL, weight=torch.FloatTensor([0.5, 2.5])).cuda() 29 | 30 | 31 | def forward_features(self, image, mask, DCT_coef, qtables, **kwargs): 32 | images, masks = self.__post_process_tensor(image, mask, DCT_coef) 33 | images, masks = images.detach(), masks.squeeze(1).detach() 34 | qtables = qtables.unsqueeze(1) 35 | outputs = self.model(images.float(), qtables.float()) 36 | return outputs, masks 37 | 38 | 39 | def forward(self, image, mask, DCT_coef, qtables, if_predcit_label=None, edge_mask=None, shape=None, *args, **kwargs): 40 | 41 | outputs, masks = self.forward_features(image, mask, DCT_coef, qtables) 42 | loss = self.loss(outputs, masks.long()) 43 | 44 | pred = F.softmax(outputs, dim=1)[:, 1].unsqueeze(1) 45 | pred = F.interpolate(pred, size=(image.shape[2], image.shape[3]), mode='bicubic') 46 | output_dict = { 47 | # loss for backward 48 | "backward_loss": loss, 49 | # predicted mask, will calculate for metrics automatically 50 | "pred_mask": pred, 51 | "pred_label": None, 52 | "visual_loss": { 53 | "predict_loss": loss, 54 | }, 55 | 56 | "visual_image": { 57 | "pred_mask": pred, 58 | } 59 | 60 | } 61 | return output_dict 62 | 63 | def __post_process_tensor(self, image_tensor, mask, DCT_coef): 64 | ignore_index = -1 65 | crop_size = (512, 512) 66 | img_RGB = torch.permute(image_tensor, (0, 2, 3, 1)) # batch_size * h * w * 3 67 | h, w = img_RGB.shape[1], img_RGB.shape[2] 68 | t_RGB = (torch.permute(img_RGB, (0, 3, 1, 2)) - 127.5) / 127.5 69 | 70 | T = 20 71 | t_DCT_vol = torch.zeros(size=(DCT_coef.shape[0], T+1, DCT_coef.shape[1], DCT_coef.shape[2])).cuda() 72 | t_DCT_vol[:, 0] += (DCT_coef.cuda() == 0).float() 73 | 74 | for i in range(1, T): 75 | t_DCT_vol[:, i] += (DCT_coef == i).float() 76 | t_DCT_vol[:, i] += (DCT_coef == -i).float() 77 | t_DCT_vol[:, T] += (DCT_coef >= T).float() 78 | t_DCT_vol[:, T] += (DCT_coef <= -T).float() 79 | tensor = torch.cat([t_RGB, t_DCT_vol], dim=1) 80 | return tensor, mask.long() 81 | 82 | 83 | 84 | -------------------------------------------------------------------------------- /setup_backup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup, find_packages 2 | 3 | from IMDLBenCo.version import __version__ 4 | 5 | 6 | """ 7 | For documentations, see here: 8 | https://github.com/pypa/sampleproject/blob/db5806e0a3204034c51b1c00dde7d5eb3fa2532e/setup.py 9 | 10 | 在工作路径下运行下列指令： 11 | pip install -e . 12 | 即可实现本地安装本包，并在更新文件时自动更新相应内容，便于调试开发。 13 | """ 14 | 15 | def readme(): 16 | with open('README.md', encoding='utf-8') as f: 17 | content = f.read() 18 | return content 19 | 20 | def requirements(): 21 | with open('requirements.txt') as f: 22 | requirements = f.read().splitlines() 23 | return requirements 24 | 25 | # exit(0) 26 | 27 | setup( 28 | name='imdlbenco', 29 | version=__version__, 30 | description="A comprehensive benchmark and code base for Image manipulation and localization.", 31 | long_description=readme(), 32 | long_description_content_type='text/markdown', 33 | url="https://github.com/scu-zjz/IMDLBenCo", 34 | author="Xiaochen Ma", # Optional 35 | author_email="xiaochen.ma.cs@gmail.com", # Optional 36 | 37 | 38 | classifiers=[ # Optional 39 | # How mature is this project? Common values are 40 | # 3 - Alpha 41 | # 4 - Beta 42 | # 5 - Production/Stable 43 | "Development Status :: 3 - Alpha", 44 | # Indicate who your project is intended for 45 | "Intended Audience :: Developers", 46 | "Intended Audience :: Science/Research", 47 | 48 | "Topic :: Security", 49 | "Topic :: Scientific/Engineering :: Artificial Intelligence", 50 | 51 | # Pick your license as you wish 52 | "License :: Free For Educational Use", 53 | # Specify the Python versions you support here. In particular, ensure 54 | # that you indicate you support Python 3. These classifiers are *not* 55 | # checked by 'pip install'. See instead 'python_requires' below. 56 | "Programming Language :: Python :: 3", 57 | "Programming Language :: Python :: 3.7", 58 | "Programming Language :: Python :: 3.8", 59 | "Programming Language :: Python :: 3.9", 60 | "Programming Language :: Python :: 3.10", 61 | "Programming Language :: Python :: 3 :: Only", 62 | ], 63 | keywords=[ 64 | "AI", 65 | "artificial intelligence", 66 | "image forensics", 67 | "image manipulation localization", 68 | "image manipulation detection" 69 | ], # Optional 70 | 71 | packages=find_packages(), 72 | python_requires=">=3.7, <4", 73 | 74 | 75 | include_package_data=True, # 76 | install_requires=requirements(), 77 | license='CC-BY-4.0', 78 | entry_points={ 79 | 'console_scripts': [ 80 | 'benco=IMDLBenCo.cli:main', # cli是你的命令行脚本的模块，main是入口函数 81 | ], 82 | }, 83 | project_urls={ # Optional 84 | "Github": "https://github.com/scu-zjz/IMDLBenCo/", 85 | "Documentation":"https://github.com/scu-zjz/IMDLBenCo-doc", 86 | "Bug Reports": "https://github.com/scu-zjz/IMDLBenCo/issues", 87 | }, 88 | ) 89 | 90 | -------------------------------------------------------------------------------- /IMDLBenCo/statics/base/mymodel.py: -------------------------------------------------------------------------------- 1 | from IMDLBenCo.registry import MODELS 2 | import torch.nn as nn 3 | import torch 4 | 5 | @MODELS.register_module() 6 | class MyModel(nn.Module): 7 | def __init__(self, MyModel_Customized_param:int, pre_trained_weights:str) -> None: 8 | """ 9 | The parameters of the `__init__` function will be automatically converted into the parameters expected by the argparser in the training and testing scripts by the framework according to their annotated types and variable names. 10 | 11 | In other words, you can directly pass in parameters with the same names and types from the `run.sh` script to initialize the model. 12 | """ 13 | super().__init__() 14 | 15 | # Useless, just an example 16 | self.MyModel_Customized_param = MyModel_Customized_param 17 | self.pre_trained_weights = pre_trained_weights 18 | 19 | # A single layer conv2d 20 | self.demo_layer = nn.Conv2d( 21 | in_channels=3, 22 | out_channels=1, 23 | kernel_size=3, 24 | stride=1, 25 | padding=1, 26 | ) 27 | 28 | # A simple loss 29 | self.loss_func_a = nn.BCEWithLogitsLoss() 30 | 31 | def forward(self, image, mask, label, *args, **kwargs): 32 | # simple forwarding 33 | pred_mask = self.demo_layer(image) 34 | 35 | # simple loss 36 | loss_a = self.loss_func_a(pred_mask, mask) 37 | loss_b = torch.abs(torch.mean(pred_mask - mask)) 38 | combined_loss = loss_a + loss_b 39 | 40 | 41 | pred_label = torch.mean(pred_mask) 42 | 43 | inverse_mask = 1 - mask 44 | 45 | # ----------Output interface-------------------------------------- 46 | output_dict = { 47 | # loss for backward 48 | "backward_loss": combined_loss, 49 | # predicted mask, will calculate for metrics automatically 50 | "pred_mask": pred_mask, 51 | # predicted binaray label, will calculate for metrics automatically 52 | "pred_label": pred_label, 53 | 54 | # ----values below is for visualization---- 55 | # automatically visualize with the key-value pairs 56 | "visual_loss": { 57 | # keys here can be customized by yourself. 58 | "predict_loss": combined_loss, 59 | 'loss_a' : loss_a, 60 | "I am loss_b :)": loss_b, 61 | }, 62 | 63 | "visual_image": { 64 | # keys here can be customized by yourself. 65 | # Various intermediate masks, heatmaps, and feature maps can be appropriately converted into RGB or single-channel images for visualization here. 66 | "pred_mask": pred_mask, 67 | "reverse_mask" : inverse_mask, 68 | } 69 | # ------------------------------------------------------------- 70 | } 71 | 72 | return output_dict 73 | 74 | 75 | if __name__ == "__main__": 76 | print(MODELS) -------------------------------------------------------------------------------- /IMDLBenCo/statics/model_zoo/configs/CAT_full.yaml: -------------------------------------------------------------------------------- 1 | CUDNN: 2 | BENCHMARK: true 3 | DETERMINISTIC: false 4 | ENABLED: true 5 | GPUS: (0,1) 6 | OUTPUT_DIR: 'output' 7 | LOG_DIR: 'log' 8 | WORKERS: 0 9 | PRINT_FREQ: 10 10 | 11 | DATASET: 12 | DATASET: splicing_dataset 13 | NUM_CLASSES: 2 14 | MODEL: 15 | NAME: network_CAT 16 | PRETRAINED_RGB: '/mnt/data0/bingkui/Cat-Net/hrnetv2_w48_imagenet_pretrained.pth' 17 | PRETRAINED_DCT: '/mnt/data0/bingkui/Cat-Net/DCT_djpeg.pth.tar' 18 | EXTRA: 19 | FINAL_CONV_KERNEL: 1 20 | STAGE1: 21 | NUM_MODULES: 1 22 | NUM_RANCHES: 1 23 | BLOCK: BOTTLENECK 24 | NUM_BLOCKS: 25 | - 4 26 | NUM_CHANNELS: 27 | - 64 28 | FUSE_METHOD: SUM 29 | STAGE2: 30 | NUM_MODULES: 1 31 | NUM_BRANCHES: 2 32 | BLOCK: BASIC 33 | NUM_BLOCKS: 34 | - 4 35 | - 4 36 | NUM_CHANNELS: 37 | - 48 38 | - 96 39 | FUSE_METHOD: SUM 40 | STAGE3: 41 | NUM_MODULES: 4 42 | NUM_BRANCHES: 3 43 | BLOCK: BASIC 44 | NUM_BLOCKS: 45 | - 4 46 | - 4 47 | - 4 48 | NUM_CHANNELS: 49 | - 48 50 | - 96 51 | - 192 52 | FUSE_METHOD: SUM 53 | STAGE4: 54 | NUM_MODULES: 3 55 | NUM_BRANCHES: 4 56 | BLOCK: BASIC 57 | NUM_BLOCKS: 58 | - 4 59 | - 4 60 | - 4 61 | - 4 62 | NUM_CHANNELS: 63 | - 48 64 | - 96 65 | - 192 66 | - 384 67 | FUSE_METHOD: SUM 68 | DC_STAGE3: 69 | NUM_MODULES: 3 # 4 70 | NUM_BRANCHES: 2 71 | BLOCK: BASIC 72 | NUM_BLOCKS: 73 | - 4 74 | - 4 75 | NUM_CHANNELS: 76 | - 96 77 | - 192 78 | FUSE_METHOD: SUM 79 | DC_STAGE4: 80 | NUM_MODULES: 2 # 3 81 | NUM_BRANCHES: 3 82 | BLOCK: BASIC 83 | NUM_BLOCKS: 84 | - 4 85 | - 4 86 | - 4 87 | NUM_CHANNELS: 88 | - 96 89 | - 192 90 | - 384 91 | FUSE_METHOD: SUM 92 | STAGE5: 93 | NUM_MODULES: 1 94 | NUM_BRANCHES: 4 # concat resolution-wise 95 | BLOCK: BASIC 96 | NUM_BLOCKS: 97 | - 4 98 | - 4 99 | - 4 100 | - 4 101 | NUM_CHANNELS: 102 | - 24 103 | - 48 104 | - 96 105 | - 192 106 | FUSE_METHOD: SUM 107 | LOSS: 108 | USE_OHEM: false 109 | OHEMTHRES: 0.9 110 | OHEMKEEP: 131072 111 | TRAIN: 112 | IMAGE_SIZE: 113 | - 512 114 | - 512 115 | BASE_SIZE: 512 116 | BATCH_SIZE_PER_GPU: 2 117 | SHUFFLE: true 118 | BEGIN_EPOCH: 0 119 | END_EPOCH: 200 120 | RESUME: true 121 | OPTIMIZER: sgd 122 | LR: 0.005 123 | WD: 0.0005 124 | MOMENTUM: 0.9 125 | NESTEROV: false 126 | FLIP: true 127 | MULTI_SCALE: true 128 | DOWNSAMPLERATE: 1 129 | IGNORE_LABEL: -1 130 | SCALE_FACTOR: 11 131 | TEST: 132 | IMAGE_SIZE: 133 | - 512 134 | - 512 135 | BASE_SIZE: 512 136 | BATCH_SIZE_PER_GPU: 2 137 | NUM_SAMPLES: 2000 138 | FLIP_TEST: false 139 | MULTI_SCALE: false 140 | -------------------------------------------------------------------------------- /tests/test_metrics/reference.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.utils.data import DataLoader 3 | from sklearn.metrics import f1_score, roc_auc_score, accuracy_score 4 | import numpy as np 5 | from tqdm import tqdm 6 | 7 | from IMDLBenCo.datasets import JsonDataset 8 | from mymodel import MyModel # 确保你已注册了这个模型 9 | 10 | # 1. 加载数据集 11 | dataset = JsonDataset( 12 | path="/mnt/data0/xiaochen/workspace/IMDLBenCo_pure/test_metrics/dataset.json", 13 | is_resizing=True, 14 | output_size=(512, 512), 15 | ) 16 | 17 | # 2. DataLoader 18 | dataloader = DataLoader(dataset, batch_size=1, shuffle=False) 19 | 20 | # 3. 加载模型 21 | model = MyModel() 22 | model.eval() 23 | 24 | # 4. 存储所有结果 25 | all_image_preds = [] 26 | all_image_labels = [] 27 | 28 | pixel_f1_list = [] 29 | pixel_auc_list = [] 30 | pixel_acc_list = [] 31 | 32 | with torch.no_grad(): 33 | for batch in tqdm(dataloader): 34 | image = batch['image'] # [B, C, H, W] 35 | mask = batch['mask'] # [B, 1, H, W] 36 | label = batch['label'] # scalar 37 | name = batch['name'] # list[str] of length B=1 38 | 39 | output = model(image, mask, name) 40 | 41 | # -------- image-level 预测 -------- 42 | pred_label = output['pred_label'].cpu().numpy() # shape: (B,) 43 | true_label = np.array([label]) # shape: (B,) 44 | 45 | all_image_preds.extend(pred_label.tolist()) 46 | all_image_labels.extend(true_label.tolist()) 47 | 48 | # -------- pixel-level 预测 -------- 49 | pred_mask = output['pred_mask'].cpu().numpy() # [B, 1, H, W] 50 | true_mask = mask.cpu().numpy() # [B, 1, H, W] 51 | 52 | # flatten 到 [B, H*W] 53 | pred_flat = pred_mask.reshape(-1) 54 | true_flat = true_mask.reshape(-1) 55 | 56 | pixel_f1 = f1_score(true_flat, np.round(pred_flat), zero_division=1) 57 | # pixel_auc = roc_auc_score(true_flat, pred_flat) 58 | pixel_acc = accuracy_score(true_flat, np.round(pred_flat)) 59 | pixel_f1_list.append(pixel_f1) 60 | # pixel_auc_list.append(pixel_auc) 61 | pixel_acc_list.append(pixel_acc) 62 | 63 | # ---------- 计算 image-level metrics ---------- 64 | print("\n=== Image-level Metrics ===") 65 | img_f1 = f1_score(all_image_labels, np.round(all_image_preds), zero_division=1) 66 | img_auc = roc_auc_score(all_image_labels, all_image_preds) 67 | img_acc = accuracy_score(all_image_labels, np.round(all_image_preds)) 68 | 69 | print(f"F1 Score: {img_f1:.4f}") 70 | print(f"AUC Score: {img_auc:.4f}") 71 | print(f"Accuracy: {img_acc:.4f}") 72 | 73 | # ---------- 计算 pixel-level metrics ---------- 74 | print("\n=== Pixel-level Metrics ===") 75 | # pix_f1 = f1_score(all_pixel_labels, np.round(all_pixel_preds), zero_division=1) 76 | # pix_auc = roc_auc_score(all_pixel_labels, all_pixel_preds) 77 | # pix_acc = accuracy_score(all_pixel_labels, np.round(all_pixel_preds)) 78 | pix_f1 = np.mean(pixel_f1_list) 79 | # pix_auc = np.mean(pixel_auc_list) 80 | pix_acc = np.mean(pixel_acc_list) 81 | 82 | 83 | print(f"F1 Score: {pix_f1:.4f}") 84 | # print(f"AUC Score: {pix_auc:.4f}") 85 | print(f"Accuracy: {pix_acc:.4f}") 86 | -------------------------------------------------------------------------------- /tests/test_tail_dataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch 3 | from torch.utils.data import Dataset, DataLoader 4 | from torch.utils.data.distributed import DistributedSampler 5 | import torch.distributed as dist 6 | 7 | device= torch.device("cuda" if torch.cuda.is_available() else "cpu") 8 | 9 | def distributed_gather_tensor(tensor: torch.Tensor, total_size: int = None): 10 | """ 11 | Gather tensors from all processes. Optionally trim to `total_size` to remove padding from the last batch. 12 | Works only for 1D or 2D tensors (along dim=0). 13 | 14 | Args: 15 | tensor (torch.Tensor): Local tensor on each process 16 | total_size (int, optional): Total valid size after gathering. Used to remove padding. 17 | 18 | Returns: 19 | gathered_tensor (torch.Tensor): Full gathered tensor on all processes. 20 | """ 21 | world_size = dist.get_world_size() 22 | tensor= tensor.to(device) 23 | # Create empty list to hold gathered tensors 24 | gather_list = [torch.zeros_like(tensor) for _ in range(world_size)] 25 | dist.all_gather(gather_list, tensor) 26 | gathered = torch.cat(gather_list, dim=0) 27 | 28 | if total_size is not None: 29 | gathered = gathered[:total_size] 30 | 31 | return gathered 32 | local_rank = int(os.environ["LOCAL_RANK"]) 33 | torch.cuda.set_device(local_rank) # 这一步是关键！！ 34 | device = torch.device("cuda", local_rank) 35 | # 初始化分布式环境 36 | dist.init_process_group(backend="nccl", init_method="env://") 37 | 38 | # 获取当前进程的 rank 和总进程数 39 | rank = dist.get_rank() 40 | world_size = dist.get_world_size() 41 | 42 | # 定义一个简单的 Dataset 43 | class SmallDataset(Dataset): 44 | def __init__(self): 45 | # 数据集只有 3 个样本 46 | self.data = [torch.tensor(i).to(torch.float32) for i in range(9)] 47 | print(self.data[0]) 48 | def __len__(self): 49 | return len(self.data) 50 | 51 | def __getitem__(self, idx): 52 | return self.data[idx] 53 | 54 | # 创建 Dataset 和 DistributedSampler 55 | dataset = SmallDataset() 56 | sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=False, drop_last=False) 57 | 58 | # 创建 DataLoader 59 | batch_size = 5 # batch size 大于数据集长度 60 | dataloader = DataLoader(dataset, batch_size=batch_size, sampler=sampler, drop_last=False, shuffle=False) 61 | """ 62 | 重要！ 63 | 当dataloader的drop_last=True时，且等效batchsize（num_tasks * batch_size）大于数据集长度时。 64 | dataloader将无法迭代，返回不了任何东西。len(dataloader) = 0 65 | """ 66 | res = [] 67 | # 测试 DataLoader 的迭代 68 | print(f"Rank {rank} iterating through DataLoader:") 69 | for batch in dataloader: 70 | batch.to(device) 71 | print(len(dataloader)) 72 | print(f"Rank {rank} batch: {batch}") 73 | res.append(batch) 74 | print(f"Rank {rank} result: {res}") 75 | 76 | local_pred = torch.cat(res, dim=0) 77 | 78 | gatherdata = distributed_gather_tensor(local_pred) 79 | print(f"Rank {rank} gatherdata: {gatherdata}") 80 | 81 | # check duplicate 82 | if rank == 0: 83 | print(f"Rank {rank} gatherdata: {gatherdata}") 84 | print(f"Rank {rank} gatherdata size: {gatherdata.size()}") 85 | print(f"Rank {rank} gatherdata shape: {gatherdata.shape}") 86 | print(f"Rank {rank} gatherdata dtype: {gatherdata.dtype}") 87 | 88 | 89 | # 清理分布式环境 90 | dist.destroy_process_group() -------------------------------------------------------------------------------- /IMDLBenCo/datasets/iml_datasets.py: -------------------------------------------------------------------------------- 1 | import json 2 | import os 3 | 4 | 5 | from ..registry import DATASETS 6 | from .abstract_dataset import AbstractDataset 7 | 8 | @DATASETS.register_module() 9 | class ManiDataset(AbstractDataset): 10 | def _init_dataset_path(self, path): 11 | self.entry_path = path 12 | tp_dir = os.path.join(path, 'Tp') 13 | gt_dir = os.path.join(path, 'Gt') 14 | 15 | assert os.path.isdir(path), NotADirectoryError(f"Get Error when loading from {self.entry_path}, the path is not a directory. Please check the path.") 16 | assert os.path.isdir(tp_dir), NotADirectoryError(f"Get Error when loading from {tp_dir}, the Tp directory is not exist. Please check the path.") 17 | assert os.path.isdir(gt_dir), NotADirectoryError(f"Get Error when loading from {gt_dir}, the Gt directory is not exist. Please check the path.") 18 | 19 | tp_list = os.listdir(tp_dir) 20 | gt_list = os.listdir(gt_dir) 21 | # Use sort mathod to keep order, to make sure the order is the same as the order in the tp_list and gt_list 22 | tp_list.sort() 23 | gt_list.sort() 24 | t_tp_list = [os.path.join(path, 'Tp', tp_list[index]) for index in range(len(tp_list))] 25 | t_gt_list = [os.path.join(path, 'Gt', gt_list[index]) for index in range(len(gt_list))] 26 | return t_tp_list, t_gt_list 27 | 28 | @DATASETS.register_module() 29 | class JsonDataset(AbstractDataset): 30 | """ init from a json file, which contains all the images path 31 | file is organized as: 32 | [ 33 | ["./Tp/6.jpg", "./Gt/6.jpg"], 34 | ["./Tp/7.jpg", "./Gt/7.jpg"], 35 | ["./Tp/8.jpg", "Negative"], 36 | ...... 37 | ] 38 | if path is "Neagative" then the image is negative sample, which means ground truths is a totally black image, and its label == 0. 39 | 40 | Args: 41 | path (_type_): _description_ 42 | transform_albu (_type_, optional): _description_. Defaults to None. 43 | mask_edge_generator (_type_, optional): _description_. Defaults to None. 44 | if_return_shape 45 | """ 46 | def _init_dataset_path(self, path): 47 | self.entry_path = path 48 | try: 49 | images = json.load(open(path, 'r')) 50 | except: 51 | raise TypeError(f"Get Error when loading from {self.entry_path}, please check the file format, it should be a json file, and the content should be like: [['./Tp/6.jpg', './Gt/6.jpg'], ['./Tp/7.jpg', './Gt/7.jpg'], ['./Tp/8.jpg', 'Negative'], ......]") 52 | tp_list = [] 53 | gt_list = [] 54 | for record in images: 55 | if os.path.isfile(record[0]): 56 | tp_list.append(record[0]) 57 | gt_list.append(record[1]) 58 | else: 59 | raise TypeError(f"Get Error when loading from {self.entry_path}, the error record is: {record[0]}, which is not a file. Please check this file or try to use absolute path instead. Otherwise if you want to use ManiDataset with a path instead of JsonDataset, please pass a path into the 'train_*.sh'. For more information please see the protocol here: https://scu-zjz.github.io/IMDLBenCo-doc/guide/quickstart/0_dataprepare.html#specific-format-definitions") 60 | return tp_list, gt_list 61 | 62 | -------------------------------------------------------------------------------- /IMDLBenCo/model_zoo/trufor/cmx/decoders/MLPDecoder.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch.nn as nn 3 | import torch 4 | 5 | import torch.nn.functional as F 6 | 7 | class MLP(nn.Module): 8 | """ 9 | Linear Embedding: 10 | """ 11 | def __init__(self, input_dim=2048, embed_dim=768): 12 | super().__init__() 13 | self.proj = nn.Linear(input_dim, embed_dim) 14 | 15 | def forward(self, x): 16 | x = x.flatten(2).transpose(1, 2) 17 | x = self.proj(x) 18 | return x 19 | 20 | 21 | class DecoderHead(nn.Module): 22 | def __init__(self, 23 | in_channels=[64, 128, 320, 512], 24 | num_classes=40, 25 | dropout_ratio=0.1, 26 | norm_layer=nn.BatchNorm2d, 27 | embed_dim=768, 28 | align_corners=False): 29 | 30 | super(DecoderHead, self).__init__() 31 | self.num_classes = num_classes 32 | self.dropout_ratio = dropout_ratio 33 | self.align_corners = align_corners 34 | 35 | self.in_channels = in_channels 36 | 37 | if dropout_ratio > 0: 38 | self.dropout = nn.Dropout2d(dropout_ratio) 39 | else: 40 | self.dropout = None 41 | 42 | c1_in_channels, c2_in_channels, c3_in_channels, c4_in_channels = self.in_channels 43 | 44 | embedding_dim = embed_dim 45 | self.linear_c4 = MLP(input_dim=c4_in_channels, embed_dim=embedding_dim) 46 | self.linear_c3 = MLP(input_dim=c3_in_channels, embed_dim=embedding_dim) 47 | self.linear_c2 = MLP(input_dim=c2_in_channels, embed_dim=embedding_dim) 48 | self.linear_c1 = MLP(input_dim=c1_in_channels, embed_dim=embedding_dim) 49 | 50 | self.linear_fuse = nn.Sequential( 51 | nn.Conv2d(in_channels=embedding_dim*4, out_channels=embedding_dim, kernel_size=1), 52 | norm_layer(embedding_dim), 53 | nn.ReLU(inplace=True) 54 | ) 55 | 56 | self.linear_pred = nn.Conv2d(embedding_dim, self.num_classes, kernel_size=1) 57 | 58 | def forward(self, inputs, return_feats=False): 59 | # len=4, 1/4,1/8,1/16,1/32 60 | c1, c2, c3, c4 = inputs 61 | 62 | ############## MLP decoder on C1-C4 ########### 63 | n, _, h, w = c4.shape 64 | 65 | _c4 = self.linear_c4(c4).permute(0,2,1).reshape(n, -1, c4.shape[2], c4.shape[3]) 66 | _c4 = F.interpolate(_c4, size=c1.size()[2:],mode='bilinear',align_corners=self.align_corners) 67 | 68 | _c3 = self.linear_c3(c3).permute(0,2,1).reshape(n, -1, c3.shape[2], c3.shape[3]) 69 | _c3 = F.interpolate(_c3, size=c1.size()[2:],mode='bilinear',align_corners=self.align_corners) 70 | 71 | _c2 = self.linear_c2(c2).permute(0,2,1).reshape(n, -1, c2.shape[2], c2.shape[3]) 72 | _c2 = F.interpolate(_c2, size=c1.size()[2:],mode='bilinear',align_corners=self.align_corners) 73 | 74 | _c1 = self.linear_c1(c1).permute(0,2,1).reshape(n, -1, c1.shape[2], c1.shape[3]) 75 | 76 | _c = torch.cat([_c4, _c3, _c2, _c1], dim=1) 77 | x = self.linear_fuse(_c) 78 | x = self.dropout(x) 79 | x = self.linear_pred(x) 80 | 81 | if return_feats: 82 | return x, _c 83 | else: 84 | return x 85 | 86 | -------------------------------------------------------------------------------- /IMDLBenCo/model_zoo/cat_net/cat_net_post_function.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from IMDLBenCo.datasets.utils import read_jpeg_from_memory 3 | from IMDLBenCo.registry import POSTFUNCS 4 | 5 | @POSTFUNCS.register_module() 6 | def cat_net_post_func(data_dict): 7 | tp_img = data_dict['image'] 8 | DCT_coef, qtables = __get_jpeg_info(tp_img) 9 | data_dict['DCT_coef'] = DCT_coef[0] 10 | data_dict['qtables'] = qtables[0] 11 | 12 | def __get_jpeg_info(image_tensor): 13 | """ 14 | :param im_path: JPEG image path 15 | :return: DCT_coef (Y,Cb,Cr), qtables (Y,Cb,Cr) 16 | """ 17 | num_channels = 1 18 | jpeg = read_jpeg_from_memory(image_tensor) 19 | 20 | # determine which axes to up-sample 21 | ci = jpeg.comp_info 22 | need_scale = [[ci[i].v_samp_factor, ci[i].h_samp_factor] for i in range(num_channels)] 23 | if num_channels == 3: 24 | if ci[0].v_samp_factor == ci[1].v_samp_factor == ci[2].v_samp_factor: 25 | need_scale[0][0] = need_scale[1][0] = need_scale[2][0] = 2 26 | if ci[0].h_samp_factor == ci[1].h_samp_factor == ci[2].h_samp_factor: 27 | need_scale[0][1] = need_scale[1][1] = need_scale[2][1] = 2 28 | else: 29 | need_scale[0][0] = 2 30 | need_scale[0][1] = 2 31 | 32 | # up-sample DCT coefficients to match image size 33 | DCT_coef = [] 34 | for i in range(num_channels): 35 | r, c = jpeg.coef_arrays[i].shape 36 | coef_view = jpeg.coef_arrays[i].reshape(r//8, 8, c//8, 8).transpose(0, 2, 1, 3) 37 | # case 1: row scale (O) and col scale (O) 38 | if need_scale[i][0]==1 and need_scale[i][1]==1: 39 | out_arr = np.zeros((r * 2, c * 2)) 40 | out_view = out_arr.reshape(r * 2 // 8, 8, c * 2 // 8, 8).transpose(0, 2, 1, 3) 41 | out_view[::2, ::2, :, :] = coef_view[:, :, :, :] 42 | out_view[1::2, ::2, :, :] = coef_view[:, :, :, :] 43 | out_view[::2, 1::2, :, :] = coef_view[:, :, :, :] 44 | out_view[1::2, 1::2, :, :] = coef_view[:, :, :, :] 45 | 46 | # case 2: row scale (O) and col scale (X) 47 | elif need_scale[i][0]==1 and need_scale[i][1]==2: 48 | out_arr = np.zeros((r * 2, c)) 49 | DCT_coef.append(out_arr) 50 | out_view = out_arr.reshape(r*2//8, 8, c // 8, 8).transpose(0, 2, 1, 3) 51 | out_view[::2, :, :, :] = coef_view[:, :, :, :] 52 | out_view[1::2, :, :, :] = coef_view[:, :, :, :] 53 | 54 | # case 3: row scale (X) and col scale (O) 55 | elif need_scale[i][0]==2 and need_scale[i][1]==1: 56 | out_arr = np.zeros((r, c * 2)) 57 | out_view = out_arr.reshape(r // 8, 8, c * 2 // 8, 8).transpose(0, 2, 1, 3) 58 | out_view[:, ::2, :, :] = coef_view[:, :, :, :] 59 | out_view[:, 1::2, :, :] = coef_view[:, :, :, :] 60 | 61 | # case 4: row scale (X) and col scale (X) 62 | elif need_scale[i][0]==2 and need_scale[i][1]==2: 63 | out_arr = np.zeros((r, c)) 64 | out_view = out_arr.reshape(r // 8, 8, c // 8, 8).transpose(0, 2, 1, 3) 65 | out_view[:, :, :, :] = coef_view[:, :, :, :] 66 | 67 | else: 68 | raise KeyError("Something wrong here.") 69 | 70 | DCT_coef.append(out_arr) 71 | 72 | # quantization tables 73 | qtables = [jpeg.quant_tables[ci[i].quant_tbl_no].astype(np.float64) for i in range(num_channels)] 74 | 75 | return DCT_coef, qtables 76 | -------------------------------------------------------------------------------- /IMDLBenCo/cli.py: -------------------------------------------------------------------------------- 1 | import os 2 | import argparse 3 | import requests 4 | 5 | from colorama import init, Fore, Style 6 | 7 | # from IMDLBenCo.utils.paths import BencoPath 8 | from IMDLBenCo.cli_funcs import cli_init, cli_env 9 | import importlib.metadata 10 | 11 | PYPI_API_URL = 'https://pypi.org/pypi/IMDLBenCo/json' 12 | from IMDLBenCo.version import __version__ 13 | 14 | def version_and_check_for_updates(): 15 | print(f'IMDLBenCo codebase version: {__version__}') 16 | try: 17 | response = requests.get(PYPI_API_URL, timeout=5) 18 | response.raise_for_status() # 如果响应码不是200，则抛出异常 19 | pypi_data = response.json() 20 | cloud_version = pypi_data['info']['version'] # 获取最新版本号 21 | # cloud_version = '0.1.21' # for debug & test 22 | print("\tChecking for updates...") 23 | print("\tLocal version: ", __version__) 24 | print("\tPyPI newest version: ", cloud_version) 25 | local_version = __version__ # 通过 importlib.metadata 获取当前安装版本 26 | 27 | if cloud_version != local_version: 28 | print(Fore.YELLOW + f"New version available: {cloud_version}. Your version: {local_version}.") 29 | print("Run 'pip install --upgrade imdlbenco' to upgrade.") 30 | else: 31 | print(Fore.GREEN + f"You are using the latest version: {local_version}.") 32 | except requests.exceptions.RequestException as e: 33 | print(Fore.RED + "Failed to check for updates from PyPI. Please check your internet connection.") 34 | print(f"Error: {e}") 35 | 36 | def main(): 37 | parser = argparse.ArgumentParser(description='Command line interface for IMDLBenCo, with codebase version: ' + __version__) 38 | 39 | # Add version argument with update check only when user requests version 40 | parser.add_argument('-v', '--version', action='store_true', help="Show the version of the tool") 41 | 42 | subparsers = parser.add_subparsers(dest='command', required=False) 43 | 44 | # init command 45 | parser_init = subparsers.add_parser('init', help='Initialize the scripts and configs in a directory') 46 | init_subparsers = parser_init.add_subparsers(dest='subcommand', required=False) 47 | 48 | # init base 49 | parser_init_base = init_subparsers.add_parser('base', help='Initialize the base environment') 50 | parser_init_base.set_defaults(subcommand='base') 51 | 52 | # init model_zoo 53 | parser_init_model_zoo = init_subparsers.add_parser('model_zoo', help='Initialize the model zoo') 54 | 55 | # init backbone 56 | parser_init_backbone = init_subparsers.add_parser('backbone', help='Initialize the backbone') 57 | 58 | 59 | # env command 60 | parser_env = subparsers.add_parser('env', help='Show environment information') 61 | 62 | parser.add_argument('--config', type=str, help='Path to the configuration file') 63 | 64 | args = parser.parse_args() 65 | if args.version: 66 | version_and_check_for_updates() 67 | 68 | if args.command == 'init': 69 | if args.subcommand is None: 70 | args.subcommand = 'base' 71 | cli_init(args.config, subcommand=args.subcommand) 72 | 73 | elif args.command == 'env': 74 | cli_env(args.config) 75 | 76 | # def train(config): 77 | # print(f'Training with config: {config}') 78 | 79 | # def evaluate(config): 80 | # print(f'Evaluating with config: {config}') 81 | 82 | if __name__ == '__main__': 83 | main() 84 | -------------------------------------------------------------------------------- /IMDLBenCo/evaluation/gradcam/grad_camera_visualize.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from typing import Union, Tuple 4 | from torch.utils.data import Dataset 5 | from .grad_cam_hack import GradCAM 6 | import numpy as np 7 | from pytorch_grad_cam.utils.image import show_cam_on_image 8 | from ...datasets.utils import denormalize 9 | import os 10 | from PIL import Image 11 | from tqdm import tqdm 12 | 13 | 14 | class OutputWrapper(torch.nn.Module): 15 | def __init__(self, model): 16 | super(OutputWrapper, self).__init__() 17 | self.model = model 18 | 19 | def forward(self, x): 20 | return self.model(**x)["pred_mask"] 21 | 22 | 23 | class SemanticSegmentationTarget: 24 | def __init__(self, mask): 25 | self.mask = torch.squeeze(mask, 0) 26 | 27 | def __call__(self, model_output): 28 | return (model_output * self.mask).sum() 29 | 30 | 31 | def grad_camera_visualize(model: nn.Module, 32 | image: Union[Tuple[str, str], Dataset], 33 | target_layers: list, 34 | output_path: str) -> None: 35 | """ 36 | visualize 37 | """ 38 | if not isinstance(image, Dataset): 39 | # TODO 40 | pass 41 | sampler = torch.utils.data.SequentialSampler(image) 42 | data_loader = torch.utils.data.DataLoader( 43 | image, 44 | sampler=sampler, 45 | batch_size=1, 46 | num_workers=1, 47 | pin_memory=False, 48 | drop_last=False, 49 | ) 50 | model = OutputWrapper(model=model) 51 | for data in tqdm(data_loader): 52 | for key, value in data.items(): 53 | if isinstance(value, torch.Tensor): 54 | data[key] = value.cuda() 55 | input_tensor = data['image'] 56 | cam = GradCAM(model=model, target_layers=target_layers) 57 | rgb_img = np.float32(torch.squeeze(denormalize(input_tensor), 0).permute(1, 2, 0).cpu()) 58 | targets = [SemanticSegmentationTarget(data['mask'])] 59 | grayscale_cam = cam(data, input_tensor=input_tensor, targets=targets)[0, :] 60 | 61 | # reshape type 62 | # target_shape = tuple(data['shape'].cpu().numpy().squeeze()) 63 | # rgb_img_resized = np.array(Image.fromarray(rgb_img.astype('uint8')).resize(target_shape, Image.BILINEAR)) 64 | # grayscale_cam_resized = np.array(Image.fromarray(grayscale_cam).resize(target_shape, Image.BILINEAR)) 65 | # cam_image = show_cam_on_image(rgb_img_resized, grayscale_cam_resized, use_rgb=True) 66 | # Image.fromarray(cam_image).save(os.path.join(output_path, data['name'][0].split('.')[0] + '.jpg')) 67 | 68 | # DePadding type 69 | target_shape = tuple(data['shape'].cpu().numpy().squeeze()) # (256, 384) for example 70 | import pdb 71 | pdb.set_trace() 72 | rgb_image_DePadding = np.array(Image.fromarray(rgb_img.astype('uint8')))[0:target_shape[0], 0:target_shape[1], :] 73 | grayscale_cam_DePadding = np.array(Image.fromarray(grayscale_cam))[0:target_shape[0], 0:target_shape[1]] 74 | cam_image = show_cam_on_image(rgb_image_DePadding, grayscale_cam_DePadding, use_rgb=True) 75 | Image.fromarray(cam_image).save(os.path.join(output_path, data['name'][0].split('.')[0] + '.jpg')) 76 | 77 | 78 | # No moving type 79 | # cam_image = show_cam_on_image(rgb_img, grayscale_cam, use_rgb=True) 80 | # Image.fromarray(cam_image).save(os.path.join(output_path, data['name'][0].split('.')[0] + '.jpg')) -------------------------------------------------------------------------------- /.github/workflows/tests.yml: -------------------------------------------------------------------------------- 1 | # This workflow will install Python dependencies, run tests and lint with a variety of Python versions 2 | # For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python 3 | 4 | name: Python tests 5 | 6 | on: 7 | push: 8 | branches: [ "main" ] 9 | pull_request: 10 | branches: [ "main" ] 11 | 12 | jobs: 13 | test: 14 | name: Run tests on ${{ matrix.os }} with Python ${{ matrix.python }} 15 | 16 | strategy: 17 | fail-fast: false 18 | matrix: 19 | os: ["ubuntu-latest", "macos-latest", "windows-latest"] 20 | torch: [{base: '1.13.0', vision: '0.14.0'}, {base: '2.5.1', vision: '0.20.1'}] 21 | python: ["3.9", "3.12"] 22 | exclude: 23 | - python: '3.12' 24 | torch: {base: '1.13.0', vision: '0.14.0'} 25 | runs-on: ${{ matrix.os }} 26 | 27 | steps: 28 | - uses: actions/checkout@v4 29 | - name: Set up Python ${{ matrix.python }} 30 | uses: actions/setup-python@v3 31 | with: 32 | python-version: ${{ matrix.python }} 33 | 34 | # install pytorch on different platform 35 | - name: Install torch on mac 36 | if: startsWith(matrix.os, 'macOS') 37 | run: pip install --no-cache-dir torch==${{ matrix.torch.base }} torchvision==${{ matrix.torch.vision }} 38 | - name: Install torch on Windows 39 | if: startsWith(matrix.os, 'windows') 40 | run: pip install --no-cache-dir torch==${{ matrix.torch.base }} torchvision==${{ matrix.torch.vision }} 41 | - name: Install torch on ubuntu 42 | if: startsWith(matrix.os, 'ubuntu') 43 | run: | 44 | sudo sed -i 's/azure\.//' /etc/apt/sources.list 45 | sudo apt update 46 | sudo apt install -y google-perftools 47 | pip install --no-cache-dir torch==${{ matrix.torch.base }}+cpu torchvision==${{ matrix.torch.vision }}+cpu --index-url https://download.pytorch.org/whl/cpu 48 | 49 | - name: Install testing dependencies 50 | run: | 51 | python -m pip install --upgrade pip 52 | pip install wheel 53 | pip install -r requirements-dev.txt 54 | pip install -r requirements.txt 55 | pip install -e . 56 | 57 | - name: Force old numpy for old torch 58 | if: ${{ matrix.torch.base == '1.13.0' }} 59 | run: pip install --upgrade 'numpy<2.0' 60 | 61 | - name: Lint with flake8 62 | run: | 63 | # stop the build if there are Python syntax errors or undefined names 64 | flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics 65 | # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide 66 | flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics 67 | - name: Run tests on Windows 68 | if: startsWith(matrix.os, 'windows') 69 | env: 70 | PYTHONDONTWRITEBYTECODE: 1 71 | # torchrun --nproc_per_node=2 -m pytest ./tests/dist 72 | run: | 73 | pytest -vv --cov=IMDLBenCo --cov-report=xml tests/pytests 74 | 75 | - name: Run '${{ matrix.testmarker }}' tests on Linux / Mac 76 | if: ${{ !startsWith(matrix.os, 'windows') }} 77 | env: 78 | # LD_PRELOAD: /usr/lib/x86_64-linux-gnu/libtcmalloc.so.4 79 | PYTHONDONTWRITEBYTECODE: 1 80 | run: | 81 | pytest -vv --cov=IMDLBenCo --cov-report=xml tests/pytests 82 | - name: Upload coverage reports to Codecov 83 | uses: codecov/codecov-action@v5 84 | with: 85 | token: ${{ secrets.CODECOV_TOKEN }} 86 | -------------------------------------------------------------------------------- /IMDLBenCo/model_zoo/span/SPAN.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from torchvision.transforms import Resize 5 | import numpy as np 6 | import os 7 | import json 8 | from . import PixelAttention as pa 9 | from . import mantranet 10 | from IMDLBenCo.registry import MODELS 11 | 12 | @MODELS.register_module() 13 | class SPAN(nn.Module): 14 | def __init__(self, weight_path, layers_steps=[1, 3, 9, 27, 81], device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')): 15 | super(SPAN, self).__init__() 16 | self.BCE_loss = nn.BCEWithLogitsLoss() 17 | self.featex = mantranet.IMTFE(device=device) # use to initialize the SRM filter. 18 | print("Model SPAN: load weight from", weight_path) 19 | self.featex.load_state_dict(torch.load(weight_path), strict=True) 20 | self.last_layer = self.Last_Layer_0725 21 | self.pixel_attention = nn.ModuleList([pa.PixelAttention(shift=step, in_channels=32, use_bn=False, use_res=True) for step in layers_steps]).to(device) 22 | self.outlier_trans = nn.Conv2d(self.featex.middle_and_last_block[-1].out_channels, 32, kernel_size=1, padding=0, bias=True).to(device) 23 | self.device = device 24 | 25 | def Last_Layer_0725(self, x): 26 | t = nn.Conv2d(32, 16, kernel_size=5, padding=2).to(x.device)(x) 27 | t = nn.ReLU().to(x.device)(t) 28 | t = nn.Conv2d(16, 8, kernel_size=5, padding=2).to(x.device)(t) 29 | t = nn.ReLU().to(x.device)(t) 30 | t = nn.Conv2d(8, 4, kernel_size=5, padding=2).to(x.device)(t) 31 | t = nn.ReLU().to(x.device)(t) 32 | t = nn.Conv2d(4, 1, kernel_size=5, padding=2).to(x.device)(t) 33 | return t 34 | 35 | def forward(self, image: torch.Tensor, mask=None, edge_mask=None, shape=None, *args, **kwargs): 36 | """denormalize image with imagenet mean and std 37 | then re-norm it to the preporcessing same with ManTra-Net (-1, 1) 38 | """ 39 | imagenet_mean=torch.tensor([0.485, 0.456, 0.406]).unsqueeze(1).unsqueeze(2).unsqueeze(0).to(self.device) 40 | imagenet_std=torch.tensor([0.229, 0.224, 0.225]).unsqueeze(1).unsqueeze(2).unsqueeze(0).to(self.device) 41 | image = image * imagenet_std # [B, 3, H, W] 42 | image = image + imagenet_mean 43 | 44 | image = image * 2 - 1 45 | x = image.to(self.device) 46 | x = self.featex(x) 47 | x = self.outlier_trans(x) 48 | for attention_layer in self.pixel_attention: 49 | x = attention_layer(x) 50 | mask_pred = self.last_layer(x) 51 | predict_loss = self.BCE_loss(mask_pred, mask) 52 | mask_pred = torch.sigmoid(mask_pred) 53 | output_dict = { 54 | "backward_loss": predict_loss, 55 | "pred_mask": mask_pred, 56 | "pred_label": None, 57 | "visual_loss": { 58 | "predict_loss": predict_loss, 59 | }, 60 | 61 | "visual_image": { 62 | "pred_mask": mask_pred 63 | } 64 | } 65 | return output_dict 66 | 67 | if __name__ == "__main__": 68 | device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 69 | model = SPAN().to(device) 70 | input_tensor = torch.randn(1, 3, 256, 256).to(device) 71 | mask_tensor = torch.randn(1, 1, 224, 224).to(device) 72 | output_dict = model(input_tensor, mask=mask_tensor) 73 | print("output_dict['pred_mask'].shape:", output_dict["pred_mask"].shape) 74 | print("output_dict['backward_loss']:", output_dict["backward_loss"]) 75 | -------------------------------------------------------------------------------- /IMDLBenCo/datasets/balanced_dataset.py: -------------------------------------------------------------------------------- 1 | import random 2 | from torch.utils.data import Dataset, DataLoader 3 | 4 | from .iml_datasets import JsonDataset, ManiDataset 5 | 6 | from ..transforms import get_albu_transforms 7 | from .utils import pil_loader, denormalize 8 | 9 | from IMDLBenCo.registry import DATASETS 10 | @DATASETS.register_module() 11 | class BalancedDataset(Dataset): 12 | """The BalancedDataset manages multiple iml_datasets, so it does not inherit from AbstractDataset. 13 | 14 | Args: 15 | Dataset (_type_): _description_ 16 | 17 | Returns: 18 | _type_: _description_ 19 | """ 20 | 21 | def __init__(self, 22 | path = None, 23 | sample_number = 1840, 24 | *args, 25 | **kwargs 26 | ) -> None: 27 | self.sample_number = sample_number 28 | if path == None: 29 | # Defalut 30 | self.settings_list = [ 31 | ['/mnt/data0/public_datasets/IML/CASIA2.0', ManiDataset], 32 | ['/mnt/data0/public_datasets/IML/FantasticReality_v1/FantasticReality.json', JsonDataset], 33 | ['/mnt/data0/public_datasets/IML/IMD_20_1024', ManiDataset], 34 | ['/mnt/data0/public_datasets/IML/tampCOCO/sp_COCO_list.json', JsonDataset], 35 | ['/mnt/data0/public_datasets/IML/tampCOCO/cm_COCO_list.json', JsonDataset], 36 | ['/mnt/data0/public_datasets/IML/tampCOCO/bcm_COCO_list.json', JsonDataset], 37 | ['/mnt/data0/public_datasets/IML/tampCOCO/bcmc_COCO_list.json', JsonDataset] 38 | ] 39 | else: 40 | import json 41 | with open(path, "r") as f: 42 | setting_json = json.load(f) 43 | # self.settings_list = path_list 44 | self.settings_list = [] 45 | for dataset_str, dataset_path in setting_json: 46 | self.settings_list.append( 47 | [ 48 | dataset_path, 49 | DATASETS.get(dataset_str) 50 | ] 51 | ) 52 | 53 | 54 | self.dataset_list = [self._get_dataset(path, dataset_type, *args, **kwargs) for path, dataset_type in self.settings_list] 55 | 56 | 57 | def _get_dataset(self, path, dataset_type, *args, **kwargs): 58 | return dataset_type(path, *args, **kwargs) 59 | 60 | def __len__(self): 61 | return self.sample_number * len(self.settings_list) 62 | 63 | def __getitem__(self, index): 64 | dataset_index = index // self.sample_number 65 | 66 | selected_dataset = self.dataset_list[dataset_index] 67 | length = len(selected_dataset) 68 | selected_item = random.randint(0, length - 1) 69 | return selected_dataset[selected_item] 70 | 71 | def __str__(self): 72 | # Print basic information about the dataset and its settings 73 | info = f"<===BalancedDataset with {len(self.settings_list)} datasets:===>\n" 74 | 75 | total_sample = 0 76 | # Iterate through the settings list and print dataset paths and types 77 | for idx, (i_dataset) in enumerate(self.dataset_list): 78 | i_len = len(i_dataset) 79 | total_sample += i_len 80 | info += f" Dataset {idx + 1}: {i_dataset}\n" 81 | 82 | # Print sample number and total number of samples 83 | info += f"Sample number per dataset: {self.sample_number:,}\n" 84 | epoch_samples = self.__len__() 85 | info += f"Samples per Epoch: {epoch_samples:,}\n" 86 | info += f"Total images in whole dataset: {total_sample:,}\n" 87 | info += f"<================================================>\n" 88 | return info 89 | -------------------------------------------------------------------------------- /IMDLBenCo/modules/backbones/unet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from IMDLBenCo.registry import MODELS 5 | from ..extractors.high_frequency_feature_extraction import FFTExtractor, DCTExtractor 6 | from ..extractors.sobel import SobelFilter 7 | from ..extractors.bayar_conv import BayerConv 8 | from ..extractors.srm_filter import SRMConv2D 9 | 10 | class UNet(nn.Module): 11 | def __init__(self, input_head=None, num_channels=3): 12 | super(UNet, self).__init__() 13 | net = torch.hub.load('milesial/Pytorch-UNet', 'unet_carvana', pretrained=True, scale=1) 14 | self.input_head = input_head 15 | if input_head != None: 16 | original_first_layer = net.inc.double_conv[0] 17 | new_first_layer = nn.Conv2d(num_channels + 3, original_first_layer.out_channels, 18 | kernel_size=original_first_layer.kernel_size, stride=original_first_layer.stride, 19 | padding=original_first_layer.padding, bias=False) 20 | new_first_layer.weight.data[:, :3, :, :] = original_first_layer.weight.data.clone()[:, :3, :, :] 21 | if num_channels > 0: 22 | new_first_layer.weight.data[:, 3:, :, :] = torch.nn.init.kaiming_normal_(new_first_layer.weight[:, 3:, :, :]) 23 | net.inc.double_conv[0] = new_first_layer 24 | last_layer = net.outc.conv 25 | new_last_layer = nn.Conv2d(last_layer.in_channels, 1, 26 | kernel_size=last_layer.kernel_size, stride=last_layer.stride, 27 | padding=last_layer.padding, bias=False) 28 | net.outc.conv = new_last_layer 29 | self.net = net 30 | self.loss_fun = nn.BCEWithLogitsLoss() 31 | 32 | def forward(self, image, mask, *args, **kwargs): 33 | if self.input_head != None: 34 | feature = self.input_head(image) 35 | input = torch.cat([image,feature],dim=1) 36 | else: 37 | input = image 38 | mask_pred = self.net(input) 39 | loss = self.loss_fun(mask_pred,mask) 40 | mask_pred = torch.sigmoid(mask_pred) 41 | output_dict = { 42 | # loss for backward 43 | "backward_loss": loss, 44 | # predicted mask, will calculate for metrics automatically 45 | "pred_mask": mask_pred, 46 | # predicted binaray label, will calculate for metrics automatically 47 | "pred_label": None, 48 | 49 | # ----values below is for visualization---- 50 | # automatically visualize with the key-value pairs 51 | "visual_loss": { 52 | "predict_loss": loss, 53 | }, 54 | 55 | "visual_image": { 56 | "pred_mask": mask_pred, 57 | } 58 | # ----------------------------------------- 59 | } 60 | return output_dict 61 | 62 | 63 | @MODELS.register_module() 64 | def unet(): 65 | return UNet(None) 66 | 67 | @MODELS.register_module() 68 | def fft_unet(): 69 | return UNet(FFTExtractor(), 3) 70 | 71 | @MODELS.register_module() 72 | def dct_unet(): 73 | return UNet(DCTExtractor(), 3) 74 | 75 | @MODELS.register_module() 76 | def sobel_unet(): 77 | return UNet(SobelFilter(), 1) 78 | 79 | @MODELS.register_module() 80 | def bayar_unet(): 81 | return UNet(BayerConv(), 3) 82 | 83 | @MODELS.register_module() 84 | def srm_unet(): 85 | return UNet(SRMConv2D(), 9) 86 | 87 | 88 | if __name__ == '__main__': 89 | from fvcore.nn import FlopCountAnalysis, parameter_count_table,flop_count_table 90 | 91 | model = UNet(None,0, 3, 1, False) 92 | flops = FlopCountAnalysis(model,(torch.randn(1,3,512,512),torch.randn(1,1,512,512))) 93 | # # 打印模型的参数信息 94 | print(flop_count_table(flops)) -------------------------------------------------------------------------------- /IMDLBenCo/modules/extractors/high_frequency_feature_extraction.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.fft 4 | import math 5 | class DCTExtractor(nn.Module): 6 | def __init__(self, alpha=0.1): 7 | super(DCTExtractor, self).__init__() 8 | if alpha <= 0 or alpha >= 1: 9 | raise ValueError("alpha must be between 0 and 1 (exclusive)") 10 | self.alpha = alpha 11 | self.dct_matrix_h = None 12 | self.dct_matrix_w = None 13 | 14 | def create_dct_matrix(self, N): 15 | n = torch.arange(N, dtype=torch.float32).reshape((1, N)) 16 | k = torch.arange(N, dtype=torch.float32).reshape((N, 1)) 17 | dct_matrix = torch.sqrt(torch.tensor(2.0 / N)) * torch.cos(math.pi * k * (2 * n + 1) / (2 * N)) 18 | dct_matrix[0, :] = 1 / math.sqrt(N) 19 | return dct_matrix 20 | 21 | def dct_2d(self, x): 22 | H, W = x.size(-2), x.size(-1) 23 | if self.dct_matrix_h is None or self.dct_matrix_h.size(0) != H: 24 | self.dct_matrix_h = self.create_dct_matrix(H).to(x.device) 25 | if self.dct_matrix_w is None or self.dct_matrix_w.size(0) != W: 26 | self.dct_matrix_w = self.create_dct_matrix(W).to(x.device) 27 | 28 | return torch.matmul(self.dct_matrix_h, torch.matmul(x, self.dct_matrix_w.t())) 29 | 30 | def idct_2d(self, x): 31 | H, W = x.size(-2), x.size(-1) 32 | if self.dct_matrix_h is None or self.dct_matrix_h.size(0) != H: 33 | self.dct_matrix_h = self.create_dct_matrix(H).to(x.device) 34 | if self.dct_matrix_w is None or self.dct_matrix_w.size(0) != W: 35 | self.dct_matrix_w = self.create_dct_matrix(W).to(x.device) 36 | 37 | return torch.matmul(self.dct_matrix_h.t(), torch.matmul(x, self.dct_matrix_w)) 38 | 39 | def high_pass_filter(self, x, alpha): 40 | h, w = x.shape[-2:] 41 | mask = torch.ones(h, w, device=x.device) 42 | alpha_h, alpha_w = int(alpha * h), int(alpha * w) 43 | mask[:alpha_h, :alpha_w] = 0 44 | return x * mask 45 | 46 | def forward(self, x): 47 | xq = self.dct_2d(x) 48 | xq_high = self.high_pass_filter(xq, self.alpha) 49 | xh = self.idct_2d(xq_high) 50 | B = xh.shape[0] 51 | min_vals = xh.reshape(B, -1).min(dim=1, keepdim=True).values.view(B, 1, 1, 1) 52 | max_vals = xh.reshape(B, -1).max(dim=1, keepdim=True).values.view(B, 1, 1, 1) 53 | xh = (xh - min_vals) / (max_vals - min_vals) 54 | return xh 55 | 56 | class FFTExtractor(nn.Module): 57 | def __init__(self, alpha=0.1): 58 | super(FFTExtractor, self).__init__() 59 | if alpha <= 0 or alpha >= 1: 60 | raise ValueError("alpha must be between 0 and 1 (exclusive)") 61 | self.alpha = alpha 62 | 63 | def fft_2d(self, x): 64 | return torch.fft.fftshift(torch.fft.fft2(x)) 65 | 66 | def ifft_2d(self, x): 67 | return torch.fft.ifft2(torch.fft.ifftshift(x)) 68 | 69 | def high_pass_filter(self, x, alpha): 70 | h, w = x.shape[-2:] 71 | mask = torch.ones(h, w, device=x.device) 72 | alpha_h, alpha_w = int(alpha * h) // 2, int(alpha * w) // 2 73 | mask_center_h = h // 2 74 | mask_center_w = w // 2 75 | mask[mask_center_h - alpha_h:mask_center_h + alpha_h, mask_center_w - alpha_w:mask_center_w + alpha_w] = 0 76 | mask = mask.expand_as(x) 77 | return x * mask 78 | 79 | def forward(self, x): 80 | xq = self.fft_2d(x) 81 | xq_high = self.high_pass_filter(xq, self.alpha) 82 | xh = self.ifft_2d(xq_high) 83 | xh = xh.real 84 | B = xh.shape[0] 85 | min_vals = xh.reshape(B, -1).min(dim=1, keepdim=True).values.reshape(B, 1, 1, 1) 86 | max_vals = xh.reshape(B, -1).max(dim=1, keepdim=True).values.reshape(B, 1, 1, 1) 87 | xh = (xh - min_vals) / (max_vals - min_vals) 88 | return xh 89 | 90 | -------------------------------------------------------------------------------- /IMDLBenCo/modules/backbones/resnet.py: -------------------------------------------------------------------------------- 1 | import timm 2 | import torch 3 | 4 | 5 | from fvcore.nn import FlopCountAnalysis,flop_count_table 6 | import torch.nn as nn 7 | from IMDLBenCo.registry import MODELS 8 | from ..extractors.high_frequency_feature_extraction import( 9 | FFTExtractor, 10 | DCTExtractor 11 | ) 12 | from ..extractors.sobel import SobelFilter 13 | from ..extractors.bayar_conv import BayerConv 14 | from ..extractors.srm_filter import SRMConv2D 15 | 16 | 17 | 18 | class ResNet(nn.Module): 19 | def __init__(self, input_head=None,num_channels=3): 20 | super(ResNet, self).__init__() 21 | model = timm.create_model('resnet152', pretrained=True) 22 | self.backbone = nn.Sequential(*list(model.children())[:7]) 23 | self.deconv_model = nn.Sequential( 24 | nn.ConvTranspose2d(1024, 256, kernel_size=4, stride=2, padding=1), # [1, 512, 32, 32] 25 | nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1), # [1, 128, 64, 64] 26 | nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1), # [1, 64, 128, 128] 27 | nn.ConvTranspose2d(64, 1, kernel_size=4, stride=2, padding=1) # [1, 1, 256, 256] 28 | ) 29 | original_first_layer = list(model.children())[0] 30 | if input_head != None: 31 | self.input_head = input_head 32 | new_first_layer = nn.Conv2d(num_channels + 3, original_first_layer.out_channels, 33 | kernel_size=original_first_layer.kernel_size, stride=original_first_layer.stride, 34 | padding=original_first_layer.padding, bias=False) 35 | new_first_layer.weight.data[:, :3, :, :] = original_first_layer.weight.data.clone()[:, :3, :, :] 36 | if num_channels > 0: 37 | new_first_layer.weight.data[:, 3:, :, :] = torch.nn.init.kaiming_normal_(new_first_layer.weight[:, 3:, :, :]) 38 | self.backbone[0] = new_first_layer 39 | else: 40 | self.input_head = None 41 | self.loss_fun = nn.BCEWithLogitsLoss() 42 | def forward(self, image, mask, *args, **kwargs): 43 | if self.input_head != None: 44 | feature = self.input_head(image) 45 | input = torch.cat([image,feature],dim=1) 46 | else: 47 | input = image 48 | mask_pred = self.deconv_model(self.backbone(input)) 49 | loss = self.loss_fun(mask_pred,mask) 50 | mask_pred = torch.sigmoid(mask_pred) 51 | output_dict = { 52 | # loss for backward 53 | "backward_loss": loss, 54 | # predicted mask, will calculate for metrics automatically 55 | "pred_mask": mask_pred, 56 | # predicted binaray label, will calculate for metrics automatically 57 | "pred_label": None, 58 | 59 | # ----values below is for visualization---- 60 | # automatically visualize with the key-value pairs 61 | "visual_loss": { 62 | "predict_loss": loss, 63 | }, 64 | 65 | "visual_image": { 66 | "pred_mask": mask_pred, 67 | } 68 | # ----------------------------------------- 69 | } 70 | return output_dict 71 | 72 | @MODELS.register_module() 73 | def resnet(): 74 | return ResNet(None,0) 75 | 76 | @MODELS.register_module() 77 | def fft_resnet(): 78 | return ResNet(FFTExtractor(),3) 79 | 80 | @MODELS.register_module() 81 | def dct_resnet(): 82 | return ResNet(DCTExtractor(),3) 83 | 84 | @MODELS.register_module() 85 | def sobel_resnet(): 86 | return ResNet(SobelFilter(), 1) 87 | 88 | @MODELS.register_module() 89 | def bayar_resnet(): 90 | return ResNet(BayerConv(), 3) 91 | 92 | @MODELS.register_module() 93 | def srm_resnet(): 94 | return ResNet(SRMConv2D(),9) 95 | 96 | 97 | if __name__ == '__main__': 98 | 99 | #创建模型 100 | model = ResNet() 101 | 102 | flops = FlopCountAnalysis(model,(torch.randn(1,3,512,512),torch.randn(1,1,512,512))) 103 | print(flop_count_table(flops)) 104 | 105 | -------------------------------------------------------------------------------- /IMDLBenCo/transforms/edge_mask_generator.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from matplotlib import pyplot as plt 3 | from torch.nn import functional as F 4 | 5 | class EdgeMaskGenerator(torch.nn.Module): 6 | """generate the 'edge bar' for a 0-1 mask Groundtruth of a image 7 | Algorithm is based on 'Morphological Dilation and Difference Reduction' 8 | 9 | Which implemented with fixed-weight Convolution layer with weight matrix looks like a cross, 10 | for example, if kernel size is 3, the weight matrix is: 11 | [[0, 1, 0], 12 | [1, 1, 1], 13 | [0, 1, 0]] 14 | 15 | """ 16 | def __init__(self, kernel_size = 3) -> None: 17 | super().__init__() 18 | self.kernel_size = kernel_size 19 | 20 | def _dilate(self, image, kernel_size=3): 21 | """Doings dilation on the image 22 | 23 | Args: 24 | image (_type_): 0-1 tensor in shape (B, C, H, W) 25 | """ 26 | assert kernel_size % 2 == 1, "Kernel size must be odd" 27 | assert image.shape[2] > kernel_size and image.shape[3] > kernel_size, "Image must be larger than kernel size" 28 | 29 | kernel = torch.zeros((1, 1, kernel_size, kernel_size)) 30 | kernel[0, 0, kernel_size // 2: kernel_size//2+1, :] = 1 31 | kernel[0, 0, :, kernel_size // 2: kernel_size//2+1] = 1 32 | kernel = kernel.float() 33 | # print(kernel) 34 | res = F.conv2d(image, kernel.view([1,1,kernel_size, kernel_size]),stride=1, padding = kernel_size // 2) 35 | return (res > 0) * 1.0 36 | 37 | 38 | def _find_edge(self, image, kernel_size=3, return_all=False): 39 | """Find 0-1 edges of the image 40 | 41 | Args: 42 | image (_type_): 0-1 ndarray in shape (B, C, H, W) 43 | """ 44 | image = torch.tensor(image).float() 45 | shape = image.shape 46 | 47 | if len(shape) == 2: 48 | image = image.reshape([1, 1, shape[0], shape[1]]) 49 | if len(shape) == 3: 50 | image = image.reshape([1, shape[0], shape[1], shape[2]]) 51 | assert image.shape[1] == 1, "Image must be single channel" 52 | 53 | img = self._dilate(image, kernel_size=kernel_size) 54 | 55 | erosion = self._dilate(1-image, kernel_size=kernel_size) 56 | 57 | diff = -torch.abs(erosion - img) + 1 58 | diff = (diff > 0) * 1.0 59 | # res = dilate(diff) 60 | diff = diff.numpy() 61 | if return_all : 62 | return diff, img, erosion 63 | else: 64 | return diff 65 | 66 | def forward(self, x, return_all=False): 67 | """ 68 | Args: 69 | image (_type_): 0-1 ndarray in shape (B, C, H, W) 70 | """ 71 | return self._find_edge(x, self.kernel_size, return_all=return_all) 72 | 73 | 74 | # """Codes below are for testing""" 75 | # if __name__ == '__main__': 76 | # lists = ['NC2016_1504.jpg', '519_mask.jpg', '526_mask.jpg', '528_mask.jpg', '534_mask.jpg'] 77 | 78 | # for i in lists: 79 | # img = plt.imread(f'./components/Edge_generator/{i}') 80 | # img = torch.tensor(img) 81 | # print(img) 82 | # img = (img > 127).float() 83 | # plt.subplot(1, 4, 1) 84 | # plt.imshow(img, cmap='gray') 85 | # print(img) 86 | # Edge = EdgeGenerator(kernel_size=11) 87 | 88 | # raw_img = img.view(1, 1, img.shape[0], img.shape[1]) 89 | 90 | 91 | # # print(img) 92 | # # plt.subplot(1,4, 2) 93 | # # plt.imshow(diff.detach().numpy()[0, 0, :, :], cmap='gray') 94 | 95 | # diff, img,erosion, = Edge(raw_img, return_all=True) 96 | 97 | # plt.subplot(1,4, 2) 98 | # plt.imshow(diff[0, 0, :, :], cmap='gray') 99 | 100 | # # plt.subplot(1,4, 3) 101 | # # plt.imshow(img.detach().numpy()[0, 0, :, :], cmap='gray') 102 | 103 | # # plt.subplot(1, 4, 4) 104 | # # plt.imshow(erosion.detach().numpy()[0, 0, :, :], cmap='gray') 105 | 106 | 107 | # print(diff.shape) 108 | # print(img.shape) 109 | # print(erosion.shape) 110 | # plt.show() -------------------------------------------------------------------------------- /IMDLBenCo/cli_funcs/cli_init.py: -------------------------------------------------------------------------------- 1 | import os 2 | import re 3 | from colorama import init, Fore, Style 4 | from IMDLBenCo.utils.paths import BencoPath 5 | from .copy_funcs import copy_files, copy_file 6 | 7 | def _copy_train_scripts(): 8 | current_dir = os.getcwd() 9 | # Copy train scripts 10 | target_dir = os.path.join(current_dir) 11 | if not os.path.exists(target_dir): 12 | os.makedirs(target_dir) 13 | copy_files(BencoPath.get_templates_dir(), target_dir) 14 | 15 | 16 | def _copy_dataset_json(): 17 | current_dir = os.getcwd() 18 | # Copy train scripts 19 | target_dir = os.path.join(current_dir) 20 | if not os.path.exists(target_dir): 21 | os.makedirs(target_dir) 22 | copy_files(BencoPath.get_dataset_json_dir(), target_dir) 23 | 24 | def _copy_init_base_files(): 25 | current_dir = os.getcwd() 26 | # Copy train scripts 27 | target_dir = os.path.join(current_dir) 28 | if not os.path.exists(target_dir): 29 | os.makedirs(target_dir) 30 | copy_files(BencoPath.get_init_base_dir(), target_dir) 31 | 32 | def _copy_demo_runs(): 33 | # Copy demo runs 34 | current_dir = os.getcwd() 35 | target_dir = os.path.join(current_dir, 'runs') 36 | if not os.path.exists(target_dir): 37 | os.makedirs(target_dir) 38 | if not os.path.exists(os.path.join(target_dir, "test_save_images")): 39 | os.makedirs(os.path.join(target_dir, "test_save_images")) 40 | if not os.path.exists(os.path.join(target_dir, "test_complexity")): 41 | os.makedirs(os.path.join(target_dir, "test_complexity")) 42 | copy_files(BencoPath.get_model_zoo_runs_dir(), target_dir) 43 | copy_files(BencoPath.get_model_zoo_runs_dir() / "test_save_images" , os.path.join(target_dir, "test_save_images")) 44 | copy_files(BencoPath.get_model_zoo_runs_dir() / "test_complexity" , os.path.join(target_dir, "test_complexity")) 45 | 46 | 47 | def _copy_demo_configs(): 48 | # Copy demo configs 49 | current_dir = os.getcwd() 50 | target_dir = os.path.join(current_dir, 'configs') 51 | if not os.path.exists(target_dir): 52 | os.makedirs(target_dir) 53 | copy_files(BencoPath.get_model_zoo_configs_dir(), target_dir) 54 | 55 | # hejack train.py test.py test_robust.py with self defined model 56 | def _inject_after_last_import(path, inject): 57 | # 读取文件内容 58 | with open(path, 'r') as file: 59 | content = file.readlines() 60 | 61 | base_file_name = os.path.basename(path) 62 | # 检查是否已经存在指定的导入语句 63 | if any(inject in line for line in content): 64 | print(f" The specified import statement already exists in {base_file_name}.") 65 | return 66 | 67 | # 正则表达式匹配 from ... import ... 行 68 | pattern = re.compile(r'^\s*from\s+\S+\s+import\s+(?!.*\()(\S+(,\s*\S+)*)\s*$') 69 | 70 | # 找到最后一个匹配的行的索引 71 | last_import_index = -1 72 | for i, line in enumerate(content): 73 | if pattern.match(line): 74 | last_import_index = i 75 | 76 | # 如果找到匹配的行，则在其后插入新的内容 77 | if last_import_index != -1: 78 | content.insert(last_import_index + 1, inject + '\n') 79 | 80 | # re-save the file 81 | with open(path, 'w') as file: 82 | file.writelines(content) 83 | print(f" Injected '{inject}' into {path}.") 84 | 85 | def cli_init(config, subcommand): 86 | print(f'{Fore.GREEN}Initializing in current working directory...{Style.RESET_ALL}') 87 | 88 | # base initialize that only contain default scripts 89 | if subcommand == "base": 90 | _copy_train_scripts() 91 | _copy_init_base_files() 92 | _copy_dataset_json() 93 | for name in ['train.py', 'test.py', 'test_robust.py', 'test_complexity.py', 'test_save_images.py']: 94 | injected_str = "from mymodel import MyModel # TODO, you need to change this line when modifying the name model" 95 | current_dir = os.getcwd() 96 | target_dir = os.path.join(current_dir, name) 97 | _inject_after_last_import(target_dir, injected_str) 98 | 99 | if subcommand == "model_zoo": 100 | _copy_train_scripts() 101 | _copy_demo_runs() 102 | _copy_demo_configs() 103 | _copy_dataset_json() 104 | # base initialize that only contain default scripts 105 | if subcommand == "backbone": 106 | _copy_train_scripts() 107 | _copy_demo_runs() 108 | _copy_demo_configs() 109 | _copy_dataset_json() 110 | print(f'{Fore.GREEN}Successfully initialized IMDLBenCo scripts.{Style.RESET_ALL}') -------------------------------------------------------------------------------- /IMDLBenCo/modules/backbones/swin.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import timm 4 | import torch.nn.functional as F 5 | from IMDLBenCo.registry import MODELS 6 | 7 | from ..extractors.high_frequency_feature_extraction import FFTExtractor, DCTExtractor 8 | from ..extractors.sobel import SobelFilter 9 | from ..extractors.bayar_conv import BayerConv 10 | from ..extractors.srm_filter import SRMConv2D 11 | 12 | class Swin(timm.models.swin_transformer.SwinTransformer): 13 | def __init__(self, input_head=None, num_channels=3,pretrained=True): 14 | super(Swin, self).__init__(img_size=512,patch_size=4, window_size=7, embed_dim=128, depths=(2, 2, 18, 2), num_heads=(4, 8, 16, 32)) 15 | if pretrained: 16 | model = timm.create_model('swin_base_patch4_window7_224', pretrained=True) 17 | self.load_state_dict(model.state_dict()) 18 | original_first_layer = self.patch_embed.proj 19 | if input_head != None: 20 | self.input_head = input_head 21 | new_first_layer = nn.Conv2d(num_channels + 3, original_first_layer.out_channels, 22 | kernel_size=original_first_layer.kernel_size, stride=original_first_layer.stride, 23 | padding=original_first_layer.padding, bias=False) 24 | new_first_layer.weight.data[:, :3, :, :] = original_first_layer.weight.data.clone()[:, :3, :, :] 25 | if num_channels > 0: 26 | new_first_layer.weight.data[:, 3:, :, :] = torch.nn.init.kaiming_normal_(new_first_layer.weight[:, 3:, :, :]) 27 | self.patch_embed.proj = new_first_layer 28 | else: 29 | self.input_head = None 30 | 31 | self.deconv_model = nn.Sequential( 32 | nn.ConvTranspose2d(1024, 256, kernel_size=3, stride=2, padding=1, output_padding=1), # 输出: 1, 256, 64, 64 33 | nn.ReLU(), 34 | nn.ConvTranspose2d(256, 128, kernel_size=3, stride=2, padding=1, output_padding=1), # 输出: 1, 128, 128, 128 35 | nn.ReLU(), 36 | nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1, output_padding=1), # 输出: 1, 64, 256, 256 37 | nn.ReLU(), 38 | nn.ConvTranspose2d(64, 32, kernel_size=3, stride=2, padding=1, output_padding=1), # 输出: 1, 1, 512, 512 39 | nn.ReLU(), 40 | nn.ConvTranspose2d(32, 1, kernel_size=3, stride=2, padding=1, output_padding=1) # 输出: 1, 1, 512, 512 41 | ) 42 | self.loss_fun = nn.BCEWithLogitsLoss() 43 | 44 | 45 | def forward(self, image, mask, *args, **kwargs): 46 | if self.input_head != None: 47 | feature = self.input_head(image) 48 | input = torch.cat([image, feature],dim=1) 49 | else: 50 | input = image 51 | # import pdb 52 | # pdb.set_trace() 53 | feature = self.forward_features(input) 54 | feature = feature.permute(0,3,1,2) 55 | mask_pred = self.deconv_model(feature) 56 | loss = self.loss_fun(mask_pred,mask) 57 | mask_pred = torch.sigmoid(mask_pred) 58 | output_dict = { 59 | # loss for backward 60 | "backward_loss": loss, 61 | # predicted mask, will calculate for metrics automatically 62 | "pred_mask": mask_pred, 63 | # predicted binaray label, will calculate for metrics automatically 64 | "pred_label": None, 65 | 66 | # ----values below is for visualization---- 67 | # automatically visualize with the key-value pairs 68 | "visual_loss": { 69 | "predict_loss": loss, 70 | }, 71 | 72 | "visual_image": { 73 | "pred_mask": mask_pred, 74 | } 75 | # ----------------------------------------- 76 | } 77 | return output_dict 78 | 79 | @MODELS.register_module() 80 | def swin(): 81 | return Swin(None) 82 | 83 | @MODELS.register_module() 84 | def fft_swin(): 85 | return Swin(FFTExtractor(), 3) 86 | 87 | @MODELS.register_module() 88 | def dct_swin(): 89 | return Swin(DCTExtractor(), 3) 90 | 91 | @MODELS.register_module() 92 | def sobel_swin(): 93 | return Swin(SobelFilter(), 1) 94 | 95 | @MODELS.register_module() 96 | def bayar_swin(): 97 | return Swin(BayerConv(), 3) 98 | 99 | @MODELS.register_module() 100 | def srm_swin(): 101 | return Swin(SRMConv2D(), 9) 102 | 103 | if __name__ == '__main__': 104 | 105 | import timm 106 | from pprint import pprint 107 | from fvcore.nn import FlopCountAnalysis, parameter_count_table,flop_count_table 108 | import torch 109 | model = Swin() 110 | flops = FlopCountAnalysis(model,(torch.randn(1,3,512,512),torch.randn(1,1,512,512))) 111 | print(flop_count_table(flops)) -------------------------------------------------------------------------------- /IMDLBenCo/model_zoo/span/PixelAttention.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | class PixelAttention(nn.Module): 6 | def __init__(self, in_channels=3, kernel_range=[3, 3], shift=1, use_bn=False, use_res=False): 7 | super(PixelAttention, self).__init__() 8 | self.in_channels = in_channels 9 | self.kernel_range = kernel_range 10 | self.shift = shift 11 | self.use_bn = use_bn 12 | self.use_res = use_res 13 | 14 | n_p = kernel_range[0] * kernel_range[1] 15 | self.k_p = nn.Conv2d(in_channels, in_channels * n_p, kernel_size=1) 16 | self.v_p = nn.Conv2d(in_channels, in_channels * n_p, kernel_size=1) 17 | self.q_p = nn.Conv2d(in_channels, in_channels, kernel_size=1) 18 | 19 | self.ff1 = nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=1) 20 | self.ff2 = nn.Conv2d(in_channels, 2 * in_channels, kernel_size=3, padding=1) 21 | self.ff3 = nn.Conv2d(2 * in_channels, in_channels, kernel_size=3, padding=1) 22 | 23 | if self.use_bn: 24 | self.bn1 = nn.BatchNorm2d(in_channels) 25 | self.bn2 = nn.BatchNorm2d(in_channels) 26 | 27 | def forward(self, x): 28 | h_half = self.kernel_range[0] // 2 29 | w_half = self.kernel_range[1] // 2 30 | s = x.shape 31 | D = s[1] 32 | # print("D:", D) 33 | x_k = self.k_p(x) 34 | x_v = self.v_p(x) 35 | x_q = self.q_p(x) 36 | # print("x_k.shape: ", x_k.shape) 37 | paddings = (w_half * self.shift, w_half * self.shift, h_half * self.shift, h_half * self.shift) 38 | # print("paddings: ", paddings) 39 | x_k = F.pad(x_k, paddings, "constant", 0) 40 | x_v = F.pad(x_v, paddings, "constant", 0) 41 | mask_x = torch.ones(s[0], 1, s[2], s[3], device=x.device) 42 | # print("x_k.shape: ", x_k.shape) 43 | # print("mask_x.shape: ", mask_x.shape) 44 | mask_pad = F.pad(mask_x, paddings, "constant", 0) 45 | 46 | k_ls = [] 47 | v_ls = [] 48 | masks = [] 49 | 50 | c_x, c_y = h_half * self.shift, w_half * self.shift 51 | layer = 0 52 | # print("s[2]: ", s[2]) 53 | # print("s[3]: ", s[3]) 54 | for i in range(-h_half, h_half + 1): 55 | for j in range(-w_half, w_half + 1): 56 | 57 | k_t = x_k[:, layer*D:(layer+1)*D, c_x + i * self.shift:c_x + i * self.shift + s[2], c_y + j * self.shift:c_y + j * self.shift + s[3]] 58 | k_ls.append(k_t) 59 | v_t = x_v[:, layer*D:(layer+1)*D, c_x + i * self.shift:c_x + i * self.shift + s[2], c_y + j * self.shift:c_y + j * self.shift + s[3]] 60 | v_ls.append(v_t) 61 | # print("mask_pad.shape:", mask_pad.shape) 62 | _m = mask_pad[:, :, c_x + i * self.shift:c_x + i * self.shift + s[2], c_y + j * self.shift:c_y + j * self.shift + s[3]] 63 | # print("v_t: ", v_t.shape) 64 | # print("_m: ", _m.shape) 65 | masks.append(_m) 66 | layer += 1 67 | # print("mask: ", masks[0].shape) 68 | m_stack = torch.stack(masks, dim=1) 69 | m_vec = m_stack.view(s[0] * s[2] * s[3], self.kernel_range[0] * self.kernel_range[1], 1) 70 | # print("len(k_ls): ", len(k_ls)) 71 | # print("k_ls[0]: ", k_ls[0].shape) 72 | k_stack = torch.stack(k_ls, dim=1) 73 | v_stack = torch.stack(v_ls, dim=1) 74 | # print("k_stack: ", k_stack.shape) 75 | k = k_stack.view(s[0] * s[2] * s[3] , self.kernel_range[0] * self.kernel_range[1], D) 76 | v = v_stack.view(s[0] * s[2] * s[3] , self.kernel_range[0] * self.kernel_range[1], D) 77 | q = x_q.view(s[0] * s[2] * s[3], 1, D) 78 | 79 | alpha = F.softmax(torch.matmul(k, q.transpose(-1, -2)) * m_vec / 8, dim=1) 80 | __res = torch.matmul(alpha.transpose(-1, -2), v) 81 | _res = __res.view(s[0], D, s[2], s[3]) 82 | 83 | if self.use_res: 84 | t = x + _res 85 | else: 86 | t = _res 87 | 88 | if self.use_bn: 89 | t = self.bn1(t) 90 | 91 | _t = t 92 | t = F.relu(self.ff1(t)) 93 | t = F.relu(self.ff2(t)) 94 | t = F.relu(self.ff3(t)) 95 | 96 | if self.use_res: 97 | t = _t + t 98 | 99 | if self.use_bn: 100 | res = self.bn2(t) 101 | else: 102 | res = t 103 | 104 | return res 105 | 106 | # 测试 PixelAttention 类 107 | if __name__ == "__main__": 108 | input_tensor = torch.randn(1, 3, 64, 64) # (batch_size, in_channels, height, width) 109 | pixel_attention = PixelAttention(3) 110 | output = pixel_attention(input_tensor) 111 | print("output:", output.shape) 112 | -------------------------------------------------------------------------------- /IMDLBenCo/modules/backbones/vit16.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import timm 4 | import torch.nn.functional as F 5 | from IMDLBenCo.registry import MODELS 6 | 7 | from ..extractors.high_frequency_feature_extraction import FFTExtractor, DCTExtractor 8 | from ..extractors.sobel import SobelFilter 9 | from ..extractors.bayar_conv import BayerConv 10 | from ..extractors.srm_filter import SRMConv2D 11 | 12 | class ViT(timm.models.vision_transformer.VisionTransformer): 13 | def __init__(self, input_head=None, num_channels=3,pretrained=True): 14 | super(ViT, self).__init__() 15 | if pretrained: 16 | model = timm.create_model('vit_base_patch16_224', pretrained=True) 17 | self.load_state_dict(model.state_dict()) 18 | original_first_layer = self.patch_embed.proj 19 | self.patch_embed.strict_img_size=False 20 | if input_head != None: 21 | self.input_head = input_head 22 | new_first_layer = nn.Conv2d(num_channels + 3, original_first_layer.out_channels, 23 | kernel_size=original_first_layer.kernel_size, stride=original_first_layer.stride, 24 | padding=original_first_layer.padding, bias=False) 25 | new_first_layer.weight.data[:, :3, :, :] = original_first_layer.weight.data.clone()[:, :3, :, :] 26 | if num_channels > 0: 27 | new_first_layer.weight.data[:, 3:, :, :] = torch.nn.init.kaiming_normal_(new_first_layer.weight[:, 3:, :, :]) 28 | self.patch_embed.proj = new_first_layer 29 | else: 30 | self.input_head = None 31 | new_length = 1024 32 | self.deconv_model = nn.Sequential( 33 | nn.ConvTranspose2d(768, 256, kernel_size=3, stride=2, padding=1, output_padding=1), # 输出: 1, 256, 64, 64 34 | nn.ReLU(), 35 | nn.ConvTranspose2d(256, 128, kernel_size=3, stride=2, padding=1, output_padding=1), # 输出: 1, 128, 128, 128 36 | nn.ReLU(), 37 | nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1, output_padding=1), # 输出: 1, 64, 256, 256 38 | nn.ReLU(), 39 | nn.ConvTranspose2d(64, 1, kernel_size=3, stride=2, padding=1, output_padding=1) # 输出: 1, 1, 512, 512 40 | ) 41 | 42 | original_pos_embed = model.pos_embed[:, 1:, :] 43 | new_pos_embed = F.interpolate(original_pos_embed.permute(0, 2, 1), size = new_length, mode='nearest').permute(0, 2, 1) 44 | 45 | # Create new positional embeddings tensor with space for class token 46 | new_pos_embed_with_cls = torch.zeros(1,new_length+1, 768, device=original_pos_embed.device) 47 | new_pos_embed_with_cls[:, 1:, :] = new_pos_embed # fill in the interpolated embeddings 48 | new_pos_embed_with_cls[:, 0, :] = model.pos_embed[:, 0, :] # copy the class token embeddings 49 | # 50 | # Replace the original positional embeddings with the new ones 51 | self.pos_embed = nn.Parameter(new_pos_embed_with_cls) 52 | self.loss_fun = nn.BCEWithLogitsLoss() 53 | 54 | 55 | def forward(self, image, mask, *args, **kwargs): 56 | if self.input_head != None: 57 | feature = self.input_head(image) 58 | input = torch.concat([image,feature],dim=1) 59 | else: 60 | input = image 61 | x = self.patch_embed(input) 62 | x = self._pos_embed(x) 63 | x = self.blocks(x) 64 | feature = self.norm(x) 65 | h = w = int((feature.shape[1])**0.5) 66 | feature = feature.permute(0,2,1)[:,:,1:1025].reshape(image.shape[0],768,h,w) 67 | mask_pred = self.deconv_model(feature) 68 | loss = self.loss_fun(mask_pred,mask) 69 | mask_pred = torch.sigmoid(mask_pred) 70 | output_dict = { 71 | # loss for backward 72 | "backward_loss": loss, 73 | # predicted mask, will calculate for metrics automatically 74 | "pred_mask": mask_pred, 75 | # predicted binaray label, will calculate for metrics automatically 76 | "pred_label": None, 77 | 78 | # ----values below is for visualization---- 79 | # automatically visualize with the key-value pairs 80 | "visual_loss": { 81 | "predict_loss": loss, 82 | }, 83 | 84 | "visual_image": { 85 | "pred_mask": mask_pred, 86 | } 87 | # ----------------------------------------- 88 | } 89 | return output_dict 90 | 91 | @MODELS.register_module() 92 | def vit16(): 93 | return ViT(None) 94 | 95 | @MODELS.register_module() 96 | def fft_vit16(): 97 | return ViT(FFTExtractor(), 3) 98 | 99 | @MODELS.register_module() 100 | def dct_vit16(): 101 | return ViT(DCTExtractor(), 3) 102 | 103 | @MODELS.register_module() 104 | def sobel_vit16(): 105 | return ViT(SobelFilter(), 1) 106 | 107 | @MODELS.register_module() 108 | def bayar_vit16(): 109 | return ViT(BayerConv(), 3) 110 | 111 | @MODELS.register_module() 112 | def srm_vit16(): 113 | return ViT(SRMConv2D(), 9) 114 | 115 | if __name__ == '__main__': 116 | from fvcore.nn import FlopCountAnalysis, parameter_count_table,flop_count_table 117 | 118 | model = ViT() 119 | flops = FlopCountAnalysis(model,(torch.randn(1,3,512,512),torch.randn(1,1,512,512))) 120 | # # 打印模型的参数信息 121 | print(flop_count_table(flops)) 122 | --------------------------------------------------------------------------------