├── README.md ├── assets ├── Hall.png ├── OFVL_overall.png ├── Parking_lot1.png ├── Parking_lot2.png ├── Room.png ├── laser_image.png └── overall.png ├── configs ├── _base_ │ └── default_runtime.py ├── ofvl_ms │ └── README.md └── ufvl_net │ ├── 12scenes.py │ ├── 7scenes.py │ └── README.md ├── model-index.yml ├── requirements ├── docs.txt ├── mminstall.txt ├── optional.txt ├── readthedocs.txt ├── runtime.txt └── tests.txt ├── setup.py ├── tools ├── integrate_params.py ├── slurm_test.sh └── test.py ├── ufvl_net.egg-info ├── PKG-INFO ├── SOURCES.txt ├── dependency_links.txt ├── not-zip-safe ├── requires.txt └── top_level.txt └── ufvl_net ├── .mim ├── configs ├── model-index.yml └── tools ├── __init__.py ├── __pycache__ ├── __init__.cpython-37.pyc └── version.cpython-37.pyc ├── apis ├── __init__.py ├── __pycache__ │ ├── __init__.cpython-37.pyc │ ├── inference.cpython-37.pyc │ ├── test.cpython-37.pyc │ └── train.cpython-37.pyc ├── inference.py ├── test.py └── train.py ├── cnn ├── __init__.py ├── __pycache__ │ ├── __init__.cpython-37.pyc │ ├── __init__.cpython-38.pyc │ ├── conv2d_share.cpython-37.pyc │ └── conv2d_share.cpython-38.pyc ├── cnn │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-37.pyc │ │ ├── __init__.cpython-38.pyc │ │ ├── conv2d_share.cpython-37.pyc │ │ └── conv2d_share.cpython-38.pyc │ └── conv2d_share.py └── conv2d_share.py ├── datasets ├── __init__.py ├── __pycache__ │ ├── __init__.cpython-37.pyc │ ├── base_dataset.cpython-37.pyc │ ├── builder.cpython-37.pyc │ ├── cambridge.cpython-37.pyc │ ├── cifar.cpython-37.pyc │ ├── cub.cpython-37.pyc │ ├── custom.cpython-37.pyc │ ├── dataset_wrappers.cpython-37.pyc │ ├── imagenet.cpython-37.pyc │ ├── imagenet21k.cpython-37.pyc │ ├── mnist.cpython-37.pyc │ ├── multi_label.cpython-37.pyc │ ├── seven_scenes.cpython-37.pyc │ ├── twelve_scenes.cpython-37.pyc │ ├── twelvescenes.cpython-37.pyc │ ├── utils.cpython-37.pyc │ └── voc.cpython-37.pyc ├── builder.py ├── dataset_wrappers.py ├── piplines │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-37.pyc │ │ ├── auto_augment.cpython-37.pyc │ │ ├── compose.cpython-37.pyc │ │ ├── formatting.cpython-37.pyc │ │ ├── loading.cpython-37.pyc │ │ └── transforms.cpython-37.pyc │ ├── auto_augment.py │ ├── compose.py │ ├── formatting.py │ ├── loading.py │ └── transforms.py ├── seven_scenes.py └── twelvescenes.py ├── models ├── __init__.py ├── __pycache__ │ ├── __init__.cpython-37.pyc │ └── builder.cpython-37.pyc ├── architecture │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-37.pyc │ │ ├── base.cpython-37.pyc │ │ └── vis_loc.cpython-37.pyc │ ├── base.py │ ├── pnpransac │ │ ├── build │ │ │ └── temp.linux-x86_64-cpython-37 │ │ │ │ └── pnpransacpy.o │ │ ├── pnpransac.cpp │ │ ├── pnpransac.cpython-37m-x86_64-linux-gnu.so │ │ ├── pnpransac.h │ │ ├── pnpransacpy.cpp │ │ ├── pnpransacpy.pxd │ │ ├── pnpransacpy.pyx │ │ └── setup.py │ └── vis_loc.py ├── backbones │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-37.pyc │ │ ├── base_backbone.cpython-37.pyc │ │ ├── resnet.cpython-37.pyc │ │ └── seresnet.cpython-37.pyc │ ├── base_backbone.py │ ├── resnet.py │ └── seresnet.py ├── builder.py ├── heads │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-37.pyc │ │ └── score_head.cpython-37.pyc │ └── score_head.py └── utils │ ├── __pycache__ │ ├── make_divisible.cpython-37.pyc │ └── se_layer.cpython-37.pyc │ ├── augment │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-37.pyc │ │ ├── augments.cpython-37.pyc │ │ ├── builder.cpython-37.pyc │ │ ├── cutmix.cpython-37.pyc │ │ ├── identity.cpython-37.pyc │ │ ├── mixup.cpython-37.pyc │ │ ├── resizemix.cpython-37.pyc │ │ └── utils.cpython-37.pyc │ ├── augments.py │ ├── builder.py │ ├── cutmix.py │ ├── identity.py │ ├── mixup.py │ ├── resizemix.py │ └── utils.py │ ├── make_divisible.py │ └── se_layer.py ├── utils ├── __init__.py ├── __pycache__ │ ├── __init__.cpython-37.pyc │ ├── collect_env.cpython-37.pyc │ ├── logger.cpython-37.pyc │ └── setup_env.cpython-37.pyc ├── collect_env.py ├── logger.py └── setup_env.py └── version.py /README.md: -------------------------------------------------------------------------------- 1 | # Visual localization with SCoRE methods 2 | *** This is a collection of our visual localization frameworks. 3 | 4 | > [**OFVL-MS**](./configs/ofvl_ms) (```@ICCV'23```): **OFVL-MS: Once for Visual Localization across Multiple Indoor Scenes** 5 | 6 | > [**UFVL-Net**](./configs/ufvl_net) (```@TIM'23```): **UFVL-Net: A Unified Framework for Visual Localization across Multiple Indoor Scenesss** 7 | 8 | ## Highlights 9 | - Once-for-multiple-scenes. 10 | Both OFVL-MS and UFVL-Net optimize visual localization tasks of various scenes collectively using a multi-task learning manner, which challenges the conventional wisdom that SCoRe typically trains a separate model for each scene. OFVL-MS realizes layer-wise parameters sharing, while UFVL-Net realizes channel-wise and kernel-wise sharing polices. 11 | 12 | - Competive performance. 13 | Both OFVL-MS and UFVL-Net deliver extraordinary performances on two benchmarks and complex real scenes. We demonstrate that once the training for our methods are done, our methods can generalize to new scenes with much fewer parameters by freezing the task-shared parameters. 14 | 15 | ## Environment Setup 16 | To set up the enviroment you can easily run the following command: 17 | - Create environment 18 | ```buildoutcfg 19 | conda create -n ufvlnet python=3.7 20 | conda activate ufvlnet 21 | ``` 22 | - Install torch, we verify UFVL-Net with pytorch 1.10.1 and cuda 11.3. 23 | ```buildoutcfg 24 | conda install pytorch==1.10.1 torchvision==0.11.2 torchaudio==0.10.1 cudatoolkit=11.3 -c pytorch -c conda-forge 25 | ``` 26 | - Subsequently, we need to build the cython module to install the PnP solver: 27 | ```buildoutcfg 28 | cd ./pnpransac 29 | rm -rf build 30 | python setup.py build_ext --inplace 31 | ``` 32 | - Install openmmlab packages. 33 | ```buildoutcfg 34 | pip install mmcv-full==1.5.0 35 | ``` 36 | - Compile the ufvl-net as a package. 37 | ```buildoutcfg 38 | cd ufvl_net 39 | pip install -e . 40 | cd .. 41 | export PYTHONPATH=./ufvl_net/ 42 | ``` 43 | ## Data Preparation 44 | We utilize two standard datasets (i.e, 7-Scenes and 12-Scenes) to evaluate our method. 45 | - 7-Scenes: The 7-Scenes dataset can be downloaded from [7-Scenes](https://www.microsoft.com/en-us/research/project/rgb-d-dataset-7-scenes/). 46 | - 12-Scenes: The 12-Scenes dataset can be downloaded from [12-Scenes](https://graphics.stanford.edu/projects/reloc/). 47 | - LIVL: The real-world LIVL dataset can be downloaded from [RealWorld-Scenes](https://drive.google.com/drive/folders/1rHILFijnb8wfQiT-5gWLvDJWbOqHMZwx). 48 | 49 | ## LIVL Dataset 50 | LIVL dataset collection equipment contains a mobile chassis, a RealSense D435 camera, and a VLP-16 laser radar. LIVL dataset records RGB-D images and corresponding camera poses of four different indoor environments. 51 | The dataset is available at [here](https://drive.google.com/drive/folders/1rHILFijnb8wfQiT-5gWLvDJWbOqHMZwx). 52 | Specifically, we utilize the ROS system to record RGB images and aligned depth images with corresponding timestamp $T_{1}$, Furthermore, we obtain point clouds with timestamp $T_{2}$ provided by VLP-16 laser radar. Then, we generate final RGB-D images and corresponding point clouds through aligning $T_{1}$ and $T_{2}$. Ultimately, We utilize the LiDAR-based SLAM system A-LOAM to compute the ground truth pose. 53 | For each scene, four sequences are recorded, in which three sequences are used for training and one sequence for testing. 54 | 55 | ![](https://github.com/mooncake199809/UFVL-Net/blob/main/assets/laser_image.png) 56 | - K544: a room spanning about $12 \times 9 m^{2}$ with $3109$ images for training and $1112$ images for testing. 57 | - Floor5: a hall spanning about $12 \times 5 m^{2}$ with $2694$ images for training and $869$ images for testing. 58 | - Parking lot1: a parking lot spanning about $8 \times 6 m^{2}$ with $2294$ images for training and $661$ images for testing. 59 | - Parking lot2: a parking lot spanning about $8 \times 8 m^{2}$ with $2415$ images for training and $875$ images for testing. 60 | -------------------------------------------------------------------------------- /assets/Hall.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/assets/Hall.png -------------------------------------------------------------------------------- /assets/OFVL_overall.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/assets/OFVL_overall.png -------------------------------------------------------------------------------- /assets/Parking_lot1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/assets/Parking_lot1.png -------------------------------------------------------------------------------- /assets/Parking_lot2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/assets/Parking_lot2.png -------------------------------------------------------------------------------- /assets/Room.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/assets/Room.png -------------------------------------------------------------------------------- /assets/laser_image.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/assets/laser_image.png -------------------------------------------------------------------------------- /assets/overall.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/assets/overall.png -------------------------------------------------------------------------------- /configs/_base_/default_runtime.py: -------------------------------------------------------------------------------- 1 | # checkpoint saving 2 | checkpoint_config = dict(interval=1) 3 | # yapf:disable 4 | log_config = dict( 5 | interval=10, 6 | hooks=[ 7 | dict(type='TextLoggerHook'), 8 | # dict(type='TensorboardLoggerHook') 9 | ]) 10 | # yapf:enable 11 | 12 | dist_params = dict(backend='nccl') 13 | log_level = 'INFO' 14 | load_from = None 15 | resume_from = None 16 | workflow = [('train', 1)] 17 | -------------------------------------------------------------------------------- /configs/ofvl_ms/README.md: -------------------------------------------------------------------------------- 1 | # OFVL-MS: Once for Visual Localization across Multiple Indoor Scenes 2 | ![overall](https://github.com/mooncake199809/UFVL-Net/blob/main/assets/OFVL_overall.png) 3 | 4 | ## Model Zoo 5 | 6 | Coming soon! -------------------------------------------------------------------------------- /configs/ufvl_net/12scenes.py: -------------------------------------------------------------------------------- 1 | _base_ = [ 2 | '../_base_/default_runtime.py' 3 | ] 4 | 5 | dataset_type = 'ufvl_net.TWESCENES' 6 | # root='/home/dk/SCORE_Methods/EAAINet/data/' 7 | root='/home/dk/SCORE_Methods/EAAINet/data/' 8 | scene='office1/lounge' 9 | all_scene = ['apt1kitchen','apt1living','apt2bed','apt2kitchen', 10 | 'apt2living','apt2luke','office1gates362', 11 | 'office1gates381','office1lounge','office1manolis', 12 | 'office25a','office25b'] 13 | share_type = "kernel" # channel or kernel 14 | 15 | test_pipeline = [] 16 | eval_pipeline = [] 17 | custom_imports = dict(imports=['ufvl_net.models']) 18 | 19 | data = dict( 20 | samples_per_gpu=1, 21 | workers_per_gpu=1, 22 | train=dict( 23 | type=dataset_type, 24 | root=root, 25 | scene=scene, 26 | split='train'), 27 | val=dict( 28 | type=dataset_type, 29 | root=root, 30 | scene=scene, 31 | split='test'), 32 | test=dict( 33 | type=dataset_type, 34 | root=root, 35 | scene=scene, 36 | split='test')) 37 | 38 | model = dict( 39 | type='ufvl_net.FDANET', 40 | backbone=dict( 41 | type='SEResNet', 42 | depth=50, 43 | stem_channels=16, 44 | expansion = 1, 45 | strides=(1, 1, 2, 2), 46 | use_maxpool=False, 47 | num_stages=4, 48 | # conv_cfg=dict(type='Conv2d_share'), 49 | out_indices=(3, ), 50 | style='pytorch', 51 | drop_path_rate=0.1, 52 | ), 53 | head=dict( 54 | type='RegHead', 55 | in_channel=512), 56 | dataset="12Scenes") 57 | 58 | scene=scene.replace("/","") -------------------------------------------------------------------------------- /configs/ufvl_net/7scenes.py: -------------------------------------------------------------------------------- 1 | _base_ = [ 2 | '../_base_/default_runtime.py' 3 | ] 4 | 5 | dataset_type = 'ufvl_net.SevenScenes' 6 | # root='/home/dk/SCORE_Methods/EAAINet/data/' 7 | root='/home/dk/SCORE_Methods/EAAINet/data/' 8 | scene='chess' 9 | all_scene = ['chess', 'fire', 'heads', 'pumpkin', 'redkitchen', 'office', 'stairs'] 10 | share_type = "channel" # channel or kernel 11 | 12 | test_pipeline = [] 13 | eval_pipeline = [] 14 | custom_imports = dict(imports=['ufvl_net.models']) 15 | 16 | data = dict( 17 | samples_per_gpu=1, 18 | workers_per_gpu=1, 19 | train=dict( 20 | type=dataset_type, 21 | root=root, 22 | scene=scene, 23 | split='train'), 24 | val=dict( 25 | type=dataset_type, 26 | root=root, 27 | scene=scene, 28 | split='test'), 29 | test=dict( 30 | type=dataset_type, 31 | root=root, 32 | scene=scene, 33 | split='test')) 34 | 35 | model = dict( 36 | type='ufvl_net.FDANET', 37 | backbone=dict( 38 | type='SEResNet', 39 | depth=34, 40 | stem_channels=16, 41 | expansion = 1, 42 | strides=(1, 1, 2, 2), 43 | use_maxpool=False, 44 | num_stages=4, 45 | # conv_cfg=dict(type='Conv2d_share'), 46 | out_indices=(3, ), 47 | style='pytorch', 48 | drop_path_rate=0.1, 49 | ), 50 | head=dict( 51 | type='RegHead', 52 | in_channel=512), 53 | dataset="7Scenes") 54 | -------------------------------------------------------------------------------- /configs/ufvl_net/README.md: -------------------------------------------------------------------------------- 1 | # UFVL-Net: A Unified Framework for Visual Localization across Multiple Indoor Scenes 2 | 3 | ![overall](https://github.com/mooncake199809/UFVL-Net/blob/main/assets/overall.png) 4 | 5 | ## Model Zoo 6 | For evaluation, we provide the checkpoints of 7-Scenes dataset in [Google Drive](https://drive.google.com/drive/folders/1l4vWMz7mo49R1gMBxl932-DdavfhxiBO). 7 | For evaluation, we also provide the checkpoints of 12-Scenes dataset in [Google Drive](https://drive.google.com/drive/folders/1Yw-DskJD7hCPo-WIXfPvHI5mP5UgRgJ9). 8 | - Note: We integrate these models into a single one. You can do the evaluation following the description in *Quick Start - Test*). 9 | 10 | ## Quick Start 11 | 12 | We provide *Test* code of UFVL-Net as follows: 13 | 14 | ### Test 15 | To test our trained models, you need to put the downloaded model in `./weights`. 16 | To test a specific model in a specific scene, you need to modify the config file in ./configs/ufvl_net/7scenes.py or ./configs/ufvl_net/12scenes.py 17 | The structure of the config file is described as follow: 18 | ```buildoutcfg 19 | dataset_type: 'ufvl_net.SevenScenes' or 'ufvl_net.TWESCENES' 20 | root: the root path of the dataset 21 | scene: the scene name that you want to test 22 | share_type: the type of weight sharing ("channel" or "kernel") 23 | data: the config of the dataset 24 | model: the config of the model 25 | ``` 26 | - If you want to test UFVL-Net-M with channel-wise sharing policy on the chess scene of 7-Scenes dataset, you need to modify the lines 8, 10, and 39 as "chess", "channel", and depth=34. Then, you could use the following command: 27 | ```buildoutcfg 28 | python tools/test.py ./configs/ufvl_net/7scenes.py ./weights/34_channel_7scenes.pth --metrics accuracy 29 | ``` -------------------------------------------------------------------------------- /model-index.yml: -------------------------------------------------------------------------------- 1 | Import: 2 | - configs/mobilenet_v2/metafile.yml 3 | - configs/resnet/metafile.yml 4 | - configs/res2net/metafile.yml 5 | - configs/resnext/metafile.yml 6 | - configs/seresnet/metafile.yml 7 | - configs/shufflenet_v1/metafile.yml 8 | - configs/shufflenet_v2/metafile.yml 9 | - configs/swin_transformer/metafile.yml 10 | - configs/vgg/metafile.yml 11 | - configs/repvgg/metafile.yml 12 | - configs/tnt/metafile.yml 13 | - configs/vision_transformer/metafile.yml 14 | - configs/t2t_vit/metafile.yml 15 | - configs/mlp_mixer/metafile.yml 16 | - configs/conformer/metafile.yml 17 | - configs/regnet/metafile.yml 18 | - configs/deit/metafile.yml 19 | - configs/twins/metafile.yml 20 | - configs/efficientnet/metafile.yml 21 | - configs/convnext/metafile.yml 22 | - configs/hrnet/metafile.yml 23 | - configs/repmlp/metafile.yml 24 | - configs/wrn/metafile.yml 25 | - configs/van/metafile.yml 26 | - configs/cspnet/metafile.yml 27 | - configs/convmixer/metafile.yml 28 | - configs/densenet/metafile.yml 29 | - configs/poolformer/metafile.yml 30 | -------------------------------------------------------------------------------- /requirements/docs.txt: -------------------------------------------------------------------------------- 1 | docutils==0.17.1 2 | myst-parser 3 | -e git+https://github.com/open-mmlab/pytorch_sphinx_theme.git#egg=pytorch_sphinx_theme 4 | sphinx==4.5.0 5 | sphinx-copybutton 6 | sphinx_markdown_tables 7 | -------------------------------------------------------------------------------- /requirements/mminstall.txt: -------------------------------------------------------------------------------- 1 | mmcv-full>=1.4.2,<1.6.0 2 | -------------------------------------------------------------------------------- /requirements/optional.txt: -------------------------------------------------------------------------------- 1 | albumentations>=0.3.2 --no-binary qudida,albumentations 2 | colorama 3 | requests 4 | rich 5 | -------------------------------------------------------------------------------- /requirements/readthedocs.txt: -------------------------------------------------------------------------------- 1 | mmcv>=1.4.2 2 | torch 3 | torchvision 4 | -------------------------------------------------------------------------------- /requirements/runtime.txt: -------------------------------------------------------------------------------- 1 | matplotlib 2 | numpy 3 | packaging 4 | -------------------------------------------------------------------------------- /requirements/tests.txt: -------------------------------------------------------------------------------- 1 | codecov 2 | flake8 3 | interrogate 4 | isort==4.3.21 5 | mmdet 6 | pytest 7 | xdoctest >= 0.10.0 8 | yapf 9 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | import os 2 | import os.path as osp 3 | import shutil 4 | import sys 5 | import warnings 6 | from setuptools import find_packages, setup 7 | 8 | 9 | def readme(): 10 | with open('README.md', encoding='utf-8') as f: 11 | content = f.read() 12 | return content 13 | 14 | 15 | def get_version(): 16 | version_file = '/home/dk/OFVL-VS2/mmcls/version.py' 17 | with open(version_file, 'r', encoding='utf-8') as f: 18 | exec(compile(f.read(), version_file, 'exec')) 19 | return locals()['__version__'] 20 | 21 | 22 | def parse_requirements(fname='requirements.txt', with_version=True): 23 | """Parse the package dependencies listed in a requirements file but strips 24 | specific versioning information. 25 | 26 | Args: 27 | fname (str): path to requirements file 28 | with_version (bool, default=False): if True include version specs 29 | 30 | Returns: 31 | List[str]: list of requirements items 32 | 33 | CommandLine: 34 | python -c "import setup; print(setup.parse_requirements())" 35 | """ 36 | import re 37 | import sys 38 | from os.path import exists 39 | require_fpath = fname 40 | 41 | def parse_line(line): 42 | """Parse information from a line in a requirements text file.""" 43 | if line.startswith('-r '): 44 | # Allow specifying requirements in other files 45 | target = line.split(' ')[1] 46 | for info in parse_require_file(target): 47 | yield info 48 | else: 49 | info = {'line': line} 50 | if line.startswith('-e '): 51 | info['package'] = line.split('#egg=')[1] 52 | else: 53 | # Remove versioning from the package 54 | pat = '(' + '|'.join(['>=', '==', '>']) + ')' 55 | parts = re.split(pat, line, maxsplit=1) 56 | parts = [p.strip() for p in parts] 57 | 58 | info['package'] = parts[0] 59 | if len(parts) > 1: 60 | op, rest = parts[1:] 61 | if ';' in rest: 62 | # Handle platform specific dependencies 63 | # http://setuptools.readthedocs.io/en/latest/setuptools.html#declaring-platform-specific-dependencies 64 | version, platform_deps = map(str.strip, 65 | rest.split(';')) 66 | info['platform_deps'] = platform_deps 67 | else: 68 | version = rest # NOQA 69 | if '--' in version: 70 | # the `extras_require` doesn't accept options. 71 | version = version.split('--')[0].strip() 72 | info['version'] = (op, version) 73 | yield info 74 | 75 | def parse_require_file(fpath): 76 | with open(fpath, 'r') as f: 77 | for line in f.readlines(): 78 | line = line.strip() 79 | if line and not line.startswith('#'): 80 | for info in parse_line(line): 81 | yield info 82 | 83 | def gen_packages_items(): 84 | if exists(require_fpath): 85 | for info in parse_require_file(require_fpath): 86 | parts = [info['package']] 87 | if with_version and 'version' in info: 88 | parts.extend(info['version']) 89 | if not sys.version.startswith('3.4'): 90 | # apparently package_deps are broken in 3.4 91 | platform_deps = info.get('platform_deps') 92 | if platform_deps is not None: 93 | parts.append(';' + platform_deps) 94 | item = ''.join(parts) 95 | yield item 96 | 97 | packages = list(gen_packages_items()) 98 | return packages 99 | 100 | 101 | def add_mim_extension(): 102 | """Add extra files that are required to support MIM into the package. 103 | 104 | These files will be added by creating a symlink to the originals if the 105 | package is installed in `editable` mode (e.g. pip install -e .), or by 106 | copying from the originals otherwise. 107 | """ 108 | 109 | # parse installment mode 110 | if 'develop' in sys.argv: 111 | # installed by `pip install -e .` 112 | mode = 'symlink' 113 | elif 'sdist' in sys.argv or 'bdist_wheel' in sys.argv: 114 | # installed by `pip install .` 115 | # or create source distribution by `python setup.py sdist` 116 | mode = 'copy' 117 | else: 118 | return 119 | 120 | filenames = ['tools', 'configs', 'model-index.yml'] 121 | repo_path = osp.dirname(__file__) 122 | mim_path = osp.join(repo_path, 'ufvl_net', '.mim') 123 | os.makedirs(mim_path, exist_ok=True) 124 | 125 | for filename in filenames: 126 | if osp.exists(filename): 127 | src_path = osp.join(repo_path, filename) 128 | tar_path = osp.join(mim_path, filename) 129 | 130 | if osp.isfile(tar_path) or osp.islink(tar_path): 131 | os.remove(tar_path) 132 | elif osp.isdir(tar_path): 133 | shutil.rmtree(tar_path) 134 | 135 | if mode == 'symlink': 136 | src_relpath = osp.relpath(src_path, osp.dirname(tar_path)) 137 | try: 138 | os.symlink(src_relpath, tar_path) 139 | except OSError: 140 | # Creating a symbolic link on windows may raise an 141 | # `OSError: [WinError 1314]` due to privilege. If 142 | # the error happens, the src file will be copied 143 | mode = 'copy' 144 | warnings.warn( 145 | f'Failed to create a symbolic link for {src_relpath}, ' 146 | f'and it will be copied to {tar_path}') 147 | else: 148 | continue 149 | 150 | if mode == 'copy': 151 | if osp.isfile(src_path): 152 | shutil.copyfile(src_path, tar_path) 153 | elif osp.isdir(src_path): 154 | shutil.copytree(src_path, tar_path) 155 | else: 156 | warnings.warn(f'Cannot copy file {src_path}.') 157 | else: 158 | raise ValueError(f'Invalid mode {mode}') 159 | 160 | 161 | if __name__ == '__main__': 162 | add_mim_extension() 163 | setup( 164 | name='ufvl_net', 165 | version=get_version(), 166 | description='OpenMMLab Image Classification Toolbox and Benchmark', 167 | long_description=readme(), 168 | long_description_content_type='text/markdown', 169 | keywords='computer vision, image classification', 170 | packages=find_packages(exclude=('configs', 'tools', 'demo')), 171 | include_package_data=True, 172 | classifiers=[ 173 | 'Development Status :: 4 - Beta', 174 | 'License :: OSI Approved :: Apache Software License', 175 | 'Operating System :: OS Independent', 176 | 'Programming Language :: Python :: 3', 177 | 'Programming Language :: Python :: 3.6', 178 | 'Programming Language :: Python :: 3.7', 179 | 'Programming Language :: Python :: 3.8', 180 | 'Programming Language :: Python :: 3.9', 181 | 'Topic :: Scientific/Engineering :: Artificial Intelligence', 182 | ], 183 | url='https://github.com/open-mmlab/mmclassification', 184 | author='MMClassification Contributors', 185 | author_email='openmmlab@gmail.com', 186 | license='Apache License 2.0', 187 | install_requires=parse_requirements('requirements/runtime.txt'), 188 | extras_require={ 189 | 'all': parse_requirements('requirements.txt'), 190 | 'tests': parse_requirements('requirements/tests.txt'), 191 | 'optional': parse_requirements('requirements/optional.txt'), 192 | 'mim': parse_requirements('requirements/mminstall.txt'), 193 | }, 194 | zip_safe=False) 195 | -------------------------------------------------------------------------------- /tools/slurm_test.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | set -x 4 | 5 | PARTITION=$1 6 | JOB_NAME=$2 7 | CONFIG=$3 8 | CHECKPOINT=$4 9 | GPUS=${GPUS:-8} 10 | GPUS_PER_NODE=${GPUS_PER_NODE:-8} 11 | CPUS_PER_TASK=${CPUS_PER_TASK:-5} 12 | PY_ARGS=${@:5} 13 | SRUN_ARGS=${SRUN_ARGS:-""} 14 | 15 | PYTHONPATH="$(dirname $0)/..":$PYTHONPATH \ 16 | srun -p ${PARTITION} \ 17 | --job-name=${JOB_NAME} \ 18 | --gres=gpu:${GPUS_PER_NODE} \ 19 | --ntasks=${GPUS} \ 20 | --ntasks-per-node=${GPUS_PER_NODE} \ 21 | --cpus-per-task=${CPUS_PER_TASK} \ 22 | --kill-on-bad-exit=1 \ 23 | ${SRUN_ARGS} \ 24 | python -u tools/test.py ${CONFIG} ${CHECKPOINT} --launcher="slurm" ${PY_ARGS} 25 | -------------------------------------------------------------------------------- /ufvl_net.egg-info/PKG-INFO: -------------------------------------------------------------------------------- 1 | Metadata-Version: 2.1 2 | Name: ufvl-net 3 | Version: 0.23.1 4 | Summary: OpenMMLab Image Classification Toolbox and Benchmark 5 | Home-page: https://github.com/open-mmlab/mmclassification 6 | Author: MMClassification Contributors 7 | Author-email: openmmlab@gmail.com 8 | License: Apache License 2.0 9 | Keywords: computer vision,image classification 10 | Classifier: Development Status :: 4 - Beta 11 | Classifier: License :: OSI Approved :: Apache Software License 12 | Classifier: Operating System :: OS Independent 13 | Classifier: Programming Language :: Python :: 3 14 | Classifier: Programming Language :: Python :: 3.6 15 | Classifier: Programming Language :: Python :: 3.7 16 | Classifier: Programming Language :: Python :: 3.8 17 | Classifier: Programming Language :: Python :: 3.9 18 | Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence 19 | Description-Content-Type: text/markdown 20 | Provides-Extra: all 21 | Provides-Extra: tests 22 | Provides-Extra: optional 23 | Provides-Extra: mim 24 | 25 |
26 | 27 | 28 |
 
