├── detection ├── gvcore │ ├── __init__.py │ ├── utils │ │ ├── __init__.py │ │ ├── types.py │ │ ├── registry.py │ │ ├── checkpoint.py │ │ ├── misc.py │ │ ├── lr_scheduler.py │ │ ├── structure.py │ │ └── config.py │ ├── evaluator │ │ └── __init__.py │ ├── model │ │ └── __init__.py │ └── dataset │ │ ├── build.py │ │ ├── __init__.py │ │ ├── transforms │ │ └── __init__.py │ │ ├── dataloader.py │ │ └── sampler.py ├── det │ ├── model │ │ ├── __init__.py │ │ ├── block │ │ │ └── __init__.py │ │ ├── head │ │ │ ├── __init__.py │ │ │ ├── unbiased_teacher │ │ │ │ ├── __init__.py │ │ │ │ ├── rpn.py │ │ │ │ └── roihead.py │ │ │ ├── virtual_category │ │ │ │ └── __init__.py │ │ │ └── fasterrcnn.py │ │ ├── loss │ │ │ ├── __init__.py │ │ │ └── functional.py │ │ ├── detector │ │ │ ├── __init__.py │ │ │ ├── unbiased_teacher.py │ │ │ ├── fasterrcnn.py │ │ │ ├── ts.py │ │ │ └── ts_cross.py │ │ ├── anchor │ │ │ └── __init__.py │ │ ├── backbone │ │ │ └── __init__.py │ │ └── utils.py │ ├── utils │ │ ├── __init__.py │ │ ├── dataset_utils │ │ │ ├── __init__.py │ │ │ ├── append_aspect_ratio.py │ │ │ └── split.py │ │ ├── box.py │ │ ├── lr_scheduler.py │ │ ├── metrics.py │ │ └── visualizer.py │ ├── evaluator │ │ ├── __init__.py │ │ └── voc.py │ ├── dataset │ │ ├── __init__.py │ │ ├── voc.py │ │ ├── coco.py │ │ ├── det_generic.py │ │ └── transforms │ │ │ └── functional.py │ ├── ext │ │ ├── fastevalapi │ │ │ ├── __init__.py │ │ │ └── csrc │ │ │ │ ├── cuda_version.cu │ │ │ │ ├── vision.cpp │ │ │ │ └── cocoeval │ │ │ │ └── cocoeval.h │ │ ├── .clang-format │ │ └── setup.py │ ├── operators │ │ ├── __init__.py │ │ ├── virtual_category_cross.py │ │ ├── virtual_category.py │ │ └── fasterrcnn.py │ ├── config │ │ ├── fasterrcnn │ │ │ ├── base.yaml │ │ │ ├── coco.yaml │ │ │ └── voc.yaml │ │ ├── unbiased_teacher │ │ │ ├── coco.yaml │ │ │ ├── base.yaml │ │ │ └── voc.yaml │ │ ├── virtual_category_cross │ │ │ ├── coco.yaml │ │ │ ├── base.yaml │ │ │ ├── voc.yaml │ │ │ └── voc_cocox.yaml │ │ └── virtual_category │ │ │ ├── coco.yaml │ │ │ ├── base.yaml │ │ │ └── voc.yaml │ └── run.py ├── setup.py └── readme.md ├── segmentation ├── gvcore │ ├── __init__.py │ ├── utils │ │ ├── __init__.py │ │ ├── types.py │ │ ├── setup.py │ │ ├── registry.py │ │ ├── checkpoint.py │ │ ├── misc.py │ │ ├── sweeper.py │ │ └── structure.py │ ├── dataset │ │ ├── transforms │ │ │ └── __init__.py │ │ ├── __init__.py │ │ ├── build.py │ │ ├── dataloader.py │ │ └── sampler.py │ ├── evaluator │ │ └── __init__.py │ ├── model │ │ └── __init__.py │ └── optim │ │ └── utils.py ├── seg │ ├── model │ │ ├── __init__.py │ │ ├── head │ │ │ ├── __init__.py │ │ │ ├── aspp_alt.py │ │ │ └── aspp.py │ │ ├── loss │ │ │ ├── __init__.py │ │ │ └── functional.py │ │ ├── segmentor │ │ │ ├── __init__.py │ │ │ ├── deeplabv3p_alt.py │ │ │ └── deeplabv3p.py │ │ ├── backbone │ │ │ └── __init__.py │ │ └── block │ │ │ ├── __init__.py │ │ │ └── block.py │ ├── utils │ │ ├── __init__.py │ │ ├── data_utils │ │ │ ├── voc_split.py │ │ │ └── cityscape_split.py │ │ └── visualizer.py │ ├── evaluator │ │ ├── __init__.py │ │ └── api.py │ ├── dataset │ │ ├── __init__.py │ │ ├── voc.py │ │ ├── cityscapes.py │ │ └── transforms │ │ │ └── __init__.py │ ├── operators │ │ ├── __init__.py │ │ ├── deeplabv3p.py │ │ └── vc.py │ ├── run.py │ └── config │ │ ├── deeplabv3p │ │ ├── base.yaml │ │ └── voc_aug.yaml │ │ ├── ts_deeplabv3p │ │ ├── voc_aug.yaml │ │ ├── base.yaml │ │ └── cityscapes.yaml │ │ └── vc_deeplabv3p │ │ ├── voc_aug.yaml │ │ ├── base.yaml │ │ └── cityscapes.yaml ├── setup.py └── readme.md ├── .gitignore └── readme.md /detection/gvcore/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /detection/det/model/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /detection/det/utils/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /segmentation/gvcore/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /detection/det/model/block/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /detection/det/model/head/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /detection/det/model/loss/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /detection/gvcore/utils/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /segmentation/gvcore/utils/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /segmentation/seg/model/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /segmentation/seg/utils/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /detection/det/model/detector/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /segmentation/seg/model/head/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /segmentation/seg/model/loss/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /detection/det/utils/dataset_utils/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /segmentation/seg/model/segmentor/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /detection/det/model/head/unbiased_teacher/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /detection/det/model/head/virtual_category/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /segmentation/seg/evaluator/__init__.py: -------------------------------------------------------------------------------- 1 | from .segmentation import * 2 | -------------------------------------------------------------------------------- /detection/det/evaluator/__init__.py: -------------------------------------------------------------------------------- 1 | from .coco import * 2 | from .voc import * 3 | -------------------------------------------------------------------------------- /detection/det/dataset/__init__.py: -------------------------------------------------------------------------------- 1 | from .transforms import * 2 | from .coco import * 3 | from .voc import * 4 | -------------------------------------------------------------------------------- /detection/det/model/anchor/__init__.py: -------------------------------------------------------------------------------- 1 | from .generator import * 2 | from .macher import * 3 | from .coder import * 4 | -------------------------------------------------------------------------------- /segmentation/gvcore/dataset/transforms/__init__.py: -------------------------------------------------------------------------------- 1 | from .transforms import * 2 | from .kornia_transforms import * 3 | -------------------------------------------------------------------------------- /segmentation/seg/dataset/__init__.py: -------------------------------------------------------------------------------- 1 | from .transforms import * 2 | from .voc import * 3 | from .cityscapes import * 4 | -------------------------------------------------------------------------------- /detection/gvcore/evaluator/__init__.py: -------------------------------------------------------------------------------- 1 | from gvcore.utils.registry import Registry 2 | 3 | EVALUATOR_REGISTRY = Registry() 4 | -------------------------------------------------------------------------------- /segmentation/gvcore/evaluator/__init__.py: -------------------------------------------------------------------------------- 1 | from gvcore.utils.registry import Registry 2 | 3 | EVALUATOR_REGISTRY = Registry() 4 | -------------------------------------------------------------------------------- /detection/det/ext/fastevalapi/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. 2 | import torch 3 | 4 | __version__ = "0.3" 5 | -------------------------------------------------------------------------------- /segmentation/seg/operators/__init__.py: -------------------------------------------------------------------------------- 1 | from .deeplabv3p import DeeplabV3pOpt 2 | from .ts import TeacherStudentOpt 3 | from .vc import VirtualCategoryOpt 4 | -------------------------------------------------------------------------------- /detection/det/operators/__init__.py: -------------------------------------------------------------------------------- 1 | from .fasterrcnn import FasterRCNNOpt 2 | from .unbiased_teacher import UnbiasedTeacherOpt 3 | from .virtual_category import VirtualCategoryOpt 4 | from .virtual_category_cross import VirtualCategoryCrossOpt -------------------------------------------------------------------------------- /segmentation/seg/model/backbone/__init__.py: -------------------------------------------------------------------------------- 1 | from model.backbone.resnet import ResNet 2 | 3 | backbones = {"resnet": ResNet} 4 | 5 | 6 | def make_backbone(cfg): 7 | assert cfg.type in backbones, "Unknown backbone type {}".format(cfg.type) 8 | 9 | return backbones[cfg.type](**cfg) 10 | -------------------------------------------------------------------------------- /detection/gvcore/utils/types.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from typing import Dict, List, Tuple 3 | 4 | from gvcore.utils.structure import GenericData 5 | 6 | 7 | TTensorList = List[torch.Tensor] 8 | TTensorDict = Dict[str, torch.Tensor] 9 | TTensorTuple = Tuple[torch.Tensor] 10 | TTensor = torch.Tensor 11 | TDataList = List[GenericData] 12 | TData = GenericData 13 | -------------------------------------------------------------------------------- /segmentation/gvcore/utils/types.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from typing import Dict, List, Tuple 3 | 4 | from gvcore.utils.structure import GenericData 5 | 6 | 7 | TTensorList = List[torch.Tensor] 8 | TTensorDict = Dict[str, torch.Tensor] 9 | TTensorTuple = Tuple[torch.Tensor] 10 | TTensor = torch.Tensor 11 | TDataList = List[GenericData] 12 | TData = GenericData 13 | -------------------------------------------------------------------------------- /detection/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | 3 | setup( 4 | name="gvcore", 5 | version="1.0", 6 | description="The core library for vision framework.", 7 | url="", 8 | author="geo", 9 | author_email="", 10 | license="MIT", 11 | packages=["gvcore"], 12 | zip_safe=False, 13 | install_requires=["torch", "torchvision", "numpy", "tabulate",], 14 | ) 15 | -------------------------------------------------------------------------------- /detection/det/dataset/voc.py: -------------------------------------------------------------------------------- 1 | from gvcore.dataset import DATASET_REGISTRY 2 | 3 | from dataset.det_generic import DetGenericDataset 4 | from utils.meta import vocmeta 5 | 6 | 7 | @DATASET_REGISTRY.register("VOC") 8 | class VOCDataset(DetGenericDataset): 9 | def __init__(self, cache, path, img_transforms, subset): 10 | super(VOCDataset, self).__init__(cache, path, img_transforms, vocmeta, subset) 11 | -------------------------------------------------------------------------------- /detection/det/dataset/coco.py: -------------------------------------------------------------------------------- 1 | from gvcore.dataset import DATASET_REGISTRY 2 | 3 | from dataset.det_generic import DetGenericDataset 4 | from utils.meta import cocometa 5 | 6 | 7 | @DATASET_REGISTRY.register("COCO") 8 | class COCODataset(DetGenericDataset): 9 | def __init__(self, cache, path, img_transforms, subset): 10 | super(COCODataset, self).__init__(cache, path, img_transforms, cocometa, subset) 11 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | detection/data/* 2 | detection/cache/* 3 | **/.vscode 4 | **/.idea 5 | **/*.jpg 6 | **/*.png 7 | **/*.pdf 8 | **/*.pkl 9 | **/*.pt 10 | **/*.pth 11 | **/.DS_Store 12 | **/*.zip 13 | **/*.pyc 14 | **/*.list 15 | **/*.npy 16 | **/__pycache__ 17 | dev_utils/* 18 | **/*.egg-info 19 | ext/build 20 | ext/dist 21 | log 22 | slurmlog 23 | wandb 24 | .dist_init 25 | .aim 26 | **/*.dict 27 | **/*.egg-info 28 | **/build/* 29 | **/dist/* -------------------------------------------------------------------------------- /segmentation/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | 3 | setup( 4 | name="gvcore", 5 | version="1.1", 6 | description="The core library for Geo's vision framework.", 7 | author="Geoffrey Chen", 8 | author_email="geoffreychen777@gmail.com", 9 | license="MIT", 10 | packages=["gvcore"], 11 | zip_safe=False, 12 | install_requires=["torch", "torchvision", "numpy", "tabulate",], 13 | entry_points={"console_scripts": ["gvrun=gvcore.utils.launcher:gvrun", "gvsubmit=gvcore.utils.launcher:gvsubmit", "gvsweep=gvcore.utils.sweeper:gvsweep"]}, 14 | ) 15 | -------------------------------------------------------------------------------- /segmentation/seg/model/block/__init__.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | import gvcore.utils.distributed as dist_utils 4 | 5 | 6 | def Normalize(type: str): 7 | if type == "BN": 8 | return nn.BatchNorm2d 9 | elif type == "GN": 10 | return nn.GroupNorm 11 | elif type == "LN": 12 | return nn.LayerNorm 13 | elif type == "SyncBN": 14 | if dist_utils.get_world_size() > 1: 15 | return nn.SyncBatchNorm 16 | else: 17 | return nn.BatchNorm2d 18 | elif type == "none": 19 | return None 20 | else: 21 | raise ValueError("Unknown norm type {}".format(type)) 22 | -------------------------------------------------------------------------------- /segmentation/gvcore/utils/setup.py: -------------------------------------------------------------------------------- 1 | from gvcore.utils.config import parse_config 2 | from gvcore.utils.logger import logger 3 | from gvcore.utils.distributed import setup_dist_running, is_distributed, get_world_size 4 | 5 | 6 | def setup(): 7 | # =================== 8 | # 1. Parse config 9 | cfg = parse_config() 10 | 11 | # =================== 12 | # 2. Init distributed running 13 | setup_dist_running(cfg) 14 | 15 | # =================== 16 | # 3. Setup logger 17 | logger.setup_logger(cfg) 18 | 19 | if is_distributed(): 20 | logger.info(f"[!] Distributed Running Initialized: world size = {get_world_size()}") 21 | 22 | return cfg 23 | 24 | -------------------------------------------------------------------------------- /detection/det/utils/dataset_utils/append_aspect_ratio.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from PIL import Image 3 | import os 4 | from tqdm import tqdm 5 | 6 | for path in os.listdir("./cache"): 7 | if path.endswith(".dict"): 8 | print(path) 9 | cache = torch.load("./cache/"+path) 10 | 11 | new_cache = {} 12 | for id, value in tqdm(cache.items()): 13 | img = Image.open(value["image_path"]) 14 | if img.size[1] > img.size[0]: 15 | group = 1 16 | else: 17 | group = 0 18 | 19 | value["aspect_ratio"] = group 20 | new_cache[id] = value 21 | 22 | torch.save(new_cache, "./cache/new/" + path) -------------------------------------------------------------------------------- /detection/gvcore/utils/registry.py: -------------------------------------------------------------------------------- 1 | class Registry: 2 | def __init__(self): 3 | self.mapper = {} 4 | 5 | def register(self, key=None): 6 | def wrapper(*args): 7 | nonlocal key 8 | class_or_func = args[0] 9 | key = class_or_func.__name__ if key is None else key 10 | self.mapper[key] = args[0] 11 | return class_or_func 12 | 13 | return wrapper 14 | 15 | def get(self, key): 16 | assert key in self.mapper, "No object named '{}' found!".format(key) 17 | class_or_func = self.mapper[key] 18 | return class_or_func 19 | 20 | def __getitem__(self, key): 21 | return self.get(key) 22 | -------------------------------------------------------------------------------- /segmentation/gvcore/utils/registry.py: -------------------------------------------------------------------------------- 1 | class Registry: 2 | def __init__(self): 3 | self.mapper = {} 4 | 5 | def register(self, key=None): 6 | def wrapper(*args): 7 | nonlocal key 8 | class_or_func = args[0] 9 | key = class_or_func.__name__ if key is None else key 10 | self.mapper[key] = args[0] 11 | return class_or_func 12 | 13 | return wrapper 14 | 15 | def get(self, key): 16 | assert key in self.mapper, "No object named '{}' found!".format(key) 17 | class_or_func = self.mapper[key] 18 | return class_or_func 19 | 20 | def __getitem__(self, key): 21 | return self.get(key) 22 | -------------------------------------------------------------------------------- /segmentation/readme.md: -------------------------------------------------------------------------------- 1 | # Virtual Category Learning: A Semi-Supervised Learning Method for Dense Prediction with Extremely Limited Labels (TPAMI under review) 2 | 3 | ## Install 4 | 5 | pytorch 1.13, python 3.8, cuda 11.7 6 | 7 | ```bash 8 | python setup.py develop 9 | ``` 10 | 11 | ## Prepare Data 12 | 13 | Download data seed and pretrained resnet: [link](https://1drv.ms/u/s!As5AmExWpCHXgfluc3OqaejrSYZN8w?e=kLjBiY) 14 | 15 | ``` 16 | |-- ... 17 | |--run.py 18 | |--data 19 | |-- cityscapes 20 | |-- VOCSegAug 21 | |-- index 22 | |-- image 23 | |-- label 24 | ``` 25 | 26 | 27 | ## Train 28 | ```bash 29 | torchrun --nproc_per_node=4 run.py train --config ./config/vc_deeplabv3p/voc_aug.yaml --num-gpus=4 30 | ``` 31 | -------------------------------------------------------------------------------- /segmentation/seg/run.py: -------------------------------------------------------------------------------- 1 | from gvcore.utils.setup import setup 2 | from gvcore.operators import OPERATOR_REGISTRY 3 | 4 | # ---------------------------------------------------------------------------- # 5 | # Register 6 | 7 | import dataset as _ 8 | import evaluator as _ 9 | import operators as _ 10 | 11 | # ---------------------------------------------------------------------------- # 12 | 13 | 14 | def main(): 15 | cfg = setup() 16 | if cfg.opt == "train": 17 | operator = OPERATOR_REGISTRY[cfg.operator](cfg) 18 | operator.train() 19 | elif cfg.opt == "test": 20 | operator = OPERATOR_REGISTRY[cfg.operator](cfg) 21 | operator.test() 22 | 23 | 24 | if __name__ == "__main__": 25 | main() 26 | -------------------------------------------------------------------------------- /detection/det/ext/fastevalapi/csrc/cuda_version.cu: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | 3 | #include 4 | 5 | namespace fastevalapi { 6 | int get_cudart_version() { 7 | // Not a ROCM platform: Either HIP is not used, or 8 | // it is used, but platform is not ROCM (i.e. it is CUDA) 9 | #if !defined(__HIP_PLATFORM_HCC__) 10 | return CUDART_VERSION; 11 | #else 12 | int version = 0; 13 | 14 | #if HIP_VERSION_MAJOR != 0 15 | // Create a convention similar to that of CUDA, as assumed by other 16 | // parts of the code. 17 | 18 | version = HIP_VERSION_MINOR; 19 | version += (HIP_VERSION_MAJOR * 100); 20 | #else 21 | hipRuntimeGetVersion(&version); 22 | #endif 23 | return version; 24 | #endif 25 | } 26 | } // namespace fastevalapi 27 | -------------------------------------------------------------------------------- /detection/gvcore/model/__init__.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | 5 | class GenericModule(nn.Module): 6 | def __init__(self, cfg): 7 | super(GenericModule, self).__init__() 8 | self.cfg = cfg 9 | self.device = torch.device("cuda") 10 | 11 | def forward(self, *args, **kwargs): 12 | if self.training: 13 | return self.forward_train(*args, **kwargs) 14 | else: 15 | return self.forward_eval(*args, **kwargs) 16 | 17 | def forward_train(self, *args, **kwargs): 18 | raise NotImplementedError 19 | 20 | def forward_eval(self, *args, **kwargs): 21 | raise NotImplementedError 22 | 23 | def get_losses(self, *args, **kwargs): 24 | raise NotImplementedError 25 | 26 | def freeze(self, if_freeze=True): 27 | for p in self.parameters(): 28 | p.requires_grad = not if_freeze 29 | -------------------------------------------------------------------------------- /detection/det/config/fasterrcnn/base.yaml: -------------------------------------------------------------------------------- 1 | dev_mode: False 2 | seed: 669 3 | operator: "fasterrcnn" 4 | evaluator: "coco" 5 | 6 | log: 7 | prefix: "fasterrcnn" 8 | comment: "full" 9 | log_dir: "./log/" 10 | summary_interval: 20 11 | wandb: "FasterRCNN" 12 | aim: 13 | 14 | model: 15 | backbone: 16 | name: "resnet50" 17 | pretrained: True 18 | trainable_layers: 3 19 | rpn: 20 | pre_nms_topk_train: 2000 21 | post_nms_topk_train: 1000 22 | pre_nms_topk_test: 1000 23 | post_nms_topk_test: 1000 24 | nms_threshold: 0.7 25 | roihead: 26 | num_classes: 80 27 | score_threshold: 0.05 28 | nms_threshold: 0.5 29 | max_detections: 100 30 | 31 | solver: 32 | iter_num: 90000 33 | ckp_interval: 5000 34 | lr: 0.02 35 | weight_decay: 0.0001 36 | momentum: 0.9 37 | lr_steps: [60000, 80000] 38 | lr_gamma: 0.1 39 | warmup: True 40 | warmup_gamma: 0.001 41 | warmup_steps: 1000 42 | 43 | distributed: 44 | use: False 45 | -------------------------------------------------------------------------------- /detection/det/model/backbone/__init__.py: -------------------------------------------------------------------------------- 1 | from model.backbone.resnet import * 2 | from model.block.block import FrozenBatchNorm2d 3 | 4 | backbones = {"resnet50": resnet50, "resnet101": resnet101} 5 | 6 | 7 | def make_backbone(cfg): 8 | backbone = backbones[cfg.model.backbone.name]( 9 | pretrained=cfg.model.backbone.pretrained, norm_layer=FrozenBatchNorm2d 10 | ) 11 | 12 | trainable_layers = cfg.model.backbone.trainable_layers 13 | assert ( 14 | 5 >= trainable_layers >= 0 and cfg.model.backbone.pretrained 15 | ), "Trainable layer can only set from 0 to 5 with pretrained = True" 16 | if cfg.model.backbone.pretrained: 17 | layers_to_train = ["res5", "res4", "res3", "res2", "stem"][:trainable_layers] 18 | for name, parameter in backbone.named_parameters(): 19 | if all([not name.startswith(layer) for layer in layers_to_train]): 20 | parameter.requires_grad_(False) 21 | 22 | return backbone 23 | -------------------------------------------------------------------------------- /readme.md: -------------------------------------------------------------------------------- 1 | # Virtual Category Learning 2 | 3 | This is the code for the virtual category learning. 4 | 5 | - [x] Detection: Semi-supervised Object Detection via Virtual Category Learning (ECCV 2022) 6 | - [x] Segmentation: Virtual Category Learning: A Semi-Supervised Learning Method for Dense Prediction with Extremely Limited Labels (T-PAMI) 7 | 8 | 9 | ## Citation 10 | 11 | ```bibtex 12 | @inproceedings{chen2022ssodvc, 13 | year = 2022, 14 | title = {Semi-supervised Object Detection via Virtual Category Learning.}, 15 | author = {Changrui Chen and Kurt Debattista and Jungong Han}, 16 | booktitle = {European Conference on Computer Vision (ECCV)} 17 | } 18 | 19 | @article{chen2024vcl, 20 | year = 2024, 21 | title = {Virtual Category Learning: A Semi-Supervised Learning Method for Dense Prediction with Extremely Limited Labels.}, 22 | author = {Changrui Chen and Kurt Debattista and Jungong Han}, 23 | journal = {IEEE Transactions on Pattern Analysis and Machine Intelligence (T-PAMI)} 24 | } 25 | ``` 26 | -------------------------------------------------------------------------------- /segmentation/seg/operators/deeplabv3p.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.nn.parallel import DistributedDataParallel 3 | 4 | from gvcore.operators import OPERATOR_REGISTRY, GenericOpt 5 | from gvcore.evaluator import EVALUATOR_REGISTRY 6 | import gvcore.utils.distributed as dist_utils 7 | 8 | from model.segmentor.deeplabv3p import DeeplabV3p 9 | 10 | 11 | @OPERATOR_REGISTRY.register("deeplabv3p") 12 | class DeeplabV3pOpt(GenericOpt): 13 | def __init__(self, cfg): 14 | super(DeeplabV3pOpt, self).__init__(cfg) 15 | 16 | def build_model(self): 17 | model = DeeplabV3p(self.cfg) 18 | model.to(self.device) 19 | if dist_utils.is_distributed(): 20 | model = DistributedDataParallel(model) 21 | self.model = model 22 | 23 | def build_evaluator(self): 24 | if self.evaluator is None: 25 | self.evaluator = EVALUATOR_REGISTRY[self.cfg.evaluator]( 26 | self.cfg.model.num_classes, distributed=dist_utils.is_distributed() 27 | ) 28 | self.evaluator.reset() 29 | -------------------------------------------------------------------------------- /detection/det/config/fasterrcnn/coco.yaml: -------------------------------------------------------------------------------- 1 | _base: "./base.yaml" 2 | 3 | log: 4 | comment: "coco" 5 | 6 | data: 7 | train: 8 | dataset: "COCO" 9 | prefetch: False 10 | infinity_sampler: True 11 | grouped_sampler: True 12 | cache: "./cache/coco17_train.dict" 13 | num_workers: 4 14 | batch_size: 4 15 | shuffle: True 16 | img_transforms: 17 | batch_transforms: 18 | batch_apply: 19 | random_horizontal_flip: 20 | resize: 21 | - [640, 672, 704, 736, 768, 800] 22 | - 1333 23 | normalize: 24 | batch_pad: 25 | test: 26 | dataset: "COCO" 27 | prefetch: False 28 | infinity_sampler: False 29 | grouped_sampler: False 30 | cache: "./cache/coco17_val.dict" 31 | json_file: "./data/COCO2017/annotations/instances_val2017.json" 32 | num_workers: 4 33 | batch_size: 1 34 | shuffle: False 35 | img_transforms: 36 | group_transforms: 37 | batch_apply: 38 | resize: 39 | - 800 40 | - 1333 41 | normalize: 42 | batch_pad: 43 | batch_transforms: 44 | -------------------------------------------------------------------------------- /detection/readme.md: -------------------------------------------------------------------------------- 1 | # Semi-supervised Object Detection via Virtual Category Learning (ECCV 2022) 2 | 3 | ## Install 4 | 5 | pytorch 1.11, torchvision 0.12, python 3.8, cuda 11.3 6 | 7 | ```bash 8 | python setup.py develop 9 | ``` 10 | 11 | ## Prepare Data 12 | 13 | Download data and cache: [OneDrive](https://1drv.ms/u/s!As5AmExWpCHXgdI2u7uNIZaFkKBGIA?e=pmFmXY) 14 | 15 | ``` 16 | |-- ... 17 | |--run.py 18 | |--data 19 | |--COCO2017 20 | |--cache 21 | |-- coco17_x_x.dict 22 | |-- ... 23 | |-- resnet50.pth 24 | ``` 25 | 26 | 27 | ## Train 28 | ```bash 29 | python run.py train --num_gpus 8 --config ./config/virtual_category/coco.yaml 30 | ``` 31 | 32 | ## Test 33 | 34 | ```bash 35 | python run.py test --num_gpus 8 --config ./config/virtual_category/coco.yaml --resume ./log/vc/01/final.pth 36 | ``` 37 | 38 | ## Citation 39 | 40 | ```bibtex 41 | @inproceedings{chen2022ssodvc, 42 | year = 2022, 43 | title = {Semi-supervised Object Detection via Virtual Category Learning.}, 44 | author = {Changrui Chen and Kurt Debattista and Jungong Han}, 45 | booktitle = {European Conference on Computer Vision (ECCV)} 46 | } 47 | ``` -------------------------------------------------------------------------------- /segmentation/seg/utils/data_utils/voc_split.py: -------------------------------------------------------------------------------- 1 | import random 2 | import torchvision.io as io 3 | import torch 4 | 5 | 6 | ratio = 128 7 | fold = 3 8 | 9 | index = [] 10 | with open("./data/VOCSegAug/index/train.txt", "r") as f: 11 | lines = f.readlines() 12 | for i in lines: 13 | index.append(i.strip()) 14 | 15 | n = int(1 / ratio * len(index)) 16 | 17 | random.shuffle(index) 18 | 19 | 20 | with open(f"./data/VOCSegAug/index/l.1.{ratio}.fold{fold}.txt", "w") as f: 21 | for i in index[:n]: 22 | f.write(i + "\n") 23 | 24 | with open(f"./data/VOCSegAug/index/u.1.{ratio}.fold{fold}.txt", "w") as f: 25 | for i in index[n:]: 26 | f.write(i + "\n") 27 | 28 | 29 | with open(f"./data/VOCSegAug/index/l.1.{ratio}.fold{fold}.txt", "r") as f: 30 | index = f.readlines() 31 | index = [i.strip() for i in index] 32 | 33 | 34 | counts = torch.zeros(256) 35 | 36 | for i in index: 37 | label = io.read_image(f"./data/VOCSegAug/label/{i}.png") 38 | cls, count = label.unique(return_counts=True) 39 | counts = torch.index_add(counts, 0, cls.long(), count.float()) 40 | 41 | print(counts[:21] / counts[:21].sum()) 42 | -------------------------------------------------------------------------------- /detection/det/run.py: -------------------------------------------------------------------------------- 1 | from gvcore.utils.config import parse_config, parse_args 2 | import gvcore.utils.distributed as dist_utils 3 | from gvcore.utils.logger import logger 4 | from gvcore.operators import OPERATOR_REGISTRY 5 | 6 | # ---------------------------------------------------------------------------- # 7 | # Register 8 | 9 | import dataset as _ 10 | import evaluator as _ 11 | import operators as _ 12 | 13 | # ---------------------------------------------------------------------------- # 14 | 15 | 16 | def launch(cfg): 17 | logger.setup_logger(cfg, dist_utils.is_main_process()) 18 | 19 | if cfg.opt == "train": 20 | operator = OPERATOR_REGISTRY[cfg.operator](cfg) 21 | operator.train() 22 | elif cfg.opt == "test": 23 | operator = OPERATOR_REGISTRY[cfg.operator](cfg) 24 | operator.test() 25 | 26 | 27 | if __name__ == "__main__": 28 | args = parse_args() 29 | cfg = parse_config(args) 30 | if not args.launch and cfg.distributed.use: 31 | dist_utils.launch_distributed(args) 32 | else: 33 | if cfg.distributed.use: 34 | dist_utils.init_distributed(cfg) 35 | launch(cfg) 36 | -------------------------------------------------------------------------------- /detection/det/model/loss/functional.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | from torch.nn import functional as F 4 | 5 | 6 | def focal_loss(x, target, mask=None, gamma=1.0, reduction="sum"): 7 | if x.shape[0] == 0: 8 | return 0 * x.sum() 9 | # 1. One-hot target 10 | if target.dim() == 1: 11 | target = F.one_hot(target, num_classes=x.shape[1]).float() 12 | 13 | # 2. Generate Mask 14 | if mask is None: 15 | mask = torch.ge(target, 0).float() 16 | else: 17 | assert mask.shape == target.shape, "Mask shape must be equal to target shape!" 18 | mask = mask.float() 19 | 20 | exp_x = x.exp() 21 | pos_term = (1 / exp_x * target * mask).sum(dim=1) 22 | neg_term = (exp_x * torch.eq(target, 0).float() * mask).sum(dim=1) 23 | 24 | CE = (1 + pos_term * neg_term).log() 25 | 26 | p = torch.exp(-CE) 27 | loss = (1 - p) ** gamma * CE 28 | 29 | if reduction == "sum": 30 | return loss.sum() 31 | elif reduction == "mean": 32 | return loss.mean() 33 | elif reduction == "none": 34 | return loss 35 | else: 36 | raise ValueError("Unsupported reduction") 37 | -------------------------------------------------------------------------------- /segmentation/seg/model/loss/functional.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | from torch.nn import functional as F 4 | 5 | 6 | def focal_loss(x, target, mask=None, gamma=1.0, reduction="sum"): 7 | if x.shape[0] == 0: 8 | return 0 * x.sum() 9 | # 1. One-hot target 10 | if target.dim() == 1: 11 | target = F.one_hot(target, num_classes=x.shape[1]).float() 12 | 13 | # 2. Generate Mask 14 | if mask is None: 15 | mask = torch.ge(target, 0).float() 16 | else: 17 | assert mask.shape == target.shape, "Mask shape must be equal to target shape!" 18 | mask = mask.float() 19 | 20 | exp_x = x.exp() 21 | pos_term = (1 / exp_x * target * mask).sum(dim=1) 22 | neg_term = (exp_x * torch.eq(target, 0).float() * mask).sum(dim=1) 23 | 24 | CE = (1 + pos_term * neg_term).log() 25 | 26 | p = torch.exp(-CE) 27 | loss = (1 - p) ** gamma * CE 28 | 29 | if reduction == "sum": 30 | return loss.sum() 31 | elif reduction == "mean": 32 | return loss.mean() 33 | elif reduction == "none": 34 | return loss 35 | else: 36 | raise ValueError("Unsupported reduction") 37 | -------------------------------------------------------------------------------- /segmentation/gvcore/model/__init__.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | from gvcore.utils.registry import Registry 5 | 6 | MODEL_REGISTRY = Registry() 7 | 8 | class GenericModule(nn.Module): 9 | def __init__(self, cfg=None): 10 | super(GenericModule, self).__init__() 11 | self.cfg = cfg 12 | self.device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") 13 | 14 | def forward(self, *args, **kwargs): 15 | if self.training: 16 | return self.forward_train(*args, **kwargs) 17 | else: 18 | return self.forward_eval(*args, **kwargs) 19 | 20 | def forward_generic(self, *args, **kwargs): 21 | raise NotImplementedError 22 | 23 | def forward_train(self, *args, **kwargs): 24 | return self.forward_generic(*args, **kwargs) 25 | 26 | def forward_eval(self, *args, **kwargs): 27 | return self.forward_generic(*args, **kwargs) 28 | 29 | def get_losses(self, *args, **kwargs): 30 | raise NotImplementedError 31 | 32 | def freeze(self, if_freeze=True): 33 | for p in self.parameters(): 34 | p.requires_grad = not if_freeze 35 | -------------------------------------------------------------------------------- /segmentation/seg/utils/data_utils/cityscape_split.py: -------------------------------------------------------------------------------- 1 | import random 2 | import torchvision.io as io 3 | import torch 4 | 5 | 6 | ratio = 80 7 | fold = 2 8 | 9 | index = [] 10 | with open("./data/cityscapes/index/train.txt", "r") as f: 11 | lines = f.readlines() 12 | for i in lines: 13 | index.append(i.strip()) 14 | 15 | n = int(1 / ratio * len(index)) 16 | 17 | random.shuffle(index) 18 | 19 | 20 | with open(f"./data/cityscapes/index/l.1.{ratio}.fold{fold}.txt", "w") as f: 21 | for i in index[:n]: 22 | f.write(i + "\n") 23 | 24 | with open(f"./data/cityscapes/index/u.1.{ratio}.fold{fold}.txt", "w") as f: 25 | for i in index[n:]: 26 | f.write(i + "\n") 27 | 28 | 29 | with open(f"./data/cityscapes/index/l.1.{ratio}.fold{fold}.txt", "r") as f: 30 | index = f.readlines() 31 | index = [i.strip() for i in index] 32 | 33 | 34 | counts = torch.zeros(256) 35 | 36 | for i in index: 37 | label = io.read_image(f"./data/cityscapes/gtFine/{i}_gtFine_labelTrainIds.png") 38 | cls, count = label.unique(return_counts=True) 39 | counts = torch.index_add(counts, 0, cls.long(), count.float()) 40 | 41 | print(counts[:19] / counts[:19].sum()) 42 | -------------------------------------------------------------------------------- /detection/det/model/detector/unbiased_teacher.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from gvcore.utils.types import TDataList 4 | 5 | from model.detector.fasterrcnn import FasterRCNN 6 | from model.head.unbiased_teacher.rpn import RPN 7 | from model.head.unbiased_teacher.roihead import ROIHead 8 | 9 | 10 | class UBTDetector(FasterRCNN): 11 | def __init__(self, cfg): 12 | super(UBTDetector, self).__init__(cfg) 13 | self.rpn = RPN(cfg) 14 | self.roi_heads = ROIHead(cfg) 15 | 16 | def forward_train(self, data_list: TDataList, labeled: bool = False): 17 | imgs = torch.stack([data.img for data in data_list], dim=0) 18 | labels = [data.label for data in data_list] 19 | img_sizes = [data.meta.cur_size for data in data_list] 20 | 21 | feats = self.backbone(imgs) 22 | feats = self.fpn(feats) 23 | self.feats = feats 24 | 25 | proposals, rpn_loss_dict = self.rpn(feats, img_sizes, labels, not labeled) 26 | roi_loss_dict = self.roi_heads(feats, proposals, labels, not labeled) 27 | 28 | loss_dict = {} 29 | loss_dict.update(rpn_loss_dict) 30 | loss_dict.update(roi_loss_dict) 31 | return loss_dict 32 | -------------------------------------------------------------------------------- /segmentation/seg/config/deeplabv3p/base.yaml: -------------------------------------------------------------------------------- 1 | dev_mode: False 2 | seed: 669 3 | operator: "deeplabv3p" 4 | evaluator: "segmentation" 5 | 6 | log: 7 | prefix: "deeplabv3p" 8 | comment: "full" 9 | log_dir: "./log/" 10 | summary_interval: 20 11 | wandb: "SEG" 12 | aim: 13 | 14 | model: 15 | backbone: 16 | type: "resnet" 17 | depth: 50 18 | strides: [1, 2, 1, 1] 19 | dilations: [1, 1, 2, 4] 20 | contract_dilation: True 21 | deep_stem: True 22 | norm_layer: "SyncBN" 23 | pretrained: "./cache/resnet50c.pth" 24 | aspp: 25 | in_channels: 2048 26 | inner_channels: 512 27 | lowlevel_in_channels: 256 28 | lowlevel_inner_channels: 48 29 | dilations: [1, 6, 12, 18] 30 | norm_layer: "SyncBN" 31 | dropout: 0 32 | 33 | solver: 34 | iter_num: 40000 35 | ckp_interval: 5000 36 | eval_interval: 1000 37 | optimizer: 38 | type: "sgd" 39 | lr: 0.001 40 | weight_decay: 0.0001 41 | weight_decay_norm: 0 42 | weight_decay_bias: 0 43 | momentum: 0.9 44 | lr_scheduler: 45 | type: "poly" 46 | iter_num: 40000 47 | power: 0.9 48 | min_lr: 0.0001 49 | warmup: True 50 | warmup_step: 500 51 | warmup_gamma: 0.001 52 | -------------------------------------------------------------------------------- /detection/det/operators/virtual_category_cross.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.nn.parallel import DistributedDataParallel 3 | 4 | import gvcore.utils.distributed as dist_utils 5 | from model.detector.virtual_category import VCDetector 6 | from model.detector.virtual_category_cross import VCCrossTSDetector 7 | from operators.virtual_category import VirtualCategoryOpt 8 | 9 | from gvcore.operators import OPERATOR_REGISTRY 10 | 11 | 12 | @OPERATOR_REGISTRY.register("virtual_category_cross") 13 | class VirtualCategoryCrossOpt(VirtualCategoryOpt): 14 | def __init__(self, cfg): 15 | super(VirtualCategoryCrossOpt, self).__init__(cfg) 16 | 17 | def build_model(self): 18 | model = VCCrossTSDetector(self.cfg, VCDetector) 19 | model.to(self.device) 20 | if dist_utils.is_distributed(): 21 | model = DistributedDataParallel(model) 22 | self.model = model 23 | 24 | def save_ckp(self, name=None): 25 | self.checkpointer.save( 26 | name=name if name is not None else self.cur_step, 27 | model=self.model, 28 | optimizer=self.optimizer, 29 | lr_scheduler=self.lr_scheduler, 30 | step=self.cur_step, 31 | ) 32 | -------------------------------------------------------------------------------- /segmentation/gvcore/dataset/__init__.py: -------------------------------------------------------------------------------- 1 | from typing import Callable, Optional 2 | import torch 3 | from gvcore.utils.structure import GenericData 4 | from gvcore.utils.registry import Registry 5 | 6 | 7 | DATASET_REGISTRY = Registry() 8 | 9 | 10 | class GenericDataset(torch.utils.data.Dataset): 11 | def __init__( 12 | self, root: str, subset: str = "train", img_transforms: Optional[Callable[[GenericData], GenericData]] = None 13 | ): 14 | super(GenericDataset, self).__init__() 15 | 16 | self._root = root 17 | self._subset = subset 18 | self._img_transforms = img_transforms 19 | 20 | self._data_ids = self._get_data_ids() 21 | 22 | def _get_data_ids(self): 23 | raise NotImplementedError 24 | 25 | def __len__(self): 26 | return len(self._data_ids) 27 | 28 | def _load(self, item_idx: int) -> GenericData: 29 | raise NotImplementedError 30 | 31 | def __getitem__(self, item_idx): 32 | sample = self._load(item_idx) 33 | if self._img_transforms is not None: 34 | sample = self._img_transforms(sample) 35 | return sample 36 | 37 | @staticmethod 38 | def collate_fn(batch): 39 | return batch 40 | -------------------------------------------------------------------------------- /detection/det/utils/dataset_utils/split.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import os.path as osp 3 | import json 4 | from tqdm import tqdm 5 | 6 | cache_path = "./cache/" 7 | 8 | coco_full_set = torch.load(osp.join(cache_path, "coco17_train.dict")) 9 | 10 | l_ratio = "10" 11 | ul_ratio = "90" 12 | 13 | str_map = { 14 | "005": "0.5", 15 | "01": "1.0", 16 | "02": "2.0", 17 | "05": "5.0", 18 | "10": "10.0", 19 | } 20 | 21 | for seed in ["1", "2", "3", "4", "5"]: 22 | l_idx = torch.load(f"./cache/ubt_seeds/l_{str_map[l_ratio]}_{seed}.pth") 23 | ul_idx = torch.load(f"./cache/ubt_seeds/ul_{str_map[l_ratio]}_{seed}.pth") 24 | 25 | print("l: ", len(l_idx)) 26 | print("ul: ", len(ul_idx)) 27 | 28 | coco_l_set = {} 29 | coco_u_set = {} 30 | 31 | for idx in tqdm(l_idx): 32 | try: 33 | coco_l_set[idx] = coco_full_set[idx] 34 | except: 35 | print(idx) 36 | 37 | for idx in tqdm(ul_idx): 38 | try: 39 | coco_u_set[idx] = coco_full_set[idx] 40 | except: 41 | print(idx) 42 | 43 | print(len(coco_l_set), len(coco_u_set)) 44 | 45 | torch.save(coco_l_set, osp.join(cache_path, f"coco17_train{l_ratio}_{seed}.dict")) 46 | torch.save(coco_u_set, osp.join(cache_path, f"coco17_train{ul_ratio}_{seed}.dict")) 47 | -------------------------------------------------------------------------------- /detection/det/config/fasterrcnn/voc.yaml: -------------------------------------------------------------------------------- 1 | _base: "./base.yaml" 2 | 3 | log: 4 | comment: "voc" 5 | 6 | evaluator: "voc" 7 | 8 | data: 9 | train: 10 | dataset: "VOC" 11 | prefetch: False 12 | infinity_sampler: True 13 | grouped_sampler: True 14 | cache: "./cache/voc07_trainval.dict" 15 | num_workers: 4 16 | batch_size: 4 17 | shuffle: True 18 | img_transforms: 19 | batch_transforms: 20 | batch_apply: 21 | random_horizontal_flip: 22 | resize: 23 | - [480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800] 24 | - 1333 25 | normalize: 26 | batch_pad: 27 | test: 28 | dataset: "VOC" 29 | prefetch: False 30 | infinity_sampler: False 31 | grouped_sampler: True 32 | cache: "./cache/voc07_test.dict" 33 | json_file: "./cache/voc07_test_cocostyle.json" 34 | num_workers: 4 35 | batch_size: 1 36 | shuffle: False 37 | img_transforms: 38 | batch_transforms: 39 | batch_apply: 40 | resize: 41 | - 600 42 | - 1000 43 | normalize: 44 | batch_pad: 45 | 46 | model: 47 | roihead: 48 | num_classes: 20 49 | 50 | solver: 51 | iter_num: 18000 52 | ckp_interval: 5000 53 | lr: 0.02 54 | weight_decay: 0.0001 55 | momentum: 0.9 56 | lr_steps: [12000, 16000] 57 | lr_gamma: 0.1 58 | warmup: True 59 | warmup_gamma: 0.001 60 | warmup_steps: 100 61 | -------------------------------------------------------------------------------- /segmentation/seg/dataset/voc.py: -------------------------------------------------------------------------------- 1 | from typing import Callable, Optional 2 | import torchvision.io as io 3 | 4 | from gvcore.utils.structure import GenericData 5 | from gvcore.dataset import DATASET_REGISTRY, GenericDataset 6 | 7 | 8 | @DATASET_REGISTRY.register("VOC") 9 | class VOCDataset(GenericDataset): 10 | def __init__( 11 | self, root: str, subset: str = "train", img_transforms: Optional[Callable[[GenericData], GenericData]] = None 12 | ): 13 | super(VOCDataset, self).__init__(root, subset, img_transforms) 14 | 15 | def _get_data_ids(self): 16 | with open(f"{self._root}/index/{self._subset}.txt", "r") as f: 17 | data_ids = [line.strip() for line in f.readlines()] 18 | 19 | return sorted(data_ids) 20 | 21 | def _load(self, item_idx): 22 | data_id = self._data_ids[item_idx] 23 | img = io.read_image(f"{self._root}/image/{data_id}.jpg") 24 | if img.shape[0] == 1: 25 | img = img.repeat(3, 1, 1) 26 | label = io.read_image(f"{self._root}/label/{data_id}.png") 27 | 28 | data = GenericData() 29 | data.set("img", img) 30 | data.set("label", label) 31 | 32 | data.set("meta", GenericData()) 33 | data.meta.set("id", data_id) 34 | 35 | img_size = data.img.shape[1:] 36 | data.meta.set("ori_size", list(img_size)) 37 | data.meta.set("cur_size", list(img_size)) 38 | 39 | return data 40 | -------------------------------------------------------------------------------- /segmentation/seg/config/deeplabv3p/voc_aug.yaml: -------------------------------------------------------------------------------- 1 | _base: "./base.yaml" 2 | 3 | log: 4 | comment: "voc" 5 | 6 | data: 7 | train: 8 | dataset: "VOC" 9 | subset: "l.1.128.fold3" 10 | prefetch: False 11 | infinity_sampler: True 12 | root: "./data/VOCSegAug" 13 | num_workers: 4 14 | batch_size: 4 15 | shuffle: True 16 | img_transforms: 17 | batch_transforms: 18 | batch_apply: 19 | resize: 20 | scale: [0.5, 0.75, 1, 1.5, 1.75, 2.0] 21 | random_crop: 22 | - [512, 512] 23 | random_horizontal_flip: 24 | normalize: 25 | - [0.485, 0.456, 0.406] 26 | - [0.229, 0.224, 0.225] 27 | - True 28 | - False 29 | batch_pad: 30 | fixed_size: [512, 512] 31 | test: 32 | dataset: "VOC" 33 | subset: "val" 34 | prefetch: False 35 | infinity_sampler: False 36 | root: "./data/VOCSegAug" 37 | num_workers: 1 38 | batch_size: 1 39 | shuffle: False 40 | img_transforms: 41 | batch_transforms: 42 | batch_apply: 43 | resize: 44 | min_size: [512, 512] 45 | max_size: 2048 46 | mode: "choice" 47 | process_label: False 48 | normalize: 49 | - [0.485, 0.456, 0.406] 50 | - [0.229, 0.224, 0.225] 51 | - True 52 | - False 53 | batch_pad: 54 | stride: 1 55 | 56 | model: 57 | num_classes: 21 58 | -------------------------------------------------------------------------------- /detection/det/config/unbiased_teacher/coco.yaml: -------------------------------------------------------------------------------- 1 | _base: "./base.yaml" 2 | 3 | log: 4 | comment: "coco" 5 | 6 | data: 7 | train_l: 8 | dataset: "COCO" 9 | prefetch: False 10 | infinity_sampler: True 11 | grouped_sampler: True 12 | cache: "./cache/coco17_train01_1.dict" 13 | num_workers: 4 14 | batch_size: 2 15 | shuffle: True 16 | img_transforms: 17 | batch_transforms: 18 | batch_apply: 19 | random_horizontal_flip: 20 | resize: 21 | - [640, 672, 704, 736, 768, 800] 22 | - 1333 23 | train_u: 24 | dataset: "COCO" 25 | prefetch: False 26 | infinity_sampler: True 27 | grouped_sampler: True 28 | cache: "./cache/coco17_train99_1.dict" 29 | num_workers: 4 30 | batch_size: 2 31 | shuffle: True 32 | img_transforms: 33 | batch_transforms: 34 | batch_apply: 35 | random_horizontal_flip: 36 | resize: 37 | - [640, 672, 704, 736, 768, 800] 38 | - 1333 39 | test: 40 | dataset: "COCO" 41 | prefetch: False 42 | infinity_sampler: False 43 | grouped_sampler: False 44 | cache: "./cache/coco17_val.dict" 45 | json_file: "./data/COCO2017/annotations/instances_val2017.json" 46 | num_workers: 4 47 | batch_size: 1 48 | shuffle: False 49 | img_transforms: 50 | batch_transforms: 51 | batch_apply: 52 | resize: 53 | - 800 54 | - 1333 55 | normalize: 56 | batch_pad: 57 | 58 | -------------------------------------------------------------------------------- /detection/det/config/virtual_category_cross/coco.yaml: -------------------------------------------------------------------------------- 1 | _base: "./base.yaml" 2 | 3 | log: 4 | comment: "coco" 5 | 6 | data: 7 | train_l: 8 | dataset: "COCO" 9 | prefetch: False 10 | infinity_sampler: True 11 | grouped_sampler: True 12 | cache: "./cache/coco17_train01_1.dict" 13 | num_workers: 4 14 | batch_size: 1 15 | shuffle: True 16 | img_transforms: 17 | batch_transforms: 18 | batch_apply: 19 | random_horizontal_flip: 20 | resize: 21 | - [400, 1200] 22 | - 1333 23 | - "range" 24 | train_u: 25 | dataset: "COCO" 26 | prefetch: False 27 | infinity_sampler: True 28 | grouped_sampler: True 29 | cache: "./cache/coco17_train99_1.dict" 30 | num_workers: 4 31 | batch_size: 4 32 | shuffle: True 33 | img_transforms: 34 | batch_transforms: 35 | batch_apply: 36 | random_horizontal_flip: 37 | resize: 38 | - [400, 1200] 39 | - 1333 40 | - "range" 41 | test: 42 | dataset: "COCO" 43 | prefetch: False 44 | infinity_sampler: False 45 | grouped_sampler: False 46 | cache: "./cache/coco17_val.dict" 47 | json_file: "./data/COCO2017/annotations/instances_val2017.json" 48 | num_workers: 4 49 | batch_size: 1 50 | shuffle: False 51 | img_transforms: 52 | batch_transforms: 53 | batch_apply: 54 | resize: 55 | - 800 56 | - 1333 57 | normalize: 58 | batch_pad: 59 | -------------------------------------------------------------------------------- /detection/det/config/virtual_category/coco.yaml: -------------------------------------------------------------------------------- 1 | _base: "./base.yaml" 2 | 3 | log: 4 | comment: "coco" 5 | 6 | data: 7 | train_l: 8 | dataset: "COCO" 9 | prefetch: False 10 | infinity_sampler: True 11 | grouped_sampler: True 12 | cache: "./cache/coco17_train01_1.dict" 13 | num_workers: 4 14 | batch_size: 1 15 | shuffle: True 16 | img_transforms: 17 | batch_transforms: 18 | batch_apply: 19 | random_horizontal_flip: 20 | resize: 21 | - [400, 1200] 22 | - 1333 23 | - "range" 24 | train_u: 25 | dataset: "COCO" 26 | prefetch: False 27 | infinity_sampler: True 28 | grouped_sampler: True 29 | cache: "./cache/coco17_train99_1.dict" 30 | num_workers: 4 31 | batch_size: 4 32 | shuffle: True 33 | img_transforms: 34 | batch_transforms: 35 | batch_apply: 36 | random_horizontal_flip: 37 | resize: 38 | - [400, 1200] 39 | - 1333 40 | - "range" 41 | test: 42 | dataset: "COCO" 43 | prefetch: False 44 | infinity_sampler: False 45 | grouped_sampler: False 46 | cache: "./cache/coco_instances_val2017.dict" 47 | json_file: "./data/COCO2017/annotations/instances_val2017.json" 48 | num_workers: 4 49 | batch_size: 1 50 | shuffle: False 51 | img_transforms: 52 | batch_transforms: 53 | batch_apply: 54 | resize: 55 | - 800 56 | - 1333 57 | normalize: 58 | batch_pad: 59 | -------------------------------------------------------------------------------- /detection/det/config/unbiased_teacher/base.yaml: -------------------------------------------------------------------------------- 1 | dev_mode: False 2 | seed: 669 3 | operator: "unbiased_teacher" 4 | evaluator: "coco" 5 | 6 | log: 7 | prefix: "unbiased_teacher" 8 | comment: 9 | log_dir: "./log/" 10 | summary_interval: 20 11 | wandb: "UBT" 12 | aim: 13 | 14 | data: 15 | weak_transforms: 16 | batch_apply: 17 | normalize: 18 | batch_pad: 19 | strong_transforms: 20 | batch_apply: 21 | random_color_jitter: 22 | random_grayscale: 23 | random_gaussian_blur: 24 | normalize: 25 | repeat: 26 | random_erase: 27 | - [[0.05, 0.2], [0.3, 3.3], 0.7] 28 | - [[0.02, 0.2], [0.1, 6.0], 0.5] 29 | - [[0.02, 0.2], [0.05, 8.0], 0.3] 30 | batch_pad: 31 | 32 | model: 33 | pretrained: 34 | teacher: 35 | score_threshold: 0.7 36 | backbone: 37 | name: "resnet50" 38 | pretrained: True 39 | trainable_layers: 3 40 | rpn: 41 | pre_nms_topk_train: 2000 42 | post_nms_topk_train: 1000 43 | pre_nms_topk_test: 1000 44 | post_nms_topk_test: 1000 45 | nms_threshold: 0.7 46 | roihead: 47 | num_classes: 80 48 | score_threshold: 0.05 49 | nms_threshold: 0.5 50 | max_detections: 100 51 | 52 | solver: 53 | burnin_iter_num: 2000 54 | iter_num: 180000 55 | ckp_interval: 500 56 | lr: 0.01 57 | weight_decay: 0.0001 58 | momentum: 0.9 59 | lr_steps: [179990, 179995] 60 | lr_gamma: 0.1 61 | warmup: True 62 | warmup_gamma: 0.001 63 | warmup_steps: 1000 64 | ema_momentum: 0.9996 65 | u_loss_weight: 4 66 | 67 | distributed: 68 | use: False 69 | -------------------------------------------------------------------------------- /segmentation/gvcore/dataset/build.py: -------------------------------------------------------------------------------- 1 | from gvcore.dataset.transforms import parse_transform_config 2 | from gvcore.dataset.dataloader import BasicDataloader 3 | from gvcore.dataset.sampler import GroupedBatchSampler, InfiniteSampler, SequentialSampler 4 | from gvcore.dataset import DATASET_REGISTRY 5 | 6 | 7 | def build_dataloader(cfg, seed=669): 8 | # 1. Construct transforms 9 | img_transforms = parse_transform_config(cfg.img_transforms) 10 | batch_transforms = parse_transform_config(cfg.batch_transforms) 11 | 12 | # 2. Construct dataset. 13 | dataset = DATASET_REGISTRY[cfg.dataset](cfg.root, subset=cfg.subset, img_transforms=img_transforms) 14 | 15 | # 3. Construct sampler 16 | if not cfg.infinity_sampler: 17 | data_sampler = SequentialSampler(size=len(dataset)) 18 | else: 19 | data_sampler = InfiniteSampler(size=len(dataset), seed=seed, shuffle=cfg.shuffle) 20 | 21 | if cfg.grouped_sampler: 22 | batch_sampler = GroupedBatchSampler(data_sampler, group_ids=dataset.group, batch_size=cfg.batch_size) 23 | else: 24 | batch_sampler = None 25 | 26 | # 4. Construct dataloader 27 | data_loader = BasicDataloader( 28 | dataset, 29 | batch_size=cfg.batch_size, 30 | sampler=data_sampler, 31 | batch_sampler=batch_sampler, 32 | num_workers=cfg.get("num_workers", 4), 33 | pin_memory=True, 34 | prefetch=cfg.get("prefetch", False), 35 | batch_transforms=batch_transforms, 36 | drop_last=cfg.get("drop_last", True), 37 | ) 38 | 39 | return data_loader 40 | -------------------------------------------------------------------------------- /detection/det/config/virtual_category/base.yaml: -------------------------------------------------------------------------------- 1 | dev_mode: False 2 | seed: 669 3 | operator: "virtual_category" 4 | evaluator: "coco" 5 | 6 | log: 7 | prefix: "virtual_category" 8 | comment: 9 | log_dir: "./log/" 10 | summary_interval: 20 11 | wandb: "VC-learning" 12 | aim: 13 | 14 | data: 15 | weak_transforms: 16 | batch_apply: 17 | normalize: 18 | batch_pad: 19 | strong_transforms: 20 | batch_apply: 21 | random_color_jitter: 22 | random_grayscale: 23 | random_gaussian_blur: 24 | normalize: 25 | repeat: 26 | random_erase: 27 | - [[0.05, 0.2], [0.3, 3.3], 0.7] 28 | - [[0.02, 0.2], [0.1, 6.0], 0.5] 29 | - [[0.02, 0.2], [0.05, 8.0], 0.3] 30 | batch_pad: 31 | 32 | model: 33 | pretrained: 34 | teacher: 35 | score_threshold: 0.7 36 | backbone: 37 | name: "resnet50" 38 | pretrained: True 39 | trainable_layers: 3 40 | rpn: 41 | pre_nms_topk_train: 2000 42 | post_nms_topk_train: 1000 43 | pre_nms_topk_test: 1000 44 | post_nms_topk_test: 1000 45 | nms_threshold: 0.7 46 | roihead: 47 | num_classes: 80 48 | temprature: 3.5 49 | loc_t: 0.05 50 | score_threshold: 0.05 51 | nms_threshold: 0.5 52 | max_detections: 100 53 | 54 | solver: 55 | burnin_iter_num: 2000 56 | iter_num: 180000 57 | ckp_interval: 5000 58 | lr: 0.01 59 | weight_decay: 0.0001 60 | momentum: 0.9 61 | lr_steps: [179990, 179995] 62 | lr_gamma: 0.1 63 | warmup: True 64 | warmup_gamma: 0.001 65 | warmup_steps: 1000 66 | ema_momentum: 0.9996 67 | u_loss_weight: 4 68 | 69 | distributed: 70 | use: False 71 | -------------------------------------------------------------------------------- /detection/det/config/virtual_category_cross/base.yaml: -------------------------------------------------------------------------------- 1 | dev_mode: False 2 | seed: 669 3 | operator: "virtual_category_cross" 4 | evaluator: "coco" 5 | 6 | log: 7 | prefix: "virtual_category" 8 | comment: 9 | log_dir: "./log/" 10 | summary_interval: 20 11 | wandb: "VC-learning" 12 | aim: 13 | 14 | data: 15 | weak_transforms: 16 | batch_apply: 17 | normalize: 18 | batch_pad: 19 | strong_transforms: 20 | batch_apply: 21 | random_color_jitter: 22 | random_grayscale: 23 | random_gaussian_blur: 24 | normalize: 25 | repeat: 26 | random_erase: 27 | - [[0.05, 0.2], [0.3, 3.3], 0.7] 28 | - [[0.02, 0.2], [0.1, 6.0], 0.5] 29 | - [[0.02, 0.2], [0.05, 8.0], 0.3] 30 | batch_pad: 31 | 32 | model: 33 | pretrained: 34 | teacher: 35 | score_threshold: 0.7 36 | backbone: 37 | name: "resnet50" 38 | pretrained: True 39 | trainable_layers: 3 40 | rpn: 41 | pre_nms_topk_train: 2000 42 | post_nms_topk_train: 1000 43 | pre_nms_topk_test: 1000 44 | post_nms_topk_test: 1000 45 | nms_threshold: 0.7 46 | roihead: 47 | num_classes: 80 48 | temprature: 3.5 49 | loc_t: 0.05 50 | score_threshold: 0.05 51 | nms_threshold: 0.5 52 | max_detections: 100 53 | 54 | solver: 55 | burnin_iter_num: 2000 56 | iter_num: 180000 57 | ckp_interval: 5000 58 | lr: 0.01 59 | weight_decay: 0.0001 60 | momentum: 0.9 61 | lr_steps: [179990, 179995] 62 | lr_gamma: 0.1 63 | warmup: True 64 | warmup_gamma: 0.001 65 | warmup_steps: 1000 66 | ema_momentum: 0.9996 67 | u_loss_weight: 4 68 | 69 | distributed: 70 | use: False 71 | -------------------------------------------------------------------------------- /segmentation/gvcore/optim/utils.py: -------------------------------------------------------------------------------- 1 | from typing import Dict, List 2 | import torch 3 | 4 | 5 | def sum_loss_list(losses_list): 6 | summed_loss = torch.zeros((1,), device="cuda") 7 | for loss in losses_list: 8 | if loss is not None: 9 | summed_loss += loss 10 | return summed_loss 11 | 12 | 13 | def sum_loss_dict(losses_dict): 14 | return sum_loss_list(list(losses_dict.values())) 15 | 16 | 17 | def sum_loss(losses, weight=None): 18 | if weight is not None: 19 | losses = weight_loss(losses, weight) 20 | 21 | if isinstance(losses, dict): 22 | return sum_loss_dict(losses) 23 | elif isinstance(losses, list): 24 | return sum_loss_list(losses) 25 | else: 26 | raise TypeError 27 | 28 | 29 | def merge_loss_dict(losses_dict_a, losses_dict_b, mean=False): 30 | losses_dict = {} 31 | a = 0.5 if mean else 1 32 | for name in losses_dict_a.keys(): 33 | losses_dict[name] = a * (losses_dict_a[name] + losses_dict_b[name]) 34 | return losses_dict 35 | 36 | 37 | def weight_loss(losses: Dict, weights: Dict): 38 | assert isinstance(losses, dict) == isinstance(weights, dict), "losses and weights must be both dict or both list" 39 | assert isinstance(losses, list) == isinstance(losses, list), "losses and weights must be both dict or both list" 40 | 41 | if isinstance(losses, dict): 42 | for name in losses.keys(): 43 | losses[name] *= weights.get(name, 1) 44 | else: 45 | assert len(losses) == len(weights), "losses and weights must be same length" 46 | for i, (loss, weight) in enumerate(zip(losses, weights)): 47 | losses[i] = loss * weight 48 | return losses 49 | -------------------------------------------------------------------------------- /segmentation/seg/config/ts_deeplabv3p/voc_aug.yaml: -------------------------------------------------------------------------------- 1 | _base: "./base.yaml" 2 | 3 | log: 4 | comment: "voc" 5 | 6 | data: 7 | train_l: 8 | dataset: "VOC" 9 | subset: "l.1.128.fold1" 10 | prefetch: False 11 | infinity_sampler: True 12 | root: "./data/VOCSegAug" 13 | num_workers: 4 14 | batch_size: 4 15 | shuffle: True 16 | img_transforms: 17 | batch_transforms: 18 | batch_apply: 19 | resize: 20 | scale: [0.5, 0.75, 1, 1.5, 1.75, 2.0] 21 | random_crop: 22 | - [512, 512] 23 | random_horizontal_flip: 24 | train_u: 25 | dataset: "VOC" 26 | subset: "u.1.128.fold1" 27 | prefetch: False 28 | infinity_sampler: True 29 | root: "./data/VOCSegAug" 30 | num_workers: 4 31 | batch_size: 4 32 | shuffle: True 33 | img_transforms: 34 | batch_transforms: 35 | batch_apply: 36 | resize: 37 | scale: [0.5, 0.75, 1, 1.5, 1.75, 2.0] 38 | random_crop: 39 | - [512, 512] 40 | random_horizontal_flip: 41 | test: 42 | dataset: "VOC" 43 | subset: "val" 44 | prefetch: False 45 | infinity_sampler: False 46 | root: "./data/VOCSegAug" 47 | num_workers: 1 48 | batch_size: 1 49 | shuffle: False 50 | img_transforms: 51 | batch_transforms: 52 | batch_apply: 53 | resize: 54 | min_size: [512, 512] 55 | max_size: 2048 56 | mode: "choice" 57 | process_label: False 58 | normalize: 59 | - [0.485, 0.456, 0.406] 60 | - [0.229, 0.224, 0.225] 61 | - True 62 | - False 63 | batch_pad: 64 | stride: 1 65 | 66 | model: 67 | num_classes: 21 68 | -------------------------------------------------------------------------------- /segmentation/seg/config/vc_deeplabv3p/voc_aug.yaml: -------------------------------------------------------------------------------- 1 | _base: "./base.yaml" 2 | 3 | log: 4 | comment: "voc" 5 | 6 | data: 7 | train_l: 8 | dataset: "VOC" 9 | subset: "l.1.128.fold3" 10 | prefetch: False 11 | infinity_sampler: True 12 | root: "./data/VOCSegAug" 13 | num_workers: 4 14 | batch_size: 4 15 | shuffle: True 16 | img_transforms: 17 | batch_transforms: 18 | batch_apply: 19 | resize: 20 | scale: [0.5, 0.75, 1, 1.5, 1.75, 2.0] 21 | random_crop: 22 | - [512, 512] 23 | random_horizontal_flip: 24 | train_u: 25 | dataset: "VOC" 26 | subset: "u.1.128.fold3" 27 | prefetch: False 28 | infinity_sampler: True 29 | root: "./data/VOCSegAug" 30 | num_workers: 4 31 | batch_size: 4 32 | shuffle: True 33 | img_transforms: 34 | batch_transforms: 35 | batch_apply: 36 | resize: 37 | scale: [0.5, 0.75, 1, 1.5, 1.75, 2.0] 38 | random_crop: 39 | - [512, 512] 40 | random_horizontal_flip: 41 | test: 42 | dataset: "VOC" 43 | subset: "val" 44 | prefetch: False 45 | infinity_sampler: False 46 | root: "./data/VOCSegAug" 47 | num_workers: 1 48 | batch_size: 1 49 | shuffle: False 50 | img_transforms: 51 | batch_transforms: 52 | batch_apply: 53 | resize: 54 | min_size: [512, 512] 55 | max_size: 2048 56 | mode: "choice" 57 | process_label: False 58 | normalize: 59 | - [0.485, 0.456, 0.406] 60 | - [0.229, 0.224, 0.225] 61 | - True 62 | - False 63 | batch_pad: 64 | stride: 1 65 | 66 | model: 67 | num_classes: 21 68 | teacher: 69 | low_score_threshold: 0.6 70 | -------------------------------------------------------------------------------- /detection/det/model/head/fasterrcnn.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | from gvcore.model.utils import xavier_fill 3 | 4 | 5 | class FasterRCNNHead(nn.Sequential): 6 | def __init__(self, fc_in_channels=12544, fc_inner_channels=1024, fc_nums=2): 7 | super(FasterRCNNHead, self).__init__() 8 | for i in range(fc_nums): 9 | if i == 0: 10 | self.add_module("flatten", nn.Flatten()) 11 | fc = nn.Linear(fc_in_channels, fc_inner_channels) 12 | self.add_module("fc{}".format(i + 1), fc) 13 | self.add_module("fc_relu{}".format(i + 1), nn.ReLU()) 14 | fc_in_channels = fc_inner_channels 15 | 16 | self._init_modules() 17 | 18 | def _init_modules(self): 19 | for layer in self: 20 | if isinstance(layer, nn.Linear): 21 | xavier_fill(layer) 22 | 23 | def forward(self, x): 24 | for layer in self: 25 | x = layer(x) 26 | return x 27 | 28 | 29 | class FastRCNNOutputLayers(nn.Module): 30 | def __init__(self, in_channels=1024, num_classes=80, cls_agnostic_bbox_reg=False): 31 | super(FastRCNNOutputLayers, self).__init__() 32 | 33 | self.cls_score = nn.Linear(in_channels, num_classes + 1) 34 | num_bbox_reg_classes = 1 if cls_agnostic_bbox_reg else num_classes 35 | self.bbox_pred = nn.Linear(in_channels, num_bbox_reg_classes * 4) 36 | 37 | self._init_modules() 38 | 39 | def _init_modules(self): 40 | nn.init.normal_(self.cls_score.weight, std=0.01) 41 | nn.init.normal_(self.bbox_pred.weight, std=0.001) 42 | for layer in [self.cls_score, self.bbox_pred]: 43 | nn.init.constant_(layer.bias, 0) 44 | 45 | def forward(self, x): 46 | return self.cls_score(x), self.bbox_pred(x) 47 | -------------------------------------------------------------------------------- /detection/det/operators/virtual_category.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.nn.parallel import DistributedDataParallel 3 | 4 | import gvcore.utils.distributed as dist_utils 5 | from model.detector.virtual_category import VCDetector, VCTSDetector 6 | from operators.unbiased_teacher import UnbiasedTeacherOpt 7 | 8 | from gvcore.operators import OPERATOR_REGISTRY 9 | 10 | 11 | @OPERATOR_REGISTRY.register("virtual_category") 12 | class VirtualCategoryOpt(UnbiasedTeacherOpt): 13 | def __init__(self, cfg): 14 | super(VirtualCategoryOpt, self).__init__(cfg) 15 | 16 | def build_model(self): 17 | model = VCTSDetector(self.cfg, VCDetector) 18 | model.to(self.device) 19 | 20 | if dist_utils.is_distributed(): 21 | model = DistributedDataParallel(model) 22 | self.model = model 23 | 24 | def resume_ckp(self, test_mode=False): 25 | super(VirtualCategoryOpt, self).resume_ckp(test_mode=test_mode) 26 | if self.cfg.resume is None: 27 | return 28 | if not test_mode: 29 | ckp = torch.load(self.cfg.resume, map_location=self.device) 30 | if "history" in ckp: 31 | history = ckp["history"] 32 | if dist_utils.is_distributed(): 33 | self.model.module.history.dict = history 34 | else: 35 | self.model.history.dict = history 36 | 37 | def save_ckp(self, name=None): 38 | self.checkpointer.save( 39 | name=name if name is not None else self.cur_step, 40 | model=self.model, 41 | optimizer=self.optimizer, 42 | lr_scheduler=self.lr_scheduler, 43 | step=self.cur_step, 44 | history=self.model.module.history.dict if dist_utils.is_distributed() else self.model.history.dict, 45 | ) 46 | -------------------------------------------------------------------------------- /detection/det/config/unbiased_teacher/voc.yaml: -------------------------------------------------------------------------------- 1 | _base: "./base.yaml" 2 | 3 | log: 4 | comment: "voc" 5 | 6 | evaluator: "voc" 7 | 8 | data: 9 | train_l: 10 | dataset: "VOC" 11 | prefetch: False 12 | infinity_sampler: True 13 | grouped_sampler: True 14 | cache: "./cache/voc07_trainval.dict" 15 | num_workers: 4 16 | batch_size: 4 17 | shuffle: True 18 | img_transforms: 19 | batch_transforms: 20 | batch_apply: 21 | random_horizontal_flip: 22 | resize: 23 | - [480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800] 24 | - 1333 25 | train_u: 26 | dataset: "VOC" 27 | prefetch: False 28 | infinity_sampler: True 29 | grouped_sampler: True 30 | cache: "./cache/voc12_trainval.dict" 31 | num_workers: 4 32 | batch_size: 4 33 | shuffle: True 34 | img_transforms: 35 | batch_transforms: 36 | batch_apply: 37 | random_horizontal_flip: 38 | resize: 39 | - [480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800] 40 | - 1333 41 | test: 42 | dataset: "VOC" 43 | prefetch: False 44 | infinity_sampler: False 45 | grouped_sampler: False 46 | cache: "./cache/voc07_test.dict" 47 | json_file: "./cache/voc07_test_cocostyle.json" 48 | num_workers: 4 49 | batch_size: 1 50 | shuffle: False 51 | img_transforms: 52 | batch_transforms: 53 | batch_apply: 54 | resize: 55 | - 600 56 | - 1000 57 | normalize: 58 | batch_pad: 59 | 60 | model: 61 | roihead: 62 | num_classes: 20 63 | 64 | solver: 65 | burnin_iter_num: 18000 66 | iter_num: 72000 67 | ckp_interval: 5000 68 | lr: 0.01 69 | weight_decay: 0.0001 70 | momentum: 0.9 71 | lr_steps: [71990, 71995] 72 | lr_gamma: 0.1 73 | warmup: True 74 | warmup_gamma: 0.001 75 | warmup_steps: 1000 76 | -------------------------------------------------------------------------------- /detection/det/config/virtual_category/voc.yaml: -------------------------------------------------------------------------------- 1 | _base: "./base.yaml" 2 | 3 | log: 4 | comment: "voc" 5 | 6 | evaluator: "voc" 7 | 8 | data: 9 | train_l: 10 | dataset: "VOC" 11 | prefetch: False 12 | infinity_sampler: True 13 | grouped_sampler: True 14 | cache: "./cache/voc07_trainval.dict" 15 | num_workers: 4 16 | batch_size: 4 17 | shuffle: True 18 | img_transforms: 19 | batch_transforms: 20 | batch_apply: 21 | random_horizontal_flip: 22 | resize: 23 | - [400, 1200] 24 | - 1333 25 | - "range" 26 | train_u: 27 | dataset: "VOC" 28 | prefetch: False 29 | infinity_sampler: True 30 | grouped_sampler: True 31 | cache: "./cache/voc12_trainval.dict" 32 | num_workers: 4 33 | batch_size: 4 34 | shuffle: True 35 | img_transforms: 36 | batch_transforms: 37 | batch_apply: 38 | random_horizontal_flip: 39 | resize: 40 | - [400, 1200] 41 | - 1333 42 | - "range" 43 | test: 44 | dataset: "VOC" 45 | prefetch: False 46 | infinity_sampler: False 47 | grouped_sampler: False 48 | cache: "./cache/voc07_test.dict" 49 | json_file: "./cache/voc07_test_cocostyle.json" 50 | num_workers: 4 51 | batch_size: 1 52 | shuffle: False 53 | img_transforms: 54 | batch_transforms: 55 | batch_apply: 56 | resize: 57 | - 600 58 | - 1000 59 | normalize: 60 | batch_pad: 61 | model: 62 | roihead: 63 | num_classes: 20 64 | 65 | solver: 66 | burnin_iter_num: 100000 67 | iter_num: 180000 68 | ckp_interval: 5000 69 | lr: 0.01 70 | weight_decay: 0.0001 71 | momentum: 0.9 72 | lr_steps: [179990, 179995] 73 | lr_gamma: 0.1 74 | warmup: True 75 | warmup_gamma: 0.001 76 | warmup_steps: 1000 77 | ema_momentum: 0.999 78 | u_loss_weight: 2 79 | -------------------------------------------------------------------------------- /detection/det/config/virtual_category_cross/voc.yaml: -------------------------------------------------------------------------------- 1 | _base: "./base.yaml" 2 | 3 | log: 4 | comment: "voc" 5 | 6 | evaluator: "voc" 7 | 8 | data: 9 | train_l: 10 | dataset: "VOC" 11 | prefetch: False 12 | infinity_sampler: True 13 | grouped_sampler: True 14 | cache: "./cache/voc07_trainval.dict" 15 | num_workers: 4 16 | batch_size: 2 17 | shuffle: True 18 | img_transforms: 19 | batch_transforms: 20 | batch_apply: 21 | random_horizontal_flip: 22 | resize: 23 | - [480, 800] 24 | - 1333 25 | - "range" 26 | train_u: 27 | dataset: "VOC" 28 | prefetch: False 29 | infinity_sampler: True 30 | grouped_sampler: True 31 | cache: "./cache/voc12_trainval.dict" 32 | num_workers: 4 33 | batch_size: 4 34 | shuffle: True 35 | img_transforms: 36 | batch_transforms: 37 | batch_apply: 38 | random_horizontal_flip: 39 | resize: 40 | - [480, 800] 41 | - 1333 42 | - "range" 43 | test: 44 | dataset: "VOC" 45 | prefetch: False 46 | infinity_sampler: False 47 | grouped_sampler: False 48 | cache: "./cache/voc07_test.dict" 49 | json_file: "./cache/voc07_test_cocostyle.json" 50 | num_workers: 4 51 | batch_size: 1 52 | shuffle: False 53 | img_transforms: 54 | batch_transforms: 55 | batch_apply: 56 | resize: 57 | - 600 58 | - 1000 59 | normalize: 60 | batch_pad: 61 | model: 62 | roihead: 63 | num_classes: 20 64 | temprature: 4 65 | 66 | solver: 67 | burnin_iter_num: 30000 68 | iter_num: 180000 69 | ckp_interval: 5000 70 | lr: 0.02 71 | weight_decay: 0.0001 72 | momentum: 0.9 73 | lr_steps: [179990, 179995] 74 | lr_gamma: 0.1 75 | warmup: True 76 | warmup_gamma: 0.001 77 | warmup_steps: 1000 78 | ema_momentum: 0.999 79 | u_loss_weight: 1 80 | -------------------------------------------------------------------------------- /detection/det/model/utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | @torch.no_grad() 5 | def match_label(boxes, labels, matcher): 6 | matches = [] 7 | match_flags = [] 8 | if isinstance(boxes, list): 9 | for b, label in zip(boxes, labels): 10 | label_box = label[:, :4] 11 | match, match_flag = matcher(b, label_box) 12 | matches.append(match) 13 | match_flags.append(match_flag) 14 | else: 15 | for label in labels: 16 | label_box = label[:, :4] 17 | match, match_flag = matcher(boxes, label_box) 18 | matches.append(match) 19 | match_flags.append(match_flag) 20 | return matches, match_flags 21 | 22 | 23 | @torch.no_grad() 24 | def sample_pos_neg(match_flags, sample_num, pos_fraction): 25 | for i, match_flag in enumerate(match_flags): 26 | positive = torch.nonzero(match_flag > 0, as_tuple=True)[0] 27 | negative = torch.nonzero(match_flag == 0, as_tuple=True)[0] 28 | 29 | num_pos = int(sample_num * pos_fraction) 30 | # protect against not enough positive examples 31 | num_pos = min(positive.numel(), num_pos) 32 | num_neg = sample_num - num_pos 33 | # protect against not enough negative examples 34 | num_neg = min(negative.numel(), num_neg) 35 | 36 | # randomly select positive and negative examples 37 | perm1 = torch.randperm(positive.numel(), device=positive.device)[:num_pos] 38 | perm2 = torch.randperm(negative.numel(), device=negative.device)[:num_neg] 39 | 40 | pos_idx = positive[perm1] 41 | neg_idx = negative[perm2] 42 | 43 | match_flag.fill_(-1) 44 | match_flag.scatter_(0, pos_idx, 1) 45 | match_flag.scatter_(0, neg_idx, 0) 46 | 47 | match_flags[i] = match_flag 48 | return match_flags 49 | 50 | 51 | def freeze(model, if_freeze): 52 | model.requires_grad_(not if_freeze) 53 | -------------------------------------------------------------------------------- /detection/det/config/virtual_category_cross/voc_cocox.yaml: -------------------------------------------------------------------------------- 1 | _base: "./base.yaml" 2 | 3 | log: 4 | comment: "voccocox" 5 | 6 | evaluator: "voc" 7 | 8 | data: 9 | train_l: 10 | dataset: "VOC" 11 | prefetch: False 12 | infinity_sampler: True 13 | grouped_sampler: True 14 | cache: "./cache/voc07_trainval.dict" 15 | num_workers: 4 16 | batch_size: 2 17 | shuffle: True 18 | img_transforms: 19 | batch_transforms: 20 | batch_apply: 21 | random_horizontal_flip: 22 | resize: 23 | - [480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800] 24 | - 1333 25 | train_u: 26 | dataset: "VOC" 27 | prefetch: False 28 | infinity_sampler: True 29 | grouped_sampler: True 30 | cache: "./cache/coco17_train_voccls_and_voc12_trainval.dict" 31 | num_workers: 4 32 | batch_size: 2 33 | shuffle: True 34 | img_transforms: 35 | batch_transforms: 36 | batch_apply: 37 | random_horizontal_flip: 38 | resize: 39 | - [480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800] 40 | - 1333 41 | test: 42 | dataset: "VOC" 43 | prefetch: False 44 | infinity_sampler: False 45 | grouped_sampler: False 46 | cache: "./cache/voc07_test.dict" 47 | json_file: "./cache/voc07_test_cocostyle.json" 48 | num_workers: 4 49 | batch_size: 1 50 | shuffle: False 51 | img_transforms: 52 | batch_transforms: 53 | batch_apply: 54 | resize: 55 | - 600 56 | - 1000 57 | normalize: 58 | batch_pad: 59 | 60 | model: 61 | roihead: 62 | num_classes: 20 63 | 64 | solver: 65 | burnin_iter_num: 30000 66 | iter_num: 180000 67 | ckp_interval: 5000 68 | lr: 0.01 69 | weight_decay: 0.0001 70 | momentum: 0.9 71 | lr_steps: [179990, 179995] 72 | lr_gamma: 0.1 73 | warmup: False 74 | warmup_gamma: 0.001 75 | warmup_steps: 1000 76 | ema_momentum: 0.999 77 | u_loss_weight: 2 78 | -------------------------------------------------------------------------------- /detection/gvcore/dataset/build.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from gvcore.dataset.transforms import parse_transform_config 4 | from gvcore.dataset.dataloader import BasicDataloader 5 | from gvcore.dataset.sampler import GroupedBatchSampler, InfiniteSampler, SequentialSampler 6 | from gvcore.dataset import DATASET_REGISTRY 7 | 8 | 9 | def build_dataloader(cfg, subset="train"): 10 | # 0. Load data cache or path 11 | data_list_cache, data_list_path = None, None 12 | if cfg.data[subset].cache is not None: 13 | data_list_cache = torch.load(cfg.data[subset].cache) 14 | elif cfg.data[subset].path is not None: 15 | data_list_path = cfg.data[subset].path 16 | else: 17 | raise ValueError("No data cache or data path specified.") 18 | 19 | datacfg = cfg.data[subset] 20 | 21 | # 1. Construct transforms 22 | img_transforms = parse_transform_config(datacfg.img_transforms) 23 | batch_transforms = parse_transform_config(datacfg.batch_transforms) 24 | 25 | # 2. Construct dataset. 26 | dataset = DATASET_REGISTRY[datacfg.dataset](data_list_cache, data_list_path, img_transforms, subset) 27 | 28 | # 3. Construct sampler 29 | if not datacfg.infinity_sampler: 30 | data_sampler = SequentialSampler(size=len(dataset)) 31 | else: 32 | data_sampler = InfiniteSampler(size=len(dataset), seed=cfg.seed) 33 | 34 | if datacfg.grouped_sampler: 35 | batch_sampler = GroupedBatchSampler(data_sampler, group_ids=dataset.group, batch_size=datacfg.batch_size) 36 | else: 37 | batch_sampler = None 38 | 39 | # 4. Construct dataloader 40 | data_loader = BasicDataloader( 41 | dataset, 42 | batch_size=datacfg.batch_size, 43 | sampler=data_sampler, 44 | batch_sampler=batch_sampler, 45 | num_workers=datacfg.num_workers, 46 | pin_memory=True, 47 | prefetch=datacfg.prefetch, 48 | batch_transforms=batch_transforms, 49 | drop_last=(subset != "test"), 50 | ) 51 | 52 | return data_loader 53 | -------------------------------------------------------------------------------- /segmentation/seg/model/segmentor/deeplabv3p_alt.py: -------------------------------------------------------------------------------- 1 | from addict import Dict 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | 6 | from gvcore.utils.types import TDataList 7 | from gvcore.model import GenericModule 8 | 9 | from model.backbone import make_backbone 10 | from model.head.aspp_alt import ASPPV3pAlt 11 | 12 | 13 | class DeeplabV3pAlt(GenericModule): 14 | def __init__(self, cfg): 15 | super(DeeplabV3pAlt, self).__init__(cfg) 16 | 17 | # Components ======== 18 | self.backbone = make_backbone(cfg.model.backbone) 19 | 20 | self.decoder = ASPPV3pAlt(cfg.model.aspp) 21 | 22 | self.dropout = nn.Dropout2d(cfg.model.dropout) 23 | self.classifier = nn.Conv2d(self.decoder.inner_channels, cfg.model.num_classes, 1) 24 | 25 | def forward_generic(self, data_list: TDataList) -> torch.Tensor: 26 | imgs = torch.stack([data.img for data in data_list], dim=0) 27 | 28 | feat_list = self.backbone(imgs) 29 | feats = self.decoder(feats=feat_list[-1], lowlevel_feats=feat_list[0]) 30 | 31 | logits = self.classifier(self.dropout(feats)) 32 | logits = F.interpolate(logits, size=imgs.shape[2:], mode="bilinear", align_corners=True) 33 | 34 | return logits 35 | 36 | def forward_train(self, data_list: TDataList) -> Dict: 37 | labels = torch.stack([data.label for data in data_list], dim=0) 38 | logits = self.forward_generic(data_list) 39 | return self.get_losses(logits, labels) 40 | 41 | def get_losses(self, logits: torch.Tensor, labels: torch.Tensor) -> Dict: 42 | loss_dict = Dict() 43 | loss = F.cross_entropy(logits, labels.long().squeeze(1), ignore_index=255) 44 | 45 | loss_dict.loss = loss 46 | 47 | return loss_dict 48 | 49 | def forward_eval(self, data_list: TDataList): 50 | logits = self.forward_generic(data_list) 51 | probs = F.softmax(logits, dim=1) 52 | preds = torch.argmax(probs, dim=1) 53 | 54 | return Dict(pred_list=preds) 55 | 56 | -------------------------------------------------------------------------------- /detection/det/operators/fasterrcnn.py: -------------------------------------------------------------------------------- 1 | from torch.nn.parallel import DistributedDataParallel 2 | from torchvision.ops import clip_boxes_to_image 3 | 4 | import gvcore.utils.distributed as dist_utils 5 | from gvcore.utils.misc import sum_loss 6 | from gvcore.operators import OPERATOR_REGISTRY, GenericOpt 7 | 8 | from model.detector.fasterrcnn import FasterRCNN 9 | 10 | 11 | @OPERATOR_REGISTRY.register("fasterrcnn") 12 | class FasterRCNNOpt(GenericOpt): 13 | def __init__(self, cfg): 14 | super(FasterRCNNOpt, self).__init__(cfg) 15 | 16 | def build_model(self): 17 | model = FasterRCNN(self.cfg) 18 | model.to(self.device) 19 | if dist_utils.is_distributed(): 20 | model = DistributedDataParallel(model) 21 | self.model = model 22 | 23 | def train_run_step(self): 24 | self.optimizer.zero_grad() 25 | data_list = self.train_loader.get_batch("cuda") 26 | loss_dict = self.model(data_list) 27 | losses = sum_loss(loss_dict) 28 | losses.backward() 29 | self.optimizer.step() 30 | 31 | metrics = {"lr": self.lr_scheduler.get_last_lr()[0]} 32 | metrics.update(loss_dict) 33 | self.summary.update(metrics) 34 | 35 | self.lr_scheduler.step() 36 | 37 | def test_run_step(self): 38 | data_list = self.test_loader.get_batch("cuda") 39 | if data_list is None: 40 | return True 41 | 42 | detected_bboxes = self.model(data_list) 43 | 44 | for data, detected_bbox in zip(data_list, detected_bboxes): 45 | meta = data.meta 46 | scale_x, scale_y = ( 47 | 1.0 * meta.ori_size[1] / meta.cur_size[1], 48 | 1.0 * meta.ori_size[0] / meta.cur_size[0], 49 | ) 50 | detected_bbox[:, [0, 2]] *= scale_x 51 | detected_bbox[:, [1, 3]] *= scale_y 52 | detected_bbox[:, :4] = clip_boxes_to_image(detected_bbox[:, :4], meta.ori_size) 53 | 54 | self.evaluator.process(data.meta.id, detected_bbox) 55 | return False 56 | -------------------------------------------------------------------------------- /detection/gvcore/dataset/__init__.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torchvision.io as io 3 | from gvcore.utils.structure import GenericData 4 | from gvcore.utils.registry import Registry 5 | 6 | 7 | DATASET_REGISTRY = Registry() 8 | TRANSFORM_REGISTRY = Registry() 9 | 10 | 11 | class GenericDataset(torch.utils.data.Dataset): 12 | def __init__(self, cache, path, img_transforms, meta=None, subset=None): 13 | super(GenericDataset, self).__init__() 14 | assert cache is not None or path is not None, "No data cache or data path specified." 15 | 16 | self._cache = cache 17 | self._path = path 18 | 19 | self._data_ids = self._get_data_ids() 20 | 21 | self._img_transforms = img_transforms 22 | 23 | self._meta = meta 24 | 25 | def _get_data_ids(self): 26 | if self._cache is not None: 27 | return sorted(list(self._cache.keys())) 28 | else: 29 | raise NotImplementedError 30 | 31 | def __len__(self): 32 | return len(self._data_ids) 33 | 34 | def _load(self, item_idx): 35 | data_id = self._data_ids[item_idx] 36 | cache = self._cache[data_id] 37 | 38 | data = GenericData() 39 | 40 | data.set("img", io.read_image(cache["image_path"])) 41 | if data.img.shape[0] == 1: 42 | data.img = data.img.repeat(3, 1, 1) 43 | label = torch.from_numpy(cache["label"]) 44 | data.set("label", label) 45 | 46 | data.set("meta", GenericData()) 47 | data.meta.set("id", data_id) 48 | data.meta.set("image_path", cache["image_path"]) 49 | 50 | img_size = data.img.shape[1:] 51 | data.meta.set("ori_size", list(img_size)) 52 | data.meta.set("cur_size", list(img_size)) 53 | data.meta.set("flip", False) 54 | 55 | return data 56 | 57 | def __getitem__(self, item_idx): 58 | sample = self._load(item_idx) 59 | if self._img_transforms is not None: 60 | sample = self._img_transforms(sample) 61 | return sample 62 | 63 | @staticmethod 64 | def collate_fn(batch): 65 | return batch 66 | -------------------------------------------------------------------------------- /segmentation/seg/config/ts_deeplabv3p/base.yaml: -------------------------------------------------------------------------------- 1 | dev_mode: False 2 | seed: 669 3 | operator: "teacher_student" 4 | evaluator: "segmentation" 5 | 6 | log: 7 | prefix: "ts_deeplabv3p" 8 | comment: "semi" 9 | log_dir: "./log/" 10 | summary_interval: 20 11 | wandb: "SEG-VC" 12 | aim: 13 | 14 | data: 15 | weak_transforms: 16 | batch_apply: 17 | normalize: 18 | - [0.485, 0.456, 0.406] 19 | - [0.229, 0.224, 0.225] 20 | - True 21 | - False 22 | batch_pad: 23 | fixed_size: [512, 512] 24 | strong_transforms: 25 | batch_apply: 26 | random_color_jitter: 27 | - 0.8 28 | - 0.5 29 | - 0.5 30 | - 0.5 31 | - 0.25 32 | random_grayscale: 33 | random_gaussian_blur: 34 | normalize: 35 | - [0.485, 0.456, 0.406] 36 | - [0.229, 0.224, 0.225] 37 | - True 38 | - False 39 | k_random_erase: 40 | scale: [0.02, 0.4] 41 | ratio: [0.3, 3.3] 42 | p: 0.5 43 | batch_pad: 44 | fixed_size: [512, 512] 45 | 46 | model: 47 | backbone: 48 | type: "resnet" 49 | depth: 50 50 | strides: [1, 2, 2, 1] 51 | dilations: [1, 1, 2, 4] 52 | multi_grid: [2, 4, 8] 53 | contract_dilation: True 54 | deep_stem: True 55 | norm_layer: "SyncBN" 56 | pretrained: "./cache/resnet50c.pth" 57 | aspp: 58 | in_channels: 2048 59 | inner_channels: 256 60 | lowlevel_in_channels: 256 61 | lowlevel_inner_channels: 48 62 | dilations: [1, 6, 12, 18] 63 | norm_layer: "SyncBN" 64 | dropout: 0. 65 | teacher: 66 | score_threshold: 0.95 67 | 68 | solver: 69 | iter_num: 40000 70 | burnin_iter_num: 1000 71 | eval_interval: 1000 72 | ckp_interval: 5000 73 | ema_momentum: 0.9996 74 | ema_bn: True 75 | u_loss_weight: 1 76 | optimizer: 77 | type: "sgd" 78 | lr: 0.001 79 | weight_decay: 0.0001 80 | weight_decay_norm: 0 81 | weight_decay_bias: 0 82 | momentum: 0.9 83 | lr_scheduler: 84 | type: "poly" 85 | iter_num: 40000 86 | power: 0.9 87 | min_lr: 0.0001 88 | warmup: True 89 | warmup_step: 500 90 | warmup_gamma: 0.001 91 | -------------------------------------------------------------------------------- /segmentation/seg/model/segmentor/deeplabv3p.py: -------------------------------------------------------------------------------- 1 | from addict import Dict 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | 6 | from gvcore.utils.types import TDataList 7 | from gvcore.model import GenericModule 8 | 9 | from model.backbone import make_backbone 10 | from model.head.aspp import ASPPV3p 11 | 12 | 13 | class DeeplabV3p(GenericModule): 14 | def __init__(self, cfg): 15 | super(DeeplabV3p, self).__init__(cfg) 16 | 17 | # Components ======== 18 | self.backbone = make_backbone(cfg.model.backbone) 19 | 20 | self.decoder = ASPPV3p(cfg.model.aspp) 21 | 22 | self.dropout = nn.Dropout2d(cfg.model.dropout) 23 | self.classifier = nn.Conv2d(self.decoder.inner_channels, cfg.model.num_classes, 1) 24 | 25 | def forward_generic(self, data_list: TDataList) -> torch.Tensor: 26 | imgs = torch.stack([data.img for data in data_list], dim=0) 27 | 28 | feat_list = self.backbone(imgs) 29 | feats = self.decoder(feats=feat_list[-1], lowlevel_feats=feat_list[0]) 30 | 31 | logits = self.classifier(self.dropout(feats)) 32 | logits = F.interpolate(logits, size=imgs.shape[2:], mode="bilinear", align_corners=False) 33 | 34 | return logits 35 | 36 | def forward_train(self, data_list: TDataList) -> Dict: 37 | labels = torch.stack([data.label for data in data_list], dim=0) 38 | logits = self.forward_generic(data_list) 39 | return self.get_losses(logits, labels) 40 | 41 | def get_losses(self, logits: torch.Tensor, labels: torch.Tensor) -> Dict: 42 | loss_dict = Dict() 43 | loss = F.cross_entropy(logits, labels.long().squeeze(1), ignore_index=255, reduction="none") 44 | loss_mask = labels.squeeze(1) != 255 45 | loss = loss * loss_mask.float() 46 | loss = loss.sum() / (loss_mask.float().sum() + 1e-10) 47 | 48 | loss_dict.loss = loss 49 | 50 | return loss_dict 51 | 52 | def forward_eval(self, data_list: TDataList): 53 | logits = self.forward_generic(data_list) 54 | probs = F.softmax(logits, dim=1) 55 | preds = torch.argmax(probs, dim=1) 56 | 57 | return Dict(pred_list=preds) 58 | -------------------------------------------------------------------------------- /detection/det/utils/box.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torchvision.ops import clip_boxes_to_image 3 | 4 | 5 | def pairwise_intersection(boxes1, boxes2) -> torch.Tensor: 6 | """ 7 | Returns: 8 | Tensor: intersection, sized [N,M]. 9 | """ 10 | width_height = torch.min(boxes1[:, None, 2:], boxes2[:, 2:]) - torch.max( 11 | boxes1[:, None, :2], boxes2[:, :2] 12 | ) # [N,M,2] 13 | 14 | width_height.clamp_(min=0) # [N,M,2] 15 | intersection = width_height.prod(dim=2) # [N,M] 16 | return intersection 17 | 18 | 19 | def pairwise_iou(boxes1, boxes2) -> torch.Tensor: 20 | """ 21 | Returns: 22 | Tensor: IoU, sized [N,M]. 23 | """ 24 | area1 = (boxes1[:, 2] - boxes1[:, 0]) * (boxes1[:, 3] - boxes1[:, 1]) 25 | area2 = (boxes2[:, 2] - boxes2[:, 0]) * (boxes2[:, 3] - boxes2[:, 1]) 26 | inter = pairwise_intersection(boxes1, boxes2) 27 | 28 | # handle empty boxes 29 | iou = torch.where( 30 | inter > 0, inter / (area1[:, None] + area2 - inter), torch.zeros(1, dtype=inter.dtype, device=inter.device) 31 | ) 32 | return iou 33 | 34 | 35 | def cxcywh2ltxywh(box): 36 | box = box.clone() 37 | box[:, :2] -= box[:, 2:4] / 2 38 | return box 39 | 40 | 41 | def xyxy2ltxywh(box): 42 | box = box.clone() 43 | box[:, 2:4] = box[:, 2:4] - box[:, :2] 44 | return box 45 | 46 | 47 | def xyxy2cxcywh(box): 48 | box = box.clone() 49 | box[:, 2:4] = box[:, 2:4] - box[:, :2] 50 | box[:, 0:2] += box[:, 2:4] / 2 51 | return box 52 | 53 | 54 | def nonempty(box, threshold=1e-5): 55 | widths = box[:, 2] - box[:, 0] 56 | heights = box[:, 3] - box[:, 1] 57 | keep = (widths > threshold) & (heights > threshold) 58 | return keep 59 | 60 | 61 | def rescale(box, meta): 62 | scale_x, scale_y = ( 63 | 1.0 * meta.ori_size[1] / (meta.cur_size[1] - meta.pad_size[2]), 64 | 1.0 * meta.ori_size[0] / (meta.cur_size[0] - meta.pad_size[3]), 65 | ) 66 | box[:, [0, 2]] *= scale_x 67 | box[:, [1, 3]] *= scale_y 68 | box[:, :4] = clip_boxes_to_image(box[:, :4], meta.ori_size) 69 | 70 | return box 71 | 72 | 73 | def area(box): 74 | return (box[:, 2] - box[:, 0]) * (box[:, 3] - box[:, 1]) 75 | -------------------------------------------------------------------------------- /segmentation/seg/config/vc_deeplabv3p/base.yaml: -------------------------------------------------------------------------------- 1 | dev_mode: False 2 | seed: 669 3 | operator: "virtual_category" 4 | evaluator: "segmentation" 5 | 6 | log: 7 | prefix: "vc_deeplabv3p" 8 | comment: "semi" 9 | log_dir: "./log/" 10 | summary_interval: 20 11 | wandb: "SEG-VC" 12 | aim: 13 | 14 | data: 15 | weak_transforms: 16 | batch_apply: 17 | normalize: 18 | - [0.485, 0.456, 0.406] 19 | - [0.229, 0.224, 0.225] 20 | - True 21 | - False 22 | batch_pad: 23 | fixed_size: [512, 512] 24 | strong_transforms: 25 | batch_apply: 26 | random_color_jitter: 27 | - 0.8 28 | - 0.5 29 | - 0.5 30 | - 0.5 31 | - 0.25 32 | random_grayscale: 33 | random_gaussian_blur: 34 | normalize: 35 | - [0.485, 0.456, 0.406] 36 | - [0.229, 0.224, 0.225] 37 | - True 38 | - False 39 | k_random_erase: 40 | scale: [0.02, 0.4] 41 | ratio: [0.3, 3.3] 42 | p: 0.5 43 | batch_pad: 44 | fixed_size: [512, 512] 45 | 46 | model: 47 | backbone: 48 | type: "resnet" 49 | depth: 50 50 | strides: [1, 2, 2, 1] 51 | dilations: [1, 1, 2, 4] 52 | multi_grid: [2, 4, 8] 53 | contract_dilation: True 54 | deep_stem: True 55 | norm_layer: "SyncBN" 56 | pretrained: "./cache/resnet50c.pth" 57 | aspp: 58 | in_channels: 2048 59 | inner_channels: 256 60 | lowlevel_in_channels: 256 61 | lowlevel_inner_channels: 48 62 | dilations: [1, 6, 12, 18] 63 | norm_layer: "SyncBN" 64 | dropout: 0. 65 | teacher: 66 | score_threshold: 0.95 67 | low_score_threshold: 0.8 68 | 69 | solver: 70 | iter_num: 40000 71 | burnin_iter_num: 1000 72 | eval_interval: 1000 73 | ckp_interval: 5000 74 | ema_momentum: 0.9996 75 | ema_bn: True 76 | u_loss_weight: 1 77 | vc_weight_norm: "adaptive_min" 78 | optimizer: 79 | type: "sgd" 80 | lr: 0.001 81 | weight_decay: 0.0001 82 | weight_decay_norm: 0 83 | weight_decay_bias: 0 84 | momentum: 0.9 85 | lr_scheduler: 86 | type: "poly" 87 | iter_num: 40000 88 | power: 0.9 89 | min_lr: 0.0001 90 | warmup: True 91 | warmup_step: 500 92 | warmup_gamma: 0.001 93 | -------------------------------------------------------------------------------- /detection/det/ext/fastevalapi/csrc/vision.cpp: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | 3 | #include 4 | #include "cocoeval/cocoeval.h" 5 | namespace fastevalapi { 6 | 7 | #if defined(WITH_CUDA) || defined(WITH_HIP) 8 | extern int get_cudart_version(); 9 | #endif 10 | 11 | std::string get_cuda_version() { 12 | #if defined(WITH_CUDA) || defined(WITH_HIP) 13 | std::ostringstream oss; 14 | 15 | #if defined(WITH_CUDA) 16 | oss << "CUDA "; 17 | #else 18 | oss << "HIP "; 19 | #endif 20 | 21 | // copied from 22 | // https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/cuda/detail/CUDAHooks.cpp#L231 23 | auto printCudaStyleVersion = [&](int v) { 24 | oss << (v / 1000) << "." << (v / 10 % 100); 25 | if (v % 10 != 0) { 26 | oss << "." << (v % 10); 27 | } 28 | }; 29 | printCudaStyleVersion(get_cudart_version()); 30 | return oss.str(); 31 | #else // neither CUDA nor HIP 32 | return std::string("not available"); 33 | #endif 34 | } 35 | 36 | // similar to 37 | // https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/Version.cpp 38 | std::string get_compiler_version() { 39 | std::ostringstream ss; 40 | #if defined(__GNUC__) 41 | #ifndef __clang__ 42 | 43 | #if ((__GNUC__ <= 4) && (__GNUC_MINOR__ <= 8)) 44 | #error "GCC >= 4.9 is required!" 45 | #endif 46 | 47 | { ss << "GCC " << __GNUC__ << "." << __GNUC_MINOR__; } 48 | #endif 49 | #endif 50 | 51 | #if defined(__clang_major__) 52 | { 53 | ss << "clang " << __clang_major__ << "." << __clang_minor__ << "." 54 | << __clang_patchlevel__; 55 | } 56 | #endif 57 | 58 | #if defined(_MSC_VER) 59 | { ss << "MSVC " << _MSC_FULL_VER; } 60 | #endif 61 | return ss.str(); 62 | } 63 | 64 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 65 | m.def("COCOevalAccumulate", &COCOeval::Accumulate, "COCOeval::Accumulate"); 66 | m.def( 67 | "COCOevalEvaluateImages", 68 | &COCOeval::EvaluateImages, 69 | "COCOeval::EvaluateImages"); 70 | pybind11::class_(m, "InstanceAnnotation") 71 | .def(pybind11::init()); 72 | pybind11::class_(m, "ImageEvaluation") 73 | .def(pybind11::init<>()); 74 | } 75 | } // namespace fastevalapi 76 | -------------------------------------------------------------------------------- /segmentation/seg/evaluator/api.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torchvision.transforms.functional as tvF 3 | 4 | 5 | def intersect_and_union(pred, label, num_classes, ignore_index, label_map=dict(), reduce_zero_label=False): 6 | """Calculate intersection and Union. 7 | 8 | Args: 9 | pred (ndarray | str): Prediction segmentation map 10 | or predict result filename. 11 | label (ndarray | str): Ground truth segmentation map 12 | or label filename. 13 | num_classes (int): Number of categories. 14 | ignore_index (int): Index that will be ignored in evaluation. 15 | label_map (dict): Mapping old labels to new labels. The parameter will 16 | work only when label is str. Default: dict(). 17 | reduce_zero_label (bool): Whether ignore zero label. The parameter will 18 | work only when label is str. Default: False. 19 | 20 | Returns: 21 | torch.Tensor: The intersection of prediction and ground truth 22 | histogram on all classes. 23 | torch.Tensor: The union of prediction and ground truth histogram on 24 | all classes. 25 | torch.Tensor: The prediction histogram on all classes. 26 | torch.Tensor: The ground truth histogram on all classes. 27 | """ 28 | if label.dim() == 3: 29 | label = label.squeeze(0) 30 | 31 | if pred.shape != label.shape: 32 | pred = tvF.resize(pred.unsqueeze(0), label.shape[-2:], interpolation=tvF.InterpolationMode.NEAREST).squeeze(0) 33 | 34 | if label_map is not None: 35 | for old_id, new_id in label_map.items(): 36 | label[label == old_id] = new_id 37 | if reduce_zero_label: 38 | label[label == 0] = 255 39 | label = label - 1 40 | label[label == 254] = 255 41 | 42 | mask = label != ignore_index 43 | pred = pred[mask] 44 | label = label[mask] 45 | 46 | intersect = pred[pred == label] 47 | area_intersect = torch.histc(intersect.float(), bins=(num_classes), min=0, max=num_classes - 1) 48 | area_pred = torch.histc(pred.float(), bins=(num_classes), min=0, max=num_classes - 1) 49 | area_label = torch.histc(label.float(), bins=(num_classes), min=0, max=num_classes - 1) 50 | area_union = area_pred + area_label - area_intersect 51 | return area_intersect, area_union, area_pred, area_label 52 | -------------------------------------------------------------------------------- /detection/det/utils/lr_scheduler.py: -------------------------------------------------------------------------------- 1 | import torch.optim.lr_scheduler as lr_s 2 | from bisect import bisect_right 3 | 4 | 5 | class PolyPolicy: 6 | def __init__(self, iter_num, power=0.9): 7 | self.power = power 8 | self.iter_num = iter_num 9 | 10 | def __call__(self, step): 11 | return (1 - step / self.iter_num) ** self.power 12 | 13 | 14 | class PolyLR(lr_s.LambdaLR): 15 | def __init__(self, optimizer, iter_num, power=0.9): 16 | self.power = power 17 | self.iter_num = iter_num 18 | poly = PolyPolicy(self.iter_num, self.power) 19 | super(PolyLR, self).__init__(optimizer, poly) 20 | 21 | 22 | class WarmupPolicy: 23 | def __init__(self, warmup_step, warmup_gamma): 24 | self.warmup_step = warmup_step 25 | self.warmup_gamma = warmup_gamma 26 | 27 | def __call__(self, step): 28 | if step >= self.warmup_step: 29 | return 1 30 | alpha = float(step) / self.warmup_step 31 | return self.warmup_gamma * (1 - alpha) + alpha 32 | 33 | 34 | class WarmupLR(lr_s.LambdaLR): 35 | def __init__(self, optimizer, warmup_step, warmup_gamma): 36 | policy = WarmupPolicy(warmup_step, warmup_gamma) 37 | super(WarmupLR, self).__init__(optimizer, policy) 38 | 39 | 40 | class MultiStepPolicy: 41 | def __init__(self, milestones, gamma, warmup=False, warmup_step=None, warmup_gamma=None): 42 | self.milestones = milestones 43 | self.gamma = gamma 44 | self.warmup = warmup 45 | if self.warmup: 46 | assert warmup_step is not None and warmup_gamma is not None, "Warmup policy need to set step and gamma!" 47 | self.warmup_policy = WarmupPolicy(warmup_step, warmup_gamma) 48 | self.warmup_step = warmup_step 49 | else: 50 | self.warmup_step = -1 51 | 52 | def __call__(self, step): 53 | if step <= self.warmup_step: 54 | return self.warmup_policy(step) 55 | else: 56 | return self.gamma ** bisect_right(self.milestones, step) 57 | 58 | 59 | class MultiStepLR(lr_s.LambdaLR): 60 | def __init__(self, optimizer, milestones, gamma=0.1, warmup=False, warmup_step=None, warmup_gamma=None): 61 | policy = MultiStepPolicy(milestones, gamma, warmup, warmup_step, warmup_gamma) 62 | super(MultiStepLR, self).__init__(optimizer, policy) 63 | -------------------------------------------------------------------------------- /detection/det/evaluator/voc.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Facebook, Inc. and its affiliates. 2 | 3 | from gvcore.evaluator import EVALUATOR_REGISTRY 4 | from utils.meta import VOCMeta 5 | from utils.box import xyxy2ltxywh 6 | from evaluator.coco import COCOEvaluator 7 | 8 | 9 | @EVALUATOR_REGISTRY.register("voc") 10 | class VOCEvaluator(COCOEvaluator): 11 | """ 12 | Evaluate AR for object proposals, AP for instance detection/segmentation, AP 13 | for keypoint detection outputs using COCO's metrics. 14 | See http://cocodataset.org/#detection-eval and 15 | http://cocodataset.org/#keypoints-eval to understand its metrics. 16 | 17 | COCO style for VOC 18 | """ 19 | 20 | def __init__(self, json_file, distributed=False, use_fast_impl=True): 21 | """ 22 | Args: 23 | distributed (True): if True, will collect results from all ranks and run evaluation 24 | in the main process. 25 | Otherwise, will only evaluate the results in the current process. 26 | use_fast_impl (bool): use a fast but **unofficial** implementation to compute AP. 27 | Although the results should be very close to the official implementation in COCO 28 | API, it is still recommended to compute results with the official API for use in 29 | papers. The faster implementation also uses more RAM. 30 | """ 31 | super(VOCEvaluator, self).__init__(json_file, distributed, use_fast_impl) 32 | self._meta = VOCMeta() 33 | 34 | def _per_cat_results(self, coco_eval, class_names=None): 35 | pass 36 | 37 | def output_to_coco_json(self, outputs, img_id): 38 | """ 39 | Dump an "Instances" object to a COCO-format json that's used for evaluation. 40 | 41 | Args: 42 | outputs (torch.Tensor): 43 | img_id (int): the image id 44 | 45 | Returns: 46 | list[dict]: list of json annotations in COCO format. 47 | """ 48 | results = [] 49 | outputs[:, :4] = xyxy2ltxywh(outputs[:, :4]) 50 | for output in outputs: 51 | result = { 52 | "image_id": img_id, 53 | "category_id": int(output[4]), 54 | "bbox": [x.item() for x in output[:4]], 55 | "score": output[5].item(), 56 | } 57 | results.append(result) 58 | return results 59 | -------------------------------------------------------------------------------- /detection/gvcore/utils/checkpoint.py: -------------------------------------------------------------------------------- 1 | from os.path import join 2 | import torch 3 | import torch.nn as nn 4 | from torch.nn.parallel import DistributedDataParallel 5 | 6 | from gvcore.utils.logger import logger 7 | from gvcore.utils.distributed import is_main_process 8 | 9 | 10 | class Checkpointer(object): 11 | def __init__(self, cfg): 12 | self.cfg = cfg 13 | self.dev_mode = cfg.dev_mode 14 | 15 | def save(self, name, **checkpointables): 16 | if not is_main_process() or self.dev_mode: 17 | return 18 | data = {} 19 | for key, obj in checkpointables.items(): 20 | if isinstance(obj, DistributedDataParallel): 21 | data[key] = obj.module.state_dict() 22 | elif isinstance(obj, (nn.Module, torch.optim.Optimizer, torch.optim.lr_scheduler._LRScheduler)): 23 | data[key] = obj.state_dict() 24 | else: 25 | data[key] = obj 26 | 27 | file_path = join(self.cfg.log.log_dir, f"{name}.pth") 28 | torch.save(data, file_path) 29 | logger.info(f"Saved checkpoint to {file_path}.") 30 | 31 | def load(self, ckp_path, **checkpointables): 32 | ckp_dict = torch.load(ckp_path, map_location="cpu") 33 | except_keys = [] 34 | for key, obj in checkpointables.items(): 35 | if isinstance(obj, DistributedDataParallel): 36 | load_obj = obj.module 37 | else: 38 | load_obj = obj 39 | if key in ckp_dict: 40 | if ckp_dict[key] is not None: 41 | if isinstance(load_obj, nn.Module): 42 | except_param_keys = load_obj.load_state_dict(ckp_dict[key], strict=False) 43 | logger.info("Unmatched params:" + str(except_param_keys)) 44 | elif isinstance(load_obj, (torch.optim.Optimizer, torch.optim.lr_scheduler._LRScheduler)): 45 | load_obj.load_state_dict(ckp_dict[key]) 46 | else: 47 | except_keys.append(key) 48 | else: 49 | except_keys.append(key) 50 | except_string = "" if len(except_keys) == 0 else f", except: {', '.join(except_keys)}" 51 | logger.info(f"Loaded checkpoint from {ckp_path}{except_string}.") 52 | 53 | if "step" in ckp_dict: 54 | return ckp_dict["step"] 55 | else: 56 | return 0 57 | -------------------------------------------------------------------------------- /detection/det/utils/metrics.py: -------------------------------------------------------------------------------- 1 | from utils.box import pairwise_iou 2 | import torch 3 | from gvcore.utils.types import TTensor, TTensorList, TTensorTuple 4 | 5 | 6 | def single_bbox_quality(pred: TTensor, label: TTensor, iou_t: int = 0.0) -> TTensorTuple: 7 | """Evaluate predicted bbox quality, return precision, recall, mean iou. 8 | """ 9 | if pred.numel() > 0: 10 | iou = pairwise_iou(torch.narrow(pred, dim=1, start=0, length=4), torch.narrow(label, dim=1, start=0, length=4)) 11 | 12 | pred_cls = pred[:, 4] 13 | label_cls = label[:, 4] 14 | 15 | # ====================== 16 | # Precision 17 | max_iou, max_idx = torch.max(iou, dim=1) 18 | 19 | mean_iou = max_iou.mean() 20 | 21 | matched_cls = label_cls[max_idx] 22 | matched_cls[max_iou < iou_t] = -1 23 | 24 | correct_flag = torch.eq(pred_cls, matched_cls) 25 | precision = correct_flag.float().sum() / matched_cls.numel() 26 | 27 | # ====================== 28 | # Recall 29 | max_iou, max_idx = torch.max(iou, dim=0) 30 | 31 | matched_cls = pred_cls[max_idx] 32 | matched_cls[max_iou < iou_t] = -1 33 | 34 | correct_flag = torch.eq(label_cls, matched_cls) 35 | recall = correct_flag.float().sum() / label_cls.numel() 36 | 37 | return precision, recall, mean_iou 38 | else: 39 | return ( 40 | torch.tensor([0,], device=pred.device), 41 | torch.tensor([0,], device=pred.device), 42 | torch.tensor([0,], device=pred.device), 43 | ) 44 | 45 | 46 | def bbox_quality(pred_list: TTensorList, label_list: TTensorList, iou_t: int = 0.0) -> TTensorTuple: 47 | """Evaluate predicted bbox quality, return precision, recall, mean iou. 48 | """ 49 | assert len(pred_list) == len(label_list), "Number of results and labels mismatch." 50 | assert len(pred_list) != 0, "Number of predictions cannot be zero." 51 | 52 | precision_list, recall_list, meaniou_list = [], [], [] 53 | 54 | for pred, label in zip(pred_list, label_list): 55 | p, r, miou = single_bbox_quality(pred, label, iou_t) 56 | precision_list.append(p) 57 | recall_list.append(r) 58 | meaniou_list.append(miou) 59 | 60 | return ( 61 | sum(precision_list) / len(precision_list), 62 | sum(recall_list) / len(recall_list), 63 | sum(meaniou_list) / len(meaniou_list), 64 | ) 65 | -------------------------------------------------------------------------------- /detection/gvcore/dataset/transforms/__init__.py: -------------------------------------------------------------------------------- 1 | from copy import deepcopy 2 | 3 | from gvcore.dataset import TRANSFORM_REGISTRY 4 | 5 | 6 | __all__ = ["parse_transform_config", "DataCopy", "Compose", "Transforms"] 7 | 8 | 9 | def parse_transform_config(cfg): 10 | transforms = [] 11 | if cfg is not None: 12 | for name, args in cfg.items(): 13 | if args is not None: 14 | if isinstance(args, (tuple, list)): 15 | transforms.append(TRANSFORM_REGISTRY[name](*args)) 16 | elif isinstance(args, dict): 17 | transforms.append(TRANSFORM_REGISTRY[name](**args)) 18 | else: 19 | raise ValueError(f"Argument for transform {name} is invalid, got {args}") 20 | else: 21 | transforms.append(TRANSFORM_REGISTRY[name]()) 22 | if len(transforms) > 1: 23 | transforms = Compose(transforms) 24 | else: 25 | transforms = transforms[0] 26 | else: 27 | transforms = None 28 | return transforms 29 | 30 | 31 | class DataCopy: 32 | """Copy data.""" 33 | 34 | def __init__(self, transforms): 35 | self.transforms = transforms 36 | 37 | def __call__(self, data): 38 | data = deepcopy(data) 39 | data = self.transforms(data) 40 | return data 41 | 42 | def __repr__(self): 43 | return "[Deep Copied] \n" + str(self.transforms) 44 | 45 | 46 | class Compose: 47 | """Composes several transforms together.""" 48 | 49 | def __init__(self, transforms): 50 | self.transforms = transforms 51 | 52 | def __call__(self, data): 53 | for t in self.transforms: 54 | data = t(data) 55 | return data 56 | 57 | def __repr__(self): 58 | format_string = self.__class__.__name__ + "(" 59 | for t in self.transforms: 60 | format_string += "\n" 61 | format_string += "{0}".format(t) 62 | format_string += "\n)" 63 | return format_string 64 | 65 | def __getitem__(self, item): 66 | return self.transforms[item] 67 | 68 | def __len__(self): 69 | return len(self.transforms) 70 | 71 | 72 | class Transforms: 73 | def __repr__(self): 74 | params = [] 75 | for k, v in self.__dict__.items(): 76 | params.append(f"{k}: {v}") 77 | params = ", ".join(params) 78 | return f"{self.__class__.__name__}({params})" 79 | 80 | -------------------------------------------------------------------------------- /detection/det/dataset/det_generic.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from tabulate import tabulate 3 | 4 | from gvcore.utils.logger import logger 5 | from gvcore.dataset import GenericDataset 6 | 7 | from utils.box import nonempty 8 | 9 | 10 | class DetGenericDataset(GenericDataset): 11 | def __init__(self, cache, path, img_transforms, meta=None, subset=None): 12 | super(DetGenericDataset, self).__init__(cache, path, img_transforms, meta) 13 | 14 | if subset != "test": 15 | self._filter_empty() 16 | self.group = self._generate_group() 17 | 18 | def _statics(self): 19 | """ 20 | To print the statistics for detection dataset such as MSCOCO. 21 | """ 22 | if self._meta is None: 23 | return 24 | 25 | _categories_count = len(self._meta) - 1 26 | 27 | categories_count = np.zeros(_categories_count) 28 | for data_id, data in self._cache.items(): 29 | category_label = data["label"][:, 4].astype(np.int64) 30 | category_label_count = np.bincount(category_label, minlength=_categories_count) 31 | categories_count += category_label_count 32 | table = [] 33 | row = [] 34 | for i, count in enumerate(categories_count): 35 | if i % 4 == 0: 36 | if len(row) > 0: 37 | table.append(row) 38 | row = [] 39 | row.append(self._meta(i)["name"]) 40 | row.append(count) 41 | 42 | logger.info( 43 | "\n" 44 | + tabulate( 45 | table, 46 | headers=["Category", "Count", "Category", "Count", "Category", "Count", "Category", "Count"], 47 | tablefmt="fancy_grid", 48 | ) 49 | ) 50 | 51 | def _filter_empty(self): 52 | """ 53 | To filter no label images and invalid labels 54 | """ 55 | for data_id, data in self._cache.items(): 56 | label = data["label"] 57 | keep = nonempty(label) 58 | data["label"] = label[keep] 59 | if data["label"].shape[0] == 0: 60 | self._data_ids.remove(data_id) 61 | 62 | def _generate_group(self): 63 | """ 64 | To generate group of images according to aspect ratios. 65 | """ 66 | group = [] 67 | for data_id in self._data_ids: 68 | data = self._cache[data_id] 69 | group.append(data["aspect_ratio"]) 70 | return group 71 | 72 | -------------------------------------------------------------------------------- /detection/det/model/head/unbiased_teacher/rpn.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as F 3 | 4 | from model.utils import match_label, sample_pos_neg 5 | from model.head.rpn import RPN as _RPN 6 | 7 | 8 | class RPN(_RPN): 9 | def __init__(self, cfg): 10 | super(RPN, self).__init__(cfg) 11 | self.rpn_positive_fraction = 0.25 12 | 13 | def forward_train(self, feats, img_sizes, labels, ignore_loc=False): 14 | self._forward(feats, img_sizes) 15 | return self.proposals, self.get_losses(labels, ignore_loc) 16 | 17 | def get_losses(self, labels, ignore_loc=False): 18 | anchors = torch.cat(self.anchors) 19 | matches, match_flags = match_label(anchors, labels, self.matcher) 20 | match_flags = sample_pos_neg(match_flags, self.rpn_num, self.rpn_positive_fraction) 21 | 22 | cls_labels, reg_labels = [], [] 23 | fg_flags, valid_flags = [], [] 24 | for label, match, match_flag in zip(labels, matches, match_flags): 25 | fg_flag = match_flag > 0 26 | valid_flag = match_flag >= 0 27 | 28 | fg_flags.append(fg_flag) 29 | valid_flags.append(valid_flag) 30 | 31 | cls_label = match_flag[valid_flag] 32 | cls_labels.append(cls_label) 33 | if not ignore_loc: 34 | if label.shape[0] == 0: 35 | matched_label = torch.zeros((1, 5), device=self.device)[match] 36 | else: 37 | matched_label = label[match] 38 | 39 | box_label = matched_label[fg_flag, :4] 40 | reg_label = self.box_coder.encode_single(box_label, anchors[fg_flag]) 41 | reg_labels.append(reg_label) 42 | 43 | cls_labels = torch.cat(cls_labels) 44 | fg_flags = torch.stack(fg_flags) 45 | valid_flags = torch.stack(valid_flags) 46 | 47 | normalizer = self.rpn_num * len(labels) 48 | 49 | rpn_cls_loss = F.binary_cross_entropy_with_logits( 50 | torch.cat(self.objectness_logits, dim=1)[valid_flags], cls_labels.float(), reduction="sum" 51 | ) 52 | if not ignore_loc: 53 | reg_labels = torch.cat(reg_labels) 54 | rpn_loc_loss = F.smooth_l1_loss( 55 | torch.cat(self.anchor_deltas, dim=1)[fg_flags], reg_labels, beta=0.0, reduction="sum" 56 | ) 57 | losses = {"rpn_cls": rpn_cls_loss / normalizer, "rpn_loc": rpn_loc_loss / normalizer} 58 | else: 59 | losses = {"rpn_cls": rpn_cls_loss / normalizer} 60 | return losses 61 | -------------------------------------------------------------------------------- /segmentation/seg/operators/vc.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import torch 3 | from torch.nn.parallel import DistributedDataParallel 4 | from addict import Dict 5 | 6 | from gvcore.optim.utils import sum_loss 7 | import gvcore.utils.distributed as dist_utils 8 | from gvcore.operators import OPERATOR_REGISTRY 9 | 10 | from gvcore.optim import make_lr_scheduler, make_optimizer 11 | from model.segmentor.deeplabv3p_alt import DeeplabV3pAlt 12 | from model.segmentor.vc import VCSegmentor 13 | from operators.ts import TeacherStudentOpt 14 | 15 | 16 | @OPERATOR_REGISTRY.register("virtual_category") 17 | class VirtualCategoryOpt(TeacherStudentOpt): 18 | def __init__(self, cfg): 19 | super(VirtualCategoryOpt, self).__init__(cfg) 20 | 21 | self.generator_optimizer = None 22 | 23 | def build_model(self): 24 | model = VCSegmentor(self.cfg, DeeplabV3pAlt) 25 | model.to(self.device) 26 | if dist_utils.is_distributed(): 27 | model = nn.SyncBatchNorm.convert_sync_batchnorm(model) 28 | model = DistributedDataParallel(model) 29 | 30 | self.model = model 31 | 32 | def build_optimizer(self): 33 | self.optimizer = make_optimizer( 34 | self.cfg.solver.optimizer, 35 | self.model.module.segmentor if dist_utils.is_distributed() else self.model.segmentor, 36 | ) 37 | self.generator_optimizer = make_optimizer( 38 | self.cfg.solver.optimizer, 39 | self.model.module.weight_generator if dist_utils.is_distributed() else self.model.weight_generator, 40 | ) 41 | 42 | self.lr_scheduler = make_lr_scheduler(self.cfg.solver.lr_scheduler, self.optimizer) 43 | 44 | def train_pre_step(self) -> Dict: 45 | self.generator_optimizer.zero_grad() 46 | return super().train_pre_step() 47 | 48 | def train_run_step(self, model_inputs: Dict, **kwargs) -> Dict: 49 | loss_dict, stat_dict = self.model(**model_inputs) 50 | 51 | losses = sum_loss(loss_dict.loss_l) + ( 52 | self.cfg.solver.u_loss_weight if self.cur_step > self.burnin_step else 0 53 | ) * sum_loss(loss_dict.loss_u) 54 | 55 | losses.backward() 56 | # gradient clipping 57 | torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=5.0) 58 | 59 | self.optimizer.step() 60 | self.generator_optimizer.step() 61 | self.lr_scheduler.step() 62 | 63 | if dist_utils.is_distributed(): 64 | self.model.module._momentum_update(self.cfg.solver.ema_momentum) 65 | else: 66 | self.model._momentum_update(self.cfg.solver.ema_momentum) 67 | 68 | return Dict(loss_dict=loss_dict, stat_dict=stat_dict) 69 | -------------------------------------------------------------------------------- /detection/gvcore/utils/misc.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import time 3 | 4 | 5 | def time_test(func): 6 | def wrapper(*args): 7 | st = time.time() 8 | out = func(*args) 9 | print(args[0].__class__.__name__, "{:.8f}".format(time.time() - st)) 10 | return out 11 | 12 | return wrapper 13 | 14 | 15 | def sum_loss_list(losses_list, weight): 16 | summed_loss = torch.zeros((1,), device="cuda") 17 | if weight is not None: 18 | for loss, w in zip(losses_list, weight): 19 | if loss is not None: 20 | summed_loss += loss * w 21 | else: 22 | for loss in losses_list: 23 | if loss is not None: 24 | summed_loss += loss 25 | return summed_loss 26 | 27 | 28 | def sum_loss_dict(losses_dict, weight): 29 | return sum_loss_list(list(losses_dict.values()), weight) 30 | 31 | 32 | def sum_loss(losses, weight=None): 33 | if isinstance(losses, dict): 34 | return sum_loss_dict(losses, weight) 35 | elif isinstance(losses, list): 36 | return sum_loss_list(losses, weight) 37 | else: 38 | raise TypeError 39 | 40 | 41 | def merge_loss_dict(losses_dict_a, losses_dict_b, mean=False): 42 | losses_dict = {} 43 | a = 0.5 if mean else 1 44 | for name in losses_dict_a.keys(): 45 | losses_dict[name] = a * (losses_dict_a[name] + losses_dict_b[name]) 46 | return losses_dict 47 | 48 | 49 | def weight_loss(losses, weights): 50 | for loss_name, loss_value in losses.items(): 51 | if loss_name in weights and loss_value is not None: 52 | losses[loss_name] *= weights[loss_name] 53 | return losses 54 | 55 | 56 | def attach_batch_idx(tensor_list): 57 | tensors = [] 58 | for i, tensor in enumerate(tensor_list): 59 | batch_idx = torch.ones((tensor.shape[0], 1), dtype=tensor.dtype, device=tensor.device) * i 60 | tensors.append(torch.cat((batch_idx, tensor), dim=1)) 61 | return torch.cat(tensors, dim=0) 62 | 63 | 64 | def print_tensor(x, device="cuda:0"): 65 | if isinstance(x, torch.Tensor): 66 | if str(x.device) == device: 67 | print(x) 68 | else: 69 | print(x) 70 | 71 | 72 | def slerp(low, high, val=0.5): 73 | omega = torch.arccos((low * high).sum(dim=1, keepdim=True)) 74 | so = torch.sin(omega) 75 | return torch.sin((1.0 - val) * omega) / so * low + torch.sin(val * omega) / so * high 76 | 77 | 78 | def entropy(x): 79 | p = torch.softmax(x, dim=1) 80 | entropy = -(p * p.log()).sum(dim=1) 81 | return entropy 82 | 83 | 84 | def split_ab(x): 85 | if len(x) <= 1: 86 | return x, x 87 | a = x[: int(len(x) / 2)] 88 | b = x[int(len(x) / 2) :] 89 | 90 | return a, b 91 | -------------------------------------------------------------------------------- /segmentation/gvcore/utils/checkpoint.py: -------------------------------------------------------------------------------- 1 | from os.path import join 2 | import os 3 | import torch 4 | import torch.nn as nn 5 | from torch.nn.parallel import DistributedDataParallel 6 | 7 | from gvcore.utils.logger import logger 8 | from gvcore.utils.distributed import is_main_process 9 | 10 | 11 | class Checkpointer(object): 12 | def __init__(self, cfg): 13 | self.cfg = cfg 14 | self.dev_mode = cfg.dev_mode 15 | 16 | def save(self, name, **checkpointables): 17 | if not is_main_process() or self.dev_mode: 18 | return 19 | 20 | if self.cfg.log.get("keep_last", 0) > 0: 21 | existing_ckps = os.listdir(self.cfg.log.log_dir) 22 | existing_ckps = [ckp for ckp in existing_ckps if ckp.endswith(".pth")] 23 | existing_ckps.sort(key=lambda x: int(x.split("_")[-1].split(".")[0])) 24 | if len(existing_ckps) >= self.cfg.log.get("keep_last", 0): 25 | os.remove(join(self.cfg.log.log_dir, existing_ckps[0])) 26 | 27 | data = {} 28 | for key, obj in checkpointables.items(): 29 | if isinstance(obj, DistributedDataParallel): 30 | data[key] = obj.module.state_dict() 31 | elif isinstance(obj, (nn.Module, torch.optim.Optimizer, torch.optim.lr_scheduler._LRScheduler)): 32 | data[key] = obj.state_dict() 33 | else: 34 | data[key] = obj 35 | 36 | file_path = join(self.cfg.log.log_dir, f"{name}.pth") 37 | torch.save(data, file_path) 38 | logger.info(f"Saved checkpoint to {file_path}.") 39 | 40 | def load(self, ckp_path, **checkpointables): 41 | ckp_dict = torch.load(ckp_path, map_location="cpu") 42 | except_keys = [] 43 | for key, obj in checkpointables.items(): 44 | if isinstance(obj, DistributedDataParallel): 45 | load_obj = obj.module 46 | else: 47 | load_obj = obj 48 | if key in ckp_dict: 49 | if ckp_dict[key] is not None: 50 | if isinstance(load_obj, nn.Module): 51 | except_param_keys = load_obj.load_state_dict(ckp_dict[key], strict=False) 52 | logger.info("Unmatched params:" + str(except_param_keys)) 53 | elif isinstance(load_obj, (torch.optim.Optimizer, torch.optim.lr_scheduler._LRScheduler)): 54 | load_obj.load_state_dict(ckp_dict[key]) 55 | else: 56 | except_keys.append(key) 57 | else: 58 | except_keys.append(key) 59 | except_string = "" if len(except_keys) == 0 else f", except: {', '.join(except_keys)}" 60 | logger.info(f"Loaded checkpoint from {ckp_path}{except_string}.") 61 | 62 | if "step" in ckp_dict: 63 | return ckp_dict["step"] 64 | else: 65 | return 0 66 | -------------------------------------------------------------------------------- /detection/det/ext/.clang-format: -------------------------------------------------------------------------------- 1 | AccessModifierOffset: -1 2 | AlignAfterOpenBracket: AlwaysBreak 3 | AlignConsecutiveAssignments: false 4 | AlignConsecutiveDeclarations: false 5 | AlignEscapedNewlinesLeft: true 6 | AlignOperands: false 7 | AlignTrailingComments: false 8 | AllowAllParametersOfDeclarationOnNextLine: false 9 | AllowShortBlocksOnASingleLine: false 10 | AllowShortCaseLabelsOnASingleLine: false 11 | AllowShortFunctionsOnASingleLine: Empty 12 | AllowShortIfStatementsOnASingleLine: false 13 | AllowShortLoopsOnASingleLine: false 14 | AlwaysBreakAfterReturnType: None 15 | AlwaysBreakBeforeMultilineStrings: true 16 | AlwaysBreakTemplateDeclarations: true 17 | BinPackArguments: false 18 | BinPackParameters: false 19 | BraceWrapping: 20 | AfterClass: false 21 | AfterControlStatement: false 22 | AfterEnum: false 23 | AfterFunction: false 24 | AfterNamespace: false 25 | AfterObjCDeclaration: false 26 | AfterStruct: false 27 | AfterUnion: false 28 | BeforeCatch: false 29 | BeforeElse: false 30 | IndentBraces: false 31 | BreakBeforeBinaryOperators: None 32 | BreakBeforeBraces: Attach 33 | BreakBeforeTernaryOperators: true 34 | BreakConstructorInitializersBeforeComma: false 35 | BreakAfterJavaFieldAnnotations: false 36 | BreakStringLiterals: false 37 | ColumnLimit: 80 38 | CommentPragmas: '^ IWYU pragma:' 39 | ConstructorInitializerAllOnOneLineOrOnePerLine: true 40 | ConstructorInitializerIndentWidth: 4 41 | ContinuationIndentWidth: 4 42 | Cpp11BracedListStyle: true 43 | DerivePointerAlignment: false 44 | DisableFormat: false 45 | ForEachMacros: [ FOR_EACH, FOR_EACH_R, FOR_EACH_RANGE, ] 46 | IncludeCategories: 47 | - Regex: '^<.*\.h(pp)?>' 48 | Priority: 1 49 | - Regex: '^<.*' 50 | Priority: 2 51 | - Regex: '.*' 52 | Priority: 3 53 | IndentCaseLabels: true 54 | IndentWidth: 2 55 | IndentWrappedFunctionNames: false 56 | KeepEmptyLinesAtTheStartOfBlocks: false 57 | MacroBlockBegin: '' 58 | MacroBlockEnd: '' 59 | MaxEmptyLinesToKeep: 1 60 | NamespaceIndentation: None 61 | ObjCBlockIndentWidth: 2 62 | ObjCSpaceAfterProperty: false 63 | ObjCSpaceBeforeProtocolList: false 64 | PenaltyBreakBeforeFirstCallParameter: 1 65 | PenaltyBreakComment: 300 66 | PenaltyBreakFirstLessLess: 120 67 | PenaltyBreakString: 1000 68 | PenaltyExcessCharacter: 1000000 69 | PenaltyReturnTypeOnItsOwnLine: 200 70 | PointerAlignment: Left 71 | ReflowComments: true 72 | SortIncludes: true 73 | SpaceAfterCStyleCast: false 74 | SpaceBeforeAssignmentOperators: true 75 | SpaceBeforeParens: ControlStatements 76 | SpaceInEmptyParentheses: false 77 | SpacesBeforeTrailingComments: 1 78 | SpacesInAngles: false 79 | SpacesInContainerLiterals: true 80 | SpacesInCStyleCastParentheses: false 81 | SpacesInParentheses: false 82 | SpacesInSquareBrackets: false 83 | Standard: Cpp11 84 | TabWidth: 8 85 | UseTab: Never 86 | -------------------------------------------------------------------------------- /segmentation/seg/model/block/block.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | from gvcore.model.utils import kaiming_init, constant_init 6 | 7 | 8 | class ConvModule(nn.Module): 9 | def __init__( 10 | self, 11 | in_channels, 12 | out_channels, 13 | kernel_size, 14 | stride=1, 15 | padding=0, 16 | dilation=1, 17 | groups=1, 18 | bias=True, 19 | padding_mode="zeros", 20 | norm=None, 21 | activation=None, 22 | ): 23 | super(ConvModule, self).__init__() 24 | 25 | self.conv = nn.Conv2d( 26 | in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias, padding_mode 27 | ) 28 | 29 | self.norm = norm 30 | self.activation = activation 31 | 32 | # self.init_weights() 33 | 34 | def init_weights(self): 35 | if not hasattr(self.conv, "init_weights"): 36 | if isinstance(self.activation, nn.LeakyReLU): 37 | nonlinearity = "leaky_relu" 38 | a = self.activation.negative_slope 39 | else: 40 | nonlinearity = "relu" 41 | a = 0 42 | kaiming_init(self.conv, a=a, nonlinearity=nonlinearity) 43 | if self.norm is not None: 44 | constant_init(self.norm, 1, bias=0) 45 | 46 | def forward(self, x: torch.Tensor) -> torch.Tensor: 47 | x = self.conv(x) 48 | if self.norm is not None: 49 | x = self.norm(x) 50 | if self.activation is not None: 51 | x = self.activation(x) 52 | return x 53 | 54 | 55 | class DepthwiseSeparableConvModule(nn.Module): 56 | def __init__( 57 | self, 58 | in_channels, 59 | out_channels, 60 | kernel_size, 61 | stride=1, 62 | padding=0, 63 | dilation=1, 64 | bias=False, 65 | dw_norm=None, 66 | pw_norm=None, 67 | activation=None, 68 | **kwargs 69 | ): 70 | super(DepthwiseSeparableConvModule, self).__init__() 71 | self.depthwise_conv = ConvModule( 72 | in_channels=in_channels, 73 | out_channels=in_channels, 74 | kernel_size=kernel_size, 75 | stride=stride, 76 | padding=padding, 77 | dilation=dilation, 78 | groups=in_channels, 79 | bias=bias, 80 | norm=dw_norm, 81 | activation=activation, 82 | ) 83 | self.pointwise_conv = ConvModule( 84 | in_channels=in_channels, 85 | out_channels=out_channels, 86 | kernel_size=1, 87 | bias=bias, 88 | norm=pw_norm, 89 | activation=activation, 90 | ) 91 | 92 | def forward(self, x: torch.Tensor) -> torch.Tensor: 93 | x = self.depthwise_conv(x) 94 | x = self.pointwise_conv(x) 95 | return x 96 | -------------------------------------------------------------------------------- /segmentation/gvcore/utils/misc.py: -------------------------------------------------------------------------------- 1 | from typing import List, Sequence 2 | import torch 3 | import time 4 | 5 | from gvcore.utils.structure import GenericData 6 | 7 | 8 | def time_test(func): 9 | def wrapper(*args): 10 | st = time.time() 11 | out = func(*args) 12 | print(args[0].__class__.__name__, "{:.8f}".format(time.time() - st)) 13 | return out 14 | 15 | return wrapper 16 | 17 | 18 | def attach_batch_idx(tensor_list): 19 | tensors = [] 20 | for i, tensor in enumerate(tensor_list): 21 | batch_idx = torch.ones((tensor.shape[0], 1), dtype=tensor.dtype, device=tensor.device) * i 22 | tensors.append(torch.cat((batch_idx, tensor), dim=1)) 23 | return torch.cat(tensors, dim=0) 24 | 25 | 26 | def print_tensor(x, device="cuda:0"): 27 | if isinstance(x, torch.Tensor): 28 | if str(x.device) == device: 29 | print(x) 30 | else: 31 | print(x) 32 | 33 | 34 | def slerp(low, high, val=0.5): 35 | omega = torch.arccos((low * high).sum(dim=1, keepdim=True)) 36 | so = torch.sin(omega) 37 | return torch.sin((1.0 - val) * omega) / so * low + torch.sin(val * omega) / so * high 38 | 39 | 40 | def entropy(x): 41 | p = torch.softmax(x, dim=1) 42 | entropy = -(p * p.log()).sum(dim=1) 43 | return entropy 44 | 45 | 46 | def split_ab(x): 47 | if len(x) <= 1: 48 | return x, x 49 | a = x[: int(len(x) / 2)] 50 | b = x[int(len(x) / 2) :] 51 | 52 | return a, b 53 | 54 | 55 | def sharpen(x: torch.Tensor, temperature: float = 0.5) -> torch.Tensor: 56 | sharpened_x = x ** (1 / temperature) 57 | return sharpened_x / sharpened_x.sum(dim=1, keepdim=True) 58 | 59 | 60 | def interleave_offsets(batch_size: int, num_unlabeled: int) -> List[int]: 61 | # TODO: scrutiny 62 | groups = [batch_size // (num_unlabeled + 1)] * (num_unlabeled + 1) 63 | for x in range(batch_size - sum(groups)): 64 | groups[-x - 1] += 1 65 | offsets = [0] 66 | for g in groups: 67 | offsets.append(offsets[-1] + g) 68 | assert offsets[-1] == batch_size 69 | return offsets 70 | 71 | 72 | def interleave(xy: Sequence[GenericData], batch_size: int) -> List[GenericData]: 73 | # TODO: scrutiny 74 | num_unlabeled = len(xy) - 1 75 | offsets = interleave_offsets(batch_size, num_unlabeled) 76 | xy = [[v[offsets[p] : offsets[p + 1]] for p in range(num_unlabeled + 1)] for v in xy] 77 | for i in range(1, num_unlabeled + 1): 78 | xy[0][i], xy[i][i] = xy[i][i], xy[0][i] 79 | 80 | outs = [] 81 | for v in xy: 82 | if torch.is_tensor(v[0]): 83 | v = torch.cat(v, dim=0) 84 | else: 85 | v = [item for subv in v for item in subv] 86 | outs.append(v) 87 | return outs 88 | 89 | 90 | def one_hot(y, num_classes): 91 | y_tensor = y.unsqueeze(1) 92 | zeros = torch.zeros([y.shape[0], num_classes] + list(y.shape[1:]), dtype=y.dtype, device=y.device) 93 | 94 | return zeros.scatter(1, y_tensor, 1) 95 | -------------------------------------------------------------------------------- /segmentation/seg/dataset/cityscapes.py: -------------------------------------------------------------------------------- 1 | import os 2 | from typing import Callable, Optional 3 | import torchvision.io as io 4 | 5 | from gvcore.utils.distributed import is_main_process, synchronize 6 | from gvcore.utils.structure import GenericData 7 | from gvcore.dataset import DATASET_REGISTRY, GenericDataset 8 | 9 | 10 | @DATASET_REGISTRY.register("Cityscapes") 11 | class CityscapeDataset(GenericDataset): 12 | def __init__( 13 | self, 14 | root: str = "./data/cityscapes", 15 | subset: str = "train", 16 | img_transforms: Optional[Callable[[GenericData], GenericData]] = None, 17 | ): 18 | super(CityscapeDataset, self).__init__(root, subset, img_transforms) 19 | 20 | def _get_data_ids(self): 21 | # 0. Create index file 22 | 23 | if is_main_process(): 24 | if not os.path.exists(f"{self._root}/index/{self._subset}.txt"): 25 | os.makedirs(f"{self._root}/index", exist_ok=True) 26 | with open(f"{self._root}/index/{self._subset}.txt", "w") as w: 27 | for p, d, f in os.walk(f"{self._root}/leftImg8bit/{self._subset}"): 28 | for file in f: 29 | if file.endswith(".png"): 30 | file_path = os.path.join(p, file) 31 | file_idx = file_path.replace("_leftImg8bit.png", "").replace( 32 | f"{self._root}/leftImg8bit/", "" 33 | ) 34 | w.write(f"{file_idx}\n") 35 | synchronize() 36 | 37 | # 1. Read index file 38 | with open(f"{self._root}/index/{self._subset}.txt", "r") as f: 39 | data_ids = [line.strip() for line in f.readlines()] 40 | 41 | return sorted(data_ids) 42 | 43 | def _load(self, item_idx): 44 | data_id = self._data_ids[item_idx] 45 | img = io.read_image(f"{self._root}/leftImg8bit/{data_id}_leftImg8bit.png") 46 | if img.shape[0] == 1: 47 | img = img.repeat(3, 1, 1) 48 | label = io.read_image(f"{self._root}/gtFine/{data_id}_gtFine_labelTrainIds.png") 49 | 50 | data = GenericData() 51 | data.set("img", img) 52 | data.set("label", label) 53 | 54 | data.set("meta", GenericData()) 55 | data.meta.set("id", data_id) 56 | 57 | img_size = data.img.shape[1:] 58 | data.meta.set("ori_size", list(img_size)) 59 | data.meta.set("cur_size", list(img_size)) 60 | 61 | return data 62 | 63 | 64 | color_map = [ 65 | (128, 64, 128), 66 | (244, 35, 232), 67 | (70, 70, 70), 68 | (102, 102, 156), 69 | (190, 153, 153), 70 | (153, 153, 153), 71 | (250, 170, 30), 72 | (220, 220, 0), 73 | (107, 142, 35), 74 | (152, 251, 152), 75 | (70, 130, 180), 76 | (220, 20, 60), 77 | (255, 0, 0), 78 | (0, 0, 142), 79 | (0, 0, 70), 80 | (0, 60, 100), 81 | (0, 80, 100), 82 | (0, 0, 230), 83 | (119, 11, 32), 84 | ] 85 | -------------------------------------------------------------------------------- /segmentation/seg/config/ts_deeplabv3p/cityscapes.yaml: -------------------------------------------------------------------------------- 1 | _base: "./base.yaml" 2 | 3 | log: 4 | comment: "city" 5 | 6 | data: 7 | weak_transforms: 8 | batch_apply: 9 | normalize: 10 | - [0.485, 0.456, 0.406] 11 | - [0.229, 0.224, 0.225] 12 | - True 13 | - False 14 | batch_pad: 15 | fixed_size: [512, 512] 16 | strong_transforms: 17 | batch_apply: 18 | random_color_jitter: 19 | - 0.8 20 | - 0.5 21 | - 0.5 22 | - 0.5 23 | - 0.25 24 | random_grayscale: 25 | random_gaussian_blur: 26 | normalize: 27 | - [0.485, 0.456, 0.406] 28 | - [0.229, 0.224, 0.225] 29 | - True 30 | - False 31 | k_random_erase: 32 | scale: [0.02, 0.4] 33 | ratio: [0.3, 3.3] 34 | p: 0.5 35 | batch_pad: 36 | fixed_size: [512, 512] 37 | train_l: 38 | dataset: "Cityscapes" 39 | subset: "l.1.72.fold3" 40 | prefetch: False 41 | infinity_sampler: True 42 | root: "./data/cityscapes" 43 | num_workers: 4 44 | batch_size: 4 45 | shuffle: True 46 | img_transforms: 47 | batch_transforms: 48 | batch_apply: 49 | resize: 50 | scale: [0.5, 0.75, 1, 1.5, 1.75, 2.0] 51 | random_crop: 52 | - [512, 512] 53 | random_horizontal_flip: 54 | train_u: 55 | dataset: "Cityscapes" 56 | subset: "u.1.72.fold3" 57 | prefetch: False 58 | infinity_sampler: True 59 | root: "./data/cityscapes" 60 | num_workers: 4 61 | batch_size: 4 62 | shuffle: True 63 | img_transforms: 64 | batch_transforms: 65 | batch_apply: 66 | resize: 67 | scale: [0.5, 0.75, 1, 1.5, 1.75, 2.0] 68 | random_crop: 69 | - [512, 512] 70 | random_horizontal_flip: 71 | test: 72 | dataset: "Cityscapes" 73 | subset: "val" 74 | prefetch: False 75 | infinity_sampler: False 76 | root: "./data/cityscapes" 77 | num_workers: 1 78 | batch_size: 1 79 | shuffle: False 80 | img_transforms: 81 | batch_transforms: 82 | batch_apply: 83 | resize: 84 | min_size: [1024, 1024] 85 | max_size: 2048 86 | mode: "choice" 87 | process_label: False 88 | normalize: 89 | - [0.485, 0.456, 0.406] 90 | - [0.229, 0.224, 0.225] 91 | - True 92 | - False 93 | batch_pad: 94 | stride: 1 95 | 96 | model: 97 | num_classes: 19 98 | 99 | solver: 100 | iter_num: 80000 101 | burnin_iter_num: 2000 102 | eval_interval: 1000 103 | ckp_interval: 5000 104 | ema_momentum: 0.9996 105 | ema_bn: True 106 | u_loss_weight: 1 107 | optimizer: 108 | type: "sgd" 109 | lr: 0.002 110 | weight_decay: 0.0001 111 | weight_decay_norm: 0 112 | weight_decay_bias: 0 113 | momentum: 0.9 114 | lr_scheduler: 115 | type: "const" 116 | warmup: True 117 | warmup_step: 500 118 | warmup_gamma: 0.001 119 | -------------------------------------------------------------------------------- /segmentation/seg/config/vc_deeplabv3p/cityscapes.yaml: -------------------------------------------------------------------------------- 1 | _base: "./base.yaml" 2 | 3 | log: 4 | comment: "city" 5 | 6 | data: 7 | weak_transforms: 8 | batch_apply: 9 | normalize: 10 | - [0.485, 0.456, 0.406] 11 | - [0.229, 0.224, 0.225] 12 | - True 13 | - False 14 | batch_pad: 15 | fixed_size: [512, 512] 16 | strong_transforms: 17 | batch_apply: 18 | random_color_jitter: 19 | - 0.8 20 | - 0.5 21 | - 0.5 22 | - 0.5 23 | - 0.25 24 | random_grayscale: 25 | random_gaussian_blur: 26 | normalize: 27 | - [0.485, 0.456, 0.406] 28 | - [0.229, 0.224, 0.225] 29 | - True 30 | - False 31 | k_random_erase: 32 | scale: [0.02, 0.4] 33 | ratio: [0.3, 3.3] 34 | p: 0.5 35 | batch_pad: 36 | fixed_size: [512, 512] 37 | train_l: 38 | dataset: "Cityscapes" 39 | subset: "l.1.72.fold2" 40 | prefetch: False 41 | infinity_sampler: True 42 | root: "./data/cityscapes" 43 | num_workers: 4 44 | batch_size: 4 45 | shuffle: True 46 | img_transforms: 47 | batch_transforms: 48 | batch_apply: 49 | resize: 50 | scale: [0.5, 0.75, 1, 1.5, 1.75, 2.0] 51 | random_crop: 52 | - [512, 512] 53 | random_horizontal_flip: 54 | train_u: 55 | dataset: "Cityscapes" 56 | subset: "u.1.72.fold2" 57 | prefetch: False 58 | infinity_sampler: True 59 | root: "./data/cityscapes" 60 | num_workers: 4 61 | batch_size: 4 62 | shuffle: True 63 | img_transforms: 64 | batch_transforms: 65 | batch_apply: 66 | resize: 67 | scale: [0.5, 0.75, 1, 1.5, 1.75, 2.0] 68 | random_crop: 69 | - [512, 512] 70 | random_horizontal_flip: 71 | test: 72 | dataset: "Cityscapes" 73 | subset: "val" 74 | prefetch: False 75 | infinity_sampler: False 76 | root: "./data/cityscapes" 77 | num_workers: 1 78 | batch_size: 1 79 | shuffle: False 80 | img_transforms: 81 | batch_transforms: 82 | batch_apply: 83 | resize: 84 | min_size: [1024, 1024] 85 | max_size: 2048 86 | mode: "choice" 87 | process_label: False 88 | normalize: 89 | - [0.485, 0.456, 0.406] 90 | - [0.229, 0.224, 0.225] 91 | - True 92 | - False 93 | batch_pad: 94 | stride: 1 95 | 96 | model: 97 | num_classes: 19 98 | 99 | solver: 100 | iter_num: 80000 101 | burnin_iter_num: 2000 102 | eval_interval: 1000 103 | ckp_interval: 5000 104 | ema_momentum: 0.9996 105 | ema_bn: True 106 | u_loss_weight: 1 107 | optimizer: 108 | type: "sgd" 109 | lr: 0.002 110 | weight_decay: 0.0001 111 | weight_decay_norm: 0 112 | weight_decay_bias: 0 113 | momentum: 0.9 114 | lr_scheduler: 115 | type: "const" 116 | warmup: True 117 | warmup_step: 500 118 | warmup_gamma: 0.001 119 | -------------------------------------------------------------------------------- /detection/det/model/detector/fasterrcnn.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from model.backbone.fpn import FPN, LastLevelMaxPool 3 | from torchvision.ops import batched_nms 4 | 5 | from gvcore.utils.types import TDataList 6 | from gvcore.model import GenericModule 7 | 8 | from model.backbone import make_backbone 9 | from model.head.rpn import RPN 10 | from model.head.roihead import ROIHead 11 | 12 | 13 | class FasterRCNN(GenericModule): 14 | def __init__(self, cfg): 15 | super(FasterRCNN, self).__init__(cfg) 16 | 17 | # Components ======== 18 | self.backbone = make_backbone(cfg) 19 | base_inc = self.backbone.inplanes // 8 20 | fpn_in_channels_list = [base_inc * 2 ** (i - 1) for i in range(1, 5)] 21 | self.fpn = FPN( 22 | in_features=["res2", "res3", "res4", "res5"], 23 | in_channels_list=fpn_in_channels_list, 24 | out_channels=256, 25 | strides=[4, 8, 16, 32], 26 | top_block=LastLevelMaxPool(), 27 | ) 28 | self.rpn = RPN(cfg) 29 | self.roi_heads = ROIHead(cfg) 30 | 31 | # == Params ==================== 32 | self.score_threshold = cfg.model.score_threshold 33 | self.nms_threshold = cfg.model.nms_threshold 34 | self.max_detections = cfg.model.max_detections 35 | 36 | self.feats = None 37 | 38 | def forward_train(self, data_list: TDataList): 39 | imgs = torch.stack([data.img for data in data_list], dim=0) 40 | labels = [data.label for data in data_list] 41 | img_sizes = [data.meta.cur_size for data in data_list] 42 | 43 | feats = self.backbone(imgs) 44 | feats = self.fpn(feats) 45 | self.feats = feats 46 | 47 | proposals, rpn_loss_dict = self.rpn(feats, img_sizes, labels) 48 | roi_loss_dict = self.roi_heads(feats, proposals, labels) 49 | 50 | loss_dict = {} 51 | loss_dict.update(rpn_loss_dict) 52 | loss_dict.update(roi_loss_dict) 53 | return loss_dict 54 | 55 | def forward_eval(self, data_list: TDataList): 56 | imgs = torch.stack([data.img for data in data_list], dim=0) 57 | img_sizes = [data.meta.cur_size for data in data_list] 58 | 59 | feats = self.backbone(imgs) 60 | feats = self.fpn(feats) 61 | self.feats = feats 62 | 63 | proposals = self.rpn(feats, img_sizes) 64 | detected_scores, detected_boxes = self.roi_heads(feats, proposals, img_sizes) 65 | 66 | detected_bboxes = [] 67 | for score, box in zip(detected_scores, detected_boxes): 68 | filter_mask = torch.gt(score, self.score_threshold) 69 | cls = torch.nonzero(filter_mask, as_tuple=True)[1] 70 | box = box[filter_mask] 71 | score = score[filter_mask] 72 | 73 | keep = batched_nms(box, score, cls, self.nms_threshold) 74 | keep = keep[: self.max_detections] 75 | box, score, cls = box[keep], score[keep], cls[keep] 76 | detected_bbox = torch.cat((box, cls.view(-1, 1), score.view(-1, 1)), dim=1) 77 | detected_bboxes.append(detected_bbox) 78 | 79 | return detected_bboxes 80 | -------------------------------------------------------------------------------- /detection/det/model/head/unbiased_teacher/roihead.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as F 3 | 4 | from model.head.roihead import ROIHead as _ROIHead 5 | from model.utils import match_label, sample_pos_neg 6 | from model.loss.functional import focal_loss 7 | 8 | 9 | class ROIHead(_ROIHead): 10 | def __init__(self, cfg): 11 | super(ROIHead, self).__init__(cfg) 12 | 13 | def forward_train(self, feats, proposals, labels, ignore_loc=False): 14 | feats = [feats[f] for f in self.in_feats] 15 | 16 | # Add ground truth to proposals 17 | proposal_boxes = [] 18 | for proposal, label in zip(proposals, labels): 19 | proposal_box = torch.cat((proposal["boxes"], label[:, :4])) 20 | proposal_boxes.append(proposal_box) 21 | 22 | matches, match_flags = match_label(proposal_boxes, labels, self.matcher) 23 | match_flags = sample_pos_neg(match_flags, self.roi_num, self.roi_positive_fraction) 24 | 25 | cls_labels, reg_labels = [], [] 26 | fg_flags = [] 27 | valid_proposals = [] 28 | for proposal_box, label, match, match_flag in zip(proposal_boxes, labels, matches, match_flags): 29 | fg_flag = match_flag > 0 30 | bg_flag = match_flag == 0 31 | valid_flag = match_flag >= 0 32 | 33 | valid_proposals.append(proposal_box[valid_flag]) 34 | fg_flags.append(fg_flag[valid_flag]) 35 | 36 | if label.shape[0] == 0: 37 | matched_label = torch.zeros((1, 5), device=self.device)[match] 38 | else: 39 | matched_label = label[match] 40 | 41 | matched_label[bg_flag, 4] = self.num_classes 42 | cls_label = matched_label[valid_flag, 4] 43 | cls_labels.append(cls_label.long()) 44 | 45 | if not ignore_loc: 46 | box_label = matched_label[fg_flag, :4] 47 | reg_label = self.box_coder.encode_single(box_label, proposal_box[fg_flag]) 48 | reg_labels.append(reg_label) 49 | 50 | # Forward head 51 | self.pooled_feats = self.pooler(feats, valid_proposals) 52 | self.box_feats = self.box_head(self.pooled_feats) 53 | self.cls_logits, self.box_deltas = self.box_predictor(self.box_feats) 54 | 55 | return self.get_losses(cls_labels, reg_labels, fg_flags, ignore_loc) 56 | 57 | def get_losses(self, cls_labels, reg_labels, fg_flags, ignore_loc=False): 58 | cls_labels = torch.cat(cls_labels) 59 | roi_cls_loss = focal_loss(self.cls_logits, cls_labels, gamma=1.5) / cls_labels.shape[0] 60 | 61 | if not ignore_loc: 62 | reg_labels = torch.cat(reg_labels) 63 | fg_flags = torch.cat(fg_flags) 64 | 65 | if self.box_deltas.shape[1] == 4: # cls-agnostic regression 66 | box_deltas = self.box_deltas[fg_flags] 67 | else: 68 | box_deltas = self.box_deltas.view(-1, self.num_classes, 4)[fg_flags, cls_labels[fg_flags]] 69 | roi_loc_loss = F.smooth_l1_loss(box_deltas, reg_labels, beta=0.0, reduction="sum") / max( 70 | 1, cls_labels.numel() 71 | ) 72 | return {"roi_cls": roi_cls_loss, "roi_loc": roi_loc_loss} 73 | else: 74 | return {"roi_cls": roi_cls_loss} 75 | -------------------------------------------------------------------------------- /detection/det/model/detector/ts.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from gvcore.utils.types import TDataList 4 | from gvcore.model import GenericModule 5 | 6 | from dataset.transforms import DataCopy, parse_transform_config 7 | 8 | 9 | class TSDetector(GenericModule): 10 | """ 11 | Teacher-Student detector for semi-supervised learning. 12 | """ 13 | 14 | def __init__(self, cfg, detector_class: GenericModule): 15 | super(TSDetector, self).__init__(cfg) 16 | 17 | # Components ======== 18 | self.detector = detector_class(cfg) 19 | self.ema_detector = detector_class(cfg) 20 | 21 | # Tools ============= 22 | self.weak_aug = DataCopy(parse_transform_config(cfg.data.weak_transforms)) 23 | self.strong_aug = DataCopy(parse_transform_config(cfg.data.strong_transforms)) 24 | 25 | # Initialize ======== 26 | for param_q, param_k in zip(self.detector.parameters(), self.ema_detector.parameters()): 27 | param_k.data.copy_(param_q.data) 28 | param_k.requires_grad = False 29 | self.ema_detector.freeze(True) 30 | self.ema_detector.eval() 31 | 32 | # Params ============== 33 | self.m = cfg.solver.ema_momentum 34 | self.pseudo_t = cfg.model.teacher.score_threshold 35 | 36 | def train(self, mode: bool = True): 37 | self.training = mode 38 | self.detector.train(mode) 39 | 40 | def forward(self, *args, **kwargs): 41 | if "stage" not in kwargs: 42 | stage = "train_l" if self.training else "eval" 43 | else: 44 | stage = kwargs.pop("stage") 45 | 46 | if stage == "train_l": 47 | return self.forward_train_l(*args, **kwargs) 48 | elif stage == "train_u": 49 | return self.forward_train_u(*args, **kwargs) 50 | elif stage == "eval": 51 | return self.forward_eval(*args, **kwargs) 52 | 53 | def forward_train_l(self, data_list_l: TDataList): 54 | weak_data_list_l = self.weak_aug(data_list_l) 55 | strong_data_list_l = self.strong_aug(data_list_l) 56 | data_list_l = weak_data_list_l + strong_data_list_l 57 | 58 | loss_dict_l = self.detector(data_list_l, labeled=True) 59 | 60 | return loss_dict_l 61 | 62 | def forward_train_u(self, data_list_u: TDataList): 63 | self._momentum_update(m=self.m) 64 | 65 | weak_data_list_u = self.weak_aug(data_list_u) 66 | strong_data_list_u = self.strong_aug(data_list_u) 67 | 68 | pseudo_labels = self.ema_detector(weak_data_list_u) 69 | for data, pseudo_label in zip(strong_data_list_u, pseudo_labels): 70 | select = pseudo_label[:, 5] > self.pseudo_t 71 | data.label = torch.narrow(pseudo_label[select], dim=1, start=0, length=5) 72 | 73 | loss_dict_u = self.detector(strong_data_list_u, labeled=False) 74 | 75 | return loss_dict_u 76 | 77 | def forward_eval(self, *args, **kwargs): 78 | selector = kwargs.pop("selector") 79 | if selector == "student": 80 | return self.detector(*args, **kwargs) 81 | elif selector == "teacher": 82 | return self.ema_detector(*args, **kwargs) 83 | 84 | @torch.no_grad() 85 | def _momentum_update(self, m: float): 86 | """ 87 | Momentum update of the key encoder 88 | """ 89 | for param_q, param_k in zip(self.detector.parameters(), self.ema_detector.parameters()): 90 | param_k.data = param_k.data * m + param_q.data * (1.0 - m) 91 | -------------------------------------------------------------------------------- /detection/gvcore/utils/lr_scheduler.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import math 3 | import torch.optim.lr_scheduler as lr_s 4 | from bisect import bisect_right 5 | 6 | 7 | class PolyPolicy: 8 | def __init__(self, iter_num, power=0.9): 9 | self.power = power 10 | self.iter_num = iter_num 11 | 12 | def __call__(self, step): 13 | return (1 - step / self.iter_num) ** self.power 14 | 15 | 16 | class PolyLR(lr_s.LambdaLR): 17 | def __init__(self, optimizer, iter_num, power=0.9): 18 | self.power = power 19 | self.iter_num = iter_num 20 | poly = PolyPolicy(self.iter_num, self.power) 21 | super(PolyLR, self).__init__(optimizer, poly) 22 | 23 | 24 | class WarmupPolicy: 25 | def __init__(self, warmup_step, warmup_gamma): 26 | self.warmup_step = warmup_step 27 | self.warmup_gamma = warmup_gamma 28 | 29 | def __call__(self, step): 30 | if step >= self.warmup_step: 31 | return 1 32 | alpha = float(step) / self.warmup_step 33 | return self.warmup_gamma * (1 - alpha) + alpha 34 | 35 | 36 | class WarmupLR(lr_s.LambdaLR): 37 | def __init__(self, optimizer, warmup_step, warmup_gamma): 38 | policy = WarmupPolicy(warmup_step, warmup_gamma) 39 | super(WarmupLR, self).__init__(optimizer, policy) 40 | 41 | 42 | class MultiStepPolicy: 43 | def __init__(self, milestones, gamma, warmup=False, warmup_step=None, warmup_gamma=None): 44 | self.milestones = milestones 45 | self.gamma = gamma 46 | self.warmup = warmup 47 | if self.warmup: 48 | assert warmup_step is not None and warmup_gamma is not None, "Warmup policy need to set step and gamma!" 49 | self.warmup_policy = WarmupPolicy(warmup_step, warmup_gamma) 50 | self.warmup_step = warmup_step 51 | else: 52 | self.warmup_step = -1 53 | 54 | def __call__(self, step): 55 | if step <= self.warmup_step: 56 | return self.warmup_policy(step) 57 | else: 58 | return self.gamma ** bisect_right(self.milestones, step) 59 | 60 | 61 | class MultiStepLR(lr_s.LambdaLR): 62 | def __init__(self, optimizer, milestones, gamma=0.1, warmup=False, warmup_step=None, warmup_gamma=None): 63 | policy = MultiStepPolicy(milestones, gamma, warmup, warmup_step, warmup_gamma) 64 | super(MultiStepLR, self).__init__(optimizer, policy) 65 | 66 | 67 | class CosinePolicy: 68 | def __init__(self, iter_num, num_cycles=7./16., warmup=False, warmup_step=None, warmup_gamma=None): 69 | self.iter_num = iter_num 70 | self.num_cycles = num_cycles 71 | self.warmup = warmup 72 | if self.warmup: 73 | assert warmup_step is not None and warmup_gamma is not None, "Warmup policy need to set step and gamma!" 74 | self.warmup_policy = WarmupPolicy(warmup_step, warmup_gamma) 75 | self.warmup_step = warmup_step 76 | else: 77 | self.warmup_step = 0 78 | 79 | def __call__(self, step): 80 | if step < self.warmup_step: 81 | return self.warmup_policy(step) 82 | else: 83 | return math.cos(math.pi * self.num_cycles * (step - self.warmup_step) / (self.iter_num - self.warmup_step)) 84 | 85 | 86 | class CosineLR(lr_s.LambdaLR): 87 | def __init__(self, optimizer, iter_num, num_cycles=7./16., warmup=False, warmup_step=None, warmup_gamma=None): 88 | policy = CosinePolicy(iter_num, num_cycles, warmup, warmup_step, warmup_gamma) 89 | super(CosineLR, self).__init__(optimizer, policy) 90 | -------------------------------------------------------------------------------- /segmentation/seg/dataset/transforms/__init__.py: -------------------------------------------------------------------------------- 1 | from typing import Iterable 2 | import torch 3 | import torchvision.transforms.functional as tvF 4 | 5 | from gvcore.dataset.transforms import TRANSFORM_REGISTRY 6 | import gvcore.dataset.transforms as GVTransforms 7 | from gvcore.utils.structure import GenericData 8 | 9 | 10 | @TRANSFORM_REGISTRY.register("resize") 11 | class Resize(GVTransforms.Resize): 12 | def _process_label(self, data: GenericData, params: any) -> GenericData: 13 | if data.has("label") and self.process_label: 14 | new_h, new_w = params[0], params[1] 15 | data.label = tvF.resize(data.label, size=[new_h, new_w], interpolation=tvF.InterpolationMode.NEAREST) 16 | return data 17 | 18 | 19 | @TRANSFORM_REGISTRY.register("random_horizontal_flip") 20 | class RandomHorizontalFlip(GVTransforms.RandomHorizontalFlip): 21 | def _process_label(self, data: GenericData, params: any) -> GenericData: 22 | if data.has("label") and params: 23 | data.label = tvF.hflip(data.label) 24 | return data 25 | 26 | 27 | @TRANSFORM_REGISTRY.register("random_crop") 28 | class RandomCrop(GVTransforms.RandomCrop): 29 | def _process_label(self, data: GenericData, params: any) -> GenericData: 30 | top, left, crop_h, crop_w = params 31 | data.label = tvF.crop(data.label, top, left, crop_h, crop_w) 32 | return data 33 | 34 | 35 | @TRANSFORM_REGISTRY.register("batch_pad") 36 | class BatchPad(GVTransforms.BatchPad): 37 | def __init__(self, fill=0, stride=32, label_fill=255, fixed_size=None): 38 | self.fill = fill 39 | self.label_fill = label_fill 40 | 41 | self.stride = stride 42 | self.fixed_size = fixed_size 43 | 44 | def _get_params(self, batch: Iterable[GenericData]): 45 | image_sizes = [data.img.shape[-2:] for data in batch] 46 | max_size = torch.tensor(image_sizes).max(0).values 47 | max_size = torch.div((max_size + (self.stride - 1)), self.stride, rounding_mode="trunc") * self.stride 48 | 49 | if self.fixed_size is not None: 50 | max_size[0] = max(max_size[0], self.fixed_size[0]) 51 | max_size[1] = max(max_size[1], self.fixed_size[1]) 52 | 53 | return max_size[0], max_size[1], image_sizes 54 | 55 | def _process_img(self, batch: Iterable[GenericData], params: any) -> Iterable[GenericData]: 56 | for i, (data, img_size) in enumerate(zip(batch, params[2])): 57 | pad_h = params[0] - img_size[0] 58 | pad_w = params[1] - img_size[1] 59 | data.img = tvF.pad(data.img, padding=[0, 0, pad_w, pad_h], fill=self.fill) 60 | data.meta.pad_size = [0, 0, int(pad_w), int(pad_h)] 61 | batch[i] = data 62 | return batch 63 | 64 | def _process_label(self, batch: Iterable[GenericData], params: any) -> Iterable[GenericData]: 65 | for i, (data, img_size) in enumerate(zip(batch, params[2])): 66 | pad_h = params[0] - img_size[0] 67 | pad_w = params[1] - img_size[1] 68 | data.label = tvF.pad(data.label, padding=[0, 0, pad_w, pad_h], fill=self.label_fill) 69 | batch[i] = data 70 | return batch 71 | 72 | 73 | @TRANSFORM_REGISTRY.register("k_random_erase") 74 | class RandomErase(GVTransforms.KRandomErase): 75 | def _process_label(self, data: GenericData, params: any) -> GenericData: 76 | if data.has("label") and params: 77 | params["values"].fill_(255.0) 78 | data.label = self.transform(data.label.float(), params=params).to(data.label.dtype) 79 | return data 80 | -------------------------------------------------------------------------------- /segmentation/seg/utils/visualizer.py: -------------------------------------------------------------------------------- 1 | from PIL import ImageFont, ImageDraw 2 | import torch 3 | from torchvision.transforms.functional import to_pil_image, resize, InterpolationMode 4 | import matplotlib.pyplot as plt 5 | import numpy as np 6 | 7 | from gvcore.dataset.transforms import Denormalize 8 | from gvcore.utils.structure import GenericData 9 | 10 | 11 | class Visualizer: 12 | def __init__(self): 13 | pass 14 | 15 | def __call__(self, img, label, pred=None, cmap=None): 16 | label = resize(label.unsqueeze(0), img.shape[1:], InterpolationMode.NEAREST).squeeze(0) 17 | if pred is not None: 18 | pred = resize(pred.unsqueeze(0), img.shape[1:], InterpolationMode.NEAREST).squeeze(0) 19 | pred = pred.expand_as(label) 20 | 21 | img = img.detach().cpu().numpy() 22 | label = label.detach().cpu().numpy() 23 | label[label == 255] = 0 24 | 25 | if pred is not None: 26 | pred = pred.detach().cpu().numpy() 27 | pred[label == 255] = 0 28 | 29 | def bitget(byteval, idx): 30 | return (byteval & (1 << idx)) != 0 31 | 32 | ## Labels to color map 33 | if cmap is None: 34 | cmap = np.zeros((21, 3), dtype=np.uint8) 35 | for i in range(21): 36 | r = g = b = 0 37 | c = i 38 | for j in range(8): 39 | r = r | (bitget(c, 0) << 7 - j) 40 | g = g | (bitget(c, 1) << 7 - j) 41 | b = b | (bitget(c, 2) << 7 - j) 42 | c = c >> 3 43 | 44 | cmap[i] = np.array([r, g, b]) 45 | label = cmap[label] 46 | if pred is not None: 47 | pred = cmap[pred] 48 | 49 | if label.ndim == 4: 50 | label = label[0] 51 | if pred is not None and pred.ndim == 4: 52 | pred = pred[0] 53 | 54 | if pred is not None: 55 | img = np.concatenate((img, label.transpose((2, 0, 1)), pred.transpose((2, 0, 1))), axis=1).transpose( 56 | (1, 2, 0) 57 | ) 58 | else: 59 | img = np.concatenate((img, label.transpose((2, 0, 1))), axis=2).transpose((1, 2, 0)) 60 | img = to_pil_image(img) 61 | 62 | return img 63 | 64 | 65 | def vis(data_list, result_list=None, out_path=None, cmap=None): 66 | 67 | visualizer = Visualizer() 68 | 69 | if isinstance(data_list, GenericData): 70 | data_list = [data_list] 71 | result_list = [result_list] 72 | elif isinstance(data_list, list): 73 | pass 74 | else: 75 | raise ValueError("data_list need to be a GenericData object or list of GenericData.") 76 | 77 | if result_list is None: 78 | result_list = [result_list] * len(data_list) 79 | 80 | imgs = [] 81 | denorm = Denormalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225], True, False) 82 | for i, (data, result) in enumerate(zip(data_list, result_list)): 83 | data = denorm(data) 84 | img = data.img.cpu().to(torch.uint8) 85 | label = data.label.cpu() 86 | if result is not None: 87 | result = result.cpu() 88 | img = visualizer(img, label, result, cmap) 89 | img = np.asarray(img) 90 | imgs.append(img) 91 | 92 | imgs = np.concatenate(imgs, axis=0) 93 | 94 | plt.figure(figsize=(imgs.shape[1] / 200, imgs.shape[0] / 200), dpi=200) 95 | plt.subplots_adjust(bottom=0, top=1, left=0, right=1) 96 | plt.imshow(imgs) 97 | 98 | if out_path is None: 99 | plt.show() 100 | else: 101 | plt.savefig(out_path) 102 | -------------------------------------------------------------------------------- /segmentation/gvcore/utils/sweeper.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import subprocess 4 | import time 5 | import wandb 6 | 7 | 8 | class Sweeper: 9 | def __init__(self, sweep_id: str, command: str, sweep_count: int = 9999, max_queue: int = 3): 10 | self.sweep_id = sweep_id 11 | self.command = command 12 | self.sweep_count = sweep_count 13 | self.max_queue = max_queue 14 | 15 | def get_waiting_job_num(self): 16 | cluster_name = os.popen("hostname").read() 17 | 18 | if "avon" in cluster_name: 19 | waiting_job_num = self._get_waiting_job_num_avon() 20 | elif "wmg" in cluster_name: 21 | waiting_job_num = self._get_waiting_job_num_wmg() 22 | 23 | return waiting_job_num 24 | 25 | def _get_waiting_job_num_avon(self): 26 | waiting_job_num = 0 27 | cmd = "squeue -u wmrkwh -h" 28 | job_list = os.popen(cmd).read().split("\n") 29 | for job in job_list: 30 | if job != "": 31 | stat = job.split()[4] 32 | if stat != "R": 33 | waiting_job_num += 1 34 | 35 | return waiting_job_num 36 | 37 | def _get_waiting_job_num_wmg(self): 38 | waiting_job_num = 0 39 | cmd = "squeue -u chen_c -h" 40 | job_list = os.popen(cmd).read().split("\n") 41 | for job in job_list: 42 | if job != "": 43 | stat = job.split()[4] 44 | if stat != "R": 45 | waiting_job_num += 1 46 | 47 | return waiting_job_num 48 | 49 | def start_agent(self): 50 | run = wandb.init() 51 | 52 | agent_confg = "" 53 | for k, v in wandb.config.items(): 54 | if v == "true": 55 | v = "True" 56 | elif v == "false": 57 | v = "False" 58 | elif isinstance(v, dict): 59 | v = f'"{v}"' 60 | v = str(v).replace(" ", "") 61 | agent_confg += f"{k}={v} " 62 | 63 | agent_confg += f"log.wandb_id={run.id}" 64 | command = self.command + f" --modify {agent_confg}" 65 | command_list = ["gvsubmit"] 66 | command_list.extend(command.split(" ")) 67 | subprocess.run(command_list) 68 | 69 | def sweep(self): 70 | for _ in range(self.sweep_count): 71 | waiting_job_num = self.get_waiting_job_num() 72 | if waiting_job_num < self.max_queue: 73 | print(f"[*] {waiting_job_num} jobs are waiting in queue. Starting sweep agent.") 74 | try: 75 | wandb.agent(self.sweep_id, function=self.start_agent, count=1) 76 | except Exception as e: 77 | print(f"[*] Sweep stopped.") 78 | exit() 79 | else: 80 | print(f"[!] {waiting_job_num} jobs are waiting in queue.") 81 | time.sleep(30) 82 | 83 | 84 | def parse_args(): 85 | parser = argparse.ArgumentParser() 86 | parser.add_argument("--sweep-id", type=str, required=True, help="Wandb sweep ID.") 87 | parser.add_argument("--sweep-count", type=int, default=9999, required=False) 88 | parser.add_argument("--max-queue", type=int, default=3, required=False) 89 | parser.add_argument("--num-gpus", type=int, required=True, help="Number of GPUs to use.") 90 | 91 | args = parser.parse_known_args() 92 | return args 93 | 94 | 95 | def gvsweep(): 96 | os.environ["WANDB_SILENT"] = "true" 97 | 98 | args = parse_args() 99 | command = " ".join(args[1]) 100 | command += f" --num-gpus {args[0].num_gpus}" 101 | sweeper = Sweeper(args[0].sweep_id, command, args[0].sweep_count, args[0].max_queue) 102 | 103 | sweeper.sweep() 104 | -------------------------------------------------------------------------------- /detection/det/dataset/transforms/functional.py: -------------------------------------------------------------------------------- 1 | from typing import Any, Callable, Tuple 2 | import torch 3 | import numpy as np 4 | from torchvision.ops.boxes import clip_boxes_to_image 5 | import torchvision.transforms.functional as tvf 6 | import torchvision.transforms as tv 7 | 8 | from gvcore.utils.structure import GenericData 9 | 10 | 11 | def resize(data: GenericData, min_size: Tuple[int], max_size: int, mode: str = "choice") -> GenericData: 12 | h, w = data.meta.cur_size 13 | if mode == "choice": 14 | size = np.random.choice(min_size) 15 | elif mode == "range": 16 | size = np.random.randint(min(min_size), max(min_size) + 1) 17 | else: 18 | raise ValueError("Unknown mode {}".format(mode)) 19 | 20 | scale = size * 1.0 / min(h, w) 21 | if h < w: 22 | new_h, new_w = size, scale * w 23 | else: 24 | new_h, new_w = scale * h, size 25 | if max(new_h, new_w) > max_size: 26 | scale = max_size * 1.0 / max(new_h, new_w) 27 | new_h = new_h * scale 28 | new_w = new_w * scale 29 | new_w = int(new_w + 0.5) 30 | new_h = int(new_h + 0.5) 31 | 32 | data.img = tvf.resize(data.img, size=[new_h, new_w]) 33 | resized_size = list(data.img.shape[1:]) 34 | data.meta.cur_size = resized_size 35 | 36 | ratios = [s / s_orig for s, s_orig in zip(resized_size, (h, w))] 37 | ratio_height, ratio_width = ratios 38 | ratio = torch.tensor([[ratio_width, ratio_height, ratio_width, ratio_height]], device=data.label.device) 39 | if data.has("label"): 40 | data.label[:, :4] *= ratio 41 | return data 42 | 43 | 44 | def normalise(data: GenericData, mean: Tuple[int], std: Tuple[int], scale: bool = False) -> GenericData: 45 | data.img = data.img[[2, 1, 0], :, :].float() 46 | if scale: 47 | data.img = data.img / 255.0 48 | if data.img.shape[0] == 1: 49 | data.img = data.img.repeat(3, 1, 1) 50 | data.img = tvf.normalize(data.img, mean, std) 51 | return data 52 | 53 | 54 | def denormalise(data: GenericData, mean: Tuple[int], std: Tuple[int], scale: bool = False) -> GenericData: 55 | data.img = tvf.normalize(data.img, mean=[0.0, 0.0, 0.0], std=[1 / x for x in std]) 56 | data.img = tvf.normalize(data.img, mean=[-1 * x for x in mean], std=[1.0, 1.0, 1.0]) 57 | data.img = data.img[[2, 1, 0], :, :] 58 | if scale: 59 | data.img *= 255.0 60 | return data 61 | 62 | 63 | def random_apply(data: Any, transform: Callable, p: float, *args, **kwargs) -> GenericData: 64 | if torch.rand(1) <= p: 65 | data = transform(data, *args, **kwargs) 66 | return data 67 | 68 | 69 | def horizontal_flip(data: GenericData) -> GenericData: 70 | data.img = tvf.hflip(data.img) 71 | h, w = data.meta.cur_size 72 | if data.has("label"): 73 | data.label[:, [0, 2]] = w - data.label[:, [2, 0]] 74 | data.meta.flip = True 75 | return data 76 | 77 | 78 | def gaussian_blur(data: GenericData, sigma: Tuple[float]) -> GenericData: 79 | kx = int(0.1 * data.meta.cur_size[1]) // 2 * 2 + 1 80 | ky = int(0.1 * data.meta.cur_size[0]) // 2 * 2 + 1 81 | data.img = tvf.gaussian_blur(data.img, [kx, ky], sigma) 82 | return data 83 | 84 | 85 | def transform_label_by_size( 86 | label: torch.Tensor, from_size: Tuple[int], to_size: Tuple[int], flip=False 87 | ) -> torch.Tensor: 88 | label = label.clone() 89 | 90 | scale_x, scale_y = ( 91 | 1.0 * to_size[1] / from_size[1], 92 | 1.0 * to_size[0] / from_size[0], 93 | ) 94 | label[:, [0, 2]] *= scale_x 95 | label[:, [1, 3]] *= scale_y 96 | label[:, :4] = clip_boxes_to_image(label[:, :4], to_size) 97 | 98 | if flip: 99 | h, w = to_size 100 | label[:, [0, 2]] = w - label[:, [2, 0]] 101 | 102 | return label 103 | -------------------------------------------------------------------------------- /detection/det/ext/fastevalapi/csrc/cocoeval/cocoeval.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) Facebook, Inc. and its affiliates. 2 | #pragma once 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | 10 | namespace py = pybind11; 11 | 12 | namespace fastevalapi { 13 | 14 | namespace COCOeval { 15 | 16 | // Annotation data for a single object instance in an image 17 | struct InstanceAnnotation { 18 | InstanceAnnotation( 19 | uint64_t id, 20 | double score, 21 | double area, 22 | bool is_crowd, 23 | bool ignore) 24 | : id{id}, score{score}, area{area}, is_crowd{is_crowd}, ignore{ignore} {} 25 | uint64_t id; 26 | double score = 0.; 27 | double area = 0.; 28 | bool is_crowd = false; 29 | bool ignore = false; 30 | }; 31 | 32 | // Stores intermediate results for evaluating detection results for a single 33 | // image that has D detected instances and G ground truth instances. This stores 34 | // matches between detected and ground truth instances 35 | struct ImageEvaluation { 36 | // For each of the D detected instances, the id of the matched ground truth 37 | // instance, or 0 if unmatched 38 | std::vector detection_matches; 39 | 40 | // The detection score of each of the D detected instances 41 | std::vector detection_scores; 42 | 43 | // Marks whether or not each of G instances was ignored from evaluation (e.g., 44 | // because it's outside area_range) 45 | std::vector ground_truth_ignores; 46 | 47 | // Marks whether or not each of D instances was ignored from evaluation (e.g., 48 | // because it's outside aRng) 49 | std::vector detection_ignores; 50 | }; 51 | 52 | template 53 | using ImageCategoryInstances = std::vector>>; 54 | 55 | // C++ implementation of COCO API cocoeval.py::COCOeval.evaluateImg(). For each 56 | // combination of image, category, area range settings, and IOU thresholds to 57 | // evaluate, it matches detected instances to ground truth instances and stores 58 | // the results into a vector of ImageEvaluation results, which will be 59 | // interpreted by the COCOeval::Accumulate() function to produce precion-recall 60 | // curves. The parameters of nested vectors have the following semantics: 61 | // image_category_ious[i][c][d][g] is the intersection over union of the d'th 62 | // detected instance and g'th ground truth instance of 63 | // category category_ids[c] in image image_ids[i] 64 | // image_category_ground_truth_instances[i][c] is a vector of ground truth 65 | // instances in image image_ids[i] of category category_ids[c] 66 | // image_category_detection_instances[i][c] is a vector of detected 67 | // instances in image image_ids[i] of category category_ids[c] 68 | std::vector EvaluateImages( 69 | const std::vector>& area_ranges, // vector of 2-tuples 70 | int max_detections, 71 | const std::vector& iou_thresholds, 72 | const ImageCategoryInstances>& image_category_ious, 73 | const ImageCategoryInstances& 74 | image_category_ground_truth_instances, 75 | const ImageCategoryInstances& 76 | image_category_detection_instances); 77 | 78 | // C++ implementation of COCOeval.accumulate(), which generates precision 79 | // recall curves for each set of category, IOU threshold, detection area range, 80 | // and max number of detections parameters. It is assumed that the parameter 81 | // evaluations is the return value of the functon COCOeval::EvaluateImages(), 82 | // which was called with the same parameter settings params 83 | py::dict Accumulate( 84 | const py::object& params, 85 | const std::vector& evalutations); 86 | 87 | } // namespace COCOeval 88 | } // namespace fastevalapi 89 | -------------------------------------------------------------------------------- /detection/gvcore/dataset/dataloader.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.utils.data.dataloader import DataLoader 3 | from collections import Iterator 4 | 5 | 6 | class Prefetcher(Iterator): 7 | def __init__(self, loader): 8 | self._loader = loader 9 | self.loader = iter(self._loader) 10 | self.stream = torch.cuda.Stream() 11 | self.next_batch = next(self.loader) 12 | self.to_device(self.next_batch, "cuda") 13 | torch.cuda.synchronize() 14 | 15 | @staticmethod 16 | def to_device(batch, device): 17 | if isinstance(batch, dict): 18 | for key in batch: 19 | batch[key] = batch[key].to(device) 20 | elif isinstance(batch, list): 21 | for i, item in enumerate(batch): 22 | batch[i] = item.to(device) 23 | else: 24 | batch = batch.to(device) 25 | return batch 26 | 27 | def preload(self): 28 | try: 29 | self.next_batch = next(self.loader) 30 | except StopIteration: 31 | self.loader = iter(self._loader) 32 | self.next_batch = next(self.loader) 33 | with torch.cuda.stream(self.stream): 34 | self.to_device(self.next_batch, "cuda") 35 | 36 | def __next__(self): 37 | torch.cuda.current_stream().wait_stream(self.stream) 38 | batch = self.next_batch 39 | self.preload() 40 | return batch 41 | 42 | 43 | class BasicDataloader: 44 | def __init__( 45 | self, 46 | dataset, 47 | batch_size=1, 48 | shuffle=False, 49 | sampler=None, 50 | batch_sampler=None, 51 | num_workers=0, 52 | collate_fn=None, 53 | pin_memory=False, 54 | drop_last=False, 55 | worker_init_fn=None, 56 | prefetch=False, 57 | batch_transforms=None, 58 | ): 59 | if sampler is not None or batch_sampler is not None: 60 | shuffle = False 61 | if batch_sampler is not None: 62 | sampler = None 63 | batch_size = 1 64 | drop_last = False 65 | self.dataset = dataset 66 | self.dataloader = DataLoader( 67 | dataset=dataset, 68 | batch_size=batch_size, 69 | shuffle=shuffle, 70 | sampler=sampler, 71 | batch_sampler=batch_sampler, 72 | num_workers=num_workers, 73 | collate_fn=dataset.collate_fn if collate_fn is None else collate_fn, 74 | pin_memory=pin_memory, 75 | drop_last=drop_last, 76 | worker_init_fn=worker_init_fn, 77 | ) 78 | 79 | self.batch_transforms = batch_transforms 80 | 81 | self.prefetch = prefetch 82 | if prefetch: 83 | self.iter_loader = Prefetcher(self.dataloader) 84 | else: 85 | self.iter_loader = iter(self.dataloader) 86 | 87 | def __len__(self): 88 | return len(self.dataset) 89 | 90 | @staticmethod 91 | def to_device(batch, device): 92 | if isinstance(batch, dict): 93 | for key in batch: 94 | batch[key] = batch[key].to(device) 95 | elif isinstance(batch, list): 96 | for i, item in enumerate(batch): 97 | batch[i] = item.to(device) 98 | else: 99 | batch = batch.to(device) 100 | return batch 101 | 102 | @torch.no_grad() 103 | def get_batch(self, device="cuda"): 104 | try: 105 | batch = next(self.iter_loader) 106 | self.to_device(batch, device) 107 | if self.batch_transforms is not None: 108 | batch = self.batch_transforms(batch) 109 | except StopIteration: 110 | batch = None 111 | self.iter_loader = iter(self.dataloader) 112 | return batch 113 | -------------------------------------------------------------------------------- /segmentation/gvcore/dataset/dataloader.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.utils.data.dataloader import DataLoader 3 | from collections import Iterator 4 | 5 | 6 | class Prefetcher(Iterator): 7 | def __init__(self, loader): 8 | self._loader = loader 9 | self.loader = iter(self._loader) 10 | self.stream = torch.cuda.Stream() 11 | self.next_batch = next(self.loader) 12 | self.to_device(self.next_batch, "cuda") 13 | torch.cuda.synchronize() 14 | 15 | @staticmethod 16 | def to_device(batch, device): 17 | if isinstance(batch, dict): 18 | for key in batch: 19 | batch[key] = batch[key].to(device) 20 | elif isinstance(batch, list): 21 | for i, item in enumerate(batch): 22 | batch[i] = item.to(device) 23 | else: 24 | batch = batch.to(device) 25 | return batch 26 | 27 | def preload(self): 28 | try: 29 | self.next_batch = next(self.loader) 30 | except StopIteration: 31 | self.loader = iter(self._loader) 32 | self.next_batch = next(self.loader) 33 | with torch.cuda.stream(self.stream): 34 | self.to_device(self.next_batch, "cuda") 35 | 36 | def __next__(self): 37 | torch.cuda.current_stream().wait_stream(self.stream) 38 | batch = self.next_batch 39 | self.preload() 40 | return batch 41 | 42 | 43 | class BasicDataloader: 44 | def __init__( 45 | self, 46 | dataset, 47 | batch_size=1, 48 | shuffle=False, 49 | sampler=None, 50 | batch_sampler=None, 51 | num_workers=0, 52 | collate_fn=None, 53 | pin_memory=False, 54 | drop_last=False, 55 | worker_init_fn=None, 56 | prefetch=False, 57 | batch_transforms=None, 58 | ): 59 | if sampler is not None or batch_sampler is not None: 60 | shuffle = False 61 | if batch_sampler is not None: 62 | sampler = None 63 | batch_size = 1 64 | drop_last = False 65 | self.dataset = dataset 66 | self.dataloader = DataLoader( 67 | dataset=dataset, 68 | batch_size=batch_size, 69 | shuffle=shuffle, 70 | sampler=sampler, 71 | batch_sampler=batch_sampler, 72 | num_workers=num_workers, 73 | collate_fn=dataset.collate_fn if collate_fn is None else collate_fn, 74 | pin_memory=pin_memory, 75 | drop_last=drop_last, 76 | worker_init_fn=worker_init_fn, 77 | ) 78 | 79 | self.batch_transforms = batch_transforms 80 | 81 | self.prefetch = prefetch 82 | if prefetch: 83 | self.iter_loader = Prefetcher(self.dataloader) 84 | else: 85 | self.iter_loader = iter(self.dataloader) 86 | 87 | def __len__(self): 88 | return len(self.dataset) 89 | 90 | @staticmethod 91 | def to_device(batch, device): 92 | if isinstance(batch, dict): 93 | for key in batch: 94 | batch[key] = batch[key].to(device) 95 | elif isinstance(batch, list): 96 | for i, item in enumerate(batch): 97 | batch[i] = item.to(device) 98 | else: 99 | batch = batch.to(device) 100 | return batch 101 | 102 | @torch.no_grad() 103 | def get_batch(self, device="cuda"): 104 | try: 105 | batch = next(self.iter_loader) 106 | self.to_device(batch, device) 107 | if self.batch_transforms is not None: 108 | batch = self.batch_transforms(batch) 109 | except StopIteration: 110 | batch = None 111 | self.iter_loader = iter(self.dataloader) 112 | return batch 113 | -------------------------------------------------------------------------------- /detection/gvcore/dataset/sampler.py: -------------------------------------------------------------------------------- 1 | import itertools 2 | import torch 3 | import numpy as np 4 | from torch.utils.data.sampler import Sampler, BatchSampler 5 | 6 | import gvcore.utils.distributed as dist_utils 7 | 8 | 9 | class InfiniteSampler(Sampler): 10 | def __init__(self, size, seed, shuffle=True): 11 | self._size = size 12 | assert size > 0 13 | self._shuffle = shuffle 14 | self._seed = int(seed) 15 | 16 | self._rank = dist_utils.get_rank() 17 | self._world_size = dist_utils.get_world_size() 18 | 19 | def __iter__(self): 20 | start = self._rank 21 | yield from itertools.islice(self._infinite_indices(), start, None, self._world_size) 22 | 23 | def _infinite_indices(self): 24 | g = torch.Generator() 25 | g.manual_seed(self._seed) 26 | while True: 27 | if self._shuffle: 28 | yield from torch.randperm(self._size, generator=g).tolist() 29 | else: 30 | yield from torch.arange(self._size).tolist() 31 | 32 | 33 | class SequentialSampler(Sampler): 34 | """ 35 | Produce indices for inference. 36 | Inference needs to run on the __exact__ set of samples, 37 | therefore when the total number of samples is not divisible by the number of workers, 38 | this sampler produces different number of samples on different workers. 39 | """ 40 | 41 | def __init__(self, size: int): 42 | """ 43 | Args: 44 | size (int): the total number of data of the underlying dataset to sample from 45 | """ 46 | self._size = size 47 | assert size > 0 48 | self._rank = dist_utils.get_rank() 49 | self._world_size = dist_utils.get_world_size() 50 | 51 | shard_size = (self._size - 1) // self._world_size + 1 52 | begin = shard_size * self._rank 53 | end = min(shard_size * (self._rank + 1), self._size) 54 | self._local_indices = range(begin, end) 55 | 56 | def __iter__(self): 57 | yield from self._local_indices 58 | 59 | def __len__(self): 60 | return len(self._local_indices) 61 | 62 | 63 | class GroupedBatchSampler(BatchSampler): 64 | """ 65 | Wraps another sampler to yield a mini-batch of indices. 66 | It enforces that the batch only contain elements from the same group. 67 | It also tries to provide mini-batches which follows an ordering which is 68 | as close as possible to the ordering from the original sampler. 69 | """ 70 | 71 | def __init__(self, sampler, group_ids, batch_size): 72 | """ 73 | Args: 74 | sampler (Sampler): Base sampler. 75 | group_ids (list[int]): If the sampler produces indices in range [0, N), 76 | `group_ids` must be a list of `N` ints which contains the group id of each sample. 77 | The group ids must be a set of integers in the range [0, num_groups). 78 | batch_size (int): Size of mini-batch. 79 | """ 80 | if not isinstance(sampler, Sampler): 81 | raise ValueError( 82 | "sampler should be an instance of " "torch.utils.data.Sampler, but got sampler={}".format(sampler) 83 | ) 84 | self.sampler = sampler 85 | self.group_ids = np.asarray(group_ids) 86 | assert self.group_ids.ndim == 1 87 | self.batch_size = batch_size 88 | groups = np.unique(self.group_ids).tolist() 89 | 90 | # buffer the indices of each group until batch size is reached 91 | self.buffer_per_group = {k: [] for k in groups} 92 | 93 | def __iter__(self): 94 | for idx in self.sampler: 95 | group_id = self.group_ids[idx] 96 | group_buffer = self.buffer_per_group[group_id] 97 | group_buffer.append(idx) 98 | if len(group_buffer) == self.batch_size: 99 | yield group_buffer[:] # yield a copy of the list 100 | del group_buffer[:] 101 | 102 | def __len__(self): 103 | raise NotImplementedError("len() of GroupedBatchSampler is not well-defined.") 104 | -------------------------------------------------------------------------------- /segmentation/seg/model/head/aspp_alt.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | from gvcore.model import GenericModule 6 | 7 | from model.block import Normalize 8 | from model.block.block import ConvModule 9 | 10 | 11 | class ASPPV3pAlt(GenericModule): 12 | def __init__(self, cfg): 13 | super(ASPPV3pAlt, self).__init__(cfg) 14 | 15 | # Params ============== 16 | self.in_channels = cfg.in_channels 17 | self.inner_channels = cfg.inner_channels 18 | 19 | self.lowlevel_in_channels = cfg.lowlevel_in_channels 20 | self.lowlevel_inner_channels = cfg.lowlevel_inner_channels 21 | 22 | self.dilations = cfg.dilations 23 | self.norm_layer = Normalize(cfg.norm_layer) 24 | 25 | # Components ========== 26 | self.aspp = nn.ModuleList( 27 | [ 28 | ConvModule( 29 | self.in_channels, 30 | self.inner_channels, 31 | 1 if dilation == 1 else 3, 32 | padding=0 if dilation == 1 else dilation, 33 | dilation=dilation, 34 | bias=False, 35 | norm=self.norm_layer(self.inner_channels), 36 | activation=nn.LeakyReLU(inplace=True), 37 | ) 38 | for dilation in self.dilations 39 | ] 40 | ) 41 | 42 | self.image_pool = nn.Sequential( 43 | nn.AdaptiveAvgPool2d(1), 44 | ConvModule( 45 | self.in_channels, 46 | self.inner_channels, 47 | 1, 48 | bias=False, 49 | norm=self.norm_layer(self.inner_channels), 50 | activation=nn.LeakyReLU(inplace=True), 51 | ), 52 | ) 53 | self.aspp_fuse = ConvModule(len(self.dilations) * self.inner_channels, self.inner_channels, 1, bias=False,) 54 | self.image_pool_fuse = nn.Sequential(ConvModule(self.inner_channels, self.inner_channels, 1, bias=False,),) 55 | self.fuse_bn = self.norm_layer(self.inner_channels) 56 | 57 | self.lowlevel_path = ConvModule( 58 | self.lowlevel_in_channels, 59 | self.lowlevel_inner_channels, 60 | 1, 61 | bias=False, 62 | norm=self.norm_layer(self.lowlevel_inner_channels), 63 | activation=nn.ReLU(inplace=True), 64 | ) 65 | 66 | self.last_conv = nn.Sequential( 67 | ConvModule( 68 | self.inner_channels + self.lowlevel_inner_channels, 69 | self.inner_channels, 70 | 3, 71 | padding=1, 72 | bias=False, 73 | norm=self.norm_layer(self.inner_channels), 74 | activation=nn.ReLU(inplace=True), 75 | ), 76 | ConvModule( 77 | self.inner_channels, 78 | self.inner_channels, 79 | 3, 80 | padding=1, 81 | bias=False, 82 | norm=self.norm_layer(self.inner_channels), 83 | activation=nn.ReLU(inplace=False), 84 | ), 85 | ) 86 | 87 | def forward(self, feats: torch.Tensor, lowlevel_feats: torch.Tensor) -> torch.Tensor: 88 | aspp_feats = [aspp(feats) for aspp in self.aspp] 89 | aspp_feats = torch.cat(aspp_feats, dim=1) 90 | aspp_feats = self.aspp_fuse(aspp_feats) 91 | 92 | gp_feats = self.image_pool(feats) 93 | gp_feats = self.image_pool_fuse(gp_feats) 94 | 95 | aspp_feats += gp_feats.repeat(1, 1, aspp_feats.shape[2], aspp_feats.shape[3]) 96 | aspp_feats = self.fuse_bn(aspp_feats) 97 | aspp_feats = F.leaky_relu(aspp_feats) 98 | 99 | lowlevel_feats = self.lowlevel_path(lowlevel_feats) 100 | aspp_feats = F.interpolate(aspp_feats, size=lowlevel_feats.shape[2:], mode="bilinear", align_corners=True) 101 | out_feats = torch.cat((aspp_feats, lowlevel_feats), dim=1) 102 | 103 | out_feats = self.last_conv(out_feats) 104 | 105 | return out_feats 106 | -------------------------------------------------------------------------------- /segmentation/gvcore/dataset/sampler.py: -------------------------------------------------------------------------------- 1 | import itertools 2 | import torch 3 | import numpy as np 4 | from torch.utils.data.sampler import Sampler, BatchSampler 5 | 6 | import gvcore.utils.distributed as dist_utils 7 | 8 | 9 | class InfiniteSampler(Sampler): 10 | def __init__(self, size, seed, shuffle=True): 11 | self._size = size 12 | assert size > 0 13 | self._shuffle = shuffle 14 | self._seed = int(seed) 15 | 16 | self._rank = dist_utils.get_rank() 17 | self._world_size = dist_utils.get_world_size() 18 | 19 | def __iter__(self): 20 | start = self._rank 21 | yield from itertools.islice(self._infinite_indices(), start, None, self._world_size) 22 | 23 | def _infinite_indices(self): 24 | g = torch.Generator() 25 | g.manual_seed(self._seed) 26 | while True: 27 | if self._shuffle: 28 | yield from torch.randperm(self._size, generator=g).tolist() 29 | else: 30 | yield from torch.arange(self._size).tolist() 31 | 32 | 33 | class SequentialSampler(Sampler): 34 | """ 35 | Produce indices for inference. 36 | Inference needs to run on the __exact__ set of samples, 37 | therefore when the total number of samples is not divisible by the number of workers, 38 | this sampler produces different number of samples on different workers. 39 | """ 40 | 41 | def __init__(self, size: int): 42 | """ 43 | Args: 44 | size (int): the total number of data of the underlying dataset to sample from 45 | """ 46 | self._size = size 47 | assert size > 0 48 | self._rank = dist_utils.get_rank() 49 | self._world_size = dist_utils.get_world_size() 50 | 51 | shard_size = (self._size - 1) // self._world_size + 1 52 | begin = shard_size * self._rank 53 | end = min(shard_size * (self._rank + 1), self._size) 54 | self._local_indices = range(begin, end) 55 | 56 | def __iter__(self): 57 | yield from self._local_indices 58 | 59 | def __len__(self): 60 | return len(self._local_indices) 61 | 62 | 63 | class GroupedBatchSampler(BatchSampler): 64 | """ 65 | Wraps another sampler to yield a mini-batch of indices. 66 | It enforces that the batch only contain elements from the same group. 67 | It also tries to provide mini-batches which follows an ordering which is 68 | as close as possible to the ordering from the original sampler. 69 | """ 70 | 71 | def __init__(self, sampler, group_ids, batch_size): 72 | """ 73 | Args: 74 | sampler (Sampler): Base sampler. 75 | group_ids (list[int]): If the sampler produces indices in range [0, N), 76 | `group_ids` must be a list of `N` ints which contains the group id of each sample. 77 | The group ids must be a set of integers in the range [0, num_groups). 78 | batch_size (int): Size of mini-batch. 79 | """ 80 | if not isinstance(sampler, Sampler): 81 | raise ValueError( 82 | "sampler should be an instance of " 83 | "torch.utils.data.Sampler, but got sampler={}".format(sampler) 84 | ) 85 | self.sampler = sampler 86 | self.group_ids = np.asarray(group_ids) 87 | assert self.group_ids.ndim == 1 88 | self.batch_size = batch_size 89 | groups = np.unique(self.group_ids).tolist() 90 | 91 | # buffer the indices of each group until batch size is reached 92 | self.buffer_per_group = {k: [] for k in groups} 93 | 94 | def __iter__(self): 95 | for idx in self.sampler: 96 | group_id = self.group_ids[idx] 97 | group_buffer = self.buffer_per_group[group_id] 98 | group_buffer.append(idx) 99 | if len(group_buffer) == self.batch_size: 100 | yield group_buffer[:] # yield a copy of the list 101 | del group_buffer[:] 102 | 103 | def __len__(self): 104 | raise NotImplementedError("len() of GroupedBatchSampler is not well-defined.") -------------------------------------------------------------------------------- /detection/gvcore/utils/structure.py: -------------------------------------------------------------------------------- 1 | from collections import deque 2 | 3 | import random 4 | import torch 5 | import copy 6 | 7 | from gvcore.utils.distributed import all_gather 8 | 9 | 10 | class GenericData: 11 | def has(self, key): 12 | return hasattr(self, key) 13 | 14 | def set(self, key, value): 15 | setattr(self, key, value) 16 | 17 | def get(self, key): 18 | if hasattr(self, key): 19 | return getattr(self, key) 20 | else: 21 | return None 22 | 23 | def remove(self, key): 24 | if hasattr(self, key): 25 | delattr(self, key) 26 | 27 | def to(self, device): 28 | for key, value in self.__dict__.items(): 29 | if isinstance(value, torch.Tensor): 30 | self.set(key, value.to(device, non_blocking=True)) 31 | elif isinstance(value, GenericData): 32 | value.to(device) 33 | return self 34 | 35 | def clone(self): 36 | return copy.deepcopy(self) 37 | 38 | def _repr(self, parent=""): 39 | repr = "" 40 | for key, value in self.__dict__.items(): 41 | if isinstance(value, torch.Tensor): 42 | value_repr = f"{parent}{key}: {value.dtype}, {list(value.size())}\n" 43 | elif isinstance(value, GenericData): 44 | value_repr = value._repr(key + "." + parent) 45 | else: 46 | value_repr = f"{parent}{key}: {value}\n" 47 | repr += value_repr 48 | return repr 49 | 50 | def __repr__(self): 51 | return "\n(\n" + self._repr() + ")\n" 52 | 53 | 54 | class TensorQueue: 55 | def __init__(self, size): 56 | self.size = size 57 | 58 | self.queue = deque(maxlen=size) 59 | 60 | def push(self, x): 61 | self.queue.append(x.to("cpu")) 62 | 63 | def sample(self, n: int): 64 | n = min(n, len(self.queue)) 65 | return random.sample(self.queue, n) 66 | 67 | def sum(self): 68 | return torch.cat(list(self.queue)).sum() 69 | 70 | def mean(self): 71 | return torch.cat(list(self.queue)).mean() 72 | 73 | 74 | class TensorDict: 75 | def __init__(self): 76 | self.dict = {} 77 | 78 | def __missing__(self, key): 79 | return self[str(key)] 80 | 81 | def __contains__(self, key): 82 | return str(key) in self.dict 83 | 84 | def __setitem__(self, key, item): 85 | ks = all_gather(key) 86 | vs = all_gather(item) 87 | 88 | for k, v in zip(ks, vs): 89 | self.dict[str(k)] = v.cpu() 90 | 91 | def __getitem__(self, key): 92 | return self.dict[str(key)] 93 | 94 | def __str__(self) -> str: 95 | return str(self.dict) 96 | 97 | 98 | class TensorMatrix: 99 | def __init__(self, size, init_value=-1, device="cpu", dtype=torch.long): 100 | self.matrix = torch.zeros(size, device=device, dtype=dtype).fill_(init_value) 101 | 102 | def __setitem__(self, idx, item): 103 | ks = all_gather(idx) 104 | vs = all_gather(item) 105 | 106 | ks = torch.cat(ks) 107 | vs = torch.cat(vs) 108 | 109 | self.matrix.scatter_(0, ks.to(self.matrix.device), vs.to(self.matrix.device)) 110 | 111 | def __getitem__(self, idx): 112 | return self.matrix[idx.to(self.matrix.device)] 113 | 114 | 115 | class ConfusionMatrix: 116 | def __init__(self, num_classes, iter_n, device="cpu") -> None: 117 | self.m = torch.zeros((iter_n, num_classes, num_classes), device=device) 118 | self.i = 0 119 | 120 | self.num_classes = num_classes 121 | 122 | @torch.no_grad() 123 | def push(self, cls_a, cls_b) -> None: 124 | cls_a_list = all_gather(cls_a) 125 | cls_b_list = all_gather(cls_b) 126 | cls_a = torch.cat(cls_a_list) 127 | cls_b = torch.cat(cls_b_list) 128 | i = self.i % self.m.shape[0] 129 | 130 | m = torch.zeros((self.num_classes, self.num_classes), device=cls_a.device) 131 | m[cls_a.long(), cls_b.long()] += 1 132 | 133 | self.m[i] = m 134 | 135 | self.i += 1 136 | 137 | @torch.no_grad() 138 | def summary(self): 139 | return self.m.sum(dim=0) 140 | -------------------------------------------------------------------------------- /detection/det/utils/visualizer.py: -------------------------------------------------------------------------------- 1 | from PIL import ImageFont, ImageDraw 2 | from dataset.transforms import Denormalize 3 | from torchvision.transforms.functional import to_pil_image 4 | from utils.meta import * 5 | from gvcore.utils.structure import GenericData 6 | import matplotlib.pyplot as plt 7 | import numpy as np 8 | 9 | 10 | class Visualizer: 11 | def __init__(self, meta=COCOMeta, form="x1,y1,x2,y2,cls,score", t=0.5, t_key="score", show_tag=True): 12 | self.meta = meta() 13 | self.form = form 14 | self.t = t 15 | self.t_key = t_key 16 | self.show_tag = show_tag 17 | 18 | @staticmethod 19 | def draw_txt(draw, x, y, text, font): 20 | draw.text((x + 1, y + 1), text, (50, 50, 50), font=font) 21 | draw.text((x, y), text, (255, 255, 255), font=font) 22 | 23 | def __call__(self, img, bboxes, form=None, denorm=True): 24 | if denorm: 25 | img = Denormalize()(img) / 255.0 26 | img = to_pil_image(img.detach().cpu()) 27 | draw = ImageDraw.Draw(img) 28 | bboxes = bboxes.detach().cpu() 29 | 30 | form = self.form if form is None else form 31 | split_form = form.split(",") 32 | form_idx = {} 33 | score_keys = [] 34 | for i, key in enumerate(split_form): 35 | form_idx[key] = i 36 | if key not in "_,x,y,w,h,x1,y1,x2,y2,cls": 37 | score_keys.append(key) 38 | 39 | for box in bboxes: 40 | if "x,y,w,h" in form: 41 | x1 = box[form_idx["x"]] - int(box[form_idx["w"]] / 2) 42 | y1 = box[form_idx["y"]] - int(box[form_idx["h"]] / 2) 43 | x2 = box[form_idx["x"]] + int(box[form_idx["w"]] / 2) 44 | y2 = box[form_idx["y"]] + int(box[form_idx["h"]] / 2) 45 | elif "x1,y1,x2,y2" in form: 46 | x1 = box[form_idx["x1"]] 47 | y1 = box[form_idx["y1"]] 48 | x2 = box[form_idx["x2"]] 49 | y2 = box[form_idx["y2"]] 50 | else: 51 | raise ValueError 52 | 53 | if "score" in form_idx: 54 | score = box[form_idx["score"]] 55 | if score < self.t: 56 | continue 57 | 58 | if "cls" in form_idx: 59 | cls = int(box[form_idx["cls"]]) 60 | else: 61 | cls = -1 62 | 63 | meta = self.meta(cls) 64 | color = tuple(meta["color"]) 65 | cls_name = meta["name"] 66 | 67 | draw.rectangle(xy=[x1, y1, x2, y2], outline=color) 68 | 69 | if self.show_tag: 70 | draw.rectangle(xy=[x1, y1, x1 + 14 + 8 * len(cls_name) + 38 * len(score_keys), y1 - 18], fill=color) 71 | monospace = ImageFont.truetype("./assets/ShareTechMono.ttf", 14) 72 | 73 | self.draw_txt(draw, x1 + 5, y1 - 17, cls_name, monospace) 74 | 75 | for i, score_key in enumerate(score_keys): 76 | self.draw_txt( 77 | draw, 78 | x1 + 10 + 8 * len(cls_name) + i * 38, 79 | y1 - 17, 80 | "{:.2f}".format(box[form_idx[score_key]]), 81 | monospace, 82 | ) 83 | 84 | return img 85 | 86 | 87 | def vis(data_list, result_list=None, out_path=None, form="x1,y1,x2,y2,cls", threshold=0, meta=COCOMeta): 88 | 89 | visualizer = Visualizer(t=threshold, meta=meta) 90 | 91 | if isinstance(data_list, GenericData): 92 | data_list = [data_list] 93 | result_list = [result_list] 94 | elif isinstance(data_list, list): 95 | pass 96 | else: 97 | raise ValueError("data_list need to be a GenericData object or list of GenericData.") 98 | 99 | if result_list is None: 100 | result_list = [result_list] * len(data_list) 101 | 102 | imgs = [] 103 | for i, (data, result) in enumerate(zip(data_list, result_list)): 104 | img = data.img.cpu() 105 | if result is not None: 106 | box = result.cpu() 107 | else: 108 | box = data.label.cpu() 109 | img = visualizer(img, box, form=form) 110 | img = np.asarray(img) 111 | imgs.append(img) 112 | 113 | imgs = np.concatenate(imgs, axis=0) 114 | 115 | plt.figure(figsize=(imgs.shape[1] / 200, imgs.shape[0] / 200), dpi=200) 116 | plt.subplots_adjust(bottom=0, top=1, left=0, right=1) 117 | plt.imshow(imgs) 118 | 119 | if out_path is None: 120 | plt.show() 121 | else: 122 | plt.savefig(out_path) 123 | -------------------------------------------------------------------------------- /segmentation/seg/model/head/aspp.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | from gvcore.model import GenericModule 6 | 7 | from model.block import Normalize 8 | from model.block.block import ConvModule, DepthwiseSeparableConvModule 9 | 10 | 11 | class ASPPV3p(GenericModule): 12 | def __init__(self, cfg): 13 | super(ASPPV3p, self).__init__(cfg) 14 | 15 | # Params ============== 16 | self.in_channels = cfg.in_channels 17 | self.inner_channels = cfg.inner_channels 18 | 19 | self.lowlevel_in_channels = cfg.lowlevel_in_channels 20 | self.lowlevel_inner_channels = cfg.lowlevel_inner_channels 21 | 22 | self.dilations = cfg.dilations 23 | self.norm_layer = Normalize(cfg.norm_layer) 24 | 25 | # Components ========== 26 | self.aspp = nn.ModuleList( 27 | [ 28 | DepthwiseSeparableConvModule( 29 | self.in_channels, 30 | self.inner_channels, 31 | 3, 32 | dilation=dilation, 33 | padding=dilation, 34 | dw_norm=self.norm_layer(self.in_channels), 35 | pw_norm=self.norm_layer(self.inner_channels), 36 | activation=nn.ReLU(inplace=True), 37 | ) 38 | for dilation in self.dilations 39 | ] 40 | ) 41 | self.aspp[0] = ConvModule( 42 | self.in_channels, 43 | self.inner_channels, 44 | 1 if self.dilations[0] == 1 else 3, 45 | dilation=self.dilations[0], 46 | padding=0 if self.dilations[0] == 1 else self.dilations[0], 47 | bias=False, 48 | norm=self.norm_layer(self.inner_channels), 49 | activation=nn.ReLU(inplace=True), 50 | ) 51 | self.image_pool = nn.Sequential( 52 | nn.AdaptiveAvgPool2d(1), 53 | ConvModule( 54 | self.in_channels, 55 | self.inner_channels, 56 | 1, 57 | bias=False, 58 | norm=self.norm_layer(self.inner_channels), 59 | activation=nn.ReLU(inplace=True), 60 | ), 61 | ) 62 | self.aspp_fuse = ConvModule( 63 | (len(self.dilations) + 1) * self.inner_channels, 64 | self.inner_channels, 65 | 3, 66 | padding=1, 67 | bias=False, 68 | norm=self.norm_layer(self.inner_channels), 69 | activation=nn.ReLU(inplace=True), 70 | ) 71 | 72 | self.lowlevel_path = ConvModule( 73 | self.lowlevel_in_channels, 74 | self.lowlevel_inner_channels, 75 | 1, 76 | bias=False, 77 | norm=self.norm_layer(self.lowlevel_inner_channels), 78 | activation=nn.ReLU(inplace=True), 79 | ) 80 | 81 | self.last_conv = nn.Sequential( 82 | DepthwiseSeparableConvModule( 83 | self.inner_channels + self.lowlevel_inner_channels, 84 | self.inner_channels, 85 | 3, 86 | padding=1, 87 | bias=False, 88 | dw_norm=self.norm_layer(self.inner_channels + self.lowlevel_inner_channels), 89 | pw_norm=self.norm_layer(self.inner_channels), 90 | activation=nn.ReLU(inplace=True), 91 | ), 92 | DepthwiseSeparableConvModule( 93 | self.inner_channels, 94 | self.inner_channels, 95 | 3, 96 | padding=1, 97 | bias=False, 98 | dw_norm=self.norm_layer(self.inner_channels), 99 | pw_norm=self.norm_layer(self.inner_channels), 100 | activation=nn.ReLU(inplace=True), 101 | ), 102 | ) 103 | 104 | def forward(self, feats: torch.Tensor, lowlevel_feats: torch.Tensor) -> torch.Tensor: 105 | aspp_feats = [] 106 | gp_feats = self.image_pool(feats) 107 | aspp_feats.append(F.interpolate(gp_feats, size=feats.shape[2:], mode="bilinear", align_corners=False)) 108 | aspp_feats.extend([aspp(feats) for aspp in self.aspp]) 109 | 110 | aspp_feats = torch.cat(aspp_feats, dim=1) 111 | aspp_feats = self.aspp_fuse(aspp_feats) 112 | 113 | lowlevel_feats = self.lowlevel_path(lowlevel_feats) 114 | aspp_feats = F.interpolate(aspp_feats, size=lowlevel_feats.shape[2:], mode="bilinear", align_corners=False) 115 | out_feats = torch.cat((aspp_feats, lowlevel_feats), dim=1) 116 | 117 | out_feats = self.last_conv(out_feats) 118 | 119 | return out_feats 120 | -------------------------------------------------------------------------------- /segmentation/gvcore/utils/structure.py: -------------------------------------------------------------------------------- 1 | from collections import deque 2 | from addict import Dict 3 | import random 4 | import torch 5 | import copy 6 | 7 | from gvcore.utils.distributed import all_gather 8 | 9 | 10 | class GenericData: 11 | img: torch.Tensor 12 | label: torch.Tensor 13 | meta: Dict 14 | 15 | def has(self, key): 16 | return hasattr(self, key) 17 | 18 | def set(self, key, value): 19 | setattr(self, key, value) 20 | 21 | def get(self, key): 22 | if hasattr(self, key): 23 | return getattr(self, key) 24 | else: 25 | return None 26 | 27 | def remove(self, key): 28 | if hasattr(self, key): 29 | delattr(self, key) 30 | 31 | def to(self, device): 32 | for key, value in self.__dict__.items(): 33 | if isinstance(value, torch.Tensor): 34 | self.set(key, value.to(device, non_blocking=True)) 35 | elif isinstance(value, GenericData): 36 | value.to(device) 37 | return self 38 | 39 | def clone(self): 40 | return copy.deepcopy(self) 41 | 42 | def _repr(self, parent=""): 43 | repr = "" 44 | for key, value in self.__dict__.items(): 45 | if isinstance(value, torch.Tensor): 46 | value_repr = f"{parent}{key}: {value.dtype}, {list(value.size())}\n" 47 | elif isinstance(value, GenericData): 48 | value_repr = value._repr(key + "." + parent) 49 | else: 50 | value_repr = f"{parent}{key}: {value}\n" 51 | repr += value_repr 52 | return repr 53 | 54 | def __repr__(self): 55 | return "\n(\n" + self._repr() + ")\n" 56 | 57 | 58 | class TensorQueue: 59 | def __init__(self, size): 60 | self.size = size 61 | 62 | self.queue = deque(maxlen=size) 63 | 64 | def push(self, x): 65 | self.queue.append(x.to("cpu")) 66 | 67 | def sample(self, n: int): 68 | n = min(n, len(self.queue)) 69 | return random.sample(self.queue, n) 70 | 71 | def sum(self): 72 | return torch.cat(list(self.queue)).sum() 73 | 74 | def mean(self): 75 | return torch.cat(list(self.queue)).mean() 76 | 77 | 78 | class TensorDict: 79 | def __init__(self): 80 | self.dict = {} 81 | 82 | def __missing__(self, key): 83 | return self[str(key)] 84 | 85 | def __contains__(self, key): 86 | return str(key) in self.dict 87 | 88 | def __setitem__(self, key, item): 89 | ks = all_gather(key) 90 | vs = all_gather(item) 91 | 92 | for k, v in zip(ks, vs): 93 | self.dict[str(k)] = v.cpu() 94 | 95 | def __getitem__(self, key): 96 | return self.dict[str(key)] 97 | 98 | def __str__(self) -> str: 99 | return str(self.dict) 100 | 101 | 102 | class TensorList: 103 | def __init__(self, size, init_value=-1, device="cpu", dtype=torch.long, all_gather=True): 104 | self._tensor = torch.zeros(size, device=device, dtype=dtype).fill_(init_value) 105 | self.all_gather = all_gather 106 | 107 | def __setitem__(self, idx, item): 108 | if self.all_gather: 109 | ks = all_gather(idx) 110 | vs = all_gather(item) 111 | 112 | ks = torch.cat(ks) 113 | vs = torch.cat(vs) 114 | else: 115 | ks = idx 116 | vs = item 117 | 118 | self._tensor.index_copy_(0, ks.to(self._tensor.device), vs.to(dtype=self._tensor.dtype, device=self._tensor.device)) 119 | 120 | def __getitem__(self, idx): 121 | return self.matrix[idx.to(self.matrix.device)] 122 | 123 | 124 | class TensorMatrix: 125 | def __init__(self, size, init_value=-1, device="cpu", dtype=torch.long): 126 | self.matrix = torch.zeros(size, device=device, dtype=dtype).fill_(init_value) 127 | 128 | def __setitem__(self, idx, item): 129 | ks = all_gather(idx) 130 | vs = all_gather(item) 131 | 132 | ks = torch.cat(ks) 133 | vs = torch.cat(vs) 134 | 135 | self.matrix.scatter_(0, ks.to(self.matrix.device), vs.to(self.matrix.device)) 136 | 137 | def __getitem__(self, idx): 138 | return self.matrix[idx.to(self.matrix.device)] 139 | 140 | 141 | class ConfusionMatrix: 142 | def __init__(self, num_classes, iter_n, device="cpu") -> None: 143 | self.m = torch.zeros((iter_n, num_classes, num_classes), device=device) 144 | self.i = 0 145 | 146 | self.num_classes = num_classes 147 | 148 | @torch.no_grad() 149 | def push(self, cls_a, cls_b) -> None: 150 | cls_a_list = all_gather(cls_a) 151 | cls_b_list = all_gather(cls_b) 152 | cls_a = torch.cat(cls_a_list) 153 | cls_b = torch.cat(cls_b_list) 154 | i = self.i % self.m.shape[0] 155 | 156 | m = torch.zeros((self.num_classes, self.num_classes), device=cls_a.device) 157 | m[cls_a.long(), cls_b.long()] += 1 158 | 159 | self.m[i] = m 160 | 161 | self.i += 1 162 | 163 | @torch.no_grad() 164 | def summary(self): 165 | return self.m.sum(dim=0) 166 | -------------------------------------------------------------------------------- /detection/gvcore/utils/config.py: -------------------------------------------------------------------------------- 1 | import yaml 2 | from easydict import EasyDict 3 | import argparse 4 | import os 5 | from mergedeep import merge, Strategy 6 | 7 | from gvcore.utils.distributed import ( 8 | _find_free_port, 9 | get_dist_url_from_slurm, 10 | get_global_rank_from_slurm, 11 | get_local_rank_from_slurm, 12 | ) 13 | 14 | 15 | def parse_yaml(yaml_path): 16 | with open(yaml_path, "r") as yaml_file: 17 | cfg = yaml.load(yaml_file, Loader=yaml.Loader) 18 | if "_base" in cfg: 19 | yaml_path = os.path.dirname(yaml_path) 20 | base_path = os.path.join(yaml_path, cfg["_base"]) 21 | with open(base_path, "r") as base_file: 22 | base_cfg = yaml.load(base_file, Loader=yaml.Loader) 23 | else: 24 | base_cfg = {} 25 | merge(base_cfg, cfg, strategy=Strategy.ADDITIVE) 26 | 27 | return EasyDict(base_cfg) 28 | 29 | 30 | def parse_config(args): 31 | assert args.config.endswith(".yaml"), f"Only support YAML config file! Wrong config path: {args.config}" 32 | 33 | cfg = parse_yaml(args.config) 34 | 35 | if args.resume is not None: 36 | assert args.resume.endswith(".pth"), f"Only support resume from .pth file! Wrong weight path: {args.resume}" 37 | cfg.resume = args.resume 38 | else: 39 | cfg.resume = None 40 | if args.dev or args.opt == "test": 41 | cfg.dev_mode = True 42 | 43 | cfg.opt = args.opt 44 | 45 | if args.modify is not None: 46 | modifies = args.modify 47 | for modify in modifies: 48 | k, v = modify.split("=") 49 | if not v.replace(".", "", 1).isdigit() and v != "True" and v != "False": 50 | exec(f"cfg.{k}='{v}'") 51 | elif v.replace(".", "", 1).isdigit(): 52 | exec(f"cfg.{k}=float({v})") 53 | else: 54 | exec(f"cfg.{k}={v}") 55 | 56 | cfg.distributed.use = args.world_size > 1 57 | if cfg.distributed.use: 58 | cfg.distributed.dist_url = args.dist_url 59 | cfg.distributed.num_gpus = args.num_gpus 60 | cfg.distributed.world_size = args.world_size 61 | cfg.distributed.dist_local_rank = args.dist_local_rank 62 | cfg.distributed.dist_global_rank = args.dist_global_rank 63 | 64 | return cfg 65 | 66 | 67 | def config_to_string(cfg_dict, level=1): 68 | all_string = "" 69 | for k, v in cfg_dict.items(): 70 | if isinstance(v, dict): 71 | all_string += f"{' | '*level}{k}:\n" 72 | all_string += config_to_string(v, level=level + 1) 73 | else: 74 | all_string += f"{' | '*level}{k}: {v}\n" 75 | return all_string 76 | 77 | 78 | def parse_args(parser=None): 79 | if parser is None: 80 | parser = argparse.ArgumentParser() 81 | parser.add_argument("opt", type=str, help="Operation.") 82 | parser.add_argument("--launch", action="store_true", help="To launch operations.") 83 | parser.add_argument("--dev", action="store_true", help="To use dev mode.") 84 | parser.add_argument("--config", type=str, default="", help="Path to config file.") 85 | parser.add_argument("--modify", nargs="+", default=None, help="Modify config file.") 86 | parser.add_argument( 87 | "--resume", type=str, default=None, help="Whether to attempt to resume from the checkpoint directory." 88 | ) 89 | parser.add_argument("--num_gpus", type=int, default=1, help="GPU nums.") 90 | parser.add_argument("--dist_url", default=None, type=str, help="URL for initilizing distributed running.") 91 | parser.add_argument("--world_size", default=-1, type=int, help="World size.") 92 | parser.add_argument("--dist_global_rank", default=-1, type=int, help="Distributed Node Rank.") 93 | parser.add_argument("--dist_local_rank", default=-1, type=int, help="Distributed Local Rank.") 94 | 95 | parser.add_argument( 96 | "--slurm", 97 | default=False, 98 | action="store_true", 99 | help="SLURM flag. If set to True, use some environment variable to complete the args.", 100 | ) 101 | parser.add_argument("--slurm_offset", default=0, type=int, help="Distributed Rank Offset.") 102 | args = parser.parse_args() 103 | 104 | # For SLURM 105 | if args.world_size == -1: 106 | args.world_size = args.num_gpus 107 | if args.slurm: 108 | if args.dist_url is None: 109 | args.dist_url = get_dist_url_from_slurm() 110 | try: 111 | if args.dist_global_rank == -1: 112 | args.dist_global_rank = get_global_rank_from_slurm() + args.slurm_offset 113 | if args.dist_local_rank == -1: 114 | args.dist_local_rank = get_local_rank_from_slurm() 115 | except Exception as e: 116 | print(f"Cannot launch distribution training: {e}") 117 | 118 | args.num_gpus = 1 119 | else: 120 | if args.dist_url is None: 121 | try: 122 | master_node = "127.0.0.1" 123 | master_port = _find_free_port() 124 | args.dist_url = f"tcp://{master_node}:{master_port}" 125 | except Exception as e: 126 | print(f"Cannot construct dist_url: {e}") 127 | return args 128 | -------------------------------------------------------------------------------- /detection/det/model/detector/ts_cross.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from gvcore.utils.types import TDataList 4 | from gvcore.utils.misc import merge_loss_dict, split_ab 5 | from gvcore.model import GenericModule 6 | 7 | from dataset.transforms import DataCopy, parse_transform_config 8 | 9 | 10 | class TSCrossDetector(GenericModule): 11 | """ 12 | Teacher-Student detector for semi-supervised learning. 13 | """ 14 | 15 | def __init__(self, cfg, detector_class: GenericModule): 16 | super(TSCrossDetector, self).__init__(cfg) 17 | 18 | # Components ======== 19 | self.detector_a = detector_class(cfg) 20 | self.ema_detector_a = detector_class(cfg) 21 | self.detector_b = detector_class(cfg) 22 | self.ema_detector_b = detector_class(cfg) 23 | 24 | # Tools ============= 25 | self.weak_aug = DataCopy(parse_transform_config(cfg.data.weak_transforms)) 26 | self.strong_aug = DataCopy(parse_transform_config(cfg.data.strong_transforms)) 27 | 28 | # Initialize ======== 29 | for param_q, param_k in zip(self.detector_a.parameters(), self.ema_detector_a.parameters()): 30 | param_k.data.copy_(param_q.data) 31 | param_k.requires_grad = False 32 | self.ema_detector_a.freeze(True) 33 | self.ema_detector_a.eval() 34 | 35 | for param_q, param_k in zip(self.detector_b.parameters(), self.ema_detector_b.parameters()): 36 | param_k.data.copy_(param_q.data) 37 | param_k.requires_grad = False 38 | self.ema_detector_b.freeze(True) 39 | self.ema_detector_b.eval() 40 | 41 | # Params ============== 42 | self.m = cfg.solver.ema_momentum 43 | self.pseudo_t = cfg.model.teacher.score_threshold 44 | 45 | def train(self, mode: bool = True): 46 | self.training = mode 47 | self.detector_a.train(mode) 48 | self.detector_b.train(mode) 49 | 50 | def forward(self, *args, **kwargs): 51 | if "stage" not in kwargs: 52 | stage = "train_l" if self.training else "eval" 53 | else: 54 | stage = kwargs.pop("stage") 55 | 56 | if stage == "train_l": 57 | return self.forward_train_l(*args, **kwargs) 58 | elif stage == "train_u": 59 | return self.forward_train_u(*args, **kwargs) 60 | elif stage == "eval": 61 | return self.forward_eval(*args, **kwargs) 62 | 63 | def forward_train_l(self, data_list_l): 64 | weak_data_list_l = self.weak_aug(data_list_l) 65 | strong_data_list_l = self.strong_aug(data_list_l) 66 | weak_data_list_l_a, weak_data_list_l_b = split_ab(weak_data_list_l) 67 | strong_data_list_l_a, strong_data_list_l_b = split_ab(strong_data_list_l) 68 | 69 | data_list_l_a = weak_data_list_l_a + strong_data_list_l_a 70 | data_list_l_b = weak_data_list_l_b + strong_data_list_l_b 71 | 72 | loss_dict_l_a = self.detector_a(data_list_l_a, labeled=True) 73 | loss_dict_l_b = self.detector_b(data_list_l_b, labeled=True) 74 | 75 | loss_dict_l = merge_loss_dict(loss_dict_l_a, loss_dict_l_b) 76 | return loss_dict_l 77 | 78 | def forward_train_u(self, data_list_u: TDataList): 79 | self._momentum_update(m=self.m) 80 | 81 | weak_data_list_u = self.weak_aug(data_list_u) 82 | strong_data_list_u = self.strong_aug(data_list_u) 83 | weak_data_list_u_a, weak_data_list_u_b = split_ab(weak_data_list_u) 84 | strong_data_list_u_a, strong_data_list_u_b = split_ab(strong_data_list_u) 85 | 86 | pseudo_labels_a = self.ema_detector_b(weak_data_list_u_a) 87 | for data, pseudo_label in zip(strong_data_list_u_a, pseudo_labels_a): 88 | select = pseudo_label[:, 5] > self.pseudo_t 89 | data.label = torch.narrow(pseudo_label[select], dim=1, start=0, length=5) 90 | pseudo_labels_b = self.ema_detector_a(weak_data_list_u_b) 91 | for data, pseudo_label in zip(strong_data_list_u_b, pseudo_labels_b): 92 | select = pseudo_label[:, 5] > self.pseudo_t 93 | data.label = torch.narrow(pseudo_label[select], dim=1, start=0, length=5) 94 | 95 | loss_dict_u_a = self.detector_a(strong_data_list_u_a, labeled=False) 96 | loss_dict_u_b = self.detector_b(strong_data_list_u_b, labeled=False) 97 | 98 | loss_dict_u = merge_loss_dict(loss_dict_u_a, loss_dict_u_b) 99 | 100 | return loss_dict_u 101 | 102 | def forward_eval(self, *args, **kwargs): 103 | selector = kwargs.pop("selector") 104 | if selector == "student": 105 | return self.detector_a(*args, **kwargs) 106 | elif selector == "teacher": 107 | return self.ema_detector_a(*args, **kwargs) 108 | 109 | @torch.no_grad() 110 | def _momentum_update(self, m: float): 111 | """ 112 | Momentum update of the key encoder 113 | """ 114 | for param_q, param_k in zip(self.detector_a.parameters(), self.ema_detector_a.parameters()): 115 | param_k.data = param_k.data * m + param_q.data * (1.0 - m) 116 | for param_q, param_k in zip(self.detector_b.parameters(), self.ema_detector_b.parameters()): 117 | param_k.data = param_k.data * m + param_q.data * (1.0 - m) 118 | -------------------------------------------------------------------------------- /detection/det/ext/setup.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # Copyright (c) Facebook, Inc. and its affiliates. 3 | 4 | import glob 5 | import os 6 | import shutil 7 | from os import path 8 | from setuptools import find_packages, setup 9 | from typing import List 10 | import torch 11 | from torch.utils.cpp_extension import CUDA_HOME, CppExtension, CUDAExtension 12 | from torch.utils.hipify import hipify_python 13 | 14 | torch_ver = [int(x) for x in torch.__version__.split(".")[:2]] 15 | assert torch_ver >= [1, 5], "Requires PyTorch >= 1.5" 16 | 17 | 18 | def get_version(): 19 | init_py_path = path.join(path.abspath(path.dirname(__file__)), "fastevalapi", "__init__.py") 20 | init_py = open(init_py_path, "r").readlines() 21 | version_line = [l.strip() for l in init_py if l.startswith("__version__")][0] 22 | version = version_line.split("=")[-1].strip().strip("'\"") 23 | 24 | # The following is used to build release packages. 25 | # Users should never use it. 26 | suffix = os.getenv("D2_VERSION_SUFFIX", "") 27 | version = version + suffix 28 | if os.getenv("BUILD_NIGHTLY", "0") == "1": 29 | from datetime import datetime 30 | 31 | date_str = datetime.today().strftime("%y%m%d") 32 | version = version + ".dev" + date_str 33 | 34 | new_init_py = [l for l in init_py if not l.startswith("__version__")] 35 | new_init_py.append('__version__ = "{}"\n'.format(version)) 36 | with open(init_py_path, "w") as f: 37 | f.write("".join(new_init_py)) 38 | return version 39 | 40 | 41 | def get_extensions(): 42 | this_dir = path.dirname(path.abspath(__file__)) 43 | extensions_dir = path.join(this_dir, "fastevalapi", "csrc") 44 | 45 | main_source = path.join(extensions_dir, "vision.cpp") 46 | sources = glob.glob(path.join(extensions_dir, "**", "*.cpp")) 47 | 48 | from torch.utils.cpp_extension import ROCM_HOME 49 | 50 | is_rocm_pytorch = True if ((torch.version.hip is not None) and (ROCM_HOME is not None)) else False 51 | 52 | if is_rocm_pytorch: 53 | hipify_python.hipify( 54 | project_directory=this_dir, 55 | output_directory=this_dir, 56 | includes="/fastevalapi/csrc/*", 57 | show_detailed=True, 58 | is_pytorch_extension=True, 59 | ) 60 | 61 | # Current version of hipify function in pytorch creates an intermediate directory 62 | # named "hip" at the same level of the path hierarchy if a "cuda" directory exists, 63 | # or modifying the hierarchy, if it doesn't. Once pytorch supports 64 | # "same directory" hipification (https://github.com/pytorch/pytorch/pull/40523), 65 | # the source_cuda will be set similarly in both cuda and hip paths, and the explicit 66 | # header file copy (below) will not be needed. 67 | source_cuda = glob.glob(path.join(extensions_dir, "**", "hip", "*.hip")) + glob.glob( 68 | path.join(extensions_dir, "hip", "*.hip") 69 | ) 70 | else: 71 | source_cuda = glob.glob(path.join(extensions_dir, "**", "*.cu")) + glob.glob(path.join(extensions_dir, "*.cu")) 72 | 73 | sources = [main_source] + sources 74 | sources = [s for s in sources if not is_rocm_pytorch or torch_ver < [1, 7] or not s.endswith("hip/vision.cpp")] 75 | 76 | extension = CppExtension 77 | 78 | extra_compile_args = {"cxx": []} 79 | define_macros = [] 80 | 81 | if (torch.cuda.is_available() and ((CUDA_HOME is not None) or is_rocm_pytorch)) or os.getenv( 82 | "FORCE_CUDA", "0" 83 | ) == "1": 84 | extension = CUDAExtension 85 | sources += source_cuda 86 | 87 | if not is_rocm_pytorch: 88 | define_macros += [("WITH_CUDA", None)] 89 | extra_compile_args["nvcc"] = [ 90 | "-O3", 91 | "-DCUDA_HAS_FP16=1", 92 | "-D__CUDA_NO_HALF_OPERATORS__", 93 | "-D__CUDA_NO_HALF_CONVERSIONS__", 94 | "-D__CUDA_NO_HALF2_OPERATORS__", 95 | ] 96 | else: 97 | define_macros += [("WITH_HIP", None)] 98 | extra_compile_args["nvcc"] = [] 99 | 100 | if torch_ver < [1, 7]: 101 | # supported by https://github.com/pytorch/pytorch/pull/43931 102 | CC = os.environ.get("CC", None) 103 | if CC is not None: 104 | extra_compile_args["nvcc"].append("-ccbin={}".format(CC)) 105 | 106 | include_dirs = [extensions_dir] 107 | 108 | ext_modules = [ 109 | extension( 110 | "fastevalapi._C", 111 | sources, 112 | include_dirs=include_dirs, 113 | define_macros=define_macros, 114 | extra_compile_args=extra_compile_args, 115 | ) 116 | ] 117 | 118 | return ext_modules 119 | 120 | 121 | setup( 122 | name="fastevalapi", 123 | version=get_version(), 124 | author="FAIR", 125 | url="https://github.com/facebookresearch/fastevalapi", 126 | description="Detectron2 is FAIR's next-generation research " "platform for object detection and segmentation.", 127 | packages=find_packages(), 128 | python_requires=">=3.6", 129 | ext_modules=get_extensions(), 130 | cmdclass={"build_ext": torch.utils.cpp_extension.BuildExtension}, 131 | ) 132 | --------------------------------------------------------------------------------