Run this command to evaluate the ViDT (Swin-nano) model on COCO :
python -m torch.distributed.launch \
- --nproc_per_node=8 \
+ --nproc_per_node=8 \
--nnodes=1 \
--use_env main.py \
--method vidt \
@@ -356,6 +575,7 @@ python -m torch.distributed.launch \
--num_workers 2 \
--aux_loss True \
--with_box_refine True \
+ --det_token_num 100 \
--coco_path /path/to/coco \
--resume /path/to/vidt_nano \
--pre_trained none \
@@ -376,6 +596,7 @@ python -m torch.distributed.launch \
--num_workers 2 \
--aux_loss True \
--with_box_refine True \
+ --det_token_num 100 \
--coco_path /path/to/coco \
--resume /path/to/vidt_tiny\
--pre_trained none \
@@ -396,6 +617,7 @@ python -m torch.distributed.launch \
--num_workers 2 \
--aux_loss True \
--with_box_refine True \
+ --det_token_num 100 \
--coco_path /path/to/coco \
--resume /path/to/vidt_small \
--pre_trained none \
@@ -416,6 +638,7 @@ python -m torch.distributed.launch \
--num_workers 2 \
--aux_loss True \
--with_box_refine True \
+ --det_token_num 100 \
--coco_path /path/to/coco \
--resume /path/to/vidt_base \
--pre_trained none \
@@ -436,6 +659,15 @@ Please consider citation if our paper is useful in your research.
}
```
+```BibTeX
+@article{song2022vidtplus,
+ title={An Extendable, Efficient and Effective Transformer-based Object Detector},
+ author={Song, Hwanjun and Sun, Deqing and Chun, Sanghyuk and Jampani, Varun and Han, Dongyoon and Heo, Byeongho and Kim, Wonjae and Yang, Ming-Hsuan},
+ journal={arXiv preprint arXiv:2204.07962},
+ year={2022}
+}
+```
+
## License
```
diff --git a/arguments.py b/arguments.py
index a23cc2a..a90aa29 100644
--- a/arguments.py
+++ b/arguments.py
@@ -7,6 +7,7 @@
import argparse
+
def str2bool(v, bool):
if isinstance(v, bool):
@@ -77,8 +78,9 @@ def get_args_parser():
parser.add_argument('--remove_difficult', action='store_true')
# * Device and Log
- parser.add_argument('--output_dir', default='',
+ parser.add_argument('--output_dir', default='/out/models',
help='path where to save, empty for no saving')
+ parser.add_argument('--tensorboard_dir', default='/out/tensorboard', help='path to save tensorboard logs')
parser.add_argument('--device', default='cuda',
help='device to use for training / testing')
parser.add_argument('--seed', default=42, type=int)
@@ -86,6 +88,8 @@ def get_args_parser():
help='start epoch')
parser.add_argument('--eval', default=False, type=lambda x: (str(x).lower() == 'true'), help='eval mode')
parser.add_argument('--resume', default='', help='resume from checkpoint')
+ parser.add_argument('--load_from', default='', help='load from checkpoint, not work with --resume')
+ parser.add_argument('--save_interval', default=100, type=int, help='the interval to save the checkpoint')
# * Training setup
parser.add_argument('--dist-url', default='tcp://127.0.0.1:3457', type=str,
@@ -137,13 +141,38 @@ def get_args_parser():
parser.add_argument('--distil_model_path', default=None, type=str, help="Distillation model path to load")
#######
- # cross-scale fusion
- parser.add_argument('--cross_scale_fusion', default=False, type=lambda x: (str(x).lower() == 'true'), help='use of scale fusion')
+ ####### For ViDT+
+
+ ## EPFF
+ parser.add_argument('--epff', default=False, type=lambda x: (str(x).lower() == 'true'),
+ help='use of EPFF module for pyramid feature fusion')
+
+ ## UQR
+ parser.add_argument('--masks', default=False, action='store_true', help="Train segmentation head if the flag is provided")
+ parser.add_argument('--with_vector', default=False, type=bool)
+ parser.add_argument('--n_keep', default=256, type=int,
+ help="Number of coeffs to be remained")
+ parser.add_argument('--gt_mask_len', default=128, type=int,
+ help="Size of target mask")
+ parser.add_argument('--vector_loss_coef', default=3.0, type=float)
+ parser.add_argument('--vector_hidden_dim', default=256, type=int,
+ help="Size of the vector embeddings (dimension of the transformer)")
+ parser.add_argument('--no_vector_loss_norm', default=False, action='store_true')
+ parser.add_argument('--activation', default='relu', type=str, help="Activation function to use")
+ parser.add_argument('--checkpoint', default=False, action='store_true')
+ parser.add_argument('--vector_start_stage', default=0, type=int)
+ parser.add_argument('--num_machines', default=1, type=int)
+ parser.add_argument('--loss_type', default='l1', type=str)
+ parser.add_argument('--dcn', default=False, action='store_true')
+
+ ## New losses
# iou-aware
- parser.add_argument('--iou_aware', default=False, type=lambda x: (str(x).lower() == 'true'), help='use of iou-aware loss')
+ parser.add_argument('--iou_aware', default=False, type=lambda x: (str(x).lower() == 'true'),
+ help='use of iou-aware loss')
parser.add_argument('--iouaware_loss_coef', default=2, type=float)
# token label
- parser.add_argument('--token_label', default=False, type=lambda x: (str(x).lower() == 'true'), help='use of token label loss')
+ parser.add_argument('--token_label', default=False, type=lambda x: (str(x).lower() == 'true'),
+ help='use of token label loss')
parser.add_argument('--token_loss_coef', default=2, type=float)
#######
diff --git a/datasets/transforms.py b/datasets/transforms.py
index 1477a84..fa4f17e 100644
--- a/datasets/transforms.py
+++ b/datasets/transforms.py
@@ -4,6 +4,7 @@
"""
import random
+import numpy as np
import PIL
import torch
import torchvision.transforms as T
@@ -11,7 +12,7 @@
from util.box_ops import box_xyxy_to_cxcywh
from util.misc import interpolate
-import numpy as np
+
def crop(image, target, region):
cropped_image = F.crop(image, *region)
@@ -74,6 +75,9 @@ def hflip(image, target):
def resize(image, target, size, max_size=None):
+ """
+ PIL image input
+ """
# size can be min_size (scalar) or (w, h) tuple
# import pdb;pdb.set_trace()
maxs = size
@@ -142,7 +146,7 @@ def get_size(image_size, size, max_size=None):
if "masks" in target:
target['masks'] = interpolate(
target['masks'][:, None].float(), size, mode="nearest")[:, 0] > 0.5
-
+
# max_size = max(rescaled_image.size)
# maxs = max(size)
# padding_size = [maxs-rescaled_image.size[0], maxs-rescaled_image.size[1]]
@@ -270,6 +274,7 @@ def __call__(self, image, target=None):
target = target.copy()
h, w = image.shape[-2:]
if "boxes" in target:
+ target['xyxy_boxes'] = target["boxes"]
boxes = target["boxes"]
boxes = box_xyxy_to_cxcywh(boxes)
boxes = boxes / torch.tensor([w, h, w, h], dtype=torch.float32)
diff --git a/datasets/voc.py b/datasets/voc.py
index 5b5ff0b..7634ba0 100644
--- a/datasets/voc.py
+++ b/datasets/voc.py
@@ -25,11 +25,12 @@
# c1, c2 = (int(box[0]), int(box[1])), (int(box[2]), int(box[3]))
# cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
# cv2.imwrite(file_path, img)
-
+
class CocoDetection(torchvision.datasets.CocoDetection):
- def __init__(self, img_folder, ann_file, transforms, return_masks):
+ def __init__(self, image_set, img_folder, ann_file, transforms, return_masks):
super(CocoDetection, self).__init__(img_folder, ann_file)
+ self.image_set = image_set
self._transforms = transforms
self.prepare = ConvertCocoPolysToMask(return_masks)
@@ -125,7 +126,7 @@ def __call__(self, image, target):
return image, target
-def make_coco_transforms(image_set):
+def make_coco_transforms(image_set, args):
normalize = T.Compose([
T.ToTensor(),
@@ -170,7 +171,7 @@ def make_coco_transforms(image_set):
if image_set == 'val':
return T.Compose([
- T.RandomResize([512], max_size=800),
+ T.RandomResize([args.eval_size], max_size=max(args.eval_size, 800)),
normalize,
])
diff --git a/datasets/ymir.py b/datasets/ymir.py
new file mode 100644
index 0000000..45b85f7
--- /dev/null
+++ b/datasets/ymir.py
@@ -0,0 +1,21 @@
+from pathlib import Path
+
+from ymir_exc.util import convert_ymir_to_coco
+
+from .voc import CocoDetection, make_coco_transforms
+
+
+def build(image_set, args):
+ root = Path(args.coco_path)
+ assert root.exists(), f'provided ymir path {root} does not exist'
+
+ data_info = convert_ymir_to_coco(cat_id_from_zero=True)
+
+ datasets = []
+ for split in image_set:
+ img_folder = data_info[split]['img_dir']
+ ann_file = data_info[split]['ann_file']
+ dataset = CocoDetection(split, img_folder, ann_file,
+ transforms=make_coco_transforms(split, args), return_masks=False)
+ datasets.append(dataset)
+ return datasets
diff --git a/figures/architecture.png b/figures/architecture.png
new file mode 100644
index 0000000..270c1dc
Binary files /dev/null and b/figures/architecture.png differ
diff --git a/figures/overview.png b/figures/overview.png
new file mode 100644
index 0000000..6804376
Binary files /dev/null and b/figures/overview.png differ
diff --git a/main.py b/main.py
index 2af0eed..dd39a69 100644
--- a/main.py
+++ b/main.py
@@ -5,25 +5,29 @@
# Additionally modified by NAVER Corp. for ViDT
# ------------------------------------------------------------------------
-import os
+import argparse
import datetime
import json
+import os
import random
+import resource
import time
from pathlib import Path
import numpy as np
import torch
+from tensorboardX import SummaryWriter
from torch.utils.data import DataLoader, DistributedSampler
-import resource
+from ymir_exc import monitor
+from ymir_exc.util import (YmirStage, get_merged_config, write_ymir_training_result, write_ymir_monitor_process)
+
import datasets
import util.misc as utils
+from arguments import get_args_parser
from datasets import build_dataset, get_coco_api_from_dataset
from engine import evaluate, train_one_epoch, train_one_epoch_with_teacher
from methods import build_model
from util.scheduler import create_scheduler
-from arguments import get_args_parser
-import argparse
def build_distil_model(args):
@@ -31,8 +35,10 @@ def build_distil_model(args):
assert args.distil_model in ['vidt_nano', 'vidt_tiny', 'vidt_small', 'vidt_base']
return build_model(args, is_teacher=True)
+
def main(args):
""" main function to train a ViDT model """
+ cfg = get_merged_config()
rlimit = resource.getrlimit(resource.RLIMIT_NOFILE)
resource.setrlimit(resource.RLIMIT_NOFILE, (4096, rlimit[1]))
@@ -40,18 +46,19 @@ def main(args):
# Gradient accumulation setup
if args.n_iter_to_acc > 1:
if args.batch_size % args.n_iter_to_acc != 0:
- print("Not supported divisor for acc grade.")
- import sys
- sys.exit(1)
+ raise Exception(
+ f"Not supported divisor for acc grade with batch size {args.batch_size} and n_iter_to_acc {args.n_iter_to_acc}"
+ )
+
print("Gradient Accumulation is applied.")
- print("The batch: ", args.batch_size, "->", int(args.batch_size / args.n_iter_to_acc),
- 'but updated every ', args.n_iter_to_acc, 'steps.')
+ print("The batch: ", args.batch_size, "->", int(args.batch_size / args.n_iter_to_acc), 'but updated every ',
+ args.n_iter_to_acc, 'steps.')
args.batch_size = args.batch_size // args.n_iter_to_acc
##
# distributed data parallel setup
utils.init_distributed_mode(args)
- print("git:\n {}\n".format(utils.get_sha()))
+ # print("git:\n {}\n".format(utils.get_sha()))
print(args)
device = torch.device(args.device)
@@ -74,11 +81,13 @@ def main(args):
if 'http' in args.distil_model_path or 'https' in args.distil_model_path:
# load from a url
- torch.hub.download_url_to_file(
- url=args.distil_model_path,
- dst="checkpoint.pth"
- )
- checkpoint = torch.load("checkpoint.pth", map_location="cpu")
+ url = args.distil_model_path
+ filename = os.path.basename(url)
+ model_dir = '/out'
+ filepath = os.path.join(model_dir, filename)
+ if not os.path.exists(filepath):
+ torch.hub.download_url_to_file(url=url, dst=filepath)
+ checkpoint = torch.load(filepath, map_location='cpu')
teacher_model.load_state_dict(checkpoint["model"])
else:
# load from a local path
@@ -91,9 +100,7 @@ def main(args):
# parallel model setup
model_without_ddp = model
if args.distributed:
- model = torch.nn.parallel.DistributedDataParallel(model,
- device_ids=[args.gpu],
- find_unused_parameters=True)
+ model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu], find_unused_parameters=True)
model_without_ddp = model.module
if teacher_model is not None:
@@ -124,9 +131,18 @@ def build_optimizer(model, args):
else:
backbone_decay.append(param)
param_dicts = [
- {"params": head},
- {"params": backbone_no_decay, "weight_decay": 0., "lr": args.lr},
- {"params": backbone_decay, "lr": args.lr},
+ {
+ "params": head
+ },
+ {
+ "params": backbone_no_decay,
+ "weight_decay": 0.,
+ "lr": args.lr
+ },
+ {
+ "params": backbone_decay,
+ "lr": args.lr
+ },
]
# print the total number of trainable params.
@@ -141,8 +157,7 @@ def build_optimizer(model, args):
lr_scheduler, _ = create_scheduler(args, optimizer)
# build data loader
- dataset_train = build_dataset(image_set='train', args=args)
- dataset_val = build_dataset(image_set='val', args=args)
+ dataset_train, dataset_val = build_dataset(image_set=['train', 'val'], args=args)
print("# train:", len(dataset_train), ", # val", len(dataset_val))
# data samplers
@@ -153,13 +168,18 @@ def build_optimizer(model, args):
sampler_train = torch.utils.data.RandomSampler(dataset_train)
sampler_val = torch.utils.data.SequentialSampler(dataset_val)
- batch_sampler_train = torch.utils.data.BatchSampler(
- sampler_train, args.batch_size, drop_last=True)
+ batch_sampler_train = torch.utils.data.BatchSampler(sampler_train, args.batch_size, drop_last=True)
- data_loader_train = DataLoader(dataset_train, batch_sampler=batch_sampler_train,
- collate_fn=utils.collate_fn, num_workers=args.num_workers)
- data_loader_val = DataLoader(dataset_val, args.batch_size, sampler=sampler_val,
- drop_last=False, collate_fn=utils.collate_fn, num_workers=args.num_workers)
+ data_loader_train = DataLoader(dataset_train,
+ batch_sampler=batch_sampler_train,
+ collate_fn=utils.collate_fn,
+ num_workers=args.num_workers)
+ data_loader_val = DataLoader(dataset_val,
+ args.batch_size,
+ sampler=sampler_val,
+ drop_last=False,
+ collate_fn=utils.collate_fn,
+ num_workers=args.num_workers)
if args.dataset_file == "coco_panoptic":
# We also evaluate AP during panoptic training, on original coco DS
@@ -170,11 +190,21 @@ def build_optimizer(model, args):
output_dir = Path(args.output_dir)
+ # finetune from a checkpoint
+ if args.load_from:
+ if args.resume:
+ raise Exception("cannot load from and resume at the same time")
+
+ if args.load_from.startswith('https'):
+ checkpoint = torch.hub.load_state_dict_from_url(args.load_from, map_location='cpu', check_hash=True)
+ else:
+ checkpoint = torch.load(args.load_from, map_location='cpu')
+ model_without_ddp.load_state_dict(checkpoint['model'], strict=False)
+ print('load a checkpoint from', args.load_from)
# resume from a checkpoint or eval with a checkpoint
- if args.resume:
+ elif args.resume:
if args.resume.startswith('https'):
- checkpoint = torch.hub.load_state_dict_from_url(
- args.resume, map_location='cpu', check_hash=True)
+ checkpoint = torch.hub.load_state_dict_from_url(args.resume, map_location='cpu', check_hash=True)
else:
checkpoint = torch.load(args.resume, map_location='cpu')
model_without_ddp.load_state_dict(checkpoint['model'], strict=False)
@@ -186,8 +216,7 @@ def build_optimizer(model, args):
# only evaluation purpose
if args.eval:
- test_stats, coco_evaluator = evaluate(model, criterion, postprocessors,
- data_loader_val, base_ds, device)
+ test_stats, coco_evaluator = evaluate(model, criterion, postprocessors, data_loader_val, base_ds, device)
if args.output_dir:
utils.save_on_master(coco_evaluator.coco_eval["bbox"].eval, output_dir / "eval.pth")
@@ -195,8 +224,10 @@ def build_optimizer(model, args):
print("Start training")
start_time = time.time()
- for epoch in range(args.start_epoch, args.epochs):
+ if args.output_dir and utils.is_main_process():
+ tb_writer = SummaryWriter(args.tensorboard_dir)
+ for epoch in range(args.start_epoch, args.epochs):
# specify the current epoch number for samplers
if args.distributed:
sampler_train.set_epoch(epoch)
@@ -217,8 +248,9 @@ def build_optimizer(model, args):
if args.output_dir:
checkpoint_paths = [output_dir / 'checkpoint.pth']
# extra checkpoint before LR drop and every 100 epochs
- if (epoch + 1) % args.lr_drop == 0 or (epoch + 1) % 100 == 0:
- checkpoint_paths.append(output_dir / f'checkpoint{epoch:04}.pth')
+ if (epoch + 1) % args.lr_drop == 0 or (args.save_interval > 0 and (epoch + 1) % args.save_interval) == 0:
+ checkpoint_paths.append(
+ output_dir / f'checkpoint{epoch:04}.pth')
for checkpoint_path in checkpoint_paths:
utils.save_on_master({
'model': model_without_ddp.state_dict(),
@@ -238,6 +270,32 @@ def build_optimizer(model, args):
'epoch': epoch,
'n_parameters': n_parameters}
if args.output_dir and utils.is_main_process():
+ percent = (epoch - args.start_epoch + 1) / (args.epochs - args.start_epoch + 1)
+ write_ymir_monitor_process(cfg, task='training', naive_stage_percent=percent, stage=YmirStage.TASK)
+
+ map50 = test_stats['coco_eval_bbox'][1]
+ if len(checkpoint_paths) == 1:
+ stage_id = 'vidt_last'
+ else:
+ stage_id = f'epoch_{epoch}'
+
+ write_ymir_training_result(cfg, map50, files=[str(checkpoint_paths[-1])], id=stage_id)
+
+ for tag in ['lr', 'loss', 'class_error', 'loss_ce', 'loss_bbox', 'loss_giou']:
+ if tag in train_stats:
+ tb_writer.add_scalar(
+ tag=f'train/{tag}', scalar_value=train_stats[tag], global_step=epoch)
+ if tag in test_stats:
+ tb_writer.add_scalar(
+ tag=f'test/{tag}', scalar_value=test_stats[tag], global_step=epoch)
+ for tag, value in zip(['map', 'map50', 'map75',
+ 'small-map', 'medium-map', 'large-map',
+ 'mar-1', 'mar-10', 'mar-100',
+ 'small-mar', 'medium-mar', 'large-mar'
+ ], test_stats['coco_eval_bbox']):
+ tb_writer.add_scalar(
+ tag=f'test/{tag}', scalar_value=value, global_step=epoch)
+
with (output_dir / "log.txt").open("a") as f:
f.write(json.dumps(log_stats) + "\n")
@@ -252,6 +310,8 @@ def build_optimizer(model, args):
torch.save(coco_evaluator.coco_eval["bbox"].eval,
output_dir / "eval" / name)
+ if args.output_dir and utils.is_main_process():
+ tb_writer.close()
total_time = time.time() - start_time
total_time_str = str(datetime.timedelta(seconds=int(total_time)))
print('Training time {}'.format(total_time_str))
@@ -259,17 +319,32 @@ def build_optimizer(model, args):
if __name__ == '__main__':
- parser = argparse.ArgumentParser('ViDT training and evaluation script', parents=[get_args_parser()])
+ parser = argparse.ArgumentParser(
+ 'ViDT training and evaluation script', parents=[get_args_parser()])
args = parser.parse_args()
''' for testing
args.method = 'vidt'
- args.backbone_name = 'swin_tiny'
- args.batch_size = 2
- args.num_workers = 2
+ args.backbone_name = 'swin_nano'
+ args.batch_size = 4
+ args.num_workers = 4
args.aux_loss = True
args.with_box_refine = True
args.output_dir = 'testing'
+ args.coco_path = '/mnt/ddn/datasets/COCO2017_Seg/train'
+
+ # seg
+ args.epff = True
+ args.token_label = False #True
+ args.iou_aware = False #True
+ args.with_vector = False #True
+ args.masks = False #True
+ args.vector_hidden_dim = 256 #1024
+ args.vector_loss_coef = 3.0
+ args.det_token_num = 300
+
+ args.resume = '/mnt/backbone-nfs/hwanjun/pami2022/optimized_checkpoints/vidt_plus_swin_nano_optimized.pth'
+ args.eval = True
'''
# set dim_feedforward differently
@@ -288,7 +363,7 @@ def build_optimizer(model, args):
args.output_dir += str(args.epochs) + '-'
args.output_dir += str(args.batch_size)
args.output_dir = args.method + '-' + args.backbone_name.upper() + '-batch-' + \
- str(args.batch_size) + '-epoch-' + str(args.epochs)
+ str(args.batch_size) + '-epoch-' + str(args.epochs)
# make log directories
if args.output_dir:
@@ -298,5 +373,3 @@ def build_optimizer(model, args):
print('log', args.output_dir)
main(args)
-
-
diff --git a/methods/__init__.py b/methods/__init__.py
index 9c34d6c..a78fc68 100644
--- a/methods/__init__.py
+++ b/methods/__init__.py
@@ -5,7 +5,7 @@
'''
from methods.vidt.detector import build as vidt_build
-from methods.vidt_wo_neck.detector import build as vidt_wo_neck_build
+
def build_model(args, is_teacher=False):
available_methods = ['vidt_wo_neck', 'vidt']
@@ -15,6 +15,3 @@ def build_model(args, is_teacher=False):
elif args.method == 'vidt':
return vidt_build(args, is_teacher=is_teacher)
-
- elif args.method == 'vidt_wo_neck':
- return vidt_wo_neck_build(args)
diff --git a/methods/swin_w_ram.py b/methods/swin_w_ram.py
index e5c80dc..7fe107d 100644
--- a/methods/swin_w_ram.py
+++ b/methods/swin_w_ram.py
@@ -8,11 +8,13 @@
# --------------------------------------------------------
import math
+import os
+
+import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint as checkpoint
-import numpy as np
from timm.models.layers import DropPath, to_2tuple, trunc_normal_
@@ -36,6 +38,7 @@ def forward(self, x):
x = self.drop(x)
return x
+
def masked_sin_pos_encoding(x, mask, num_pos_feats, temperature=10000, scale=2 * math.pi):
""" Masked Sinusoidal Positional Encoding
@@ -66,8 +69,10 @@ def masked_sin_pos_encoding(x, mask, num_pos_feats, temperature=10000, scale=2 *
pos_x = x_embed[:, :, :, None] / dim_t
pos_y = y_embed[:, :, :, None] / dim_t
- pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3)
- pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3)
+ pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(),
+ pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3)
+ pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(),
+ pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3)
pos = torch.cat((pos_y, pos_x), dim=3)
return pos
@@ -83,8 +88,10 @@ def window_partition(x, window_size):
windows: (num_windows*B, window_size, window_size, C)
"""
B, H, W, C = x.shape
- x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
- windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
+ x = x.view(B, H // window_size, window_size,
+ W // window_size, window_size, C)
+ windows = x.permute(0, 1, 3, 2, 4, 5).contiguous(
+ ).view(-1, window_size, window_size, C)
return windows
@@ -100,7 +107,8 @@ def window_reverse(windows, window_size, H, W):
x: (B, H, W, C)
"""
B = int(windows.shape[0] / (H * W / window_size / window_size))
- x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
+ x = windows.view(B, H // window_size, W // window_size,
+ window_size, window_size, -1)
x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
return x
@@ -143,13 +151,17 @@ def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, at
coords_w = torch.arange(self.window_size[1])
coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, Wh, Ww
coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww
- relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww
- relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2
- relative_coords[:, :, 0] += self.window_size[0] - 1 # shift to start from 0
+ relative_coords = coords_flatten[:, :, None] - \
+ coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww
+ relative_coords = relative_coords.permute(
+ 1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2
+ relative_coords[:, :, 0] += self.window_size[0] - \
+ 1 # shift to start from 0
relative_coords[:, :, 1] += self.window_size[1] - 1
relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
relative_position_index = relative_coords.sum(-1) # Wh*Ww, Wh*Ww
- self.register_buffer("relative_position_index", relative_position_index)
+ self.register_buffer("relative_position_index",
+ relative_position_index)
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
self.attn_drop = nn.Dropout(attn_drop)
@@ -159,7 +171,6 @@ def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, at
trunc_normal_(self.relative_position_bias_table, std=.02)
self.softmax = nn.Softmax(dim=-1)
-
def forward(self, x, det, mask=None, cross_attn=False, cross_attn_mask=None):
""" Forward function.
RAM module receives [Patch] and [DET] tokens and returns their calibrated ones
@@ -201,13 +212,16 @@ def forward(self, x, det, mask=None, cross_attn=False, cross_attn_mask=None):
x = torch.cat([shifted_x, cross_x, det], dim=1)
full_qkv = self.qkv(x)
patch_qkv, cross_patch_qkv, det_qkv = \
- full_qkv[:, :N, :], full_qkv[:, N:N + ori_N, :], full_qkv[:, N + ori_N:, :]
+ full_qkv[:, :N, :], full_qkv[:, N:N +
+ ori_N, :], full_qkv[:, N + ori_N:, :]
patch_qkv = patch_qkv.view(B, H, W, -1)
# window partitioning for [PATCH] tokens
- patch_qkv = window_partition(patch_qkv, window_size) # nW*B, window_size, window_size, C
+ # nW*B, window_size, window_size, C
+ patch_qkv = window_partition(patch_qkv, window_size)
B_ = patch_qkv.shape[0]
- patch_qkv = patch_qkv.reshape(B_, window_size * window_size, 3, self.num_heads, C // self.num_heads)
+ patch_qkv = patch_qkv.reshape(
+ B_, window_size * window_size, 3, self.num_heads, C // self.num_heads)
_patch_qkv = patch_qkv.permute(2, 0, 3, 1, 4)
patch_q, patch_k, patch_v = _patch_qkv[0], _patch_qkv[1], _patch_qkv[2]
@@ -217,19 +231,22 @@ def forward(self, x, det, mask=None, cross_attn=False, cross_attn_mask=None):
# add relative pos bias for [patch x patch] self-attention
relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1) # Wh*Ww,Wh*Ww,nH
- relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww
+ relative_position_bias = relative_position_bias.permute(
+ 2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww
patch_attn = patch_attn + relative_position_bias.unsqueeze(0)
# if shifted window is used, it needs to apply the mask
if mask is not None:
nW = mask.shape[0]
patch_attn = patch_attn.view(B_ // nW, nW, self.num_heads, local_map_size, local_map_size) + \
- mask.unsqueeze(1).unsqueeze(0)
- patch_attn = patch_attn.view(-1, self.num_heads, local_map_size, local_map_size)
+ mask.unsqueeze(1).unsqueeze(0)
+ patch_attn = patch_attn.view(-1, self.num_heads,
+ local_map_size, local_map_size)
patch_attn = self.softmax(patch_attn)
patch_attn = self.attn_drop(patch_attn)
- patch_x = (patch_attn @ patch_v).transpose(1, 2).reshape(B_, window_size, window_size, C)
+ patch_x = (patch_attn @ patch_v).transpose(1,
+ 2).reshape(B_, window_size, window_size, C)
# extract qkv for [DET] tokens
det_qkv = det_qkv.view(B, -1, 3, self.num_heads, C // self.num_heads)
@@ -240,15 +257,18 @@ def forward(self, x, det, mask=None, cross_attn=False, cross_attn_mask=None):
if cross_attn:
# reconstruct the spatial form of [PATCH] tokens for global [DET x PATCH] attention
- cross_patch_qkv = cross_patch_qkv.view(B, ori_H, ori_W, 3, self.num_heads, C // self.num_heads)
- patch_kv = cross_patch_qkv[:, :, :, 1:, :, :].permute(3, 0, 4, 1, 2, 5).contiguous()
+ cross_patch_qkv = cross_patch_qkv.view(
+ B, ori_H, ori_W, 3, self.num_heads, C // self.num_heads)
+ patch_kv = cross_patch_qkv[:, :, :, 1:, :, :].permute(
+ 3, 0, 4, 1, 2, 5).contiguous()
patch_kv = patch_kv.view(2, B, self.num_heads, ori_H * ori_W, -1)
# extract "key and value" of [PATCH] tokens for cross-attention
cross_patch_k, cross_patch_v = patch_kv[0], patch_kv[1]
# bind key and value of [PATCH] and [DET] tokens for [DET X [PATCH, DET]] attention
- det_k, det_v = torch.cat([cross_patch_k, det_k], dim=2), torch.cat([cross_patch_v, det_v], dim=2)
+ det_k, det_v = torch.cat([cross_patch_k, det_k], dim=2), torch.cat(
+ [cross_patch_v, det_v], dim=2)
# [DET x DET] self-attention or binded [DET x [PATCH, DET]] attention
det_q = det_q * self.scale
@@ -309,10 +329,12 @@ def __init__(self, dim, num_heads, window_size=7, shift_size=0,
dim, window_size=to_2tuple(self.window_size), num_heads=num_heads,
qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
- self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+ self.drop_path = DropPath(
+ drop_path) if drop_path > 0. else nn.Identity()
self.norm2 = norm_layer(dim)
mlp_hidden_dim = int(dim * mlp_ratio)
- self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+ self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim,
+ act_layer=act_layer, drop=drop)
self.H = None
self.W = None
@@ -358,7 +380,8 @@ def forward(self, x, mask_matrix, pos, cross_attn, cross_attn_mask):
# cyclic shift
if self.shift_size > 0:
- shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
+ shifted_x = torch.roll(
+ x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
attn_mask = mask_matrix
else:
shifted_x = x
@@ -385,7 +408,8 @@ def forward(self, x, mask_matrix, pos, cross_attn, cross_attn_mask):
# reverse cyclic shift
if self.shift_size > 0:
- x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
+ x = torch.roll(shifted_x, shifts=(
+ self.shift_size, self.shift_size), dims=(1, 2))
else:
x = shifted_x
@@ -518,17 +542,18 @@ def __init__(self,
qk_scale=qk_scale,
drop=drop,
attn_drop=attn_drop,
- drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
+ drop_path=drop_path[i] if isinstance(
+ drop_path, list) else drop_path,
norm_layer=norm_layer)
for i in range(depth)])
# patch merging layer
if downsample is not None:
- self.downsample = downsample(dim=dim, norm_layer=norm_layer, expand=(not last))
+ self.downsample = downsample(
+ dim=dim, norm_layer=norm_layer, expand=(not last))
else:
self.downsample = None
-
def forward(self, x, H, W, det_pos, input_mask, cross_attn=False):
""" Forward function.
@@ -560,16 +585,19 @@ def forward(self, x, H, W, det_pos, input_mask, cross_attn=False):
# mask for cyclic shift
mask_windows = window_partition(img_mask, self.window_size)
- mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
+ mask_windows = mask_windows.view(-1,
+ self.window_size * self.window_size)
attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
- attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
+ attn_mask = attn_mask.masked_fill(
+ attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
# compute sinusoidal pos encoding and cross-attn mask here to avoid redundant computation
if cross_attn:
_H, _W = input_mask.shape[1:]
if not (_H == H and _W == W):
- input_mask = F.interpolate(input_mask[None].float(), size=(H, W)).to(torch.bool)[0]
+ input_mask = F.interpolate(
+ input_mask[None].float(), size=(H, W)).to(torch.bool)[0]
# sinusoidal pos encoding for [PATCH] tokens used in cross-attention
patch_pos = masked_sin_pos_encoding(x, input_mask, self.dim)
@@ -581,8 +609,10 @@ def forward(self, x, H, W, det_pos, input_mask, cross_attn=False):
masked_fill(cross_attn_mask == 0.0, float(0.0))
# pad for detection token (this padding is required to process the binded [PATCH, DET] attention
- cross_attn_mask = cross_attn_mask.view(B, H * W).unsqueeze(1).unsqueeze(2)
- cross_attn_mask = F.pad(cross_attn_mask, (0, self.det_token_num), value=0)
+ cross_attn_mask = cross_attn_mask.view(
+ B, H * W).unsqueeze(1).unsqueeze(2)
+ cross_attn_mask = F.pad(
+ cross_attn_mask, (0, self.det_token_num), value=0)
else:
patch_pos = None
@@ -598,7 +628,7 @@ def forward(self, x, H, W, det_pos, input_mask, cross_attn=False):
if cross_attn:
_cross_attn = True
_cross_attn_mask = cross_attn_mask
- _pos = pos # i.e., (patch_pos, det_pos)
+ _pos = pos # i.e., (patch_pos, det_pos)
else:
_cross_attn = False
_cross_attn_mask = None
@@ -645,7 +675,8 @@ def __init__(self, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None):
self.in_chans = in_chans
self.embed_dim = embed_dim
- self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+ self.proj = nn.Conv2d(in_chans, embed_dim,
+ kernel_size=patch_size, stride=patch_size)
if norm_layer is not None:
self.norm = norm_layer(embed_dim)
else:
@@ -659,7 +690,8 @@ def forward(self, x):
if W % self.patch_size[1] != 0:
x = F.pad(x, (0, self.patch_size[1] - W % self.patch_size[1]))
if H % self.patch_size[0] != 0:
- x = F.pad(x, (0, 0, 0, self.patch_size[0] - H % self.patch_size[0]))
+ x = F.pad(
+ x, (0, 0, 0, self.patch_size[0] - H % self.patch_size[0]))
x = self.proj(x) # B C Wh Ww
if self.norm is not None:
@@ -717,7 +749,8 @@ def __init__(self,
norm_layer=nn.LayerNorm,
ape=False,
patch_norm=True,
- out_indices=[1, 2, 3], # not used in the current version, please ignore.
+ # not used in the current version, please ignore.
+ out_indices=[1, 2, 3],
frozen_stages=-1,
use_checkpoint=False):
super().__init__()
@@ -739,15 +772,18 @@ def __init__(self,
if self.ape:
pretrain_img_size = to_2tuple(pretrain_img_size)
patch_size = to_2tuple(patch_size)
- patches_resolution = [pretrain_img_size[0] // patch_size[0], pretrain_img_size[1] // patch_size[1]]
+ patches_resolution = [
+ pretrain_img_size[0] // patch_size[0], pretrain_img_size[1] // patch_size[1]]
- self.absolute_pos_embed = nn.Parameter(torch.zeros(1, embed_dim, patches_resolution[0], patches_resolution[1]))
+ self.absolute_pos_embed = nn.Parameter(torch.zeros(
+ 1, embed_dim, patches_resolution[0], patches_resolution[1]))
trunc_normal_(self.absolute_pos_embed, std=.02)
self.pos_drop = nn.Dropout(p=drop_rate)
# stochastic depth
- dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] # stochastic depth decay rule
+ dpr = [x.item() for x in torch.linspace(0, drop_path_rate,
+ sum(depths))] # stochastic depth decay rule
# build layers
self.layers = nn.ModuleList()
@@ -765,13 +801,15 @@ def __init__(self,
drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])],
norm_layer=norm_layer,
# modified by ViDT
- downsample=PatchMerging if (i_layer < self.num_layers) else None,
+ downsample=PatchMerging if (
+ i_layer < self.num_layers) else None,
last=None if (i_layer < self.num_layers-1) else True,
#
use_checkpoint=use_checkpoint)
self.layers.append(layer)
- num_features = [int(embed_dim * 2 ** i) for i in range(self.num_layers)]
+ num_features = [int(embed_dim * 2 ** i)
+ for i in range(self.num_layers)]
self.num_features = num_features
# add a norm layer for each output
@@ -821,7 +859,8 @@ def finetune_det(self, method, det_token_num=100, pos_dim=256, cross_indices=[3]
# how many object we detect?
self.det_token_num = det_token_num
- self.det_token = nn.Parameter(torch.zeros(1, det_token_num, self.num_features[0]))
+ self.det_token = nn.Parameter(torch.zeros(
+ 1, det_token_num, self.num_features[0]))
self.det_token = trunc_normal_(self.det_token, std=.02)
# dim size of pos encoding
@@ -833,9 +872,11 @@ def finetune_det(self, method, det_token_num=100, pos_dim=256, cross_indices=[3]
self.det_pos_embed = torch.nn.Parameter(det_pos_embed)
# info for detection
- self.num_channels = [self.num_features[i+1] for i in range(len(self.num_features)-1)]
+ self.num_channels = [self.num_features[i+1]
+ for i in range(len(self.num_features)-1)]
if method == 'vidt':
- self.num_channels.append(self.pos_dim) # default: 256 (same to the default pos_dim)
+ # default: 256 (same to the default pos_dim)
+ self.num_channels.append(self.pos_dim)
self.cross_indices = cross_indices
# divisor to reduce the spatial size of the mask
self.mask_divisor = 2 ** (len(self.layers) - len(self.cross_indices))
@@ -887,7 +928,7 @@ def forward(self, x, mask):
# prepare a mask for cross attention
mask = F.interpolate(mask[None].float(),
- size=(Wh // self.mask_divisor, Ww // self.mask_divisor)).to(torch.bool)[0]
+ size=(Wh // self.mask_divisor, Ww // self.mask_divisor)).to(torch.bool)[0]
patch_outs = []
for stage in range(self.num_layers):
@@ -906,11 +947,13 @@ def forward(self, x, mask):
det_pos=det_pos,
cross_attn=cross_attn)
- x, det_token = x[:, :-self.det_token_num, :], x[:, -self.det_token_num:, :]
+ x, det_token = x[:, :-self.det_token_num,
+ :], x[:, -self.det_token_num:, :]
# Aggregate intermediate outputs
if stage > 0:
- patch_out = x_out[:, :-self.det_token_num, :].view(B, H, W, -1).permute(0, 3, 1, 2)
+ patch_out = x_out[:, :-self.det_token_num,
+ :].view(B, H, W, -1).permute(0, 3, 1, 2)
patch_outs.append(patch_out)
# patch token reduced from last stage output
@@ -935,11 +978,26 @@ def flops(self):
flops += self.patch_embed.flops()
for i, layer in enumerate(self.layers):
flops += layer.flops()
- flops += self.num_features * self.patches_resolution[0] * self.patches_resolution[1] // (2 ** self.num_layers)
+ flops += self.num_features * \
+ self.patches_resolution[0] * \
+ self.patches_resolution[1] // (2 ** self.num_layers)
flops += self.num_features * self.num_classes
return flops
+def load_state_dict_from_url(url, model_dir, map_location):
+ """
+ for pytorch 1.6, the torch.hub.load_state_dict_from_url not work,
+ view https://github.com/open-mmlab/mmsegmentation/issues/687 for details
+ """
+ filename = os.path.basename(url)
+ filepath = os.path.join(model_dir, filename)
+
+ if not os.path.exists(filepath):
+ torch.hub.download_url_to_file(url=url, dst=filepath)
+ return torch.load(filepath, map_location=map_location)
+
+
def swin_nano(pretrained=None, **kwargs):
model = SwinTransformer(pretrain_img_size=[224, 224], embed_dim=48, depths=[2, 2, 6, 2],
num_heads=[3, 6, 12, 24], window_size=7, drop_path_rate=0.0, **kwargs)
@@ -949,11 +1007,11 @@ def swin_nano(pretrained=None, **kwargs):
if pretrained is not None:
if pretrained == 'imagenet':
- torch.hub.download_url_to_file(
- url="https://github.com/naver-ai/vidt/releases/download/v0.1-swin/swin_nano_patch4_window7_224.pth",
- dst="checkpoint.pth"
- )
- checkpoint = torch.load("checkpoint.pth", map_location="cpu")
+ url = "https://github.com/naver-ai/vidt/releases/download/v0.1-swin/swin_nano_patch4_window7_224.pth"
+ checkpoint = load_state_dict_from_url(url=url,
+ model_dir="/out",
+ map_location="cpu")
+
model.load_state_dict(checkpoint["model"], strict=False)
print('Load the backbone pretrained on ImageNet 1K')
@@ -974,11 +1032,11 @@ def swin_tiny(pretrained=None, **kwargs):
if pretrained is not None:
if pretrained == 'imagenet':
- torch.hub.download_url_to_file(
- url="https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_tiny_patch4_window7_224.pth",
- dst="checkpoint.pth"
- )
- checkpoint = torch.load("checkpoint.pth", map_location="cpu")
+ url = "https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_tiny_patch4_window7_224.pth"
+
+ checkpoint = load_state_dict_from_url(url=url,
+ model_dir="/out",
+ map_location="cpu")
model.load_state_dict(checkpoint["model"], strict=False)
print('Load the backbone pretrained on ImageNet 1K')
else:
@@ -997,11 +1055,10 @@ def swin_small(pretrained=None, **kwargs):
if pretrained is not None:
if pretrained == 'imagenet':
- torch.hub.download_url_to_file(
- url="https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_small_patch4_window7_224.pth",
- dst="checkpoint.pth"
- )
- checkpoint = torch.load("checkpoint.pth", map_location="cpu")
+ url = "https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_small_patch4_window7_224.pth"
+ checkpoint = load_state_dict_from_url(url=url,
+ model_dir="/out",
+ map_location="cpu")
model.load_state_dict(checkpoint["model"], strict=False)
print('Load the backbone pretrained on ImageNet 1K')
else:
@@ -1020,11 +1077,10 @@ def swin_base_win7(pretrained=None, **kwargs):
if pretrained is not None:
if pretrained == 'imagenet':
- torch.hub.download_url_to_file(
- url="https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window7_224_22k.pth",
- dst="checkpoint.pth"
- )
- checkpoint = torch.load("checkpoint.pth", map_location="cpu")
+ url = "https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window7_224_22k.pth"
+ checkpoint = load_state_dict_from_url(url=url,
+ model_dir="/out",
+ map_location="cpu")
model.load_state_dict(checkpoint["model"], strict=False)
print('Load the backbone pretrained on ImageNet 22K')
else:
@@ -1043,11 +1099,11 @@ def swin_large_win7(pretrained=None, **kwargs):
if pretrained is not None:
if pretrained == 'imagenet':
- torch.hub.download_url_to_file(
- url="https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_large_patch4_window7_224_22k.pth",
- dst="checkpoint.pth"
- )
- checkpoint = torch.load("checkpoint.pth", map_location="cpu")
+ url = "https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_large_patch4_window7_224_22k.pth"
+
+ checkpoint = load_state_dict_from_url(url=url,
+ model_dir="/out",
+ map_location="cpu")
model.load_state_dict(checkpoint["model"], strict=False)
print('Load the backbone pretrained on ImageNet 22K')
else:
diff --git a/methods/vidt/criterion.py b/methods/vidt/criterion.py
index 58ecb19..c942b90 100644
--- a/methods/vidt/criterion.py
+++ b/methods/vidt/criterion.py
@@ -19,6 +19,9 @@
is_dist_avail_and_initialized)
from methods.segmentation import (dice_loss, sigmoid_focal_loss)
import copy
+from util.detectron2.structures.masks import BitMasks
+import cv2
+import numpy as np
class SetCriterion(nn.Module):
@@ -28,7 +31,13 @@ class SetCriterion(nn.Module):
2) we supervise each pair of matched ground-truth / prediction (supervise class and box)
"""
- def __init__(self, num_classes, matcher, weight_dict, losses, focal_alpha=0.25):
+ def __init__(self, num_classes, matcher, weight_dict, losses, focal_alpha=0.25,
+ # For UQR module for instance segmentation
+ with_vector=False,
+ processor_dct=None,
+ vector_loss_coef=0.7,
+ no_vector_loss_norm=False,
+ vector_start_stage=0):
""" Create the criterion.
Parameters:
num_classes: number of object categories, omitting the special no-object category
@@ -44,6 +53,18 @@ def __init__(self, num_classes, matcher, weight_dict, losses, focal_alpha=0.25):
self.weight_dict = weight_dict
self.losses = losses
self.focal_alpha = focal_alpha
+ #
+ self.with_vector = with_vector
+ self.processor_dct = processor_dct
+ self.vector_loss_coef = vector_loss_coef
+ self.no_vector_loss_norm = no_vector_loss_norm
+ self.vector_start_stage = vector_start_stage
+
+ if self.with_vector is True:
+ print(f'Training with {6-self.vector_start_stage} vector stages.')
+ print(f"Training with vector_loss_coef {self.vector_loss_coef}.")
+ if not self.no_vector_loss_norm:
+ print('Training with vector_loss_norm.')
def loss_labels(self, outputs, targets, indices, num_boxes, log=True):
"""Classification loss (NLL)
@@ -113,30 +134,52 @@ def loss_masks(self, outputs, targets, indices, num_boxes):
"""Compute the losses related to the masks: the focal loss and the dice loss.
targets dicts must contain the key "masks" containing a tensor of dim [nb_target_boxes, h, w]
"""
-
- assert "pred_masks" in outputs
+ assert "pred_vectors" in outputs
src_idx = self._get_src_permutation_idx(indices)
tgt_idx = self._get_tgt_permutation_idx(indices)
- src_masks = outputs["pred_masks"]
+ src_masks = outputs["pred_vectors"]
+ src_boxes = outputs['pred_boxes']
# TODO use valid to mask invalid areas due to padding in loss
+ target_boxes = torch.cat([t['xyxy_boxes'][i] for t, (_, i) in zip(targets, indices)], dim=0)
target_masks, valid = nested_tensor_from_tensor_list([t["masks"] for t in targets]).decompose()
target_masks = target_masks.to(src_masks)
+ src_vectors = src_masks[src_idx]
+ src_boxes = src_boxes[src_idx]
+ target_masks = target_masks[tgt_idx]
+
+ # crop gt_masks
+ n_keep, gt_mask_len = self.processor_dct.n_keep, self.processor_dct.gt_mask_len
+ gt_masks = BitMasks(target_masks)
+ gt_masks = gt_masks.crop_and_resize(target_boxes, gt_mask_len).to(device=src_masks.device).float()
+ target_masks = gt_masks
+
+ if target_masks.shape[0] == 0:
+ losses = {
+ "loss_vector": src_vectors.sum() * 0
+ }
+ return losses
+
+ # perform dct transform
+ target_vectors = []
+ for i in range(target_masks.shape[0]):
+ gt_mask_i = ((target_masks[i,:,:] >= 0.5)* 1).to(dtype=torch.uint8)
+ gt_mask_i = gt_mask_i.cpu().numpy().astype(np.float32)
+ coeffs = cv2.dct(gt_mask_i)
+ coeffs = torch.from_numpy(coeffs).flatten()
+ coeffs = coeffs[torch.tensor(self.processor_dct.zigzag_table)]
+ gt_label = coeffs.unsqueeze(0)
+ target_vectors.append(gt_label)
+
+ target_vectors = torch.cat(target_vectors, dim=0).to(device=src_vectors.device)
+ losses = {}
+ if self.no_vector_loss_norm:
+ losses['loss_vector'] = self.vector_loss_coef * F.l1_loss(src_vectors, target_vectors, reduction='none').sum() / num_boxes
+ else:
+ losses['loss_vector'] = self.vector_loss_coef * F.l1_loss(src_vectors, target_vectors, reduction='mean')
- src_masks = src_masks[src_idx]
- # upsample predictions to the target size
- src_masks = interpolate(src_masks[:, None], size=target_masks.shape[-2:],
- mode="bilinear", align_corners=False)
- src_masks = src_masks[:, 0].flatten(1)
-
- target_masks = target_masks[tgt_idx].flatten(1)
-
- losses = {
- "loss_mask": sigmoid_focal_loss(src_masks, target_masks, num_boxes),
- "loss_dice": dice_loss(src_masks, target_masks, num_boxes),
- }
return losses
def loss_iouaware(self, outputs, targets, indices, num_boxes):
@@ -220,13 +263,13 @@ def forward(self, outputs, targets, distil_tokens=None):
outputs: dict of tensors, see the output specification of the model for the format
targets: list of dicts, such that len(targets) == batch_size.
The expected keys in each dict depends on the losses applied, see each loss' doc
+ distil_tokens: for token distillation
"""
outputs_without_aux = {k: v for k, v in outputs.items() if k != 'aux_outputs' and k != 'enc_outputs'}
# Retrieve the matching between the outputs of the last layer and the targets
indices = self.matcher(outputs_without_aux, targets)
- _indices = indices
# Compute the average number of target boxes accross all nodes, for normalization purposes
num_boxes = sum(len(t["labels"]) for t in targets)
@@ -246,7 +289,7 @@ def forward(self, outputs, targets, distil_tokens=None):
for i, aux_outputs in enumerate(outputs['aux_outputs']):
indices = self.matcher(aux_outputs, targets)
for loss in self.losses:
- if loss == 'masks':
+ if loss == 'masks' and i < self.vector_start_stage:
# Intermediate masks losses are too costly to compute, we ignore them.
continue
kwargs = {}
@@ -255,7 +298,6 @@ def forward(self, outputs, targets, distil_tokens=None):
kwargs['log'] = False
l_dict = self.get_loss(loss, aux_outputs, targets, indices, num_boxes, **kwargs)
l_dict = {k + f'_{i}': v for k, v in l_dict.items()}
-
losses.update(l_dict)
if 'enc_outputs' in outputs:
diff --git a/methods/vidt/dct.py b/methods/vidt/dct.py
new file mode 100644
index 0000000..7fc48c3
--- /dev/null
+++ b/methods/vidt/dct.py
@@ -0,0 +1,194 @@
+# ------------------------------------------------------------------------
+# Copyright (c) 2021 megvii-model. All Rights Reserved.
+# ------------------------------------------------------------------------
+import numpy as np
+import functools
+
+print = functools.partial(print, flush=True)
+
+class ProcessorDCT(object):
+ def __init__(self, n_keep, gt_mask_len):
+ self.n_keep = n_keep
+ self.gt_mask_len = gt_mask_len
+
+ inputs = np.zeros((self.gt_mask_len, self.gt_mask_len))
+ _, zigzag_table = self.zigzag(inputs)
+ self.zigzag_table = zigzag_table[:self.n_keep]
+
+ def sigmoid(self, x):
+ """Apply the sigmoid operation.
+ """
+ y = 1. / (1. + 1. / np.exp(x))
+ dy = y * (1 - y)
+ return y
+
+ def inverse_sigmoid(self, x):
+ """Apply the inverse sigmoid operation.
+ y = -ln(1-x/x)
+ """
+ y = -1 * np.log((1 - x) / x)
+ return y
+
+ def zigzag(self, input, gt=None):
+ """
+ Zigzag scan of a matrix
+ Argument is a two-dimensional matrix of any size,
+ not strictly a square one.
+ Function returns a 1-by-(m*n) array,
+ where m and n are sizes of an input matrix,
+ consisting of its items scanned by a zigzag method.
+
+ Args:
+ input (np.array): shape [h,w], value belong to [-127, 128], transformed from gt.
+ gt (np.array): shape [h,w], value belong to {0,1}, original instance segmentation gt mask.
+ Returns:
+ output (np.array): zig-zag encoded values, shape [h*w].
+ indicator (np.array): positive sample indicator, shape [h,w].
+ """
+ # initializing the variables
+ h = 0
+ v = 0
+
+ vmin = 0
+ hmin = 0
+
+ vmax = input.shape[0]
+ hmax = input.shape[1]
+ assert vmax == hmax
+
+ i = 0
+ output = np.zeros((vmax * hmax))
+ indicator = []
+
+ while ((v < vmax) and (h < hmax)):
+ if ((h + v) % 2) == 0: # going up
+ if (v == vmin):
+ output[i] = input[v, h] # if we got to the first line
+ indicator.append(v * vmax + h)
+ if (h == hmax):
+ v = v + 1
+ else:
+ h = h + 1
+ i = i + 1
+ elif ((h == hmax - 1) and (v < vmax)): # if we got to the last column
+ output[i] = input[v, h]
+ indicator.append(v * vmax + h)
+ v = v + 1
+ i = i + 1
+ elif ((v > vmin) and (h < hmax - 1)): # all other cases
+ output[i] = input[v, h]
+ indicator.append(v * vmax + h)
+ v = v - 1
+ h = h + 1
+ i = i + 1
+ else: # going down
+ if ((v == vmax - 1) and (h <= hmax - 1)): # if we got to the last line
+ output[i] = input[v, h]
+ indicator.append(v * vmax + h)
+ h = h + 1
+ i = i + 1
+ elif (h == hmin): # if we got to the first column
+ output[i] = input[v, h]
+ indicator.append(v * vmax + h)
+ if (v == vmax - 1):
+ h = h + 1
+ else:
+ v = v + 1
+ i = i + 1
+ elif ((v < vmax - 1) and (h > hmin)): # all other cases
+ output[i] = input[v, h]
+ indicator.append(v * vmax + h)
+ v = v + 1
+ h = h - 1
+ i = i + 1
+ if ((v == vmax - 1) and (h == hmax - 1)): # bottom right element
+ output[i] = input[v, h]
+ indicator.append(v * vmax + h)
+ break
+ return output, indicator
+
+ def inverse_zigzag(self, input, vmax, hmax):
+ """
+ Zigzag scan of a matrix
+ Argument is a two-dimensional matrix of any size,
+ not strictly a square one.
+ Function returns a 1-by-(m*n) array,
+ where m and n are sizes of an input matrix,
+ consisting of its items scanned by a zigzag method.
+ """
+ # initializing the variables
+ h = 0
+ v = 0
+
+ vmin = 0
+ hmin = 0
+
+ output = np.zeros((vmax, hmax))
+ i = 0
+ while ((v < vmax) and (h < hmax)):
+ if ((h + v) % 2) == 0: # going up
+ if (v == vmin):
+ output[v, h] = input[i] # if we got to the first line
+ if (h == hmax):
+ v = v + 1
+ else:
+ h = h + 1
+ i = i + 1
+ elif ((h == hmax - 1) and (v < vmax)): # if we got to the last column
+ output[v, h] = input[i]
+ v = v + 1
+ i = i + 1
+ elif ((v > vmin) and (h < hmax - 1)): # all other cases
+ output[v, h] = input[i]
+ v = v - 1
+ h = h + 1
+ i = i + 1
+ else: # going down
+ if ((v == vmax - 1) and (h <= hmax - 1)): # if we got to the last line
+ output[v, h] = input[i]
+ h = h + 1
+ i = i + 1
+ elif (h == hmin): # if we got to the first column
+ output[v, h] = input[i]
+ if (v == vmax - 1):
+ h = h + 1
+ else:
+ v = v + 1
+ i = i + 1
+ elif ((v < vmax - 1) and (h > hmin)): # all other cases
+ output[v, h] = input[i]
+ v = v + 1
+ h = h - 1
+ i = i + 1
+ if ((v == vmax - 1) and (h == hmax - 1)): # bottom right element
+ output[v, h] = input[i]
+ break
+ return output
+
+'''
+if __name__ == "__main__":
+ img_path = '/data/dongbin/projects/SparseR-CNN/projects/SparseRCNN/debug/tgt_1.jpg'
+
+ # complete API
+ img_dir = '/data/dongbin/projects/Deformable-DETR-main/exps/paper_figs'
+ name = 'gt_4.jpg'
+ img_path = os.path.join(img_dir, name)
+ gt_mask_len = 128
+ n_keep = 256
+ processor_dct = ProcessorDCT(n_keep=n_keep, gt_mask_len=gt_mask_len)
+ mask = cv2.imread(img_path, 0).astype(np.float32)
+ coeffs = cv2.dct(mask)
+ cv2.imwrite(os.path.join(img_dir, '{}_coeffs.png'.format(name.split('.')[0])), coeffs)
+
+ idct = np.zeros((gt_mask_len ** 2))
+ vectors = torch.from_numpy(coeffs).flatten()
+ vectors = vectors[torch.tensor(processor_dct.zigzag_table)]
+ idct[:n_keep] = vectors.cpu().numpy()
+ idct = processor_dct.inverse_zigzag(idct, gt_mask_len, gt_mask_len)
+ cv2.imwrite(os.path.join(img_dir, '{}_i_coeffs.png'.format(name.split('.')[0])), idct)
+ re_mask = cv2.idct(idct)
+ max_v = np.max(re_mask)
+ min_v = np.min(re_mask)
+ re_mask = np.where(re_mask > (max_v + min_v) / 2., 255, 0)
+ cv2.imwrite(os.path.join(img_dir, '{}_recover.png'.format(name.split('.')[0])), re_mask)
+'''
\ No newline at end of file
diff --git a/methods/vidt/detector.py b/methods/vidt/detector.py
index 9e75010..18627c3 100644
--- a/methods/vidt/detector.py
+++ b/methods/vidt/detector.py
@@ -17,12 +17,12 @@
from methods.coat_w_ram import coat_lite_tiny, coat_lite_mini, coat_lite_small
from .matcher import build_matcher
from .criterion import SetCriterion
-from .postprocessor import PostProcess
+from .postprocessor import PostProcess, PostProcessSegm
from .deformable_transformer import build_deforamble_transformer
from methods.vidt.fpn_fusion import FPNFusionModule
import copy
import math
-
+from .dct import ProcessorDCT
def _get_clones(module, N):
""" Clone a moudle N times """
@@ -35,9 +35,10 @@ class Detector(nn.Module):
def __init__(self, backbone, transformer, num_classes, num_queries,
aux_loss=False, with_box_refine=False,
- # The three techniques were not used in ViDT paper.
- # After submitting our paper, we saw the ViDT performance could be further enhanced with them.
- cross_scale_fusion=None, iou_aware=False, token_label=False,
+ # The three additional techniques for ViDT+
+ epff=None,# (1) Efficient Pyramid Feature Fusion Module
+ with_vector=False, processor_dct=None, vector_hidden_dim=256,# (2) UQR Module
+ iou_aware=False, token_label=False, # (3) Additional losses
distil=False):
""" Initializes the model.
Parameters:
@@ -48,7 +49,7 @@ def __init__(self, backbone, transformer, num_classes, num_queries,
DETR can detect in a single image. For COCO, we recommend 100 queries.
aux_loss: True if auxiliary decoding losses (loss at each decoder layer) are to be used.
with_box_refine: iterative bounding box refinement
- cross_scale_fusion: None or fusion module available
+ epff: None or fusion module available
iou_aware: True if iou_aware is to be used.
see the original paper https://arxiv.org/abs/1912.05992
token_label: True if token_label is to be used.
@@ -68,16 +69,24 @@ def __init__(self, backbone, transformer, num_classes, num_queries,
self.aux_loss = aux_loss
self.with_box_refine = with_box_refine
- # three additional techniques not used in the ViDT paper
- # optional use, we will revise our paper for the below techniques
+
+ # For UQR module for ViDT+
+ self.with_vector = with_vector
+ self.processor_dct = processor_dct
+ if self.with_vector:
+ print(f'Training with vector_hidden_dim {vector_hidden_dim}.', flush=True)
+ self.vector_embed = MLP(hidden_dim, vector_hidden_dim, self.processor_dct.n_keep, 3)
+
+ ############ Modified for ViDT+
+ # For two additional losses for ViDT+
self.iou_aware = iou_aware
self.token_label = token_label
# distillation
self.distil = distil
- # [PATCH] token channel reduction for the input to transformer decoder
- if cross_scale_fusion is None:
+ # For EPFF module for ViDT+
+ if epff is None:
num_backbone_outs = len(backbone.num_channels)
input_proj_list = []
for _ in range(num_backbone_outs):
@@ -96,7 +105,8 @@ def __init__(self, backbone, transformer, num_classes, num_queries,
self.fusion = None
else:
# the cross scale fusion module has its own reduction layers
- self.fusion = cross_scale_fusion
+ self.fusion = epff
+ ############
# channel dim reduction for [DET] tokens
self.tgt_proj = nn.Sequential(
@@ -125,6 +135,12 @@ def __init__(self, backbone, transformer, num_classes, num_queries,
nn.init.xavier_uniform_(self.query_pos_proj[0].weight, gain=1)
nn.init.constant_(self.query_pos_proj[0].bias, 0)
+ ############ Added for UQR
+ if self.with_vector:
+ nn.init.constant_(self.vector_embed.layers[-1].weight.data, 0)
+ nn.init.constant_(self.vector_embed.layers[-1].bias.data, 0)
+ ############
+
# the prediction is made for each decoding layers + the standalone detector (Swin with RAM)
num_pred = transformer.decoder.num_layers + 1
@@ -141,6 +157,12 @@ def __init__(self, backbone, transformer, num_classes, num_queries,
self.bbox_embed = nn.ModuleList([self.bbox_embed for _ in range(num_pred)])
self.transformer.decoder.bbox_embed = None
+ ############ Added for UQR
+ if self.with_vector:
+ nn.init.constant_(self.vector_embed.layers[-1].bias.data[2:], -2.0)
+ self.vector_embed = nn.ModuleList([self.vector_embed for _ in range(num_pred)])
+ ############
+
if self.iou_aware:
self.iou_embed = MLP(hidden_dim, hidden_dim, 1, 3)
if with_box_refine:
@@ -170,8 +192,7 @@ def forward(self, samples: NestedTensor):
If iou_aware is True, "pred_ious" is also returns as one of the key in "aux_outputs"
- "enc_tokens": If token_label is True, "enc_tokens" is returned to be used
- Note that aux_loss and box refinement is used in ViDT in default. The detailed ablation of using
- the cross_scale_fusion, iou_aware & token_lablel loss will be discussed in a later version
+ Note that aux_loss and box refinement is used in ViDT in default.
"""
if isinstance(samples, (list, torch.Tensor)):
@@ -197,7 +218,7 @@ def forward(self, samples: NestedTensor):
for l, src in enumerate(features):
srcs.append(self.input_proj[l](src))
else:
- # multi-scale fusion is used if fusion is not None
+ # EPFF (multi-scale fusion) is used if fusion is activated
srcs = self.fusion(features)
masks = []
@@ -237,12 +258,29 @@ def forward(self, samples: NestedTensor):
outputs_class = torch.stack(outputs_classes)
outputs_coord = torch.stack(outputs_coords)
+ ############ Added for UQR
+ outputs_vector = None
+ if self.with_vector:
+ outputs_vectors = []
+ for lvl in range(hs.shape[0]):
+ outputs_vector = self.vector_embed[lvl](hs[lvl])
+ outputs_vectors.append(outputs_vector)
+ outputs_vector = torch.stack(outputs_vectors)
+ ############
+
# final prediction is made the last decoding layer
out = {'pred_logits': outputs_class[-1], 'pred_boxes': outputs_coord[-1]}
+ ############ Added for UQR
+ if self.with_vector:
+ out.update({'pred_vectors': outputs_vector[-1]})
+
+ ############
+
# aux loss is defined by using the rest predictions
if self.aux_loss and self.transformer.decoder.num_layers > 0:
- out['aux_outputs'] = self._set_aux_loss(outputs_class, outputs_coord)
+ out['aux_outputs'] = self._set_aux_loss(outputs_class, outputs_coord, outputs_vector)
+
# iou awareness loss is defined for each decoding layer similar to auxiliary decoding loss
if self.iou_aware:
@@ -268,13 +306,17 @@ def forward(self, samples: NestedTensor):
return out
@torch.jit.unused
- def _set_aux_loss(self, outputs_class, outputs_coord):
+ def _set_aux_loss(self, outputs_class, outputs_coord, outputs_vector):
# this is a workaround to make torchscript happy, as torchscript
# doesn't support dictionary with non-homogeneous values, such
# as a dict having both a Tensor and a list.
- return [{'pred_logits': a, 'pred_boxes': b}
- for a, b in zip(outputs_class[:-1], outputs_coord[:-1])]
+ if outputs_vector is None:
+ return [{'pred_logits': a, 'pred_boxes': b}
+ for a, b in zip(outputs_class[:-1], outputs_coord[:-1])]
+ else:
+ return [{'pred_logits': a, 'pred_boxes': b, 'pred_vectors': c}
+ for a, b, c in zip(outputs_class[:-1], outputs_coord[:-1], outputs_vector[:-1])]
class MLP(nn.Module):
@@ -293,10 +335,14 @@ def forward(self, x):
def build(args, is_teacher=False):
- # a teacher model for distilation
+ # distillation is deprecated
if is_teacher:
- return build_teacher(args)
- #
+ print('Token Distillation is deprecated in this version. Please use the previous version of ViDT.')
+ assert is_teacher is False
+ # a teacher model for distilation
+ #if is_teacher:
+ # return build_teacher(args)
+ ############################
if args.dataset_file == 'coco':
num_classes = 91
@@ -332,12 +378,16 @@ def build(args, is_teacher=False):
pos_dim=args.reduced_dim,
cross_indices=args.cross_indices)
- cross_scale_fusion = None
- if args.cross_scale_fusion:
- cross_scale_fusion = FPNFusionModule(backbone.num_channels, fuse_dim=args.reduced_dim)
+ epff = None
+ if args.epff:
+ epff = FPNFusionModule(backbone.num_channels, fuse_dim=args.reduced_dim)
deform_transformers = build_deforamble_transformer(args)
+ # Added for UQR module
+ if args.with_vector:
+ processor_dct = ProcessorDCT(args.n_keep, args.gt_mask_len)
+
model = Detector(
backbone,
deform_transformers,
@@ -346,10 +396,15 @@ def build(args, is_teacher=False):
# two essential techniques used in ViDT
aux_loss=args.aux_loss,
with_box_refine=args.with_box_refine,
- # three additional techniques (optionally)
- cross_scale_fusion=cross_scale_fusion,
+ # an epff module for ViDT+
+ epff=epff,
+ # an UQR module for ViDT+
+ with_vector=args.with_vector,
+ processor_dct=processor_dct if args.with_vector else None,
+ # two additional losses for VIDT+
iou_aware=args.iou_aware,
token_label=args.token_label,
+ vector_hidden_dim=args.vector_hidden_dim,
# distil
distil=False if args.distil_model is None else True,
)
@@ -366,6 +421,10 @@ def build(args, is_teacher=False):
weight_dict['loss_token_focal'] = args.token_loss_coef
weight_dict['loss_token_dice'] = args.token_loss_coef
+ # For UQR module
+ if args.masks:
+ weight_dict["loss_vector"] = 1
+
if args.distil_model is not None:
weight_dict['loss_distil'] = args.distil_loss_coef
@@ -381,19 +440,36 @@ def build(args, is_teacher=False):
if args.iou_aware:
losses += ['iouaware']
+ # For UQR
+ if args.masks:
+ losses += ["masks"]
+
# num_classes, matcher, weight_dict, losses, focal_alpha=0.25
- criterion = SetCriterion(num_classes, matcher, weight_dict, losses, focal_alpha=args.focal_alpha)
+ criterion = SetCriterion(num_classes, matcher, weight_dict, losses,
+ focal_alpha=args.focal_alpha,
+ # For UQR
+ with_vector=args.with_vector,
+ processor_dct=processor_dct if args.with_vector else None,
+ vector_loss_coef=args.vector_loss_coef,
+ no_vector_loss_norm=args.no_vector_loss_norm,
+ vector_start_stage=args.vector_start_stage)
criterion.to(device)
- postprocessors = {'bbox': PostProcess(args.dataset_file)}
+ # HER -=> post도 고쳐야함.
+ postprocessors = {'bbox': PostProcess(processor_dct=processor_dct if (args.with_vector) else None)}
+ if args.masks:
+ postprocessors['segm'] = PostProcessSegm(processor_dct=processor_dct if args.with_vector else None)
return model, criterion, postprocessors
-
+''' deprecated
def build_teacher(args):
if args.dataset_file == 'coco':
num_classes = 91
+ if args.dataset_file == 'ymir':
+ num_classes = args.num_classes
+
if args.dataset_file == "coco_panoptic":
num_classes = 250
@@ -436,3 +512,4 @@ def build_teacher(args):
)
return model
+'''
\ No newline at end of file
diff --git a/methods/vidt/postprocessor.py b/methods/vidt/postprocessor.py
index d0cd8f0..3f7eb2b 100644
--- a/methods/vidt/postprocessor.py
+++ b/methods/vidt/postprocessor.py
@@ -11,18 +11,22 @@
from torch import nn
from util import box_ops
import torch.nn.functional as F
+from util.detectron2.utils.memory import retry_if_cuda_oom
+from util.detectron2.layers.mask_ops import paste_masks_in_image
+import numpy as np
class PostProcess(nn.Module):
- """ This module converts the model's output into the format expected by the coco api"""
+ """ This module converts the model's output into the format expected by the coco api"""
- def __init__(self, dataset_file):
- super().__init__()
- self.dataset_file = dataset_file
+ def __init__(self, processor_dct=None):
+ super().__init__()
+ # For instance segmentation using UQR module
+ self.processor_dct = processor_dct
- @torch.no_grad()
- def forward(self, outputs, target_sizes, target_boxes=None):
- """ Perform the computation
+ @torch.no_grad()
+ def forward(self, outputs, target_sizes):
+ """ Perform the computation
Parameters:
outputs: raw outputs of the model
@@ -31,48 +35,83 @@ def forward(self, outputs, target_sizes, target_boxes=None):
For visualization, this should be the image size after data augment, but before padding
"""
- out_logits, out_bbox = outputs['pred_logits'], outputs['pred_boxes']
-
- assert len(out_logits) == len(target_sizes)
- assert target_sizes.shape[1] == 2
-
- if self.dataset_file == 'coco':
- prob = out_logits.sigmoid()
- topk_values, topk_indexes = torch.topk(prob.view(out_logits.shape[0], -1), 100, dim=1)
- scores = topk_values
- topk_boxes = topk_indexes // out_logits.shape[2]
- labels = topk_indexes % out_logits.shape[2]
- boxes = box_ops.box_cxcywh_to_xyxy(out_bbox)
- boxes = torch.gather(boxes, 1, topk_boxes.unsqueeze(-1).repeat(1, 1, 4))
-
- # and from relative [0, 1] to absolute [0, height] coordinates
- img_h, img_w = target_sizes.unbind(1)
- scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1).to(torch.float32)
- boxes = boxes * scale_fct[:, None, :]
-
- results = [{'scores': s, 'labels': l, 'boxes': b} for s, l, b in zip(scores, labels, boxes)]
-
- return results
-
- elif self.dataset_file == 'voc':
-
- prob = F.softmax(out_logits, -1)
- scores, labels = prob.max(-1)
-
- # convert to [x0, y0, x1, y1] format
- boxes = box_ops.box_cxcywh_to_xyxy(out_bbox)
- # and from relative [0, 1] to absolute [0, height] coordinates
- img_h, img_w = target_sizes.unbind(1)
- scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1).to(torch.float32)
- boxes = boxes * scale_fct[:, None, :]
-
- results = [{'scores': s, 'labels': l, 'boxes': b} for s, l, b in zip(scores, labels, boxes)]
-
- true_boxes = []
- for id in range(len(target_boxes)):
- true_box = box_ops.box_cxcywh_to_xyxy(target_boxes[id])
- scale_fct = torch.stack([img_w[id], img_h[id], img_w[id], img_h[id]], dim=0).to(torch.float32)
- true_boxes.append(true_box * scale_fct)
-
- return results, true_boxes
-
+ if self.processor_dct is not None:
+ out_logits, out_bbox, out_vector = outputs['pred_logits'], outputs['pred_boxes'], outputs['pred_vectors']
+ else:
+ out_logits, out_bbox = outputs['pred_logits'], outputs['pred_boxes']
+
+ assert len(out_logits) == len(target_sizes)
+ assert target_sizes.shape[1] == 2
+
+ prob = out_logits.sigmoid()
+ topk_values, topk_indexes = torch.topk(prob.view(out_logits.shape[0], -1), 100, dim=1)
+ scores = topk_values
+ topk_boxes = topk_indexes // out_logits.shape[2]
+ labels = topk_indexes % out_logits.shape[2]
+ boxes = box_ops.box_cxcywh_to_xyxy(out_bbox)
+ boxes = torch.gather(boxes, 1, topk_boxes.unsqueeze(-1).repeat(1, 1, 4))
+
+ # Added for Instance Segmentation
+ if self.processor_dct is not None:
+ n_keep = self.processor_dct.n_keep
+ vectors = torch.gather(out_vector, 1, topk_boxes.unsqueeze(-1).repeat(1, 1, n_keep))
+ ###########
+
+ # and from relative [0, 1] to absolute [0, height] coordinates
+ img_h, img_w = target_sizes.unbind(1)
+ scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1).to(torch.float32)
+ boxes = boxes * scale_fct[:, None, :]
+
+ # Added for Instance Segmentation
+ if self.processor_dct is not None:
+ masks = []
+ n_keep, gt_mask_len = self.processor_dct.n_keep, self.processor_dct.gt_mask_len
+ b, r, c = vectors.shape
+ for bi in range(b):
+ outputs_masks_per_image = []
+ for ri in range(r):
+ # here visual for training
+ idct = np.zeros((gt_mask_len**2))
+ idct[:n_keep] = vectors[bi, ri].cpu().numpy()
+ idct = self.processor_dct.inverse_zigzag(idct, gt_mask_len, gt_mask_len)
+ re_mask = cv2.idct(idct)
+ max_v = np.max(re_mask)
+ min_v = np.min(re_mask)
+ re_mask = np.where(re_mask > (max_v + min_v) / 2., 1, 0)
+ re_mask = torch.from_numpy(re_mask)[None].float()
+ outputs_masks_per_image.append(re_mask)
+ outputs_masks_per_image = torch.cat(outputs_masks_per_image, dim=0).to(out_vector.device)
+ # here padding local mask to global mask
+ outputs_masks_per_image = retry_if_cuda_oom(paste_masks_in_image)(
+ outputs_masks_per_image, # N, 1, M, M
+ boxes[bi],
+ (img_h[bi], img_w[bi]),
+ threshold=0.5,
+ )
+ outputs_masks_per_image = outputs_masks_per_image.unsqueeze(1).cpu()
+ masks.append(outputs_masks_per_image)
+ ###########
+
+ if self.processor_dct is None:
+ results = [{'scores': s, 'labels': l, 'boxes': b} for s, l, b in zip(scores, labels, boxes)]
+ else:
+ results = [{
+ 'scores': s,
+ 'labels': l,
+ 'boxes': b,
+ 'masks': m
+ } for s, l, b, m in zip(scores, labels, boxes, masks)]
+
+ return results
+
+
+class PostProcessSegm(nn.Module):
+
+ def __init__(self, threshold=0.5, processor_dct=None):
+ super().__init__()
+ self.threshold = threshold
+ self.processor_dct = processor_dct
+
+ @torch.no_grad()
+ def forward(self, results, outputs, orig_target_sizes, max_target_sizes):
+ return results
diff --git a/methods/vidt_wo_neck/detector.py b/methods/vidt_wo_neck/detector.py
index c3b223a..f130a12 100644
--- a/methods/vidt_wo_neck/detector.py
+++ b/methods/vidt_wo_neck/detector.py
@@ -110,6 +110,9 @@ def build(args):
if args.dataset_file == 'coco':
num_classes = 91
+ if args.dataset_file == 'ymir':
+ num_classes = args.num_classes
+
if args.dataset_file == "coco_panoptic":
num_classes = 250
device = torch.device(args.device)
diff --git a/mypy.ini b/mypy.ini
new file mode 100644
index 0000000..1e256fc
--- /dev/null
+++ b/mypy.ini
@@ -0,0 +1,15 @@
+[mypy]
+allow_untyped_defs = True
+warn_unused_configs = True
+ignore_missing_imports = True
+
+exclude = (?x)(
+ ^methods
+ | ^ops
+ | ^util
+ | ^datasets
+ | main.py
+ | engine.py
+ | arguments.py
+ | fps_calculator.py
+ )
diff --git a/ops/setup.py b/ops/setup.py
index a0131bc..a084995 100644
--- a/ops/setup.py
+++ b/ops/setup.py
@@ -6,17 +6,12 @@
# Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
# ------------------------------------------------------------------------------------------------
-import os
import glob
+import os
import torch
-
-from torch.utils.cpp_extension import CUDA_HOME
-from torch.utils.cpp_extension import CppExtension
-from torch.utils.cpp_extension import CUDAExtension
-
-from setuptools import find_packages
-from setuptools import setup
+from setuptools import find_packages, setup
+from torch.utils.cpp_extension import CUDA_HOME, CppExtension, CUDAExtension
requirements = ["torch", "torchvision"]
@@ -33,7 +28,8 @@ def get_extensions():
extra_compile_args = {"cxx": []}
define_macros = []
- if torch.cuda.is_available() and CUDA_HOME is not None:
+ # remove torch.cuda.is_available() for docker build
+ if CUDA_HOME is not None:
extension = CUDAExtension
sources += source_cuda
define_macros += [("WITH_CUDA", None)]
diff --git a/requirements.txt b/requirements.txt
index 3f24c42..9970147 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,7 +1,9 @@
scipy
pycocotools
cython
-onnx
-onnxruntime
+# onnx
+# onnxruntime
timm
-einops
\ No newline at end of file
+einops
+opencv_python
+nptyping
diff --git a/util/detectron2/layers/mask_ops.py b/util/detectron2/layers/mask_ops.py
new file mode 100644
index 0000000..ce83929
--- /dev/null
+++ b/util/detectron2/layers/mask_ops.py
@@ -0,0 +1,275 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+import numpy as np
+from typing import Tuple
+import torch
+from PIL import Image
+from torch.nn import functional as F
+
+__all__ = ["paste_masks_in_image"]
+
+
+BYTES_PER_FLOAT = 4
+# TODO: This memory limit may be too much or too little. It would be better to
+# determine it based on available resources.
+GPU_MEM_LIMIT = 1024 ** 3 # 1 GB memory limit
+
+
+def _do_paste_mask(masks, boxes, img_h: int, img_w: int, skip_empty: bool = True):
+ """
+ Args:
+ masks: N, 1, H, W
+ boxes: N, 4
+ img_h, img_w (int):
+ skip_empty (bool): only paste masks within the region that
+ tightly bound all boxes, and returns the results this region only.
+ An important optimization for CPU.
+
+ Returns:
+ if skip_empty == False, a mask of shape (N, img_h, img_w)
+ if skip_empty == True, a mask of shape (N, h', w'), and the slice
+ object for the corresponding region.
+ """
+ # On GPU, paste all masks together (up to chunk size)
+ # by using the entire image to sample the masks
+ # Compared to pasting them one by one,
+ # this has more operations but is faster on COCO-scale dataset.
+ device = masks.device
+
+ if skip_empty and not torch.jit.is_scripting():
+ x0_int, y0_int = torch.clamp(boxes.min(dim=0).values.floor()[:2] - 1, min=0).to(
+ dtype=torch.int32
+ )
+ x1_int = torch.clamp(boxes[:, 2].max().ceil() + 1, max=img_w).to(dtype=torch.int32)
+ y1_int = torch.clamp(boxes[:, 3].max().ceil() + 1, max=img_h).to(dtype=torch.int32)
+ else:
+ x0_int, y0_int = 0, 0
+ x1_int, y1_int = img_w, img_h
+ x0, y0, x1, y1 = torch.split(boxes, 1, dim=1) # each is Nx1
+
+ N = masks.shape[0]
+
+ img_y = torch.arange(y0_int, y1_int, device=device, dtype=torch.float32) + 0.5
+ img_x = torch.arange(x0_int, x1_int, device=device, dtype=torch.float32) + 0.5
+ img_y = (img_y - y0) / (y1 - y0) * 2 - 1
+ img_x = (img_x - x0) / (x1 - x0) * 2 - 1
+ # img_x, img_y have shapes (N, w), (N, h)
+
+ gx = img_x[:, None, :].expand(N, img_y.size(1), img_x.size(1))
+ gy = img_y[:, :, None].expand(N, img_y.size(1), img_x.size(1))
+ grid = torch.stack([gx, gy], dim=3)
+
+ if not torch.jit.is_scripting():
+ if not masks.dtype.is_floating_point:
+ masks = masks.float()
+ img_masks = F.grid_sample(masks, grid.to(masks.dtype), align_corners=False)
+
+ if skip_empty and not torch.jit.is_scripting():
+ return img_masks[:, 0], (slice(y0_int, y1_int), slice(x0_int, x1_int))
+ else:
+ return img_masks[:, 0], ()
+
+
+# Annotate boxes as Tensor (but not Boxes) in order to use scripting
+#@torch.jit.script_if_tracing
+def paste_masks_in_image(
+ masks: torch.Tensor, boxes: torch.Tensor, image_shape: Tuple[int, int], threshold: float = 0.5
+):
+ """
+ Paste a set of masks that are of a fixed resolution (e.g., 28 x 28) into an image.
+ The location, height, and width for pasting each mask is determined by their
+ corresponding bounding boxes in boxes.
+
+ Note:
+ This is a complicated but more accurate implementation. In actual deployment, it is
+ often enough to use a faster but less accurate implementation.
+ See :func:`paste_mask_in_image_old` in this file for an alternative implementation.
+
+ Args:
+ masks (tensor): Tensor of shape (Bimg, Hmask, Wmask), where Bimg is the number of
+ detected object instances in the image and Hmask, Wmask are the mask width and mask
+ height of the predicted mask (e.g., Hmask = Wmask = 28). Values are in [0, 1].
+ boxes (Boxes or Tensor): A Boxes of length Bimg or Tensor of shape (Bimg, 4).
+ boxes[i] and masks[i] correspond to the same object instance.
+ image_shape (tuple): height, width
+ threshold (float): A threshold in [0, 1] for converting the (soft) masks to
+ binary masks.
+
+ Returns:
+ img_masks (Tensor): A tensor of shape (Bimg, Himage, Wimage), where Bimg is the
+ number of detected object instances and Himage, Wimage are the image width
+ and height. img_masks[i] is a binary mask for object instance i.
+ """
+
+ assert masks.shape[-1] == masks.shape[-2], "Only square mask predictions are supported"
+ N = len(masks)
+ if N == 0:
+ return masks.new_empty((0,) + image_shape, dtype=torch.uint8)
+ if not isinstance(boxes, torch.Tensor):
+ boxes = boxes.tensor
+ device = boxes.device
+ assert len(boxes) == N, boxes.shape
+
+ img_h, img_w = image_shape
+
+ # The actual implementation split the input into chunks,
+ # and paste them chunk by chunk.
+ if device.type == "cpu" or torch.jit.is_scripting():
+ # CPU is most efficient when they are pasted one by one with skip_empty=True
+ # so that it performs minimal number of operations.
+ num_chunks = N
+ else:
+ # GPU benefits from parallelism for larger chunks, but may have memory issue
+ # int(img_h) because shape may be tensors in tracing
+ num_chunks = int(np.ceil(N * int(img_h) * int(img_w) * BYTES_PER_FLOAT / GPU_MEM_LIMIT))
+ assert (
+ num_chunks <= N
+ ), "Default GPU_MEM_LIMIT in mask_ops.py is too small; try increasing it"
+ chunks = torch.chunk(torch.arange(N, device=device), num_chunks)
+
+ img_masks = torch.zeros(
+ N, img_h, img_w, device=device, dtype=torch.bool if threshold >= 0 else torch.uint8
+ )
+ for inds in chunks:
+ masks_chunk, spatial_inds = _do_paste_mask(
+ masks[inds, None, :, :], boxes[inds], img_h, img_w, skip_empty=device.type == "cpu"
+ )
+
+ if threshold >= 0:
+ masks_chunk = (masks_chunk >= threshold).to(dtype=torch.bool)
+ else:
+ # for visualization and debugging
+ masks_chunk = (masks_chunk * 255).to(dtype=torch.uint8)
+
+ if torch.jit.is_scripting(): # Scripting does not use the optimized codepath
+ img_masks[inds] = masks_chunk
+ else:
+ img_masks[(inds,) + spatial_inds] = masks_chunk
+ return img_masks
+
+
+# The below are the original paste function (from Detectron1) which has
+# larger quantization error.
+# It is faster on CPU, while the aligned one is faster on GPU thanks to grid_sample.
+
+
+def paste_mask_in_image_old(mask, box, img_h, img_w, threshold):
+ """
+ Paste a single mask in an image.
+ This is a per-box implementation of :func:`paste_masks_in_image`.
+ This function has larger quantization error due to incorrect pixel
+ modeling and is not used any more.
+
+ Args:
+ mask (Tensor): A tensor of shape (Hmask, Wmask) storing the mask of a single
+ object instance. Values are in [0, 1].
+ box (Tensor): A tensor of shape (4, ) storing the x0, y0, x1, y1 box corners
+ of the object instance.
+ img_h, img_w (int): Image height and width.
+ threshold (float): Mask binarization threshold in [0, 1].
+
+ Returns:
+ im_mask (Tensor):
+ The resized and binarized object mask pasted into the original
+ image plane (a tensor of shape (img_h, img_w)).
+ """
+ # Conversion from continuous box coordinates to discrete pixel coordinates
+ # via truncation (cast to int32). This determines which pixels to paste the
+ # mask onto.
+ box = box.to(dtype=torch.int32) # Continuous to discrete coordinate conversion
+ # An example (1D) box with continuous coordinates (x0=0.7, x1=4.3) will map to
+ # a discrete coordinates (x0=0, x1=4). Note that box is mapped to 5 = x1 - x0 + 1
+ # pixels (not x1 - x0 pixels).
+ samples_w = box[2] - box[0] + 1 # Number of pixel samples, *not* geometric width
+ samples_h = box[3] - box[1] + 1 # Number of pixel samples, *not* geometric height
+
+ # Resample the mask from it's original grid to the new samples_w x samples_h grid
+ mask = Image.fromarray(mask.cpu().numpy())
+ mask = mask.resize((samples_w, samples_h), resample=Image.BILINEAR)
+ mask = np.array(mask, copy=False)
+
+ if threshold >= 0:
+ mask = np.array(mask > threshold, dtype=np.uint8)
+ mask = torch.from_numpy(mask)
+ else:
+ # for visualization and debugging, we also
+ # allow it to return an unmodified mask
+ mask = torch.from_numpy(mask * 255).to(torch.uint8)
+
+ im_mask = torch.zeros((img_h, img_w), dtype=torch.uint8)
+ x_0 = max(box[0], 0)
+ x_1 = min(box[2] + 1, img_w)
+ y_0 = max(box[1], 0)
+ y_1 = min(box[3] + 1, img_h)
+
+ im_mask[y_0:y_1, x_0:x_1] = mask[
+ (y_0 - box[1]) : (y_1 - box[1]), (x_0 - box[0]) : (x_1 - box[0])
+ ]
+ return im_mask
+
+
+# Our pixel modeling requires extrapolation for any continuous
+# coordinate < 0.5 or > length - 0.5. When sampling pixels on the masks,
+# we would like this extrapolation to be an interpolation between boundary values and zero,
+# instead of using absolute zero or boundary values.
+# Therefore `paste_mask_in_image_old` is often used with zero padding around the masks like this:
+# masks, scale = pad_masks(masks[:, 0, :, :], 1)
+# boxes = scale_boxes(boxes.tensor, scale)
+
+
+def pad_masks(masks, padding):
+ """
+ Args:
+ masks (tensor): A tensor of shape (B, M, M) representing B masks.
+ padding (int): Number of cells to pad on all sides.
+
+ Returns:
+ The padded masks and the scale factor of the padding size / original size.
+ """
+ B = masks.shape[0]
+ M = masks.shape[-1]
+ pad2 = 2 * padding
+ scale = float(M + pad2) / M
+ padded_masks = masks.new_zeros((B, M + pad2, M + pad2))
+ padded_masks[:, padding:-padding, padding:-padding] = masks
+ return padded_masks, scale
+
+
+def scale_boxes(boxes, scale):
+ """
+ Args:
+ boxes (tensor): A tensor of shape (B, 4) representing B boxes with 4
+ coords representing the corners x0, y0, x1, y1,
+ scale (float): The box scaling factor.
+
+ Returns:
+ Scaled boxes.
+ """
+ w_half = (boxes[:, 2] - boxes[:, 0]) * 0.5
+ h_half = (boxes[:, 3] - boxes[:, 1]) * 0.5
+ x_c = (boxes[:, 2] + boxes[:, 0]) * 0.5
+ y_c = (boxes[:, 3] + boxes[:, 1]) * 0.5
+
+ w_half *= scale
+ h_half *= scale
+
+ scaled_boxes = torch.zeros_like(boxes)
+ scaled_boxes[:, 0] = x_c - w_half
+ scaled_boxes[:, 2] = x_c + w_half
+ scaled_boxes[:, 1] = y_c - h_half
+ scaled_boxes[:, 3] = y_c + h_half
+ return scaled_boxes
+
+
+#@torch.jit.script_if_tracing
+def _paste_masks_tensor_shape(
+ masks: torch.Tensor,
+ boxes: torch.Tensor,
+ image_shape: Tuple[torch.Tensor, torch.Tensor],
+ threshold: float = 0.5,
+):
+ """
+ A wrapper of paste_masks_in_image where image_shape is Tensor.
+ During tracing, shapes might be tensors instead of ints. The Tensor->int
+ conversion should be scripted rather than traced.
+ """
+ return paste_masks_in_image(masks, boxes, (int(image_shape[0]), int(image_shape[1])), threshold)
\ No newline at end of file
diff --git a/util/detectron2/layers/roi_align.py b/util/detectron2/layers/roi_align.py
new file mode 100644
index 0000000..a511afd
--- /dev/null
+++ b/util/detectron2/layers/roi_align.py
@@ -0,0 +1,74 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+from torch import nn
+from torchvision.ops import roi_align
+
+
+# NOTE: torchvision's RoIAlign has a different default aligned=False
+class ROIAlign(nn.Module):
+ def __init__(self, output_size, spatial_scale, sampling_ratio, aligned=True):
+ """
+ Args:
+ output_size (tuple): h, w
+ spatial_scale (float): scale the input boxes by this number
+ sampling_ratio (int): number of inputs samples to take for each output
+ sample. 0 to take samples densely.
+ aligned (bool): if False, use the legacy implementation in
+ Detectron. If True, align the results more perfectly.
+
+ Note:
+ The meaning of aligned=True:
+
+ Given a continuous coordinate c, its two neighboring pixel indices (in our
+ pixel model) are computed by floor(c - 0.5) and ceil(c - 0.5). For example,
+ c=1.3 has pixel neighbors with discrete indices [0] and [1] (which are sampled
+ from the underlying signal at continuous coordinates 0.5 and 1.5). But the original
+ roi_align (aligned=False) does not subtract the 0.5 when computing neighboring
+ pixel indices and therefore it uses pixels with a slightly incorrect alignment
+ (relative to our pixel model) when performing bilinear interpolation.
+
+ With `aligned=True`,
+ we first appropriately scale the ROI and then shift it by -0.5
+ prior to calling roi_align. This produces the correct neighbors; see
+ detectron2/tests/test_roi_align.py for verification.
+
+ The difference does not make a difference to the model's performance if
+ ROIAlign is used together with conv layers.
+ """
+ super().__init__()
+ self.output_size = output_size
+ self.spatial_scale = spatial_scale
+ self.sampling_ratio = sampling_ratio
+ self.aligned = aligned
+
+ from torchvision import __version__
+
+ version = tuple(int(x) for x in __version__.split(".")[:2])
+ # https://github.com/pytorch/vision/pull/2438
+ #assert version >= (0, 7), "Require torchvision >= 0.7"
+
+ def forward(self, input, rois):
+ """
+ Args:
+ input: NCHW images
+ rois: Bx5 boxes. First column is the index into N. The other 4 columns are xyxy.
+ """
+ assert rois.dim() == 2 and rois.size(1) == 5
+ if input.is_quantized:
+ input = input.dequantize()
+ return roi_align(
+ input,
+ rois.to(dtype=input.dtype),
+ self.output_size,
+ self.spatial_scale,
+ self.sampling_ratio,
+ self.aligned,
+ )
+
+ def __repr__(self):
+ tmpstr = self.__class__.__name__ + "("
+ tmpstr += "output_size=" + str(self.output_size)
+ tmpstr += ", spatial_scale=" + str(self.spatial_scale)
+ tmpstr += ", sampling_ratio=" + str(self.sampling_ratio)
+ tmpstr += ", aligned=" + str(self.aligned)
+ tmpstr += ")"
+ return tmpstr
\ No newline at end of file
diff --git a/util/detectron2/structures/boxes.py b/util/detectron2/structures/boxes.py
new file mode 100644
index 0000000..8b0c005
--- /dev/null
+++ b/util/detectron2/structures/boxes.py
@@ -0,0 +1,423 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+import math
+import numpy as np
+from enum import IntEnum, unique
+from typing import List, Tuple, Union
+import torch
+from torch import device
+
+_RawBoxType = Union[List[float], Tuple[float, ...], torch.Tensor, np.ndarray]
+
+
+@unique
+class BoxMode(IntEnum):
+ """
+ Enum of different ways to represent a box.
+ """
+
+ XYXY_ABS = 0
+ """
+ (x0, y0, x1, y1) in absolute floating points coordinates.
+ The coordinates in range [0, width or height].
+ """
+ XYWH_ABS = 1
+ """
+ (x0, y0, w, h) in absolute floating points coordinates.
+ """
+ XYXY_REL = 2
+ """
+ Not yet supported!
+ (x0, y0, x1, y1) in range [0, 1]. They are relative to the size of the image.
+ """
+ XYWH_REL = 3
+ """
+ Not yet supported!
+ (x0, y0, w, h) in range [0, 1]. They are relative to the size of the image.
+ """
+ XYWHA_ABS = 4
+ """
+ (xc, yc, w, h, a) in absolute floating points coordinates.
+ (xc, yc) is the center of the rotated box, and the angle a is in degrees ccw.
+ """
+
+ @staticmethod
+ def convert(box: _RawBoxType, from_mode: "BoxMode", to_mode: "BoxMode") -> _RawBoxType:
+ """
+ Args:
+ box: can be a k-tuple, k-list or an Nxk array/tensor, where k = 4 or 5
+ from_mode, to_mode (BoxMode)
+
+ Returns:
+ The converted box of the same type.
+ """
+ if from_mode == to_mode:
+ return box
+
+ original_type = type(box)
+ is_numpy = isinstance(box, np.ndarray)
+ single_box = isinstance(box, (list, tuple))
+ if single_box:
+ assert len(box) == 4 or len(box) == 5, (
+ "BoxMode.convert takes either a k-tuple/list or an Nxk array/tensor,"
+ " where k == 4 or 5"
+ )
+ arr = torch.tensor(box)[None, :]
+ else:
+ # avoid modifying the input box
+ if is_numpy:
+ arr = torch.from_numpy(np.asarray(box)).clone()
+ else:
+ arr = box.clone()
+
+ assert to_mode not in [BoxMode.XYXY_REL, BoxMode.XYWH_REL] and from_mode not in [
+ BoxMode.XYXY_REL,
+ BoxMode.XYWH_REL,
+ ], "Relative mode not yet supported!"
+
+ if from_mode == BoxMode.XYWHA_ABS and to_mode == BoxMode.XYXY_ABS:
+ assert (
+ arr.shape[-1] == 5
+ ), "The last dimension of input shape must be 5 for XYWHA format"
+ original_dtype = arr.dtype
+ arr = arr.double()
+
+ w = arr[:, 2]
+ h = arr[:, 3]
+ a = arr[:, 4]
+ c = torch.abs(torch.cos(a * math.pi / 180.0))
+ s = torch.abs(torch.sin(a * math.pi / 180.0))
+ # This basically computes the horizontal bounding rectangle of the rotated box
+ new_w = c * w + s * h
+ new_h = c * h + s * w
+
+ # convert center to top-left corner
+ arr[:, 0] -= new_w / 2.0
+ arr[:, 1] -= new_h / 2.0
+ # bottom-right corner
+ arr[:, 2] = arr[:, 0] + new_w
+ arr[:, 3] = arr[:, 1] + new_h
+
+ arr = arr[:, :4].to(dtype=original_dtype)
+ elif from_mode == BoxMode.XYWH_ABS and to_mode == BoxMode.XYWHA_ABS:
+ original_dtype = arr.dtype
+ arr = arr.double()
+ arr[:, 0] += arr[:, 2] / 2.0
+ arr[:, 1] += arr[:, 3] / 2.0
+ angles = torch.zeros((arr.shape[0], 1), dtype=arr.dtype)
+ arr = torch.cat((arr, angles), axis=1).to(dtype=original_dtype)
+ else:
+ if to_mode == BoxMode.XYXY_ABS and from_mode == BoxMode.XYWH_ABS:
+ arr[:, 2] += arr[:, 0]
+ arr[:, 3] += arr[:, 1]
+ elif from_mode == BoxMode.XYXY_ABS and to_mode == BoxMode.XYWH_ABS:
+ arr[:, 2] -= arr[:, 0]
+ arr[:, 3] -= arr[:, 1]
+ else:
+ raise NotImplementedError(
+ "Conversion from BoxMode {} to {} is not supported yet".format(
+ from_mode, to_mode
+ )
+ )
+
+ if single_box:
+ return original_type(arr.flatten().tolist())
+ if is_numpy:
+ return arr.numpy()
+ else:
+ return arr
+
+
+class Boxes:
+ """
+ This structure stores a list of boxes as a Nx4 torch.Tensor.
+ It supports some common methods about boxes
+ (`area`, `clip`, `nonempty`, etc),
+ and also behaves like a Tensor
+ (support indexing, `to(device)`, `.device`, and iteration over all boxes)
+
+ Attributes:
+ tensor (torch.Tensor): float matrix of Nx4. Each row is (x1, y1, x2, y2).
+ """
+
+ def __init__(self, tensor: torch.Tensor):
+ """
+ Args:
+ tensor (Tensor[float]): a Nx4 matrix. Each row is (x1, y1, x2, y2).
+ """
+ device = tensor.device if isinstance(tensor, torch.Tensor) else torch.device("cpu")
+ tensor = torch.as_tensor(tensor, dtype=torch.float32, device=device)
+ if tensor.numel() == 0:
+ # Use reshape, so we don't end up creating a new tensor that does not depend on
+ # the inputs (and consequently confuses jit)
+ tensor = tensor.reshape((-1, 4)).to(dtype=torch.float32, device=device)
+ assert tensor.dim() == 2 and tensor.size(-1) == 4, tensor.size()
+
+ self.tensor = tensor
+
+ def clone(self) -> "Boxes":
+ """
+ Clone the Boxes.
+
+ Returns:
+ Boxes
+ """
+ return Boxes(self.tensor.clone())
+
+ def to(self, device: torch.device):
+ # Boxes are assumed float32 and does not support to(dtype)
+ return Boxes(self.tensor.to(device=device))
+
+ def area(self) -> torch.Tensor:
+ """
+ Computes the area of all the boxes.
+
+ Returns:
+ torch.Tensor: a vector with areas of each box.
+ """
+ box = self.tensor
+ area = (box[:, 2] - box[:, 0]) * (box[:, 3] - box[:, 1])
+ return area
+
+ def clip(self, box_size: Tuple[int, int]) -> None:
+ """
+ Clip (in place) the boxes by limiting x coordinates to the range [0, width]
+ and y coordinates to the range [0, height].
+
+ Args:
+ box_size (height, width): The clipping box's size.
+ """
+ assert torch.isfinite(self.tensor).all(), "Box tensor contains infinite or NaN!"
+ h, w = box_size
+ x1 = self.tensor[:, 0].clamp(min=0, max=w)
+ y1 = self.tensor[:, 1].clamp(min=0, max=h)
+ x2 = self.tensor[:, 2].clamp(min=0, max=w)
+ y2 = self.tensor[:, 3].clamp(min=0, max=h)
+ self.tensor = torch.stack((x1, y1, x2, y2), dim=-1)
+
+ def nonempty(self, threshold: float = 0.0) -> torch.Tensor:
+ """
+ Find boxes that are non-empty.
+ A box is considered empty, if either of its side is no larger than threshold.
+
+ Returns:
+ Tensor:
+ a binary vector which represents whether each box is empty
+ (False) or non-empty (True).
+ """
+ box = self.tensor
+ widths = box[:, 2] - box[:, 0]
+ heights = box[:, 3] - box[:, 1]
+ keep = (widths > threshold) & (heights > threshold)
+ return keep
+
+ def __getitem__(self, item) -> "Boxes":
+ """
+ Args:
+ item: int, slice, or a BoolTensor
+
+ Returns:
+ Boxes: Create a new :class:`Boxes` by indexing.
+
+ The following usage are allowed:
+
+ 1. `new_boxes = boxes[3]`: return a `Boxes` which contains only one box.
+ 2. `new_boxes = boxes[2:10]`: return a slice of boxes.
+ 3. `new_boxes = boxes[vector]`, where vector is a torch.BoolTensor
+ with `length = len(boxes)`. Nonzero elements in the vector will be selected.
+
+ Note that the returned Boxes might share storage with this Boxes,
+ subject to Pytorch's indexing semantics.
+ """
+ if isinstance(item, int):
+ return Boxes(self.tensor[item].view(1, -1))
+ b = self.tensor[item]
+ assert b.dim() == 2, "Indexing on Boxes with {} failed to return a matrix!".format(item)
+ return Boxes(b)
+
+ def __len__(self) -> int:
+ return self.tensor.shape[0]
+
+ def __repr__(self) -> str:
+ return "Boxes(" + str(self.tensor) + ")"
+
+ def inside_box(self, box_size: Tuple[int, int], boundary_threshold: int = 0) -> torch.Tensor:
+ """
+ Args:
+ box_size (height, width): Size of the reference box.
+ boundary_threshold (int): Boxes that extend beyond the reference box
+ boundary by more than boundary_threshold are considered "outside".
+
+ Returns:
+ a binary vector, indicating whether each box is inside the reference box.
+ """
+ height, width = box_size
+ inds_inside = (
+ (self.tensor[..., 0] >= -boundary_threshold)
+ & (self.tensor[..., 1] >= -boundary_threshold)
+ & (self.tensor[..., 2] < width + boundary_threshold)
+ & (self.tensor[..., 3] < height + boundary_threshold)
+ )
+ return inds_inside
+
+ def get_centers(self) -> torch.Tensor:
+ """
+ Returns:
+ The box centers in a Nx2 array of (x, y).
+ """
+ return (self.tensor[:, :2] + self.tensor[:, 2:]) / 2
+
+ def scale(self, scale_x: float, scale_y: float) -> None:
+ """
+ Scale the box with horizontal and vertical scaling factors
+ """
+ self.tensor[:, 0::2] *= scale_x
+ self.tensor[:, 1::2] *= scale_y
+
+ @classmethod
+ def cat(cls, boxes_list: List["Boxes"]) -> "Boxes":
+ """
+ Concatenates a list of Boxes into a single Boxes
+
+ Arguments:
+ boxes_list (list[Boxes])
+
+ Returns:
+ Boxes: the concatenated Boxes
+ """
+ assert isinstance(boxes_list, (list, tuple))
+ if len(boxes_list) == 0:
+ return cls(torch.empty(0))
+ assert all([isinstance(box, Boxes) for box in boxes_list])
+
+ # use torch.cat (v.s. layers.cat) so the returned boxes never share storage with input
+ cat_boxes = cls(torch.cat([b.tensor for b in boxes_list], dim=0))
+ return cat_boxes
+
+ @property
+ def device(self) -> device:
+ return self.tensor.device
+
+ # type "Iterator[torch.Tensor]", yield, and iter() not supported by torchscript
+ # https://github.com/pytorch/pytorch/issues/18627
+ @torch.jit.unused
+ def __iter__(self):
+ """
+ Yield a box as a Tensor of shape (4,) at a time.
+ """
+ yield from self.tensor
+
+
+def pairwise_intersection(boxes1: Boxes, boxes2: Boxes) -> torch.Tensor:
+ """
+ Given two lists of boxes of size N and M,
+ compute the intersection area between __all__ N x M pairs of boxes.
+ The box order must be (xmin, ymin, xmax, ymax)
+
+ Args:
+ boxes1,boxes2 (Boxes): two `Boxes`. Contains N & M boxes, respectively.
+
+ Returns:
+ Tensor: intersection, sized [N,M].
+ """
+ boxes1, boxes2 = boxes1.tensor, boxes2.tensor
+ width_height = torch.min(boxes1[:, None, 2:], boxes2[:, 2:]) - torch.max(
+ boxes1[:, None, :2], boxes2[:, :2]
+ ) # [N,M,2]
+
+ width_height.clamp_(min=0) # [N,M,2]
+ intersection = width_height.prod(dim=2) # [N,M]
+ return intersection
+
+
+# implementation from https://github.com/kuangliu/torchcv/blob/master/torchcv/utils/box.py
+# with slight modifications
+def pairwise_iou(boxes1: Boxes, boxes2: Boxes) -> torch.Tensor:
+ """
+ Given two lists of boxes of size N and M, compute the IoU
+ (intersection over union) between **all** N x M pairs of boxes.
+ The box order must be (xmin, ymin, xmax, ymax).
+
+ Args:
+ boxes1,boxes2 (Boxes): two `Boxes`. Contains N & M boxes, respectively.
+
+ Returns:
+ Tensor: IoU, sized [N,M].
+ """
+ area1 = boxes1.area() # [N]
+ area2 = boxes2.area() # [M]
+ inter = pairwise_intersection(boxes1, boxes2)
+
+ # handle empty boxes
+ iou = torch.where(
+ inter > 0,
+ inter / (area1[:, None] + area2 - inter),
+ torch.zeros(1, dtype=inter.dtype, device=inter.device),
+ )
+ return iou
+
+
+def pairwise_ioa(boxes1: Boxes, boxes2: Boxes) -> torch.Tensor:
+ """
+ Similar to :func:`pariwise_iou` but compute the IoA (intersection over boxes2 area).
+
+ Args:
+ boxes1,boxes2 (Boxes): two `Boxes`. Contains N & M boxes, respectively.
+
+ Returns:
+ Tensor: IoA, sized [N,M].
+ """
+ area2 = boxes2.area() # [M]
+ inter = pairwise_intersection(boxes1, boxes2)
+
+ # handle empty boxes
+ ioa = torch.where(
+ inter > 0, inter / area2, torch.zeros(1, dtype=inter.dtype, device=inter.device)
+ )
+ return ioa
+
+
+def pairwise_point_box_distance(points: torch.Tensor, boxes: Boxes):
+ """
+ Pairwise distance between N points and M boxes. The distance between a
+ point and a box is represented by the distance from the point to 4 edges
+ of the box. Distances are all positive when the point is inside the box.
+
+ Args:
+ points: Nx2 coordinates. Each row is (x, y)
+ boxes: M boxes
+
+ Returns:
+ Tensor: distances of size (N, M, 4). The 4 values are distances from
+ the point to the left, top, right, bottom of the box.
+ """
+ x, y = points.unsqueeze(dim=2).unbind(dim=1) # (N, 1)
+ x0, y0, x1, y1 = boxes.tensor.unsqueeze(dim=0).unbind(dim=2) # (1, M)
+ return torch.stack([x - x0, y - y0, x1 - x, y1 - y], dim=2)
+
+
+def matched_pairwise_iou(boxes1: Boxes, boxes2: Boxes) -> torch.Tensor:
+ """
+ Compute pairwise intersection over union (IOU) of two sets of matched
+ boxes that have the same number of boxes.
+ Similar to :func:`pairwise_iou`, but computes only diagonal elements of the matrix.
+
+ Args:
+ boxes1 (Boxes): bounding boxes, sized [N,4].
+ boxes2 (Boxes): same length as boxes1
+ Returns:
+ Tensor: iou, sized [N].
+ """
+ assert len(boxes1) == len(
+ boxes2
+ ), "boxlists should have the same" "number of entries, got {}, {}".format(
+ len(boxes1), len(boxes2)
+ )
+ area1 = boxes1.area() # [N]
+ area2 = boxes2.area() # [N]
+ box1, box2 = boxes1.tensor, boxes2.tensor
+ lt = torch.max(box1[:, :2], box2[:, :2]) # [N,2]
+ rb = torch.min(box1[:, 2:], box2[:, 2:]) # [N,2]
+ wh = (rb - lt).clamp(min=0) # [N,2]
+ inter = wh[:, 0] * wh[:, 1] # [N]
+ iou = inter / (area1 + area2 - inter) # [N]
+ return iou
\ No newline at end of file
diff --git a/util/detectron2/structures/masks.py b/util/detectron2/structures/masks.py
new file mode 100644
index 0000000..6c40b20
--- /dev/null
+++ b/util/detectron2/structures/masks.py
@@ -0,0 +1,531 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+import copy
+import itertools
+import numpy as np
+from typing import Any, Iterator, List, Union
+import pycocotools.mask as mask_util
+import torch
+from torch import device
+from ..layers.roi_align import ROIAlign
+from ..utils.memory import retry_if_cuda_oom
+
+from .boxes import Boxes
+
+
+def polygon_area(x, y):
+ # Using the shoelace formula
+ # https://stackoverflow.com/questions/24467972/calculate-area-of-polygon-given-x-y-coordinates
+ return 0.5 * np.abs(np.dot(x, np.roll(y, 1)) - np.dot(y, np.roll(x, 1)))
+
+
+def polygons_to_bitmask(polygons: List[np.ndarray], height: int, width: int) -> np.ndarray:
+ """
+ Args:
+ polygons (list[ndarray]): each array has shape (Nx2,)
+ height, width (int)
+
+ Returns:
+ ndarray: a bool mask of shape (height, width)
+ """
+ if len(polygons) == 0:
+ # COCOAPI does not support empty polygons
+ return np.zeros((height, width)).astype(np.bool)
+ rles = mask_util.frPyObjects(polygons, height, width)
+ rle = mask_util.merge(rles)
+ return mask_util.decode(rle).astype(np.bool)
+
+
+def rasterize_polygons_within_box(
+ polygons: List[np.ndarray], box: np.ndarray, mask_size: int
+) -> torch.Tensor:
+ """
+ Rasterize the polygons into a mask image and
+ crop the mask content in the given box.
+ The cropped mask is resized to (mask_size, mask_size).
+
+ This function is used when generating training targets for mask head in Mask R-CNN.
+ Given original ground-truth masks for an image, new ground-truth mask
+ training targets in the size of `mask_size x mask_size`
+ must be provided for each predicted box. This function will be called to
+ produce such targets.
+
+ Args:
+ polygons (list[ndarray[float]]): a list of polygons, which represents an instance.
+ box: 4-element numpy array
+ mask_size (int):
+
+ Returns:
+ Tensor: BoolTensor of shape (mask_size, mask_size)
+ """
+ # 1. Shift the polygons w.r.t the boxes
+ w, h = box[2] - box[0], box[3] - box[1]
+
+ polygons = copy.deepcopy(polygons)
+ for p in polygons:
+ p[0::2] = p[0::2] - box[0]
+ p[1::2] = p[1::2] - box[1]
+
+ # 2. Rescale the polygons to the new box size
+ # max() to avoid division by small number
+ ratio_h = mask_size / max(h, 0.1)
+ ratio_w = mask_size / max(w, 0.1)
+
+ if ratio_h == ratio_w:
+ for p in polygons:
+ p *= ratio_h
+ else:
+ for p in polygons:
+ p[0::2] *= ratio_w
+ p[1::2] *= ratio_h
+
+ # 3. Rasterize the polygons with coco api
+ mask = polygons_to_bitmask(polygons, mask_size, mask_size)
+ mask = torch.from_numpy(mask)
+ return mask
+
+
+class BitMasks:
+ """
+ This class stores the segmentation masks for all objects in one image, in
+ the form of bitmaps.
+
+ Attributes:
+ tensor: bool Tensor of N,H,W, representing N instances in the image.
+ """
+
+ def __init__(self, tensor: Union[torch.Tensor, np.ndarray]):
+ """
+ Args:
+ tensor: bool Tensor of N,H,W, representing N instances in the image.
+ """
+ device = tensor.device if isinstance(tensor, torch.Tensor) else torch.device("cpu")
+ tensor = torch.as_tensor(tensor, dtype=torch.bool, device=device)
+ assert tensor.dim() == 3, tensor.size()
+ self.image_size = tensor.shape[1:]
+ self.tensor = tensor
+
+ @torch.jit.unused
+ def to(self, *args: Any, **kwargs: Any) -> "BitMasks":
+ return BitMasks(self.tensor.to(*args, **kwargs))
+
+ @property
+ def device(self) -> torch.device:
+ return self.tensor.device
+
+ @torch.jit.unused
+ def __getitem__(self, item: Union[int, slice, torch.BoolTensor]) -> "BitMasks":
+ """
+ Returns:
+ BitMasks: Create a new :class:`BitMasks` by indexing.
+
+ The following usage are allowed:
+
+ 1. `new_masks = masks[3]`: return a `BitMasks` which contains only one mask.
+ 2. `new_masks = masks[2:10]`: return a slice of masks.
+ 3. `new_masks = masks[vector]`, where vector is a torch.BoolTensor
+ with `length = len(masks)`. Nonzero elements in the vector will be selected.
+
+ Note that the returned object might share storage with this object,
+ subject to Pytorch's indexing semantics.
+ """
+ if isinstance(item, int):
+ return BitMasks(self.tensor[item].unsqueeze(0))
+ m = self.tensor[item]
+ assert m.dim() == 3, "Indexing on BitMasks with {} returns a tensor with shape {}!".format(
+ item, m.shape
+ )
+ return BitMasks(m)
+
+ @torch.jit.unused
+ def __iter__(self) -> torch.Tensor:
+ yield from self.tensor
+
+ @torch.jit.unused
+ def __repr__(self) -> str:
+ s = self.__class__.__name__ + "("
+ s += "num_instances={})".format(len(self.tensor))
+ return s
+
+ def __len__(self) -> int:
+ return self.tensor.shape[0]
+
+ def nonempty(self) -> torch.Tensor:
+ """
+ Find masks that are non-empty.
+
+ Returns:
+ Tensor: a BoolTensor which represents
+ whether each mask is empty (False) or non-empty (True).
+ """
+ return self.tensor.flatten(1).any(dim=1)
+
+ @staticmethod
+ def from_polygon_masks(
+ polygon_masks: Union["PolygonMasks", List[List[np.ndarray]]], height: int, width: int
+ ) -> "BitMasks":
+ """
+ Args:
+ polygon_masks (list[list[ndarray]] or PolygonMasks)
+ height, width (int)
+ """
+ if isinstance(polygon_masks, PolygonMasks):
+ polygon_masks = polygon_masks.polygons
+ masks = [polygons_to_bitmask(p, height, width) for p in polygon_masks]
+ if len(masks):
+ return BitMasks(torch.stack([torch.from_numpy(x) for x in masks]))
+ else:
+ return BitMasks(torch.empty(0, height, width, dtype=torch.bool))
+
+ @staticmethod
+ def from_roi_masks(roi_masks: "ROIMasks", height: int, width: int) -> "BitMasks":
+ """
+ Args:
+ roi_masks:
+ height, width (int):
+ """
+ return roi_masks.to_bitmasks(height, width)
+
+ def crop_and_resize(self, boxes: torch.Tensor, mask_size: int) -> torch.Tensor:
+ """
+ Crop each bitmask by the given box, and resize results to (mask_size, mask_size).
+ This can be used to prepare training targets for Mask R-CNN.
+ It has less reconstruction error compared to rasterization with polygons.
+ However we observe no difference in accuracy,
+ but BitMasks requires more memory to store all the masks.
+
+ Args:
+ boxes (Tensor): Nx4 tensor storing the boxes for each mask
+ mask_size (int): the size of the rasterized mask.
+
+ Returns:
+ Tensor:
+ A bool tensor of shape (N, mask_size, mask_size), where
+ N is the number of predicted boxes for this image.
+ """
+ assert len(boxes) == len(self), "{} != {}".format(len(boxes), len(self))
+ device = self.tensor.device
+
+ batch_inds = torch.arange(len(boxes), device=device).to(dtype=boxes.dtype)[:, None]
+ rois = torch.cat([batch_inds, boxes], dim=1) # Nx5
+
+ bit_masks = self.tensor.to(dtype=torch.float32)
+ rois = rois.to(device=device)
+ output = (
+ ROIAlign((mask_size, mask_size), 1.0, 0, aligned=True)
+ .forward(bit_masks[:, None, :, :], rois)
+ .squeeze(1)
+ )
+ output = output >= 0.5
+ return output
+
+ def get_bounding_boxes(self) -> Boxes:
+ """
+ Returns:
+ Boxes: tight bounding boxes around bitmasks.
+ If a mask is empty, it's bounding box will be all zero.
+ """
+ boxes = torch.zeros(self.tensor.shape[0], 4, dtype=torch.float32)
+ x_any = torch.any(self.tensor, dim=1)
+ y_any = torch.any(self.tensor, dim=2)
+ for idx in range(self.tensor.shape[0]):
+ x = torch.where(x_any[idx, :])[0]
+ y = torch.where(y_any[idx, :])[0]
+ if len(x) > 0 and len(y) > 0:
+ boxes[idx, :] = torch.as_tensor(
+ [x[0], y[0], x[-1] + 1, y[-1] + 1], dtype=torch.float32
+ )
+ return Boxes(boxes)
+
+ @staticmethod
+ def cat(bitmasks_list: List["BitMasks"]) -> "BitMasks":
+ """
+ Concatenates a list of BitMasks into a single BitMasks
+
+ Arguments:
+ bitmasks_list (list[BitMasks])
+
+ Returns:
+ BitMasks: the concatenated BitMasks
+ """
+ assert isinstance(bitmasks_list, (list, tuple))
+ assert len(bitmasks_list) > 0
+ assert all(isinstance(bitmask, BitMasks) for bitmask in bitmasks_list)
+
+ cat_bitmasks = type(bitmasks_list[0])(torch.cat([bm.tensor for bm in bitmasks_list], dim=0))
+ return cat_bitmasks
+
+
+class PolygonMasks:
+ """
+ This class stores the segmentation masks for all objects in one image, in the form of polygons.
+
+ Attributes:
+ polygons: list[list[ndarray]]. Each ndarray is a float64 vector representing a polygon.
+ """
+
+ def __init__(self, polygons: List[List[Union[torch.Tensor, np.ndarray]]]):
+ """
+ Arguments:
+ polygons (list[list[np.ndarray]]): The first
+ level of the list correspond to individual instances,
+ the second level to all the polygons that compose the
+ instance, and the third level to the polygon coordinates.
+ The third level array should have the format of
+ [x0, y0, x1, y1, ..., xn, yn] (n >= 3).
+ """
+ if not isinstance(polygons, list):
+ raise ValueError(
+ "Cannot create PolygonMasks: Expect a list of list of polygons per image. "
+ "Got '{}' instead.".format(type(polygons))
+ )
+
+ def _make_array(t: Union[torch.Tensor, np.ndarray]) -> np.ndarray:
+ # Use float64 for higher precision, because why not?
+ # Always put polygons on CPU (self.to is a no-op) since they
+ # are supposed to be small tensors.
+ # May need to change this assumption if GPU placement becomes useful
+ if isinstance(t, torch.Tensor):
+ t = t.cpu().numpy()
+ return np.asarray(t).astype("float64")
+
+ def process_polygons(
+ polygons_per_instance: List[Union[torch.Tensor, np.ndarray]]
+ ) -> List[np.ndarray]:
+ if not isinstance(polygons_per_instance, list):
+ raise ValueError(
+ "Cannot create polygons: Expect a list of polygons per instance. "
+ "Got '{}' instead.".format(type(polygons_per_instance))
+ )
+ # transform each polygon to a numpy array
+ polygons_per_instance = [_make_array(p) for p in polygons_per_instance]
+ for polygon in polygons_per_instance:
+ if len(polygon) % 2 != 0 or len(polygon) < 6:
+ raise ValueError(f"Cannot create a polygon from {len(polygon)} coordinates.")
+ return polygons_per_instance
+
+ self.polygons: List[List[np.ndarray]] = [
+ process_polygons(polygons_per_instance) for polygons_per_instance in polygons
+ ]
+
+ def to(self, *args: Any, **kwargs: Any) -> "PolygonMasks":
+ return self
+
+ @property
+ def device(self) -> torch.device:
+ return torch.device("cpu")
+
+ def get_bounding_boxes(self) -> Boxes:
+ """
+ Returns:
+ Boxes: tight bounding boxes around polygon masks.
+ """
+ boxes = torch.zeros(len(self.polygons), 4, dtype=torch.float32)
+ for idx, polygons_per_instance in enumerate(self.polygons):
+ minxy = torch.as_tensor([float("inf"), float("inf")], dtype=torch.float32)
+ maxxy = torch.zeros(2, dtype=torch.float32)
+ for polygon in polygons_per_instance:
+ coords = torch.from_numpy(polygon).view(-1, 2).to(dtype=torch.float32)
+ minxy = torch.min(minxy, torch.min(coords, dim=0).values)
+ maxxy = torch.max(maxxy, torch.max(coords, dim=0).values)
+ boxes[idx, :2] = minxy
+ boxes[idx, 2:] = maxxy
+ return Boxes(boxes)
+
+ def nonempty(self) -> torch.Tensor:
+ """
+ Find masks that are non-empty.
+
+ Returns:
+ Tensor:
+ a BoolTensor which represents whether each mask is empty (False) or not (True).
+ """
+ keep = [1 if len(polygon) > 0 else 0 for polygon in self.polygons]
+ return torch.from_numpy(np.asarray(keep, dtype=np.bool))
+
+ def __getitem__(self, item: Union[int, slice, List[int], torch.BoolTensor]) -> "PolygonMasks":
+ """
+ Support indexing over the instances and return a `PolygonMasks` object.
+ `item` can be:
+
+ 1. An integer. It will return an object with only one instance.
+ 2. A slice. It will return an object with the selected instances.
+ 3. A list[int]. It will return an object with the selected instances,
+ correpsonding to the indices in the list.
+ 4. A vector mask of type BoolTensor, whose length is num_instances.
+ It will return an object with the instances whose mask is nonzero.
+ """
+ if isinstance(item, int):
+ selected_polygons = [self.polygons[item]]
+ elif isinstance(item, slice):
+ selected_polygons = self.polygons[item]
+ elif isinstance(item, list):
+ selected_polygons = [self.polygons[i] for i in item]
+ elif isinstance(item, torch.Tensor):
+ # Polygons is a list, so we have to move the indices back to CPU.
+ if item.dtype == torch.bool:
+ assert item.dim() == 1, item.shape
+ item = item.nonzero().squeeze(1).cpu().numpy().tolist()
+ elif item.dtype in [torch.int32, torch.int64]:
+ item = item.cpu().numpy().tolist()
+ else:
+ raise ValueError("Unsupported tensor dtype={} for indexing!".format(item.dtype))
+ selected_polygons = [self.polygons[i] for i in item]
+ return PolygonMasks(selected_polygons)
+
+ def __iter__(self) -> Iterator[List[np.ndarray]]:
+ """
+ Yields:
+ list[ndarray]: the polygons for one instance.
+ Each Tensor is a float64 vector representing a polygon.
+ """
+ return iter(self.polygons)
+
+ def __repr__(self) -> str:
+ s = self.__class__.__name__ + "("
+ s += "num_instances={})".format(len(self.polygons))
+ return s
+
+ def __len__(self) -> int:
+ return len(self.polygons)
+
+ def crop_and_resize(self, boxes: torch.Tensor, mask_size: int) -> torch.Tensor:
+ """
+ Crop each mask by the given box, and resize results to (mask_size, mask_size).
+ This can be used to prepare training targets for Mask R-CNN.
+
+ Args:
+ boxes (Tensor): Nx4 tensor storing the boxes for each mask
+ mask_size (int): the size of the rasterized mask.
+
+ Returns:
+ Tensor: A bool tensor of shape (N, mask_size, mask_size), where
+ N is the number of predicted boxes for this image.
+ """
+ assert len(boxes) == len(self), "{} != {}".format(len(boxes), len(self))
+
+ device = boxes.device
+ # Put boxes on the CPU, as the polygon representation is not efficient GPU-wise
+ # (several small tensors for representing a single instance mask)
+ boxes = boxes.to(torch.device("cpu"))
+
+ results = [
+ rasterize_polygons_within_box(poly, box.numpy(), mask_size)
+ for poly, box in zip(self.polygons, boxes)
+ ]
+ """
+ poly: list[list[float]], the polygons for one instance
+ box: a tensor of shape (4,)
+ """
+ if len(results) == 0:
+ return torch.empty(0, mask_size, mask_size, dtype=torch.bool, device=device)
+ return torch.stack(results, dim=0).to(device=device)
+
+ def area(self):
+ """
+ Computes area of the mask.
+ Only works with Polygons, using the shoelace formula:
+ https://stackoverflow.com/questions/24467972/calculate-area-of-polygon-given-x-y-coordinates
+
+ Returns:
+ Tensor: a vector, area for each instance
+ """
+
+ area = []
+ for polygons_per_instance in self.polygons:
+ area_per_instance = 0
+ for p in polygons_per_instance:
+ area_per_instance += polygon_area(p[0::2], p[1::2])
+ area.append(area_per_instance)
+
+ return torch.tensor(area)
+
+ @staticmethod
+ def cat(polymasks_list: List["PolygonMasks"]) -> "PolygonMasks":
+ """
+ Concatenates a list of PolygonMasks into a single PolygonMasks
+
+ Arguments:
+ polymasks_list (list[PolygonMasks])
+
+ Returns:
+ PolygonMasks: the concatenated PolygonMasks
+ """
+ assert isinstance(polymasks_list, (list, tuple))
+ assert len(polymasks_list) > 0
+ assert all(isinstance(polymask, PolygonMasks) for polymask in polymasks_list)
+
+ cat_polymasks = type(polymasks_list[0])(
+ list(itertools.chain.from_iterable(pm.polygons for pm in polymasks_list))
+ )
+ return cat_polymasks
+
+
+class ROIMasks:
+ """
+ Represent masks by N smaller masks defined in some ROIs. Once ROI boxes are given,
+ full-image bitmask can be obtained by "pasting" the mask on the region defined
+ by the corresponding ROI box.
+ """
+
+ def __init__(self, tensor: torch.Tensor):
+ """
+ Args:
+ tensor: (N, M, M) mask tensor that defines the mask within each ROI.
+ """
+ if tensor.dim() != 3:
+ raise ValueError("ROIMasks must take a masks of 3 dimension.")
+ self.tensor = tensor
+
+ def to(self, device: torch.device) -> "ROIMasks":
+ return ROIMasks(self.tensor.to(device))
+
+ @property
+ def device(self) -> device:
+ return self.tensor.device
+
+ def __len__(self):
+ return self.tensor.shape[0]
+
+ def __getitem__(self, item) -> "ROIMasks":
+ """
+ Returns:
+ ROIMasks: Create a new :class:`ROIMasks` by indexing.
+
+ The following usage are allowed:
+
+ 1. `new_masks = masks[2:10]`: return a slice of masks.
+ 2. `new_masks = masks[vector]`, where vector is a torch.BoolTensor
+ with `length = len(masks)`. Nonzero elements in the vector will be selected.
+
+ Note that the returned object might share storage with this object,
+ subject to Pytorch's indexing semantics.
+ """
+ t = self.tensor[item]
+ if t.dim() != 3:
+ raise ValueError(
+ f"Indexing on ROIMasks with {item} returns a tensor with shape {t.shape}!"
+ )
+ return ROIMasks(t)
+
+ @torch.jit.unused
+ def __repr__(self) -> str:
+ s = self.__class__.__name__ + "("
+ s += "num_instances={})".format(len(self.tensor))
+ return s
+
+ @torch.jit.unused
+ def to_bitmasks(self, boxes: torch.Tensor, height, width, threshold=0.5):
+ """
+ Args: see documentation of :func:`paste_masks_in_image`.
+ """
+ from detectron2.layers.mask_ops import paste_masks_in_image, _paste_masks_tensor_shape
+
+ if torch.jit.is_tracing():
+ if isinstance(height, torch.Tensor):
+ paste_func = _paste_masks_tensor_shape
+ else:
+ paste_func = paste_masks_in_image
+ else:
+ paste_func = retry_if_cuda_oom(paste_masks_in_image)
+ bitmasks = paste_func(self.tensor, boxes.tensor, (height, width), threshold=threshold)
+ return BitMasks(bitmasks)
\ No newline at end of file
diff --git a/util/detectron2/utils/memory.py b/util/detectron2/utils/memory.py
new file mode 100644
index 0000000..3ecd7d7
--- /dev/null
+++ b/util/detectron2/utils/memory.py
@@ -0,0 +1,84 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+import logging
+from contextlib import contextmanager
+from functools import wraps
+import torch
+
+__all__ = ["retry_if_cuda_oom"]
+
+
+@contextmanager
+def _ignore_torch_cuda_oom():
+ """
+ A context which ignores CUDA OOM exception from pytorch.
+ """
+ try:
+ yield
+ except RuntimeError as e:
+ # NOTE: the string may change?
+ if "CUDA out of memory. " in str(e):
+ pass
+ else:
+ raise
+
+
+def retry_if_cuda_oom(func):
+ """
+ Makes a function retry itself after encountering
+ pytorch's CUDA OOM error.
+ It will first retry after calling `torch.cuda.empty_cache()`.
+
+ If that still fails, it will then retry by trying to convert inputs to CPUs.
+ In this case, it expects the function to dispatch to CPU implementation.
+ The return values may become CPU tensors as well and it's user's
+ responsibility to convert it back to CUDA tensor if needed.
+
+ Args:
+ func: a stateless callable that takes tensor-like objects as arguments
+
+ Returns:
+ a callable which retries `func` if OOM is encountered.
+
+ Examples:
+ ::
+ output = retry_if_cuda_oom(some_torch_function)(input1, input2)
+ # output may be on CPU even if inputs are on GPU
+
+ Note:
+ 1. When converting inputs to CPU, it will only look at each argument and check
+ if it has `.device` and `.to` for conversion. Nested structures of tensors
+ are not supported.
+
+ 2. Since the function might be called more than once, it has to be
+ stateless.
+ """
+
+ def maybe_to_cpu(x):
+ try:
+ like_gpu_tensor = x.device.type == "cuda" and hasattr(x, "to")
+ except AttributeError:
+ like_gpu_tensor = False
+ if like_gpu_tensor:
+ return x.to(device="cpu")
+ else:
+ return x
+
+ @wraps(func)
+ def wrapped(*args, **kwargs):
+ with _ignore_torch_cuda_oom():
+ return func(*args, **kwargs)
+
+ # Clear cache and retry
+ torch.cuda.empty_cache()
+ with _ignore_torch_cuda_oom():
+ return func(*args, **kwargs)
+
+ # Try on CPU. This slows down the code significantly, therefore print a notice.
+ logger = logging.getLogger(__name__)
+ logger.info("Attempting to copy inputs of {} to CPU due to CUDA OOM".format(str(func)))
+ new_args = (maybe_to_cpu(x) for x in args)
+ new_kwargs = {k: maybe_to_cpu(v) for k, v in kwargs.items()}
+ return func(*new_args, **new_kwargs)
+
+ return wrapped
\ No newline at end of file
diff --git a/util/misc.py b/util/misc.py
index acae144..964f722 100644
--- a/util/misc.py
+++ b/util/misc.py
@@ -18,7 +18,8 @@
# needed due to empty tensor bug in pytorch and torchvision 0.5
import torchvision
-if float(torchvision.__version__[:3]) < 0.7:
+if float(torchvision.__version__[:3]) < 0.7 and float(torchvision.__version__[:3]) != 0.1:
+#if float(torchvision.__version__[:3]) < 0.7:
from torchvision.ops import _new_empty_tensor
from torchvision.ops.misc import _output_size
@@ -470,7 +471,7 @@ def interpolate(input, size=None, scale_factor=None, mode="nearest", align_corne
This will eventually be supported natively by PyTorch, and this
class can go away.
"""
- if float(torchvision.__version__[:3]) < 0.7:
+ if float(torchvision.__version__[:3]) < 0.7 and float(torchvision.__version__[:3]) != 0.1:
if input.numel() > 0:
return torch.nn.functional.interpolate(
input, size, scale_factor, mode, align_corners
diff --git a/ymir/docker/cuda102.dockerfile b/ymir/docker/cuda102.dockerfile
index ff40015..d7c5568 100644
--- a/ymir/docker/cuda102.dockerfile
+++ b/ymir/docker/cuda102.dockerfile
@@ -1,15 +1,38 @@
-FROM pytorch/pytorch:1.6.0-cuda10.1-cudnn7-devel
+ARG PYTORCH="1.6.0"
+ARG CUDA="10.1"
+ARG CUDNN="7"
-ENV CUDA_HOME=/usr/local/cuda
+FROM pytorch/pytorch:${PYTORCH}-cuda${CUDA}-cudnn${CUDNN}-devel
+
+ENV TORCH_CUDA_ARCH_LIST="6.0 6.1 7.0+PTX"
+ENV TORCH_NVCC_FLAGS="-Xfatbin -compress-all"
+ENV CMAKE_PREFIX_PATH="$(dirname $(which conda))/../"
+ENV PYTHONPATH=.
+ENV FORCE_CUDA="1"
+ENV MKL_SERVICE_FORCE_INTEL="1"
+ENV MKL_THREADING_LAYER="GNU"
+
+LABEL pytorch="1.6.0"
+LABEL cuda="10.1"
+LABEL cudnn="7"
+LABEL ymir="1.1.0"
+
+# To fix GPG key error when running apt-get update
+# apt-key adv --keyserver keyserver.ubuntu.com --recv-keys A4B469963BF863CC
+RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/3bf863cc.pub
+RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1804/x86_64/7fa2af80.pub
+
+RUN sed -i 's/archive.ubuntu.com/mirrors.tuna.tsinghua.edu.cn/g' /etc/apt/sources.list && \
+pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple && \
+apt update && apt install -y git vim libgl1-mesa-glx ffmpeg libsm6 libxext6 ninja-build libglib2.0-0 libsm6 libxrender-dev libxext6 && rm -rf /var/lib/apt/lists/* && \
+mkdir -p /img-man && echo "python3 ymir/start.py" > /usr/bin/start.sh
COPY . /app
WORKDIR /app
-RUN < /usr/bin/start.sh
+
+COPY ./requirements.txt /workspace
+RUN pip install -r /workspace/requirements.txt
+
+COPY . /app
+WORKDIR /app
+
+RUN cd /app/ops && bash make.sh && \
+ mv /app/ymir/img-man/*.yaml /img-man && \
+ pip install "git+https://github.com/modelai/ymir-executor-sdk.git@ymir1.3.0"
+
+CMD bash /usr/bin/start.sh
diff --git a/ymir/img-man/infer-template.yaml b/ymir/img-man/infer-template.yaml
new file mode 100644
index 0000000..30694a1
--- /dev/null
+++ b/ymir/img-man/infer-template.yaml
@@ -0,0 +1,9 @@
+# infer template for your executor app
+# after build image, it should at /img-man/infer-template.yaml
+# key: gpu_id, task_id, model_params_path, class_names should be preserved
+
+# gpu_id: '0'
+# task_id: 'default-infer-task'
+# model_params_path: []
+# class_names: []
+conf_threshold: 0.2
diff --git a/ymir/img-man/mining-template.yaml b/ymir/img-man/mining-template.yaml
new file mode 100644
index 0000000..a82c8f3
--- /dev/null
+++ b/ymir/img-man/mining-template.yaml
@@ -0,0 +1,10 @@
+# mining template for your executor app
+# after build image, it should at /img-man/mining-template.yaml
+# key: gpu_id, task_id, model_params_path, class_names should be preserved
+
+# gpu_id: '0'
+# task_id: 'default-mining-task'
+# model_params_path: []
+# class_names: []
+
+conf_threshold: 0.2
diff --git a/ymir/img-man/training-template.yaml b/ymir/img-man/training-template.yaml
new file mode 100644
index 0000000..8e3ce8b
--- /dev/null
+++ b/ymir/img-man/training-template.yaml
@@ -0,0 +1,18 @@
+# training template for your executor app
+# after build image, it should at /img-man/training-template.yaml
+# key: gpu_id, task_id, pretrained_model_params, class_names should be preserved
+
+# gpu_id: '0'
+# task_id: 'default-training-task'
+# pretrained_model_params: []
+# class_names: []
+
+batch_size_per_gpu: 2
+num_workers_per_gpu: 2
+epochs: 50
+learning_rate: 0.0001
+eval_size: 640
+backbone_name: 'swin_nano'
+weight_save_interval: 100
+args_options: ''
+export_format: 'ark:raw'
diff --git a/ymir/readme.md b/ymir/readme.md
index f8dddc1..b2c4c4e 100644
--- a/ymir/readme.md
+++ b/ymir/readme.md
@@ -4,3 +4,112 @@
```
onnxruntime 1.12.1 has requirement numpy>=1.21.0
```
+
+## ymir supports
+
+### ymir dataset supports
+add new option dataset_file=ymir and num_classes for args, modify related files.
+
+### ymir training supports
+
+- ymir/ymir_training.py, entry point for ymir training task,
+- ymir/img-man/training-template.py hyper-param template
+```
+cmd = f"""
+ python -m torch.distributed.launch \
+ --nproc_per_node {num_gpus} \
+ --nnodes 1 \
+ --use_env main.py \
+ --method vidt \
+ --backbone_name {backbone_name} \
+ --epochs {epochs} \
+ --lr {learning_rate} \
+ --min-lr 1e-7 \
+ --batch_size {batch_size} \
+ --num_workers {num_workers} \
+ --eval_size {eval_size} \
+ --aux_loss True \
+ --with_box_refine True \
+ --dataset_file ymir \
+ --num_classes {num_classes} \
+ --coco_path /in \
+ --output_dir {models_dir}
+ """
+```
+
+- save weight and map50, monitor process write tensorboard logs
+
+```
+from tensorboardX import SummaryWriter
+from ymir_exc import monitor
+from ymir_exc.util import (YmirStage, get_merged_config, get_ymir_process,
+ write_ymir_training_result)
+
+# save weight, map and monitor process for main process.
+tb_writer = SummaryWriter(args.tensorboard_dir)
+for epoch in range(args.start_epoch, args.epochs):
+ # monitor process
+ percent = get_ymir_process(stage=YmirStage.TASK, p=(epoch - args.start_epoch + 1) / (args.epochs - args.start_epoch + 1))
+ monitor.write_monitor_logger(percent=percent)
+
+ # save weight and map
+ map50 = test_stats['coco_eval_bbox'][1]
+ write_ymir_training_result(cfg, map50, files=checkpoint_paths, id=str(epoch))
+
+ # writer tensorboard logs
+ tb_writer_add_scaler(tag='train/loss', scalar_value=0.1, global_step=epoch)
+```
+
+- finetune
+
+```
+# finetune from offered weight file
+weight_files = get_weight_files(cfg, suffix=('.pth'))
+if weight_files:
+ latest_weight_file = max(weight_files, key=osp.getctime)
+
+ # auto finetune if not designed by user.
+ if args_options.find('--load_from')== -1 and args_options.find('--resume')==-1:
+ cmd = cmd + " --load_from " + latest_weight_file
+logging.info(f"Running command: {cmd}")
+subprocess.run(cmd.split(), check=True)
+
+if args.load_from:
+ if args.resume:
+ raise Exception("cannot load from and resume at the same time")
+
+ if args.load_from.startswith('https'):
+ checkpoint = torch.hub.load_state_dict_from_url(
+ args.load_from, map_location='cpu', check_hash=True)
+ else:
+ checkpoint = torch.load(args.load_from, map_location='cpu')
+ model_without_ddp.load_state_dict(checkpoint['model'], strict=False)
+ print('load a checkpoint from', args.load_from)
+```
+
+### ymir infer support
+
+- view `ymir/ymir_infer.py` for detail
+
+- init model (view class `YmirModel.init_detector()`)
+ - build backbone
+ - build detector
+ - load weight file
+ - preprocess
+ - postprocess
+
+- infer (view class `YmirModel.infer()`)
+ - read image with `PIL.Image`
+ - obtain `conf_threshold` and `class_names` from `get_merged_config()`
+ - convert result to ymir format `rw.Anotation` and save infer result with `rw.write_infer_result()`
+
+### ymir mining support
+
+- view `ymir/ymir_mining.py` for detail
+
+- save mining result with `rw.write_mining_result()`
+
+## thanks to
+
+- [detr](https://github.com/facebookresearch/detr)
+- [vidt](https://github.com/naver-ai/vidt)
diff --git a/ymir/start.py b/ymir/start.py
new file mode 100644
index 0000000..c142f9b
--- /dev/null
+++ b/ymir/start.py
@@ -0,0 +1,57 @@
+import logging
+import subprocess
+import sys
+
+from ymir_exc import monitor
+from easydict import EasyDict as edict
+from ymir_exc.util import YmirStage, get_merged_config, write_ymir_monitor_process
+
+
+def start() -> int:
+ cfg = get_merged_config()
+
+ logging.info(f'merged config: {cfg}')
+
+ if cfg.ymir.run_training:
+ _run_training()
+ elif cfg.ymir.run_mining or cfg.ymir.run_infer:
+ # for multiple tasks, run mining first, infer later.
+
+ if cfg.ymir.run_mining:
+ _run_mining(cfg)
+ if cfg.ymir.run_infer:
+ _run_infer(cfg)
+ else:
+ logging.warning('no task running')
+
+ return 0
+
+
+def _run_training() -> None:
+ command = 'python3 ymir/ymir_training.py'
+ logging.info(f'training: {command}')
+ subprocess.run(command.split(), check=True)
+ monitor.write_monitor_logger(percent=1.0)
+
+
+def _run_mining(cfg: edict) -> None:
+ command = 'python3 ymir/ymir_mining.py'
+ logging.info(f'mining: {command}')
+ subprocess.run(command.split(), check=True)
+ write_ymir_monitor_process(cfg, task='mining', naive_stage_percent=1.0, stage=YmirStage.POSTPROCESS)
+
+
+def _run_infer(cfg: edict) -> None:
+ command = 'python3 ymir/ymir_infer.py'
+ logging.info(f'infer: {command}')
+ subprocess.run(command.split(), check=True)
+ write_ymir_monitor_process(cfg, task='infer', naive_stage_percent=1.0, stage=YmirStage.POSTPROCESS)
+
+
+if __name__ == '__main__':
+ logging.basicConfig(stream=sys.stdout,
+ format='%(levelname)-8s: [%(asctime)s] %(message)s',
+ datefmt='%Y%m%d-%H:%M:%S',
+ level=logging.INFO)
+
+ sys.exit(start())
diff --git a/ymir/ymir_infer.py b/ymir/ymir_infer.py
new file mode 100644
index 0000000..4e2a833
--- /dev/null
+++ b/ymir/ymir_infer.py
@@ -0,0 +1,219 @@
+"""
+ymir infer task entry point
+"""
+import os.path as osp
+import sys
+from typing import List
+
+import datasets.transforms as T
+import torch
+from easydict import EasyDict as edict
+from methods.coat_w_ram import coat_lite_mini, coat_lite_small, coat_lite_tiny
+from methods.swin_w_ram import swin_base_win7, swin_large_win7, swin_nano, swin_small, swin_tiny
+from PIL import Image
+from ymir_exc import dataset_reader as dr
+from ymir_exc import env
+from ymir_exc import result_writer as rw
+from ymir_exc.util import YmirStage, get_merged_config, get_weight_files, write_ymir_monitor_process
+
+
+def build_model(args):
+ if args.backbone_name == 'swin_nano':
+ backbone, hidden_dim = swin_nano(pretrained=None)
+ elif args.backbone_name == 'swin_tiny':
+ backbone, hidden_dim = swin_tiny(pretrained=None)
+ elif args.backbone_name == 'swin_small':
+ backbone, hidden_dim = swin_small(pretrained=None)
+ elif args.backbone_name == 'swin_base_win7_22k':
+ backbone, hidden_dim = swin_base_win7(pretrained=None)
+ elif args.backbone_name == 'swin_large_win7_22k':
+ backbone, hidden_dim = swin_large_win7(pretrained=None)
+ elif args.backbone_name == 'coat_lite_tiny':
+ backbone, hidden_dim = coat_lite_tiny(pretrained=None)
+ elif args.backbone_name == 'coat_lite_mini':
+ backbone, hidden_dim = coat_lite_mini(pretrained=None)
+ elif args.backbone_name == 'coat_lite_small':
+ backbone, hidden_dim = coat_lite_small(pretrained=None)
+ else:
+ raise ValueError(f'backbone {args.backbone_name} not supported')
+
+ if args.method == 'vidt':
+ return build_vidt_model(args, backbone)
+ elif args.method == 'vidt_wo_neck':
+ return build_vidt_wo_neck_model(args, backbone)
+ else:
+ available_methods = ['vidt_wo_neck', 'vidt']
+ raise ValueError(f'method {args.method} is not in {available_methods}')
+
+
+def build_vidt_model(args, backbone):
+ from methods.vidt.deformable_transformer import build_deforamble_transformer
+ from methods.vidt.detector import Detector
+ from methods.vidt.fpn_fusion import FPNFusionModule
+
+ backbone.finetune_det(method=args.method,
+ det_token_num=args.det_token_num,
+ pos_dim=args.reduced_dim,
+ cross_indices=args.cross_indices)
+
+ cross_scale_fusion = None
+ if args.cross_scale_fusion:
+ cross_scale_fusion = FPNFusionModule(backbone.num_channels,
+ fuse_dim=args.reduced_dim)
+
+ deform_transformers = build_deforamble_transformer(args)
+
+ model = Detector(
+ backbone,
+ deform_transformers,
+ num_classes=args.num_classes,
+ num_queries=args.det_token_num,
+ # two essential techniques used in ViDT
+ aux_loss=args.aux_loss,
+ with_box_refine=args.with_box_refine,
+ # three additional techniques (optionally)
+ cross_scale_fusion=cross_scale_fusion,
+ iou_aware=args.iou_aware,
+ token_label=args.token_label,
+ # distil
+ distil=False if args.distil_model is None else True,
+ )
+
+ return model
+
+
+def build_vidt_wo_neck_model(args, backbone):
+ from methods.vidt_wo_neck.detector import Detector
+
+ backbone.finetune_det(method=args.method,
+ det_token_num=args.det_token_num,
+ pos_dim=args.pos_dim,
+ cross_indices=args.cross_indices)
+
+ model = Detector(
+ backbone,
+ reduced_dim=args.reduced_dim,
+ num_classes=args.num_classes,
+ )
+
+ return model
+
+
+def get_postprocessor(args):
+ if args.method == 'vidt':
+ from methods.vidt.postprocessor import PostProcess as postprocess_vidt
+ return postprocess_vidt(args.dataset_file)
+ elif args.method == 'vidt_wo_neck':
+ from methods.vidt_wo_neck.postprocessor import PostProcess as postprocess_vidt_wo_neck
+ return postprocess_vidt_wo_neck()
+ else:
+ available_methods = ['vidt_wo_neck', 'vidt']
+ raise ValueError(f'method {args.method} is not in {available_methods}')
+
+
+class YmirModel(object):
+
+ def __init__(self, cfg: edict):
+ """model for ymir infer task
+
+ build model and load weight, inference
+
+ Args:
+ cfg: the ymir merged config, view get_merged_config()
+
+ Raises:
+ Exception: No weight files specified
+ """
+ self.cfg = cfg
+ self.init_detector()
+
+ self.conf_threshold = self.cfg.param.conf_threshold
+ self.class_names = self.cfg.param.class_names
+
+ def init_detector(self):
+ """
+ build model and load weight
+ """
+ weight_files = get_weight_files(self.cfg, suffix=('.pth'))
+ weight_files = [
+ f for f in weight_files if osp.basename(f).startswith('checkpoint')
+ ]
+
+ if len(weight_files) == 0:
+ raise Exception('No weight files specified')
+
+ latest_weight_file = max(weight_files, key=osp.getctime)
+
+ # save "model", "optimizer", "lr_scheduler", "epoch" and "args" in weight_file
+ load_data = torch.load(latest_weight_file, map_location='cpu')
+
+ args = load_data['args']
+ self.model = build_model(args)
+ self.model.load_state_dict(load_data['model'], strict=False)
+ self.device = torch.device(args.device)
+ self.model.to(self.device)
+ self.model.eval()
+
+ self.preprocess = T.Compose([
+ T.RandomResize([512], max_size=800),
+ T.ToTensor(),
+ T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+ ])
+ self.postprocess = get_postprocessor(args)
+
+ def infer(self, img: Image) -> List[rw.Annotation]:
+ sample, _ = self.preprocess(img, None)
+ img_w, img_h = img.size
+ orig_target_sizes = torch.tensor([[img_h, img_w]], device=self.device)
+ outputs = self.model([sample.to(self.device)])
+ # results = [{'scores': s, 'labels': l, 'boxes': b} for s, l, b in zip(scores, labels, boxes)]
+ results = self.postprocess(outputs, orig_target_sizes)
+
+ anns = []
+ # for batch
+ for r in results:
+ # for bbox
+ for conf, cls, (xmin, ymin, xmax,
+ ymax) in zip(r['scores'], r['labels'], r['boxes']):
+ if conf < self.conf_threshold:
+ continue
+
+ ann = rw.Annotation(class_name=self.class_names[int(cls)],
+ score=conf,
+ box=rw.Box(x=int(xmin),
+ y=int(ymin),
+ w=int(xmax - xmin),
+ h=int(ymax - ymin)))
+
+ anns.append(ann)
+ return anns
+
+
+def main() -> int:
+ cfg = get_merged_config()
+ model = YmirModel(cfg)
+ write_ymir_monitor_process(cfg, task='infer', naive_stage_percent=1.0, stage=YmirStage.PREPROCESS)
+
+ N = dr.items_count(env.DatasetType.CANDIDATE)
+ infer_result = {}
+
+ idx = -1
+
+ monitor_gap = max(1, N // 1000)
+ for asset_path, _ in dr.item_paths(dataset_type=env.DatasetType.CANDIDATE):
+ # img = cv2.imread(asset_path)
+ img = Image.open(asset_path).convert('RGB')
+ result = model.infer(img)
+ infer_result[asset_path] = result
+ idx += 1
+
+ if idx % monitor_gap == 0:
+ write_ymir_monitor_process(cfg, task='infer', naive_stage_percent=idx / N, stage=YmirStage.TASK)
+
+ rw.write_infer_result(infer_result=infer_result)
+ write_ymir_monitor_process(cfg, task='infer', naive_stage_percent=1.0, stage=YmirStage.POSTPROCESS)
+ return 0
+
+
+if __name__ == '__main__':
+ sys.exit(main())
diff --git a/ymir/ymir_mining.py b/ymir/ymir_mining.py
new file mode 100644
index 0000000..9f68276
--- /dev/null
+++ b/ymir/ymir_mining.py
@@ -0,0 +1,374 @@
+"""
+data augmentations for CALD method, including horizontal_flip, rotate(5'), cutout
+official code: https://github.com/we1pingyu/CALD/blob/master/cald/cald_helper.py
+"""
+import random
+import sys
+from typing import Any, Callable, Dict, List, Tuple
+
+import cv2
+import numpy as np
+from easydict import EasyDict as edict
+from nptyping import NDArray, Shape, UInt8
+from PIL import Image
+from scipy.stats import entropy
+from tqdm import tqdm
+from ymir.ymir_infer import YmirModel
+from ymir_exc import dataset_reader as dr
+from ymir_exc import env
+from ymir_exc import result_writer as rw
+from ymir_exc.util import YmirStage, get_merged_config, write_ymir_monitor_process
+
+BBOX = NDArray[Shape['*,4'], Any]
+CV_IMAGE = NDArray[Shape['*,*,3'], UInt8]
+
+
+def intersect(boxes1: BBOX, boxes2: BBOX) -> NDArray:
+ '''
+ Find intersection of every box combination between two sets of box
+ boxes1: bounding boxes 1, a tensor of dimensions (n1, 4)
+ boxes2: bounding boxes 2, a tensor of dimensions (n2, 4)
+
+ Out: Intersection each of boxes1 with respect to each of boxes2,
+ a tensor of dimensions (n1, n2)
+ '''
+ n1 = boxes1.shape[0]
+ n2 = boxes2.shape[0]
+ max_xy = np.minimum(
+ np.expand_dims(boxes1[:, 2:], axis=1).repeat(n2, axis=1),
+ np.expand_dims(boxes2[:, 2:], axis=0).repeat(n1, axis=0))
+
+ min_xy = np.maximum(
+ np.expand_dims(boxes1[:, :2], axis=1).repeat(n2, axis=1),
+ np.expand_dims(boxes2[:, :2], axis=0).repeat(n1, axis=0))
+ inter = np.clip(max_xy - min_xy, a_min=0, a_max=None) # (n1, n2, 2)
+ return inter[:, :, 0] * inter[:, :, 1] # (n1, n2)
+
+
+def horizontal_flip(image: CV_IMAGE, bbox: BBOX) \
+ -> Tuple[CV_IMAGE, BBOX]:
+ """
+ image: opencv image, [height,width,channels]
+ bbox: numpy.ndarray, [N,4] --> [x1,y1,x2,y2]
+ """
+ image = image.copy()
+
+ width = image.shape[1]
+ # Flip image horizontally
+ image = image[:, ::-1, :]
+ if len(bbox) > 0:
+ bbox = bbox.copy()
+ # Flip bbox horizontally
+ bbox[:, [0, 2]] = width - bbox[:, [2, 0]]
+ return image, bbox
+
+
+def cutout(image: CV_IMAGE,
+ bbox: BBOX,
+ cut_num: int = 2,
+ fill_val: int = 0,
+ bbox_remove_thres: float = 0.4,
+ bbox_min_thres: float = 0.1) -> Tuple[CV_IMAGE, BBOX]:
+ '''
+ Cutout augmentation
+ image: A PIL image
+ boxes: bounding boxes, a tensor of dimensions (#objects, 4)
+ labels: labels of object, a tensor of dimensions (#objects)
+ fill_val: Value filled in cut out
+ bbox_remove_thres: Theshold to remove bbox cut by cutout
+
+ Out: new image, new_boxes, new_labels
+ '''
+ image = image.copy()
+ bbox = bbox.copy()
+
+ if len(bbox) == 0:
+ return image, bbox
+
+ original_h, original_w, original_channel = image.shape
+ count = 0
+ for _ in range(50):
+ # Random cutout size: [0.15, 0.5] of original dimension
+ cutout_size_h = random.uniform(0.05 * original_h, 0.2 * original_h)
+ cutout_size_w = random.uniform(0.05 * original_w, 0.2 * original_w)
+
+ # Random position for cutout
+ left = random.uniform(0, original_w - cutout_size_w)
+ right = left + cutout_size_w
+ top = random.uniform(0, original_h - cutout_size_h)
+ bottom = top + cutout_size_h
+ cutout = np.array([[float(left), float(top), float(right), float(bottom)]])
+
+ # Calculate intersect between cutout and bounding boxes
+ overlap_size = intersect(cutout, bbox)
+ area_boxes = (bbox[:, 2] - bbox[:, 0]) * (bbox[:, 3] - bbox[:, 1])
+ ratio = overlap_size / (area_boxes + 1e-14)
+ # If all boxes have Iou greater than bbox_remove_thres, try again
+ if ratio.max() > bbox_remove_thres or ratio.max() < bbox_min_thres:
+ continue
+
+ image[int(top):int(bottom), int(left):int(right), :] = fill_val
+ count += 1
+ if count >= cut_num:
+ break
+ return image, bbox
+
+
+def rotate(image: CV_IMAGE, bbox: BBOX, rot: float = 5) -> Tuple[CV_IMAGE, BBOX]:
+ image = image.copy()
+ bbox = bbox.copy()
+ h, w, c = image.shape
+ center = np.array([w / 2.0, h / 2.0])
+ s = max(h, w) * 1.0
+ trans = get_affine_transform(center, s, rot, [w, h])
+ if len(bbox) > 0:
+ for i in range(bbox.shape[0]):
+ x1, y1 = affine_transform(bbox[i, :2], trans)
+ x2, y2 = affine_transform(bbox[i, 2:], trans)
+ x3, y3 = affine_transform(bbox[i, [2, 1]], trans)
+ x4, y4 = affine_transform(bbox[i, [0, 3]], trans)
+ bbox[i, :2] = [min(x1, x2, x3, x4), min(y1, y2, y3, y4)]
+ bbox[i, 2:] = [max(x1, x2, x3, x4), max(y1, y2, y3, y4)]
+ image = cv2.warpAffine(image, trans, (w, h), flags=cv2.INTER_LINEAR)
+ return image, bbox
+
+
+def get_3rd_point(a: NDArray, b: NDArray) -> NDArray:
+ direct = a - b
+ return b + np.array([-direct[1], direct[0]], dtype=np.float32)
+
+
+def get_dir(src_point: NDArray, rot_rad: float) -> List:
+ sn, cs = np.sin(rot_rad), np.cos(rot_rad)
+
+ src_result = [0, 0]
+ src_result[0] = src_point[0] * cs - src_point[1] * sn
+ src_result[1] = src_point[0] * sn + src_point[1] * cs
+
+ return src_result
+
+
+def transform_preds(coords: NDArray, center: NDArray, scale: Any, rot: float, output_size: List) -> NDArray:
+ trans = get_affine_transform(center, scale, rot, output_size, inv=True)
+ target_coords = affine_transform(coords, trans)
+ return target_coords
+
+
+def get_affine_transform(center: NDArray,
+ scale: Any,
+ rot: float,
+ output_size: List,
+ shift: NDArray = np.array([0, 0], dtype=np.float32),
+ inv: bool = False) -> NDArray:
+ if not isinstance(scale, np.ndarray) and not isinstance(scale, list):
+ scale = np.array([scale, scale], dtype=np.float32)
+
+ scale_tmp = scale
+ src_w = scale_tmp[0]
+ dst_w = output_size[0]
+ dst_h = output_size[1]
+
+ rot_rad = np.pi * rot / 180
+ src_dir = get_dir(np.array([0, src_w * -0.5], np.float32), rot_rad)
+ dst_dir = np.array([0, dst_w * -0.5], np.float32)
+
+ src = np.zeros((3, 2), dtype=np.float32)
+ dst = np.zeros((3, 2), dtype=np.float32)
+ src[0, :] = center + scale_tmp * shift
+ src[1, :] = center + src_dir + scale_tmp * shift
+ dst[0, :] = [dst_w * 0.5, dst_h * 0.5]
+ dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5], np.float32) + dst_dir
+
+ src[2:, :] = get_3rd_point(src[0, :], src[1, :])
+ dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :])
+
+ if inv:
+ trans = cv2.getAffineTransform(np.float32(dst), np.float32(src))
+ else:
+ trans = cv2.getAffineTransform(np.float32(src), np.float32(dst))
+
+ return trans
+
+
+def affine_transform(pt: NDArray, t: NDArray) -> NDArray:
+ new_pt = np.array([pt[0], pt[1], 1.], dtype=np.float32).T
+ new_pt = np.dot(t, new_pt)
+ return new_pt[:2]
+
+
+def resize(img: CV_IMAGE, boxes: BBOX, ratio: float = 0.8) -> Tuple[CV_IMAGE, BBOX]:
+ """
+ ratio: <= 1.0
+ """
+ assert ratio <= 1.0, f'resize ratio {ratio} must <= 1.0'
+
+ h, w, _ = img.shape
+ ow = int(w * ratio)
+ oh = int(h * ratio)
+ resize_img = cv2.resize(img, (ow, oh))
+ new_img = np.zeros_like(img)
+ new_img[:oh, :ow] = resize_img
+
+ if len(boxes) == 0:
+ return new_img, boxes
+ else:
+ return new_img, boxes * ratio
+
+
+def get_ious(boxes1: BBOX, boxes2: BBOX) -> NDArray:
+ """
+ args:
+ boxes1: np.array, (N, 4), xyxy
+ boxes2: np.array, (M, 4), xyxy
+ return:
+ iou: np.array, (N, M)
+ """
+ area1 = (boxes1[:, 2] - boxes1[:, 0]) * (boxes1[:, 3] - boxes1[:, 1])
+ area2 = (boxes2[:, 2] - boxes2[:, 0]) * (boxes2[:, 3] - boxes2[:, 1])
+ iner_area = intersect(boxes1, boxes2)
+ area1 = area1.reshape(-1, 1).repeat(area2.shape[0], axis=1)
+ area2 = area2.reshape(1, -1).repeat(area1.shape[0], axis=0)
+ iou = iner_area / (area1 + area2 - iner_area + 1e-14)
+ return iou
+
+
+def split_result(result: NDArray) -> Tuple[BBOX, NDArray, NDArray]:
+ if len(result) > 0:
+ bboxes = result[:, :4].astype(np.int32)
+ conf = result[:, 4]
+ class_id = result[:, 5]
+ else:
+ bboxes = np.zeros(shape=(0, 4), dtype=np.int32)
+ conf = np.zeros(shape=(0, 1), dtype=np.float32)
+ class_id = np.zeros(shape=(0, 1), dtype=np.int32)
+
+ return bboxes, conf, class_id
+
+
+class YmirMining(YmirModel):
+
+ def __init__(self, cfg: edict):
+ super().__init__(cfg)
+ # for multiple tasks, mining first, then infer
+ if cfg.ymir.run_mining and cfg.ymir.run_infer:
+ mining_task_idx = 0
+ task_num = 2
+ else:
+ mining_task_idx = 0
+ task_num = 1
+ self.task_idx = mining_task_idx
+ self.task_num = task_num
+
+ def mining(self):
+ N = dr.items_count(env.DatasetType.CANDIDATE)
+ monitor_gap = max(1, N // 1000)
+ idx = -1
+ beta = 1.3
+ mining_result = []
+ for asset_path, _ in tqdm(dr.item_paths(dataset_type=env.DatasetType.CANDIDATE)):
+ # img = cv2.imread(asset_path)
+ img = cv2.imread(asset_path)
+ # xyxy,conf,cls
+ result = self.predict(img)
+ bboxes, conf, _ = split_result(result)
+ if len(result) == 0:
+ # no result for the image without augmentation
+ mining_result.append((asset_path, -beta))
+ continue
+
+ consistency = 0.0
+ aug_bboxes_dict, aug_results_dict = self.aug_predict(img, bboxes)
+ for key in aug_results_dict:
+ # no result for the image with augmentation f'{key}'
+ if len(aug_results_dict[key]) == 0:
+ consistency += beta
+ continue
+
+ bboxes_key, conf_key, _ = split_result(aug_results_dict[key])
+ cls_scores_aug = 1 - conf_key
+ cls_scores = 1 - conf
+
+ consistency_per_aug = 2.0
+ ious = get_ious(bboxes_key, aug_bboxes_dict[key])
+ aug_idxs = np.argmax(ious, axis=0)
+ for origin_idx, aug_idx in enumerate(aug_idxs):
+ max_iou = ious[aug_idx, origin_idx]
+ if max_iou == 0:
+ consistency_per_aug = min(consistency_per_aug, beta)
+ p = cls_scores_aug[aug_idx]
+ q = cls_scores[origin_idx]
+ m = (p + q) / 2.
+ js = 0.5 * entropy([p, 1 - p], [m, 1 - m]) + 0.5 * entropy([q, 1 - q], [m, 1 - m])
+ if js < 0:
+ js = 0
+ consistency_box = max_iou
+ consistency_cls = 0.5 * \
+ (conf[origin_idx] + conf_key[aug_idx]) * (1 - js)
+ consistency_per_inst = abs(consistency_box + consistency_cls - beta)
+ consistency_per_aug = min(consistency_per_aug, consistency_per_inst.item())
+
+ consistency += consistency_per_aug
+
+ consistency /= len(aug_results_dict)
+
+ mining_result.append((asset_path, consistency))
+ idx += 1
+
+ if idx % monitor_gap == 0:
+ write_ymir_monitor_process(self.cfg, task='mining', naive_stage_percent=idx / N, stage=YmirStage.TASK)
+
+ return mining_result
+
+ def predict(self, img: CV_IMAGE) -> NDArray:
+ """
+ predict single image and return bbox information
+ img: opencv BGR, uint8 format
+ """
+ pillow_img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
+ anns = self.infer(pillow_img)
+
+ xyxy_conf_idx_list = []
+ for idx, ann in enumerate(anns):
+ bbox = ann.box
+ score = ann.score
+ x1, y1, x2, y2 = bbox.x, bbox.y, bbox.x + bbox.w, bbox.y + bbox.h
+ idx = self.class_names.index(ann.class_name)
+ xyxy_conf_idx_list.append([x1, y1, x2, y2, score, idx])
+
+ if len(xyxy_conf_idx_list) == 0:
+ return np.zeros(shape=(0, 6), dtype=np.float32)
+ else:
+ return np.array(xyxy_conf_idx_list, dtype=np.float32)
+
+ def aug_predict(self, image: CV_IMAGE, bboxes: BBOX) -> Tuple[Dict[str, BBOX], Dict[str, NDArray]]:
+ """
+ for different augmentation methods: flip, cutout, rotate and resize
+ augment the image and bbox and use model to predict them.
+
+ return the predict result and augment bbox.
+ """
+ aug_dict: Dict[str, Callable] = dict(flip=horizontal_flip, cutout=cutout, rotate=rotate, resize=resize)
+
+ aug_bboxes = dict()
+ aug_results = dict()
+ for key in aug_dict:
+ aug_img, aug_bbox = aug_dict[key](image, bboxes)
+
+ aug_result = self.predict(aug_img)
+ aug_bboxes[key] = aug_bbox
+ aug_results[key] = aug_result
+
+ return aug_bboxes, aug_results
+
+
+def main():
+ cfg = get_merged_config()
+ miner = YmirMining(cfg)
+ mining_result = miner.mining()
+ rw.write_mining_result(mining_result=mining_result)
+ write_ymir_monitor_process(cfg, task='mining', naive_stage_percent=1, stage=YmirStage.POSTPROCESS)
+ return 0
+
+
+if __name__ == "__main__":
+ sys.exit(main())
diff --git a/ymir/ymir_training.py b/ymir/ymir_training.py
new file mode 100644
index 0000000..d428b91
--- /dev/null
+++ b/ymir/ymir_training.py
@@ -0,0 +1,72 @@
+import logging
+import os
+import os.path as osp
+import subprocess
+import sys
+
+from easydict import EasyDict as edict
+from ymir_exc.util import get_merged_config, get_weight_files, write_ymir_training_result
+
+
+def main(cfg: edict) -> int:
+ num_gpus = len(cfg.param.gpu_id.split(','))
+ models_dir = cfg.ymir.output.models_dir
+ tensorboard_dir = cfg.ymir.output.tensorboard_dir
+ args_options = cfg.param.get('args_options', '')
+ num_classes = len(cfg.param.class_names)
+ # gpu_count = len(cfg.param.get('gpu_id', '0').split(','))
+ batch_size = int(cfg.param.batch_size_per_gpu)
+ eval_size = int(cfg.param.eval_size)
+ epochs = int(cfg.param.epochs)
+ num_workers = int(cfg.param.num_workers_per_gpu)
+ learning_rate = float(cfg.param.learning_rate)
+ backbone_name = str(cfg.param.backbone_name)
+ weight_save_interval = int(cfg.param.weight_save_interval)
+
+ cmd = f"""
+ python -m torch.distributed.launch \
+ --nproc_per_node {num_gpus} \
+ --nnodes 1 \
+ --use_env main.py \
+ --backbone_name {backbone_name} \
+ --epochs {epochs} \
+ --lr {learning_rate} \
+ --batch_size {batch_size} \
+ --num_workers {num_workers} \
+ --eval_size {eval_size} \
+ --dataset_file ymir \
+ --num_classes {num_classes} \
+ --coco_path /in \
+ --output_dir {models_dir} \
+ --save_interval {weight_save_interval} \
+ --tensorboard_dir {tensorboard_dir}
+ """
+
+ if args_options:
+ cmd = cmd + " " + args_options
+
+ # finetune from offered weight file
+ weight_files = get_weight_files(cfg, suffix=('.pth'))
+ weight_files = [
+ f for f in weight_files if osp.basename(f).startswith('checkpoint')
+ ]
+ if weight_files:
+ latest_weight_file = max(weight_files, key=osp.getctime)
+
+ # auto finetune if not specified by user.
+ if args_options.find('--load_from') == -1 and args_options.find(
+ '--resume') == -1:
+ cmd = cmd + " --load_from " + latest_weight_file
+
+ logging.info(f"Running command: {cmd}")
+ subprocess.run(cmd.split(), check=True)
+
+ write_ymir_training_result(cfg, map50=0, files=[], id='last')
+
+ return 0
+
+
+if __name__ == '__main__':
+ cfg = get_merged_config()
+ os.environ.setdefault('PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION', 'python')
+ sys.exit(main(cfg))
+