├── cvcore ├── model │ ├── __init__.py │ └── model_zoo.py ├── config │ ├── __init__.py │ ├── multi_unet_Diceloss.yaml │ ├── multi_unet++_b0_diceloss.yaml │ ├── multi_FPN_b0_diceloss.yaml │ ├── multi_unet_b0_diceloss.yaml │ └── default.py ├── solver │ ├── __init__.py │ ├── build.py │ └── lr_scheduler.py ├── tools │ ├── __init__.py │ ├── args.py │ ├── train_tool.py │ └── valid_tool.py ├── utils │ ├── __init__.py │ ├── seed.py │ ├── comm.py │ ├── weight_init.py │ ├── logging.py │ └── checkpoint.py ├── unused │ └── batchstat_normalize.py └── data │ └── multi_rib_dataset.py ├── LICENSE ├── README.md └── main.py /cvcore/model/__init__.py: -------------------------------------------------------------------------------- 1 | from .model_zoo import build_model -------------------------------------------------------------------------------- /cvcore/config/__init__.py: -------------------------------------------------------------------------------- 1 | from .default import get_cfg_defaults -------------------------------------------------------------------------------- /cvcore/solver/__init__.py: -------------------------------------------------------------------------------- 1 | from .lr_scheduler import WarmupCyclicalLR, WarmupMultiStepLR 2 | from .build import make_optimizer -------------------------------------------------------------------------------- /cvcore/tools/__init__.py: -------------------------------------------------------------------------------- 1 | from .args import parse_args 2 | from .train_tool import train_loop 3 | from .valid_tool import valid_model 4 | -------------------------------------------------------------------------------- /cvcore/utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .checkpoint import save_checkpoint, load_checkpoint 2 | from .comm import synchronize 3 | from .logging import AverageMeter, setup_logger 4 | from .seed import setup_determinism 5 | from .weight_init import _initialize_weights -------------------------------------------------------------------------------- /cvcore/utils/seed.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import random 3 | import torch 4 | 5 | 6 | def setup_determinism(seed): 7 | torch.backends.cudnn.deterministic = True 8 | torch.backends.cudnn.benchmark = False 9 | torch.manual_seed(seed) 10 | np.random.seed(seed) 11 | random.seed(seed) -------------------------------------------------------------------------------- /cvcore/utils/comm.py: -------------------------------------------------------------------------------- 1 | import torch.distributed as dist 2 | 3 | 4 | def synchronize(): 5 | """ 6 | Helper function to synchronize (barrier) among all processes when 7 | using distributed training. 8 | """ 9 | if not dist.is_available(): 10 | return 11 | if not dist.is_initialized(): 12 | return 13 | world_size = dist.get_world_size() 14 | if world_size == 1: 15 | return 16 | dist.barrier() -------------------------------------------------------------------------------- /cvcore/unused/batchstat_normalize.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | def batchstat_norm(x): 4 | ''' 5 | Normalize tensor x of size (N,3,H,W) with batch mean and std. 6 | ''' 7 | mean = [] 8 | std = [] 9 | for i in range(x.shape[1]): 10 | xi = x[:,i,:,:] 11 | mean.append(xi.mean().repeat(x.shape[0], x.shape[2], x.shape[3]).unsqueeze(1)) 12 | std.append(xi.std().repeat(x.shape[0], x.shape[2], x.shape[3]).unsqueeze(1)) 13 | mean = torch.cat(mean, 1) 14 | std = torch.cat(std, 1) 15 | x = (x - mean) / std 16 | 17 | return x -------------------------------------------------------------------------------- /cvcore/config/multi_unet_Diceloss.yaml: -------------------------------------------------------------------------------- 1 | NAME: 'multi_unet_diceloss.yaml' 2 | 3 | 4 | 5 | DATA: 6 | TYPE: "multilabel" 7 | JSON: 8 | TRAIN: "data/train/VinDr_VinCXR_train_mask.json" 9 | VAL: "data/val/VinDr_VinCXR_val_mask.json" 10 | INP_CHANNEL: 1 11 | 12 | MODEL: 13 | NAME: "unet()" 14 | CLS_HEAD: "linear" 15 | NUM_CLASSES: 20 16 | 17 | LOSS: 18 | NAME: "dice" 19 | METRIC: 20 | NAME: "dice" 21 | SOLVER: 22 | ADAM_EPS: 1e-6 23 | WARMUP_LENGTH: 5 24 | GD_STEPS: 1 25 | SCHEDULER: 'none' 26 | MIN_LR: 1e-3 27 | BASE_LR: 1e-3 28 | 29 | TRAIN: 30 | BATCH_SIZE: 8 31 | EPOCHES: [200] 32 | 33 | SYSTEM: 34 | NUM_WORKERS: 8 35 | SEED: 27 36 | -------------------------------------------------------------------------------- /cvcore/config/multi_unet++_b0_diceloss.yaml: -------------------------------------------------------------------------------- 1 | NAME: 'multi_unet++(b0)_DiceLoss.yaml' 2 | 3 | 4 | 5 | DATA: 6 | TYPE: "multilabel" 7 | JSON: 8 | TRAIN: "data/train/VinDr_VinCXR_train_mask.json" 9 | VAL: "data/val/VinDr_VinCXR_val_mask.json" 10 | INP_CHANNEL: 1 11 | MODEL: 12 | NAME: "unet++(b0)" 13 | CLS_HEAD: "linear" 14 | NUM_CLASSES: 20 15 | 16 | LOSS: 17 | NAME: "dice" 18 | METRIC: 19 | NAME: "dice" 20 | SOLVER: 21 | ADAM_EPS: 1e-6 22 | WARMUP_LENGTH: 5 23 | GD_STEPS: 1 24 | SCHEDULER: 'step' 25 | MIN_LR: 1e-3 26 | BASE_LR: 1e-3 27 | 28 | TRAIN: 29 | BATCH_SIZE: 8 30 | EPOCHES: [200] 31 | 32 | SYSTEM: 33 | NUM_WORKERS: 8 34 | SEED: 27 35 | -------------------------------------------------------------------------------- /cvcore/config/multi_FPN_b0_diceloss.yaml: -------------------------------------------------------------------------------- 1 | NAME: 'multi_FPN_b0_diceloss.yaml' 2 | 3 | 4 | 5 | DATA: 6 | TYPE: "multilabel" 7 | JSON: 8 | TRAIN: "data/train/VinDr_VinCXR_train_mask.json" 9 | VAL: "data/val/VinDr_VinCXR_val_mask.json" 10 | INP_CHANNEL: 1 11 | 12 | MODEL: 13 | NAME: "fpn(b0)" 14 | CLS_HEAD: "linear" 15 | NUM_CLASSES: 20 16 | DROPOUT: 0.5 17 | 18 | LOSS: 19 | NAME: "dice" 20 | METRIC: 21 | NAME: "dice" 22 | SOLVER: 23 | ADAM_EPS: 1e-6 24 | WARMUP_LENGTH: 5 25 | GD_STEPS: 1 26 | SCHEDULER: 'step' 27 | MIN_LR: 1e-3 28 | BASE_LR: 1e-3 29 | 30 | TRAIN: 31 | BATCH_SIZE: 8 32 | EPOCHES: [200] 33 | 34 | SYSTEM: 35 | NUM_WORKERS: 8 36 | SEED: 27 37 | -------------------------------------------------------------------------------- /cvcore/utils/weight_init.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | 5 | def _initialize_weights(module): 6 | for m in module.modules(): 7 | if isinstance(m, nn.Conv2d): 8 | nn.init.xavier_uniform_(m.weight.data) 9 | if m.bias is not None: 10 | m.bias.data.zero_() 11 | elif isinstance(m, nn.BatchNorm2d): 12 | m.weight.data.fill_(1) 13 | m.bias.data.zero_() 14 | elif isinstance(m, nn.BatchNorm1d): 15 | m.weight.data.fill_(1) 16 | m.bias.data.zero_() 17 | elif isinstance(m, nn.Linear): 18 | nn.init.xavier_uniform_(m.weight.data) 19 | if m.bias is not None: 20 | m.bias.data.zero_() -------------------------------------------------------------------------------- /cvcore/solver/build.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | def make_optimizer(cfg, model): 5 | """ 6 | Create optimizer with per-layer learning rate and weight decay. 7 | """ 8 | params = [] 9 | for key, value in model.named_parameters(): 10 | if not value.requires_grad: 11 | continue 12 | lr = cfg.SOLVER.BASE_LR 13 | weight_decay = cfg.SOLVER.WEIGHT_DECAY 14 | if "bias" in key: 15 | weight_decay = cfg.SOLVER.WEIGHT_DECAY_BIAS 16 | params += [{"params": [value], "lr": lr, "weight_decay": weight_decay,'initial_lr':lr}] 17 | 18 | if cfg.SOLVER.OPTIMIZER == "adamw": 19 | optimizer = torch.optim.AdamW(params, lr, eps=cfg.SOLVER.ADAM_EPS) 20 | elif cfg.SOLVER.OPTIMIZER == "sgd": 21 | optimizer = torch.optim.SGD(params, lr, momentum=0.9, nesterov=True) 22 | return optimizer -------------------------------------------------------------------------------- /cvcore/config/multi_unet_b0_diceloss.yaml: -------------------------------------------------------------------------------- 1 | NAME: 'multi_unet_b0_DiceLoss.yaml' 2 | 3 | 4 | 5 | DATA: 6 | TYPE: "multilabel" 7 | JSON: 8 | TRAIN: "data/train/VinDr_VinCXR_train_mask.json" 9 | VAL: "data/val/VinDr_VinCXR_val_mask.json" 10 | # RANDAUG: 11 | # N: 2 12 | # M: 27 13 | 14 | #ALBU: 15 | #CONFIG: 3 16 | # IMG_SIZE: (240, 240) 17 | INP_CHANNEL: 1 18 | # CROP: 19 | # ENABLED: True 20 | # CROPSIZE: (448, 448) 21 | 22 | MODEL: 23 | NAME: "unet(b0)" 24 | CLS_HEAD: "linear" 25 | NUM_CLASSES: 20 26 | DROPOUT: 0.5 27 | 28 | LOSS: 29 | NAME: "dice" 30 | METRIC: 31 | NAME: "dice" 32 | SOLVER: 33 | # OPTIMIZER: "sgd" 34 | ADAM_EPS: 1e-6 35 | WARMUP_LENGTH: 5 36 | GD_STEPS: 1 37 | SCHEDULER: 'step' 38 | MIN_LR: 1e-3 39 | BASE_LR: 1e-3 40 | 41 | TRAIN: 42 | BATCH_SIZE: 8 43 | EPOCHES: [200] 44 | 45 | SYSTEM: 46 | NUM_WORKERS: 8 47 | SEED: 27 48 | -------------------------------------------------------------------------------- /cvcore/tools/args.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | 4 | def parse_args(): 5 | parser = argparse.ArgumentParser() 6 | parser.add_argument("--config", type=str, default="", 7 | help="config yaml path") 8 | parser.add_argument("--load", type=str, default="", 9 | help="path to model weight") 10 | parser.add_argument("--mode", type=str, default="train", 11 | help="model running mode (train/valid/test)") 12 | parser.add_argument("--reset", action="store_true", 13 | help="reset epoch") 14 | parser.add_argument("--clear", action="store_true", 15 | help="clear best metric") 16 | parser.add_argument("opts", default=None, 17 | help="Modify config options using the command-line", 18 | nargs=argparse.REMAINDER, 19 | ) 20 | 21 | args = parser.parse_args() 22 | 23 | return args 24 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2021 VinBigdata Medical 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /cvcore/utils/logging.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import os 3 | import sys 4 | 5 | 6 | def setup_logger(name, save_dir, distributed_rank, filename="log.txt"): 7 | """ 8 | Source: 9 | 10 | https://github.com/facebookresearch/maskrcnn-benchmark/blob/master/maskrcnn_benchmark/utils/logger.py 11 | """ 12 | logger = logging.getLogger(name) 13 | logger.setLevel(logging.DEBUG) 14 | # don't log results for the non-master process 15 | if distributed_rank > 0: 16 | return logger 17 | ch = logging.StreamHandler(stream=sys.stdout) 18 | ch.setLevel(logging.DEBUG) 19 | formatter = logging.Formatter("%(asctime)s %(name)s %(levelname)s: %(message)s") 20 | ch.setFormatter(formatter) 21 | logger.addHandler(ch) 22 | 23 | if save_dir: 24 | fh = logging.FileHandler(os.path.join(save_dir, filename)) 25 | fh.setLevel(logging.DEBUG) 26 | fh.setFormatter(formatter) 27 | logger.addHandler(fh) 28 | 29 | return logger 30 | 31 | 32 | class AverageMeter(object): 33 | """ 34 | Computes and stores the average and current value. 35 | """ 36 | def __init__(self): 37 | self.reset() 38 | 39 | def reset(self): 40 | self.val = 0 41 | self.avg = 0 42 | self.sum = 0 43 | self.count = 0 44 | 45 | def update(self, val, n=1): 46 | self.val = val 47 | self.sum += val * n 48 | self.count += n 49 | self.avg = self.sum / self.count -------------------------------------------------------------------------------- /cvcore/utils/checkpoint.py: -------------------------------------------------------------------------------- 1 | import os 2 | import shutil 3 | import torch 4 | import numpy as np 5 | 6 | 7 | def save_checkpoint(state, is_best, root, filename): 8 | """ 9 | Saves checkpoint and best checkpoint (optionally) 10 | """ 11 | torch.save(state, os.path.join(root, filename)) 12 | # if is_best: 13 | # shutil.copyfile( 14 | # os.path.join( 15 | # root, filename), os.path.join( 16 | # root, 'best_' + filename)) 17 | 18 | def load_checkpoint(args, log, model): 19 | if args.load != "": 20 | if os.path.isfile(args.load): 21 | log(f"=> loading checkpoint {args.load}") 22 | ckpt = torch.load(args.load, "cpu") 23 | model.load_state_dict(ckpt.pop('state_dict')) 24 | start_epoch, best_metric = ckpt['epoch'], ckpt['best_metric'] 25 | try: 26 | scheduler = ckpt['scheduler'] 27 | except: 28 | scheduler = None 29 | if args.reset: 30 | start_epoch = 0 31 | scheduler = None 32 | if args.clear: 33 | best_metric = np.inf 34 | log( 35 | f"=> loaded checkpoint '{args.load}' \ 36 | (epoch {ckpt['epoch']}, best_metric: {ckpt['best_metric']})") 37 | else: 38 | log(f"=> no checkpoint found at '{args.load}'") 39 | 40 | else: 41 | start_epoch = 0 42 | best_metric = 0 43 | scheduler = None 44 | 45 | return model, start_epoch, best_metric, scheduler -------------------------------------------------------------------------------- /cvcore/tools/train_tool.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn.functional as F 4 | import torch.nn as nn 5 | from torch.cuda.amp import autocast 6 | from tqdm import tqdm 7 | import os 8 | 9 | from cvcore.utils import AverageMeter, save_checkpoint 10 | from cvcore.solver import WarmupCyclicalLR, WarmupMultiStepLR 11 | from cvcore.model import build_model 12 | 13 | 14 | def train_loop(_print, cfg, model, train_loader, 15 | criterion, optimizer, scheduler, epoch, scaler): 16 | _print(f"\nEpoch {epoch + 1}") 17 | losses = AverageMeter() 18 | model.train() 19 | tbar = tqdm(train_loader) 20 | for i, (image, target) in enumerate(tbar): 21 | target[target!=0]=1 22 | target=target.long() 23 | image = image.to(device='cuda',dtype=torch.float) 24 | target = target.to(device='cuda') 25 | with autocast(): 26 | loss = criterion(model(image), target) 27 | # gradient accumulation 28 | loss = loss / cfg.SOLVER.GD_STEPS 29 | scaler.scale(loss).backward() 30 | # lr scheduler and optim. step 31 | if (i + 1) % cfg.SOLVER.GD_STEPS == 0: 32 | # optimizer.step() 33 | scaler.step(optimizer) 34 | optimizer.zero_grad() 35 | scaler.update() 36 | if isinstance(scheduler, WarmupCyclicalLR): 37 | scheduler(optimizer, i, epoch) 38 | elif isinstance(scheduler, WarmupMultiStepLR): 39 | scheduler.step() 40 | # record loss 41 | losses.update(loss.item() * cfg.SOLVER.GD_STEPS, target.size(0)) 42 | tbar.set_description("Train loss: %.5f, learning rate: %.6f" % ( 43 | losses.avg, optimizer.param_groups[-1]['lr'])) 44 | 45 | _print("Train loss: %.5f, learning rate: %.6f" % 46 | (losses.avg, optimizer.param_groups[-1]['lr'])) 47 | 48 | 49 | -------------------------------------------------------------------------------- /cvcore/tools/valid_tool.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn.functional as F 4 | from tqdm import tqdm 5 | import os 6 | from cvcore.utils import save_checkpoint 7 | from sklearn.metrics import f1_score, recall_score, precision_score 8 | def valid_model(_print, cfg, model, valid_loader, 9 | loss_function, metric_function, epoch, 10 | best_metric=None, checkpoint=False): 11 | # switch to evaluate mode 12 | model.eval() 13 | preds = [] 14 | labels = [] 15 | tbar = tqdm(valid_loader) 16 | 17 | with torch.no_grad(): 18 | for i, (image, label) in enumerate(tbar): 19 | image = image.to(device='cuda',dtype=torch.float) 20 | w_output = model(image)>0.5 21 | preds.append(w_output.cpu()) 22 | labels.append(label.cpu()) 23 | metric_name= cfg.METRIC.NAME 24 | preds, labels = torch.cat(preds, 0), torch.cat(labels, 0) 25 | val_loss = loss_function(preds.float(), labels) 26 | scores = [] 27 | final_score = metric_function(preds, labels.long())[0].item() 28 | print(final_score) 29 | 30 | _print(f"Validation {metric_name}: {final_score:04f}, val loss:{val_loss:05f} best: {best_metric:04f}\n") 31 | 32 | 33 | # checkpoint 34 | 35 | if checkpoint: 36 | is_best = final_score > best_metric 37 | best_metric = max(final_score, best_metric) 38 | save_dict = {"epoch": epoch + 1, 39 | "arch": cfg.NAME, 40 | "state_dict": model.state_dict(), 41 | "best_metric": best_metric} 42 | save_filename = f"{cfg.NAME}.pth" 43 | if is_best: # only save best checkpoint, no need resume 44 | print("score improved, saving new checkpoint...") 45 | save_checkpoint(save_dict, is_best, 46 | root=cfg.DIRS.WEIGHTS, filename=save_filename) 47 | return val_loss, best_metric -------------------------------------------------------------------------------- /cvcore/model/model_zoo.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from torch.nn.parameter import Parameter 5 | from torch.cuda.amp import autocast 6 | from collections import OrderedDict 7 | import math 8 | import os 9 | 10 | import torchvision 11 | import timm 12 | from timm.models.layers.adaptive_avgmax_pool import SelectAdaptivePool2d 13 | import segmentation_models_pytorch as smp 14 | from monai.networks.nets import Unet,BasicUnet 15 | 16 | def build_model(cfg): 17 | 18 | if 'unet()'== cfg.MODEL.NAME: 19 | model= Unet(dimensions=2,in_channels=cfg.DATA.INP_CHANNEL,out_channels=cfg.MODEL.NUM_CLASSES,\ 20 | channels=(16, 32, 64, 128, 256),strides=(2, 2, 2, 2),num_res_units=2,dropout=cfg.MODEL.DROPOUT) 21 | elif 'unet(resnet18)'== cfg.MODEL.NAME: 22 | model= smp.Unet('resnet18',classes=cfg.MODEL.NUM_CLASSES,encoder_weights='imagenet',in_channels=cfg.DATA.INP_CHANNEL) 23 | elif 'unet(resnet50)'== cfg.MODEL.NAME: 24 | model= smp.Unet('resnet50',classes=cfg.MODEL.NUM_CLASSES,encoder_weights='imagenet',in_channels=cfg.DATA.INP_CHANNEL) 25 | elif 'unet(resnet101)'== cfg.MODEL.NAME: 26 | model= smp.Unet('resnet101',classes=cfg.MODEL.NUM_CLASSES,encoder_weights='imagenet',in_channels=cfg.DATA.INP_CHANNEL) 27 | elif 'unet(densenet169)'== cfg.MODEL.NAME: 28 | model= smp.Unet('densenet169',classes=cfg.MODEL.NUM_CLASSES,encoder_weights='imagenet',in_channels=cfg.DATA.INP_CHANNEL) 29 | elif 'unet(densenet121)'== cfg.MODEL.NAME: 30 | model= smp.Unet('densenet121',classes=cfg.MODEL.NUM_CLASSES,encoder_weights='imagenet',in_channels=cfg.DATA.INP_CHANNEL) 31 | elif 'deeplabv3(resnet50)' == cfg.MODEL.NAME: 32 | model=smp.DeepLabV3('resnet50',classes=cfg.MODEL.NUM_CLASSES,in_channels=cfg.DATA.INP_CHANNEL) 33 | elif 'unet(b3)'== cfg.MODEL.NAME: 34 | model= smp.Unet('efficientnet-b3',classes=cfg.MODEL.NUM_CLASSES,encoder_weights='imagenet'\ 35 | ,in_channels=cfg.DATA.INP_CHANNEL) 36 | elif 'unet++()'== cfg.MODEL.NAME: 37 | model=NestedUNet(cfg) 38 | elif 'unet++(resnet101)'== cfg.MODEL.NAME: 39 | model=smp.UnetPlusPlus('resnet101',classes=cfg.MODEL.NUM_CLASSES,encoder_weights='imagenet'\ 40 | ,in_channels=cfg.DATA.INP_CHANNEL) 41 | elif 'unet++(b0)' ==cfg.MODEL.NAME: 42 | model=smp.UnetPlusPlus('efficientnet-b0',classes=cfg.MODEL.NUM_CLASSES,encoder_weights='imagenet'\ 43 | ,in_channels=cfg.DATA.INP_CHANNEL) 44 | elif('fpn(b0)'==cfg.MODEL.NAME): 45 | model=smp.FPN('efficientnet-b0',classes=cfg.MODEL.NUM_CLASSES,encoder_weights='imagenet',\ 46 | in_channels=cfg.DATA.INP_CHANNEL,decoder_dropout=cfg.MODEL.DROPOUT ) 47 | elif 'unet(b0)'== cfg.MODEL.NAME: 48 | model= smp.Unet('efficientnet-b0',classes=cfg.MODEL.NUM_CLASSES,encoder_weights='imagenet'\ 49 | ,in_channels=cfg.DATA.INP_CHANNEL) 50 | elif 'unet++(b3)'== cfg.MODEL.NAME: 51 | model= smp.UnetPlusPlus('efficientnet-b3',classes=cfg.MODEL.NUM_CLASSES,encoder_weights='imagenet'\ 52 | ,in_channels=cfg.DATA.INP_CHANNEL) 53 | return model -------------------------------------------------------------------------------- /cvcore/data/multi_rib_dataset.py: -------------------------------------------------------------------------------- 1 | import albumentations 2 | from albumentations import Compose, HorizontalFlip, Normalize, VerticalFlip, Rotate, Resize, ShiftScaleRotate, OneOf, GridDistortion, OpticalDistortion, \ 3 | ElasticTransform, IAAAdditiveGaussianNoise, GaussNoise, MedianBlur, Blur, CoarseDropout,RandomBrightnessContrast 4 | from albumentations.pytorch import ToTensorV2 5 | import torch.nn as nn 6 | import matplotlib.pyplot as plt 7 | import torch 8 | import os 9 | import numpy as np 10 | from pathlib import Path 11 | from PIL import Image 12 | import pandas as pd 13 | from torch.utils.data import ConcatDataset, Dataset, DataLoader, Subset 14 | import torchvision.transforms as transforms 15 | import torch.optim as optim 16 | import torchvision 17 | import random 18 | import cv2 19 | import time 20 | class multi_ribs_dataset(Dataset): 21 | def __init__(self,df,transforms,mode,list_label=[]): 22 | self.df=df 23 | self.transforms= transforms 24 | self.mode=mode 25 | self.list_label=list_label 26 | def __len__(self): 27 | return len(self.df['img']) 28 | def __getitem__(self,index): 29 | img_path=self.df['img'][index] 30 | img=Image.open(img_path) 31 | img=img.convert('L') 32 | img=np.asarray(img,dtype=np.float32) 33 | label0=[] 34 | for name in self.list_label: 35 | pts=self.df[name][index] 36 | label= np.zeros((img.shape[:2]),dtype=np.uint8) 37 | if pts!='None': 38 | pts= np.array([[[int(pt['x']),int(pt['y'])]] for pt in pts ]) 39 | label=cv2.fillPoly(label,[pts],255) 40 | label0.append(label) 41 | label0=np.stack(label0) 42 | img=img/255 43 | label0=label0/255 44 | label0=label0.transpose((1,2,0)) 45 | img=np.expand_dims(img,axis=2) 46 | dic=self.transforms[self.mode]\ 47 | ( image=img,mask=label0) 48 | img=dic['image'] 49 | mask=(dic['mask'].permute(2,0,1)) 50 | return img, mask 51 | def make_multi_ribs_dataloader(cfg,mode='train'): 52 | df_train= pd.read_json(cfg.DATA.JSON.TRAIN) 53 | df_val=pd.read_json(cfg.DATA.JSON.VAL) 54 | list_label=['R1','R2','R3','R4','R5','R6','R7','R8','R9','R10',\ 55 | 'L1','L2','L3','L4','L5','L6','L7','L8','L9','L10'] 56 | data_transform= { 57 | 'train': Compose([ 58 | Resize(512,512), 59 | HorizontalFlip(), 60 | ShiftScaleRotate(rotate_limit=10), 61 | RandomBrightnessContrast(), 62 | ToTensorV2() 63 | ]), 64 | 'val': Compose([ 65 | Resize(512,512), 66 | ToTensorV2() 67 | ]) 68 | } 69 | if mode=='train': 70 | ribs_dataset_train= multi_ribs_dataset(df_train,data_transform,mode,list_label) 71 | return (DataLoader(dataset=ribs_dataset_train,batch_size=cfg.TRAIN.BATCH_SIZE,shuffle=True)) 72 | elif mode=='val': 73 | ribs_dataset_val= multi_ribs_dataset(df_val,data_transform,mode,list_label) 74 | return (DataLoader(dataset=ribs_dataset_val,batch_size=1,shuffle=False)) -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # VinDr-RibCXR: A Benchmark Dataset for Automatic Segmentation and Labeling of Individual Ribs on Chest X-rays 2 | 3 | This repository contains the training code for our paper entitled "VinDr-RibCXR: A Benchmark Dataset for Automatic Segmentation and Labeling of Individual Ribs on Chest X-rays", which was submitted and under review by [Medical Imaging with Deep Learning 2021 (MIDL2021)](https://2021.midl.io/) 4 | # Abstract: 5 | We introduce a new benchmark dataset, namely VinDr-RibCXR, for automatic segmentation and labeling of individual ribs from chest X-ray (CXR) scans. The VinDr-RibCXR contains 245 CXRs with corresponding ground truth annotations provided by human experts. A set of state-of-the-art segmentation models are trained on 196 images from the VinDr-RibCXR to segment and label 20 individual ribs. Our best performing model obtains a Dice score of 0.834 (95% CI, 0.810{0.853) on an independent test set of 49 images. Our study, therefore, serves as a proof of concept and baseline performance for future research.\ 6 | Keywords: Rib segmentation, CXR, Benchmark dataset, Deep learning. 7 | 8 | # Data: 9 | - The dataset contains 245 images and split into 196 training images and 49 validation images. To download the VinDr-RibCXR dataset, please sign our [Data Use Agreement](https://drive.google.com/file/d/1Wr3iI7-OwZHD4eWtCALpRZKvhJuemKAs/view?fbclid=IwAR2lmFoe5JCqpkCVApIc_oXDnldJ21BGpib1PebC3GysrEkjfnqn-Wh2NE8) (DUA) and send the signed DUA to Ha Nguyen (nguyenquyha@gmail.com) for obtaining the downloadable link. 10 | - To record the label, we use json file, the json file have this format: 11 | ``` 12 | { 13 | "img": { 14 | "0": "data/train/img/VinDr_RibCXR_train_000.png", 15 | .... 16 | } 17 | "R1": { 18 | "0":[ 19 | { 20 | "x": 1027.8114701859, 21 | "y": 470.688105676 22 | }, 23 | .... 24 | ] 25 | } 26 | .... 27 | } 28 | ``` 29 | - To view the label, you can go to vinlab (https://lab.vindr.ai) or use the notebook visualize.ipynb 30 | - For further description, please go to this (https://vindr.ai/datasets/ribcxr) 31 | # Setting: 32 | We use Pytorch 1.7.1 for this project 33 | 34 | # Model: 35 | In this work, we train 4 baselines model, that is vanila U-net and U-net, FPN and U-net plus plus with imagenet pretrained encoder efficientet B0 from (https://github.com/qubvel/segmentation_models.pytorch) 36 | # Train model by your self: 37 | - We provide 4 configs file conresponding to 4 setting in our model, you can also write your config for your self 38 | - You need to put the data like this: 39 | ``` 40 | data 41 | ├── train 42 | | ├──img/ 43 | | └──Vindr_RibCXR_train_mask.json 44 | └── val 45 | ├──img/ 46 | └──Vindr_RibCXR_val_mask.json 47 | ``` 48 | - For training use this command: 49 | ``` 50 | python main.py --config/cvcore/multi_unet_b0_diceloss.yaml 51 | ``` 52 | 53 | # Result: 54 | 55 | 56 | | Model | Dice | 95% HD | Sensitivity | Specificity | 57 | |----------------------------|------------------|------------------------|------------------|------------------| 58 | | U-Net | .765 (.737-.788) | 28.038 (22.449-34.604) | .773 (.738-.791) | .996 (.996-.997) | 59 | | U-Net w. EfficientNet-B0 | .829 (.808-.847) | 16.807 (14.372-19.539) | .844 (.818-.858) | .998 (.997-.998) | 60 | | FPN w. EfficientNet-B0 | .807 (.773-.824) | 15.049 (13.190-16.953) | .808 (.773-.828) | .997 (.997-.998) | 61 | | U-Net++ w. EfficientNet-B0 | .834 (.810-.853) | 15.453 (13.340-17.450) | .841 (.812-.858) | .998 (.997-.998) | 62 | 63 | # Acknowledgements: 64 | This research was supported by the Vingroup Big Data Institute. We are especially thankful 65 | to Tien D. Phan, Dung T. Le, and Chau T.B. Pham for their help during the data annotation 66 | process -------------------------------------------------------------------------------- /cvcore/solver/lr_scheduler.py: -------------------------------------------------------------------------------- 1 | from bisect import bisect_right 2 | import math 3 | 4 | from torch.optim.lr_scheduler import _LRScheduler 5 | from torch.optim.lr_scheduler import MultiStepLR 6 | 7 | 8 | 9 | class WarmupMultiStepLR(MultiStepLR): 10 | """ 11 | Source: 12 | 13 | https://github.com/facebookresearch/maskrcnn-benchmark/blob/master/maskrcnn_benchmark/solver/lr_scheduler.py 14 | """ 15 | def __init__( 16 | self, 17 | optimizer, 18 | milestones, 19 | iter_per_epoch, 20 | gamma=0.5, 21 | warmup_factor=1.0 / 3, 22 | warmup_iters=500, 23 | warmup_method="linear", 24 | last_epoch=-1, 25 | ): 26 | 27 | 28 | if not list(milestones) == sorted(milestones): 29 | raise ValueError( 30 | "Milestones should be a list of" " increasing integers. Got {}", 31 | milestones, 32 | ) 33 | 34 | if warmup_method not in ("constant", "linear"): 35 | raise ValueError( 36 | "Only 'constant' or 'linear' warmup_method accepted" 37 | "got {}".format(warmup_method) 38 | ) 39 | self.milestones = [m * iter_per_epoch for m in milestones] 40 | self.gamma = gamma 41 | self.warmup_factor = warmup_factor 42 | self.warmup_iters = warmup_iters 43 | self.warmup_method = warmup_method 44 | super().__init__(optimizer=optimizer,milestones=milestones,last_epoch= last_epoch) 45 | 46 | def get_lr(self): 47 | warmup_factor = 1 48 | if self.last_epoch < self.warmup_iters: 49 | if self.warmup_method == "constant": 50 | warmup_factor = self.warmup_factor 51 | elif self.warmup_method == "linear": 52 | alpha = float(self.last_epoch) / self.warmup_iters 53 | warmup_factor = self.warmup_factor * (1 - alpha) + alpha 54 | return [ 55 | base_lr 56 | * warmup_factor 57 | * self.gamma ** bisect_right(self.milestones, self.last_epoch) 58 | for base_lr in self.base_lrs 59 | ] 60 | 61 | 62 | class WarmupCyclicalLR(object): 63 | """ 64 | Cyclical learning rate scheduler with linear warm-up. E.g.: 65 | 66 | Step mode: ``lr = base_lr * 0.1 ^ {floor(epoch-1 / lr_step)}``. 67 | 68 | Cosine mode: ``lr = base_lr * 0.5 * (1 + cos(iter/maxiter))``. 69 | 70 | Poly mode: ``lr = base_lr * (1 - iter/maxiter) ^ 0.9``. 71 | 72 | Arguments: 73 | mode (str): one of ('cos', 'poly', 'step'). 74 | base_lr (float): base optimizer's learning rate. 75 | num_epochs (int): number of epochs. 76 | iters_per_epoch (int): number of iterations (updates) per epoch. 77 | warmup_epochs (int): number of epochs to gradually increase learning rate from zero to base_lr. 78 | """ 79 | 80 | def __init__(self, mode, base_lr, num_epochs, iters_per_epoch=0, 81 | lr_step=0, warmup_epochs=0,min_lr=0.): 82 | self.mode = mode 83 | assert self.mode in ('cos', 'poly', 'step'), "Unsupported learning rate scheduler" 84 | 85 | self.lr = base_lr 86 | if mode == 'step': 87 | assert lr_step 88 | self.lr_step = lr_step 89 | self.iters_per_epoch = iters_per_epoch 90 | self.N = num_epochs * iters_per_epoch 91 | self.epoch = -1 92 | self.warmup_iters = int(warmup_epochs * iters_per_epoch) 93 | self.min_lr=min_lr 94 | def __call__(self, optimizer, i, epoch): 95 | T = epoch * self.iters_per_epoch + i 96 | if self.mode == 'cos': 97 | lr = self.min_lr+0.5 * (self.lr-self.min_lr) * (1 + math.cos(1.0 * T / self.N * math.pi)) 98 | elif self.mode == 'poly': 99 | lr = self.lr * pow((1 - 1.0 * T / self.N), 0.9) 100 | elif self.mode == 'step': 101 | lr = self.lr * (0.1 ** (epoch // self.lr_step)) 102 | 103 | # warm-up lr scheduler 104 | if self.warmup_iters > 0 and T < self.warmup_iters: 105 | lr = lr * 1.0 * T / self.warmup_iters 106 | 107 | assert lr >= 0 108 | self._adjust_learning_rate(optimizer, lr) 109 | 110 | def _adjust_learning_rate(self, optimizer, lr): 111 | for i in range(len(optimizer.param_groups)): 112 | optimizer.param_groups[i]['lr'] = lr 113 | -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | import gc 2 | import os 3 | import sys 4 | import time 5 | 6 | import pandas as pd 7 | import numpy as np 8 | import torch 9 | import torch.nn as nn 10 | import torch.nn.functional as F 11 | import torch.optim as optim 12 | from torch.cuda.amp import GradScaler 13 | from sklearn.metrics import accuracy_score 14 | from cvcore.config import get_cfg_defaults 15 | from cvcore.model import build_model 16 | from cvcore.solver import make_optimizer, WarmupCyclicalLR, WarmupMultiStepLR 17 | from cvcore.utils import setup_determinism, setup_logger, load_checkpoint 18 | from cvcore.tools import parse_args, train_loop, valid_model 19 | from cvcore.data.multi_rib_dataset import make_multi_ribs_dataloader, multi_ribs_dataset 20 | import segmentation_models_pytorch as smp 21 | from monai.metrics import DiceMetric 22 | from monai.losses import DiceLoss 23 | scaler = GradScaler() 24 | 25 | def main(args, cfg): 26 | # Set logger 27 | logger = setup_logger( 28 | args.mode, 29 | cfg.DIRS.LOGS, 30 | 0, 31 | filename=f"{cfg.NAME}.txt") 32 | # Declare variables 33 | best_metric = 0 34 | start_epoch = 0 35 | 36 | # Define model 37 | model = build_model(cfg) 38 | # Define optimizer 39 | optimizer = make_optimizer(cfg, model) 40 | 41 | # Define loss 42 | if cfg.LOSS.NAME == "ce": 43 | train_criterion = nn.CrossEntropyLoss() 44 | elif cfg.LOSS.NAME=='Bce': 45 | train_criterion= nn.BCEWithLogitsLoss() 46 | elif cfg.LOSS.NAME=='dice': 47 | train_criterion= DiceLoss(sigmoid=True) 48 | model = model.to(device='cuda') 49 | model = nn.DataParallel(model) 50 | 51 | # Load checkpoint 52 | model, start_epoch, best_metric, scheduler = load_checkpoint(args, logger.info, model) 53 | 54 | # Load and split data 55 | if cfg.DATA.TYPE== 'multilabel': 56 | if args.mode in ("train", "val"): 57 | valid_loader= make_multi_ribs_dataloader(cfg,mode='val') 58 | if args.mode=='train': 59 | train_loader= make_multi_ribs_dataloader(cfg,mode='train') 60 | elif args.mode == "test": 61 | #TODO: Write test dataloader. 62 | pass 63 | if args.mode == "train" and scheduler == None: 64 | if cfg.SOLVER.SCHEDULER == "cyclical": 65 | scheduler = WarmupCyclicalLR("cos", cfg.SOLVER.BASE_LR, cfg.TRAIN.EPOCHES[-1], 66 | iters_per_epoch=len(train_loader), 67 | warmup_epochs=cfg.SOLVER.WARMUP_LENGTH,min_lr=cfg.SOLVER.MIN_LR) 68 | elif cfg.SOLVER.SCHEDULER == "step": 69 | scheduler = WarmupMultiStepLR( 70 | optimizer=optimizer, 71 | milestones=cfg.TRAIN.EPOCHES[:-1], 72 | iter_per_epoch=len(train_loader), 73 | warmup_factor=cfg.SOLVER.BASE_LR/(cfg.SOLVER.WARMUP_LENGTH * len(train_loader)), 74 | warmup_iters=cfg.SOLVER.WARMUP_LENGTH * len(train_loader), 75 | last_epoch=start_epoch if start_epoch else -1) 76 | else: 77 | scheduler = None 78 | 79 | if args.mode != "test": 80 | valid_criterion = train_criterion 81 | if cfg.METRIC.NAME == 'dice': 82 | valid_metric = DiceMetric(include_background=True,reduction= 'mean') 83 | if args.mode == "train": 84 | for epoch in range(start_epoch, cfg.TRAIN.EPOCHES[-1]): 85 | train_loop(logger.info, cfg, model, 86 | train_loader, train_criterion, optimizer, 87 | scheduler, epoch, scaler) 88 | _, best_metric = valid_model(logger.info, cfg, model, 89 | valid_loader, valid_criterion, 90 | valid_metric, epoch, best_metric, True) 91 | elif args.mode == "val": 92 | valid_model(logger.info, cfg, model, 93 | valid_loader, valid_criterion, 94 | valid_metric, start_epoch) 95 | elif args.mode == "test": 96 | #TODO: Write test function. 97 | pass 98 | 99 | 100 | if __name__ == "__main__": 101 | args = parse_args() 102 | cfg = get_cfg_defaults() 103 | 104 | if args.config != "": 105 | cfg.merge_from_file(args.config) 106 | if args.opts != "": 107 | cfg.merge_from_list(args.opts) 108 | 109 | # make dirs 110 | for _dir in ["WEIGHTS", "OUTPUTS", "LOGS"]: 111 | if not os.path.isdir(cfg.DIRS[_dir]): 112 | os.mkdir(cfg.DIRS[_dir]) 113 | # seed, run 114 | setup_determinism(cfg.SYSTEM.SEED) 115 | main(args, cfg) -------------------------------------------------------------------------------- /cvcore/config/default.py: -------------------------------------------------------------------------------- 1 | from yacs.config import CfgNode as CN 2 | 3 | 4 | # Create root config node 5 | _C = CN() 6 | # Config name 7 | _C.NAME = "" 8 | # Config version to manage version of configuration names and default 9 | _C.VERSION = "0.1" 10 | 11 | 12 | # ---------------------------------------- 13 | # System config 14 | # ---------------------------------------- 15 | _C.SYSTEM = CN() 16 | 17 | # Number of workers for dataloader 18 | _C.SYSTEM.NUM_WORKERS = 8 19 | # Use GPU for training and inference. Default is True 20 | _C.SYSTEM.CUDA = True 21 | # Random seed for seeding everything (NumPy, Torch,...) 22 | _C.SYSTEM.SEED = 0 23 | # Use half floating point precision 24 | _C.SYSTEM.FP16 = True 25 | # FP16 Optimization level. See more at: https://nvidia.github.io/apex/amp.html#opt-levels 26 | _C.SYSTEM.OPT_L = "O2" 27 | 28 | 29 | # ---------------------------------------- 30 | # Directory name config 31 | # ---------------------------------------- 32 | _C.DIRS = CN() 33 | 34 | # Train, Validation and Testing image folders 35 | _C.DIRS.TRAIN_IMAGES = "" 36 | _C.DIRS.VALIDATION_IMAGES = "" 37 | _C.DIRS.TEST_IMAGES = "" 38 | # Trained weights folder 39 | _C.DIRS.WEIGHTS = "./weights/" 40 | # Inference output folder 41 | _C.DIRS.OUTPUTS = "./outputs/" 42 | # Training log folder 43 | _C.DIRS.LOGS = "./logs/" 44 | 45 | 46 | # ---------------------------------------- 47 | # Datasets config 48 | # ---------------------------------------- 49 | _C.DATA = CN() 50 | # Create small subset to debug 51 | _C.DATA.DEBUG = False 52 | # Datasets problem (multiclass / multilabel) 53 | _C.DATA.TYPE = "" 54 | # Image size for training 55 | _C.DATA.IMG_SIZE = (224, 224) 56 | # Image input channel for training 57 | _C.DATA.INP_CHANNEL = 3 58 | # For CSV loading dataset style 59 | # If dataset is contructed as folders with one class for each folder, see ImageFolder dataset style 60 | # Train, Validation and Test CSV files 61 | _C.DATA.JSON = CN() 62 | _C.DATA.JSON.TRAIN = "" 63 | _C.DATA.JSON.VAL = "" 64 | _C.DATA.JSON.TEST = "" 65 | 66 | # ---------------------------------------- 67 | # Training config 68 | # ---------------------------------------- 69 | _C.TRAIN = CN() 70 | 71 | # Number of training cycles 72 | _C.TRAIN.NUM_CYCLES = 1 73 | # Number of epoches for each cycle. Length of epoches list must equals number of cycle 74 | _C.TRAIN.EPOCHES = [50] 75 | # Training batchsize 76 | _C.TRAIN.BATCH_SIZE = 32 77 | 78 | 79 | 80 | # ---------------------------------------- 81 | # Solver config 82 | # ---------------------------------------- 83 | _C.SOLVER = CN() 84 | 85 | # Solver algorithm 86 | _C.SOLVER.OPTIMIZER = "adamw" 87 | _C.SOLVER.ADAM_EPS = 1e-8 88 | # Solver scheduler (constant / step / cyclical) 89 | _C.SOLVER.SCHEDULER = "cyclical" 90 | # Warmup length. Set 0 if do not want to use 91 | _C.SOLVER.WARMUP_LENGTH = 0 92 | # Use gradient accumulation. If not used, step equals 1 93 | _C.SOLVER.GD_STEPS = 1 94 | # Starting learning rate (after warmup, if used) 95 | _C.SOLVER.BASE_LR = 1e-3 96 | _C.SOLVER.MIN_LR=0. 97 | # Weight decay coeffs 98 | _C.SOLVER.WEIGHT_DECAY = 1e-2 99 | _C.SOLVER.WEIGHT_DECAY_BIAS = 0.0 100 | 101 | # ---------------------------------------- 102 | # Loss function config 103 | # ---------------------------------------- 104 | _C.LOSS = CN() 105 | 106 | # Loss function (ce / focal / dice) 107 | _C.LOSS.NAME = "ce" 108 | _C.LOSS.MSE_CE_WEIGHTS = [1, 1] 109 | 110 | # ---------------------------------------- 111 | # Metric config 112 | # ---------------------------------------- 113 | _C.METRIC=CN() 114 | _C.METRIC.NAME='dice' 115 | # ---------------------------------------- 116 | # Model config 117 | # ---------------------------------------- 118 | _C.MODEL = CN() 119 | 120 | # Classification model arch 121 | _C.MODEL.NAME = "resnet50" 122 | # Load ImageNet pretrained weights 123 | _C.MODEL.PRETRAINED = True 124 | # Classification head 125 | _C.MODEL.CLS_HEAD = 'linear' 126 | # Number of classification class 127 | _C.MODEL.ISUP_NUM_CLASSES = 6 128 | _C.MODEL.G_ONE_NUM_CLASSES = 6 129 | _C.MODEL.G_TWO_NUM_CLASSES = 6 130 | _C.MODEL.POOL = "adaptive_pooling" 131 | _C.MODEL.ORDINAL = CN({"ENABLED": False}) 132 | _C.MODEL.DROPOUT = 0. 133 | _C.MODEL.DROPPATH = 0. 134 | _C.MODEL.CHANNEL_MULTIPLIER = 1. 135 | _C.MODEL.NUM_CLASSES=2 136 | def get_cfg_defaults(): 137 | """Get a yacs CfgNode object with default values for my_project.""" 138 | # Return a clone so that the defaults will not be altered 139 | # This is for the "local variable" use pattern 140 | return _C.clone() 141 | 142 | # Alternatively, provide a way to import the defaults as 143 | # a global singleton: 144 | # cfg = _C # users can `from config import cfg` 145 | --------------------------------------------------------------------------------