29 |
30 | OpenMMLab website 31 | 32 | 33 | HOT 34 | 35 | 36 |      37 | OpenMMLab platform 38 | 39 | 40 | TRY IT OUT 41 | 42 | 43 |
44 |
 
45 | 46 | [![PyPI](https://img.shields.io/pypi/v/mmcls)](https://pypi.org/project/mmcls) 47 | [![Docs](https://img.shields.io/badge/docs-latest-blue)](https://mmclassification.readthedocs.io/en/latest/) 48 | [![Build Status](https://github.com/open-mmlab/mmclassification/workflows/build/badge.svg)](https://github.com/open-mmlab/mmclassification/actions) 49 | [![codecov](https://codecov.io/gh/open-mmlab/mmclassification/branch/master/graph/badge.svg)](https://codecov.io/gh/open-mmlab/mmclassification) 50 | [![license](https://img.shields.io/github/license/open-mmlab/mmclassification.svg)](https://github.com/open-mmlab/mmclassification/blob/master/LICENSE) 51 | [![open issues](https://isitmaintained.com/badge/open/open-mmlab/mmclassification.svg)](https://github.com/open-mmlab/mmclassification/issues) 52 | [![issue resolution](https://isitmaintained.com/badge/resolution/open-mmlab/mmclassification.svg)](https://github.com/open-mmlab/mmclassification/issues) 53 | 54 | [📘 Documentation](https://mmclassification.readthedocs.io/en/latest/) | 55 | [🛠️ Installation](https://mmclassification.readthedocs.io/en/latest/install.html) | 56 | [👀 Model Zoo](https://mmclassification.readthedocs.io/en/latest/model_zoo.html) | 57 | [🆕 Update News](https://mmclassification.readthedocs.io/en/latest/changelog.html) | 58 | [🤔 Reporting Issues](https://github.com/open-mmlab/mmclassification/issues/new/choose) 59 | 60 |
61 | 62 | ## Introduction 63 | 64 | English | [简体中文](/README_zh-CN.md) 65 | 66 | MMClassification is an open source image classification toolbox based on PyTorch. It is 67 | a part of the [OpenMMLab](https://openmmlab.com/) project. 68 | 69 | The master branch works with **PyTorch 1.5+**. 70 | 71 |
72 | 73 |
74 | 75 | ### Major features 76 | 77 | - Various backbones and pretrained models 78 | - Bag of training tricks 79 | - Large-scale training configs 80 | - High efficiency and extensibility 81 | - Powerful toolkits 82 | 83 | ## What's new 84 | 85 | v0.23.0 was released in 1/5/2022. 86 | Highlights of the new version: 87 | 88 | - Support **DenseNet**, **VAN** and **PoolFormer**, and provide pre-trained models. 89 | - Support training on IPU. 90 | - New style API docs, welcome [view it](https://mmclassification.readthedocs.io/en/master/api/models.html). 91 | 92 | v0.22.0 was released in 30/3/2022. 93 | 94 | Highlights of the new version: 95 | 96 | - Support a series of **CSP Network**, such as CSP-ResNet, CSP-ResNeXt and CSP-DarkNet. 97 | - A new `CustomDataset` class to help you **build dataset of yourself**! 98 | - Support new backbones - **ConvMixer**, **RepMLP** and new dataset - **CUB dataset**. 99 | 100 | Please refer to [changelog.md](docs/en/changelog.md) for more details and other release history. 101 | 102 | ## Installation 103 | 104 | Below are quick steps for installation: 105 | 106 | ```shell 107 | conda create -n open-mmlab python=3.8 pytorch=1.10 cudatoolkit=11.3 torchvision -c pytorch -y 108 | conda activate open-mmlab 109 | pip3 install openmim 110 | mim install mmcv-full 111 | git clone https://github.com/open-mmlab/mmclassification.git 112 | cd mmclassification 113 | pip3 install -e . 114 | ``` 115 | 116 | Please refer to [install.md](https://mmclassification.readthedocs.io/en/latest/install.html) for more detailed installation and dataset preparation. 117 | 118 | ## Getting Started 119 | 120 | Please see [Getting Started](https://mmclassification.readthedocs.io/en/latest/getting_started.html) for the basic usage of MMClassification. There are also tutorials: 121 | 122 | - [Learn about Configs](https://mmclassification.readthedocs.io/en/latest/tutorials/config.html) 123 | - [Fine-tune Models](https://mmclassification.readthedocs.io/en/latest/tutorials/finetune.html) 124 | - [Add New Dataset](https://mmclassification.readthedocs.io/en/latest/tutorials/new_dataset.html) 125 | - [Customizie Data Pipeline](https://mmclassification.readthedocs.io/en/latest/tutorials/data_pipeline.html) 126 | - [Add New Modules](https://mmclassification.readthedocs.io/en/latest/tutorials/new_modules.html) 127 | - [Customizie Schedule](https://mmclassification.readthedocs.io/en/latest/tutorials/schedule.html) 128 | - [Customizie Runtime Settings](https://mmclassification.readthedocs.io/en/latest/tutorials/runtime.html) 129 | 130 | Colab tutorials are also provided: 131 | 132 | - Learn about MMClassification **Python API**: [Preview the notebook](https://github.com/open-mmlab/mmclassification/blob/master/docs/en/tutorials/MMClassification_python.ipynb) or directly [run on Colab](https://colab.research.google.com/github/open-mmlab/mmclassification/blob/master/docs/en/tutorials/MMClassification_python.ipynb). 133 | - Learn about MMClassification **CLI tools**: [Preview the notebook](https://github.com/open-mmlab/mmclassification/blob/master/docs/en/tutorials/MMClassification_tools.ipynb) or directly [run on Colab](https://colab.research.google.com/github/open-mmlab/mmclassification/blob/master/docs/en/tutorials/MMClassification_tools.ipynb). 134 | 135 | ## Model zoo 136 | 137 | Results and models are available in the [model zoo](https://mmclassification.readthedocs.io/en/latest/model_zoo.html). 138 | 139 |
140 | Supported backbones 141 | 142 | - [x] [VGG](https://github.com/open-mmlab/mmclassification/tree/master/configs/vgg) 143 | - [x] [ResNet](https://github.com/open-mmlab/mmclassification/tree/master/configs/resnet) 144 | - [x] [ResNeXt](https://github.com/open-mmlab/mmclassification/tree/master/configs/resnext) 145 | - [x] [SE-ResNet](https://github.com/open-mmlab/mmclassification/tree/master/configs/seresnet) 146 | - [x] [SE-ResNeXt](https://github.com/open-mmlab/mmclassification/tree/master/configs/seresnet) 147 | - [x] [RegNet](https://github.com/open-mmlab/mmclassification/tree/master/configs/regnet) 148 | - [x] [ShuffleNetV1](https://github.com/open-mmlab/mmclassification/tree/master/configs/shufflenet_v1) 149 | - [x] [ShuffleNetV2](https://github.com/open-mmlab/mmclassification/tree/master/configs/shufflenet_v2) 150 | - [x] [MobileNetV2](https://github.com/open-mmlab/mmclassification/tree/master/configs/mobilenet_v2) 151 | - [x] [MobileNetV3](https://github.com/open-mmlab/mmclassification/tree/master/configs/mobilenet_v3) 152 | - [x] [Swin-Transformer](https://github.com/open-mmlab/mmclassification/tree/master/configs/swin_transformer) 153 | - [x] [RepVGG](https://github.com/open-mmlab/mmclassification/tree/master/configs/repvgg) 154 | - [x] [Vision-Transformer](https://github.com/open-mmlab/mmclassification/tree/master/configs/vision_transformer) 155 | - [x] [Transformer-in-Transformer](https://github.com/open-mmlab/mmclassification/tree/master/configs/tnt) 156 | - [x] [Res2Net](https://github.com/open-mmlab/mmclassification/tree/master/configs/res2net) 157 | - [x] [MLP-Mixer](https://github.com/open-mmlab/mmclassification/tree/master/configs/mlp_mixer) 158 | - [x] [DeiT](https://github.com/open-mmlab/mmclassification/tree/master/configs/deit) 159 | - [x] [Conformer](https://github.com/open-mmlab/mmclassification/tree/master/configs/conformer) 160 | - [x] [T2T-ViT](https://github.com/open-mmlab/mmclassification/tree/master/configs/t2t_vit) 161 | - [x] [Twins](https://github.com/open-mmlab/mmclassification/tree/master/configs/twins) 162 | - [x] [EfficientNet](https://github.com/open-mmlab/mmclassification/tree/master/configs/efficientnet) 163 | - [x] [ConvNeXt](https://github.com/open-mmlab/mmclassification/tree/master/configs/convnext) 164 | - [x] [HRNet](https://github.com/open-mmlab/mmclassification/tree/master/configs/hrnet) 165 | - [x] [VAN](https://github.com/open-mmlab/mmclassification/tree/master/configs/van) 166 | - [x] [ConvMixer](https://github.com/open-mmlab/mmclassification/tree/master/configs/convmixer) 167 | - [x] [CSPNet](https://github.com/open-mmlab/mmclassification/tree/master/configs/cspnet) 168 | - [x] [PoolFormer](https://github.com/open-mmlab/mmclassification/tree/master/configs/poolformer) 169 | 170 |
171 | 172 | ## Contributing 173 | 174 | We appreciate all contributions to improve MMClassification. 175 | Please refer to [CONTRUBUTING.md](https://mmclassification.readthedocs.io/en/latest/community/CONTRIBUTING.html) for the contributing guideline. 176 | 177 | ## Acknowledgement 178 | 179 | MMClassification is an open source project that is contributed by researchers and engineers from various colleges and companies. We appreciate all the contributors who implement their methods or add new features, as well as users who give valuable feedbacks. 180 | We wish that the toolbox and benchmark could serve the growing research community by providing a flexible toolkit to reimplement existing methods and develop their own new classifiers. 181 | 182 | ## Citation 183 | 184 | If you find this project useful in your research, please consider cite: 185 | 186 | ```BibTeX 187 | @misc{2020mmclassification, 188 | title={OpenMMLab's Image Classification Toolbox and Benchmark}, 189 | author={MMClassification Contributors}, 190 | howpublished = {\url{https://github.com/open-mmlab/mmclassification}}, 191 | year={2020} 192 | } 193 | ``` 194 | 195 | ## License 196 | 197 | This project is released under the [Apache 2.0 license](LICENSE). 198 | 199 | ## Projects in OpenMMLab 200 | 201 | - [MMCV](https://github.com/open-mmlab/mmcv): OpenMMLab foundational library for computer vision. 202 | - [MIM](https://github.com/open-mmlab/mim): MIM installs OpenMMLab packages. 203 | - [MMClassification](https://github.com/open-mmlab/mmclassification): OpenMMLab image classification toolbox and benchmark. 204 | - [MMDetection](https://github.com/open-mmlab/mmdetection): OpenMMLab detection toolbox and benchmark. 205 | - [MMDetection3D](https://github.com/open-mmlab/mmdetection3d): OpenMMLab's next-generation platform for general 3D object detection. 206 | - [MMRotate](https://github.com/open-mmlab/mmrotate): OpenMMLab rotated object detection toolbox and benchmark. 207 | - [MMSegmentation](https://github.com/open-mmlab/mmsegmentation): OpenMMLab semantic segmentation toolbox and benchmark. 208 | - [MMOCR](https://github.com/open-mmlab/mmocr): OpenMMLab text detection, recognition, and understanding toolbox. 209 | - [MMPose](https://github.com/open-mmlab/mmpose): OpenMMLab pose estimation toolbox and benchmark. 210 | - [MMHuman3D](https://github.com/open-mmlab/mmhuman3d): OpenMMLab 3D human parametric model toolbox and benchmark. 211 | - [MMSelfSup](https://github.com/open-mmlab/mmselfsup): OpenMMLab self-supervised learning toolbox and benchmark. 212 | - [MMRazor](https://github.com/open-mmlab/mmrazor): OpenMMLab model compression toolbox and benchmark. 213 | - [MMFewShot](https://github.com/open-mmlab/mmfewshot): OpenMMLab fewshot learning toolbox and benchmark. 214 | - [MMAction2](https://github.com/open-mmlab/mmaction2): OpenMMLab's next-generation action understanding toolbox and benchmark. 215 | - [MMTracking](https://github.com/open-mmlab/mmtracking): OpenMMLab video perception toolbox and benchmark. 216 | - [MMFlow](https://github.com/open-mmlab/mmflow): OpenMMLab optical flow toolbox and benchmark. 217 | - [MMEditing](https://github.com/open-mmlab/mmediting): OpenMMLab image and video editing toolbox. 218 | - [MMGeneration](https://github.com/open-mmlab/mmgeneration): OpenMMLab image and video generative models toolbox. 219 | - [MMDeploy](https://github.com/open-mmlab/mmdeploy): OpenMMLab model deployment framework. 220 | -------------------------------------------------------------------------------- /ufvl_net.egg-info/SOURCES.txt: -------------------------------------------------------------------------------- 1 | README.md 2 | setup.py 3 | ufvl_net/__init__.py 4 | ufvl_net/version.py 5 | ufvl_net.egg-info/PKG-INFO 6 | ufvl_net.egg-info/SOURCES.txt 7 | ufvl_net.egg-info/dependency_links.txt 8 | ufvl_net.egg-info/not-zip-safe 9 | ufvl_net.egg-info/requires.txt 10 | ufvl_net.egg-info/top_level.txt 11 | ufvl_net/apis/__init__.py 12 | ufvl_net/apis/inference.py 13 | ufvl_net/apis/test.py 14 | ufvl_net/apis/train.py 15 | ufvl_net/datasets/__init__.py 16 | ufvl_net/datasets/builder.py 17 | ufvl_net/datasets/seven_scenes.py 18 | ufvl_net/models/__init__.py 19 | ufvl_net/models/builder.py 20 | ufvl_net/models/architecture/__init__.py 21 | ufvl_net/models/architecture/vis_loc.py 22 | ufvl_net/models/backbones/__init__.py 23 | ufvl_net/models/backbones/resnet.py 24 | ufvl_net/models/backbones/seresnet.py 25 | ufvl_net/models/heads/__init__.py 26 | ufvl_net/models/heads/score_head.py 27 | ufvl_net/utils/__init__.py 28 | ufvl_net/utils/collect_env.py 29 | ufvl_net/utils/logger.py 30 | ufvl_net/utils/setup_env.py -------------------------------------------------------------------------------- /ufvl_net.egg-info/dependency_links.txt: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /ufvl_net.egg-info/not-zip-safe: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /ufvl_net.egg-info/requires.txt: -------------------------------------------------------------------------------- 1 | matplotlib 2 | numpy 3 | packaging 4 | 5 | [all] 6 | 7 | [mim] 8 | mmcv-full<1.6.0,>=1.4.2 9 | 10 | [optional] 11 | albumentations>=0.3.2 12 | colorama 13 | requests 14 | rich 15 | 16 | [tests] 17 | codecov 18 | flake8 19 | interrogate 20 | isort==4.3.21 21 | mmdet 22 | pytest 23 | xdoctest>=0.10.0 24 | yapf 25 | -------------------------------------------------------------------------------- /ufvl_net.egg-info/top_level.txt: -------------------------------------------------------------------------------- 1 | ufvl_net 2 | -------------------------------------------------------------------------------- /ufvl_net/.mim/configs: -------------------------------------------------------------------------------- 1 | ../../configs -------------------------------------------------------------------------------- /ufvl_net/.mim/model-index.yml: -------------------------------------------------------------------------------- 1 | ../../model-index.yml -------------------------------------------------------------------------------- /ufvl_net/.mim/tools: -------------------------------------------------------------------------------- 1 | ../../tools -------------------------------------------------------------------------------- /ufvl_net/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) OpenMMLab. All rights reserved. 2 | import warnings 3 | 4 | import mmcv 5 | from packaging.version import parse 6 | 7 | from .version import __version__ 8 | 9 | 10 | def digit_version(version_str: str, length: int = 4): 11 | """Convert a version string into a tuple of integers. 12 | 13 | This method is usually used for comparing two versions. For pre-release 14 | versions: alpha < beta < rc. 15 | 16 | Args: 17 | version_str (str): The version string. 18 | length (int): The maximum number of version levels. Default: 4. 19 | 20 | Returns: 21 | tuple[int]: The version info in digits (integers). 22 | """ 23 | version = parse(version_str) 24 | assert version.release, f'failed to parse version {version_str}' 25 | release = list(version.release) 26 | release = release[:length] 27 | if len(release) < length: 28 | release = release + [0] * (length - len(release)) 29 | if version.is_prerelease: 30 | mapping = {'a': -3, 'b': -2, 'rc': -1} 31 | val = -4 32 | # version.pre can be None 33 | if version.pre: 34 | if version.pre[0] not in mapping: 35 | warnings.warn(f'unknown prerelease version {version.pre[0]}, ' 36 | 'version checking may go wrong') 37 | else: 38 | val = mapping[version.pre[0]] 39 | release.extend([val, version.pre[-1]]) 40 | else: 41 | release.extend([val, 0]) 42 | 43 | elif version.is_postrelease: 44 | release.extend([1, version.post]) 45 | else: 46 | release.extend([0, 0]) 47 | return tuple(release) 48 | 49 | 50 | mmcv_minimum_version = '1.4.2' 51 | mmcv_maximum_version = '1.6.0' 52 | mmcv_version = digit_version(mmcv.__version__) 53 | 54 | 55 | assert (mmcv_version >= digit_version(mmcv_minimum_version) 56 | and mmcv_version <= digit_version(mmcv_maximum_version)), \ 57 | f'MMCV=={mmcv.__version__} is used but incompatible. ' \ 58 | f'Please install mmcv>={mmcv_minimum_version}, <={mmcv_maximum_version}.' 59 | 60 | __all__ = ['__version__', 'digit_version'] 61 | -------------------------------------------------------------------------------- /ufvl_net/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/__pycache__/version.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/__pycache__/version.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/apis/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) OpenMMLab. All rights reserved. 2 | from .inference import inference_model, init_model, show_result_pyplot 3 | from .test import multi_gpu_test, single_gpu_test 4 | from .train import init_random_seed, set_random_seed 5 | 6 | __all__ = [ 7 | 'set_random_seed', 'init_model', 'inference_model', 8 | 'multi_gpu_test', 'single_gpu_test', 'show_result_pyplot', 9 | 'init_random_seed' 10 | ] 11 | -------------------------------------------------------------------------------- /ufvl_net/apis/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/apis/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/apis/__pycache__/inference.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/apis/__pycache__/inference.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/apis/__pycache__/test.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/apis/__pycache__/test.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/apis/__pycache__/train.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/apis/__pycache__/train.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/apis/inference.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) OpenMMLab. All rights reserved. 2 | import warnings 3 | 4 | import mmcv 5 | import numpy as np 6 | import torch 7 | from mmcv.parallel import collate, scatter 8 | from mmcv.runner import load_checkpoint 9 | 10 | from ufvl_net.datasets.piplines import Compose 11 | from ufvl_net.models import build_architecture 12 | 13 | 14 | def init_model(config, checkpoint=None, device='cuda:0', options=None): 15 | """Initialize a classifier from config file. 16 | 17 | Args: 18 | config (str or :obj:`mmcv.Config`): Config file path or the config 19 | object. 20 | checkpoint (str, optional): Checkpoint path. If left as None, the model 21 | will not load any weights. 22 | options (dict): Options to override some settings in the used config. 23 | 24 | Returns: 25 | nn.Module: The constructed classifier. 26 | """ 27 | if isinstance(config, str): 28 | config = mmcv.Config.fromfile(config) 29 | elif not isinstance(config, mmcv.Config): 30 | raise TypeError('config must be a filename or Config object, ' 31 | f'but got {type(config)}') 32 | if options is not None: 33 | config.merge_from_dict(options) 34 | config.model.pretrained = None 35 | model = build_classifier(config.model) 36 | if checkpoint is not None: 37 | # Mapping the weights to GPU may cause unexpected video memory leak 38 | # which refers to https://github.com/open-mmlab/mmdetection/pull/6405 39 | checkpoint = load_checkpoint(model, checkpoint, map_location='cpu') 40 | if 'CLASSES' in checkpoint.get('meta', {}): 41 | model.CLASSES = checkpoint['meta']['CLASSES'] 42 | else: 43 | from mmcls.datasets import ImageNet 44 | warnings.simplefilter('once') 45 | warnings.warn('Class names are not saved in the checkpoint\'s ' 46 | 'meta data, use imagenet by default.') 47 | model.CLASSES = ImageNet.CLASSES 48 | model.cfg = config # save the config in the model for convenience 49 | model.to(device) 50 | model.eval() 51 | return model 52 | 53 | 54 | def inference_model(model, img): 55 | """Inference image(s) with the classifier. 56 | 57 | Args: 58 | model (nn.Module): The loaded classifier. 59 | img (str/ndarray): The image filename or loaded image. 60 | 61 | Returns: 62 | result (dict): The classification results that contains 63 | `class_name`, `pred_label` and `pred_score`. 64 | """ 65 | cfg = model.cfg 66 | device = next(model.parameters()).device # model device 67 | # build the data pipeline 68 | if isinstance(img, str): 69 | if cfg.data.test.pipeline[0]['type'] != 'LoadImageFromFile': 70 | cfg.data.test.pipeline.insert(0, dict(type='LoadImageFromFile')) 71 | data = dict(img_info=dict(filename=img), img_prefix=None) 72 | else: 73 | if cfg.data.test.pipeline[0]['type'] == 'LoadImageFromFile': 74 | cfg.data.test.pipeline.pop(0) 75 | data = dict(img=img) 76 | test_pipeline = Compose(cfg.data.test.pipeline) 77 | data = test_pipeline(data) 78 | data = collate([data], samples_per_gpu=1) 79 | if next(model.parameters()).is_cuda: 80 | # scatter to specified GPU 81 | data = scatter(data, [device])[0] 82 | 83 | # forward the model 84 | with torch.no_grad(): 85 | scores = model(return_loss=False, **data) 86 | pred_score = np.max(scores, axis=1)[0] 87 | pred_label = np.argmax(scores, axis=1)[0] 88 | result = {'pred_label': pred_label, 'pred_score': float(pred_score)} 89 | result['pred_class'] = model.CLASSES[result['pred_label']] 90 | return result 91 | 92 | 93 | def show_result_pyplot(model, 94 | img, 95 | result, 96 | fig_size=(15, 10), 97 | title='result', 98 | wait_time=0): 99 | """Visualize the classification results on the image. 100 | 101 | Args: 102 | model (nn.Module): The loaded classifier. 103 | img (str or np.ndarray): Image filename or loaded image. 104 | result (list): The classification result. 105 | fig_size (tuple): Figure size of the pyplot figure. 106 | Defaults to (15, 10). 107 | title (str): Title of the pyplot figure. 108 | Defaults to 'result'. 109 | wait_time (int): How many seconds to display the image. 110 | Defaults to 0. 111 | """ 112 | if hasattr(model, 'module'): 113 | model = model.module 114 | model.show_result( 115 | img, 116 | result, 117 | show=True, 118 | fig_size=fig_size, 119 | win_name=title, 120 | wait_time=wait_time) 121 | -------------------------------------------------------------------------------- /ufvl_net/apis/test.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) OpenMMLab. All rights reserved. 2 | import os.path as osp 3 | import pickle 4 | import shutil 5 | import tempfile 6 | import time 7 | 8 | import mmcv 9 | import numpy as np 10 | import torch 11 | import torch.distributed as dist 12 | from mmcv.image import tensor2imgs 13 | from mmcv.runner import get_dist_info 14 | 15 | 16 | def single_gpu_test(model, 17 | data_loader, 18 | show=False, 19 | out_dir=None, 20 | **show_kwargs): 21 | """Test model with local single gpu. 22 | 23 | This method tests model with a single gpu and supports showing results. 24 | 25 | Args: 26 | model (:obj:`torch.nn.Module`): Model to be tested. 27 | data_loader (:obj:`torch.utils.data.DataLoader`): Pytorch data loader. 28 | show (bool): Whether to show the test results. Defaults to False. 29 | out_dir (str): The output directory of result plots of all samples. 30 | Defaults to None, which means not to write output files. 31 | **show_kwargs: Any other keyword arguments for showing results. 32 | 33 | Returns: 34 | list: The prediction results. 35 | """ 36 | model.eval() 37 | results = [] 38 | transl_err_list = [] 39 | rot_err_list = [] 40 | time_list = [] 41 | for i, data in enumerate(data_loader): 42 | with torch.no_grad(): 43 | start_time = time.time() 44 | result = model(return_loss=False, **data) 45 | time_list.append(time.time() - start_time) 46 | print(np.min(np.array(time_list))) 47 | 48 | transl_err_list.append(result['trans_error_med']) 49 | rot_err_list.append(result['rot_err_med']) 50 | res_ = np.array([transl_err_list, rot_err_list]).T 51 | print('Median pose error: {}m, {}\u00b0'.format(np.median(res_[:, 0]), 52 | np.median(res_[:, 1]))) 53 | print('Average pose error: {}m, {}\u00b0'.format(np.mean(res_[:, 0]), 54 | np.mean(res_[:, 1]))) 55 | print('stddev: {}m, {}\u00b0'.format(np.std(res_[:, 0], ddof=1), 56 | np.std(res_[:, 1], ddof=1))) 57 | print("ACC:{}", np.sum((res_[:, 0] <= 0.050) * (res_[:, 1] <= 5)) * 1. / len(res_) * 100) 58 | results.append(result) 59 | return results 60 | 61 | 62 | def multi_gpu_test(model, data_loader, tmpdir=None, gpu_collect=False): 63 | """Test model with multiple gpus. 64 | 65 | This method tests model with multiple gpus and collects the results 66 | under two different modes: gpu and cpu modes. By setting 'gpu_collect=True' 67 | it encodes results to gpu tensors and use gpu communication for results 68 | collection. On cpu mode it saves the results on different gpus to 'tmpdir' 69 | and collects them by the rank 0 worker. 70 | 71 | Args: 72 | model (nn.Module): Model to be tested. 73 | data_loader (nn.Dataloader): Pytorch data loader. 74 | tmpdir (str): Path of directory to save the temporary results from 75 | different gpus under cpu mode. 76 | gpu_collect (bool): Option to use either gpu or cpu to collect results. 77 | 78 | Returns: 79 | list: The prediction results. 80 | """ 81 | model.eval() 82 | results = [] 83 | dataset = data_loader.dataset 84 | rank, world_size = get_dist_info() 85 | if rank == 0: 86 | # Check if tmpdir is valid for cpu_collect 87 | if (not gpu_collect) and (tmpdir is not None and osp.exists(tmpdir)): 88 | raise OSError((f'The tmpdir {tmpdir} already exists.', 89 | ' Since tmpdir will be deleted after testing,', 90 | ' please make sure you specify an empty one.')) 91 | prog_bar = mmcv.ProgressBar(len(dataset)) 92 | time.sleep(2) 93 | dist.barrier() 94 | for i, data in enumerate(data_loader): 95 | with torch.no_grad(): 96 | result = model(return_loss=False, **data) 97 | if isinstance(result, list): 98 | results.extend(result) 99 | else: 100 | results.append(result) 101 | 102 | if rank == 0: 103 | batch_size = data['img'].size(0) 104 | for _ in range(batch_size * world_size): 105 | prog_bar.update() 106 | 107 | # collect results from all ranks 108 | if gpu_collect: 109 | results = collect_results_gpu(results, len(dataset)) 110 | else: 111 | results = collect_results_cpu(results, len(dataset), tmpdir) 112 | return results 113 | 114 | 115 | def collect_results_cpu(result_part, size, tmpdir=None): 116 | rank, world_size = get_dist_info() 117 | # create a tmp dir if it is not specified 118 | if tmpdir is None: 119 | MAX_LEN = 512 120 | # 32 is whitespace 121 | dir_tensor = torch.full((MAX_LEN, ), 122 | 32, 123 | dtype=torch.uint8, 124 | device='cuda') 125 | if rank == 0: 126 | mmcv.mkdir_or_exist('.dist_test') 127 | tmpdir = tempfile.mkdtemp(dir='.dist_test') 128 | tmpdir = torch.tensor( 129 | bytearray(tmpdir.encode()), dtype=torch.uint8, device='cuda') 130 | dir_tensor[:len(tmpdir)] = tmpdir 131 | dist.broadcast(dir_tensor, 0) 132 | tmpdir = dir_tensor.cpu().numpy().tobytes().decode().rstrip() 133 | else: 134 | mmcv.mkdir_or_exist(tmpdir) 135 | # dump the part result to the dir 136 | mmcv.dump(result_part, osp.join(tmpdir, f'part_{rank}.pkl')) 137 | dist.barrier() 138 | # collect all parts 139 | if rank != 0: 140 | return None 141 | else: 142 | # load results of all parts from tmp dir 143 | part_list = [] 144 | for i in range(world_size): 145 | part_file = osp.join(tmpdir, f'part_{i}.pkl') 146 | part_result = mmcv.load(part_file) 147 | part_list.append(part_result) 148 | # sort the results 149 | ordered_results = [] 150 | for res in zip(*part_list): 151 | ordered_results.extend(list(res)) 152 | # the dataloader may pad some samples 153 | ordered_results = ordered_results[:size] 154 | # remove tmp dir 155 | shutil.rmtree(tmpdir) 156 | return ordered_results 157 | 158 | 159 | def collect_results_gpu(result_part, size): 160 | rank, world_size = get_dist_info() 161 | # dump result part to tensor with pickle 162 | part_tensor = torch.tensor( 163 | bytearray(pickle.dumps(result_part)), dtype=torch.uint8, device='cuda') 164 | # gather all result part tensor shape 165 | shape_tensor = torch.tensor(part_tensor.shape, device='cuda') 166 | shape_list = [shape_tensor.clone() for _ in range(world_size)] 167 | dist.all_gather(shape_list, shape_tensor) 168 | # padding result part tensor to max length 169 | shape_max = torch.tensor(shape_list).max() 170 | part_send = torch.zeros(shape_max, dtype=torch.uint8, device='cuda') 171 | part_send[:shape_tensor[0]] = part_tensor 172 | part_recv_list = [ 173 | part_tensor.new_zeros(shape_max) for _ in range(world_size) 174 | ] 175 | # gather all result part 176 | dist.all_gather(part_recv_list, part_send) 177 | 178 | if rank == 0: 179 | part_list = [] 180 | for recv, shape in zip(part_recv_list, shape_list): 181 | part_result = pickle.loads(recv[:shape[0]].cpu().numpy().tobytes()) 182 | part_list.append(part_result) 183 | # sort the results 184 | ordered_results = [] 185 | for res in zip(*part_list): 186 | ordered_results.extend(list(res)) 187 | # the dataloader may pad some samples 188 | ordered_results = ordered_results[:size] 189 | return ordered_results 190 | -------------------------------------------------------------------------------- /ufvl_net/apis/train.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) OpenMMLab. All rights reserved. 2 | import random 3 | import warnings 4 | 5 | import numpy as np 6 | import torch 7 | import torch.distributed as dist 8 | from mmcv.parallel import MMDataParallel, MMDistributedDataParallel 9 | from mmcv.runner import (DistSamplerSeedHook, Fp16OptimizerHook, 10 | build_optimizer, build_runner, get_dist_info) 11 | from mmcv.runner.hooks import DistEvalHook, EvalHook 12 | 13 | from ufvl_net.datasets import build_dataloader, build_dataset 14 | from ufvl_net.utils import get_root_logger 15 | 16 | 17 | def init_random_seed(seed=None, device='cuda'): 18 | """Initialize random seed. 19 | 20 | If the seed is not set, the seed will be automatically randomized, 21 | and then broadcast to all processes to prevent some potential bugs. 22 | 23 | Args: 24 | seed (int, Optional): The seed. Default to None. 25 | device (str): The device where the seed will be put on. 26 | Default to 'cuda'. 27 | 28 | Returns: 29 | int: Seed to be used. 30 | """ 31 | if seed is not None: 32 | return seed 33 | 34 | # Make sure all ranks share the same random seed to prevent 35 | # some potential bugs. Please refer to 36 | # https://github.com/open-mmlab/mmdetection/issues/6339 37 | rank, world_size = get_dist_info() 38 | seed = np.random.randint(2**31) 39 | if world_size == 1: 40 | return seed 41 | 42 | if rank == 0: 43 | random_num = torch.tensor(seed, dtype=torch.int32, device=device) 44 | else: 45 | random_num = torch.tensor(0, dtype=torch.int32, device=device) 46 | dist.broadcast(random_num, src=0) 47 | return random_num.item() 48 | 49 | 50 | def set_random_seed(seed, deterministic=False): 51 | """Set random seed. 52 | 53 | Args: 54 | seed (int): Seed to be used. 55 | deterministic (bool): Whether to set the deterministic option for 56 | CUDNN backend, i.e., set `torch.backends.cudnn.deterministic` 57 | to True and `torch.backends.cudnn.benchmark` to False. 58 | Default: False. 59 | """ 60 | random.seed(seed) 61 | np.random.seed(seed) 62 | torch.manual_seed(seed) 63 | torch.cuda.manual_seed_all(seed) 64 | if deterministic: 65 | torch.backends.cudnn.deterministic = True 66 | torch.backends.cudnn.benchmark = False 67 | 68 | -------------------------------------------------------------------------------- /ufvl_net/cnn/__init__.py: -------------------------------------------------------------------------------- 1 | from .conv2d_share import Conv2d_share, BatchNorm2d_share 2 | 3 | __all__ = [ 4 | 'Conv2d_share', 'BatchNorm2d_share' 5 | ] 6 | -------------------------------------------------------------------------------- /ufvl_net/cnn/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/cnn/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/cnn/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/cnn/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /ufvl_net/cnn/__pycache__/conv2d_share.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/cnn/__pycache__/conv2d_share.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/cnn/__pycache__/conv2d_share.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/cnn/__pycache__/conv2d_share.cpython-38.pyc -------------------------------------------------------------------------------- /ufvl_net/cnn/cnn/__init__.py: -------------------------------------------------------------------------------- 1 | from .conv2d_share import Conv2d_share, BatchNorm2d_share 2 | 3 | __all__ = [ 4 | 'Conv2d_share', 'BatchNorm2d_share' 5 | ] 6 | -------------------------------------------------------------------------------- /ufvl_net/cnn/cnn/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/cnn/cnn/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/cnn/cnn/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/cnn/cnn/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /ufvl_net/cnn/cnn/__pycache__/conv2d_share.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/cnn/cnn/__pycache__/conv2d_share.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/cnn/cnn/__pycache__/conv2d_share.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/cnn/cnn/__pycache__/conv2d_share.cpython-38.pyc -------------------------------------------------------------------------------- /ufvl_net/cnn/cnn/conv2d_share.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from mmcv.cnn import CONV_LAYERS, NORM_LAYERS 5 | 6 | THRESHOLD = 0.5 7 | 8 | 9 | @CONV_LAYERS.register_module() 10 | class Conv2d_share(nn.Conv2d): 11 | def __init__(self, in_channels, out_channels, kernel_size, share_type=None, 12 | stride=1, padding=0, dilation=1, groups=1, bias=True, 13 | padding_mode='zeros'): 14 | super().__init__(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias, padding_mode 15 | ) 16 | self.share_type = 'channel_wise' 17 | self.register_parameter('specific_weight', nn.Parameter(self.weight.clone())) 18 | if self.bias is not None: 19 | self.register_parameter('specific_bias', nn.Parameter(self.bias.clone())) 20 | 21 | # TODO: ablate one, zero and random 22 | # import pdb; pdb.set_trace() 23 | if self.share_type == 'channel_wise': 24 | self.score = nn.Parameter(torch.rand(self.weight.size(1)).cuda()) 25 | elif self.share_type == 'kernel_wise': 26 | self.score = nn.Parameter(torch.rand((self.weight.size(2), self.weight.size(3))).cuda()) 27 | else: 28 | self.score = nn.Parameter(torch.rand(1).cuda()) 29 | 30 | def forward(self, input): 31 | if self.share_type == 'channel_wise': 32 | score = binarizer_fn(self.score).unsqueeze(0).unsqueeze(-1).unsqueeze(-1) 33 | elif self.share_type == 'kernel_wise': 34 | score = binarizer_fn(self.score).unsqueeze(0).unsqueeze(0) 35 | else: 36 | score = binarizer_fn(self.score) 37 | 38 | weight = score * self.weight + (1 - score) * self.specific_weight 39 | if self.bias is not None: 40 | bias = score * self.bias + (1 - score) * self.specific_bias 41 | else: 42 | bias = None 43 | return self._conv_forward(input, weight) 44 | 45 | 46 | @NORM_LAYERS.register_module() 47 | class BatchNorm2d_share(nn.BatchNorm2d): 48 | def __init__(self, num_features, eps=0.00001, momentum=0.1, affine=True, 49 | track_running_stats=True, device=None, dtype=None): 50 | super().__init__(num_features, eps, momentum, affine, track_running_stats, device, dtype) 51 | self.register_parameter('specific_weight', nn.Parameter(self.weight.clone())) 52 | if self.bias is not None: 53 | self.register_parameter('specific_bias', nn.Parameter(self.bias.clone())) 54 | self.score = nn.Parameter(torch.rand(1).cuda()) 55 | 56 | def forward(self, input): 57 | self._check_input_dim(input) 58 | 59 | # exponential_average_factor is set to self.momentum 60 | # (when it is available) only so that it gets updated 61 | # in ONNX graph when this node is exported to ONNX. 62 | if self.momentum is None: 63 | exponential_average_factor = 0.0 64 | else: 65 | exponential_average_factor = self.momentum 66 | 67 | if self.training and self.track_running_stats: 68 | # TODO: if statement only here to tell the jit to skip emitting this when it is None 69 | if self.num_batches_tracked is not None: # type: ignore[has-type] 70 | self.num_batches_tracked = self.num_batches_tracked + 1 # type: ignore[has-type] 71 | if self.momentum is None: # use cumulative moving average 72 | exponential_average_factor = 1.0 / float(self.num_batches_tracked) 73 | else: # use exponential moving average 74 | exponential_average_factor = self.momentum 75 | 76 | r""" 77 | Decide whether the mini-batch stats should be used for normalization rather than the buffers. 78 | Mini-batch stats are used in training mode, and in eval mode when buffers are None. 79 | """ 80 | if self.training: 81 | bn_training = True 82 | else: 83 | bn_training = (self.running_mean is None) and (self.running_var is None) 84 | 85 | r""" 86 | Buffers are only updated if they are to be tracked and we are in training mode. Thus they only need to be 87 | passed when the update should occur (i.e. in training mode when they are tracked), or when buffer stats are 88 | used for normalization (i.e. in eval mode when buffers are not None). 89 | """ 90 | 91 | score = binarizer_fn(self.score) 92 | 93 | weight = score * self.weight + (1 - score) * self.specific_weight 94 | if self.bias is not None: 95 | bias = score * self.bias + (1 - score) * self.specific_bias 96 | else: 97 | bias = None 98 | 99 | return F.batch_norm( 100 | input, 101 | # If buffers are not to be tracked, ensure that they won't be updated 102 | self.running_mean 103 | if not self.training or self.track_running_stats 104 | else None, 105 | self.running_var if not self.training or self.track_running_stats else None, 106 | weight, 107 | bias, 108 | bn_training, 109 | exponential_average_factor, 110 | self.eps, 111 | ) 112 | 113 | 114 | class BinarizerFn(torch.autograd.Function): 115 | """Binarizes {0, 1} a real valued tensor.""" 116 | 117 | @staticmethod 118 | def forward(ctx, inputs, threshold=THRESHOLD): 119 | outputs = inputs.clone() 120 | outputs[inputs.le(threshold)] = 1 121 | outputs[inputs.gt(threshold)] = 0 122 | return outputs 123 | 124 | @staticmethod 125 | def backward(self, gradOutput): 126 | return gradOutput, None 127 | 128 | 129 | binarizer_fn = BinarizerFn.apply -------------------------------------------------------------------------------- /ufvl_net/cnn/conv2d_share.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from mmcv.cnn import CONV_LAYERS, NORM_LAYERS 5 | 6 | THRESHOLD = 0.5 7 | 8 | 9 | @CONV_LAYERS.register_module() 10 | class Conv2d_share(nn.Conv2d): 11 | def __init__(self, in_channels, out_channels, kernel_size, share_type=None, 12 | stride=1, padding=0, dilation=1, groups=1, bias=True, 13 | padding_mode='zeros'): 14 | super().__init__(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias, padding_mode 15 | ) 16 | self.share_type = 'channel_wise' 17 | self.register_parameter('specific_weight', nn.Parameter(self.weight.clone())) 18 | if self.bias is not None: 19 | self.register_parameter('specific_bias', nn.Parameter(self.bias.clone())) 20 | 21 | # TODO: ablate one, zero and random 22 | # import pdb; pdb.set_trace() 23 | if self.share_type == 'channel_wise': 24 | self.score = nn.Parameter(torch.rand(self.weight.size(1)).cuda()) 25 | elif self.share_type == 'kernel_wise': 26 | self.score = nn.Parameter(torch.rand((self.weight.size(2), self.weight.size(3))).cuda()) 27 | else: 28 | self.score = nn.Parameter(torch.rand(1).cuda()) 29 | 30 | def forward(self, input): 31 | if self.share_type == 'channel_wise': 32 | score = binarizer_fn(self.score).unsqueeze(0).unsqueeze(-1).unsqueeze(-1) 33 | elif self.share_type == 'kernel_wise': 34 | score = binarizer_fn(self.score).unsqueeze(0).unsqueeze(0) 35 | else: 36 | score = binarizer_fn(self.score) 37 | 38 | weight = score * self.weight + (1 - score) * self.specific_weight 39 | if self.bias is not None: 40 | bias = score * self.bias + (1 - score) * self.specific_bias 41 | else: 42 | bias = None 43 | return self._conv_forward(input, weight) 44 | 45 | 46 | @NORM_LAYERS.register_module() 47 | class BatchNorm2d_share(nn.BatchNorm2d): 48 | def __init__(self, num_features, eps=0.00001, momentum=0.1, affine=True, 49 | track_running_stats=True, device=None, dtype=None): 50 | super().__init__(num_features, eps, momentum, affine, track_running_stats, device, dtype) 51 | self.register_parameter('specific_weight', nn.Parameter(self.weight.clone())) 52 | if self.bias is not None: 53 | self.register_parameter('specific_bias', nn.Parameter(self.bias.clone())) 54 | self.score = nn.Parameter(torch.rand(1).cuda()) 55 | 56 | def forward(self, input): 57 | self._check_input_dim(input) 58 | 59 | # exponential_average_factor is set to self.momentum 60 | # (when it is available) only so that it gets updated 61 | # in ONNX graph when this node is exported to ONNX. 62 | if self.momentum is None: 63 | exponential_average_factor = 0.0 64 | else: 65 | exponential_average_factor = self.momentum 66 | 67 | if self.training and self.track_running_stats: 68 | # TODO: if statement only here to tell the jit to skip emitting this when it is None 69 | if self.num_batches_tracked is not None: # type: ignore[has-type] 70 | self.num_batches_tracked = self.num_batches_tracked + 1 # type: ignore[has-type] 71 | if self.momentum is None: # use cumulative moving average 72 | exponential_average_factor = 1.0 / float(self.num_batches_tracked) 73 | else: # use exponential moving average 74 | exponential_average_factor = self.momentum 75 | 76 | r""" 77 | Decide whether the mini-batch stats should be used for normalization rather than the buffers. 78 | Mini-batch stats are used in training mode, and in eval mode when buffers are None. 79 | """ 80 | if self.training: 81 | bn_training = True 82 | else: 83 | bn_training = (self.running_mean is None) and (self.running_var is None) 84 | 85 | r""" 86 | Buffers are only updated if they are to be tracked and we are in training mode. Thus they only need to be 87 | passed when the update should occur (i.e. in training mode when they are tracked), or when buffer stats are 88 | used for normalization (i.e. in eval mode when buffers are not None). 89 | """ 90 | 91 | score = binarizer_fn(self.score) 92 | 93 | weight = score * self.weight + (1 - score) * self.specific_weight 94 | if self.bias is not None: 95 | bias = score * self.bias + (1 - score) * self.specific_bias 96 | else: 97 | bias = None 98 | 99 | return F.batch_norm( 100 | input, 101 | # If buffers are not to be tracked, ensure that they won't be updated 102 | self.running_mean 103 | if not self.training or self.track_running_stats 104 | else None, 105 | self.running_var if not self.training or self.track_running_stats else None, 106 | weight, 107 | bias, 108 | bn_training, 109 | exponential_average_factor, 110 | self.eps, 111 | ) 112 | 113 | 114 | class BinarizerFn(torch.autograd.Function): 115 | """Binarizes {0, 1} a real valued tensor.""" 116 | 117 | @staticmethod 118 | def forward(ctx, inputs, threshold=THRESHOLD): 119 | outputs = inputs.clone() 120 | outputs[inputs.le(threshold)] = 1 121 | outputs[inputs.gt(threshold)] = 0 122 | return outputs 123 | 124 | @staticmethod 125 | def backward(ctx, gradOutput): 126 | return gradOutput, None 127 | 128 | 129 | binarizer_fn = BinarizerFn.apply 130 | -------------------------------------------------------------------------------- /ufvl_net/datasets/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) OpenMMLab. All rights reserved. 2 | from .builder import (DATASETS, PIPELINES, SAMPLERS, build_dataloader, 3 | build_dataset, build_sampler) 4 | from .seven_scenes import SevenScenes 5 | from .twelvescenes import TWESCENES 6 | 7 | __all__ = ['SevenScenes', 'TWESCENES', 'DATASETS', 'PIPELINES', 'SAMPLERS', 8 | 'build_dataloader', 'build_dataset', 'build_sampler'] 9 | -------------------------------------------------------------------------------- /ufvl_net/datasets/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/datasets/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/datasets/__pycache__/base_dataset.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/datasets/__pycache__/base_dataset.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/datasets/__pycache__/builder.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/datasets/__pycache__/builder.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/datasets/__pycache__/cambridge.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/datasets/__pycache__/cambridge.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/datasets/__pycache__/cifar.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/datasets/__pycache__/cifar.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/datasets/__pycache__/cub.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/datasets/__pycache__/cub.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/datasets/__pycache__/custom.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/datasets/__pycache__/custom.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/datasets/__pycache__/dataset_wrappers.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/datasets/__pycache__/dataset_wrappers.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/datasets/__pycache__/imagenet.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/datasets/__pycache__/imagenet.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/datasets/__pycache__/imagenet21k.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/datasets/__pycache__/imagenet21k.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/datasets/__pycache__/mnist.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/datasets/__pycache__/mnist.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/datasets/__pycache__/multi_label.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/datasets/__pycache__/multi_label.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/datasets/__pycache__/seven_scenes.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/datasets/__pycache__/seven_scenes.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/datasets/__pycache__/twelve_scenes.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/datasets/__pycache__/twelve_scenes.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/datasets/__pycache__/twelvescenes.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/datasets/__pycache__/twelvescenes.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/datasets/__pycache__/utils.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/datasets/__pycache__/utils.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/datasets/__pycache__/voc.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/datasets/__pycache__/voc.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/datasets/builder.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) OpenMMLab. All rights reserved. 2 | import copy 3 | import platform 4 | import random 5 | from functools import partial 6 | 7 | import numpy as np 8 | import torch 9 | from mmcv.parallel import collate 10 | from mmcv.runner import get_dist_info 11 | from mmcv.utils import Registry, build_from_cfg, digit_version 12 | from torch.utils.data import DataLoader 13 | 14 | try: 15 | from mmcv.utils import IS_IPU_AVAILABLE 16 | except ImportError: 17 | IS_IPU_AVAILABLE = False 18 | 19 | if platform.system() != 'Windows': 20 | # https://github.com/pytorch/pytorch/issues/973 21 | import resource 22 | rlimit = resource.getrlimit(resource.RLIMIT_NOFILE) 23 | hard_limit = rlimit[1] 24 | soft_limit = min(4096, hard_limit) 25 | resource.setrlimit(resource.RLIMIT_NOFILE, (soft_limit, hard_limit)) 26 | 27 | DATASETS = Registry('dataset') 28 | PIPELINES = Registry('pipeline') 29 | SAMPLERS = Registry('sampler') 30 | 31 | 32 | def build_dataset(cfg, default_args=None): 33 | from .dataset_wrappers import (ClassBalancedDataset, ConcatDataset, 34 | KFoldDataset, RepeatDataset) 35 | if isinstance(cfg, (list, tuple)): 36 | dataset = ConcatDataset([build_dataset(c, default_args) for c in cfg]) 37 | elif cfg['type'] == 'ConcatDataset': 38 | dataset = ConcatDataset( 39 | [build_dataset(c, default_args) for c in cfg['datasets']], 40 | separate_eval=cfg.get('separate_eval', True)) 41 | elif cfg['type'] == 'RepeatDataset': 42 | dataset = RepeatDataset( 43 | build_dataset(cfg['dataset'], default_args), cfg['times']) 44 | elif cfg['type'] == 'ClassBalancedDataset': 45 | dataset = ClassBalancedDataset( 46 | build_dataset(cfg['dataset'], default_args), cfg['oversample_thr']) 47 | elif cfg['type'] == 'KFoldDataset': 48 | cp_cfg = copy.deepcopy(cfg) 49 | if cp_cfg.get('test_mode', None) is None: 50 | cp_cfg['test_mode'] = (default_args or {}).pop('test_mode', False) 51 | cp_cfg['dataset'] = build_dataset(cp_cfg['dataset'], default_args) 52 | cp_cfg.pop('type') 53 | dataset = KFoldDataset(**cp_cfg) 54 | else: 55 | dataset = build_from_cfg(cfg, DATASETS, default_args) 56 | 57 | return dataset 58 | 59 | 60 | def build_dataloader(dataset, 61 | samples_per_gpu, 62 | workers_per_gpu, 63 | num_gpus=1, 64 | dist=True, 65 | shuffle=True, 66 | round_up=True, 67 | seed=None, 68 | pin_memory=True, 69 | persistent_workers=True, 70 | sampler_cfg=None, 71 | **kwargs): 72 | """Build PyTorch DataLoader. 73 | 74 | In distributed training, each GPU/process has a dataloader. 75 | In non-distributed training, there is only one dataloader for all GPUs. 76 | 77 | Args: 78 | dataset (Dataset): A PyTorch dataset. 79 | samples_per_gpu (int): Number of training samples on each GPU, i.e., 80 | batch size of each GPU. 81 | workers_per_gpu (int): How many subprocesses to use for data loading 82 | for each GPU. 83 | num_gpus (int): Number of GPUs. Only used in non-distributed training. 84 | dist (bool): Distributed training/test or not. Default: True. 85 | shuffle (bool): Whether to shuffle the data at every epoch. 86 | Default: True. 87 | round_up (bool): Whether to round up the length of dataset by adding 88 | extra samples to make it evenly divisible. Default: True. 89 | pin_memory (bool): Whether to use pin_memory in DataLoader. 90 | Default: True 91 | persistent_workers (bool): If True, the data loader will not shutdown 92 | the worker processes after a dataset has been consumed once. 93 | This allows to maintain the workers Dataset instances alive. 94 | The argument also has effect in PyTorch>=1.7.0. 95 | Default: True 96 | sampler_cfg (dict): sampler configuration to override the default 97 | sampler 98 | kwargs: any keyword argument to be used to initialize DataLoader 99 | 100 | Returns: 101 | DataLoader: A PyTorch dataloader. 102 | """ 103 | rank, world_size = get_dist_info() 104 | 105 | # Custom sampler logic 106 | if sampler_cfg: 107 | # shuffle=False when val and test 108 | sampler_cfg.update(shuffle=shuffle) 109 | sampler = build_sampler( 110 | sampler_cfg, 111 | default_args=dict( 112 | dataset=dataset, num_replicas=world_size, rank=rank, 113 | seed=seed)) 114 | # Default sampler logic 115 | elif dist: 116 | sampler = build_sampler( 117 | dict( 118 | type='DistributedSampler', 119 | dataset=dataset, 120 | num_replicas=world_size, 121 | rank=rank, 122 | shuffle=shuffle, 123 | round_up=round_up, 124 | seed=seed)) 125 | else: 126 | sampler = None 127 | 128 | # If sampler exists, turn off dataloader shuffle 129 | if sampler is not None: 130 | shuffle = False 131 | 132 | if dist: 133 | batch_size = samples_per_gpu 134 | num_workers = workers_per_gpu 135 | else: 136 | batch_size = num_gpus * samples_per_gpu 137 | num_workers = num_gpus * workers_per_gpu 138 | 139 | init_fn = partial( 140 | worker_init_fn, num_workers=num_workers, rank=rank, 141 | seed=seed) if seed is not None else None 142 | 143 | if digit_version(torch.__version__) >= digit_version('1.8.0'): 144 | kwargs['persistent_workers'] = persistent_workers 145 | if IS_IPU_AVAILABLE: 146 | from mmcv.device.ipu import IPUDataLoader 147 | data_loader = IPUDataLoader( 148 | dataset, 149 | None, 150 | batch_size=samples_per_gpu, 151 | num_workers=num_workers, 152 | shuffle=shuffle, 153 | worker_init_fn=init_fn, 154 | **kwargs) 155 | else: 156 | data_loader = DataLoader( 157 | dataset, 158 | batch_size=batch_size, 159 | sampler=sampler, 160 | num_workers=num_workers, 161 | collate_fn=partial(collate, samples_per_gpu=samples_per_gpu), 162 | pin_memory=pin_memory, 163 | shuffle=shuffle, 164 | worker_init_fn=init_fn, 165 | **kwargs) 166 | 167 | return data_loader 168 | 169 | 170 | def worker_init_fn(worker_id, num_workers, rank, seed): 171 | # The seed of each worker equals to 172 | # num_worker * rank + worker_id + user_seed 173 | worker_seed = num_workers * rank + worker_id + seed 174 | np.random.seed(worker_seed) 175 | random.seed(worker_seed) 176 | torch.manual_seed(worker_seed) 177 | 178 | 179 | def build_sampler(cfg, default_args=None): 180 | if cfg is None: 181 | return None 182 | else: 183 | return build_from_cfg(cfg, SAMPLERS, default_args=default_args) 184 | -------------------------------------------------------------------------------- /ufvl_net/datasets/dataset_wrappers.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) OpenMMLab. All rights reserved. 2 | import bisect 3 | import math 4 | from collections import defaultdict 5 | 6 | import numpy as np 7 | from mmcv.utils import print_log 8 | from torch.utils.data.dataset import ConcatDataset as _ConcatDataset 9 | 10 | from .builder import DATASETS 11 | 12 | 13 | @DATASETS.register_module() 14 | class ConcatDataset(_ConcatDataset): 15 | """A wrapper of concatenated dataset. 16 | 17 | Same as :obj:`torch.utils.data.dataset.ConcatDataset`, but 18 | add `get_cat_ids` function. 19 | 20 | Args: 21 | datasets (list[:obj:`BaseDataset`]): A list of datasets. 22 | separate_eval (bool): Whether to evaluate the results 23 | separately if it is used as validation dataset. 24 | Defaults to True. 25 | """ 26 | 27 | def __init__(self, datasets, separate_eval=True): 28 | super(ConcatDataset, self).__init__(datasets) 29 | self.separate_eval = separate_eval 30 | 31 | self.CLASSES = datasets[0].CLASSES 32 | 33 | if not separate_eval: 34 | if len(set([type(ds) for ds in datasets])) != 1: 35 | raise NotImplementedError( 36 | 'To evaluate a concat dataset non-separately, ' 37 | 'all the datasets should have same types') 38 | 39 | def get_cat_ids(self, idx): 40 | if idx < 0: 41 | if -idx > len(self): 42 | raise ValueError( 43 | 'absolute value of index should not exceed dataset length') 44 | idx = len(self) + idx 45 | dataset_idx = bisect.bisect_right(self.cumulative_sizes, idx) 46 | if dataset_idx == 0: 47 | sample_idx = idx 48 | else: 49 | sample_idx = idx - self.cumulative_sizes[dataset_idx - 1] 50 | return self.datasets[dataset_idx].get_cat_ids(sample_idx) 51 | 52 | def evaluate(self, results, *args, indices=None, logger=None, **kwargs): 53 | """Evaluate the results. 54 | 55 | Args: 56 | results (list[list | tuple]): Testing results of the dataset. 57 | indices (list, optional): The indices of samples corresponding to 58 | the results. It's unavailable on ConcatDataset. 59 | Defaults to None. 60 | logger (logging.Logger | str, optional): Logger used for printing 61 | related information during evaluation. Defaults to None. 62 | 63 | Returns: 64 | dict[str: float]: AP results of the total dataset or each separate 65 | dataset if `self.separate_eval=True`. 66 | """ 67 | if indices is not None: 68 | raise NotImplementedError( 69 | 'Use indices to evaluate speific samples in a ConcatDataset ' 70 | 'is not supported by now.') 71 | 72 | assert len(results) == len(self), \ 73 | ('Dataset and results have different sizes: ' 74 | f'{len(self)} v.s. {len(results)}') 75 | 76 | # Check whether all the datasets support evaluation 77 | for dataset in self.datasets: 78 | assert hasattr(dataset, 'evaluate'), \ 79 | f"{type(dataset)} haven't implemented the evaluate function." 80 | 81 | if self.separate_eval: 82 | total_eval_results = dict() 83 | for dataset_idx, dataset in enumerate(self.datasets): 84 | start_idx = 0 if dataset_idx == 0 else \ 85 | self.cumulative_sizes[dataset_idx-1] 86 | end_idx = self.cumulative_sizes[dataset_idx] 87 | 88 | results_per_dataset = results[start_idx:end_idx] 89 | print_log( 90 | f'Evaluateing dataset-{dataset_idx} with ' 91 | f'{len(results_per_dataset)} images now', 92 | logger=logger) 93 | 94 | eval_results_per_dataset = dataset.evaluate( 95 | results_per_dataset, *args, logger=logger, **kwargs) 96 | for k, v in eval_results_per_dataset.items(): 97 | total_eval_results.update({f'{dataset_idx}_{k}': v}) 98 | 99 | return total_eval_results 100 | else: 101 | original_data_infos = self.datasets[0].data_infos 102 | self.datasets[0].data_infos = sum( 103 | [dataset.data_infos for dataset in self.datasets], []) 104 | eval_results = self.datasets[0].evaluate( 105 | results, logger=logger, **kwargs) 106 | self.datasets[0].data_infos = original_data_infos 107 | return eval_results 108 | 109 | 110 | @DATASETS.register_module() 111 | class RepeatDataset(object): 112 | """A wrapper of repeated dataset. 113 | 114 | The length of repeated dataset will be `times` larger than the original 115 | dataset. This is useful when the data loading time is long but the dataset 116 | is small. Using RepeatDataset can reduce the data loading time between 117 | epochs. 118 | 119 | Args: 120 | dataset (:obj:`BaseDataset`): The dataset to be repeated. 121 | times (int): Repeat times. 122 | """ 123 | 124 | def __init__(self, dataset, times): 125 | self.dataset = dataset 126 | self.times = times 127 | self.CLASSES = dataset.CLASSES 128 | 129 | self._ori_len = len(self.dataset) 130 | 131 | def __getitem__(self, idx): 132 | return self.dataset[idx % self._ori_len] 133 | 134 | def get_cat_ids(self, idx): 135 | return self.dataset.get_cat_ids(idx % self._ori_len) 136 | 137 | def __len__(self): 138 | return self.times * self._ori_len 139 | 140 | def evaluate(self, *args, **kwargs): 141 | raise NotImplementedError( 142 | 'evaluate results on a repeated dataset is weird. ' 143 | 'Please inference and evaluate on the original dataset.') 144 | 145 | def __repr__(self): 146 | """Print the number of instance number.""" 147 | dataset_type = 'Test' if self.test_mode else 'Train' 148 | result = ( 149 | f'\n{self.__class__.__name__} ({self.dataset.__class__.__name__}) ' 150 | f'{dataset_type} dataset with total number of samples {len(self)}.' 151 | ) 152 | return result 153 | 154 | 155 | # Modified from https://github.com/facebookresearch/detectron2/blob/41d475b75a230221e21d9cac5d69655e3415e3a4/detectron2/data/samplers/distributed_sampler.py#L57 # noqa 156 | @DATASETS.register_module() 157 | class ClassBalancedDataset(object): 158 | r"""A wrapper of repeated dataset with repeat factor. 159 | 160 | Suitable for training on class imbalanced datasets like LVIS. Following the 161 | sampling strategy in `this paper`_, in each epoch, an image may appear 162 | multiple times based on its "repeat factor". 163 | 164 | .. _this paper: https://arxiv.org/pdf/1908.03195.pdf 165 | 166 | The repeat factor for an image is a function of the frequency the rarest 167 | category labeled in that image. The "frequency of category c" in [0, 1] 168 | is defined by the fraction of images in the training set (without repeats) 169 | in which category c appears. 170 | 171 | The dataset needs to implement :func:`self.get_cat_ids` to support 172 | ClassBalancedDataset. 173 | 174 | The repeat factor is computed as followed. 175 | 176 | 1. For each category c, compute the fraction :math:`f(c)` of images that 177 | contain it. 178 | 2. For each category c, compute the category-level repeat factor 179 | 180 | .. math:: 181 | r(c) = \max(1, \sqrt{\frac{t}{f(c)}}) 182 | 183 | 3. For each image I and its labels :math:`L(I)`, compute the image-level 184 | repeat factor 185 | 186 | .. math:: 187 | r(I) = \max_{c \in L(I)} r(c) 188 | 189 | Args: 190 | dataset (:obj:`BaseDataset`): The dataset to be repeated. 191 | oversample_thr (float): frequency threshold below which data is 192 | repeated. For categories with ``f_c`` >= ``oversample_thr``, there 193 | is no oversampling. For categories with ``f_c`` < 194 | ``oversample_thr``, the degree of oversampling following the 195 | square-root inverse frequency heuristic above. 196 | """ 197 | 198 | def __init__(self, dataset, oversample_thr): 199 | self.dataset = dataset 200 | self.oversample_thr = oversample_thr 201 | self.CLASSES = dataset.CLASSES 202 | 203 | repeat_factors = self._get_repeat_factors(dataset, oversample_thr) 204 | repeat_indices = [] 205 | for dataset_index, repeat_factor in enumerate(repeat_factors): 206 | repeat_indices.extend([dataset_index] * math.ceil(repeat_factor)) 207 | self.repeat_indices = repeat_indices 208 | 209 | flags = [] 210 | if hasattr(self.dataset, 'flag'): 211 | for flag, repeat_factor in zip(self.dataset.flag, repeat_factors): 212 | flags.extend([flag] * int(math.ceil(repeat_factor))) 213 | assert len(flags) == len(repeat_indices) 214 | self.flag = np.asarray(flags, dtype=np.uint8) 215 | 216 | def _get_repeat_factors(self, dataset, repeat_thr): 217 | # 1. For each category c, compute the fraction # of images 218 | # that contain it: f(c) 219 | category_freq = defaultdict(int) 220 | num_images = len(dataset) 221 | for idx in range(num_images): 222 | cat_ids = set(self.dataset.get_cat_ids(idx)) 223 | for cat_id in cat_ids: 224 | category_freq[cat_id] += 1 225 | for k, v in category_freq.items(): 226 | assert v > 0, f'caterogy {k} does not contain any images' 227 | category_freq[k] = v / num_images 228 | 229 | # 2. For each category c, compute the category-level repeat factor: 230 | # r(c) = max(1, sqrt(t/f(c))) 231 | category_repeat = { 232 | cat_id: max(1.0, math.sqrt(repeat_thr / cat_freq)) 233 | for cat_id, cat_freq in category_freq.items() 234 | } 235 | 236 | # 3. For each image I and its labels L(I), compute the image-level 237 | # repeat factor: 238 | # r(I) = max_{c in L(I)} r(c) 239 | repeat_factors = [] 240 | for idx in range(num_images): 241 | cat_ids = set(self.dataset.get_cat_ids(idx)) 242 | repeat_factor = max( 243 | {category_repeat[cat_id] 244 | for cat_id in cat_ids}) 245 | repeat_factors.append(repeat_factor) 246 | 247 | return repeat_factors 248 | 249 | def __getitem__(self, idx): 250 | ori_index = self.repeat_indices[idx] 251 | return self.dataset[ori_index] 252 | 253 | def __len__(self): 254 | return len(self.repeat_indices) 255 | 256 | def evaluate(self, *args, **kwargs): 257 | raise NotImplementedError( 258 | 'evaluate results on a class-balanced dataset is weird. ' 259 | 'Please inference and evaluate on the original dataset.') 260 | 261 | def __repr__(self): 262 | """Print the number of instance number.""" 263 | dataset_type = 'Test' if self.test_mode else 'Train' 264 | result = ( 265 | f'\n{self.__class__.__name__} ({self.dataset.__class__.__name__}) ' 266 | f'{dataset_type} dataset with total number of samples {len(self)}.' 267 | ) 268 | return result 269 | 270 | 271 | @DATASETS.register_module() 272 | class KFoldDataset: 273 | """A wrapper of dataset for K-Fold cross-validation. 274 | 275 | K-Fold cross-validation divides all the samples in groups of samples, 276 | called folds, of almost equal sizes. And we use k-1 of folds to do training 277 | and use the fold left to do validation. 278 | 279 | Args: 280 | dataset (:obj:`BaseDataset`): The dataset to be divided. 281 | fold (int): The fold used to do validation. Defaults to 0. 282 | num_splits (int): The number of all folds. Defaults to 5. 283 | test_mode (bool): Use the training dataset or validation dataset. 284 | Defaults to False. 285 | seed (int, optional): The seed to shuffle the dataset before splitting. 286 | If None, not shuffle the dataset. Defaults to None. 287 | """ 288 | 289 | def __init__(self, 290 | dataset, 291 | fold=0, 292 | num_splits=5, 293 | test_mode=False, 294 | seed=None): 295 | self.dataset = dataset 296 | self.CLASSES = dataset.CLASSES 297 | self.test_mode = test_mode 298 | self.num_splits = num_splits 299 | 300 | length = len(dataset) 301 | indices = list(range(length)) 302 | if isinstance(seed, int): 303 | rng = np.random.default_rng(seed) 304 | rng.shuffle(indices) 305 | 306 | test_start = length * fold // num_splits 307 | test_end = length * (fold + 1) // num_splits 308 | if test_mode: 309 | self.indices = indices[test_start:test_end] 310 | else: 311 | self.indices = indices[:test_start] + indices[test_end:] 312 | 313 | def get_cat_ids(self, idx): 314 | return self.dataset.get_cat_ids(self.indices[idx]) 315 | 316 | def get_gt_labels(self): 317 | dataset_gt_labels = self.dataset.get_gt_labels() 318 | gt_labels = np.array([dataset_gt_labels[idx] for idx in self.indices]) 319 | return gt_labels 320 | 321 | def __getitem__(self, idx): 322 | return self.dataset[self.indices[idx]] 323 | 324 | def __len__(self): 325 | return len(self.indices) 326 | 327 | def evaluate(self, *args, **kwargs): 328 | kwargs['indices'] = self.indices 329 | return self.dataset.evaluate(*args, **kwargs) 330 | -------------------------------------------------------------------------------- /ufvl_net/datasets/piplines/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) OpenMMLab. All rights reserved. 2 | from .auto_augment import (AutoAugment, AutoContrast, Brightness, 3 | ColorTransform, Contrast, Cutout, Equalize, Invert, 4 | Posterize, RandAugment, Rotate, Sharpness, Shear, 5 | Solarize, SolarizeAdd, Translate) 6 | from .compose import Compose 7 | from .formatting import (Collect, ImageToTensor, ToNumpy, ToPIL, ToTensor, 8 | Transpose, to_tensor) 9 | from .loading import LoadImageFromFile 10 | from .transforms import (CenterCrop, ColorJitter, Lighting, Normalize, Pad, 11 | RandomCrop, RandomErasing, RandomFlip, 12 | RandomGrayscale, RandomResizedCrop, Resize) 13 | 14 | __all__ = [ 15 | 'Compose', 'to_tensor', 'ToTensor', 'ImageToTensor', 'ToPIL', 'ToNumpy', 16 | 'Transpose', 'Collect', 'LoadImageFromFile', 'Resize', 'CenterCrop', 17 | 'RandomFlip', 'Normalize', 'RandomCrop', 'RandomResizedCrop', 18 | 'RandomGrayscale', 'Shear', 'Translate', 'Rotate', 'Invert', 19 | 'ColorTransform', 'Solarize', 'Posterize', 'AutoContrast', 'Equalize', 20 | 'Contrast', 'Brightness', 'Sharpness', 'AutoAugment', 'SolarizeAdd', 21 | 'Cutout', 'RandAugment', 'Lighting', 'ColorJitter', 'RandomErasing', 'Pad' 22 | ] 23 | -------------------------------------------------------------------------------- /ufvl_net/datasets/piplines/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/datasets/piplines/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/datasets/piplines/__pycache__/auto_augment.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/datasets/piplines/__pycache__/auto_augment.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/datasets/piplines/__pycache__/compose.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/datasets/piplines/__pycache__/compose.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/datasets/piplines/__pycache__/formatting.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/datasets/piplines/__pycache__/formatting.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/datasets/piplines/__pycache__/loading.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/datasets/piplines/__pycache__/loading.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/datasets/piplines/__pycache__/transforms.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/datasets/piplines/__pycache__/transforms.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/datasets/piplines/compose.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) OpenMMLab. All rights reserved. 2 | from collections.abc import Sequence 3 | 4 | from mmcv.utils import build_from_cfg 5 | 6 | from ..builder import PIPELINES 7 | 8 | 9 | @PIPELINES.register_module() 10 | class Compose(object): 11 | """Compose a data pipeline with a sequence of transforms. 12 | 13 | Args: 14 | transforms (list[dict | callable]): 15 | Either config dicts of transforms or transform objects. 16 | """ 17 | 18 | def __init__(self, transforms): 19 | assert isinstance(transforms, Sequence) 20 | self.transforms = [] 21 | for transform in transforms: 22 | if isinstance(transform, dict): 23 | transform = build_from_cfg(transform, PIPELINES) 24 | self.transforms.append(transform) 25 | elif callable(transform): 26 | self.transforms.append(transform) 27 | else: 28 | raise TypeError('transform must be callable or a dict, but got' 29 | f' {type(transform)}') 30 | 31 | def __call__(self, data): 32 | for t in self.transforms: 33 | data = t(data) 34 | if data is None: 35 | return None 36 | return data 37 | 38 | def __repr__(self): 39 | format_string = self.__class__.__name__ + '(' 40 | for t in self.transforms: 41 | format_string += f'\n {t}' 42 | format_string += '\n)' 43 | return format_string 44 | -------------------------------------------------------------------------------- /ufvl_net/datasets/piplines/formatting.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) OpenMMLab. All rights reserved. 2 | from collections.abc import Sequence 3 | 4 | import mmcv 5 | import numpy as np 6 | import torch 7 | from mmcv.parallel import DataContainer as DC 8 | from PIL import Image 9 | 10 | from ..builder import PIPELINES 11 | 12 | 13 | def to_tensor(data): 14 | """Convert objects of various python types to :obj:`torch.Tensor`. 15 | 16 | Supported types are: :class:`numpy.ndarray`, :class:`torch.Tensor`, 17 | :class:`Sequence`, :class:`int` and :class:`float`. 18 | """ 19 | if isinstance(data, torch.Tensor): 20 | return data 21 | elif isinstance(data, np.ndarray): 22 | return torch.from_numpy(data) 23 | elif isinstance(data, Sequence) and not mmcv.is_str(data): 24 | return torch.tensor(data) 25 | elif isinstance(data, int): 26 | return torch.LongTensor([data]) 27 | elif isinstance(data, float): 28 | return torch.FloatTensor([data]) 29 | else: 30 | raise TypeError( 31 | f'Type {type(data)} cannot be converted to tensor.' 32 | 'Supported types are: `numpy.ndarray`, `torch.Tensor`, ' 33 | '`Sequence`, `int` and `float`') 34 | 35 | 36 | @PIPELINES.register_module() 37 | class ToTensor(object): 38 | 39 | def __init__(self, keys): 40 | self.keys = keys 41 | 42 | def __call__(self, results): 43 | for key in self.keys: 44 | results[key] = to_tensor(results[key]) 45 | return results 46 | 47 | def __repr__(self): 48 | return self.__class__.__name__ + f'(keys={self.keys})' 49 | 50 | 51 | @PIPELINES.register_module() 52 | class ImageToTensor(object): 53 | 54 | def __init__(self, keys): 55 | self.keys = keys 56 | 57 | def __call__(self, results): 58 | for key in self.keys: 59 | img = results[key] 60 | if len(img.shape) < 3: 61 | img = np.expand_dims(img, -1) 62 | results[key] = to_tensor(img.transpose(2, 0, 1)) 63 | return results 64 | 65 | def __repr__(self): 66 | return self.__class__.__name__ + f'(keys={self.keys})' 67 | 68 | 69 | @PIPELINES.register_module() 70 | class Transpose(object): 71 | 72 | def __init__(self, keys, order): 73 | self.keys = keys 74 | self.order = order 75 | 76 | def __call__(self, results): 77 | for key in self.keys: 78 | results[key] = results[key].transpose(self.order) 79 | return results 80 | 81 | def __repr__(self): 82 | return self.__class__.__name__ + \ 83 | f'(keys={self.keys}, order={self.order})' 84 | 85 | 86 | @PIPELINES.register_module() 87 | class ToPIL(object): 88 | 89 | def __init__(self): 90 | pass 91 | 92 | def __call__(self, results): 93 | results['img'] = Image.fromarray(results['img']) 94 | return results 95 | 96 | 97 | @PIPELINES.register_module() 98 | class ToNumpy(object): 99 | 100 | def __init__(self): 101 | pass 102 | 103 | def __call__(self, results): 104 | results['img'] = np.array(results['img'], dtype=np.float32) 105 | return results 106 | 107 | 108 | @PIPELINES.register_module() 109 | class Collect(object): 110 | """Collect data from the loader relevant to the specific task. 111 | 112 | This is usually the last stage of the data loader pipeline. Typically keys 113 | is set to some subset of "img" and "gt_label". 114 | 115 | Args: 116 | keys (Sequence[str]): Keys of results to be collected in ``data``. 117 | meta_keys (Sequence[str], optional): Meta keys to be converted to 118 | ``mmcv.DataContainer`` and collected in ``data[img_metas]``. 119 | Default: ('filename', 'ori_shape', 'img_shape', 'flip', 120 | 'flip_direction', 'img_norm_cfg') 121 | 122 | Returns: 123 | dict: The result dict contains the following keys 124 | 125 | - keys in ``self.keys`` 126 | - ``img_metas`` if available 127 | """ 128 | 129 | def __init__(self, 130 | keys, 131 | meta_keys=('filename', 'ori_filename', 'ori_shape', 132 | 'img_shape', 'flip', 'flip_direction', 133 | 'img_norm_cfg')): 134 | self.keys = keys 135 | self.meta_keys = meta_keys 136 | 137 | def __call__(self, results): 138 | data = {} 139 | img_meta = {} 140 | for key in self.meta_keys: 141 | if key in results: 142 | img_meta[key] = results[key] 143 | data['img_metas'] = DC(img_meta, cpu_only=True) 144 | for key in self.keys: 145 | data[key] = results[key] 146 | return data 147 | 148 | def __repr__(self): 149 | return self.__class__.__name__ + \ 150 | f'(keys={self.keys}, meta_keys={self.meta_keys})' 151 | 152 | 153 | @PIPELINES.register_module() 154 | class WrapFieldsToLists(object): 155 | """Wrap fields of the data dictionary into lists for evaluation. 156 | 157 | This class can be used as a last step of a test or validation 158 | pipeline for single image evaluation or inference. 159 | 160 | Example: 161 | >>> test_pipeline = [ 162 | >>> dict(type='LoadImageFromFile'), 163 | >>> dict(type='Normalize', 164 | mean=[123.675, 116.28, 103.53], 165 | std=[58.395, 57.12, 57.375], 166 | to_rgb=True), 167 | >>> dict(type='ImageToTensor', keys=['img']), 168 | >>> dict(type='Collect', keys=['img']), 169 | >>> dict(type='WrapIntoLists') 170 | >>> ] 171 | """ 172 | 173 | def __call__(self, results): 174 | # Wrap dict fields into lists 175 | for key, val in results.items(): 176 | results[key] = [val] 177 | return results 178 | 179 | def __repr__(self): 180 | return f'{self.__class__.__name__}()' 181 | 182 | 183 | @PIPELINES.register_module() 184 | class ToHalf(object): 185 | 186 | def __init__(self, keys): 187 | self.keys = keys 188 | 189 | def __call__(self, results): 190 | for k in self.keys: 191 | if isinstance(results[k], torch.Tensor): 192 | results[k] = results[k].to(torch.half) 193 | else: 194 | results[k] = results[k].astype(np.float16) 195 | return results 196 | -------------------------------------------------------------------------------- /ufvl_net/datasets/piplines/loading.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) OpenMMLab. All rights reserved. 2 | import os.path as osp 3 | 4 | import mmcv 5 | import numpy as np 6 | 7 | from ..builder import PIPELINES 8 | 9 | 10 | @PIPELINES.register_module() 11 | class LoadImageFromFile(object): 12 | """Load an image from file. 13 | 14 | Required keys are "img_prefix" and "img_info" (a dict that must contain the 15 | key "filename"). Added or updated keys are "filename", "img", "img_shape", 16 | "ori_shape" (same as `img_shape`) and "img_norm_cfg" (means=0 and stds=1). 17 | 18 | Args: 19 | to_float32 (bool): Whether to convert the loaded image to a float32 20 | numpy array. If set to False, the loaded image is an uint8 array. 21 | Defaults to False. 22 | color_type (str): The flag argument for :func:`mmcv.imfrombytes()`. 23 | Defaults to 'color'. 24 | file_client_args (dict): Arguments to instantiate a FileClient. 25 | See :class:`mmcv.fileio.FileClient` for details. 26 | Defaults to ``dict(backend='disk')``. 27 | """ 28 | 29 | def __init__(self, 30 | to_float32=False, 31 | color_type='color', 32 | file_client_args=dict(backend='disk')): 33 | self.to_float32 = to_float32 34 | self.color_type = color_type 35 | self.file_client_args = file_client_args.copy() 36 | self.file_client = None 37 | 38 | def __call__(self, results): 39 | if self.file_client is None: 40 | self.file_client = mmcv.FileClient(**self.file_client_args) 41 | 42 | if results['img_prefix'] is not None: 43 | filename = osp.join(results['img_prefix'], 44 | results['img_info']['filename']) 45 | else: 46 | filename = results['img_info']['filename'] 47 | 48 | img_bytes = self.file_client.get(filename) 49 | img = mmcv.imfrombytes(img_bytes, flag=self.color_type) 50 | if self.to_float32: 51 | img = img.astype(np.float32) 52 | 53 | results['filename'] = filename 54 | results['ori_filename'] = results['img_info']['filename'] 55 | results['img'] = img 56 | results['img_shape'] = img.shape 57 | results['ori_shape'] = img.shape 58 | num_channels = 1 if len(img.shape) < 3 else img.shape[2] 59 | results['img_norm_cfg'] = dict( 60 | mean=np.zeros(num_channels, dtype=np.float32), 61 | std=np.ones(num_channels, dtype=np.float32), 62 | to_rgb=False) 63 | return results 64 | 65 | def __repr__(self): 66 | repr_str = (f'{self.__class__.__name__}(' 67 | f'to_float32={self.to_float32}, ' 68 | f"color_type='{self.color_type}', " 69 | f'file_client_args={self.file_client_args})') 70 | return repr_str 71 | -------------------------------------------------------------------------------- /ufvl_net/datasets/twelvescenes.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | import imgaug 4 | import torch 5 | import numpy as np 6 | import random 7 | from imgaug import augmenters as iaa 8 | import cv2 9 | from torch.utils.data import Dataset 10 | 11 | from .builder import DATASETS 12 | 13 | 14 | @DATASETS.register_module() 15 | class TWESCENES(Dataset): 16 | def __init__(self, root, dataset='7S', scene='heads', split='train', 17 | model='fdanet', aug='False', **kwargs): 18 | self.intrinsics_color = np.array([[572.0, 0.0, 320.0], 19 | [0.0, 572.0, 240.0], 20 | [0.0, 0.0, 1.0]]) 21 | 22 | self.intrinsics_color_inv = np.linalg.inv(self.intrinsics_color) 23 | 24 | self.model = model 25 | self.dataset = dataset 26 | self.aug = aug 27 | self.root = os.path.join(root, '12Scenes') 28 | self.scene = scene 29 | 30 | self.split = split 31 | self.obj_suffixes = ['.color.jpg', '.pose.txt', '.depth.png', 32 | '.label.png'] 33 | self.obj_keys = ['color', 'pose', 'depth', 'label'] 34 | # 这里设定了训练/测试的图片 35 | if self.dataset == '12S' or self.split == 'test': 36 | with open(os.path.join(self.root, self.scene, 37 | '{}{}'.format('test', '.txt')), 'r') as f: 38 | self.frames = f.readlines() 39 | else: 40 | self.frames = [] 41 | with open(os.path.join(self.root, scene, 42 | '{}{}'.format(self.split, '.txt')), 'r') as f: 43 | frames = f.readlines() 44 | frames = [scene + ' ' + frame for frame in frames ] 45 | self.frames.extend(frames) 46 | 47 | def __len__(self): 48 | return len(self.frames) 49 | 50 | def __getitem__(self, index): 51 | # 遍历每一张图片 52 | frame = self.frames[index].rstrip('\n') 53 | seq_id, frame_id = frame.split('-') 54 | objs = {} 55 | objs['color'] = '/mnt/share/sda-2T/xietao/12Scenes/' + self.scene + '/data/' + seq_id + '-' + frame_id + '.color.jpg' 56 | objs['depth'] = '/mnt/share/sda-2T/xietao/12Scenes/' + self.scene + '/data/' + seq_id + '-' + frame_id + '.depth.png' # Twc 57 | objs['pose'] = '/mnt/share/sda-2T/xietao/12Scenes/' + self.scene + '/data/' + seq_id + '-' + frame_id + '.pose.txt' 58 | 59 | img = cv2.imread(objs['color']) 60 | img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) 61 | img = cv2.resize(img, (640, 480)) 62 | 63 | pose = np.loadtxt(objs['pose']) 64 | 65 | if self.split == 'test': 66 | img, pose = to_tensor_query(img, pose) 67 | ret = dict(img=img, gt_lables=pose) 68 | return ret # 返回torch类型的图片和位姿 69 | 70 | depth = cv2.imread(objs['depth'], -1) 71 | pose[0:3, 3] = pose[0:3, 3] * 1000 72 | coord, mask = get_coord(depth, pose, self.intrinsics_color_inv, self.dataset) 73 | img, coord, mask = data_aug(img, coord, mask, self.aug) 74 | 75 | if self.model == 'hscnet': 76 | coord = coord 77 | coord = coord[4::8, 4::8, :] # [60 80] 78 | mask = mask[4::8, 4::8].astype(np.float16) 79 | img, coord, mask = to_tensor(img, coord, mask) 80 | 81 | ret = dict(img=img, gt_lables=coord, mask=mask) 82 | 83 | return ret 84 | 85 | def evaluate(self, results, *args, **kwargs): 86 | transl_err_list = list() 87 | rot_err_list = list() 88 | for i in range(len(results)): 89 | transl_err_list.append(results[i]['trans_error_med']) 90 | rot_err_list.append(results[i]['rot_err_med']) 91 | res_ = np.array([transl_err_list, rot_err_list]).T 92 | 93 | return dict(median_trans_error=np.median(res_[:, 0]), 94 | median_rot_error=np.median(res_[:, 1]), 95 | accuracy=np.sum((res_[:, 0] <= 0.050) * (res_[:, 1] <= 5)) * 1. / len(res_) 96 | ) 97 | 98 | 99 | # depth：深度图[480 640] 100 | # 返回深度图对齐到RGB图后，RGB相应的深度信息 101 | def get_depth(depth, calibration_extrinsics, intrinsics_color, 102 | intrinsics_depth_inv): 103 | """Return the calibrated depth image (7-Scenes). 104 | Calibration parameters from DSAC (https://github.com/cvlab-dresden/DSAC) 105 | are used. 106 | """ 107 | ''' 108 | 利用深度摄像头内参矩阵把深度平面坐标（深度图坐标）转换到深度摄像头空间坐标， 109 | 再利用外参计算旋转矩阵和平移矩阵，把深度摄像头空间坐标转换到RGB摄像头空间坐标， 110 | 最后利用RGB摄像头内参矩阵把RGB摄像头空间坐标转换到RGB平面坐标（RGB图坐标）。 111 | 这里只记录一下最终测试程序的思路： 112 | ''' 113 | img_height, img_width = depth.shape[0], depth.shape[1] 114 | depth_ = np.zeros_like(depth) # [480 640] 115 | x = np.linspace(0, img_width - 1, img_width) # 640 116 | y = np.linspace(0, img_height - 1, img_height) # 480 117 | 118 | xx, yy = np.meshgrid(x, y) # 坐标网格化[img_width img_height] 119 | xx = np.reshape(xx, (1, -1)) # [1, img_width*img_height] 120 | yy = np.reshape(yy, (1, -1)) # [1, img_width*img_height] 121 | ones = np.ones_like(xx) # [1, img_width*img_height] 122 | 123 | pcoord_depth = np.concatenate((xx, yy, ones), axis=0) # [3, img_width*img_height], 像素坐标 124 | depth = np.reshape(depth, (1, img_height * img_width)) # [1, img_width*img_height] 125 | 126 | ccoord_depth = np.dot(intrinsics_depth_inv, pcoord_depth) * depth # 像素坐标-->归一化坐标-->相机坐标[3, img_width*img_height] 127 | 128 | ccoord_depth[1, :] = - ccoord_depth[1, :] 129 | ccoord_depth[2, :] = - ccoord_depth[2, :] 130 | 131 | ccoord_depth = np.concatenate((ccoord_depth, ones), axis=0) # [4, img_width*img_height] 132 | ccoord_color = np.dot(calibration_extrinsics, ccoord_depth) # [3, img_width*img_height],RGB相机坐标 133 | 134 | ccoord_color = ccoord_color[0:3, :] 135 | ccoord_color[1, :] = - ccoord_color[1, :] 136 | ccoord_color[2, :] = depth 137 | 138 | pcoord_color = np.dot(intrinsics_color, ccoord_color) # RGB像素坐标*Z 139 | pcoord_color = pcoord_color[:, pcoord_color[2, :] != 0] 140 | 141 | pcoord_color[0, :] = pcoord_color[0, :] / pcoord_color[2, :] + 0.5 # RGB像素坐标 142 | pcoord_color[0, :] = pcoord_color[0, :].astype(int) 143 | pcoord_color[1, :] = pcoord_color[1, :] / pcoord_color[2, :] + 0.5 144 | pcoord_color[1, :] = pcoord_color[1, :].astype(int) 145 | pcoord_color = pcoord_color[:, pcoord_color[0, :] >= 0] 146 | pcoord_color = pcoord_color[:, pcoord_color[1, :] >= 0] 147 | 148 | pcoord_color = pcoord_color[:, pcoord_color[0, :] < img_width] 149 | pcoord_color = pcoord_color[:, pcoord_color[1, :] < img_height] 150 | 151 | depth_[pcoord_color[1, :].astype(int), 152 | pcoord_color[0, :].astype(int)] = pcoord_color[2, :] 153 | return depth_ 154 | 155 | 156 | def get_coord(depth, pose, intrinsics_color_inv, dataset): 157 | """Generate the ground truth scene coordinates from depth and pose. 158 | """ 159 | img_height, img_width = depth.shape[0], depth.shape[1] # 480 640 160 | mask = np.ones_like(depth) 161 | mask[depth == 0] = 0 # 深度为0处，数值为0，否则为1 162 | mask = np.reshape(mask, (img_height, img_width, 1)) # [480 640 1] 163 | x = np.linspace(0, img_width - 1, img_width) 164 | y = np.linspace(0, img_height - 1, img_height) 165 | 166 | xx, yy = np.meshgrid(x, y) 167 | 168 | xx = np.reshape(xx, (1, -1)) # [1, 640*480] 169 | yy = np.reshape(yy, (1, -1)) # [1, 640*480] 170 | ones = np.ones_like(xx) # [1, 640*480] 171 | pcoord = np.concatenate((xx, yy, ones), axis=0) # [3, 640*480],像素坐标 172 | 173 | depth = np.reshape(depth, (1, img_height * img_width)) # [1, 640*480] 174 | ccoord = np.dot(intrinsics_color_inv, pcoord) * depth # 相机坐标 [3 640*480] 175 | ccoord = np.concatenate((ccoord, ones), axis=0) # 相机坐标 [4 640*480] 176 | 177 | # if dataset == 'my': 178 | # scoord = np.dot(np.swapaxes(ccoord,0,1), pose) 179 | # else: 180 | scoord = np.dot(pose, ccoord) # 世界坐标 [3 640*480] 181 | scoord = np.swapaxes(scoord, 0, 1) # 世界坐标 [640*480 3] 182 | 183 | scoord = scoord[:, 0:3] 184 | scoord = np.reshape(scoord, (img_height, img_width, 3)) # 世界坐标 [480 640 3] 185 | scoord = scoord * mask 186 | mask = np.reshape(mask, (img_height, img_width)) # [480 640] 187 | 188 | return scoord, mask 189 | 190 | 191 | # 数据增强操作 192 | def data_aug(img, coord, mask, aug=True, sp_coords=None): 193 | img_h, img_w = img.shape[0:2] 194 | if aug: 195 | trans_x = random.uniform(-0.2, 0.2) # 平移 196 | trans_y = random.uniform(-0.2, 0.2) 197 | 198 | aug_add = iaa.Add(random.randint(-20, 20)) 199 | 200 | scale = random.uniform(0.7, 1.5) # 缩放 201 | rotate = random.uniform(-30, 30) # 旋转 202 | shear = random.uniform(-10, 10) # 裁剪 203 | 204 | aug_affine = iaa.Affine(scale=scale, rotate=rotate, 205 | shear=shear, translate_percent={"x": trans_x, "y": trans_y}) 206 | aug_affine_lbl = iaa.Affine(scale=scale, rotate=rotate, 207 | shear=shear, translate_percent={"x": trans_x, "y": trans_y}, 208 | order=0, cval=1) 209 | img = aug_add.augment_image(img) 210 | else: 211 | trans_x = random.randint(-3, 4) 212 | trans_y = random.randint(-3, 4) 213 | 214 | aug_affine = iaa.Affine(translate_px={"x": trans_x, "y": trans_y}) 215 | 216 | padding = torch.randint(0, 255, size=(img_h, img_w, 3)).data.numpy().astype(np.uint8) 217 | padding_mask = np.ones((img_h, img_w)).astype(np.uint8) 218 | 219 | img = aug_affine.augment_image(img) 220 | coord = aug_affine.augment_image(coord) 221 | mask = aug_affine.augment_image(mask) 222 | mask = np.round(mask) 223 | padding_mask = aug_affine.augment_image(padding_mask) 224 | img = img + (1 - np.expand_dims(padding_mask, axis=2)) * padding 225 | 226 | if isinstance(sp_coords, np.ndarray): 227 | ia_kpts = [] 228 | out_kpts = [] 229 | for i in range(sp_coords.shape[0]): 230 | # if np.isnan(sp_coords[i][0]): 231 | # ia_kpts.append(imgaug.Keypoint(x=0, y=0)) 232 | ia_kpts.append(imgaug.Keypoint(x=sp_coords[i][0], y=sp_coords[i][1])) 233 | ia_kpts = imgaug.KeypointsOnImage(ia_kpts, shape=img.shape) 234 | ia_kpts = aug_affine_lbl.augment_keypoints(ia_kpts) 235 | for i in range(len(ia_kpts)): 236 | out_kpts.append(np.array((ia_kpts[i].x, ia_kpts[i].y))) 237 | out_kpts = np.stack(out_kpts, axis=0) 238 | return img, coord, mask, out_kpts 239 | else: 240 | return img, coord, mask 241 | 242 | 243 | # img [480 640 3] 244 | # coord_img [60, 80, 3] 245 | # mask [60 80] 246 | def to_tensor(img, coord_img, mask): 247 | img = img.transpose(2, 0, 1) 248 | coord_img = coord_img.transpose(2, 0, 1) 249 | img = img / 255. 250 | img = img * 2. - 1. 251 | coord_img = coord_img / 1000. 252 | img = torch.from_numpy(img).float() 253 | coord_img = torch.from_numpy(coord_img).float() 254 | mask = torch.from_numpy(mask).float() 255 | return img, coord_img, mask 256 | 257 | 258 | def to_tensor_query(img, pose): 259 | img = img.transpose(2, 0, 1) 260 | img = img / 255. 261 | img = img * 2. - 1. 262 | img = torch.from_numpy(img).float() 263 | pose = torch.from_numpy(pose).float() 264 | return img, pose 265 | 266 | 267 | def data_aug_label(img, coord, mask, lbl, aug=True): 268 | img_h, img_w = img.shape[0:2] 269 | if aug: 270 | trans_x = random.uniform(-0.2, 0.2) 271 | trans_y = random.uniform(-0.2, 0.2) 272 | 273 | aug_add = iaa.Add(random.randint(-20, 20)) 274 | 275 | scale = random.uniform(0.7, 1.5) 276 | rotate = random.uniform(-30, 30) 277 | shear = random.uniform(-10, 10) 278 | 279 | aug_affine = iaa.Affine(scale=scale, rotate=rotate, 280 | shear=shear, translate_percent={"x": trans_x, "y": trans_y}) 281 | aug_affine_lbl = iaa.Affine(scale=scale, rotate=rotate, 282 | shear=shear, translate_percent={"x": trans_x, "y": trans_y}, 283 | order=0, cval=1) 284 | img = aug_add.augment_image(img) 285 | else: 286 | trans_x = random.randint(-3, 4) 287 | trans_y = random.randint(-3, 4) 288 | 289 | aug_affine = iaa.Affine(translate_px={"x": trans_x, "y": trans_y}) 290 | aug_affine_lbl = iaa.Affine(translate_px={"x": trans_x, "y": trans_y}, 291 | order=0, cval=1) 292 | 293 | padding = torch.randint(0, 255, size=(img_h, 294 | img_w, 3)).data.numpy().astype(np.uint8) 295 | padding_mask = np.ones((img_h, img_w)).astype(np.uint8) 296 | 297 | img = aug_affine.augment_image(img) 298 | coord = aug_affine.augment_image(coord) 299 | mask = aug_affine.augment_image(mask) 300 | mask = np.round(mask) 301 | lbl = aug_affine_lbl.augment_image(lbl) 302 | padding_mask = aug_affine.augment_image(padding_mask) 303 | img = img + (1 - np.expand_dims(padding_mask, axis=2)) * padding 304 | 305 | return img, coord, mask, lbl 306 | 307 | 308 | def one_hot(x, N=25): 309 | one_hot = torch.FloatTensor(N, x.size(0), x.size(1)).zero_() 310 | one_hot = one_hot.scatter_(0, x.unsqueeze(0), 1) 311 | return one_hot 312 | 313 | 314 | def to_tensor_label(img, coord_img, mask, lbl, N1=25, N2=25): 315 | img = img.transpose(2, 0, 1) 316 | coord_img = coord_img.transpose(2, 0, 1) 317 | 318 | img = img / 255. 319 | img = img * 2. - 1. 320 | 321 | coord_img = coord_img / 1000. 322 | 323 | img = torch.from_numpy(img).float() 324 | coord_img = torch.from_numpy(coord_img).float() 325 | mask = torch.from_numpy(mask).float() 326 | 327 | lbl = torch.from_numpy(lbl / 1.0).long() 328 | lbl_oh = one_hot(lbl, N=N1) 329 | return img, coord_img, mask, lbl, lbl_oh 330 | 331 | 332 | def to_tensor_query_label(img, pose): 333 | img = img.transpose(2, 0, 1) 334 | img = img / 255. 335 | img = img * 2. - 1. 336 | img = torch.from_numpy(img).float() 337 | pose = torch.from_numpy(pose).float() 338 | 339 | return img, pose 340 | -------------------------------------------------------------------------------- /ufvl_net/models/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) OpenMMLab. All rights reserved. 2 | from .backbones import * # noqa: F401,F403 3 | from .builder import (BACKBONES, ARCHITECTURE, HEADS, 4 | build_backbone, build_architecture, build_head) 5 | from .architecture import * # noqa: F401,F403 6 | from .heads import * # noqa: F401,F403 7 | 8 | __all__ = [ 9 | 'BACKBONES', 'HEADS', 'ARCHITECTURE', 'build_backbone', 10 | 'build_head', 'build_architecture' 11 | ] 12 | -------------------------------------------------------------------------------- /ufvl_net/models/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/models/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/models/__pycache__/builder.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/models/__pycache__/builder.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/models/architecture/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) OpenMMLab. All rights reserved. 2 | from .vis_loc import FDANET 3 | 4 | __all__ = ['FDANET'] 5 | -------------------------------------------------------------------------------- /ufvl_net/models/architecture/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/models/architecture/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/models/architecture/__pycache__/base.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/models/architecture/__pycache__/base.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/models/architecture/__pycache__/vis_loc.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/models/architecture/__pycache__/vis_loc.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/models/architecture/base.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) OpenMMLab. All rights reserved. 2 | from abc import ABCMeta, abstractmethod 3 | from collections import OrderedDict 4 | from typing import Sequence 5 | 6 | import mmcv 7 | import torch 8 | import torch.distributed as dist 9 | from mmcv.runner import BaseModule, auto_fp16 10 | 11 | class BaseClassifier(BaseModule, metaclass=ABCMeta): 12 | """Base class for classifiers.""" 13 | 14 | def __init__(self, init_cfg=None): 15 | super(BaseClassifier, self).__init__(init_cfg) 16 | self.fp16_enabled = False 17 | 18 | @property 19 | def with_neck(self): 20 | return hasattr(self, 'neck') and self.neck is not None 21 | 22 | @property 23 | def with_head(self): 24 | return hasattr(self, 'head') and self.head is not None 25 | 26 | @abstractmethod 27 | def extract_feat(self, imgs, stage=None): 28 | pass 29 | 30 | def extract_feats(self, imgs, stage=None): 31 | assert isinstance(imgs, Sequence) 32 | kwargs = {} if stage is None else {'stage': stage} 33 | for img in imgs: 34 | yield self.extract_feat(img, **kwargs) 35 | 36 | @abstractmethod 37 | def forward_train(self, imgs, **kwargs): 38 | """ 39 | Args: 40 | img (list[Tensor]): List of tensors of shape (1, C, H, W). 41 | Typically these should be mean centered and std scaled. 42 | kwargs (keyword arguments): Specific to concrete implementation. 43 | """ 44 | pass 45 | 46 | @abstractmethod 47 | def simple_test(self, img, **kwargs): 48 | pass 49 | 50 | def forward_test(self, imgs, **kwargs): 51 | """ 52 | Args: 53 | imgs (List[Tensor]): the outer list indicates test-time 54 | augmentations and inner Tensor should have a shape NxCxHxW, 55 | which contains all images in the batch. 56 | """ 57 | if isinstance(imgs, torch.Tensor): 58 | imgs = [imgs] 59 | for var, name in [(imgs, 'imgs')]: 60 | if not isinstance(var, list): 61 | raise TypeError(f'{name} must be a list, but got {type(var)}') 62 | 63 | if len(imgs) == 1: 64 | return self.simple_test(imgs[0], **kwargs) 65 | else: 66 | raise NotImplementedError('aug_test has not been implemented') 67 | 68 | @auto_fp16(apply_to=('img', )) 69 | def forward(self, img, return_loss=True, **kwargs): 70 | """Calls either forward_train or forward_test depending on whether 71 | return_loss=True. 72 | 73 | Note this setting will change the expected inputs. When 74 | `return_loss=True`, img and img_meta are single-nested (i.e. Tensor and 75 | List[dict]), and when `resturn_loss=False`, img and img_meta should be 76 | double nested (i.e. List[Tensor], List[List[dict]]), with the outer 77 | list indicating test time augmentations. 78 | """ 79 | if return_loss: 80 | return self.forward_train(img, **kwargs) 81 | else: 82 | return self.forward_test(img, **kwargs) 83 | 84 | def _parse_losses(self, losses): 85 | log_vars = OrderedDict() 86 | for loss_name, loss_value in losses.items(): 87 | if isinstance(loss_value, torch.Tensor): 88 | log_vars[loss_name] = loss_value.mean() 89 | elif isinstance(loss_value, list): 90 | log_vars[loss_name] = sum(_loss.mean() for _loss in loss_value) 91 | elif isinstance(loss_value, dict): 92 | for name, value in loss_value.items(): 93 | log_vars[name] = value 94 | else: 95 | raise TypeError( 96 | f'{loss_name} is not a tensor or list of tensors') 97 | 98 | loss = sum(_value for _key, _value in log_vars.items() 99 | if 'loss' in _key) 100 | 101 | log_vars['loss'] = loss 102 | for loss_name, loss_value in log_vars.items(): 103 | # reduce loss when distributed training 104 | if dist.is_available() and dist.is_initialized(): 105 | loss_value = loss_value.data.clone() 106 | dist.all_reduce(loss_value.div_(dist.get_world_size())) 107 | log_vars[loss_name] = loss_value.item() 108 | 109 | return loss, log_vars 110 | 111 | def train_step(self, data, optimizer=None, **kwargs): 112 | """The iteration step during training. 113 | 114 | This method defines an iteration step during training, except for the 115 | back propagation and optimizer updating, which are done in an optimizer 116 | hook. Note that in some complicated cases or models, the whole process 117 | including back propagation and optimizer updating are also defined in 118 | this method, such as GAN. 119 | 120 | Args: 121 | data (dict): The output of dataloader. 122 | optimizer (:obj:`torch.optim.Optimizer` | dict, optional): The 123 | optimizer of runner is passed to ``train_step()``. This 124 | argument is unused and reserved. 125 | 126 | Returns: 127 | dict: Dict of outputs. The following fields are contained. 128 | - loss (torch.Tensor): A tensor for back propagation, which \ 129 | can be a weighted sum of multiple losses. 130 | - log_vars (dict): Dict contains all the variables to be sent \ 131 | to the logger. 132 | - num_samples (int): Indicates the batch size (when the model \ 133 | is DDP, it means the batch size on each GPU), which is \ 134 | used for averaging the logs. 135 | """ 136 | losses = self(**data) 137 | loss, log_vars = self._parse_losses(losses) 138 | 139 | outputs = dict( 140 | loss=loss, log_vars=log_vars, num_samples=len(data['img'].data)) 141 | 142 | return outputs 143 | 144 | def val_step(self, data, optimizer=None, **kwargs): 145 | """The iteration step during validation. 146 | 147 | This method shares the same signature as :func:`train_step`, but used 148 | during val epochs. Note that the evaluation after training epochs is 149 | not implemented with this method, but an evaluation hook. 150 | 151 | Args: 152 | data (dict): The output of dataloader. 153 | optimizer (:obj:`torch.optim.Optimizer` | dict, optional): The 154 | optimizer of runner is passed to ``train_step()``. This 155 | argument is unused and reserved. 156 | 157 | Returns: 158 | dict: Dict of outputs. The following fields are contained. 159 | - loss (torch.Tensor): A tensor for back propagation, which \ 160 | can be a weighted sum of multiple losses. 161 | - log_vars (dict): Dict contains all the variables to be sent \ 162 | to the logger. 163 | - num_samples (int): Indicates the batch size (when the model \ 164 | is DDP, it means the batch size on each GPU), which is \ 165 | used for averaging the logs. 166 | """ 167 | losses = self(**data) 168 | loss, log_vars = self._parse_losses(losses) 169 | 170 | outputs = dict( 171 | loss=loss, log_vars=log_vars, num_samples=len(data['img'].data)) 172 | 173 | return outputs 174 | 175 | -------------------------------------------------------------------------------- /ufvl_net/models/architecture/pnpransac/build/temp.linux-x86_64-cpython-37/pnpransacpy.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/models/architecture/pnpransac/build/temp.linux-x86_64-cpython-37/pnpransacpy.o -------------------------------------------------------------------------------- /ufvl_net/models/architecture/pnpransac/pnpransac.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | PnP-RANSAC implementation based on DSAC++ 3 | Code: https://github.com/vislearn/LessMore 4 | Paper: https://arxiv.org/abs/1711.10228 5 | */ 6 | 7 | /* 8 | Copyright (c) 2016, TU Dresden 9 | Copyright (c) 2017, Heidelberg University 10 | All rights reserved. 11 | Redistribution and use in source and binary forms, with or without 12 | modification, are permitted provided that the following conditions are met: 13 | * Redistributions of source code must retain the above copyright 14 | notice, this list of conditions and the following disclaimer. 15 | * Redistributions in binary form must reproduce the above copyright 16 | notice, this list of conditions and the following disclaimer in the 17 | documentation and/or other materials provided with the distribution. 18 | * Neither the name of the TU Dresden, Heidelberg University nor the 19 | names of its contributors may be used to endorse or promote products 20 | derived from this software without specific prior written permission. 21 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND 22 | ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED 23 | WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 24 | DISCLAIMED. IN NO EVENT SHALL TU DRESDEN OR HEIDELBERG UNIVERSITY BE LIABLE FOR ANY 25 | DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES 26 | (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; 27 | LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND 28 | ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 29 | (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS 30 | SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 31 | */ 32 | 33 | #include "pnpransac.h" 34 | 35 | std::vector ThreadRand::generators; 36 | bool ThreadRand::initialised = false; 37 | 38 | void ThreadRand::forceInit(unsigned seed) 39 | { 40 | initialised = false; 41 | init(seed); 42 | } 43 | 44 | void ThreadRand::init(unsigned seed) 45 | { 46 | #pragma omp critical 47 | { 48 | if(!initialised) 49 | { 50 | unsigned nThreads = omp_get_max_threads(); 51 | 52 | for(unsigned i = 0; i < nThreads; i++) 53 | { 54 | generators.push_back(std::mt19937()); 55 | generators[i].seed(i+seed); 56 | } 57 | 58 | initialised = true; 59 | } 60 | } 61 | } 62 | 63 | int ThreadRand::irand(int min, int max, int tid) 64 | { 65 | std::uniform_int_distribution dist(min, max); 66 | 67 | unsigned threadID = omp_get_thread_num(); 68 | if(tid >= 0) threadID = tid; 69 | 70 | if(!initialised) init(); 71 | 72 | return dist(ThreadRand::generators[threadID]); 73 | } 74 | 75 | double ThreadRand::drand(double min, double max, int tid) 76 | { 77 | std::uniform_real_distribution dist(min, max); 78 | 79 | unsigned threadID = omp_get_thread_num(); 80 | if(tid >= 0) threadID = tid; 81 | 82 | if(!initialised) init(); 83 | 84 | return dist(ThreadRand::generators[threadID]); 85 | } 86 | 87 | double ThreadRand::dgauss(double mean, double stdDev, int tid) 88 | { 89 | std::normal_distribution dist(mean, stdDev); 90 | 91 | unsigned threadID = omp_get_thread_num(); 92 | if(tid >= 0) threadID = tid; 93 | 94 | if(!initialised) init(); 95 | 96 | return dist(ThreadRand::generators[threadID]); 97 | } 98 | 99 | int irand(int incMin, int excMax, int tid) 100 | { 101 | return ThreadRand::irand(incMin, excMax - 1, tid); 102 | } 103 | 104 | double drand(double incMin, double incMax,int tid) 105 | { 106 | return ThreadRand::drand(incMin, incMax, tid); 107 | } 108 | 109 | int igauss(int mean, int stdDev, int tid) 110 | { 111 | return (int) ThreadRand::dgauss(mean, stdDev, tid); 112 | } 113 | 114 | double dgauss(double mean, double stdDev, int tid) 115 | { 116 | return ThreadRand::dgauss(mean, stdDev, tid); 117 | } 118 | 119 | namespace poseSolver { 120 | 121 | std::pair getInvHyp(const std::pair& hyp) 122 | { 123 | cv::Mat_ hypR, trans = cv::Mat_::eye(4, 4); 124 | cv::Rodrigues(hyp.first, hypR); 125 | 126 | hypR.copyTo(trans.rowRange(0,3).colRange(0,3)); 127 | trans(0, 3) = hyp.second.at(0, 0); 128 | trans(1, 3) = hyp.second.at(0, 1); 129 | trans(2, 3) = hyp.second.at(0, 2); 130 | 131 | trans = trans.inv(); 132 | 133 | std::pair invHyp; 134 | cv::Rodrigues(trans.rowRange(0,3).colRange(0,3), invHyp.first); 135 | invHyp.second = cv::Mat_(1, 3); 136 | invHyp.second.at(0, 0) = trans(0, 3); 137 | invHyp.second.at(0, 1) = trans(1, 3); 138 | invHyp.second.at(0, 2) = trans(2, 3); 139 | 140 | return invHyp; 141 | } 142 | 143 | double calcAngularDistance(const std::pair & h1, const std::pair & h2) 144 | { 145 | cv::Mat r1, r2; 146 | cv::Rodrigues(h1.first, r1); 147 | cv::Rodrigues(h2.first, r2); 148 | 149 | cv::Mat rotDiff= r2 * r1.t(); 150 | double trace = cv::trace(rotDiff)[0]; 151 | 152 | trace = std::min(3.0, std::max(-1.0, trace)); 153 | return 180*acos((trace-1.0)/2.0)/CV_PI; 154 | } 155 | 156 | double maxLoss(const std::pair & h1, const std::pair & h2) 157 | { 158 | // measure loss of inverted poses (camera pose instead of scene pose) 159 | std::pair invH1 = getInvHyp(h1); 160 | std::pair invH2 = getInvHyp(h2); 161 | 162 | double rotErr = calcAngularDistance(invH1, invH2); 163 | double tErr = cv::norm(invH1.second - invH2.second); 164 | 165 | return std::max(rotErr, tErr * 100); 166 | } 167 | 168 | inline bool safeSolvePnP( 169 | const std::vector& objPts, 170 | const std::vector& imgPts, 171 | const cv::Mat& camMat, 172 | const cv::Mat& distCoeffs, 173 | cv::Mat& rot, 174 | cv::Mat& trans, 175 | bool extrinsicGuess, 176 | int methodFlag) 177 | { 178 | if(rot.type() == 0) rot = cv::Mat_::zeros(1, 3); 179 | if(trans.type() == 0) trans= cv::Mat_::zeros(1, 3); 180 | 181 | if(!cv::solvePnP(objPts, imgPts, camMat, distCoeffs, rot, trans, extrinsicGuess,methodFlag)) 182 | { 183 | rot = cv::Mat_::zeros(1, 3); 184 | trans = cv::Mat_::zeros(1, 3); 185 | return false; 186 | } 187 | return true; 188 | } 189 | 190 | PnPRANSAC::PnPRANSAC () { 191 | this->camMat = cv::Mat_::eye(3, 3); 192 | } 193 | 194 | PnPRANSAC::PnPRANSAC (float fx, float fy, float cx, float cy) { 195 | this->camMat = cv::Mat_::eye(3, 3); 196 | this->camMat(0,0) = fx; 197 | this->camMat(1,1) = fy; 198 | this->camMat(0,2) = cx; 199 | this->camMat(1,2) = cy; 200 | } 201 | 202 | PnPRANSAC::~PnPRANSAC () {} 203 | 204 | void PnPRANSAC::camMatUpdate(float fx, float fy, float cx, float cy){ 205 | this->camMat = cv::Mat_::eye(3, 3); 206 | this->camMat(0,0) = fx; 207 | this->camMat(1,1) = fy; 208 | this->camMat(0,2) = cx; 209 | this->camMat(1,2) = cy; 210 | } 211 | 212 | double* PnPRANSAC::RANSACLoop( 213 | float* imgPts_, 214 | float* objPts_, 215 | int nPts, 216 | int objHyps) 217 | { 218 | int inlierThreshold2D = 10; 219 | int refSteps = 100; 220 | 221 | std::vector imgPts(nPts); 222 | std::vector objPts(nPts); 223 | #pragma omp parallel for 224 | for(unsigned i=0; i> sampledImgPts(objHyps); 231 | std::vector> sampledObjPts(objHyps); 232 | std::vector> rotHyp(objHyps); 233 | std::vector> tHyp(objHyps); 234 | std::vector scores(objHyps); 235 | std::vector> reproDiff(objHyps); 236 | 237 | // sample hypotheses 238 | #pragma omp parallel for 239 | for(int h = 0; h < objHyps; h++) 240 | while(true) 241 | { 242 | std::vector projections; 243 | std::vector alreadyChosen(nPts,0); 244 | sampledImgPts[h].clear(); 245 | sampledObjPts[h].clear(); 246 | 247 | for(int j = 0; j < 4; j++) 248 | { 249 | int idx = irand(0, nPts); 250 | 251 | if(alreadyChosen[idx] > 0) 252 | { 253 | j--; 254 | continue; 255 | } 256 | 257 | alreadyChosen[idx] = 1; 258 | 259 | sampledImgPts[h].push_back(imgPts[idx]); // 2D location in the original RGB image 260 | sampledObjPts[h].push_back(objPts[idx]); // 3D object coordinate 261 | 262 | } 263 | 264 | if(!safeSolvePnP(sampledObjPts[h], sampledImgPts[h], this->camMat, cv::Mat(), rotHyp[h], tHyp[h], false, CV_P3P)) 265 | { 266 | continue; 267 | 268 | } 269 | 270 | cv::projectPoints(sampledObjPts[h], rotHyp[h], tHyp[h], this->camMat, cv::Mat(), projections); 271 | 272 | // check reconstruction, 4 sampled points should be reconstructed perfectly 273 | bool foundOutlier = false; 274 | for(unsigned j = 0; j < sampledImgPts[h].size(); j++) 275 | { 276 | if(cv::norm(sampledImgPts[h][j] - projections[j]) < inlierThreshold2D) 277 | continue; 278 | foundOutlier = true; 279 | break; 280 | } 281 | if(foundOutlier) 282 | continue; 283 | else{ 284 | // compute reprojection error and hypothesis score 285 | std::vector projections; 286 | cv::projectPoints(objPts, rotHyp[h], tHyp[h], this->camMat, cv::Mat(), projections); 287 | std::vector diff(nPts); 288 | float score = 0.; 289 | for(unsigned pt = 0; pt < imgPts.size(); pt++) 290 | { 291 | float err = cv::norm(imgPts[pt] - projections[pt]); 292 | diff[pt] = err; 293 | score = score + (1. / (1. + std::exp(-(0.5*(err-inlierThreshold2D))))); 294 | } 295 | reproDiff[h] = diff; 296 | scores[h] = score; 297 | break; 298 | } 299 | } 300 | 301 | int hypIdx = std::min_element(scores.begin(),scores.end()) - scores.begin(); // select winning hypothesis 302 | 303 | double convergenceThresh = 0.01; // stop refinement if 6D pose vector converges 304 | 305 | std::vector localDiff = reproDiff[hypIdx]; 306 | 307 | for(int rStep = 0; rStep < refSteps; rStep++) 308 | { 309 | // collect inliers 310 | std::vector localImgPts; 311 | std::vector localObjPts; 312 | 313 | for(int pt = 0; pt < nPts; pt++) 314 | { 315 | if(localDiff[pt] < inlierThreshold2D) 316 | { 317 | localImgPts.push_back(imgPts[pt]); 318 | localObjPts.push_back(objPts[pt]); 319 | } 320 | } 321 | 322 | if(localImgPts.size() < 4) 323 | break; 324 | 325 | // recalculate pose 326 | cv::Mat_ rotNew = rotHyp[hypIdx].clone(); 327 | cv::Mat_ tNew = tHyp[hypIdx].clone(); 328 | 329 | if(!safeSolvePnP(localObjPts, localImgPts, this->camMat, cv::Mat(), rotNew, tNew, true, (localImgPts.size() > 4) ? CV_ITERATIVE : CV_P3P)) 330 | break; //abort if PnP fails 331 | std::pair hypNew; 332 | std::pair hypOld; 333 | hypNew.first = rotNew; 334 | hypNew.second = tNew; 335 | 336 | hypOld.first = rotHyp[hypIdx]; 337 | hypOld.second = tHyp[hypIdx]; 338 | if(maxLoss(hypNew, hypOld) < convergenceThresh) 339 | break; // convergned 340 | 341 | rotHyp[hypIdx] = rotNew; 342 | tHyp[hypIdx] = tNew; 343 | 344 | // recalculate pose errors 345 | std::vector projections; 346 | cv::projectPoints(objPts, rotHyp[hypIdx], tHyp[hypIdx], this->camMat, cv::Mat(), projections); 347 | std::vector diff(nPts); 348 | 349 | #pragma omp parallel for 350 | for(unsigned pt = 0; pt < imgPts.size(); pt++) 351 | { 352 | float err = cv::norm(imgPts[pt] - projections[pt]); 353 | diff[pt] = err; 354 | } 355 | localDiff = diff; 356 | } 357 | 358 | static double pose[6]; 359 | for (int i = 0; i < 3; i++) 360 | pose[i] = rotHyp[hypIdx](0,i); 361 | for (int i = 3; i < 6; i++) 362 | pose[i] = tHyp[hypIdx](0,i-3); 363 | return pose; 364 | } 365 | } 366 | 367 | -------------------------------------------------------------------------------- /ufvl_net/models/architecture/pnpransac/pnpransac.cpython-37m-x86_64-linux-gnu.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/models/architecture/pnpransac/pnpransac.cpython-37m-x86_64-linux-gnu.so -------------------------------------------------------------------------------- /ufvl_net/models/architecture/pnpransac/pnpransac.h: -------------------------------------------------------------------------------- 1 | /* 2 | PnP-RANSAC implementation based on DSAC++ 3 | Code: https://github.com/vislearn/LessMore 4 | Paper: https://arxiv.org/abs/1711.10228 5 | */ 6 | 7 | /* 8 | Copyright (c) 2016, TU Dresden 9 | Copyright (c) 2017, Heidelberg University 10 | All rights reserved. 11 | Redistribution and use in source and binary forms, with or without 12 | modification, are permitted provided that the following conditions are met: 13 | * Redistributions of source code must retain the above copyright 14 | notice, this list of conditions and the following disclaimer. 15 | * Redistributions in binary form must reproduce the above copyright 16 | notice, this list of conditions and the following disclaimer in the 17 | documentation and/or other materials provided with the distribution. 18 | * Neither the name of the TU Dresden, Heidelberg University nor the 19 | names of its contributors may be used to endorse or promote products 20 | derived from this software without specific prior written permission. 21 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND 22 | ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED 23 | WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 24 | DISCLAIMED. IN NO EVENT SHALL TU DRESDEN OR HEIDELBERG UNIVERSITY BE LIABLE FOR ANY 25 | DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES 26 | (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; 27 | LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND 28 | ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 29 | (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS 30 | SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 31 | */ 32 | 33 | #pragma once 34 | 35 | #include 36 | #include 37 | #include 38 | #include 39 | #include 40 | #include "opencv2/opencv.hpp" 41 | 42 | /** Classes and methods for generating random numbers in multi-threaded programs. */ 43 | 44 | /** 45 | * @brief Provides random numbers for multiple threads. 46 | * 47 | * Singelton class. Holds a random number generator for each thread and gives random numbers for the current thread. 48 | */ 49 | class ThreadRand 50 | { 51 | public: 52 | /** 53 | * @brief Returns a random integer (uniform distribution). 54 | * 55 | * @param min Minimum value of the random integer (inclusive). 56 | * @param max Maximum value of the random integer (exclusive). 57 | * @param tid Optional parameter. ID of the thread to use. If not given, the method will obtain the thread ID itself. 58 | * @return int Random integer value. 59 | */ 60 | static int irand(int min, int max, int tid = -1); 61 | 62 | /** 63 | * @brief Returns a random double value (uniform distribution). 64 | * 65 | * @param min Minimum value of the random double (inclusive). 66 | * @param max Maximum value of the random double (inclusive). 67 | * @param tid Optional parameter. ID of the thread to use. If not given, the method will obtain the thread ID itself. 68 | * @return double Random double value. 69 | */ 70 | static double drand(double min, double max, int tid = -1); 71 | 72 | /** 73 | * @brief Returns a random double value (Gauss distribution). 74 | * 75 | * @param mean Mean of the Gauss distribution to sample from. 76 | * @param stdDev Standard deviation of the Gauss distribution to sample from. 77 | * @param tid Optional parameter. ID of the thread to use. If not given, the method will obtain the thread ID itself. 78 | * @return double Random double value. 79 | */ 80 | static double dgauss(double mean, double stdDev, int tid = -1); 81 | 82 | /** 83 | * @brief Re-Initialize the object with the given seed. 84 | * 85 | * @param seed Seed to initialize the random number generators (seed is incremented by one for each generator). 86 | * @return void 87 | */ 88 | static void forceInit(unsigned seed); 89 | 90 | private: 91 | /** 92 | * @brief List of random number generators. One for each thread. 93 | * 94 | */ 95 | static std::vector generators; 96 | /** 97 | * @brief True if the class has been initialized already 98 | */ 99 | static bool initialised; 100 | /** 101 | * @brief Initialize class with the given seed. 102 | * 103 | * Method will create a random number generator for each thread. The given seed 104 | * will be incremented by one for each generator. This methods is automatically 105 | * called when this calss is used the first time. 106 | * 107 | * @param seed Optional parameter. Seed to be used when initializing the generators. Will be incremented by one for each generator. 108 | * @return void 109 | */ 110 | static void init(unsigned seed = 1305); 111 | }; 112 | 113 | /** 114 | * @brief Returns a random integer (uniform distribution). 115 | * 116 | * This method used the ThreadRand class. 117 | * 118 | * @param min Minimum value of the random integer (inclusive). 119 | * @param max Maximum value of the random integer (exclusive). 120 | * @param tid Optional parameter. ID of the thread to use. If not given, the method will obtain the thread ID itself. 121 | * @return int Random integer value. 122 | */ 123 | int irand(int incMin, int excMax, int tid = -1); 124 | /** 125 | * @brief Returns a random double value (uniform distribution). 126 | * 127 | * This method used the ThreadRand class. 128 | * 129 | * @param min Minimum value of the random double (inclusive). 130 | * @param max Maximum value of the random double (inclusive). 131 | * @param tid Optional parameter. ID of the thread to use. If not given, the method will obtain the thread ID itself. 132 | * @return double Random double value. 133 | */ 134 | double drand(double incMin, double incMax, int tid = -1); 135 | 136 | /** 137 | * @brief Returns a random integer value (Gauss distribution). 138 | * 139 | * This method used the ThreadRand class. 140 | * 141 | * @param mean Mean of the Gauss distribution to sample from. 142 | * @param stdDev Standard deviation of the Gauss distribution to sample from. 143 | * @param tid Optional parameter. ID of the thread to use. If not given, the method will obtain the thread ID itself. 144 | * @return double Random integer value. 145 | */ 146 | int igauss(int mean, int stdDev, int tid = -1); 147 | 148 | /** 149 | * @brief Returns a random double value (Gauss distribution). 150 | * 151 | * This method used the ThreadRand class. 152 | * 153 | * @param mean Mean of the Gauss distribution to sample from. 154 | * @param stdDev Standard deviation of the Gauss distribution to sample from. 155 | * @param tid Optional parameter. ID of the thread to use. If not given, the method will obtain the thread ID itself. 156 | * @return double Random double value. 157 | */ 158 | double dgauss(double mean, double stdDev, int tid = -1); 159 | 160 | namespace poseSolver { 161 | 162 | /** 163 | * @brief Inverts a given transformation. 164 | * @param hyp Input transformation. 165 | * @return Inverted transformation. 166 | */ 167 | std::pair getInvHyp(const std::pair& hyp); 168 | 169 | /** 170 | * @brief Maximum of translational error (cm) and rotational error (deg) between two pose hypothesis. 171 | * @param h1 Pose 1. 172 | * @param h2 Pose 2. 173 | * @return Loss. 174 | */ 175 | double maxLoss(const std::pair& h1, const std::pair& h2); 176 | 177 | /** 178 | * @brief Calculates the rotational distance in degree between two transformations. 179 | * Translation will be ignored. 180 | * 181 | * @param h1 Transformation 1. 182 | * @param h2 Transformation 2. 183 | * @return Angle in degree. 184 | */ 185 | double calcAngularDistance(const std::pair& h1, const std::pair& h2); 186 | 187 | /** 188 | * @brief Wrapper around the OpenCV PnP function that returns a zero pose in case PnP fails. See also documentation of cv::solvePnP. 189 | * @param objPts List of 3D points. 190 | * @param imgPts Corresponding 2D points. 191 | * @param camMat Calibration matrix of the camera. 192 | * @param distCoeffs Distortion coefficients. 193 | * @param rot Output parameter. Camera rotation. 194 | * @param trans Output parameter. Camera translation. 195 | * @param extrinsicGuess If true uses input rot and trans as initialization. 196 | * @param methodFlag Specifies the PnP algorithm to be used. 197 | * @return True if PnP succeeds. 198 | */ 199 | inline bool safeSolvePnP( 200 | const std::vector& objPts, 201 | const std::vector& imgPts, 202 | const cv::Mat& camMat, 203 | const cv::Mat& distCoeffs, 204 | cv::Mat& rot, 205 | cv::Mat& trans, 206 | bool extrinsicGuess, 207 | int methodFlag); 208 | 209 | class PnPRANSAC{ 210 | public: 211 | cv::Mat_ camMat; 212 | PnPRANSAC(); 213 | 214 | PnPRANSAC(float fx, float fy, float cx, float cy); 215 | 216 | ~PnPRANSAC(); 217 | 218 | void camMatUpdate(float fx, float fy, float cx, float cy); 219 | 220 | double* RANSACLoop(float* imgPts, float* objPts, int nPts, int objHyps); 221 | }; 222 | 223 | } 224 | 225 | -------------------------------------------------------------------------------- /ufvl_net/models/architecture/pnpransac/pnpransacpy.pxd: -------------------------------------------------------------------------------- 1 | from libcpp cimport bool 2 | from libcpp.vector cimport vector 3 | 4 | cdef extern from "pnpransac.cpp": 5 | pass 6 | 7 | cdef extern from "pnpransac.h" namespace "poseSolver": 8 | cdef cppclass PnPRANSAC: 9 | PnPRANSAC() except + 10 | PnPRANSAC(float, float, float, float) except + 11 | void camMatUpdate(float, float, float, float) 12 | double* RANSACLoop(float*, float*, int, int) -------------------------------------------------------------------------------- /ufvl_net/models/architecture/pnpransac/pnpransacpy.pyx: -------------------------------------------------------------------------------- 1 | # distutils: language = c++ 2 | 3 | import numpy as np 4 | cimport numpy as np 5 | 6 | from pnpransacpy cimport PnPRANSAC 7 | 8 | cdef class pnpransac: 9 | cdef PnPRANSAC c_pnpransac 10 | 11 | def __cinit__(self, float fx, float fy, float cx, float cy): 12 | self.c_pnpransac = PnPRANSAC(fx, fy, cx, cy) 13 | 14 | def update_camMat(self, float fx, float fy, float cx, float cy): 15 | self.c_pnpransac.camMatUpdate(fx, fy, cx, cy) 16 | 17 | def RANSAC_loop(self, np.ndarray[double, ndim=2, mode="c"] img_pts, 18 | np.ndarray[double, ndim=2, mode="c"] obj_pts, int n_hyp): 19 | cdef float[:, :] img_pts_ = img_pts.astype(np.float32) 20 | cdef float[:, :] obj_pts_ = obj_pts.astype(np.float32) 21 | cdef int n_pts 22 | n_pts = img_pts_.shape[0] 23 | assert img_pts_.shape[0] == obj_pts_.shape[0] 24 | cdef double* pose 25 | pose = self.c_pnpransac.RANSACLoop(&img_pts_[0,0], &obj_pts_[0,0], 26 | n_pts, n_hyp) 27 | rot = np.array([pose[0],pose[1],pose[2]]) 28 | transl = np.array([pose[3],pose[4],pose[5]]) 29 | return rot, transl -------------------------------------------------------------------------------- /ufvl_net/models/architecture/pnpransac/setup.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import os 3 | from distutils.core import setup 4 | from distutils.extension import Extension 5 | 6 | from Cython.Distutils import build_ext 7 | from Cython.Build import cythonize 8 | 9 | # where to find opencv headers and libraries 10 | cv_include_dir = os.path.join(sys.prefix, 'include') 11 | cv_library_dir = os.path.join(sys.prefix, 'lib') 12 | 13 | ext_modules = [ 14 | Extension( 15 | "pnpransac", 16 | sources=["pnpransacpy.pyx"], 17 | language="c++", 18 | include_dirs=[cv_include_dir,'/home/dk/.conda/envs/HSCXT/lib/python3.7/site-packages/numpy/core/include/'], 19 | library_dirs=[cv_library_dir], 20 | libraries=['opencv_core','opencv_calib3d'], 21 | extra_compile_args=['-fopenmp','-std=c++11'], 22 | ) 23 | ] 24 | 25 | setup( 26 | name='pnpransac', 27 | cmdclass={'build_ext': build_ext}, 28 | ext_modules=cythonize(ext_modules), 29 | ) -------------------------------------------------------------------------------- /ufvl_net/models/architecture/vis_loc.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) OpenMMLab. All rights reserved. 2 | import os 3 | import sys 4 | import cv2 5 | import numpy as np 6 | 7 | import torch 8 | import torch.nn as nn 9 | from mmcv.runner import BaseModule 10 | 11 | from ..builder import ARCHITECTURE, build_backbone, build_head 12 | from ..utils.augment import Augments 13 | from .base import BaseClassifier 14 | 15 | @ARCHITECTURE.register_module() 16 | class FDANET(BaseClassifier): 17 | def __init__(self, 18 | backbone, 19 | neck=None, 20 | head=None, 21 | pretrained=None, 22 | train_cfg=None, 23 | init_cfg=None, 24 | dataset="7Scenes"): 25 | super(FDANET, self).__init__(init_cfg) 26 | 27 | self.backbone = build_backbone(backbone) 28 | self.head = build_head(head) 29 | self.loss = EuclideanLoss_with_Uncertainty() 30 | self.dataset = dataset 31 | 32 | def forward_test(self, img, gt_lables=None, **kwargs): 33 | import numpy as np 34 | import sys 35 | import os 36 | import cv2 37 | sys.path.append("/home/dk/ufvl_net/ufvl_net/models/architecture/pnpransac") 38 | import pnpransac 39 | if self.dataset == "7Scenes": 40 | intrinsics_color = np.array([[525.0, 0.0, 320.0], 41 | [0.0, 525.0, 240.0], 42 | [0.0, 0.0, 1.0]]) 43 | elif self.dataset == "12Scenes": 44 | intrinsics_color = np.array([[572.0, 0.0, 320.0], 45 | [0.0, 572.0, 240.0], 46 | [0.0, 0.0, 1.0]]) 47 | pose_solver = pnpransac.pnpransac(intrinsics_color[0, 0], intrinsics_color[1, 1], intrinsics_color[0, 2], 48 | intrinsics_color[1, 2]) 49 | 50 | def get_pose_err(pose_gt, pose_est): 51 | transl_err = np.linalg.norm(pose_gt[0:3, 3] - pose_est[0:3, 3]) 52 | rot_err = pose_est[0:3, 0:3].T.dot(pose_gt[0:3, 0:3]) 53 | rot_err = cv2.Rodrigues(rot_err)[0] # 旋转向量 [3 1] 54 | rot_err = np.reshape(rot_err, (1, 3)) # 旋转向量 [1 3] 55 | rot_err = np.reshape(np.linalg.norm(rot_err, axis=1), -1) / np.pi * 180. # 二范数即转角 56 | return transl_err, rot_err[0] 57 | 58 | x = np.linspace(4, 640 - 4, 80) 59 | y = np.linspace(4, 480 - 4, 60) 60 | 61 | xx, yy = np.meshgrid(x, y) # [60 80] 62 | pcoord = np.concatenate((np.expand_dims(xx, axis=2), np.expand_dims(yy, axis=2)), axis=2) 63 | 64 | x = self.backbone(img) 65 | coord, uncertainty = self.head(x[0]) 66 | coord = np.transpose(coord.cpu().data.numpy()[0, :, :, :], (1, 2, 0)) # [3 60 80]->[60 80 3] 67 | uncertainty = np.transpose(uncertainty[0].cpu().data.numpy(), (1, 2, 0)) 68 | coord = np.concatenate([coord, uncertainty], axis=2) # [60 80 4] 69 | coord = np.ascontiguousarray(coord) 70 | pcoord = np.ascontiguousarray(pcoord) 71 | 72 | pcoord = pcoord.reshape(-1, 2) 73 | coords = coord[:, :, 0:3].reshape(-1, 3) 74 | confidences = coord[:, :, 3].flatten().tolist() 75 | 76 | coords_filtered = [] 77 | coords_filtered_2D = [] 78 | for i in range(len(confidences)): 79 | if confidences[i] > 0: 80 | coords_filtered.append(coords[i]) 81 | coords_filtered_2D.append(pcoord[i]) 82 | 83 | coords_filtered = np.vstack(coords_filtered) 84 | coords_filtered_2D = np.vstack(coords_filtered_2D) 85 | 86 | rot, transl = pose_solver.RANSAC_loop(coords_filtered_2D.astype(np.float64), coords_filtered.astype(np.float64), 87 | 256) # 预测结果,每次取256组点进行PNP Tcw 88 | pose_gt = gt_lables.cpu().numpy()[0, :, :] # [4 4] 89 | pose_est = np.eye(4) # [4 4] 90 | pose_est[0:3, 0:3] = cv2.Rodrigues(rot)[0].T # Rwc 91 | pose_est[0:3, 3] = -np.dot(pose_est[0:3, 0:3], transl) # twc 92 | 93 | transl_err, rot_err = get_pose_err(pose_gt, pose_est) 94 | 95 | return dict(trans_error_med=transl_err, rot_err_med=rot_err) 96 | 97 | def forward_train(self, img, gt_lables, mask, **kwargs): 98 | x = self.backbone(img) 99 | coord, uncer = self.head(x[0]) 100 | loss, accuracy = self.loss(coord, gt_lables, mask, uncer) 101 | losses = dict(loss=loss, accuracy=accuracy) 102 | return losses 103 | 104 | def forward(self, img, gt_lables, mask=None, return_loss=True, **kwargs): 105 | """Calls either forward_train or forward_test depending on whether 106 | return_loss=True. 107 | 108 | Note this setting will change the expected inputs. When 109 | `return_loss=True`, img and img_meta are single-nested (i.e. Tensor and 110 | List[dict]), and when `resturn_loss=False`, img and img_meta should be 111 | double nested (i.e. List[Tensor], List[List[dict]]), with the outer 112 | list indicating test time augmentations. 113 | """ 114 | if return_loss: 115 | assert mask is not None 116 | return self.forward_train(img, gt_lables, mask, **kwargs) 117 | else: 118 | return self.forward_test(img, gt_lables, **kwargs) 119 | 120 | def extract_feat(self, img, stage='neck'): 121 | pass 122 | 123 | def simple_test(self, img, img_metas=None, **kwargs): 124 | pass 125 | 126 | 127 | class EuclideanLoss_with_Uncertainty(nn.Module): 128 | def __init__(self): 129 | super(EuclideanLoss_with_Uncertainty, self).__init__() 130 | self.pdist = nn.PairwiseDistance(p=2) 131 | 132 | def forward(self, pred, target, mask, certainty): 133 | loss_reg = self.pdist(pred.permute(0,2,3,1), target.permute(0,2,3,1)) 134 | certainty_map = torch.clamp(certainty, 1e-6) 135 | loss_map = 3 * torch.log(certainty_map) + loss_reg / (2 * certainty_map.pow(2)) 136 | 137 | loss_map = loss_map * mask 138 | loss =torch.sum(loss_map) / mask.sum() 139 | 140 | if mask is not None: 141 | valid_pixel = mask.sum() + 1 142 | diff_coord_map = mask * loss_reg 143 | 144 | thres_coord_map = torch.clamp(diff_coord_map - 0.05, 0) 145 | num_accurate = valid_pixel - thres_coord_map.nonzero().shape[0] 146 | accuracy = num_accurate / valid_pixel 147 | loss1 = torch.sum(loss_reg*mask) / mask.sum() 148 | return loss, accuracy 149 | -------------------------------------------------------------------------------- /ufvl_net/models/backbones/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) OpenMMLab. All rights reserved. 2 | from .resnet import ResNet, ResNetV1c, ResNetV1d 3 | from .seresnet import SEResNet 4 | 5 | __all__ = [ 6 | 'ResNet', 'ResNetV1c', 'ResNetV1d', 'SEResNet' 7 | ] 8 | -------------------------------------------------------------------------------- /ufvl_net/models/backbones/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/models/backbones/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/models/backbones/__pycache__/base_backbone.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/models/backbones/__pycache__/base_backbone.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/models/backbones/__pycache__/resnet.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/models/backbones/__pycache__/resnet.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/models/backbones/__pycache__/seresnet.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/models/backbones/__pycache__/seresnet.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/models/backbones/base_backbone.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) OpenMMLab. All rights reserved. 2 | from abc import ABCMeta, abstractmethod 3 | 4 | from mmcv.runner import BaseModule 5 | 6 | 7 | class BaseBackbone(BaseModule, metaclass=ABCMeta): 8 | """Base backbone. 9 | 10 | This class defines the basic functions of a backbone. Any backbone that 11 | inherits this class should at least define its own `forward` function. 12 | """ 13 | 14 | def __init__(self, init_cfg=None): 15 | super(BaseBackbone, self).__init__(init_cfg) 16 | 17 | @abstractmethod 18 | def forward(self, x): 19 | """Forward computation. 20 | 21 | Args: 22 | x (tensor | tuple[tensor]): x could be a Torch.tensor or a tuple of 23 | Torch.tensor, containing input data for forward computation. 24 | """ 25 | pass 26 | 27 | def train(self, mode=True): 28 | """Set module status before forward computation. 29 | 30 | Args: 31 | mode (bool): Whether it is train_mode or test_mode 32 | """ 33 | super(BaseBackbone, self).train(mode) 34 | -------------------------------------------------------------------------------- /ufvl_net/models/backbones/seresnet.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) OpenMMLab. All rights reserved. 2 | import torch.utils.checkpoint as cp 3 | 4 | from ..builder import BACKBONES 5 | from ..utils.se_layer import SELayer 6 | from .resnet import BasicBlock, Bottleneck, ResLayer, ResNet 7 | 8 | 9 | class SEBasicBlock(BasicBlock): 10 | def __init__(self, in_channels, out_channels, se_ratio=16, **kwargs) : 11 | super(SEBasicBlock, self).__init__(in_channels, out_channels, **kwargs) 12 | self.se_layer = SELayer(out_channels, ratio=se_ratio) 13 | 14 | def forward(self, x): 15 | 16 | def _inner_forward(x): 17 | identity = x 18 | 19 | out = self.conv1(x) 20 | out = self.norm1(out) 21 | out = self.relu(out) 22 | 23 | out = self.conv2(out) 24 | out = self.norm2(out) 25 | 26 | out = self.se_layer(out) 27 | if self.downsample is not None: 28 | identity = self.downsample(x) 29 | 30 | out = self.drop_path(out) 31 | 32 | out += identity 33 | 34 | return out 35 | 36 | if self.with_cp and x.requires_grad: 37 | out = cp.checkpoint(_inner_forward, x) 38 | else: 39 | out = _inner_forward(x) 40 | 41 | out = self.relu(out) 42 | 43 | return out 44 | 45 | class SEBottleneck(Bottleneck): 46 | """SEBottleneck block for SEResNet. 47 | 48 | Args: 49 | in_channels (int): The input channels of the SEBottleneck block. 50 | out_channels (int): The output channel of the SEBottleneck block. 51 | se_ratio (int): Squeeze ratio in SELayer. Default: 16 52 | """ 53 | 54 | def __init__(self, in_channels, out_channels, se_ratio=16, **kwargs): 55 | super(SEBottleneck, self).__init__(in_channels, out_channels, **kwargs) 56 | self.se_layer = SELayer(out_channels, ratio=se_ratio) 57 | 58 | def forward(self, x): 59 | 60 | def _inner_forward(x): 61 | identity = x 62 | 63 | out = self.conv1(x) 64 | out = self.norm1(out) 65 | out = self.relu(out) 66 | 67 | out = self.conv2(out) 68 | out = self.norm2(out) 69 | out = self.relu(out) 70 | 71 | out = self.conv3(out) 72 | out = self.norm3(out) 73 | 74 | out = self.se_layer(out) 75 | 76 | if self.downsample is not None: 77 | identity = self.downsample(x) 78 | 79 | out += identity 80 | 81 | return out 82 | 83 | if self.with_cp and x.requires_grad: 84 | out = cp.checkpoint(_inner_forward, x) 85 | else: 86 | out = _inner_forward(x) 87 | 88 | out = self.relu(out) 89 | 90 | return out 91 | 92 | 93 | @BACKBONES.register_module() 94 | class SEResNet(ResNet): 95 | """SEResNet backbone. 96 | 97 | Please refer to the `paper `__ for 98 | details. 99 | 100 | Args: 101 | depth (int): Network depth, from {50, 101, 152}. 102 | se_ratio (int): Squeeze ratio in SELayer. Default: 16. 103 | in_channels (int): Number of input image channels. Default: 3. 104 | stem_channels (int): Output channels of the stem layer. Default: 64. 105 | num_stages (int): Stages of the network. Default: 4. 106 | strides (Sequence[int]): Strides of the first block of each stage. 107 | Default: ``(1, 2, 2, 2)``. 108 | dilations (Sequence[int]): Dilation of each stage. 109 | Default: ``(1, 1, 1, 1)``. 110 | out_indices (Sequence[int]): Output from which stages. If only one 111 | stage is specified, a single tensor (feature map) is returned, 112 | otherwise multiple stages are specified, a tuple of tensors will 113 | be returned. Default: ``(3, )``. 114 | style (str): `pytorch` or `caffe`. If set to "pytorch", the stride-two 115 | layer is the 3x3 conv layer, otherwise the stride-two layer is 116 | the first 1x1 conv layer. 117 | deep_stem (bool): Replace 7x7 conv in input stem with 3 3x3 conv. 118 | Default: False. 119 | avg_down (bool): Use AvgPool instead of stride conv when 120 | downsampling in the bottleneck. Default: False. 121 | frozen_stages (int): Stages to be frozen (stop grad and set eval mode). 122 | -1 means not freezing any parameters. Default: -1. 123 | conv_cfg (dict | None): The config dict for conv layers. Default: None. 124 | norm_cfg (dict): The config dict for norm layers. 125 | norm_eval (bool): Whether to set norm layers to eval mode, namely, 126 | freeze running stats (mean and var). Note: Effect on Batch Norm 127 | and its variants only. Default: False. 128 | with_cp (bool): Use checkpoint or not. Using checkpoint will save some 129 | memory while slowing down the training speed. Default: False. 130 | zero_init_residual (bool): Whether to use zero init for last norm layer 131 | in resblocks to let them behave as identity. Default: True. 132 | 133 | Example: 134 | >>> from mmcls.models import SEResNet 135 | >>> import torch 136 | >>> self = SEResNet(depth=50) 137 | >>> self.eval() 138 | >>> inputs = torch.rand(1, 3, 224, 224) 139 | >>> level_outputs = self.forward(inputs) 140 | >>> for level_out in level_outputs: 141 | ... print(tuple(level_out.shape)) 142 | (1, 64, 56, 56) 143 | (1, 128, 28, 28) 144 | (1, 256, 14, 14) 145 | (1, 512, 7, 7) 146 | """ 147 | 148 | arch_settings = { 149 | 18: (SEBasicBlock, (2, 2, 2, 2)), 150 | 34: (SEBasicBlock, (3, 4, 6, 3)), 151 | 50: (SEBottleneck, (3, 4, 6, 3)), 152 | 101: (SEBottleneck, (3, 4, 23, 3)), 153 | 152: (SEBottleneck, (3, 8, 36, 3)) 154 | } 155 | 156 | def __init__(self, depth, se_ratio=16, **kwargs): 157 | if depth not in self.arch_settings: 158 | raise KeyError(f'invalid depth {depth} for SEResNet') 159 | self.se_ratio = se_ratio 160 | super(SEResNet, self).__init__(depth, **kwargs) 161 | 162 | def make_res_layer(self, **kwargs): 163 | return ResLayer(se_ratio=self.se_ratio, **kwargs) 164 | -------------------------------------------------------------------------------- /ufvl_net/models/builder.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) OpenMMLab. All rights reserved. 2 | from mmcv.cnn import MODELS as MMCV_MODELS 3 | from mmcv.cnn.bricks.registry import ATTENTION as MMCV_ATTENTION 4 | from mmcv.utils import Registry 5 | 6 | MODELS = Registry('models', parent=MMCV_MODELS) 7 | 8 | BACKBONES = MODELS 9 | HEADS = MODELS 10 | ARCHITECTURE = MODELS 11 | 12 | ATTENTION = Registry('attention', parent=MMCV_ATTENTION) 13 | 14 | 15 | def build_backbone(cfg): 16 | """Build backbone.""" 17 | return BACKBONES.build(cfg) 18 | 19 | 20 | def build_head(cfg): 21 | """Build head.""" 22 | return HEADS.build(cfg) 23 | 24 | 25 | def build_architecture(cfg): 26 | return ARCHITECTURE.build(cfg) 27 | -------------------------------------------------------------------------------- /ufvl_net/models/heads/__init__.py: -------------------------------------------------------------------------------- 1 | from .score_head import RegHead 2 | 3 | __all__ = [ 4 | 'RegHead' 5 | ] 6 | -------------------------------------------------------------------------------- /ufvl_net/models/heads/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/models/heads/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/models/heads/__pycache__/score_head.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/models/heads/__pycache__/score_head.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/models/heads/score_head.py: -------------------------------------------------------------------------------- 1 | import math 2 | import numbers 3 | 4 | import torch 5 | import torch.nn as nn 6 | import torch.nn.functional as F 7 | from einops import rearrange 8 | from mmcv.cnn import ConvModule 9 | from mmcv.runner import BaseModule 10 | 11 | from ..builder import HEADS 12 | 13 | 14 | @HEADS.register_module() 15 | class RegHead(BaseModule): 16 | def __init__(self, 17 | in_channel=2048, 18 | norm_cfg=dict(type='BN'), 19 | init_cfg=None): 20 | super(RegHead, self).__init__(init_cfg=init_cfg) 21 | 22 | self.in_channel = in_channel 23 | self.conv_reg1 = ConvModule(self.in_channel, 24 | 512, 25 | kernel_size=3, 26 | padding=1, 27 | norm_cfg=norm_cfg) 28 | self.conv_reg2 = ConvModule(512, 29 | 256, 30 | kernel_size=3, 31 | padding=1, 32 | norm_cfg=norm_cfg) 33 | self.conv_reg3 = ConvModule(256, 34 | 128, 35 | kernel_size=3, 36 | padding=1, 37 | norm_cfg=norm_cfg) 38 | self.coord_conv = ConvModule(128, 39 | 64, 40 | kernel_size=3, 41 | norm_cfg=norm_cfg, 42 | padding=1) 43 | 44 | self.coord_reg = torch.nn.Conv2d(64, 3, kernel_size=1) 45 | 46 | self.uncer_conv = ConvModule(128, 47 | 64, 48 | kernel_size=3, 49 | norm_cfg=norm_cfg, 50 | padding=1) 51 | 52 | self.uncer_reg = torch.nn.Conv2d(64, 1, kernel_size=1) 53 | 54 | 55 | def forward(self, feat, **kwargs): 56 | 57 | feat = self.conv_reg3(self.conv_reg2(self.conv_reg1(feat))) 58 | coord = self.coord_reg(self.coord_conv(feat)) 59 | uncer = self.uncer_reg(self.uncer_conv(feat)) 60 | uncer = torch.sigmoid(uncer) 61 | return coord, uncer -------------------------------------------------------------------------------- /ufvl_net/models/utils/__pycache__/make_divisible.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/models/utils/__pycache__/make_divisible.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/models/utils/__pycache__/se_layer.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/models/utils/__pycache__/se_layer.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/models/utils/augment/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) OpenMMLab. All rights reserved. 2 | from .augments import Augments 3 | from .cutmix import BatchCutMixLayer 4 | from .identity import Identity 5 | from .mixup import BatchMixupLayer 6 | from .resizemix import BatchResizeMixLayer 7 | 8 | __all__ = ('Augments', 'BatchCutMixLayer', 'Identity', 'BatchMixupLayer', 9 | 'BatchResizeMixLayer') 10 | -------------------------------------------------------------------------------- /ufvl_net/models/utils/augment/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/models/utils/augment/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/models/utils/augment/__pycache__/augments.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/models/utils/augment/__pycache__/augments.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/models/utils/augment/__pycache__/builder.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/models/utils/augment/__pycache__/builder.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/models/utils/augment/__pycache__/cutmix.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/models/utils/augment/__pycache__/cutmix.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/models/utils/augment/__pycache__/identity.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/models/utils/augment/__pycache__/identity.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/models/utils/augment/__pycache__/mixup.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/models/utils/augment/__pycache__/mixup.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/models/utils/augment/__pycache__/resizemix.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/models/utils/augment/__pycache__/resizemix.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/models/utils/augment/__pycache__/utils.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/models/utils/augment/__pycache__/utils.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/models/utils/augment/augments.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) OpenMMLab. All rights reserved. 2 | import random 3 | 4 | import numpy as np 5 | 6 | from .builder import build_augment 7 | 8 | 9 | class Augments(object): 10 | """Data augments. 11 | 12 | We implement some data augmentation methods, such as mixup, cutmix. 13 | 14 | Args: 15 | augments_cfg (list[`mmcv.ConfigDict`] | obj:`mmcv.ConfigDict`): 16 | Config dict of augments 17 | 18 | Example: 19 | >>> augments_cfg = [ 20 | dict(type='BatchCutMix', alpha=1., num_classes=10, prob=0.5), 21 | dict(type='BatchMixup', alpha=1., num_classes=10, prob=0.3) 22 | ] 23 | >>> augments = Augments(augments_cfg) 24 | >>> imgs = torch.randn(16, 3, 32, 32) 25 | >>> label = torch.randint(0, 10, (16, )) 26 | >>> imgs, label = augments(imgs, label) 27 | 28 | To decide which augmentation within Augments block is used 29 | the following rule is applied. 30 | We pick augmentation based on the probabilities. In the example above, 31 | we decide if we should use BatchCutMix with probability 0.5, 32 | BatchMixup 0.3. As Identity is not in augments_cfg, we use Identity with 33 | probability 1 - 0.5 - 0.3 = 0.2. 34 | """ 35 | 36 | def __init__(self, augments_cfg): 37 | super(Augments, self).__init__() 38 | 39 | if isinstance(augments_cfg, dict): 40 | augments_cfg = [augments_cfg] 41 | 42 | assert len(augments_cfg) > 0, \ 43 | 'The length of augments_cfg should be positive.' 44 | self.augments = [build_augment(cfg) for cfg in augments_cfg] 45 | self.augment_probs = [aug.prob for aug in self.augments] 46 | 47 | has_identity = any([cfg['type'] == 'Identity' for cfg in augments_cfg]) 48 | if has_identity: 49 | assert sum(self.augment_probs) == 1.0,\ 50 | 'The sum of augmentation probabilities should equal to 1,' \ 51 | ' but got {:.2f}'.format(sum(self.augment_probs)) 52 | else: 53 | assert sum(self.augment_probs) <= 1.0,\ 54 | 'The sum of augmentation probabilities should less than or ' \ 55 | 'equal to 1, but got {:.2f}'.format(sum(self.augment_probs)) 56 | identity_prob = 1 - sum(self.augment_probs) 57 | if identity_prob > 0: 58 | num_classes = self.augments[0].num_classes 59 | self.augments += [ 60 | build_augment( 61 | dict( 62 | type='Identity', 63 | num_classes=num_classes, 64 | prob=identity_prob)) 65 | ] 66 | self.augment_probs += [identity_prob] 67 | 68 | def __call__(self, img, gt_label): 69 | if self.augments: 70 | random_state = np.random.RandomState(random.randint(0, 2**32 - 1)) 71 | aug = random_state.choice(self.augments, p=self.augment_probs) 72 | return aug(img, gt_label) 73 | return img, gt_label 74 | -------------------------------------------------------------------------------- /ufvl_net/models/utils/augment/builder.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) OpenMMLab. All rights reserved. 2 | from mmcv.utils import Registry, build_from_cfg 3 | 4 | AUGMENT = Registry('augment') 5 | 6 | 7 | def build_augment(cfg, default_args=None): 8 | return build_from_cfg(cfg, AUGMENT, default_args) 9 | -------------------------------------------------------------------------------- /ufvl_net/models/utils/augment/cutmix.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) OpenMMLab. All rights reserved. 2 | from abc import ABCMeta, abstractmethod 3 | 4 | import numpy as np 5 | import torch 6 | 7 | from .builder import AUGMENT 8 | from .utils import one_hot_encoding 9 | 10 | 11 | class BaseCutMixLayer(object, metaclass=ABCMeta): 12 | """Base class for CutMixLayer. 13 | 14 | Args: 15 | alpha (float): Parameters for Beta distribution. Positive(>0) 16 | num_classes (int): The number of classes 17 | prob (float): MixUp probability. It should be in range [0, 1]. 18 | Default to 1.0 19 | cutmix_minmax (List[float], optional): cutmix min/max image ratio. 20 | (as percent of image size). When cutmix_minmax is not None, we 21 | generate cutmix bounding-box using cutmix_minmax instead of alpha 22 | correct_lam (bool): Whether to apply lambda correction when cutmix bbox 23 | clipped by image borders. Default to True 24 | """ 25 | 26 | def __init__(self, 27 | alpha, 28 | num_classes, 29 | prob=1.0, 30 | cutmix_minmax=None, 31 | correct_lam=True): 32 | super(BaseCutMixLayer, self).__init__() 33 | 34 | assert isinstance(alpha, float) and alpha > 0 35 | assert isinstance(num_classes, int) 36 | assert isinstance(prob, float) and 0.0 <= prob <= 1.0 37 | 38 | self.alpha = alpha 39 | self.num_classes = num_classes 40 | self.prob = prob 41 | self.cutmix_minmax = cutmix_minmax 42 | self.correct_lam = correct_lam 43 | 44 | def rand_bbox_minmax(self, img_shape, count=None): 45 | """Min-Max CutMix bounding-box Inspired by Darknet cutmix 46 | implementation. It generates a random rectangular bbox based on min/max 47 | percent values applied to each dimension of the input image. 48 | 49 | Typical defaults for minmax are usually in the .2-.3 for min and 50 | .8-.9 range for max. 51 | 52 | Args: 53 | img_shape (tuple): Image shape as tuple 54 | count (int, optional): Number of bbox to generate. Default to None 55 | """ 56 | assert len(self.cutmix_minmax) == 2 57 | img_h, img_w = img_shape[-2:] 58 | cut_h = np.random.randint( 59 | int(img_h * self.cutmix_minmax[0]), 60 | int(img_h * self.cutmix_minmax[1]), 61 | size=count) 62 | cut_w = np.random.randint( 63 | int(img_w * self.cutmix_minmax[0]), 64 | int(img_w * self.cutmix_minmax[1]), 65 | size=count) 66 | yl = np.random.randint(0, img_h - cut_h, size=count) 67 | xl = np.random.randint(0, img_w - cut_w, size=count) 68 | yu = yl + cut_h 69 | xu = xl + cut_w 70 | return yl, yu, xl, xu 71 | 72 | def rand_bbox(self, img_shape, lam, margin=0., count=None): 73 | """Standard CutMix bounding-box that generates a random square bbox 74 | based on lambda value. This implementation includes support for 75 | enforcing a border margin as percent of bbox dimensions. 76 | 77 | Args: 78 | img_shape (tuple): Image shape as tuple 79 | lam (float): Cutmix lambda value 80 | margin (float): Percentage of bbox dimension to enforce as margin 81 | (reduce amount of box outside image). Default to 0. 82 | count (int, optional): Number of bbox to generate. Default to None 83 | """ 84 | ratio = np.sqrt(1 - lam) 85 | img_h, img_w = img_shape[-2:] 86 | cut_h, cut_w = int(img_h * ratio), int(img_w * ratio) 87 | margin_y, margin_x = int(margin * cut_h), int(margin * cut_w) 88 | cy = np.random.randint(0 + margin_y, img_h - margin_y, size=count) 89 | cx = np.random.randint(0 + margin_x, img_w - margin_x, size=count) 90 | yl = np.clip(cy - cut_h // 2, 0, img_h) 91 | yh = np.clip(cy + cut_h // 2, 0, img_h) 92 | xl = np.clip(cx - cut_w // 2, 0, img_w) 93 | xh = np.clip(cx + cut_w // 2, 0, img_w) 94 | return yl, yh, xl, xh 95 | 96 | def cutmix_bbox_and_lam(self, img_shape, lam, count=None): 97 | """Generate bbox and apply lambda correction. 98 | 99 | Args: 100 | img_shape (tuple): Image shape as tuple 101 | lam (float): Cutmix lambda value 102 | count (int, optional): Number of bbox to generate. Default to None 103 | """ 104 | if self.cutmix_minmax is not None: 105 | yl, yu, xl, xu = self.rand_bbox_minmax(img_shape, count=count) 106 | else: 107 | yl, yu, xl, xu = self.rand_bbox(img_shape, lam, count=count) 108 | if self.correct_lam or self.cutmix_minmax is not None: 109 | bbox_area = (yu - yl) * (xu - xl) 110 | lam = 1. - bbox_area / float(img_shape[-2] * img_shape[-1]) 111 | return (yl, yu, xl, xu), lam 112 | 113 | @abstractmethod 114 | def cutmix(self, imgs, gt_label): 115 | pass 116 | 117 | 118 | @AUGMENT.register_module(name='BatchCutMix') 119 | class BatchCutMixLayer(BaseCutMixLayer): 120 | r"""CutMix layer for a batch of data. 121 | 122 | CutMix is a method to improve the network's generalization capability. It's 123 | proposed in `CutMix: Regularization Strategy to Train Strong Classifiers 124 | with Localizable Features ` 125 | 126 | With this method, patches are cut and pasted among training images where 127 | the ground truth labels are also mixed proportionally to the area of the 128 | patches. 129 | 130 | Args: 131 | alpha (float): Parameters for Beta distribution to generate the 132 | mixing ratio. It should be a positive number. More details 133 | can be found in :class:`BatchMixupLayer`. 134 | num_classes (int): The number of classes 135 | prob (float): The probability to execute cutmix. It should be in 136 | range [0, 1]. Defaults to 1.0. 137 | cutmix_minmax (List[float], optional): The min/max area ratio of the 138 | patches. If not None, the bounding-box of patches is uniform 139 | sampled within this ratio range, and the ``alpha`` will be ignored. 140 | Otherwise, the bounding-box is generated according to the 141 | ``alpha``. Defaults to None. 142 | correct_lam (bool): Whether to apply lambda correction when cutmix bbox 143 | clipped by image borders. Defaults to True. 144 | 145 | Note: 146 | If the ``cutmix_minmax`` is None, how to generate the bounding-box of 147 | patches according to the ``alpha``? 148 | 149 | First, generate a :math:`\lambda`, details can be found in 150 | :class:`BatchMixupLayer`. And then, the area ratio of the bounding-box 151 | is calculated by: 152 | 153 | .. math:: 154 | \text{ratio} = \sqrt{1-\lambda} 155 | """ 156 | 157 | def __init__(self, *args, **kwargs): 158 | super(BatchCutMixLayer, self).__init__(*args, **kwargs) 159 | 160 | def cutmix(self, img, gt_label): 161 | one_hot_gt_label = one_hot_encoding(gt_label, self.num_classes) 162 | lam = np.random.beta(self.alpha, self.alpha) 163 | batch_size = img.size(0) 164 | index = torch.randperm(batch_size) 165 | 166 | (bby1, bby2, bbx1, 167 | bbx2), lam = self.cutmix_bbox_and_lam(img.shape, lam) 168 | img[:, :, bby1:bby2, bbx1:bbx2] = \ 169 | img[index, :, bby1:bby2, bbx1:bbx2] 170 | mixed_gt_label = lam * one_hot_gt_label + ( 171 | 1 - lam) * one_hot_gt_label[index, :] 172 | return img, mixed_gt_label 173 | 174 | def __call__(self, img, gt_label): 175 | return self.cutmix(img, gt_label) 176 | -------------------------------------------------------------------------------- /ufvl_net/models/utils/augment/identity.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) OpenMMLab. All rights reserved. 2 | from .builder import AUGMENT 3 | from .utils import one_hot_encoding 4 | 5 | 6 | @AUGMENT.register_module(name='Identity') 7 | class Identity(object): 8 | """Change gt_label to one_hot encoding and keep img as the same. 9 | 10 | Args: 11 | num_classes (int): The number of classes. 12 | prob (float): MixUp probability. It should be in range [0, 1]. 13 | Default to 1.0 14 | """ 15 | 16 | def __init__(self, num_classes, prob=1.0): 17 | super(Identity, self).__init__() 18 | 19 | assert isinstance(num_classes, int) 20 | assert isinstance(prob, float) and 0.0 <= prob <= 1.0 21 | 22 | self.num_classes = num_classes 23 | self.prob = prob 24 | 25 | def one_hot(self, gt_label): 26 | return one_hot_encoding(gt_label, self.num_classes) 27 | 28 | def __call__(self, img, gt_label): 29 | return img, self.one_hot(gt_label) 30 | -------------------------------------------------------------------------------- /ufvl_net/models/utils/augment/mixup.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) OpenMMLab. All rights reserved. 2 | from abc import ABCMeta, abstractmethod 3 | 4 | import numpy as np 5 | import torch 6 | 7 | from .builder import AUGMENT 8 | from .utils import one_hot_encoding 9 | 10 | 11 | class BaseMixupLayer(object, metaclass=ABCMeta): 12 | """Base class for MixupLayer. 13 | 14 | Args: 15 | alpha (float): Parameters for Beta distribution to generate the 16 | mixing ratio. It should be a positive number. 17 | num_classes (int): The number of classes. 18 | prob (float): MixUp probability. It should be in range [0, 1]. 19 | Default to 1.0 20 | """ 21 | 22 | def __init__(self, alpha, num_classes, prob=1.0): 23 | super(BaseMixupLayer, self).__init__() 24 | 25 | assert isinstance(alpha, float) and alpha > 0 26 | assert isinstance(num_classes, int) 27 | assert isinstance(prob, float) and 0.0 <= prob <= 1.0 28 | 29 | self.alpha = alpha 30 | self.num_classes = num_classes 31 | self.prob = prob 32 | 33 | @abstractmethod 34 | def mixup(self, imgs, gt_label): 35 | pass 36 | 37 | 38 | @AUGMENT.register_module(name='BatchMixup') 39 | class BatchMixupLayer(BaseMixupLayer): 40 | r"""Mixup layer for a batch of data. 41 | 42 | Mixup is a method to reduces the memorization of corrupt labels and 43 | increases the robustness to adversarial examples. It's 44 | proposed in `mixup: Beyond Empirical Risk Minimization 45 | ` 46 | 47 | This method simply linearly mix pairs of data and their labels. 48 | 49 | Args: 50 | alpha (float): Parameters for Beta distribution to generate the 51 | mixing ratio. It should be a positive number. More details 52 | are in the note. 53 | num_classes (int): The number of classes. 54 | prob (float): The probability to execute mixup. It should be in 55 | range [0, 1]. Default sto 1.0. 56 | 57 | Note: 58 | The :math:`\alpha` (``alpha``) determines a random distribution 59 | :math:`Beta(\alpha, \alpha)`. For each batch of data, we sample 60 | a mixing ratio (marked as :math:`\lambda`, ``lam``) from the random 61 | distribution. 62 | """ 63 | 64 | def __init__(self, *args, **kwargs): 65 | super(BatchMixupLayer, self).__init__(*args, **kwargs) 66 | 67 | def mixup(self, img, gt_label): 68 | one_hot_gt_label = one_hot_encoding(gt_label, self.num_classes) 69 | lam = np.random.beta(self.alpha, self.alpha) 70 | batch_size = img.size(0) 71 | index = torch.randperm(batch_size) 72 | 73 | mixed_img = lam * img + (1 - lam) * img[index, :] 74 | mixed_gt_label = lam * one_hot_gt_label + ( 75 | 1 - lam) * one_hot_gt_label[index, :] 76 | 77 | return mixed_img, mixed_gt_label 78 | 79 | def __call__(self, img, gt_label): 80 | return self.mixup(img, gt_label) 81 | -------------------------------------------------------------------------------- /ufvl_net/models/utils/augment/resizemix.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) OpenMMLab. All rights reserved. 2 | import numpy as np 3 | import torch 4 | import torch.nn.functional as F 5 | 6 | from ufvl_net.models.utils.augment.builder import AUGMENT 7 | from .cutmix import BatchCutMixLayer 8 | from .utils import one_hot_encoding 9 | 10 | 11 | @AUGMENT.register_module(name='BatchResizeMix') 12 | class BatchResizeMixLayer(BatchCutMixLayer): 13 | r"""ResizeMix Random Paste layer for a batch of data. 14 | 15 | The ResizeMix will resize an image to a small patch and paste it on another 16 | image. It's proposed in `ResizeMix: Mixing Data with Preserved Object 17 | Information and True Labels `_ 18 | 19 | Args: 20 | alpha (float): Parameters for Beta distribution to generate the 21 | mixing ratio. It should be a positive number. More details 22 | can be found in :class:`BatchMixupLayer`. 23 | num_classes (int): The number of classes. 24 | lam_min(float): The minimum value of lam. Defaults to 0.1. 25 | lam_max(float): The maximum value of lam. Defaults to 0.8. 26 | interpolation (str): algorithm used for upsampling: 27 | 'nearest' | 'linear' | 'bilinear' | 'bicubic' | 'trilinear' | 28 | 'area'. Default to 'bilinear'. 29 | prob (float): The probability to execute resizemix. It should be in 30 | range [0, 1]. Defaults to 1.0. 31 | cutmix_minmax (List[float], optional): The min/max area ratio of the 32 | patches. If not None, the bounding-box of patches is uniform 33 | sampled within this ratio range, and the ``alpha`` will be ignored. 34 | Otherwise, the bounding-box is generated according to the 35 | ``alpha``. Defaults to None. 36 | correct_lam (bool): Whether to apply lambda correction when cutmix bbox 37 | clipped by image borders. Defaults to True 38 | **kwargs: Any other parameters accpeted by :class:`BatchCutMixLayer`. 39 | 40 | Note: 41 | The :math:`\lambda` (``lam``) is the mixing ratio. It's a random 42 | variable which follows :math:`Beta(\alpha, \alpha)` and is mapped 43 | to the range [``lam_min``, ``lam_max``]. 44 | 45 | .. math:: 46 | \lambda = \frac{Beta(\alpha, \alpha)} 47 | {\lambda_{max} - \lambda_{min}} + \lambda_{min} 48 | 49 | And the resize ratio of source images is calculated by :math:`\lambda`: 50 | 51 | .. math:: 52 | \text{ratio} = \sqrt{1-\lambda} 53 | """ 54 | 55 | def __init__(self, 56 | alpha, 57 | num_classes, 58 | lam_min: float = 0.1, 59 | lam_max: float = 0.8, 60 | interpolation='bilinear', 61 | prob=1.0, 62 | cutmix_minmax=None, 63 | correct_lam=True, 64 | **kwargs): 65 | super(BatchResizeMixLayer, self).__init__( 66 | alpha=alpha, 67 | num_classes=num_classes, 68 | prob=prob, 69 | cutmix_minmax=cutmix_minmax, 70 | correct_lam=correct_lam, 71 | **kwargs) 72 | self.lam_min = lam_min 73 | self.lam_max = lam_max 74 | self.interpolation = interpolation 75 | 76 | def cutmix(self, img, gt_label): 77 | one_hot_gt_label = one_hot_encoding(gt_label, self.num_classes) 78 | 79 | lam = np.random.beta(self.alpha, self.alpha) 80 | lam = lam * (self.lam_max - self.lam_min) + self.lam_min 81 | batch_size = img.size(0) 82 | index = torch.randperm(batch_size) 83 | 84 | (bby1, bby2, bbx1, 85 | bbx2), lam = self.cutmix_bbox_and_lam(img.shape, lam) 86 | 87 | img[:, :, bby1:bby2, bbx1:bbx2] = F.interpolate( 88 | img[index], 89 | size=(bby2 - bby1, bbx2 - bbx1), 90 | mode=self.interpolation) 91 | mixed_gt_label = lam * one_hot_gt_label + ( 92 | 1 - lam) * one_hot_gt_label[index, :] 93 | return img, mixed_gt_label 94 | -------------------------------------------------------------------------------- /ufvl_net/models/utils/augment/utils.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) OpenMMLab. All rights reserved. 2 | import torch.nn.functional as F 3 | 4 | 5 | def one_hot_encoding(gt, num_classes): 6 | """Change gt_label to one_hot encoding. 7 | 8 | If the shape has 2 or more 9 | dimensions, return it without encoding. 10 | Args: 11 | gt (Tensor): The gt label with shape (N,) or shape (N, */). 12 | num_classes (int): The number of classes. 13 | Return: 14 | Tensor: One hot gt label. 15 | """ 16 | if gt.ndim == 1: 17 | # multi-class classification 18 | return F.one_hot(gt, num_classes=num_classes) 19 | else: 20 | # binary classification 21 | # example. [[0], [1], [1]] 22 | # multi-label classification 23 | # example. [[0, 1, 1], [1, 0, 0], [1, 1, 1]] 24 | return gt 25 | -------------------------------------------------------------------------------- /ufvl_net/models/utils/make_divisible.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) OpenMMLab. All rights reserved. 2 | def make_divisible(value, divisor, min_value=None, min_ratio=0.9): 3 | """Make divisible function. 4 | 5 | This function rounds the channel number down to the nearest value that can 6 | be divisible by the divisor. 7 | 8 | Args: 9 | value (int): The original channel number. 10 | divisor (int): The divisor to fully divide the channel number. 11 | min_value (int, optional): The minimum value of the output channel. 12 | Default: None, means that the minimum value equal to the divisor. 13 | min_ratio (float): The minimum ratio of the rounded channel 14 | number to the original channel number. Default: 0.9. 15 | Returns: 16 | int: The modified output channel number 17 | """ 18 | 19 | if min_value is None: 20 | min_value = divisor 21 | new_value = max(min_value, int(value + divisor / 2) // divisor * divisor) 22 | # Make sure that round down does not go down by more than (1-min_ratio). 23 | if new_value < min_ratio * value: 24 | new_value += divisor 25 | return new_value 26 | -------------------------------------------------------------------------------- /ufvl_net/models/utils/se_layer.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) OpenMMLab. All rights reserved. 2 | import mmcv 3 | import torch.nn as nn 4 | from mmcv.cnn import ConvModule 5 | from mmcv.runner import BaseModule 6 | 7 | from .make_divisible import make_divisible 8 | 9 | 10 | class SELayer(BaseModule): 11 | """Squeeze-and-Excitation Module. 12 | 13 | Args: 14 | channels (int): The input (and output) channels of the SE layer. 15 | squeeze_channels (None or int): The intermediate channel number of 16 | SElayer. Default: None, means the value of ``squeeze_channels`` 17 | is ``make_divisible(channels // ratio, divisor)``. 18 | ratio (int): Squeeze ratio in SELayer, the intermediate channel will 19 | be ``make_divisible(channels // ratio, divisor)``. Only used when 20 | ``squeeze_channels`` is None. Default: 16. 21 | divisor(int): The divisor to true divide the channel number. Only 22 | used when ``squeeze_channels`` is None. Default: 8. 23 | conv_cfg (None or dict): Config dict for convolution layer. Default: 24 | None, which means using conv2d. 25 | return_weight(bool): Whether to return the weight. Default: False. 26 | act_cfg (dict or Sequence[dict]): Config dict for activation layer. 27 | If act_cfg is a dict, two activation layers will be configurated 28 | by this dict. If act_cfg is a sequence of dicts, the first 29 | activation layer will be configurated by the first dict and the 30 | second activation layer will be configurated by the second dict. 31 | Default: (dict(type='ReLU'), dict(type='Sigmoid')) 32 | """ 33 | 34 | def __init__(self, 35 | channels, 36 | squeeze_channels=None, 37 | ratio=16, 38 | divisor=8, 39 | bias='auto', 40 | conv_cfg=None, 41 | act_cfg=(dict(type='ReLU'), dict(type='Sigmoid')), 42 | return_weight=False, 43 | init_cfg=None): 44 | super(SELayer, self).__init__(init_cfg) 45 | if isinstance(act_cfg, dict): 46 | act_cfg = (act_cfg, act_cfg) 47 | assert len(act_cfg) == 2 48 | assert mmcv.is_tuple_of(act_cfg, dict) 49 | self.global_avgpool = nn.AdaptiveAvgPool2d(1) 50 | if squeeze_channels is None: 51 | squeeze_channels = make_divisible(channels // ratio, divisor) 52 | assert isinstance(squeeze_channels, int) and squeeze_channels > 0, \ 53 | '"squeeze_channels" should be a positive integer, but get ' + \ 54 | f'{squeeze_channels} instead.' 55 | self.return_weight = return_weight 56 | self.conv1 = ConvModule( 57 | in_channels=channels, 58 | out_channels=squeeze_channels, 59 | kernel_size=1, 60 | stride=1, 61 | bias=bias, 62 | conv_cfg=conv_cfg, 63 | act_cfg=act_cfg[0]) 64 | self.conv2 = ConvModule( 65 | in_channels=squeeze_channels, 66 | out_channels=channels, 67 | kernel_size=1, 68 | stride=1, 69 | bias=bias, 70 | conv_cfg=conv_cfg, 71 | act_cfg=act_cfg[1]) 72 | 73 | def forward(self, x): 74 | out = self.global_avgpool(x) 75 | out = self.conv1(out) 76 | out = self.conv2(out) 77 | if self.return_weight: 78 | return out 79 | else: 80 | return x * out 81 | -------------------------------------------------------------------------------- /ufvl_net/utils/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) OpenMMLab. All rights reserved. 2 | from .collect_env import collect_env 3 | from .logger import get_root_logger, load_json_log 4 | from .setup_env import setup_multi_processes 5 | 6 | __all__ = [ 7 | 'collect_env', 'get_root_logger', 'load_json_log', 'setup_multi_processes' 8 | ] 9 | -------------------------------------------------------------------------------- /ufvl_net/utils/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/utils/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/utils/__pycache__/collect_env.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/utils/__pycache__/collect_env.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/utils/__pycache__/logger.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/utils/__pycache__/logger.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/utils/__pycache__/setup_env.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mooncake199809/UFVL-Net/b0976db816c6bf61797603149b9fc96871dd8d9e/ufvl_net/utils/__pycache__/setup_env.cpython-37.pyc -------------------------------------------------------------------------------- /ufvl_net/utils/collect_env.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) OpenMMLab. All rights reserved. 2 | from mmcv.utils import collect_env as collect_base_env 3 | from mmcv.utils import get_git_hash 4 | 5 | import ufvl_net 6 | 7 | 8 | def collect_env(): 9 | """Collect the information of the running environments.""" 10 | env_info = collect_base_env() 11 | env_info['MMClassification'] = mmcls.__version__ + '+' + get_git_hash()[:7] 12 | return env_info 13 | 14 | 15 | if __name__ == '__main__': 16 | for name, val in collect_env().items(): 17 | print(f'{name}: {val}') 18 | -------------------------------------------------------------------------------- /ufvl_net/utils/logger.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) OpenMMLab. All rights reserved. 2 | import json 3 | import logging 4 | from collections import defaultdict 5 | 6 | from mmcv.utils import get_logger 7 | 8 | 9 | def get_root_logger(log_file=None, log_level=logging.INFO): 10 | """Get root logger. 11 | 12 | Args: 13 | log_file (str, optional): File path of log. Defaults to None. 14 | log_level (int, optional): The level of logger. 15 | Defaults to :obj:`logging.INFO`. 16 | 17 | Returns: 18 | :obj:`logging.Logger`: The obtained logger 19 | """ 20 | return get_logger('mmcls', log_file, log_level) 21 | 22 | 23 | def load_json_log(json_log): 24 | """load and convert json_logs to log_dicts. 25 | 26 | Args: 27 | json_log (str): The path of the json log file. 28 | 29 | Returns: 30 | dict[int, dict[str, list]]: 31 | Key is the epoch, value is a sub dict. The keys in each sub dict 32 | are different metrics, e.g. memory, bbox_mAP, and the value is a 33 | list of corresponding values in all iterations in this epoch. 34 | 35 | .. code-block:: python 36 | 37 | # An example output 38 | { 39 | 1: {'iter': [100, 200, 300], 'loss': [6.94, 6.73, 6.53]}, 40 | 2: {'iter': [100, 200, 300], 'loss': [6.33, 6.20, 6.07]}, 41 | ... 42 | } 43 | """ 44 | log_dict = dict() 45 | with open(json_log, 'r') as log_file: 46 | for line in log_file: 47 | log = json.loads(line.strip()) 48 | # skip lines without `epoch` field 49 | if 'epoch' not in log: 50 | continue 51 | epoch = log.pop('epoch') 52 | if epoch not in log_dict: 53 | log_dict[epoch] = defaultdict(list) 54 | for k, v in log.items(): 55 | log_dict[epoch][k].append(v) 56 | return log_dict 57 | -------------------------------------------------------------------------------- /ufvl_net/utils/setup_env.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) OpenMMLab. All rights reserved. 2 | import os 3 | import platform 4 | import warnings 5 | 6 | import cv2 7 | import torch.multiprocessing as mp 8 | 9 | 10 | def setup_multi_processes(cfg): 11 | """Setup multi-processing environment variables.""" 12 | # set multi-process start method as `fork` to speed up the training 13 | if platform.system() != 'Windows': 14 | mp_start_method = cfg.get('mp_start_method', 'fork') 15 | current_method = mp.get_start_method(allow_none=True) 16 | if current_method is not None and current_method != mp_start_method: 17 | warnings.warn( 18 | f'Multi-processing start method `{mp_start_method}` is ' 19 | f'different from the previous setting `{current_method}`.' 20 | f'It will be force set to `{mp_start_method}`. You can change ' 21 | f'this behavior by changing `mp_start_method` in your config.') 22 | mp.set_start_method(mp_start_method, force=True) 23 | 24 | # disable opencv multithreading to avoid system being overloaded 25 | opencv_num_threads = cfg.get('opencv_num_threads', 0) 26 | cv2.setNumThreads(opencv_num_threads) 27 | 28 | # setup OMP threads 29 | # This code is referred from https://github.com/pytorch/pytorch/blob/master/torch/distributed/run.py # noqa 30 | if 'OMP_NUM_THREADS' not in os.environ and cfg.data.workers_per_gpu > 1: 31 | omp_num_threads = 1 32 | warnings.warn( 33 | f'Setting OMP_NUM_THREADS environment variable for each process ' 34 | f'to be {omp_num_threads} in default, to avoid your system being ' 35 | f'overloaded, please further tune the variable for optimal ' 36 | f'performance in your application as needed.') 37 | os.environ['OMP_NUM_THREADS'] = str(omp_num_threads) 38 | 39 | # setup MKL threads 40 | if 'MKL_NUM_THREADS' not in os.environ and cfg.data.workers_per_gpu > 1: 41 | mkl_num_threads = 1 42 | warnings.warn( 43 | f'Setting MKL_NUM_THREADS environment variable for each process ' 44 | f'to be {mkl_num_threads} in default, to avoid your system being ' 45 | f'overloaded, please further tune the variable for optimal ' 46 | f'performance in your application as needed.') 47 | os.environ['MKL_NUM_THREADS'] = str(mkl_num_threads) 48 | -------------------------------------------------------------------------------- /ufvl_net/version.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) OpenMMLab. All rights reserved 2 | 3 | __version__ = '0.1.0' 4 | 5 | 6 | def parse_version_info(version_str): 7 | """Parse a version string into a tuple. 8 | 9 | Args: 10 | version_str (str): The version string. 11 | Returns: 12 | tuple[int | str]: The version info, e.g., "1.3.0" is parsed into 13 | (1, 3, 0), and "2.0.0rc1" is parsed into (2, 0, 0, 'rc1'). 14 | """ 15 | version_info = [] 16 | for x in version_str.split('.'): 17 | if x.isdigit(): 18 | version_info.append(int(x)) 19 | elif x.find('rc') != -1: 20 | patch_version = x.split('rc') 21 | version_info.append(int(patch_version[0])) 22 | version_info.append(f'rc{patch_version[1]}') 23 | return tuple(version_info) 24 | 25 | 26 | version_info = parse_version_info(__version__) 27 | 28 | __all__ = ['__version__', 'version_info', 'parse_version_info'] 29 | --------------------------------------------------------------------------------